chore: integrate develop into hare/develop candidate
This commit is contained in:
@@ -22,7 +22,10 @@ use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use agen::Item;
|
||||
use agen::interceptor::{Interceptor, PreRequestAction, PreToolAction, ToolCallInfo};
|
||||
use agen::interceptor::{
|
||||
Interceptor, InterceptorResult, PreLlmRequestContext, PreRequestAction, PreToolAction,
|
||||
ToolCallInfo,
|
||||
};
|
||||
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput, ToolResult};
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
@@ -397,15 +400,19 @@ impl CompactWorkerInterceptor {
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Interceptor for CompactWorkerInterceptor {
|
||||
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
|
||||
impl<A: Send + Sync> Interceptor<A> for CompactWorkerInterceptor {
|
||||
async fn pre_llm_request(
|
||||
&self,
|
||||
context: PreLlmRequestContext<'_, A>,
|
||||
) -> InterceptorResult<PreRequestAction> {
|
||||
let context = context.items;
|
||||
let records = self.usage_tracker.records();
|
||||
let estimate = agen::token_counter::total_tokens(context, &records);
|
||||
if estimate.tokens > self.max_input_tokens {
|
||||
return PreRequestAction::Cancel(format!(
|
||||
return Ok(PreRequestAction::Cancel(format!(
|
||||
"compact worker input occupancy exceeded {} tokens",
|
||||
self.max_input_tokens
|
||||
));
|
||||
)));
|
||||
}
|
||||
|
||||
let remaining = self.max_input_tokens.saturating_sub(estimate.tokens);
|
||||
@@ -413,25 +420,28 @@ impl Interceptor for CompactWorkerInterceptor {
|
||||
.store(remaining, Ordering::Release);
|
||||
if let Some(item) = self.maybe_emit_warning(remaining) {
|
||||
self.usage_tracker.note_request(context.len() + 1);
|
||||
return PreRequestAction::ContinueWith(vec![item]);
|
||||
return Ok(PreRequestAction::ContinueWith(vec![item]));
|
||||
}
|
||||
|
||||
self.usage_tracker.note_request(context.len());
|
||||
PreRequestAction::Continue
|
||||
Ok(PreRequestAction::Continue)
|
||||
}
|
||||
|
||||
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction {
|
||||
async fn pre_tool_call(
|
||||
&self,
|
||||
info: &mut ToolCallInfo<'_, A>,
|
||||
) -> InterceptorResult<PreToolAction> {
|
||||
if self.final_reserve_tokens == 0 || info.call.name == "write_summary" {
|
||||
return PreToolAction::Continue;
|
||||
return Ok(PreToolAction::Continue);
|
||||
}
|
||||
let remaining = self.last_remaining_tokens.load(Ordering::Acquire);
|
||||
if remaining > self.final_reserve_tokens {
|
||||
return PreToolAction::Continue;
|
||||
return Ok(PreToolAction::Continue);
|
||||
}
|
||||
PreToolAction::SyntheticResult(ToolResult::error(
|
||||
Ok(PreToolAction::SyntheticResult(ToolResult::error(
|
||||
info.call.id.clone(),
|
||||
"compact worker final reserve reached; do not perform more exploratory tool reads. Call `write_summary` now.",
|
||||
))
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -467,13 +477,27 @@ mod tests {
|
||||
let mut context = vec![Item::user_message("hello")];
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
interceptor
|
||||
.pre_llm_request(PreLlmRequestContext::<()> {
|
||||
invocation: Default::default(),
|
||||
items: &mut context,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap(),
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
tracker.record_usage(&make_usage(100));
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
interceptor
|
||||
.pre_llm_request(PreLlmRequestContext::<()> {
|
||||
invocation: Default::default(),
|
||||
items: &mut context,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap(),
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
tracker.record_usage(&make_usage(100));
|
||||
@@ -481,7 +505,14 @@ mod tests {
|
||||
// Two 100-token requests would exceed a cumulative 150-token cap, but
|
||||
// current occupancy is still the latest 100-token measurement.
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
interceptor
|
||||
.pre_llm_request(PreLlmRequestContext::<()> {
|
||||
invocation: Default::default(),
|
||||
items: &mut context,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap(),
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
}
|
||||
@@ -503,13 +534,27 @@ mod tests {
|
||||
let mut context = vec![Item::user_message("hello")];
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
interceptor
|
||||
.pre_llm_request(PreLlmRequestContext::<()> {
|
||||
invocation: Default::default(),
|
||||
items: &mut context,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap(),
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
tracker.record_usage(&make_usage(100));
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
interceptor
|
||||
.pre_llm_request(PreLlmRequestContext::<()> {
|
||||
invocation: Default::default(),
|
||||
items: &mut context,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap(),
|
||||
PreRequestAction::ContinueWith(items)
|
||||
if items.len() == 1 && items[0].as_text().unwrap_or_default().contains("write_summary")
|
||||
));
|
||||
@@ -523,13 +568,27 @@ mod tests {
|
||||
let mut context = vec![Item::user_message("hello")];
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
interceptor
|
||||
.pre_llm_request(PreLlmRequestContext::<()> {
|
||||
invocation: Default::default(),
|
||||
items: &mut context,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap(),
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
tracker.record_usage(&make_usage(100));
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
interceptor
|
||||
.pre_llm_request(PreLlmRequestContext::<()> {
|
||||
invocation: Default::default(),
|
||||
items: &mut context,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap(),
|
||||
PreRequestAction::Cancel(message) if message.contains("occupancy")
|
||||
));
|
||||
}
|
||||
|
||||
@@ -197,8 +197,8 @@ async fn finish_controller_run<C, St>(
|
||||
{
|
||||
// history / user_segments are no longer mirrored on WorkerSharedState —
|
||||
// clients reconstruct them from `Event::Snapshot` + live
|
||||
// `Event::Entry` deliveries driven by the session-log sink. We
|
||||
// flip the status and kick post-run memory jobs here.
|
||||
// `Event::Entry` deliveries driven by the session-log sink. The
|
||||
// lifecycle hook/task registry observes the terminal commit separately.
|
||||
//
|
||||
// In-flight blocks are run-local streaming state, not durable transcript.
|
||||
// Any block not cleared by a committed AssistantItem must be discarded at
|
||||
@@ -206,7 +206,6 @@ async fn finish_controller_run<C, St>(
|
||||
// partial text/tool arguments after newer entries.
|
||||
worker.clear_in_flight_events();
|
||||
set_controller_status(shared_state, runtime_dir, working_event_tx, new_status).await;
|
||||
worker.spawn_post_run_memory_jobs();
|
||||
}
|
||||
|
||||
/// Pending turn launch staged by an event handler for the next outer-loop
|
||||
@@ -938,7 +937,6 @@ where
|
||||
let local_filesystem = worker.local_working_directory().cloned();
|
||||
let local_workspace_root = local_filesystem.as_ref().map(|local| local.root.clone());
|
||||
let task_feature = worker.task_feature();
|
||||
let memory_config = worker.manifest().memory.clone();
|
||||
let web_config = worker.manifest().web.clone();
|
||||
let mcp_config = worker.manifest().mcp.clone();
|
||||
let spawner_name = worker.manifest().worker.name.clone();
|
||||
@@ -995,6 +993,41 @@ where
|
||||
let worker_enabled = feature_config.worker.enabled;
|
||||
let sub_worker_enabled = feature_config.sub_worker.enabled;
|
||||
let mut feature_registry = FeatureRegistryBuilder::new();
|
||||
let memory_install_plan = crate::feature::builtin::memory::MemoryFeatureInstallPlan::prepare(
|
||||
worker.manifest(),
|
||||
worker.workspace_client_handle(),
|
||||
worker.prompts().load_full(),
|
||||
)
|
||||
.await?;
|
||||
let memory_prompt_contribution = memory_install_plan.as_ref().map(|plan| {
|
||||
(
|
||||
plan.resident_summary.clone(),
|
||||
plan.system_prompt_override.clone(),
|
||||
)
|
||||
});
|
||||
let memory_lifecycle_config = memory_install_plan
|
||||
.as_ref()
|
||||
.map(|plan| plan.resolved_config.clone());
|
||||
if let Some(plan) = memory_install_plan {
|
||||
feature_registry.add_module(plan.module);
|
||||
}
|
||||
if let Some(memory_config) = memory_lifecycle_config
|
||||
&& let Some(memory_lifecycle) =
|
||||
crate::feature::builtin::memory_lifecycle::MemoryLifecycleFeature::from_resolved_config(
|
||||
worker.manifest_lifecycle_features_enabled(),
|
||||
memory_config,
|
||||
worker.committed_session_capture_handle(),
|
||||
worker.session_extension_handle(),
|
||||
worker.workspace_client_handle(),
|
||||
spawner_manifest.clone(),
|
||||
worker.llm_client_handle(),
|
||||
prompts.clone(),
|
||||
spawner_workspace_context.clone(),
|
||||
worker.working_event_sender(),
|
||||
)?
|
||||
{
|
||||
feature_registry.add_module(memory_lifecycle);
|
||||
}
|
||||
if sub_worker_enabled && !worker_enabled {
|
||||
feature_registry.add_module(
|
||||
crate::feature::builtin::manage_worker::sub_worker_control_feature(
|
||||
@@ -1148,38 +1181,6 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
// Memory tools require explicit feature exposure. Workspace memory access
|
||||
// is authority-bound to the Backend Workspace API; the Worker must not
|
||||
// register local filesystem memory tools even when it has local cwd/root
|
||||
// authority for shell/file tools.
|
||||
if feature_config.memory.enabled {
|
||||
let _mem = memory_config.as_ref().ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"[feature.memory].enabled = true requires a [memory] configuration section",
|
||||
)
|
||||
})?;
|
||||
if workspace_client.is_available() && workspace_client.workspace_id().is_some() {
|
||||
let definitions = if feature_config.memory.staging {
|
||||
crate::feature::builtin::memory::workspace_http_memory_consolidation_tools(
|
||||
workspace_client.clone(),
|
||||
)
|
||||
} else {
|
||||
crate::feature::builtin::memory::workspace_http_memory_tools(
|
||||
workspace_client.clone(),
|
||||
)
|
||||
};
|
||||
for definition in definitions {
|
||||
engine.register_tool(definition);
|
||||
}
|
||||
} else {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"memory tools require Backend Workspace API authority",
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let mut observation_providers: Vec<
|
||||
Arc<dyn crate::feature::builtin::worker_observation::WorkerObservationProvider>,
|
||||
> = Vec::new();
|
||||
@@ -1235,6 +1236,9 @@ where
|
||||
),
|
||||
));
|
||||
}
|
||||
if let Some((resident_summary, system_prompt_override)) = memory_prompt_contribution {
|
||||
worker.install_system_prompt_contribution(resident_summary, system_prompt_override);
|
||||
}
|
||||
if let Some(tracker) = tracker {
|
||||
worker.attach_tracker(tracker);
|
||||
}
|
||||
@@ -1590,7 +1594,9 @@ async fn controller_loop<C, St>(
|
||||
&working_event_tx,
|
||||
target,
|
||||
expected_head_entries,
|
||||
) {
|
||||
)
|
||||
.await
|
||||
{
|
||||
worker.clear_in_flight_events();
|
||||
shared_state.set_status(WorkerStatus::Idle);
|
||||
let _ = working_event_tx.send(Event::Status {
|
||||
@@ -1711,10 +1717,9 @@ async fn controller_loop<C, St>(
|
||||
tracing::warn!(%error, "Worker runtime socket cleanup failed");
|
||||
}
|
||||
|
||||
// Background memory jobs own extract/consolidate workers after a
|
||||
// turn completes. Join them before closing the Workdir session so no
|
||||
// Worker-owned task can outlive its operation attachment.
|
||||
worker.wait_for_memory_jobs().await;
|
||||
// Feature callbacks and tasks share the Worker scope. Stop them before
|
||||
// Memory/Workdir teardown so they cannot observe a partially closed Worker.
|
||||
worker.stop_feature_runtime("controller shutdown").await;
|
||||
|
||||
if let Some(session) = worker.workdir_session()
|
||||
&& let Err(error) = session.close().await
|
||||
@@ -1993,7 +1998,7 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_rewind<C, St>(
|
||||
async fn apply_rewind<C, St>(
|
||||
worker: &mut Worker<C, St>,
|
||||
working_event_tx: &broadcast::Sender<Event>,
|
||||
target: RewindTargetId,
|
||||
@@ -2003,7 +2008,7 @@ where
|
||||
C: LlmClient + 'static,
|
||||
St: Store,
|
||||
{
|
||||
match worker.rewind_to(target, expected_head_entries) {
|
||||
match worker.rewind_to(target, expected_head_entries).await {
|
||||
Ok(applied) => {
|
||||
let session =
|
||||
session_store::public_snapshot::project_current_session_snapshot(&applied.entries);
|
||||
|
||||
+392
-148
@@ -23,7 +23,14 @@ use agen::tool::ToolDefinition;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::hook::{Hook, HookRegistryBuilder, OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall};
|
||||
use crate::hook::{
|
||||
BeforeSessionRewrite, Hook, HookExecutionPolicy, HookRegistryBuilder, OnPromptSubmit,
|
||||
OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall, RunCommitted, RunExit, WorkerStopping,
|
||||
};
|
||||
use background::{
|
||||
BackgroundTaskSpec, FeatureBackgroundTask, FeatureBackgroundTaskRegistry,
|
||||
FeatureBackgroundTaskRegistryBuilder,
|
||||
};
|
||||
|
||||
/// Stable source-qualified identifier for a feature module.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
|
||||
@@ -253,10 +260,15 @@ pub enum FeatureRuntimeKind {
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum FeatureHookPoint {
|
||||
PreRequest,
|
||||
PromptSubmit,
|
||||
PreLlmRequest,
|
||||
PreToolCall,
|
||||
ToolResult,
|
||||
TurnEnd,
|
||||
PostToolCall,
|
||||
AssistantTurnEnd,
|
||||
RunExit,
|
||||
RunCommitted,
|
||||
BeforeSessionRewrite,
|
||||
WorkerStopping,
|
||||
}
|
||||
|
||||
/// Serializable declaration of a tool contribution. The executable factory is
|
||||
@@ -379,16 +391,17 @@ impl FeatureInstructionContribution {
|
||||
}
|
||||
}
|
||||
|
||||
/// Background task lifecycle phase represented by this registry slice.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
/// Background tasks are always Worker-managed and execute inside the owning
|
||||
/// feature scope. Report-only and detached host-managed declarations are not
|
||||
/// accepted by the current contract.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BackgroundTaskLifecycle {
|
||||
DescriptorOnly,
|
||||
HostManaged,
|
||||
WorkerManaged,
|
||||
}
|
||||
|
||||
/// Declaration for a feature-provided background task.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
/// Declaration for a feature-provided executable background task.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
|
||||
pub struct BackgroundTaskDeclaration {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
@@ -396,11 +409,11 @@ pub struct BackgroundTaskDeclaration {
|
||||
}
|
||||
|
||||
impl BackgroundTaskDeclaration {
|
||||
pub fn descriptor_only(name: impl Into<String>, description: impl Into<String>) -> Self {
|
||||
pub fn worker_managed(name: impl Into<String>, description: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
description: description.into(),
|
||||
lifecycle: BackgroundTaskLifecycle::DescriptorOnly,
|
||||
lifecycle: BackgroundTaskLifecycle::WorkerManaged,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -772,6 +785,15 @@ impl FeatureInstallReport {
|
||||
}
|
||||
}
|
||||
|
||||
fn clear_installed_contributions(&mut self) {
|
||||
self.installed = false;
|
||||
self.installed_tools.clear();
|
||||
self.installed_hooks.clear();
|
||||
self.installed_instructions.clear();
|
||||
self.declared_background_tasks.clear();
|
||||
self.provided_services.clear();
|
||||
}
|
||||
|
||||
fn mark_skipped(
|
||||
&mut self,
|
||||
kind: FeatureContributionKind,
|
||||
@@ -881,46 +903,6 @@ fn reject_undeclared_contribution(
|
||||
error
|
||||
}
|
||||
|
||||
/// Model-visible durable notification sink skeleton. The first slice exposes
|
||||
/// the boundary without implementing a new event channel.
|
||||
pub struct FeatureNotificationSink<'a> {
|
||||
report: &'a mut FeatureInstallReport,
|
||||
}
|
||||
|
||||
impl FeatureNotificationSink<'_> {
|
||||
pub fn notify_model(&mut self, message: impl Into<String>) -> Result<(), FeatureInstallError> {
|
||||
let message = message.into();
|
||||
self.report.diagnostics.push(FeatureDiagnostic::warning(format!(
|
||||
"model notification requested during feature installation but no durable Notify host is attached: {message}"
|
||||
)));
|
||||
self.report.mark_skipped(
|
||||
FeatureContributionKind::Notification,
|
||||
"notify_model",
|
||||
"durable Notify/SystemItem host is not connected during feature installation",
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Transient human-facing alert sink skeleton.
|
||||
pub struct FeatureAlertSink<'a> {
|
||||
report: &'a mut FeatureInstallReport,
|
||||
}
|
||||
|
||||
impl FeatureAlertSink<'_> {
|
||||
pub fn alert(&mut self, message: impl Into<String>) {
|
||||
let message = message.into();
|
||||
self.report
|
||||
.diagnostics
|
||||
.push(FeatureDiagnostic::info(format!("feature alert: {message}")));
|
||||
self.report.mark_skipped(
|
||||
FeatureContributionKind::Alert,
|
||||
"alert",
|
||||
"transient alert host is not connected during feature installation",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Diagnostic sink available to feature installers.
|
||||
pub struct FeatureDiagnosticSink<'a> {
|
||||
report: &'a mut FeatureInstallReport,
|
||||
@@ -1042,15 +1024,74 @@ impl HookContributionRegistrar<'_> {
|
||||
))
|
||||
}
|
||||
|
||||
fn record(&mut self, declaration: HookDeclaration) {
|
||||
if !self.report.installed_hooks.contains(&declaration) {
|
||||
self.report.installed_hooks.push(declaration);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_prompt_submit(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<OnPromptSubmit> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::PromptSubmit);
|
||||
self.require_declared(&declaration)?;
|
||||
self.hook_builder
|
||||
.add_named_on_prompt_submit(
|
||||
format!("{}:{}", self.feature_id, declaration.name),
|
||||
policy,
|
||||
hook,
|
||||
)
|
||||
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
|
||||
self.record(declaration);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn add_pre_llm_request(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<PreLlmRequest> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::PreLlmRequest);
|
||||
self.require_declared(&declaration)?;
|
||||
self.hook_builder
|
||||
.add_named_pre_llm_request(
|
||||
format!("{}:{}", self.feature_id, declaration.name),
|
||||
policy,
|
||||
hook,
|
||||
)
|
||||
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
|
||||
self.record(declaration);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn add_pre_request(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
hook: impl Hook<PreLlmRequest> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::PreRequest);
|
||||
self.add_pre_llm_request(name, HookExecutionPolicy::fail_closed(), hook)
|
||||
}
|
||||
|
||||
pub fn add_pre_tool_call_with_policy(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<PreToolCall> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::PreToolCall);
|
||||
self.require_declared(&declaration)?;
|
||||
self.hook_builder.add_pre_llm_request(hook);
|
||||
self.report.installed_hooks.push(declaration);
|
||||
self.hook_builder
|
||||
.add_named_pre_tool_call(
|
||||
format!("{}:{}", self.feature_id, declaration.name),
|
||||
policy,
|
||||
hook,
|
||||
)
|
||||
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
|
||||
self.record(declaration);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1059,10 +1100,25 @@ impl HookContributionRegistrar<'_> {
|
||||
name: impl Into<String>,
|
||||
hook: impl Hook<PreToolCall> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::PreToolCall);
|
||||
self.add_pre_tool_call_with_policy(name, HookExecutionPolicy::fail_closed(), hook)
|
||||
}
|
||||
|
||||
pub fn add_post_tool_call(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<PostToolCall> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::PostToolCall);
|
||||
self.require_declared(&declaration)?;
|
||||
self.hook_builder.add_pre_tool_call(hook);
|
||||
self.report.installed_hooks.push(declaration);
|
||||
self.hook_builder
|
||||
.add_named_post_tool_call(
|
||||
format!("{}:{}", self.feature_id, declaration.name),
|
||||
policy,
|
||||
hook,
|
||||
)
|
||||
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
|
||||
self.record(declaration);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1071,10 +1127,25 @@ impl HookContributionRegistrar<'_> {
|
||||
name: impl Into<String>,
|
||||
hook: impl Hook<PostToolCall> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::ToolResult);
|
||||
self.add_post_tool_call(name, HookExecutionPolicy::fail_closed(), hook)
|
||||
}
|
||||
|
||||
pub fn add_assistant_turn_end(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<OnTurnEnd> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::AssistantTurnEnd);
|
||||
self.require_declared(&declaration)?;
|
||||
self.hook_builder.add_post_tool_call(hook);
|
||||
self.report.installed_hooks.push(declaration);
|
||||
self.hook_builder
|
||||
.add_named_on_turn_end(
|
||||
format!("{}:{}", self.feature_id, declaration.name),
|
||||
policy,
|
||||
hook,
|
||||
)
|
||||
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
|
||||
self.record(declaration);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1083,10 +1154,82 @@ impl HookContributionRegistrar<'_> {
|
||||
name: impl Into<String>,
|
||||
hook: impl Hook<OnTurnEnd> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::TurnEnd);
|
||||
self.add_assistant_turn_end(name, HookExecutionPolicy::fail_closed(), hook)
|
||||
}
|
||||
|
||||
pub fn add_run_exit(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<RunExit> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::RunExit);
|
||||
self.require_declared(&declaration)?;
|
||||
self.hook_builder.add_on_turn_end(hook);
|
||||
self.report.installed_hooks.push(declaration);
|
||||
self.hook_builder
|
||||
.add_named_run_exit(
|
||||
format!("{}:{}", self.feature_id, declaration.name),
|
||||
policy,
|
||||
hook,
|
||||
)
|
||||
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
|
||||
self.record(declaration);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn add_run_committed(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<RunCommitted> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::RunCommitted);
|
||||
self.require_declared(&declaration)?;
|
||||
self.hook_builder
|
||||
.add_named_run_committed(
|
||||
format!("{}:{}", self.feature_id, declaration.name),
|
||||
policy,
|
||||
hook,
|
||||
)
|
||||
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
|
||||
self.record(declaration);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn add_before_session_rewrite(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<BeforeSessionRewrite> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::BeforeSessionRewrite);
|
||||
self.require_declared(&declaration)?;
|
||||
self.hook_builder
|
||||
.add_named_before_session_rewrite(
|
||||
format!("{}:{}", self.feature_id, declaration.name),
|
||||
policy,
|
||||
hook,
|
||||
)
|
||||
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
|
||||
self.record(declaration);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn add_worker_stopping(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<WorkerStopping> + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
let declaration = HookDeclaration::new(name, FeatureHookPoint::WorkerStopping);
|
||||
self.require_declared(&declaration)?;
|
||||
self.hook_builder
|
||||
.add_named_worker_stopping(
|
||||
format!("{}:{}", self.feature_id, declaration.name),
|
||||
policy,
|
||||
hook,
|
||||
)
|
||||
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
|
||||
self.record(declaration);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1124,33 +1267,40 @@ impl FeatureInstructionRegistrar<'_> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Background task registrar for descriptor/report-only contributions.
|
||||
/// Registrar for executable, Worker-managed background task contributions.
|
||||
pub struct BackgroundTaskRegistrar<'a> {
|
||||
feature_id: &'a FeatureId,
|
||||
declarations: &'a FeatureContributionDeclarations,
|
||||
registry: &'a mut FeatureBackgroundTaskRegistryBuilder,
|
||||
report: &'a mut FeatureInstallReport,
|
||||
}
|
||||
|
||||
impl BackgroundTaskRegistrar<'_> {
|
||||
pub fn declare(
|
||||
pub fn register(
|
||||
&mut self,
|
||||
declaration: BackgroundTaskDeclaration,
|
||||
spec: BackgroundTaskSpec,
|
||||
task: impl FeatureBackgroundTask + 'static,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
if !self.declarations.contains_background_task(&declaration) {
|
||||
if !self
|
||||
.declarations
|
||||
.contains_background_task(&spec.declaration)
|
||||
{
|
||||
return Err(reject_undeclared_contribution(
|
||||
self.feature_id,
|
||||
self.report,
|
||||
FeatureContributionKind::BackgroundTask,
|
||||
declaration.name,
|
||||
spec.declaration.name,
|
||||
));
|
||||
}
|
||||
self.registry
|
||||
.register(self.feature_id.clone(), spec.clone(), task)?;
|
||||
if !self
|
||||
.report
|
||||
.declared_background_tasks
|
||||
.iter()
|
||||
.any(|task| task.name == declaration.name)
|
||||
.any(|task| task.name == spec.declaration.name)
|
||||
{
|
||||
self.report.declared_background_tasks.push(declaration);
|
||||
self.report.declared_background_tasks.push(spec.declaration);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -1330,15 +1480,17 @@ impl ProtocolProviderRegistrar<'_> {
|
||||
}
|
||||
}
|
||||
|
||||
for task in background_tasks {
|
||||
if !self
|
||||
.report
|
||||
.declared_background_tasks
|
||||
.iter()
|
||||
.any(|declared| declared.name == task.name)
|
||||
{
|
||||
self.report.declared_background_tasks.push(task);
|
||||
}
|
||||
if let Some(task) = background_tasks.first() {
|
||||
let reason = format!(
|
||||
"protocol provider background task `{}` has no executable Worker-managed handler",
|
||||
task.name
|
||||
);
|
||||
self.report.mark_skipped(
|
||||
FeatureContributionKind::BackgroundTask,
|
||||
task.name.clone(),
|
||||
reason.clone(),
|
||||
);
|
||||
return Err(FeatureInstallError::InvalidDescriptor(reason));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -1352,6 +1504,7 @@ pub struct FeatureInstallContext<'a> {
|
||||
pending_tools: &'a mut Vec<ToolDefinition>,
|
||||
installed_tool_names: &'a mut HashMap<String, FeatureId>,
|
||||
hook_builder: &'a mut HookRegistryBuilder,
|
||||
background_task_builder: &'a mut FeatureBackgroundTaskRegistryBuilder,
|
||||
service_registry: &'a mut FeatureServiceRegistry,
|
||||
report: &'a mut FeatureInstallReport,
|
||||
}
|
||||
@@ -1392,6 +1545,7 @@ impl FeatureInstallContext<'_> {
|
||||
BackgroundTaskRegistrar {
|
||||
feature_id: self.feature_id,
|
||||
declarations: self.declarations,
|
||||
registry: self.background_task_builder,
|
||||
report: self.report,
|
||||
}
|
||||
}
|
||||
@@ -1416,18 +1570,6 @@ impl FeatureInstallContext<'_> {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn notifications(&mut self) -> FeatureNotificationSink<'_> {
|
||||
FeatureNotificationSink {
|
||||
report: self.report,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn alerts(&mut self) -> FeatureAlertSink<'_> {
|
||||
FeatureAlertSink {
|
||||
report: self.report,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn diagnostics(&mut self) -> FeatureDiagnosticSink<'_> {
|
||||
FeatureDiagnosticSink {
|
||||
report: self.report,
|
||||
@@ -1440,6 +1582,7 @@ impl FeatureInstallContext<'_> {
|
||||
pub struct FeatureRegistryInstallReport {
|
||||
pub reports: Vec<FeatureInstallReport>,
|
||||
pub services: FeatureServiceRegistry,
|
||||
pub background_tasks: FeatureBackgroundTaskRegistry,
|
||||
pub plan_error: Option<FeaturePlanError>,
|
||||
}
|
||||
|
||||
@@ -1795,7 +1938,7 @@ impl FeatureRegistryBuilder {
|
||||
}
|
||||
|
||||
/// Install modules into the existing Engine tool path and hook builder.
|
||||
pub(crate) fn install_into_engine<C: LlmClient, A>(
|
||||
pub(crate) fn install_into_engine<C: LlmClient, A: Send + Sync>(
|
||||
self,
|
||||
worker: &mut Engine<C, Mutable, A>,
|
||||
hook_builder: &mut HookRegistryBuilder,
|
||||
@@ -1861,12 +2004,16 @@ impl FeatureRegistryBuilder {
|
||||
return FeatureRegistryInstallReport {
|
||||
reports,
|
||||
services: FeatureServiceRegistry::default(),
|
||||
background_tasks: FeatureBackgroundTaskRegistry::default(),
|
||||
plan_error: Some(error),
|
||||
};
|
||||
}
|
||||
};
|
||||
let mut service_registry = FeatureServiceRegistry::default();
|
||||
let mut background_task_builder = FeatureBackgroundTaskRegistryBuilder::default();
|
||||
let mut reports = Vec::with_capacity(plan.ordered_indices.len());
|
||||
let install_hook_checkpoint = hook_builder.checkpoint();
|
||||
let install_tool_checkpoint = pending_tools.len();
|
||||
let mut modules = self.modules.into_iter().map(Some).collect::<Vec<_>>();
|
||||
let ordered_modules = plan
|
||||
.ordered_indices
|
||||
@@ -1884,6 +2031,11 @@ impl FeatureRegistryBuilder {
|
||||
for (module, descriptor) in ordered_modules {
|
||||
let declarations = FeatureContributionDeclarations::from_descriptor(&descriptor);
|
||||
let mut report = FeatureInstallReport::new(&descriptor);
|
||||
let hook_checkpoint = hook_builder.checkpoint();
|
||||
let background_checkpoint = background_task_builder.checkpoint();
|
||||
let service_checkpoint = service_registry.clone();
|
||||
let tool_checkpoint = pending_tools.len();
|
||||
let installed_tool_checkpoint = installed_tool_names.clone();
|
||||
|
||||
let mut required_service_failed = false;
|
||||
for requirement in descriptor.requires_services.iter().cloned() {
|
||||
@@ -1920,10 +2072,6 @@ impl FeatureRegistryBuilder {
|
||||
continue;
|
||||
}
|
||||
|
||||
for background_task in descriptor.background_tasks.iter().cloned() {
|
||||
report.declared_background_tasks.push(background_task);
|
||||
}
|
||||
|
||||
let install_result = {
|
||||
let mut context = FeatureInstallContext {
|
||||
feature_id: &descriptor.id,
|
||||
@@ -1931,6 +2079,7 @@ impl FeatureRegistryBuilder {
|
||||
pending_tools,
|
||||
installed_tool_names: &mut installed_tool_names,
|
||||
hook_builder,
|
||||
background_task_builder: &mut background_task_builder,
|
||||
service_registry: &mut service_registry,
|
||||
report: &mut report,
|
||||
};
|
||||
@@ -1940,18 +2089,81 @@ impl FeatureRegistryBuilder {
|
||||
match install_result {
|
||||
Ok(()) => report.installed = true,
|
||||
Err(error) => {
|
||||
hook_builder.rollback_to(hook_checkpoint);
|
||||
background_task_builder.rollback_to(&background_checkpoint);
|
||||
service_registry = service_checkpoint.clone();
|
||||
pending_tools.truncate(tool_checkpoint);
|
||||
installed_tool_names = installed_tool_checkpoint.clone();
|
||||
report.clear_installed_contributions();
|
||||
report
|
||||
.diagnostics
|
||||
.push(FeatureDiagnostic::error(error.to_string()));
|
||||
}
|
||||
}
|
||||
if report.installed {
|
||||
for hook in &descriptor.hooks {
|
||||
if !report.installed_hooks.contains(hook) {
|
||||
report.diagnostics.push(FeatureDiagnostic::error(format!(
|
||||
"feature `{}` declared hook `{}` at {:?} but did not register it",
|
||||
descriptor.id, hook.name, hook.point
|
||||
)));
|
||||
}
|
||||
}
|
||||
for task in &descriptor.background_tasks {
|
||||
if !report.declared_background_tasks.contains(task) {
|
||||
report.diagnostics.push(FeatureDiagnostic::error(format!(
|
||||
"feature `{}` declared background task `{}` but did not register an executable handler",
|
||||
descriptor.id, task.name
|
||||
)));
|
||||
}
|
||||
}
|
||||
if report
|
||||
.diagnostics
|
||||
.iter()
|
||||
.any(|diagnostic| diagnostic.severity == FeatureDiagnosticSeverity::Error)
|
||||
{
|
||||
hook_builder.rollback_to(hook_checkpoint);
|
||||
background_task_builder.rollback_to(&background_checkpoint);
|
||||
service_registry = service_checkpoint.clone();
|
||||
pending_tools.truncate(tool_checkpoint);
|
||||
installed_tool_names = installed_tool_checkpoint.clone();
|
||||
report.clear_installed_contributions();
|
||||
report.clear_installed_contributions();
|
||||
}
|
||||
}
|
||||
reports.push(report);
|
||||
}
|
||||
|
||||
FeatureRegistryInstallReport {
|
||||
reports,
|
||||
services: service_registry,
|
||||
plan_error: None,
|
||||
let failed = reports.iter().any(|report| {
|
||||
report
|
||||
.diagnostics
|
||||
.iter()
|
||||
.any(|diagnostic| diagnostic.severity == FeatureDiagnosticSeverity::Error)
|
||||
});
|
||||
if failed {
|
||||
hook_builder.rollback_to(install_hook_checkpoint);
|
||||
pending_tools.truncate(install_tool_checkpoint);
|
||||
for report in &mut reports {
|
||||
if report.installed {
|
||||
report.clear_installed_contributions();
|
||||
report.diagnostics.push(FeatureDiagnostic::warning(
|
||||
"feature scope rolled back because another contribution failed",
|
||||
));
|
||||
}
|
||||
}
|
||||
FeatureRegistryInstallReport {
|
||||
reports,
|
||||
services: FeatureServiceRegistry::default(),
|
||||
background_tasks: FeatureBackgroundTaskRegistry::default(),
|
||||
plan_error: None,
|
||||
}
|
||||
} else {
|
||||
FeatureRegistryInstallReport {
|
||||
reports,
|
||||
services: service_registry,
|
||||
background_tasks: background_task_builder.build(),
|
||||
plan_error: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1996,9 +2208,11 @@ pub enum FeatureInstallError {
|
||||
Install(String),
|
||||
}
|
||||
|
||||
pub mod background;
|
||||
pub mod builtin;
|
||||
pub mod mcp;
|
||||
pub mod plugin;
|
||||
pub(crate) mod session;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
@@ -2398,13 +2612,9 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn descriptor_contributions_are_recorded() {
|
||||
fn executable_contributions_are_recorded() {
|
||||
let descriptor = FeatureDescriptor::builtin("dummy", "Dummy")
|
||||
.with_tool(ToolDeclaration::new("Dummy", "dummy tool"))
|
||||
.with_background_task(BackgroundTaskDeclaration::descriptor_only(
|
||||
"daily",
|
||||
"descriptor-only background task",
|
||||
));
|
||||
.with_tool(ToolDeclaration::new("Dummy", "dummy tool"));
|
||||
let mut hook_builder = HookRegistryBuilder::default();
|
||||
let mut pending_tools = Vec::new();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
@@ -2420,7 +2630,7 @@ mod tests {
|
||||
let feature_report = &report.reports[0];
|
||||
assert!(feature_report.installed);
|
||||
assert_eq!(feature_report.installed_tools, vec!["Dummy"]);
|
||||
assert_eq!(feature_report.declared_background_tasks[0].name, "daily");
|
||||
assert!(feature_report.declared_background_tasks.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -2480,8 +2690,9 @@ mod tests {
|
||||
})
|
||||
.install_into_pending(&mut pending_tools, &mut hook_builder);
|
||||
|
||||
assert_eq!(pending_tools.len(), 1);
|
||||
assert!(report.reports[0].installed);
|
||||
assert!(pending_tools.is_empty());
|
||||
assert!(!report.reports[0].installed);
|
||||
assert!(report.reports[0].installed_tools.is_empty());
|
||||
assert!(!report.reports[1].installed);
|
||||
assert!(
|
||||
report.reports[1]
|
||||
@@ -2558,7 +2769,7 @@ mod tests {
|
||||
"1.0.0",
|
||||
"startup-discovered service",
|
||||
))
|
||||
.with_background_task(BackgroundTaskDeclaration::descriptor_only(
|
||||
.with_background_task(BackgroundTaskDeclaration::worker_managed(
|
||||
"provider-poller",
|
||||
"provider lifecycle poller",
|
||||
))
|
||||
@@ -2568,7 +2779,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn protocol_provider_registers_startup_discovered_contributions_through_worker_path() {
|
||||
fn protocol_provider_report_only_background_task_is_rejected_atomically() {
|
||||
let provider = ProtocolProviderDeclaration::new(
|
||||
ProviderId::builtin("dynamic-provider"),
|
||||
"test-protocol",
|
||||
@@ -2599,30 +2810,18 @@ mod tests {
|
||||
.collect();
|
||||
let feature_report = &report.reports[0];
|
||||
|
||||
assert!(feature_report.installed);
|
||||
assert_eq!(feature_report.installed_tools, vec!["DynamicTool"]);
|
||||
assert_eq!(tool_names, vec!["DynamicTool"]);
|
||||
assert!(!feature_report.installed);
|
||||
assert!(feature_report.installed_tools.is_empty());
|
||||
assert!(tool_names.is_empty());
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(feature_report.provided_services.len(), 1);
|
||||
assert_eq!(
|
||||
feature_report.provided_services[0].id,
|
||||
ServiceId::builtin("dynamic-service")
|
||||
);
|
||||
assert_eq!(
|
||||
feature_report.declared_background_tasks[0].name,
|
||||
"provider-poller"
|
||||
);
|
||||
assert!(feature_report.provided_services.is_empty());
|
||||
assert!(feature_report.declared_background_tasks.is_empty());
|
||||
assert_eq!(feature_report.protocol_providers.len(), 1);
|
||||
assert_eq!(
|
||||
feature_report.protocol_providers[0].state,
|
||||
ProtocolProviderLifecycleState::Ready
|
||||
);
|
||||
assert!(
|
||||
feature_report
|
||||
.diagnostics
|
||||
.iter()
|
||||
.any(|diagnostic| diagnostic.message.contains("startup discovery completed"))
|
||||
);
|
||||
assert!(feature_report.diagnostics.iter().any(|diagnostic| {
|
||||
diagnostic
|
||||
.message
|
||||
.contains("has no executable Worker-managed handler")
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -2779,8 +2978,8 @@ mod tests {
|
||||
async fn call(
|
||||
&self,
|
||||
_input: &crate::hook::ToolCallSummary,
|
||||
) -> crate::hook::HookPreToolAction {
|
||||
crate::hook::HookPreToolAction::Continue
|
||||
) -> Result<crate::hook::HookPreToolAction, crate::hook::HookError> {
|
||||
Ok(crate::hook::HookPreToolAction::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2804,6 +3003,19 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
struct NoopBackgroundTask;
|
||||
|
||||
#[async_trait]
|
||||
impl FeatureBackgroundTask for NoopBackgroundTask {
|
||||
async fn run(
|
||||
&self,
|
||||
_context: background::BackgroundTaskContext,
|
||||
_cancellation: background::BackgroundTaskCancellation,
|
||||
) -> Result<(), crate::hook::HookError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
struct BackgroundFeature {
|
||||
descriptor: FeatureDescriptor,
|
||||
task_name: &'static str,
|
||||
@@ -2818,12 +3030,22 @@ mod tests {
|
||||
&self,
|
||||
context: &mut FeatureInstallContext<'_>,
|
||||
) -> Result<(), FeatureInstallError> {
|
||||
context
|
||||
.background_tasks()
|
||||
.declare(BackgroundTaskDeclaration::descriptor_only(
|
||||
self.task_name,
|
||||
"runtime background task",
|
||||
))
|
||||
let declaration = self
|
||||
.descriptor
|
||||
.background_tasks
|
||||
.iter()
|
||||
.find(|task| task.name == self.task_name)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| {
|
||||
BackgroundTaskDeclaration::worker_managed(
|
||||
self.task_name,
|
||||
"undeclared background task",
|
||||
)
|
||||
});
|
||||
context.background_tasks().register(
|
||||
BackgroundTaskSpec::single_flight(declaration, std::time::Duration::from_secs(1)),
|
||||
NoopBackgroundTask,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2985,25 +3207,44 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn background_task_declaration_is_descriptor_contribution() {
|
||||
#[tokio::test]
|
||||
async fn executable_background_task_is_registered_in_worker_scope() {
|
||||
let descriptor = FeatureDescriptor::builtin("background", "Background")
|
||||
.with_background_task(BackgroundTaskDeclaration::descriptor_only(
|
||||
.with_background_task(BackgroundTaskDeclaration::worker_managed(
|
||||
"declared-task",
|
||||
"descriptor contribution",
|
||||
));
|
||||
let mut hook_builder = HookRegistryBuilder::default();
|
||||
let mut pending_tools = Vec::new();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(ServiceFeature { descriptor })
|
||||
.with_module(BackgroundFeature {
|
||||
descriptor,
|
||||
task_name: "declared-task",
|
||||
})
|
||||
.install_into_pending(&mut pending_tools, &mut hook_builder);
|
||||
|
||||
assert!(report.reports[0].installed);
|
||||
assert_eq!(
|
||||
report.reports[0].declared_background_tasks[0].name,
|
||||
"declared-task"
|
||||
);
|
||||
assert!(report.reports[0].skipped.is_empty());
|
||||
assert!(matches!(
|
||||
report
|
||||
.background_tasks
|
||||
.start(
|
||||
&FeatureId::builtin("background"),
|
||||
"declared-task",
|
||||
crate::hook::HookInvocationContext::default(),
|
||||
)
|
||||
.unwrap(),
|
||||
background::BackgroundTaskStart::Started { .. }
|
||||
));
|
||||
report.background_tasks.shutdown().await.unwrap();
|
||||
assert!(matches!(
|
||||
report.background_tasks.diagnostics()[0].outcome,
|
||||
background::BackgroundTaskOutcome::Completed
|
||||
| background::BackgroundTaskOutcome::Cancelled
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -3118,7 +3359,10 @@ mod tests {
|
||||
assert_eq!(descriptor.runtime, FeatureRuntimeKind::Builtin);
|
||||
assert_eq!(
|
||||
hook_points,
|
||||
vec![FeatureHookPoint::PreRequest, FeatureHookPoint::PreToolCall]
|
||||
vec![
|
||||
FeatureHookPoint::PreLlmRequest,
|
||||
FeatureHookPoint::PreToolCall
|
||||
]
|
||||
);
|
||||
assert!(descriptor.background_tasks.is_empty());
|
||||
assert!(descriptor.provides_services.is_empty());
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -8,7 +8,8 @@ pub mod flow_transition;
|
||||
pub mod manage_workdir;
|
||||
pub mod manage_worker;
|
||||
pub mod memory;
|
||||
pub mod memory_extract;
|
||||
pub(crate) mod memory_lifecycle;
|
||||
pub mod memory_staging_output;
|
||||
pub mod merge_request;
|
||||
pub mod objective;
|
||||
pub mod orchestration;
|
||||
@@ -19,8 +20,6 @@ pub mod ticket;
|
||||
pub mod worker_observation;
|
||||
pub mod workspace_worker_discovery;
|
||||
|
||||
pub(crate) use memory_extract::{MemoryExtractFeature, MemoryExtractState, render_extract_input};
|
||||
pub(crate) use session_explore::{SessionExploreFeature, SessionExploreState};
|
||||
pub use task::{TaskFeature, task_tools_feature};
|
||||
pub use ticket::{
|
||||
TicketFeature, TicketFeatureAccess, ticket_tools_feature, ticket_tools_feature_with_access,
|
||||
|
||||
@@ -18,6 +18,10 @@ use schemars::JsonSchema;
|
||||
use serde::de::DeserializeOwned;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::feature::{
|
||||
FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ToolContribution,
|
||||
ToolDeclaration,
|
||||
};
|
||||
use crate::worker::{
|
||||
WorkspaceClient, WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod,
|
||||
};
|
||||
@@ -338,6 +342,151 @@ fn query_schema() -> serde_json::Value {
|
||||
})
|
||||
}
|
||||
|
||||
pub struct MemoryFeatureInstallPlan {
|
||||
pub module: MemoryToolsFeature,
|
||||
pub resident_summary: Option<String>,
|
||||
pub system_prompt_override: Option<String>,
|
||||
pub(crate) resolved_config: manifest::ResolvedMemoryFeatureConfig,
|
||||
}
|
||||
|
||||
impl MemoryFeatureInstallPlan {
|
||||
pub async fn prepare(
|
||||
manifest: &manifest::WorkerManifest,
|
||||
client: Arc<dyn WorkspaceClient>,
|
||||
prompts: Arc<crate::prompt::catalog::PromptCatalog>,
|
||||
) -> std::io::Result<Option<Self>> {
|
||||
Self::prepare_resolved(
|
||||
manifest.feature.memory.clone(),
|
||||
client,
|
||||
prompts,
|
||||
manifest.profile.clone(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn prepare_resolved(
|
||||
config: manifest::ResolvedMemoryFeatureConfig,
|
||||
client: Arc<dyn WorkspaceClient>,
|
||||
prompts: Arc<crate::prompt::catalog::PromptCatalog>,
|
||||
profile: Option<manifest::ProfileManifestSnapshot>,
|
||||
) -> std::io::Result<Option<Self>> {
|
||||
let memory_consolidation_worker = profile.as_ref().is_some_and(|snapshot| {
|
||||
matches!(
|
||||
&snapshot.source,
|
||||
manifest::ProfileSource::Registry {
|
||||
source: manifest::ProfileRegistrySource::Builtin,
|
||||
name,
|
||||
..
|
||||
} if name == "memory-consolidation"
|
||||
)
|
||||
});
|
||||
config
|
||||
.validate_execution()
|
||||
.map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidInput, message))?;
|
||||
if !config.profile.enabled {
|
||||
return Ok(None);
|
||||
}
|
||||
let workspace_id = client.workspace_id().ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"Memory tools require Backend Workspace API authority",
|
||||
)
|
||||
})?;
|
||||
let settings = config
|
||||
.workspace_settings()
|
||||
.expect("validated enabled Memory config has Workspace settings");
|
||||
if settings.workspace_id != workspace_id {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"Memory settings belong to {} instead of {}",
|
||||
settings.workspace_id, workspace_id
|
||||
),
|
||||
));
|
||||
}
|
||||
|
||||
let resident_summary = if config.profile.resident.inject_summary {
|
||||
match client
|
||||
.execute_memory_backend_operation(
|
||||
memory::backend::MemoryBackendOperation::ResidentSummary(
|
||||
memory::backend::MemoryResidentSummaryOperation::default(),
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(memory::backend::MemoryBackendOperationResult::ToolOutput(output)) => {
|
||||
output.content
|
||||
}
|
||||
Ok(other) => {
|
||||
tracing::debug!(?other, "unexpected resident Memory Backend result");
|
||||
None
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::debug!(%error, "resident Memory summary unavailable");
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let system_prompt_override = if memory_consolidation_worker {
|
||||
let language = settings.language;
|
||||
Some(
|
||||
prompts
|
||||
.memory_consolidation_system(&language)
|
||||
.map_err(|error| std::io::Error::other(error.to_string()))?,
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Ok(Some(Self {
|
||||
module: MemoryToolsFeature::new(client, config.profile.staging_tools),
|
||||
resident_summary,
|
||||
system_prompt_override,
|
||||
resolved_config: config,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MemoryToolsFeature {
|
||||
tools: Vec<ToolDefinition>,
|
||||
}
|
||||
|
||||
impl MemoryToolsFeature {
|
||||
pub fn new(client: Arc<dyn WorkspaceClient>, staging_tools: bool) -> Self {
|
||||
let tools = if staging_tools {
|
||||
workspace_http_memory_consolidation_tools(client)
|
||||
} else {
|
||||
workspace_http_memory_tools(client)
|
||||
};
|
||||
Self { tools }
|
||||
}
|
||||
}
|
||||
|
||||
impl FeatureModule for MemoryToolsFeature {
|
||||
fn descriptor(&self) -> FeatureDescriptor {
|
||||
let mut descriptor = FeatureDescriptor::builtin("memory", "Memory")
|
||||
.with_description("Workspace Memory document, query, and staging tools.");
|
||||
for tool in &self.tools {
|
||||
let (meta, _) = tool();
|
||||
descriptor = descriptor.with_tool(ToolDeclaration::new(meta.name, meta.description));
|
||||
}
|
||||
descriptor
|
||||
}
|
||||
|
||||
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
|
||||
for tool in &self.tools {
|
||||
let (meta, _) = tool();
|
||||
context
|
||||
.tools()
|
||||
.register(ToolContribution::new(meta.name, tool.clone()))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -350,6 +499,39 @@ mod tests {
|
||||
))
|
||||
}
|
||||
|
||||
fn resident_client(content: &str) -> Arc<dyn WorkspaceClient> {
|
||||
use std::io::{Read, Write};
|
||||
use std::net::TcpListener;
|
||||
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let content = content.to_string();
|
||||
std::thread::spawn(move || {
|
||||
let (mut stream, _) = listener.accept().unwrap();
|
||||
let mut request = [0_u8; 1024];
|
||||
let _ = stream.read(&mut request).unwrap();
|
||||
let body = serde_json::json!({
|
||||
"status": "ok",
|
||||
"result": {
|
||||
"kind": "tool_output",
|
||||
"summary": "resident Memory summary collected",
|
||||
"content": content,
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
body.len(),
|
||||
body
|
||||
);
|
||||
stream.write_all(response.as_bytes()).unwrap();
|
||||
});
|
||||
Arc::new(crate::worker::TestWorkspaceHttpClient::new(
|
||||
"workspace",
|
||||
format!("http://{addr}"),
|
||||
))
|
||||
}
|
||||
|
||||
fn tool_names(definitions: Vec<ToolDefinition>) -> Vec<String> {
|
||||
let mut names = definitions
|
||||
.into_iter()
|
||||
@@ -368,6 +550,132 @@ mod tests {
|
||||
.input_schema
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn memory_install_plan_is_the_fail_closed_config_boundary() {
|
||||
let prompts = crate::prompt::catalog::PromptCatalog::builtins_only().unwrap();
|
||||
let disabled = MemoryFeatureInstallPlan::prepare_resolved(
|
||||
manifest::ResolvedMemoryFeatureConfig::default(),
|
||||
test_client(),
|
||||
prompts.clone(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(disabled.is_none());
|
||||
|
||||
let mut enabled = manifest::ResolvedMemoryFeatureConfig::default();
|
||||
enabled.profile.enabled = true;
|
||||
enabled.profile.resident.inject_summary = false;
|
||||
assert!(
|
||||
MemoryFeatureInstallPlan::prepare_resolved(
|
||||
enabled.clone(),
|
||||
test_client(),
|
||||
prompts.clone(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
enabled
|
||||
.bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot {
|
||||
workspace_id: "workspace".to_string(),
|
||||
settings_revision: 1,
|
||||
language: "English".to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
let mut foreign = enabled.clone();
|
||||
foreign.workspace_settings.as_mut().unwrap().workspace_id = "other-workspace".to_string();
|
||||
assert!(
|
||||
MemoryFeatureInstallPlan::prepare_resolved(
|
||||
foreign,
|
||||
test_client(),
|
||||
prompts.clone(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let plan = MemoryFeatureInstallPlan::prepare_resolved(
|
||||
enabled.clone(),
|
||||
test_client(),
|
||||
prompts.clone(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(plan.resident_summary.is_none());
|
||||
assert!(plan.system_prompt_override.is_none());
|
||||
|
||||
enabled.profile.resident.inject_summary = true;
|
||||
let plan = MemoryFeatureInstallPlan::prepare_resolved(
|
||||
enabled,
|
||||
resident_client("# Durable Memory"),
|
||||
prompts,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(plan.resident_summary.as_deref(), Some("# Durable Memory"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn memory_prompt_contribution_rereads_resident_summary_for_each_install() {
|
||||
let prompts = crate::prompt::catalog::PromptCatalog::builtins_only().unwrap();
|
||||
let mut config = manifest::ResolvedMemoryFeatureConfig::default();
|
||||
config.profile.enabled = true;
|
||||
config
|
||||
.bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot {
|
||||
workspace_id: "workspace".to_string(),
|
||||
settings_revision: 1,
|
||||
language: "English".to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let first = MemoryFeatureInstallPlan::prepare_resolved(
|
||||
config.clone(),
|
||||
resident_client("first resident summary"),
|
||||
prompts.clone(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let restored = MemoryFeatureInstallPlan::prepare_resolved(
|
||||
config,
|
||||
resident_client("updated resident summary"),
|
||||
prompts,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
first.resident_summary.as_deref(),
|
||||
Some("first resident summary")
|
||||
);
|
||||
assert_eq!(
|
||||
restored.resident_summary.as_deref(),
|
||||
Some("updated resident summary")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memory_feature_owns_normal_and_staging_tool_surfaces() {
|
||||
let normal = MemoryToolsFeature::new(test_client(), false);
|
||||
let normal_names = tool_names(normal.tools);
|
||||
assert!(normal_names.contains(&"MemoryQuery".to_string()));
|
||||
assert!(!normal_names.contains(&"MemoryStagingList".to_string()));
|
||||
|
||||
let staging = MemoryToolsFeature::new(test_client(), true);
|
||||
assert_eq!(staging.descriptor().id.as_str(), "builtin:memory");
|
||||
let staging_names = tool_names(staging.tools);
|
||||
assert!(staging_names.contains(&"MemoryQuery".to_string()));
|
||||
assert!(staging_names.contains(&"MemoryStagingList".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normal_workspace_memory_tools_do_not_include_staging_tools() {
|
||||
let names = tool_names(workspace_http_memory_tools(test_client()));
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+17
-19
@@ -28,7 +28,7 @@ const FINISH_DESCRIPTION: &str =
|
||||
"Finish Memory extraction after validating the number of candidates staged during this run.";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MemoryExtractState {
|
||||
pub(crate) struct MemoryStagingOutputState {
|
||||
view: Arc<SessionCapture>,
|
||||
workspace_client: Arc<dyn WorkspaceClient>,
|
||||
source: SourceRef,
|
||||
@@ -37,7 +37,7 @@ pub(crate) struct MemoryExtractState {
|
||||
finished: Arc<Mutex<Option<FinishMemoryExtractionParams>>>,
|
||||
}
|
||||
|
||||
impl MemoryExtractState {
|
||||
impl MemoryStagingOutputState {
|
||||
pub(crate) fn new(
|
||||
view: SessionCapture,
|
||||
workspace_client: Arc<dyn WorkspaceClient>,
|
||||
@@ -70,22 +70,20 @@ impl MemoryExtractState {
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MemoryExtractFeature {
|
||||
state: MemoryExtractState,
|
||||
pub(crate) struct MemoryStagingOutputFeature {
|
||||
state: MemoryStagingOutputState,
|
||||
}
|
||||
|
||||
impl MemoryExtractFeature {
|
||||
pub(crate) fn new(state: MemoryExtractState) -> Self {
|
||||
impl MemoryStagingOutputFeature {
|
||||
pub(crate) fn new(state: MemoryStagingOutputState) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
impl FeatureModule for MemoryExtractFeature {
|
||||
impl FeatureModule for MemoryStagingOutputFeature {
|
||||
fn descriptor(&self) -> FeatureDescriptor {
|
||||
FeatureDescriptor::builtin("memory-extract", "Memory Extract")
|
||||
.with_description(
|
||||
"Memory staging and extraction completion, independent from session exploration.",
|
||||
)
|
||||
FeatureDescriptor::builtin("memory-staging-output", "Memory Staging Output")
|
||||
.with_description("Restricted Memory staging output for an extraction Internal Worker.")
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"StageMemoryCandidate",
|
||||
STAGE_DESCRIPTION,
|
||||
@@ -109,7 +107,7 @@ impl FeatureModule for MemoryExtractFeature {
|
||||
}
|
||||
}
|
||||
|
||||
fn stage_definition(state: MemoryExtractState) -> ToolDefinition {
|
||||
fn stage_definition(state: MemoryStagingOutputState) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(StageMemoryCandidateParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
@@ -123,7 +121,7 @@ fn stage_definition(state: MemoryExtractState) -> ToolDefinition {
|
||||
})
|
||||
}
|
||||
|
||||
fn finish_definition(state: MemoryExtractState) -> ToolDefinition {
|
||||
fn finish_definition(state: MemoryStagingOutputState) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(FinishMemoryExtractionParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
@@ -157,7 +155,7 @@ struct FinishMemoryExtractionParams {
|
||||
}
|
||||
|
||||
struct StageMemoryCandidateTool {
|
||||
state: MemoryExtractState,
|
||||
state: MemoryStagingOutputState,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -252,7 +250,7 @@ impl Tool for StageMemoryCandidateTool {
|
||||
}
|
||||
|
||||
struct FinishMemoryExtractionTool {
|
||||
state: MemoryExtractState,
|
||||
state: MemoryStagingOutputState,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -431,8 +429,8 @@ mod tests {
|
||||
|
||||
use super::*;
|
||||
|
||||
fn state() -> MemoryExtractState {
|
||||
MemoryExtractState::new(
|
||||
fn state() -> MemoryStagingOutputState {
|
||||
MemoryStagingOutputState::new(
|
||||
SessionCapture::new("segment-1", vec![Item::user_message("durable decision")]),
|
||||
crate::worker::marker_workspace_client(None, "test-backend"),
|
||||
SourceRef {
|
||||
@@ -445,8 +443,8 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn memory_extract_declares_only_memory_mutation_tools() {
|
||||
let descriptor = MemoryExtractFeature::new(state()).descriptor();
|
||||
assert_eq!(descriptor.id.as_str(), "builtin:memory-extract");
|
||||
let descriptor = MemoryStagingOutputFeature::new(state()).descriptor();
|
||||
assert_eq!(descriptor.id.as_str(), "builtin:memory-staging-output");
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.tools
|
||||
@@ -125,7 +125,7 @@ impl FeatureModule for TaskFeature {
|
||||
))
|
||||
.with_hook(HookDeclaration::new(
|
||||
"task-reminder-pre-request",
|
||||
FeatureHookPoint::PreRequest,
|
||||
FeatureHookPoint::PreLlmRequest,
|
||||
))
|
||||
.with_hook(HookDeclaration::new(
|
||||
"task-reminder-tool-usage",
|
||||
@@ -209,24 +209,27 @@ struct TaskReminderPreRequestHook {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for TaskReminderPreRequestHook {
|
||||
async fn call(&self, input: &PreRequestContext) -> HookPreRequestAction {
|
||||
async fn call(
|
||||
&self,
|
||||
input: &PreRequestContext,
|
||||
) -> Result<HookPreRequestAction, crate::hook::HookError> {
|
||||
let tasks = self.state.task_store.list();
|
||||
if tasks.is_empty() {
|
||||
return HookPreRequestAction::Continue;
|
||||
return Ok(HookPreRequestAction::Continue);
|
||||
}
|
||||
|
||||
let (since_task_management, since_reminder) = self.state.reminder_state.note_request();
|
||||
if since_task_management < TASK_REMINDER_REQUEST_THRESHOLD
|
||||
|| since_reminder < TASK_REMINDER_COOLDOWN_REQUESTS
|
||||
{
|
||||
return HookPreRequestAction::Continue;
|
||||
return Ok(HookPreRequestAction::Continue);
|
||||
}
|
||||
|
||||
if let Some(system_items) = input.system_items() {
|
||||
self.state.reminder_state.note_reminder();
|
||||
system_items.append_task_reminder(render_task_reminder_body(&tasks));
|
||||
}
|
||||
HookPreRequestAction::Continue
|
||||
Ok(HookPreRequestAction::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -236,11 +239,14 @@ struct TaskReminderToolUsageHook {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreToolCall> for TaskReminderToolUsageHook {
|
||||
async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction {
|
||||
async fn call(
|
||||
&self,
|
||||
input: &ToolCallSummary,
|
||||
) -> Result<HookPreToolAction, crate::hook::HookError> {
|
||||
if is_task_management_tool(&input.tool_name) {
|
||||
self.state.reminder_state.note_task_management();
|
||||
}
|
||||
HookPreToolAction::Continue
|
||||
Ok(HookPreToolAction::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2115,8 +2115,7 @@ read exit_notification || true
|
||||
meta.name
|
||||
})
|
||||
.collect();
|
||||
assert!(!names.iter().any(|name| name == "Mcp_demo_search_files"));
|
||||
assert!(names.iter().any(|name| name == "Mcp_demo_unique"));
|
||||
assert!(names.is_empty());
|
||||
}
|
||||
|
||||
fn shell_tool_server(response: &str) -> McpStdioServerSpec {
|
||||
|
||||
@@ -7621,7 +7621,7 @@ mod tests {
|
||||
)))
|
||||
.install_into_pending(&mut pending, &mut hooks);
|
||||
|
||||
assert_eq!(pending.len(), 1);
|
||||
assert!(pending.is_empty());
|
||||
assert_eq!(skipped_count(&report), 1);
|
||||
assert!(has_diagnostic(&report, "duplicate tool contribution"));
|
||||
}
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use agen::{HistoryEntry, UsageRecord};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::session_history::SessionHistoryMetadata;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum CommittedRunExit {
|
||||
Finished,
|
||||
NonFinal,
|
||||
Interrupted,
|
||||
}
|
||||
|
||||
/// Immutable projection of one durably committed session-log location.
|
||||
///
|
||||
/// Feature code receives this value only after the host has committed the
|
||||
/// terminal run record. The projection deliberately carries annotated history
|
||||
/// rather than the public flattened transcript so provenance-sensitive
|
||||
/// features can construct their own bounded views.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CommittedSessionCapture {
|
||||
pub(crate) session_id: String,
|
||||
pub(crate) segment_id: String,
|
||||
/// Monotonic committed-log revision for the captured Segment.
|
||||
pub(crate) session_revision: u64,
|
||||
pub(crate) entry_count: usize,
|
||||
pub(crate) run_exit: CommittedRunExit,
|
||||
pub(crate) history: Vec<HistoryEntry<SessionHistoryMetadata>>,
|
||||
pub(crate) usage_history: Vec<UsageRecord>,
|
||||
pub(crate) extensions: Vec<(String, Value)>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(crate) struct CommittedSessionLocation {
|
||||
pub(crate) session_id: String,
|
||||
pub(crate) segment_id: String,
|
||||
/// Monotonic committed-log revision for the captured Segment.
|
||||
pub(crate) session_revision: u64,
|
||||
pub(crate) entry_count: usize,
|
||||
}
|
||||
|
||||
impl CommittedSessionCapture {
|
||||
pub(crate) fn location(&self) -> CommittedSessionLocation {
|
||||
CommittedSessionLocation {
|
||||
session_id: self.session_id.clone(),
|
||||
segment_id: self.segment_id.clone(),
|
||||
session_revision: self.session_revision,
|
||||
entry_count: self.entry_count,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(crate) enum FeatureSessionError {
|
||||
#[error("read committed session failed: {0}")]
|
||||
Capture(String),
|
||||
#[error("append session extension failed: {0}")]
|
||||
Extension(String),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CommittedSessionCaptureHandle {
|
||||
capture: Arc<
|
||||
dyn Fn() -> Result<CommittedSessionCapture, FeatureSessionError> + Send + Sync + 'static,
|
||||
>,
|
||||
}
|
||||
|
||||
impl CommittedSessionCaptureHandle {
|
||||
pub(crate) fn new(
|
||||
capture: impl Fn() -> Result<CommittedSessionCapture, FeatureSessionError>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
) -> Self {
|
||||
Self {
|
||||
capture: Arc::new(capture),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn capture(&self) -> Result<CommittedSessionCapture, FeatureSessionError> {
|
||||
(self.capture)()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct SessionExtensionHandle {
|
||||
append: Arc<
|
||||
dyn Fn(&CommittedSessionLocation, &str, Value) -> Result<bool, FeatureSessionError>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
>,
|
||||
}
|
||||
|
||||
impl SessionExtensionHandle {
|
||||
pub(crate) fn new(
|
||||
append: impl Fn(&CommittedSessionLocation, &str, Value) -> Result<bool, FeatureSessionError>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
) -> Self {
|
||||
Self {
|
||||
append: Arc::new(append),
|
||||
}
|
||||
}
|
||||
|
||||
/// Appends an extension only while the committed session is still at the
|
||||
/// exact location captured by the feature. `Ok(false)` is a stale-write
|
||||
/// fence, not an I/O failure.
|
||||
pub(crate) fn append_if_current(
|
||||
&self,
|
||||
expected: &CommittedSessionLocation,
|
||||
domain: &str,
|
||||
payload: Value,
|
||||
) -> Result<bool, FeatureSessionError> {
|
||||
(self.append)(expected, domain, payload)
|
||||
}
|
||||
}
|
||||
+645
-52
@@ -23,8 +23,105 @@ use agen::interceptor::{
|
||||
};
|
||||
use agen::tool::{ToolOutput, ToolResult};
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use session_store::{SystemItem, SystemReminder};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::SessionEntryRef;
|
||||
|
||||
const HOOK_DIAGNOSTIC_MAX_BYTES: usize = 1_024;
|
||||
|
||||
/// Failure category exposed by the safe Worker hook boundary.
|
||||
///
|
||||
/// Categories are intentionally closed and payload-free so extensions cannot
|
||||
/// smuggle provider, credential, or history data into diagnostics.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum HookErrorCategory {
|
||||
InvalidInput,
|
||||
Dependency,
|
||||
Timeout,
|
||||
Cancelled,
|
||||
Trap,
|
||||
ScopeDisposed,
|
||||
Internal,
|
||||
}
|
||||
|
||||
/// Bounded hook callback failure. Raw tool arguments, output, prompts, and
|
||||
/// credentials must never be placed in `diagnostic`.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Error)]
|
||||
#[error("{category:?}: {diagnostic}")]
|
||||
pub struct HookError {
|
||||
pub category: HookErrorCategory,
|
||||
pub diagnostic: String,
|
||||
}
|
||||
|
||||
impl HookError {
|
||||
pub fn new(category: HookErrorCategory, diagnostic: impl Into<String>) -> Self {
|
||||
Self {
|
||||
category,
|
||||
diagnostic: bounded_utf8(diagnostic.into(), HOOK_DIAGNOSTIC_MAX_BYTES),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Failure behavior declared when a hook is registered.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum HookFailurePolicy {
|
||||
/// Gate the current operation when the hook cannot decide safely.
|
||||
FailClosed,
|
||||
/// Keep the already-authorized operation moving and emit a diagnostic.
|
||||
FailOpenWithDiagnostic,
|
||||
/// Keep committed state intact and mark the failure for operator attention.
|
||||
AttentionRequired,
|
||||
}
|
||||
|
||||
const DEFAULT_HOOK_TIMEOUT_MS: u64 = 30_000;
|
||||
const MAX_HOOK_TIMEOUT_MS: u64 = 120_000;
|
||||
|
||||
/// Host-enforced execution and failure budget for one hook registration.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct HookExecutionPolicy {
|
||||
pub failure: HookFailurePolicy,
|
||||
pub timeout_ms: u64,
|
||||
}
|
||||
|
||||
impl HookExecutionPolicy {
|
||||
pub const fn new(failure: HookFailurePolicy, timeout_ms: u64) -> Self {
|
||||
Self {
|
||||
failure,
|
||||
timeout_ms,
|
||||
}
|
||||
}
|
||||
|
||||
pub const fn fail_closed() -> Self {
|
||||
Self::new(HookFailurePolicy::FailClosed, DEFAULT_HOOK_TIMEOUT_MS)
|
||||
}
|
||||
|
||||
fn validate(self) -> Result<Self, HookError> {
|
||||
if self.timeout_ms == 0 || self.timeout_ms > MAX_HOOK_TIMEOUT_MS {
|
||||
return Err(HookError::new(
|
||||
HookErrorCategory::InvalidInput,
|
||||
format!("hook timeout_ms must be within 1..={MAX_HOOK_TIMEOUT_MS}"),
|
||||
));
|
||||
}
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
fn bounded_utf8(mut value: String, max_bytes: usize) -> String {
|
||||
if value.len() <= max_bytes {
|
||||
return value;
|
||||
}
|
||||
let mut end = max_bytes;
|
||||
while end > 0 && !value.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
value.truncate(end);
|
||||
value
|
||||
}
|
||||
|
||||
/// Hook-facing prompt-submit action.
|
||||
///
|
||||
@@ -285,12 +382,6 @@ pub struct TurnEndInfo {
|
||||
pub final_text_preview: String,
|
||||
}
|
||||
|
||||
/// Information passed to `OnAbort` hooks.
|
||||
pub struct AbortInfo {
|
||||
/// Reason supplied by the aborter.
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Hook Event Kinds
|
||||
// =============================================================================
|
||||
@@ -315,10 +406,8 @@ pub struct PreLlmRequest;
|
||||
pub struct PreToolCall;
|
||||
/// After each tool completes; observational except it may abort the run.
|
||||
pub struct PostToolCall;
|
||||
/// When a turn ends with no tool calls; observational except it may pause.
|
||||
/// After every terminal assistant response is committed; observational except it may pause.
|
||||
pub struct OnTurnEnd;
|
||||
/// When execution is interrupted; observational only.
|
||||
pub struct OnAbort;
|
||||
|
||||
impl HookEventKind for OnPromptSubmit {
|
||||
type Input = PromptSubmitInfo;
|
||||
@@ -345,11 +434,6 @@ impl HookEventKind for OnTurnEnd {
|
||||
type Output = HookTurnEndAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for OnAbort {
|
||||
type Input = AbortInfo;
|
||||
type Output = ();
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Hook Trait
|
||||
// =============================================================================
|
||||
@@ -362,25 +446,194 @@ impl HookEventKind for OnAbort {
|
||||
/// short-circuit on the first non-continue action.
|
||||
#[async_trait]
|
||||
pub trait Hook<E: HookEventKind>: Send + Sync {
|
||||
async fn call(&self, input: &E::Input) -> E::Output;
|
||||
async fn call(&self, input: &E::Input) -> Result<E::Output, HookError>;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Hook Registry
|
||||
// =============================================================================
|
||||
|
||||
/// Bounded identity-only view of the committed history visible at a lifecycle
|
||||
/// boundary. Hook context never carries history payloads; a Feature that needs
|
||||
/// content must use a separately granted session-exploration service.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct HookHistoryRange {
|
||||
pub first_entry_ref: Option<SessionEntryRef>,
|
||||
pub last_entry_ref: Option<SessionEntryRef>,
|
||||
pub entry_count: usize,
|
||||
}
|
||||
|
||||
/// Stable provenance attached to every Worker lifecycle callback.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct HookInvocationContext {
|
||||
pub workspace_id: Option<String>,
|
||||
pub worker_id: String,
|
||||
pub session_id: String,
|
||||
pub session_revision: u64,
|
||||
pub run_id: Option<String>,
|
||||
pub turn_index: Option<usize>,
|
||||
pub call_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum RunCommittedExit {
|
||||
Finished,
|
||||
Paused,
|
||||
Yielded,
|
||||
Interrupted,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RunExitContext {
|
||||
pub invocation: HookInvocationContext,
|
||||
pub exit: RunCommittedExit,
|
||||
pub history_len: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct RunCommittedContext {
|
||||
pub invocation: HookInvocationContext,
|
||||
pub exit: RunCommittedExit,
|
||||
pub committed_history: HookHistoryRange,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum SessionRewriteKind {
|
||||
Rewind,
|
||||
Compact,
|
||||
Fork,
|
||||
Restore,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct BeforeSessionRewriteContext {
|
||||
pub invocation: HookInvocationContext,
|
||||
pub kind: SessionRewriteKind,
|
||||
pub current_history: HookHistoryRange,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum BeforeSessionRewriteAction {
|
||||
Continue,
|
||||
Deny(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct WorkerStoppingContext {
|
||||
pub invocation: HookInvocationContext,
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
pub struct RunExit;
|
||||
pub struct RunCommitted;
|
||||
pub struct BeforeSessionRewrite;
|
||||
pub struct WorkerStopping;
|
||||
|
||||
impl HookEventKind for RunExit {
|
||||
type Input = RunExitContext;
|
||||
type Output = ();
|
||||
}
|
||||
|
||||
impl HookEventKind for RunCommitted {
|
||||
type Input = RunCommittedContext;
|
||||
type Output = ();
|
||||
}
|
||||
|
||||
impl HookEventKind for BeforeSessionRewrite {
|
||||
type Input = BeforeSessionRewriteContext;
|
||||
type Output = BeforeSessionRewriteAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for WorkerStopping {
|
||||
type Input = WorkerStoppingContext;
|
||||
type Output = ();
|
||||
}
|
||||
|
||||
pub(crate) struct RegisteredHook<E: HookEventKind> {
|
||||
owner: String,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: Box<dyn Hook<E>>,
|
||||
}
|
||||
|
||||
impl<E: HookEventKind> RegisteredHook<E> {
|
||||
pub(crate) async fn call(&self, input: &E::Input) -> Result<E::Output, HookExecutionError> {
|
||||
let result = tokio::time::timeout(
|
||||
std::time::Duration::from_millis(self.policy.timeout_ms),
|
||||
self.hook.call(input),
|
||||
)
|
||||
.await
|
||||
.unwrap_or_else(|_| {
|
||||
Err(HookError::new(
|
||||
HookErrorCategory::Timeout,
|
||||
"hook exceeded its host-enforced execution deadline",
|
||||
))
|
||||
});
|
||||
result.map_err(|source| HookExecutionError {
|
||||
owner: self.owner.clone(),
|
||||
policy: self.policy.failure,
|
||||
source,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn call_optional(
|
||||
&self,
|
||||
input: &E::Input,
|
||||
) -> Result<Option<E::Output>, HookExecutionError> {
|
||||
match self.call(input).await {
|
||||
Ok(output) => Ok(Some(output)),
|
||||
Err(error) if error.policy == HookFailurePolicy::FailOpenWithDiagnostic => {
|
||||
tracing::warn!(owner = %error.owner, error = %error.source, "inline hook failed open");
|
||||
Ok(None)
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Error)]
|
||||
#[error("hook `{owner}` failed under {policy:?}: {source}")]
|
||||
pub struct HookExecutionError {
|
||||
pub owner: String,
|
||||
pub policy: HookFailurePolicy,
|
||||
pub source: HookError,
|
||||
}
|
||||
|
||||
/// Builder for constructing a frozen `HookRegistry`.
|
||||
///
|
||||
/// Hooks are added during setup, then `build()` produces an immutable
|
||||
/// registry that can be shared via `Arc`.
|
||||
#[derive(Default)]
|
||||
pub struct HookRegistryBuilder {
|
||||
on_prompt_submit: Vec<Box<dyn Hook<OnPromptSubmit>>>,
|
||||
pre_llm_request: Vec<Box<dyn Hook<PreLlmRequest>>>,
|
||||
pre_tool_call: Vec<Box<dyn Hook<PreToolCall>>>,
|
||||
post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
|
||||
on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
|
||||
on_abort: Vec<Box<dyn Hook<OnAbort>>>,
|
||||
on_prompt_submit: Vec<RegisteredHook<OnPromptSubmit>>,
|
||||
pre_llm_request: Vec<RegisteredHook<PreLlmRequest>>,
|
||||
pre_tool_call: Vec<RegisteredHook<PreToolCall>>,
|
||||
post_tool_call: Vec<RegisteredHook<PostToolCall>>,
|
||||
on_turn_end: Vec<RegisteredHook<OnTurnEnd>>,
|
||||
run_exit: Vec<RegisteredHook<RunExit>>,
|
||||
run_committed: Vec<RegisteredHook<RunCommitted>>,
|
||||
before_session_rewrite: Vec<RegisteredHook<BeforeSessionRewrite>>,
|
||||
worker_stopping: Vec<RegisteredHook<WorkerStopping>>,
|
||||
}
|
||||
|
||||
macro_rules! add_hook_methods {
|
||||
($default:ident, $named:ident, $field:ident, $event:ty) => {
|
||||
pub fn $default(&mut self, hook: impl Hook<$event> + 'static) {
|
||||
self.$named("worker.host", HookExecutionPolicy::fail_closed(), hook)
|
||||
.expect("host hook policy is valid");
|
||||
}
|
||||
|
||||
pub fn $named(
|
||||
&mut self,
|
||||
owner: impl Into<String>,
|
||||
policy: HookExecutionPolicy,
|
||||
hook: impl Hook<$event> + 'static,
|
||||
) -> Result<(), HookError> {
|
||||
let policy = policy.validate()?;
|
||||
self.$field.push(RegisteredHook {
|
||||
owner: owner.into(),
|
||||
policy,
|
||||
hook: Box::new(hook),
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
impl HookRegistryBuilder {
|
||||
@@ -388,31 +641,82 @@ impl HookRegistryBuilder {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn add_on_prompt_submit(&mut self, hook: impl Hook<OnPromptSubmit> + 'static) {
|
||||
self.on_prompt_submit.push(Box::new(hook));
|
||||
add_hook_methods!(
|
||||
add_on_prompt_submit,
|
||||
add_named_on_prompt_submit,
|
||||
on_prompt_submit,
|
||||
OnPromptSubmit
|
||||
);
|
||||
add_hook_methods!(
|
||||
add_pre_llm_request,
|
||||
add_named_pre_llm_request,
|
||||
pre_llm_request,
|
||||
PreLlmRequest
|
||||
);
|
||||
add_hook_methods!(
|
||||
add_pre_tool_call,
|
||||
add_named_pre_tool_call,
|
||||
pre_tool_call,
|
||||
PreToolCall
|
||||
);
|
||||
add_hook_methods!(
|
||||
add_post_tool_call,
|
||||
add_named_post_tool_call,
|
||||
post_tool_call,
|
||||
PostToolCall
|
||||
);
|
||||
add_hook_methods!(
|
||||
add_on_turn_end,
|
||||
add_named_on_turn_end,
|
||||
on_turn_end,
|
||||
OnTurnEnd
|
||||
);
|
||||
add_hook_methods!(add_run_exit, add_named_run_exit, run_exit, RunExit);
|
||||
add_hook_methods!(
|
||||
add_run_committed,
|
||||
add_named_run_committed,
|
||||
run_committed,
|
||||
RunCommitted
|
||||
);
|
||||
add_hook_methods!(
|
||||
add_before_session_rewrite,
|
||||
add_named_before_session_rewrite,
|
||||
before_session_rewrite,
|
||||
BeforeSessionRewrite
|
||||
);
|
||||
add_hook_methods!(
|
||||
add_worker_stopping,
|
||||
add_named_worker_stopping,
|
||||
worker_stopping,
|
||||
WorkerStopping
|
||||
);
|
||||
|
||||
pub(crate) fn checkpoint(&self) -> [usize; 9] {
|
||||
[
|
||||
self.on_prompt_submit.len(),
|
||||
self.pre_llm_request.len(),
|
||||
self.pre_tool_call.len(),
|
||||
self.post_tool_call.len(),
|
||||
self.on_turn_end.len(),
|
||||
self.run_exit.len(),
|
||||
self.run_committed.len(),
|
||||
self.before_session_rewrite.len(),
|
||||
self.worker_stopping.len(),
|
||||
]
|
||||
}
|
||||
|
||||
pub fn add_pre_llm_request(&mut self, hook: impl Hook<PreLlmRequest> + 'static) {
|
||||
self.pre_llm_request.push(Box::new(hook));
|
||||
pub(crate) fn rollback_to(&mut self, checkpoint: [usize; 9]) {
|
||||
self.on_prompt_submit.truncate(checkpoint[0]);
|
||||
self.pre_llm_request.truncate(checkpoint[1]);
|
||||
self.pre_tool_call.truncate(checkpoint[2]);
|
||||
self.post_tool_call.truncate(checkpoint[3]);
|
||||
self.on_turn_end.truncate(checkpoint[4]);
|
||||
self.run_exit.truncate(checkpoint[5]);
|
||||
self.run_committed.truncate(checkpoint[6]);
|
||||
self.before_session_rewrite.truncate(checkpoint[7]);
|
||||
self.worker_stopping.truncate(checkpoint[8]);
|
||||
}
|
||||
|
||||
pub fn add_pre_tool_call(&mut self, hook: impl Hook<PreToolCall> + 'static) {
|
||||
self.pre_tool_call.push(Box::new(hook));
|
||||
}
|
||||
|
||||
pub fn add_post_tool_call(&mut self, hook: impl Hook<PostToolCall> + 'static) {
|
||||
self.post_tool_call.push(Box::new(hook));
|
||||
}
|
||||
|
||||
pub fn add_on_turn_end(&mut self, hook: impl Hook<OnTurnEnd> + 'static) {
|
||||
self.on_turn_end.push(Box::new(hook));
|
||||
}
|
||||
|
||||
pub fn add_on_abort(&mut self, hook: impl Hook<OnAbort> + 'static) {
|
||||
self.on_abort.push(Box::new(hook));
|
||||
}
|
||||
|
||||
/// Freeze the builder into an immutable registry.
|
||||
pub fn build(self) -> HookRegistry {
|
||||
HookRegistry {
|
||||
on_prompt_submit: self.on_prompt_submit,
|
||||
@@ -420,19 +724,140 @@ impl HookRegistryBuilder {
|
||||
pre_tool_call: self.pre_tool_call,
|
||||
post_tool_call: self.post_tool_call,
|
||||
on_turn_end: self.on_turn_end,
|
||||
on_abort: self.on_abort,
|
||||
run_exit: self.run_exit,
|
||||
run_committed: self.run_committed,
|
||||
before_session_rewrite: self.before_session_rewrite,
|
||||
worker_stopping: self.worker_stopping,
|
||||
diagnostics: std::sync::Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Frozen registry of hooks. Constructed via [`HookRegistryBuilder::build()`].
|
||||
pub struct HookRegistry {
|
||||
pub(crate) on_prompt_submit: Vec<Box<dyn Hook<OnPromptSubmit>>>,
|
||||
pub(crate) pre_llm_request: Vec<Box<dyn Hook<PreLlmRequest>>>,
|
||||
pub(crate) pre_tool_call: Vec<Box<dyn Hook<PreToolCall>>>,
|
||||
pub(crate) post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
|
||||
pub(crate) on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
|
||||
pub(crate) on_abort: Vec<Box<dyn Hook<OnAbort>>>,
|
||||
pub(crate) on_prompt_submit: Vec<RegisteredHook<OnPromptSubmit>>,
|
||||
pub(crate) pre_llm_request: Vec<RegisteredHook<PreLlmRequest>>,
|
||||
pub(crate) pre_tool_call: Vec<RegisteredHook<PreToolCall>>,
|
||||
pub(crate) post_tool_call: Vec<RegisteredHook<PostToolCall>>,
|
||||
pub(crate) on_turn_end: Vec<RegisteredHook<OnTurnEnd>>,
|
||||
run_exit: Vec<RegisteredHook<RunExit>>,
|
||||
run_committed: Vec<RegisteredHook<RunCommitted>>,
|
||||
before_session_rewrite: Vec<RegisteredHook<BeforeSessionRewrite>>,
|
||||
worker_stopping: Vec<RegisteredHook<WorkerStopping>>,
|
||||
diagnostics: std::sync::Mutex<Vec<HookExecutionError>>,
|
||||
}
|
||||
|
||||
impl HookRegistry {
|
||||
fn record_diagnostic(&self, error: HookExecutionError) {
|
||||
let mut diagnostics = self.diagnostics.lock().expect("hook diagnostics poisoned");
|
||||
diagnostics.push(error);
|
||||
if diagnostics.len() > 128 {
|
||||
let remove = diagnostics.len() - 128;
|
||||
diagnostics.drain(..remove);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn record_chain_timeout(&self, lifecycle: &str) {
|
||||
self.record_diagnostic(HookExecutionError {
|
||||
owner: format!("worker.{lifecycle}"),
|
||||
policy: HookFailurePolicy::AttentionRequired,
|
||||
source: HookError::new(
|
||||
HookErrorCategory::Timeout,
|
||||
"hook chain exceeded its host-enforced lifecycle deadline",
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
pub fn diagnostics(&self) -> Vec<HookExecutionError> {
|
||||
self.diagnostics
|
||||
.lock()
|
||||
.expect("hook diagnostics poisoned")
|
||||
.clone()
|
||||
}
|
||||
pub async fn on_run_exit(&self, context: &RunExitContext) -> Result<(), HookExecutionError> {
|
||||
for registration in &self.run_exit {
|
||||
if let Err(error) = registration.call(context).await {
|
||||
self.record_diagnostic(error.clone());
|
||||
match error.policy {
|
||||
HookFailurePolicy::FailOpenWithDiagnostic => {
|
||||
tracing::warn!(owner = %error.owner, error = %error.source, "run-exit hook failed open");
|
||||
}
|
||||
HookFailurePolicy::FailClosed | HookFailurePolicy::AttentionRequired => {
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn on_run_committed(
|
||||
&self,
|
||||
context: &RunCommittedContext,
|
||||
) -> Result<(), HookExecutionError> {
|
||||
for registration in &self.run_committed {
|
||||
if let Err(error) = registration.call(context).await {
|
||||
self.record_diagnostic(error.clone());
|
||||
match error.policy {
|
||||
HookFailurePolicy::FailOpenWithDiagnostic => {
|
||||
tracing::warn!(owner = %error.owner, error = %error.source, "run-committed hook failed open");
|
||||
}
|
||||
HookFailurePolicy::FailClosed | HookFailurePolicy::AttentionRequired => {
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn before_session_rewrite(
|
||||
&self,
|
||||
context: &BeforeSessionRewriteContext,
|
||||
) -> Result<BeforeSessionRewriteAction, HookExecutionError> {
|
||||
let mut denials = Vec::new();
|
||||
for registration in &self.before_session_rewrite {
|
||||
match registration.call(context).await {
|
||||
Ok(BeforeSessionRewriteAction::Continue) => {}
|
||||
Ok(BeforeSessionRewriteAction::Deny(reason)) => {
|
||||
denials.push((registration.owner.clone(), reason));
|
||||
}
|
||||
Err(error) if error.policy == HookFailurePolicy::FailOpenWithDiagnostic => {
|
||||
self.record_diagnostic(error.clone());
|
||||
tracing::warn!(owner = %error.owner, error = %error.source, "session-rewrite hook failed open");
|
||||
}
|
||||
Err(error) => {
|
||||
self.record_diagnostic(error.clone());
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
denials.sort_by(|left, right| left.0.cmp(&right.0));
|
||||
Ok(denials
|
||||
.into_iter()
|
||||
.next()
|
||||
.map(|(_, reason)| BeforeSessionRewriteAction::Deny(reason))
|
||||
.unwrap_or(BeforeSessionRewriteAction::Continue))
|
||||
}
|
||||
|
||||
pub async fn on_worker_stopping(
|
||||
&self,
|
||||
context: &WorkerStoppingContext,
|
||||
) -> Result<(), HookExecutionError> {
|
||||
for registration in &self.worker_stopping {
|
||||
if let Err(error) = registration.call(context).await {
|
||||
self.record_diagnostic(error.clone());
|
||||
match error.policy {
|
||||
HookFailurePolicy::FailOpenWithDiagnostic
|
||||
| HookFailurePolicy::AttentionRequired => {
|
||||
tracing::warn!(owner = %error.owner, error = %error.source, "worker-stopping hook requires attention");
|
||||
}
|
||||
HookFailurePolicy::FailClosed => return Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -494,4 +919,172 @@ mod tests {
|
||||
let pause_action = HookPreToolAction::Pause.into_worker_action("call_4".into());
|
||||
assert!(matches!(pause_action, PreToolAction::Pause));
|
||||
}
|
||||
|
||||
struct RewriteHook {
|
||||
action: BeforeSessionRewriteAction,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<BeforeSessionRewrite> for RewriteHook {
|
||||
async fn call(
|
||||
&self,
|
||||
_input: &BeforeSessionRewriteContext,
|
||||
) -> Result<BeforeSessionRewriteAction, HookError> {
|
||||
Ok(self.action.clone())
|
||||
}
|
||||
}
|
||||
|
||||
struct FailingRewriteHook;
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<BeforeSessionRewrite> for FailingRewriteHook {
|
||||
async fn call(
|
||||
&self,
|
||||
_input: &BeforeSessionRewriteContext,
|
||||
) -> Result<BeforeSessionRewriteAction, HookError> {
|
||||
Err(HookError::new(
|
||||
HookErrorCategory::Dependency,
|
||||
"provider unavailable",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn rewrite_context() -> BeforeSessionRewriteContext {
|
||||
BeforeSessionRewriteContext {
|
||||
invocation: HookInvocationContext {
|
||||
workspace_id: Some("workspace".into()),
|
||||
worker_id: "worker".into(),
|
||||
session_id: "session".into(),
|
||||
session_revision: 4,
|
||||
run_id: None,
|
||||
turn_index: None,
|
||||
call_id: None,
|
||||
},
|
||||
kind: SessionRewriteKind::Compact,
|
||||
current_history: HookHistoryRange::default(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rewrite_denials_are_resolved_by_owner_not_registration_order() {
|
||||
let mut builder = HookRegistryBuilder::new();
|
||||
builder
|
||||
.add_named_before_session_rewrite(
|
||||
"z-feature",
|
||||
HookExecutionPolicy::fail_closed(),
|
||||
RewriteHook {
|
||||
action: BeforeSessionRewriteAction::Deny("z denied".into()),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
builder
|
||||
.add_named_before_session_rewrite(
|
||||
"a-feature",
|
||||
HookExecutionPolicy::fail_closed(),
|
||||
RewriteHook {
|
||||
action: BeforeSessionRewriteAction::Deny("a denied".into()),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
builder
|
||||
.build()
|
||||
.before_session_rewrite(&rewrite_context())
|
||||
.await
|
||||
.unwrap(),
|
||||
BeforeSessionRewriteAction::Deny("a denied".into())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hook_failure_policy_is_applied_at_the_registry_boundary() {
|
||||
let mut fail_open = HookRegistryBuilder::new();
|
||||
fail_open
|
||||
.add_named_before_session_rewrite(
|
||||
"feature",
|
||||
HookExecutionPolicy::new(HookFailurePolicy::FailOpenWithDiagnostic, 30_000),
|
||||
FailingRewriteHook,
|
||||
)
|
||||
.unwrap();
|
||||
let fail_open = fail_open.build();
|
||||
assert_eq!(
|
||||
fail_open
|
||||
.before_session_rewrite(&rewrite_context())
|
||||
.await
|
||||
.unwrap(),
|
||||
BeforeSessionRewriteAction::Continue
|
||||
);
|
||||
assert_eq!(fail_open.diagnostics().len(), 1);
|
||||
|
||||
let mut fail_closed = HookRegistryBuilder::new();
|
||||
fail_closed
|
||||
.add_named_before_session_rewrite(
|
||||
"feature",
|
||||
HookExecutionPolicy::fail_closed(),
|
||||
FailingRewriteHook,
|
||||
)
|
||||
.unwrap();
|
||||
let error = fail_closed
|
||||
.build()
|
||||
.before_session_rewrite(&rewrite_context())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(error.source.category, HookErrorCategory::Dependency);
|
||||
}
|
||||
|
||||
struct NeverReturns;
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<BeforeSessionRewrite> for NeverReturns {
|
||||
async fn call(
|
||||
&self,
|
||||
_input: &BeforeSessionRewriteContext,
|
||||
) -> Result<BeforeSessionRewriteAction, HookError> {
|
||||
std::future::pending().await
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn hook_execution_deadline_cancels_a_non_cooperative_callback() {
|
||||
let mut builder = HookRegistryBuilder::new();
|
||||
builder
|
||||
.add_named_before_session_rewrite(
|
||||
"feature",
|
||||
HookExecutionPolicy::new(HookFailurePolicy::FailClosed, 10),
|
||||
NeverReturns,
|
||||
)
|
||||
.unwrap();
|
||||
let registry = builder.build();
|
||||
|
||||
let error = tokio::time::timeout(
|
||||
std::time::Duration::from_millis(100),
|
||||
registry.before_session_rewrite(&rewrite_context()),
|
||||
)
|
||||
.await
|
||||
.expect("host hook deadline must terminate the callback")
|
||||
.unwrap_err();
|
||||
assert_eq!(error.source.category, HookErrorCategory::Timeout);
|
||||
assert_eq!(registry.diagnostics().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_hook_execution_budget_is_rejected_before_registration() {
|
||||
let mut builder = HookRegistryBuilder::new();
|
||||
let error = builder
|
||||
.add_named_before_session_rewrite(
|
||||
"feature",
|
||||
HookExecutionPolicy::new(HookFailurePolicy::FailClosed, 0),
|
||||
NeverReturns,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert_eq!(error.category, HookErrorCategory::InvalidInput);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hook_diagnostics_are_utf8_bounded() {
|
||||
let error = HookError::new(HookErrorCategory::Internal, "界".repeat(1_000));
|
||||
assert!(error.diagnostic.len() <= HOOK_DIAGNOSTIC_MAX_BYTES);
|
||||
assert!(error.diagnostic.is_char_boundary(error.diagnostic.len()));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -123,14 +123,14 @@ where
|
||||
|
||||
// Internal identities are run-scoped and never enter the public Runtime Worker catalog.
|
||||
manifest.worker.name = format!("internal-{}-{}", identity.kind, identity.run_id);
|
||||
// Internal jobs only receive features supplied below. A parent manifest must not accidentally
|
||||
// grant its normal public tool surface or recursively schedule memory work.
|
||||
// Internal jobs only receive the explicitly supplied Feature set below. A
|
||||
// parent manifest cannot accidentally grant its normal public tool surface
|
||||
// or recursively schedule Feature-owned background work.
|
||||
manifest.feature = Default::default();
|
||||
manifest.plugins = Default::default();
|
||||
manifest.mcp = Default::default();
|
||||
manifest.skills = None;
|
||||
manifest.compaction = None;
|
||||
manifest.memory = None;
|
||||
|
||||
let last_usage = Arc::new(Mutex::new(None::<UsageEvent>));
|
||||
let usage_slot = last_usage.clone();
|
||||
@@ -164,6 +164,7 @@ where
|
||||
identity: identity.clone(),
|
||||
history_entries: 0,
|
||||
})?;
|
||||
worker.disable_manifest_lifecycle_features();
|
||||
if let Some(session) = inherited_workdir_session {
|
||||
worker.bind_workdir_session(Some(session));
|
||||
}
|
||||
@@ -210,7 +211,7 @@ where
|
||||
let segment_id = worker.segment_id();
|
||||
on_cancel_sender(worker.engine_mut().cancel_sender());
|
||||
|
||||
match worker.run_text(&input).await {
|
||||
let outcome = match worker.run_text(&input).await {
|
||||
Ok(lifecycle @ WorkerRunResult::Finished)
|
||||
| Ok(lifecycle @ WorkerRunResult::Paused)
|
||||
| Ok(lifecycle @ WorkerRunResult::RolledBack) => Ok(InternalWorkerResult {
|
||||
@@ -239,7 +240,11 @@ where
|
||||
identity,
|
||||
history_entries: store.entries_count(session_id, segment_id),
|
||||
}),
|
||||
}
|
||||
};
|
||||
worker
|
||||
.stop_feature_runtime("internal Worker terminal outcome")
|
||||
.await;
|
||||
outcome
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -544,7 +549,6 @@ pub(crate) async fn spawn_internal_worker_session(
|
||||
authority,
|
||||
} = spec;
|
||||
manifest.worker.name = format!("internal-{}-{}", identity.kind, identity.run_id);
|
||||
manifest.memory = None;
|
||||
|
||||
let last_usage = Arc::new(Mutex::new(None::<UsageEvent>));
|
||||
let usage_slot = last_usage.clone();
|
||||
@@ -575,6 +579,7 @@ pub(crate) async fn spawn_internal_worker_session(
|
||||
.map_err(|source| InternalWorkerSessionError::Build {
|
||||
message: source.to_string(),
|
||||
})?;
|
||||
worker.disable_manifest_lifecycle_features();
|
||||
if let Some(session) = inherited_workdir_session {
|
||||
worker.bind_workdir_session(Some(session));
|
||||
}
|
||||
@@ -645,7 +650,6 @@ pub(crate) fn prepare_internal_worker_from_spec(
|
||||
manifest.mcp = Default::default();
|
||||
manifest.skills = None;
|
||||
manifest.compaction = None;
|
||||
manifest.memory = None;
|
||||
|
||||
let mut engine =
|
||||
Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated(client)
|
||||
@@ -669,6 +673,7 @@ pub(crate) fn prepare_internal_worker_from_spec(
|
||||
.map_err(|source| InternalWorkerSessionError::Build {
|
||||
message: source.to_string(),
|
||||
})?;
|
||||
worker.disable_manifest_lifecycle_features();
|
||||
if let Some(session) = inherited_workdir_session {
|
||||
worker.bind_workdir_session(Some(session));
|
||||
}
|
||||
@@ -782,7 +787,8 @@ pub(crate) async fn prepare_internal_worker_session(
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
while let Some(command) = command_rx.recv().await {
|
||||
let mut stop_done = None;
|
||||
'actor: while let Some(command) = command_rx.recv().await {
|
||||
match command {
|
||||
InternalWorkerSessionCommand::Run(input) => {
|
||||
actor_in_flight.clear();
|
||||
@@ -825,21 +831,16 @@ pub(crate) async fn prepare_internal_worker_session(
|
||||
Some(InternalWorkerSessionCommand::Stop(done)) => {
|
||||
let _ = cancel_sender.send(()).await;
|
||||
let _ = (&mut run).await;
|
||||
actor_in_flight.clear();
|
||||
status.store(InternalWorkerSessionStatus::Stopped.encode(), std::sync::atomic::Ordering::Release);
|
||||
let _ = event_tx.send(Event::Status { status: WorkerStatus::Stopped });
|
||||
let _ = event_tx.send(Event::Shutdown);
|
||||
state_changed.notify_waiters();
|
||||
let _ = done.send(());
|
||||
return;
|
||||
stop_done = Some(done);
|
||||
break 'actor;
|
||||
}
|
||||
Some(InternalWorkerSessionCommand::Run(_)) => {
|
||||
// `send` reserves Running atomically, so a second Run cannot be enqueued.
|
||||
}
|
||||
None => {
|
||||
let _ = cancel_sender.send(()).await;
|
||||
actor_in_flight.clear();
|
||||
return;
|
||||
let _ = (&mut run).await;
|
||||
break 'actor;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -847,22 +848,27 @@ pub(crate) async fn prepare_internal_worker_session(
|
||||
}
|
||||
}
|
||||
InternalWorkerSessionCommand::Stop(done) => {
|
||||
actor_in_flight.clear();
|
||||
status.store(
|
||||
InternalWorkerSessionStatus::Stopped.encode(),
|
||||
std::sync::atomic::Ordering::Release,
|
||||
);
|
||||
let _ = event_tx.send(Event::Status {
|
||||
status: WorkerStatus::Stopped,
|
||||
});
|
||||
let _ = event_tx.send(Event::Shutdown);
|
||||
state_changed.notify_waiters();
|
||||
let _ = done.send(());
|
||||
return;
|
||||
stop_done = Some(done);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
worker
|
||||
.stop_feature_runtime("internal Worker session stopped")
|
||||
.await;
|
||||
actor_in_flight.clear();
|
||||
status.store(
|
||||
InternalWorkerSessionStatus::Stopped.encode(),
|
||||
std::sync::atomic::Ordering::Release,
|
||||
);
|
||||
let _ = event_tx.send(Event::Status {
|
||||
status: WorkerStatus::Stopped,
|
||||
});
|
||||
let _ = event_tx.send(Event::Shutdown);
|
||||
state_changed.notify_waiters();
|
||||
if let Some(done) = stop_done {
|
||||
let _ = done.send(());
|
||||
}
|
||||
});
|
||||
|
||||
Ok(handle)
|
||||
|
||||
@@ -15,7 +15,9 @@ use std::sync::{Arc, Mutex};
|
||||
use agen::Item;
|
||||
use agen::UsageRecord;
|
||||
use agen::interceptor::{
|
||||
Interceptor, PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo,
|
||||
AssistantTurnEndContext, Interceptor, InterceptorError, InterceptorErrorCategory,
|
||||
InterceptorResult, PendingHistoryAppendsContext, PostToolAction, PreLlmRequestContext,
|
||||
PreRequestAction, PreToolAction, PromptAction, PromptSubmitContext, ToolCallInfo,
|
||||
ToolResultInfo, TurnEndAction,
|
||||
};
|
||||
use agen::tool::ToolOutput;
|
||||
@@ -28,9 +30,9 @@ use crate::compact::usage_tracker::UsageTracker;
|
||||
use session_store::SystemItem;
|
||||
|
||||
use crate::hook::{
|
||||
AbortInfo, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction,
|
||||
HookEventKind, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookPromptAction,
|
||||
HookRegistry, HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo,
|
||||
SystemItemAppendHandle, ToolCallSummary, ToolResultSummary, TurnEndInfo,
|
||||
RegisteredHook, SystemItemAppendHandle, ToolCallSummary, ToolResultSummary, TurnEndInfo,
|
||||
};
|
||||
use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance};
|
||||
use crate::prompt::catalog::PromptCatalog;
|
||||
@@ -41,6 +43,32 @@ use agen::token_counter::total_tokens;
|
||||
|
||||
/// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`.
|
||||
const FINAL_TEXT_PREVIEW_LIMIT: usize = 512;
|
||||
const INLINE_HOOK_CHAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
|
||||
|
||||
async fn call_hook_before_deadline<E: HookEventKind>(
|
||||
hook: &RegisteredHook<E>,
|
||||
input: &E::Input,
|
||||
deadline: tokio::time::Instant,
|
||||
) -> Result<Option<E::Output>, InterceptorError> {
|
||||
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
|
||||
if remaining.is_zero() {
|
||||
return Err(InterceptorError::new(
|
||||
InterceptorErrorCategory::Policy,
|
||||
"Worker hook chain exceeded its host-enforced lifecycle deadline",
|
||||
));
|
||||
}
|
||||
tokio::time::timeout(remaining, hook.call_optional(input))
|
||||
.await
|
||||
.map_err(|_| {
|
||||
InterceptorError::new(
|
||||
InterceptorErrorCategory::Policy,
|
||||
"Worker hook chain exceeded its host-enforced lifecycle deadline",
|
||||
)
|
||||
})?
|
||||
.map_err(|error| {
|
||||
InterceptorError::new(InterceptorErrorCategory::Dependency, error.to_string())
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) struct WorkerInterceptor {
|
||||
registry: Arc<HookRegistry>,
|
||||
@@ -231,8 +259,12 @@ impl WorkerInterceptor {
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Interceptor for WorkerInterceptor {
|
||||
async fn on_prompt_submit(&self, item: &mut Item) -> PromptAction {
|
||||
impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
|
||||
async fn on_prompt_submit(
|
||||
&self,
|
||||
context: PromptSubmitContext<'_, SessionHistoryMetadata>,
|
||||
) -> InterceptorResult<PromptAction> {
|
||||
let item = context.item;
|
||||
let turn_index = self.next_turn_index.fetch_add(1, Ordering::Relaxed);
|
||||
self.tool_calls_this_turn.store(0, Ordering::Relaxed);
|
||||
|
||||
@@ -240,19 +272,27 @@ impl Interceptor for WorkerInterceptor {
|
||||
input_text: extract_message_text(item).unwrap_or_default(),
|
||||
turn_index,
|
||||
};
|
||||
let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT;
|
||||
let mut cancellations = Vec::new();
|
||||
for hook in &self.registry.on_prompt_submit {
|
||||
let action = hook.call(&info).await;
|
||||
if !matches!(action, HookPromptAction::Continue) {
|
||||
return action.into();
|
||||
let Some(action) = call_hook_before_deadline(hook, &info, deadline).await? else {
|
||||
continue;
|
||||
};
|
||||
if let HookPromptAction::Cancel(reason) = action {
|
||||
cancellations.push(reason);
|
||||
}
|
||||
}
|
||||
cancellations.sort();
|
||||
if let Some(reason) = cancellations.into_iter().next() {
|
||||
return Ok(PromptAction::Cancel(reason));
|
||||
}
|
||||
let mut extras: Vec<SystemItem> = std::mem::take(
|
||||
&mut *self
|
||||
.pending_attachments
|
||||
.lock()
|
||||
.expect("pending_attachments poisoned"),
|
||||
);
|
||||
if extras.is_empty() {
|
||||
Ok(if extras.is_empty() {
|
||||
PromptAction::Continue
|
||||
} else {
|
||||
// Commit the typed system items first, then hand the
|
||||
@@ -266,10 +306,13 @@ impl Interceptor for WorkerInterceptor {
|
||||
Ok(()) => PromptAction::ContinueWith(items),
|
||||
Err(error) => PromptAction::Cancel(format!("session persistence failed: {error}")),
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn pending_history_appends(&self) -> Result<Vec<Item>, String> {
|
||||
async fn pending_history_appends(
|
||||
&self,
|
||||
_context: PendingHistoryAppendsContext<'_, SessionHistoryMetadata>,
|
||||
) -> InterceptorResult<Vec<Item>> {
|
||||
let drained = self.pending_notifies.drain();
|
||||
if drained.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
@@ -295,7 +338,10 @@ impl Interceptor for WorkerInterceptor {
|
||||
Ok(system_item) => system_item,
|
||||
Err(error) => {
|
||||
self.pending_notifies.requeue_front(drained);
|
||||
return Err(format!("failed to render notify_wrapper: {error}"));
|
||||
return Err(InterceptorError::new(
|
||||
InterceptorErrorCategory::Dependency,
|
||||
format!("failed to render notify_wrapper: {error}"),
|
||||
));
|
||||
}
|
||||
};
|
||||
items.push(system_item.to_history_item());
|
||||
@@ -303,15 +349,22 @@ impl Interceptor for WorkerInterceptor {
|
||||
}
|
||||
if let Err(error) = self.commit_system_items(&system_items) {
|
||||
self.pending_notifies.requeue_front(drained);
|
||||
return Err(format!("session persistence failed: {error}"));
|
||||
return Err(InterceptorError::new(
|
||||
InterceptorErrorCategory::Dependency,
|
||||
format!("session persistence failed: {error}"),
|
||||
));
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
|
||||
async fn pre_llm_request(
|
||||
&self,
|
||||
context: PreLlmRequestContext<'_, SessionHistoryMetadata>,
|
||||
) -> InterceptorResult<PreRequestAction> {
|
||||
let context = context.items;
|
||||
let initial_tokens = self.estimated_tokens(context);
|
||||
if self.request_threshold_exceeded(initial_tokens, context) {
|
||||
return PreRequestAction::Yield;
|
||||
return Ok(PreRequestAction::Yield);
|
||||
}
|
||||
let info = PreRequestInfo {
|
||||
item_count: context.len(),
|
||||
@@ -325,12 +378,28 @@ impl Interceptor for WorkerInterceptor {
|
||||
.as_ref()
|
||||
.map(|_| SystemItemAppendHandle::new(Arc::clone(&pending_hook_system_items)));
|
||||
let hook_context = PreRequestContext::new(info, system_item_sink);
|
||||
let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT;
|
||||
let mut cancellations = Vec::new();
|
||||
let mut should_yield = false;
|
||||
for hook in &self.registry.pre_llm_request {
|
||||
let action = hook.call(&hook_context).await;
|
||||
if !matches!(action, HookPreRequestAction::Continue) {
|
||||
return action.into();
|
||||
let Some(action) = call_hook_before_deadline(hook, &hook_context, deadline).await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match action {
|
||||
HookPreRequestAction::Continue => {}
|
||||
HookPreRequestAction::Yield => should_yield = true,
|
||||
HookPreRequestAction::Cancel(reason) => cancellations.push(reason),
|
||||
}
|
||||
}
|
||||
cancellations.sort();
|
||||
if let Some(reason) = cancellations.into_iter().next() {
|
||||
return Ok(PreRequestAction::Cancel(reason));
|
||||
}
|
||||
if should_yield {
|
||||
return Ok(PreRequestAction::Yield);
|
||||
}
|
||||
|
||||
let mut system_items: Vec<SystemItem> = std::mem::take(
|
||||
&mut *pending_hook_system_items
|
||||
@@ -353,44 +422,73 @@ impl Interceptor for WorkerInterceptor {
|
||||
|
||||
if self.request_threshold_exceeded(current_tokens, effective_context.as_ref()) {
|
||||
if let Err(error) = self.commit_system_items(&system_items) {
|
||||
return PreRequestAction::Cancel(format!("session persistence failed: {error}"));
|
||||
return Ok(PreRequestAction::Cancel(format!(
|
||||
"session persistence failed: {error}"
|
||||
)));
|
||||
}
|
||||
return if appended_items.is_empty() {
|
||||
return Ok(if appended_items.is_empty() {
|
||||
PreRequestAction::Yield
|
||||
} else {
|
||||
PreRequestAction::YieldWith(appended_items)
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(usage_tracker) = self.usage_tracker.as_ref() {
|
||||
usage_tracker.note_request(effective_context.len());
|
||||
}
|
||||
if system_items.is_empty() {
|
||||
return PreRequestAction::Continue;
|
||||
return Ok(PreRequestAction::Continue);
|
||||
}
|
||||
match self.commit_system_items(&system_items) {
|
||||
Ok(match self.commit_system_items(&system_items) {
|
||||
Ok(()) => PreRequestAction::ContinueWith(appended_items),
|
||||
Err(error) => PreRequestAction::Cancel(format!("session persistence failed: {error}")),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction {
|
||||
async fn pre_tool_call(
|
||||
&self,
|
||||
info: &mut ToolCallInfo<'_, SessionHistoryMetadata>,
|
||||
) -> InterceptorResult<PreToolAction> {
|
||||
let summary = ToolCallSummary {
|
||||
call_id: info.call.id.clone(),
|
||||
tool_name: info.call.name.clone(),
|
||||
arguments: info.call.input.clone(),
|
||||
};
|
||||
let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT;
|
||||
let mut aborts = Vec::new();
|
||||
let mut should_pause = false;
|
||||
let mut denials = Vec::new();
|
||||
for hook in &self.registry.pre_tool_call {
|
||||
let action = hook.call(&summary).await;
|
||||
if !matches!(action, HookPreToolAction::Continue) {
|
||||
return action.into_worker_action(summary.call_id.clone());
|
||||
let Some(action) = call_hook_before_deadline(hook, &summary, deadline).await? else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match action {
|
||||
HookPreToolAction::Continue => {}
|
||||
HookPreToolAction::Pause => should_pause = true,
|
||||
HookPreToolAction::Deny(reason) => denials.push(reason),
|
||||
HookPreToolAction::Abort(reason) => aborts.push(reason),
|
||||
}
|
||||
}
|
||||
aborts.sort();
|
||||
if let Some(reason) = aborts.into_iter().next() {
|
||||
return Ok(HookPreToolAction::Abort(reason).into_worker_action(summary.call_id.clone()));
|
||||
}
|
||||
if should_pause {
|
||||
return Ok(PreToolAction::Pause);
|
||||
}
|
||||
denials.sort();
|
||||
if let Some(reason) = denials.into_iter().next() {
|
||||
return Ok(HookPreToolAction::Deny(reason).into_worker_action(summary.call_id.clone()));
|
||||
}
|
||||
self.tool_calls_this_turn.fetch_add(1, Ordering::Relaxed);
|
||||
PreToolAction::Continue
|
||||
Ok(PreToolAction::Continue)
|
||||
}
|
||||
|
||||
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction {
|
||||
async fn post_tool_call(
|
||||
&self,
|
||||
info: &ToolResultInfo<'_, SessionHistoryMetadata>,
|
||||
) -> InterceptorResult<PostToolAction> {
|
||||
let summary = ToolResultSummary {
|
||||
call_id: info.result.tool_use_id.clone(),
|
||||
tool_name: info.call.name.clone(),
|
||||
@@ -402,21 +500,34 @@ impl Interceptor for WorkerInterceptor {
|
||||
attachments: Vec::new(),
|
||||
},
|
||||
};
|
||||
let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT;
|
||||
let mut aborts = Vec::new();
|
||||
for hook in &self.registry.post_tool_call {
|
||||
let action = hook.call(&summary).await;
|
||||
if !matches!(action, HookPostToolAction::Continue) {
|
||||
return action.into();
|
||||
let Some(action) = call_hook_before_deadline(hook, &summary, deadline).await? else {
|
||||
continue;
|
||||
};
|
||||
|
||||
if let HookPostToolAction::Abort(reason) = action {
|
||||
aborts.push(reason);
|
||||
}
|
||||
}
|
||||
PostToolAction::Continue
|
||||
aborts.sort();
|
||||
if let Some(reason) = aborts.into_iter().next() {
|
||||
return Ok(PostToolAction::Abort(reason));
|
||||
}
|
||||
Ok(PostToolAction::Continue)
|
||||
}
|
||||
|
||||
async fn on_turn_end(&self, history: &[Item]) -> TurnEndAction {
|
||||
async fn on_assistant_turn_end(
|
||||
&self,
|
||||
context: AssistantTurnEndContext<'_, SessionHistoryMetadata>,
|
||||
) -> InterceptorResult<TurnEndAction> {
|
||||
let history = context.history;
|
||||
let final_text_preview = history
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|i| i.is_assistant_message())
|
||||
.and_then(extract_message_text)
|
||||
.find(|entry| entry.item.is_assistant_message())
|
||||
.and_then(|entry| extract_message_text(&entry.item))
|
||||
.map(|t| preview(&t, FINAL_TEXT_PREVIEW_LIMIT))
|
||||
.unwrap_or_default();
|
||||
let info = TurnEndInfo {
|
||||
@@ -424,22 +535,20 @@ impl Interceptor for WorkerInterceptor {
|
||||
tool_calls_count: self.tool_calls_this_turn.load(Ordering::Relaxed),
|
||||
final_text_preview,
|
||||
};
|
||||
let deadline = tokio::time::Instant::now() + INLINE_HOOK_CHAIN_TIMEOUT;
|
||||
let mut should_pause = false;
|
||||
for hook in &self.registry.on_turn_end {
|
||||
let action = hook.call(&info).await;
|
||||
if !matches!(action, HookTurnEndAction::Finish) {
|
||||
return action.into();
|
||||
let Some(action) = call_hook_before_deadline(hook, &info, deadline).await? else {
|
||||
continue;
|
||||
};
|
||||
if matches!(action, HookTurnEndAction::Pause) {
|
||||
should_pause = true;
|
||||
}
|
||||
}
|
||||
TurnEndAction::Finish
|
||||
}
|
||||
|
||||
async fn on_abort(&self, reason: &str) {
|
||||
let info = AbortInfo {
|
||||
reason: reason.to_string(),
|
||||
};
|
||||
for hook in &self.registry.on_abort {
|
||||
hook.call(&info).await;
|
||||
if should_pause {
|
||||
return Ok(TurnEndAction::Pause);
|
||||
}
|
||||
Ok(TurnEndAction::Finish)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -509,6 +618,7 @@ mod tests {
|
||||
Hook, HookPostToolAction, HookPreRequestAction, HookPreToolAction, HookRegistryBuilder,
|
||||
HookTurnEndAction, OnTurnEnd, PostToolCall, PreLlmRequest, PreToolCall,
|
||||
};
|
||||
use crate::session_history::{WorkerHistoryProvenance, history_entry};
|
||||
|
||||
fn test_prompts() -> Arc<ArcSwap<PromptCatalog>> {
|
||||
Arc::new(ArcSwap::from(PromptCatalog::builtins_only().unwrap()))
|
||||
@@ -518,9 +628,12 @@ mod tests {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for CountingHook {
|
||||
async fn call(&self, _info: &PreRequestContext) -> HookPreRequestAction {
|
||||
async fn call(
|
||||
&self,
|
||||
_info: &PreRequestContext,
|
||||
) -> Result<HookPreRequestAction, crate::hook::HookError> {
|
||||
self.0.fetch_add(1, Ordering::Relaxed);
|
||||
HookPreRequestAction::Continue
|
||||
Ok(HookPreRequestAction::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -559,16 +672,22 @@ mod tests {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for AppendingPreRequestHook {
|
||||
async fn call(&self, input: &PreRequestContext) -> HookPreRequestAction {
|
||||
async fn call(
|
||||
&self,
|
||||
input: &PreRequestContext,
|
||||
) -> Result<HookPreRequestAction, crate::hook::HookError> {
|
||||
if let Some(system_items) = input.system_items() {
|
||||
self.saw_handle.store(true, Ordering::Relaxed);
|
||||
system_items.append_task_reminder("hook reminder");
|
||||
}
|
||||
HookPreRequestAction::Continue
|
||||
Ok(HookPreRequestAction::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
fn task_tool_call_info(name: &str, input: serde_json::Value) -> ToolCallInfo {
|
||||
fn task_tool_call_info(
|
||||
name: &str,
|
||||
input: serde_json::Value,
|
||||
) -> ToolCallInfo<'static, SessionHistoryMetadata> {
|
||||
let def = crate::feature::builtin::task::task_tools(
|
||||
crate::feature::builtin::task::TaskStore::new(),
|
||||
)
|
||||
@@ -580,6 +699,8 @@ mod tests {
|
||||
.expect("task tool definition");
|
||||
let (meta, tool) = def();
|
||||
ToolCallInfo {
|
||||
invocation: Default::default(),
|
||||
history: &[],
|
||||
call: agen::tool::ToolCall {
|
||||
id: "call-id".into(),
|
||||
name: name.into(),
|
||||
@@ -623,7 +744,14 @@ mod tests {
|
||||
None,
|
||||
);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Yield));
|
||||
// Hook must not run when an internal mechanism short-circuits first.
|
||||
@@ -655,7 +783,14 @@ mod tests {
|
||||
})),
|
||||
);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
match action {
|
||||
PreRequestAction::YieldWith(items) => assert_eq!(items.len(), 1),
|
||||
@@ -692,7 +827,14 @@ mod tests {
|
||||
)
|
||||
.with_usage_tracker(usage_tracker);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Yield));
|
||||
}
|
||||
@@ -716,7 +858,14 @@ mod tests {
|
||||
None,
|
||||
);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
@@ -757,7 +906,14 @@ mod tests {
|
||||
None,
|
||||
);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
@@ -784,7 +940,14 @@ mod tests {
|
||||
None,
|
||||
);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
@@ -805,7 +968,14 @@ mod tests {
|
||||
None,
|
||||
);
|
||||
let mut ctx: Vec<Item> = Vec::new();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
@@ -834,7 +1004,14 @@ mod tests {
|
||||
);
|
||||
|
||||
let mut ctx: Vec<Item> = Vec::new();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(saw_handle.load(Ordering::Relaxed));
|
||||
let PreRequestAction::ContinueWith(items) = action else {
|
||||
@@ -881,7 +1058,14 @@ mod tests {
|
||||
);
|
||||
|
||||
let mut ctx: Vec<Item> = Vec::new();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!saw_handle.load(Ordering::Relaxed));
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
@@ -891,33 +1075,42 @@ mod tests {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for AbortingHook {
|
||||
async fn call(&self, _info: &PreRequestContext) -> HookPreRequestAction {
|
||||
async fn call(
|
||||
&self,
|
||||
_info: &PreRequestContext,
|
||||
) -> Result<HookPreRequestAction, crate::hook::HookError> {
|
||||
self.0.store(true, Ordering::Relaxed);
|
||||
HookPreRequestAction::Cancel("nope".into())
|
||||
Ok(HookPreRequestAction::Cancel("nope".into()))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn public_pre_tool_hook_deny_becomes_synthetic_error_and_short_circuits() {
|
||||
async fn public_pre_tool_hook_denials_compose_without_short_circuiting() {
|
||||
struct DenyToolHook(Arc<AtomicUsize>);
|
||||
struct CountingToolHook(Arc<AtomicUsize>);
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreToolCall> for DenyToolHook {
|
||||
async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction {
|
||||
async fn call(
|
||||
&self,
|
||||
input: &ToolCallSummary,
|
||||
) -> Result<HookPreToolAction, crate::hook::HookError> {
|
||||
self.0.fetch_add(1, Ordering::Relaxed);
|
||||
assert_eq!(input.call_id, "call-id");
|
||||
assert_eq!(input.tool_name, "TaskList");
|
||||
assert_eq!(input.arguments, serde_json::json!({"scope": "all"}));
|
||||
HookPreToolAction::Deny("blocked by public hook".into())
|
||||
Ok(HookPreToolAction::Deny("blocked by public hook".into()))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreToolCall> for CountingToolHook {
|
||||
async fn call(&self, _input: &ToolCallSummary) -> HookPreToolAction {
|
||||
async fn call(
|
||||
&self,
|
||||
_input: &ToolCallSummary,
|
||||
) -> Result<HookPreToolAction, crate::hook::HookError> {
|
||||
self.0.fetch_add(1, Ordering::Relaxed);
|
||||
HookPreToolAction::Continue
|
||||
Ok(HookPreToolAction::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -938,7 +1131,7 @@ mod tests {
|
||||
);
|
||||
let mut info = task_tool_call_info("TaskList", serde_json::json!({"scope": "all"}));
|
||||
|
||||
let action = interceptor.pre_tool_call(&mut info).await;
|
||||
let action = interceptor.pre_tool_call(&mut info).await.unwrap();
|
||||
|
||||
match action {
|
||||
PreToolAction::SyntheticResult(result) => {
|
||||
@@ -950,7 +1143,7 @@ mod tests {
|
||||
other => panic!("expected synthetic denial, got {other:?}"),
|
||||
}
|
||||
assert_eq!(first_count.load(Ordering::Relaxed), 1);
|
||||
assert_eq!(second_count.load(Ordering::Relaxed), 0);
|
||||
assert_eq!(second_count.load(Ordering::Relaxed), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -959,14 +1152,17 @@ mod tests {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PostToolCall> for AbortAfterToolHook {
|
||||
async fn call(&self, input: &ToolResultSummary) -> HookPostToolAction {
|
||||
async fn call(
|
||||
&self,
|
||||
input: &ToolResultSummary,
|
||||
) -> Result<HookPostToolAction, crate::hook::HookError> {
|
||||
self.0.fetch_add(1, Ordering::Relaxed);
|
||||
assert_eq!(input.call_id, "call-id");
|
||||
assert_eq!(input.tool_name, "TaskList");
|
||||
assert!(!input.is_error);
|
||||
assert_eq!(input.output.summary, "ok");
|
||||
assert_eq!(input.output.content.as_deref(), Some("full"));
|
||||
HookPostToolAction::Abort("post tool abort".into())
|
||||
Ok(HookPostToolAction::Abort("post tool abort".into()))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -984,7 +1180,9 @@ mod tests {
|
||||
None,
|
||||
);
|
||||
let info = task_tool_call_info("TaskList", serde_json::json!({}));
|
||||
let mut result_info = ToolResultInfo {
|
||||
let result_info = ToolResultInfo {
|
||||
invocation: Default::default(),
|
||||
history: &[],
|
||||
call: info.call,
|
||||
result: agen::tool::ToolResult::from_output(
|
||||
"call-id",
|
||||
@@ -1000,7 +1198,7 @@ mod tests {
|
||||
context: info.context,
|
||||
};
|
||||
|
||||
let action = interceptor.post_tool_call(&mut result_info).await;
|
||||
let action = interceptor.post_tool_call(&result_info).await.unwrap();
|
||||
|
||||
assert_eq!(action, PostToolAction::Abort("post tool abort".to_string()));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
@@ -1012,12 +1210,15 @@ mod tests {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<OnTurnEnd> for PauseTurnEndHook {
|
||||
async fn call(&self, input: &TurnEndInfo) -> HookTurnEndAction {
|
||||
async fn call(
|
||||
&self,
|
||||
input: &TurnEndInfo,
|
||||
) -> Result<HookTurnEndAction, crate::hook::HookError> {
|
||||
self.0.fetch_add(1, Ordering::Relaxed);
|
||||
assert_eq!(input.turn_index, 0);
|
||||
assert_eq!(input.tool_calls_count, 0);
|
||||
assert_eq!(input.final_text_preview, "done");
|
||||
HookTurnEndAction::Pause
|
||||
Ok(HookTurnEndAction::Pause)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1034,9 +1235,25 @@ mod tests {
|
||||
test_prompts(),
|
||||
None,
|
||||
);
|
||||
let history = vec![Item::user_message("hi"), Item::assistant_message("done")];
|
||||
|
||||
let action = interceptor.on_turn_end(&history).await;
|
||||
let history = vec![
|
||||
history_entry(
|
||||
Item::user_message("hi"),
|
||||
WorkerHistoryProvenance::LegacyUnknown,
|
||||
),
|
||||
history_entry(
|
||||
Item::assistant_message("done"),
|
||||
WorkerHistoryProvenance::LegacyUnknown,
|
||||
),
|
||||
];
|
||||
let action = interceptor
|
||||
.on_assistant_turn_end(AssistantTurnEndContext {
|
||||
invocation: Default::default(),
|
||||
assistant_entries: &history[1..],
|
||||
history: &history,
|
||||
tool_calls: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(action, TurnEndAction::Pause));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
@@ -1073,7 +1290,14 @@ mod tests {
|
||||
let ctx_items = vec![Item::user_message("hi")];
|
||||
for _ in 0..23 {
|
||||
let mut ctx = ctx_items.clone();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
usage_tracker.record_usage(&agen::event::UsageEvent {
|
||||
input_tokens: Some(10),
|
||||
@@ -1085,7 +1309,14 @@ mod tests {
|
||||
}
|
||||
|
||||
let mut ctx = ctx_items.clone();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let appended_len = match action {
|
||||
PreRequestAction::ContinueWith(items) => items.len(),
|
||||
other => panic!("expected reminder append, got {other:?}"),
|
||||
@@ -1158,7 +1389,13 @@ mod tests {
|
||||
));
|
||||
|
||||
buffer.push_notify("updated".to_string(), false);
|
||||
let appends = interceptor.pending_history_appends().await.unwrap();
|
||||
let appends = interceptor
|
||||
.pending_history_appends(PendingHistoryAppendsContext {
|
||||
invocation: Default::default(),
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(appends.len(), 1);
|
||||
assert!(format!("{:?}", appends[0]).contains("CURRENT-PROJECTION updated"));
|
||||
let committed = committed.lock().unwrap();
|
||||
@@ -1208,9 +1445,19 @@ mod tests {
|
||||
));
|
||||
buffer.push_notify("must persist".to_string(), false);
|
||||
|
||||
let error = interceptor.pending_history_appends().await.unwrap_err();
|
||||
let error = interceptor
|
||||
.pending_history_appends(PendingHistoryAppendsContext {
|
||||
invocation: Default::default(),
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(error.contains("failed to render notify_wrapper"));
|
||||
assert!(
|
||||
error
|
||||
.diagnostic()
|
||||
.contains("failed to render notify_wrapper")
|
||||
);
|
||||
let requeued = buffer.drain();
|
||||
assert_eq!(requeued.len(), 1);
|
||||
}
|
||||
@@ -1232,7 +1479,13 @@ mod tests {
|
||||
None,
|
||||
);
|
||||
|
||||
let items = interceptor.pending_history_appends().await.unwrap();
|
||||
let items = interceptor
|
||||
.pending_history_appends(PendingHistoryAppendsContext {
|
||||
invocation: Default::default(),
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(items.len(), 2);
|
||||
let first = items[0].as_text().unwrap_or_default();
|
||||
let second = items[1].as_text().unwrap_or_default();
|
||||
@@ -1246,7 +1499,13 @@ mod tests {
|
||||
);
|
||||
|
||||
// Empty buffer → empty Vec (no synthesised items).
|
||||
let again = interceptor.pending_history_appends().await.unwrap();
|
||||
let again = interceptor
|
||||
.pending_history_appends(PendingHistoryAppendsContext {
|
||||
invocation: Default::default(),
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(again.is_empty());
|
||||
}
|
||||
|
||||
@@ -1269,7 +1528,14 @@ mod tests {
|
||||
None,
|
||||
);
|
||||
let mut ctx: Vec<Item> = vec![Item::user_message("hi")];
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(ctx.len(), 1, "pre_llm_request must not append notifies");
|
||||
@@ -1299,10 +1565,17 @@ mod tests {
|
||||
None,
|
||||
);
|
||||
let mut ctx: Vec<Item> = Vec::new();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
let action = interceptor
|
||||
.pre_llm_request(PreLlmRequestContext {
|
||||
invocation: Default::default(),
|
||||
items: &mut ctx,
|
||||
history: &[],
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Cancel(_)));
|
||||
assert!(first_called.load(Ordering::Relaxed));
|
||||
assert_eq!(second_count.load(Ordering::Relaxed), 0);
|
||||
assert_eq!(second_count.load(Ordering::Relaxed), 1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,6 +50,7 @@ pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTem
|
||||
pub use protocol::{ErrorCode, Event, Method, TurnResult, WorkerStatus};
|
||||
pub use runtime::dir::RuntimeDir;
|
||||
pub use segment_log_sink::SegmentLogSink;
|
||||
pub use session_capture::SessionEntryRef;
|
||||
pub use session_history::{
|
||||
SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata,
|
||||
WorkerHistoryProvenance, WorkerSubjectSnapshot,
|
||||
|
||||
@@ -45,14 +45,17 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreToolCall> for PermissionHook {
|
||||
async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction {
|
||||
match self.action_for(input) {
|
||||
async fn call(
|
||||
&self,
|
||||
input: &ToolCallSummary,
|
||||
) -> Result<HookPreToolAction, crate::hook::HookError> {
|
||||
Ok(match self.action_for(input) {
|
||||
ToolPermissionAction::Allow => HookPreToolAction::Continue,
|
||||
ToolPermissionAction::Deny => HookPreToolAction::Deny(permission_denied_message(input)),
|
||||
ToolPermissionAction::Ask => {
|
||||
HookPreToolAction::Deny(permission_ask_unsupported_message(input))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -174,7 +177,7 @@ mod tests {
|
||||
))
|
||||
.await;
|
||||
match denied {
|
||||
HookPreToolAction::Deny(message) => {
|
||||
Ok(HookPreToolAction::Deny(message)) => {
|
||||
assert!(message.contains("permission denied"));
|
||||
assert!(message.contains("Bash"));
|
||||
}
|
||||
@@ -192,7 +195,7 @@ mod tests {
|
||||
))
|
||||
.await;
|
||||
match asked {
|
||||
HookPreToolAction::Deny(message) => {
|
||||
Ok(HookPreToolAction::Deny(message)) => {
|
||||
assert!(message.contains("permission ask unsupported"));
|
||||
assert!(message.contains("denied fail-closed"));
|
||||
}
|
||||
|
||||
@@ -23,14 +23,14 @@ const OVERVIEW_ANCHOR_STRIDE: usize = 8;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub(crate) struct SessionEntryRef(String);
|
||||
pub struct SessionEntryRef(String);
|
||||
|
||||
impl SessionEntryRef {
|
||||
pub(crate) fn from_history_entry_id(entry_id: &crate::SessionHistoryEntryId) -> Self {
|
||||
pub fn from_history_entry_id(entry_id: &crate::SessionHistoryEntryId) -> Self {
|
||||
Self(format!("E{}", entry_id.0))
|
||||
}
|
||||
|
||||
pub(crate) fn parse(value: &str) -> Option<Self> {
|
||||
pub fn parse(value: &str) -> Option<Self> {
|
||||
let suffix = value.strip_prefix('E')?;
|
||||
if suffix.is_empty()
|
||||
|| suffix.len() > 64
|
||||
@@ -43,7 +43,7 @@ impl SessionEntryRef {
|
||||
Some(Self(value.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) fn as_str(&self) -> &str {
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
|
||||
@@ -70,9 +70,12 @@ impl TicketIntakeReadyShutdownHook {
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PostToolCall> for TicketIntakeReadyShutdownHook {
|
||||
async fn call(&self, info: &ToolResultSummary) -> HookPostToolAction {
|
||||
async fn call(
|
||||
&self,
|
||||
info: &ToolResultSummary,
|
||||
) -> Result<HookPostToolAction, crate::hook::HookError> {
|
||||
self.observe_tool_result(info);
|
||||
HookPostToolAction::Continue
|
||||
Ok(HookPostToolAction::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+28
-125
@@ -2,127 +2,11 @@ use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::worker::WorkspaceClient;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SkillDiagnosticSeverity {
|
||||
Error,
|
||||
Warning,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct SkillDiagnostic {
|
||||
pub severity: SkillDiagnosticSeverity,
|
||||
pub code: String,
|
||||
pub message: String,
|
||||
/// Path-free authority/provenance label such as `builtin:foo` or `workspace:foo`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub source: Option<String>,
|
||||
}
|
||||
|
||||
impl SkillDiagnostic {
|
||||
pub fn error(
|
||||
code: impl Into<String>,
|
||||
message: impl Into<String>,
|
||||
source: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
severity: SkillDiagnosticSeverity::Error,
|
||||
code: code.into(),
|
||||
message: message.into(),
|
||||
source,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn warning(
|
||||
code: impl Into<String>,
|
||||
message: impl Into<String>,
|
||||
source: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
severity: SkillDiagnosticSeverity::Warning,
|
||||
code: code.into(),
|
||||
message: message.into(),
|
||||
source,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SkillSourceKind {
|
||||
Builtin,
|
||||
Workspace,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct SkillProvenance {
|
||||
pub kind: SkillSourceKind,
|
||||
/// Stable id: `builtin:<name>` or `workspace:<name>`.
|
||||
pub id: String,
|
||||
/// Virtual config/resource path. Never an absolute host filesystem path.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub virtual_path: Option<String>,
|
||||
/// Active Workspace config revision for Workspace Skills.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub revision: Option<u64>,
|
||||
/// Digest of the immutable `SKILL.md` source.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub source_digest: Option<String>,
|
||||
/// Digest of the active virtual config tree snapshot.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tree_digest: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct SkillResourceRef {
|
||||
pub kind: String,
|
||||
/// Skill-relative resource name/path. Never an absolute filesystem path.
|
||||
pub name: String,
|
||||
pub supported: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub diagnostic: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct SkillCatalogEntry {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub provenance: SkillProvenance,
|
||||
#[serde(default)]
|
||||
pub overrides: Vec<SkillProvenance>,
|
||||
#[serde(default)]
|
||||
pub diagnostics: Vec<SkillDiagnostic>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct SkillCatalogResponse {
|
||||
/// Authority label for diagnostics; callers must not interpret it as a path.
|
||||
pub authority: String,
|
||||
#[serde(default)]
|
||||
pub entries: Vec<SkillCatalogEntry>,
|
||||
#[serde(default)]
|
||||
pub diagnostics: Vec<SkillDiagnostic>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct SkillDetailResponse {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub provenance: SkillProvenance,
|
||||
#[serde(default)]
|
||||
pub overrides: Vec<SkillProvenance>,
|
||||
#[serde(default)]
|
||||
pub diagnostics: Vec<SkillDiagnostic>,
|
||||
/// Imported Markdown content with YAML frontmatter delimiters removed.
|
||||
/// This is intentionally omitted from catalog responses.
|
||||
pub body: String,
|
||||
#[serde(default)]
|
||||
pub allowed_tools: Vec<String>,
|
||||
/// Explicitly documents that allowed-tools is parsed only as an experimental hint.
|
||||
pub allowed_tools_status: String,
|
||||
#[serde(default)]
|
||||
pub resources: Vec<SkillResourceRef>,
|
||||
}
|
||||
pub use workspace_api::{
|
||||
SkillActivationStatus, SkillCatalogEntry, SkillCatalogResponse, SkillDetailResponse,
|
||||
SkillDiagnostic, SkillDiagnosticSeverity, SkillProjectionIdentity, SkillProjectionStatus,
|
||||
SkillProvenance, SkillResourceRef, SkillSourceKind,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct SkillActivationResponse {
|
||||
@@ -144,6 +28,8 @@ pub enum SkillClientError {
|
||||
Request(#[from] crate::worker::WorkspaceClientError),
|
||||
#[error("Skill API response JSON is invalid: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
#[error("Skill API response violates the shared contract: {0}")]
|
||||
InvalidResponse(#[from] workspace_api::SkillApiValidationError),
|
||||
#[error("Skill API returned HTTP {status}: {body}")]
|
||||
Http {
|
||||
status: reqwest::StatusCode,
|
||||
@@ -155,11 +41,15 @@ pub enum SkillClientError {
|
||||
|
||||
impl dyn WorkspaceClient + '_ {
|
||||
pub fn list_skills(&self) -> Result<SkillCatalogResponse, SkillClientError> {
|
||||
self.get_skill_json("skills")
|
||||
let response: SkillCatalogResponse = self.get_skill_json("skills")?;
|
||||
response.validate()?;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
pub fn read_skill(&self, name: &str) -> Result<SkillDetailResponse, SkillClientError> {
|
||||
self.get_skill_json(&format!("skills/{name}"))
|
||||
let response: SkillDetailResponse = self.get_skill_json(&format!("skills/{name}"))?;
|
||||
response.validate()?;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
pub fn activate_skill(&self, name: &str) -> Result<SkillActivationResponse, SkillClientError> {
|
||||
@@ -229,11 +119,24 @@ mod tests {
|
||||
assert_eq!(worker_header, None);
|
||||
assert_eq!(authorization, None);
|
||||
let body = serde_json::json!({
|
||||
"authority": "workspace-backend-skills-v0",
|
||||
"authority": "workspace-config-skills-v1",
|
||||
"projection": {
|
||||
"config_revision": 7,
|
||||
"tree_digest": "tree-digest"
|
||||
},
|
||||
"entries": [{
|
||||
"name": "triage-errors",
|
||||
"description": "Use when triaging errors.",
|
||||
"provenance": { "kind": "workspace", "id": "workspace:triage-errors" },
|
||||
"activation_status": "active",
|
||||
"projection_status": "valid",
|
||||
"provenance": {
|
||||
"kind": "workspace",
|
||||
"id": "workspace:triage-errors",
|
||||
"virtual_path": "skills/triage-errors/SKILL.md",
|
||||
"revision": 7,
|
||||
"source_digest": "source-digest",
|
||||
"tree_digest": "tree-digest"
|
||||
},
|
||||
"overrides": [],
|
||||
"diagnostics": []
|
||||
}],
|
||||
|
||||
@@ -425,6 +425,8 @@ impl Tool for SubWorkerSpawnTool {
|
||||
WorkerManifestConfig::resolution_defaults().merge(child_config),
|
||||
)
|
||||
.map_err(|error| ToolError::ExecutionFailed(format!("resolve child manifest: {error}")))?;
|
||||
bind_child_memory_settings(&self.spawner_manifest, &mut child_manifest)
|
||||
.map_err(ToolError::ExecutionFailed)?;
|
||||
// Delegated children stay bound to their scoped session and cannot use
|
||||
// Workspace attachment tools to replace it with parent-level authority.
|
||||
child_manifest.feature.manage_workdir.enabled = false;
|
||||
@@ -827,6 +829,33 @@ fn profile_error_with_available(error: ProfileError, available: &AvailableProfil
|
||||
)
|
||||
}
|
||||
|
||||
fn bind_child_memory_settings(
|
||||
parent: &manifest::WorkerManifest,
|
||||
child: &mut manifest::WorkerManifest,
|
||||
) -> Result<(), String> {
|
||||
if !child.feature.memory.profile.enabled {
|
||||
return child
|
||||
.feature
|
||||
.memory
|
||||
.validate_execution()
|
||||
.map_err(str::to_string);
|
||||
}
|
||||
let workspace_settings = parent.feature.memory.workspace_settings().ok_or_else(|| {
|
||||
"enabled child Memory feature requires the parent's trusted Workspace settings snapshot"
|
||||
.to_string()
|
||||
})?;
|
||||
child
|
||||
.feature
|
||||
.memory
|
||||
.bind_workspace_settings(workspace_settings)
|
||||
.map_err(str::to_string)?;
|
||||
child
|
||||
.feature
|
||||
.memory
|
||||
.validate_execution()
|
||||
.map_err(str::to_string)
|
||||
}
|
||||
|
||||
fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfig {
|
||||
WorkerManifestConfig {
|
||||
worker: WorkerMetaConfig {
|
||||
@@ -894,7 +923,6 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi
|
||||
model: c.model.clone(),
|
||||
}),
|
||||
web: manifest.web.clone(),
|
||||
memory: manifest.memory.clone(),
|
||||
skills: manifest.skills.clone(),
|
||||
}
|
||||
}
|
||||
@@ -1091,10 +1119,7 @@ enabled = true
|
||||
thread = true
|
||||
|
||||
[feature.memory]
|
||||
enabled = true
|
||||
|
||||
[memory]
|
||||
extract_threshold = 4000
|
||||
enabled = false
|
||||
"#;
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1526,6 +1551,33 @@ extract_threshold = 4000
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn child_memory_inherits_only_the_parents_trusted_settings_snapshot() {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let mut parent = parent_manifest(temp.path(), None);
|
||||
parent.feature.memory.profile.enabled = true;
|
||||
parent
|
||||
.feature
|
||||
.memory
|
||||
.bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot {
|
||||
workspace_id: "workspace-1".to_string(),
|
||||
settings_revision: 4,
|
||||
language: "日本語".to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
let mut child = parent.clone();
|
||||
child.feature.memory.workspace_settings = None;
|
||||
|
||||
bind_child_memory_settings(&parent, &mut child).unwrap();
|
||||
assert_eq!(
|
||||
child.feature.memory.workspace_settings(),
|
||||
parent.feature.memory.workspace_settings()
|
||||
);
|
||||
|
||||
child.feature.memory.profile.enabled = false;
|
||||
assert!(bind_child_memory_settings(&parent, &mut child).is_err());
|
||||
}
|
||||
|
||||
fn write_project_profile_registry(
|
||||
project: &Path,
|
||||
default: Option<&str>,
|
||||
|
||||
+514
-1457
File diff suppressed because it is too large
Load Diff
@@ -578,138 +578,6 @@ async fn mid_turn_compact_success_broadcasts_start_and_done() {
|
||||
assert_eq!(new_id_in_event, Some(worker.segment_id()));
|
||||
}
|
||||
|
||||
/// Regression: `Worker::compact()` must reset the in-memory
|
||||
/// `extract_pointer` so extract keeps firing on the new compacted
|
||||
/// session.
|
||||
///
|
||||
/// Without the reset, the pointer's `processed_through_history_len`
|
||||
/// holds the old (typically large) item count, while the new compacted
|
||||
/// session starts with a much shorter history (`[summary, ...]`).
|
||||
/// `cumulative_input_tokens_since` would then filter every new
|
||||
/// usage record out (their `history_len` is below the stale pointer)
|
||||
/// and extract would never re-fire for the rest of the process.
|
||||
const EXTRACT_PLUS_COMPACT_MANIFEST: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
pwd = "./"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[memory]
|
||||
workspace_id = "test-workspace"
|
||||
settings_revision = 1
|
||||
language = "English"
|
||||
extract_threshold = 1
|
||||
|
||||
[compaction]
|
||||
compact_threshold = 1
|
||||
compact_retained_tokens = 0
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
fn finish_memory_extraction_tool_use_events(call_id: &str) -> Vec<LlmEvent> {
|
||||
let input = serde_json::json!({
|
||||
"staged_count": 0,
|
||||
"no_candidates_reason": "test run has no durable candidates"
|
||||
})
|
||||
.to_string();
|
||||
vec![
|
||||
LlmEvent::tool_use_start(0, call_id, "FinishMemoryExtraction"),
|
||||
LlmEvent::tool_input_delta(0, input),
|
||||
LlmEvent::tool_use_stop(0),
|
||||
LlmEvent::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
]
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compact_resets_extract_pointer_so_extract_can_fire_again() {
|
||||
// Mock LLM responses, in call order:
|
||||
// [0] first run with usage(1000) so extract threshold (=1) fires.
|
||||
// [1] extract worker invokes FinishMemoryExtraction with empty output.
|
||||
// [2] extract worker closes after the tool result.
|
||||
// [3] compact worker invokes write_summary.
|
||||
// [4] compact worker closes after the tool result.
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("hi", 1000),
|
||||
finish_memory_extraction_tool_use_events("ec1"),
|
||||
single_text_events("done"),
|
||||
write_summary_tool_use_events("sc1", "summary"),
|
||||
single_text_events("done"),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_PLUS_COMPACT_MANIFEST, client).await;
|
||||
|
||||
worker.run_text("first").await.unwrap();
|
||||
|
||||
// extract fires; pointer becomes Some.
|
||||
worker.try_post_run_extract().await.unwrap();
|
||||
assert!(
|
||||
worker.extract_pointer().is_some(),
|
||||
"extract_pointer should be Some after a successful extract"
|
||||
);
|
||||
|
||||
// Compact runs. Without the fix the in-memory pointer would still
|
||||
// reference the old Segment's history_len.
|
||||
worker.try_pre_run_compact().await;
|
||||
assert!(
|
||||
worker.extract_pointer().is_none(),
|
||||
"extract_pointer must be reset to None after compact (matches cold-restore on the new Segment)"
|
||||
);
|
||||
}
|
||||
|
||||
/// `extract_threshold = 0` is treated as "disabled" — without this, a
|
||||
/// raw `>=` comparison against `tokens_since` would fire extract on
|
||||
/// every post-run regardless of activity. Mirrors the consolidation
|
||||
/// zero-threshold convention so users have a single way to opt out
|
||||
/// without removing the `[memory]` section.
|
||||
const EXTRACT_THRESHOLD_ZERO_MANIFEST: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
pwd = "./"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[memory]
|
||||
extract_threshold = 0
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
#[tokio::test]
|
||||
async fn extract_threshold_zero_is_disabled() {
|
||||
// Mock provides exactly one response — the first run. If extract
|
||||
// were treated as "fire on any change" because of `tokens_since >= 0`,
|
||||
// it would call into the extract worker and exhaust the mock.
|
||||
let client = MockClient::new(vec![text_events_with_usage("hi", 1000)]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_THRESHOLD_ZERO_MANIFEST, client).await;
|
||||
|
||||
worker.run_text("first").await.unwrap();
|
||||
worker
|
||||
.try_post_run_extract()
|
||||
.await
|
||||
.expect("extract_threshold=0 must skip silently, not fail");
|
||||
assert!(
|
||||
worker.extract_pointer().is_none(),
|
||||
"no extract should have run — pointer must remain None"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_run_compact_failure_broadcasts_start_and_failed() {
|
||||
// Only the first run has a response. Compaction will run the
|
||||
@@ -746,112 +614,6 @@ async fn pre_run_compact_failure_broadcasts_start_and_failed() {
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Detached post-run memory jobs (`spawn_post_run_memory_jobs` /
|
||||
// `wait_for_memory_jobs`). Covers the detach round-trip and the structural
|
||||
// invariant that the cloned memory-task Worker shares `SegmentState` with the
|
||||
// source Worker, so that `save_extension` from the background extract does not
|
||||
// leave the next turn's `save_user_input` looking at a stale session pointer.
|
||||
|
||||
const EXTRACT_NO_COMPACT_MANIFEST: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
pwd = "./"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[memory]
|
||||
workspace_id = "test-workspace"
|
||||
settings_revision = 1
|
||||
language = "English"
|
||||
extract_threshold = 1
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
#[tokio::test]
|
||||
async fn extract_large_unprocessed_range_does_not_abort_on_input_occupancy() {
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("recorded", 1000),
|
||||
finish_memory_extraction_tool_use_events("ec-large"),
|
||||
single_text_events("done"),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
|
||||
|
||||
let large_request = format!("remember this large slice: {}", "x ".repeat(200_000));
|
||||
worker.run_text(&large_request).await.unwrap();
|
||||
|
||||
worker.try_post_run_extract().await.expect(
|
||||
"large unprocessed extract ranges must reach the extract worker, not abort locally",
|
||||
);
|
||||
assert!(
|
||||
worker.extract_pointer().is_some(),
|
||||
"successful extract should advance the pointer even when the input range is large"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn spawn_and_wait_drives_extract_to_completion() {
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("hi", 1000),
|
||||
finish_memory_extraction_tool_use_events("ec1"),
|
||||
single_text_events("done"),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
|
||||
|
||||
worker.run_text("first").await.unwrap();
|
||||
assert!(
|
||||
worker.extract_pointer().is_none(),
|
||||
"extract has not run yet — pointer must be None"
|
||||
);
|
||||
|
||||
worker.spawn_post_run_memory_jobs();
|
||||
worker.wait_for_memory_jobs().await;
|
||||
|
||||
assert!(
|
||||
worker.extract_pointer().is_some(),
|
||||
"spawn + wait must complete extract; pointer should be set"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn detached_extract_does_not_fork_session_log() {
|
||||
// Source worker and the cloned memory-task worker share `SegmentState` via
|
||||
// `Arc<_>`. The detached extract advances the entry tally through
|
||||
// `save_extension`; the next `run` must see that same tally so
|
||||
// `ensure_head_or_fork` does not spawn a new session.
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("hi", 1000),
|
||||
finish_memory_extraction_tool_use_events("ec1"),
|
||||
single_text_events("done"),
|
||||
text_events_with_usage("ok", 1000),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
|
||||
|
||||
worker.run_text("first").await.unwrap();
|
||||
let session_before = worker.segment_id();
|
||||
|
||||
worker.spawn_post_run_memory_jobs();
|
||||
worker.wait_for_memory_jobs().await;
|
||||
|
||||
worker.run_text("second").await.unwrap();
|
||||
let session_after = worker.segment_id();
|
||||
|
||||
assert_eq!(
|
||||
session_before, session_after,
|
||||
"detached extract's save_extension and the next turn's save_user_input \
|
||||
must share the entry tally through SegmentState — a fork here means the \
|
||||
clone carried its own counter"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn controller_compact_method_emits_start_and_done() {
|
||||
let client = MockClient::new(vec![
|
||||
|
||||
Reference in New Issue
Block a user