diff --git a/crates/agen/src/engine.rs b/crates/agen/src/engine.rs index 4031d3514..356d59e70 100644 --- a/crates/agen/src/engine.rs +++ b/crates/agen/src/engine.rs @@ -1480,15 +1480,7 @@ impl Engine { fn get_unanswered_tool_calls(&self, history: &History) -> (Vec, Vec) { let mut pending_calls = Vec::new(); let mut started_call_ids = Vec::new(); - let mut answered_call_ids = std::collections::HashSet::new(); - - for item in history.items() { - if let Item::ToolResult { call_id, .. } = item { - answered_call_ids.insert(call_id.clone()); - } - } - - for item in history.items() { + for entry in unanswered_tool_call_entries(history) { if let Item::ToolCall { call_id, name, @@ -1496,8 +1488,7 @@ impl Engine { execution_id, status, .. - } = item - && !answered_call_ids.contains(call_id) + } = &entry.item { if execution_id.is_some() || *status == Some(ItemStatus::InProgress) { started_call_ids.push(call_id.clone()); @@ -1639,6 +1630,30 @@ impl Engine { } = info; tool_call = call; context.call_id = tool_call.id.clone(); + if let PreToolAction::SyntheticResult(result) = &action + && result.tool_use_id != expected_tool_use_id + { + return Err(InterceptorFailure::new( + InterceptorPhase::PreToolCall, + InterceptorError::new( + InterceptorErrorCategory::ContractViolation, + "synthetic tool result changed immutable tool call identity", + ), + ) + .into()); + } + let execution_id = if matches!(action, PreToolAction::Pause) { + None + } else { + Some(context.execution_id()) + }; + self.append_early_tool_call(history, annotate, &tool_call, call_index, execution_id)?; + // Policy runs before this call exists in history. Only the successful + // append's annotation may bind the identity, overriding hook changes. + context.logical_call_id = history + .entries() + .last() + .and_then(|entry| self.interceptor.logical_tool_call_id(&entry.annotation)); batch.call_info_map.insert( tool_call.id.clone(), (tool_call.clone(), meta, tool.clone(), context.clone()), @@ -1647,13 +1662,6 @@ impl Engine { match action { PreToolAction::Continue => { let attempt_id = context.execution_id(); - self.append_early_tool_call( - history, - annotate, - &tool_call, - call_index, - Some(attempt_id.clone()), - )?; batch .attempt_fence .register(tool_call.id.clone(), attempt_id.clone()); @@ -1695,34 +1703,10 @@ impl Engine { batch.dispatched_any = true; } PreToolAction::SyntheticResult(result) => { - if result.tool_use_id != expected_tool_use_id { - return Err(InterceptorFailure::new( - InterceptorPhase::PreToolCall, - InterceptorError::new( - InterceptorErrorCategory::ContractViolation, - "synthetic tool result changed immutable tool call identity", - ), - ) - .into()); - } - self.append_early_tool_call( - history, - annotate, - &tool_call, - call_index, - Some(context.execution_id()), - )?; batch.synthetic_results.push_back(result); batch.dispatched_any = true; } PreToolAction::Skip => { - self.append_early_tool_call( - history, - annotate, - &tool_call, - call_index, - Some(context.execution_id()), - )?; batch.synthetic_results.push_back(ToolResult::error( &tool_call.id, "Tool execution skipped by interceptor", @@ -1730,13 +1714,6 @@ impl Engine { batch.dispatched_any = true; } PreToolAction::Abort(reason) => { - self.append_early_tool_call( - history, - annotate, - &tool_call, - call_index, - Some(context.execution_id()), - )?; batch.synthetic_results.push_back(ToolResult::error( &tool_call.id, format!("Tool execution aborted: {reason}"), @@ -1745,7 +1722,6 @@ impl Engine { return Err(EngineError::Aborted(reason)); } PreToolAction::Pause => { - self.append_early_tool_call(history, annotate, &tool_call, call_index, None)?; batch.dispatched_any = true; return Ok(EarlyToolAdmission::Pause); } @@ -2010,10 +1986,24 @@ impl Engine { let batch_id = format!("tool-batch-{}", self.tool_execution_batch_count); self.tool_execution_batch_count += 1; + let mut logical_call_ids: HashMap>> = HashMap::new(); + for entry in unanswered_tool_call_entries(history) { + if let Item::ToolCall { call_id, .. } = &entry.item { + logical_call_ids + .entry(call_id.clone()) + .or_default() + .push_back(self.interceptor.logical_tool_call_id(&entry.annotation)); + } + } + // Phase 1: Apply pre_tool_call interceptor (determine skip/abort/synthetic result) let mut approved_calls = Vec::new(); for (call_index, mut tool_call) in tool_calls.into_iter().enumerate() { let expected_tool_use_id = tool_call.id.clone(); + let logical_call_id = logical_call_ids + .get_mut(&tool_call.id) + .and_then(VecDeque::pop_front) + .flatten(); let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index); if !tool_call.input.is_object() { synthetic_results.push(ToolResult::error( @@ -2079,6 +2069,7 @@ impl Engine { let tool_call = info.call; let mut context = info.context; context.call_id = tool_call.id.clone(); + context.logical_call_id = logical_call_id; call_info_map.insert( tool_call.id.clone(), (tool_call, info.meta.clone(), info.tool.clone(), context), @@ -2098,6 +2089,7 @@ impl Engine { tool_call = info.call; let mut context = info.context; context.call_id = tool_call.id.clone(); + context.logical_call_id = logical_call_id; call_info_map.insert( tool_call.id.clone(), @@ -2111,7 +2103,8 @@ impl Engine { approved_calls.push((tool_call, context, Some(info.tool))); } else { // Unknown tools go into approved list as-is (will error at execution) - let context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index); + let mut context = ToolExecutionContext::new(&tool_call.id, &batch_id, call_index); + context.logical_call_id = logical_call_id; approved_calls.push((tool_call, context, None)); } } @@ -4333,6 +4326,38 @@ fn item_kind(item: &Item) -> &'static str { } } +/// Pair each result with one preceding unanswered occurrence, never with a later +/// call that happens to reuse the provider ID. Keep the exact annotated entries +/// so dispatch and restoration resolve the same host-owned identity. +fn unanswered_tool_call_entries(history: &History) -> Vec<&HistoryEntry> { + let mut calls = Vec::new(); + let mut by_provider_id: HashMap<&str, VecDeque> = HashMap::new(); + for entry in history.entries() { + match &entry.item { + Item::ToolCall { call_id, .. } => { + by_provider_id + .entry(call_id.as_str()) + .or_default() + .push_back(calls.len()); + calls.push(Some(entry)); + } + Item::ToolResult { call_id, .. } => { + if let Some(queue) = by_provider_id.get_mut(call_id.as_str()) + && let Some(index) = queue.pop_front() + { + calls[index] = None; + } + } + _ => {} + } + } + calls.into_iter().flatten().collect() +} + +#[cfg(test)] +#[path = "tests/tool_call_identity.rs"] +mod tool_call_identity_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/crates/agen/src/interceptor.rs b/crates/agen/src/interceptor.rs index 56117c06d..11f580b9e 100644 --- a/crates/agen/src/interceptor.rs +++ b/crates/agen/src/interceptor.rs @@ -309,7 +309,7 @@ pub struct ToolCallInfo<'a, A = ()> { pub meta: ToolMeta, /// Tool instance (for state access). pub tool: Arc, - /// Response-local execution context for this call. + /// Execution context. The Engine binds durable identity after pre-tool policy. pub context: ToolExecutionContext, } @@ -324,7 +324,7 @@ pub struct ToolResultInfo<'a, A = ()> { pub meta: ToolMeta, /// Tool instance (for state access). pub tool: Arc, - /// Response-local execution context for this call. + /// Execution context. The Engine binds durable identity after pre-tool policy. pub context: ToolExecutionContext, } @@ -344,6 +344,14 @@ pub struct ToolResultInfo<'a, A = ()> { /// flows, permission checks, and other trusted host adaptation. #[async_trait] pub trait Interceptor: Send + Sync { + /// Pure projection of a committed ToolCall's host annotation into an opaque, + /// durable call identity. Agen neither interprets storage metadata nor derives + /// this identity from provider IDs. Early pre-tool policy runs before commit; + /// the Engine binds this value after policy, before execution. + fn logical_tool_call_id(&self, _annotation: &A) -> Option { + None + } + /// Whether side effects may start before [`Self::on_assistant_turn_end`]. /// /// The conservative default preserves the historical whole-response gate. diff --git a/crates/agen/src/tests/tool_call_identity.rs b/crates/agen/src/tests/tool_call_identity.rs new file mode 100644 index 000000000..73419dcd4 --- /dev/null +++ b/crates/agen/src/tests/tool_call_identity.rs @@ -0,0 +1,250 @@ +//! Engine component tests with a recording tool and scripted host annotations. +//! Persistence is represented by the append callback; Server/IPC storage is not +//! exercised here. Run with `cargo test -p agen --lib tool_call_identity_tests`. + +use super::*; +use async_trait::async_trait; +use std::sync::Mutex; + +#[derive(Clone)] +struct NoLlm; + +#[async_trait] +impl LlmClient for NoLlm { + async fn stream(&self, _: Request) -> Result { + panic!("identity tests must not request a model response") + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +struct HostIdentity { + early: bool, +} + +#[async_trait] +impl Interceptor for HostIdentity { + fn logical_tool_call_id(&self, annotation: &String) -> Option { + Some(annotation.clone()) + } + + async fn pre_tool_call( + &self, + info: &mut ToolCallInfo<'_, String>, + ) -> Result { + if self.early { + assert!(info.history.is_empty()); + } + // Pre-tool hooks may replace the context, but cannot replace persisted identity. + info.context.logical_call_id = Some("hook-forged-identity".into()); + Ok(PreToolAction::Continue) + } +} + +struct RecordingTool(Arc>>); + +#[async_trait] +impl crate::tool::Tool for RecordingTool { + async fn execute( + &self, + _: &str, + context: ToolExecutionContext, + ) -> Result { + self.0.lock().unwrap().push(context); + Ok("recorded".to_string().into()) + } +} + +fn engine( + early: bool, +) -> ( + Engine, + Arc>>, +) { + let mut engine = Engine::new_annotated(NoLlm); + engine.set_interceptor(HostIdentity { early }); + engine.max_turns = Some(0); + engine.start_logical_run(); + let observed = Arc::new(Mutex::new(Vec::new())); + let tool = Arc::new(RecordingTool(observed.clone())); + engine.register_tool(Arc::new(move || { + (crate::tool::ToolMeta::new("record"), tool.clone()) + })); + engine.tool_server_handle().flush_pending(); + (engine, observed) +} + +fn entry(item: Item, id: &str) -> HistoryEntry { + HistoryEntry::new(item, id.to_string()) +} + +fn call(id: &str) -> Item { + Item::tool_call_json(id, "record", json!({})) +} + +#[test] +fn results_answer_only_preceding_occurrences_of_reused_provider_ids() { + let history = History::from_entries(vec![ + entry( + Item::tool_result("orphan", "no preceding call"), + "orphan-result", + ), + entry(call("reused"), "host-first"), + entry(call("sibling"), "host-sibling"), + entry( + Item::tool_result("sibling", "out of order"), + "sibling-result", + ), + entry(Item::tool_result("reused", "answered"), "first-result"), + entry(call("reused"), "host-pending"), + entry(call("orphan"), "host-after-orphan"), + ]); + let pending = unanswered_tool_call_entries(&history); + assert_eq!( + pending + .iter() + .map(|entry| entry.annotation.as_str()) + .collect::>(), + ["host-pending", "host-after-orphan"] + ); +} + +#[tokio::test] +async fn early_dispatch_binds_identity_only_after_successful_append() { + for reject_append in [false, true] { + let (mut engine, observed) = engine(true); + let mut history = History::new(); + let mut batch = EarlyToolExecutionBatch::new("live-batch".into(), 0); + let result = engine + .admit_early_tool_calls( + &mut history, + &mut |_| { + if reject_append { + Err("storage unavailable".into()) + } else { + Ok("host-committed-entry".into()) + } + }, + &mut batch, + vec![crate::timeline::CollectedToolCall { + call_index: 0, + call: ToolCall { + id: "provider-call".into(), + name: "record".into(), + input: json!({}), + }, + }], + ) + .await; + if reject_append { + assert!(matches!(result, Err(EngineError::HistoryAppend(_)))); + assert!(history.is_empty()); + assert!(batch.execution_handles.is_empty()); + assert!(observed.lock().unwrap().is_empty()); + } else { + assert!(matches!(result, Ok(EarlyToolAdmission::Continue))); + let (_, terminal) = batch.futures.next().await.unwrap(); + assert_eq!(terminal.tool_use_id, "provider-call"); + let observed = observed.lock().unwrap(); + assert_eq!(observed.len(), 1); + assert_eq!( + observed[0].logical_call_id.as_deref(), + Some(history.entries()[0].annotation.as_str()) + ); + assert_eq!( + observed[0].logical_call_id.as_deref(), + Some("host-committed-entry") + ); + assert_eq!(batch.call_info_map["provider-call"].3, observed[0]); + } + } +} + +#[tokio::test] +async fn normal_dispatch_and_unstarted_restore_use_the_exact_unanswered_annotation() { + for restore in [false, true] { + let (mut engine, observed) = engine(false); + let mut history = History::from_entries(vec![ + entry(call("reused"), "host-answered"), + entry( + Item::tool_result("reused", "already completed"), + "answered-result", + ), + ]); + engine + .append_history_items(&mut history, [call("reused")], &mut |_| { + Ok("host-pending".into()) + }) + .unwrap(); + if restore { + let persisted = serde_json::to_string(&history).unwrap(); + history = serde_json::from_str(&persisted).unwrap(); + } + let (pending, started) = engine.get_unanswered_tool_calls(&history); + assert_eq!(pending.len(), 1); + assert!(started.is_empty()); + if restore { + engine + .run_turn_loop(&mut history, &mut |_| Ok("terminal-result".into())) + .await + .unwrap(); + } else { + let result = engine + .execute_tools(&mut history, &mut |_| Ok("terminal-result".into()), pending) + .await + .unwrap(); + assert!(matches!(result, ToolExecutionResult::Completed)); + } + assert_eq!(observed.lock().unwrap().len(), 1); + let context = observed.lock().unwrap()[0].clone(); + assert_eq!(context.call_id, "reused"); + assert_eq!(context.logical_call_id.as_deref(), Some("host-pending")); + // A subsequent resume sees the terminal result, not another live attempt. + engine + .run_turn_loop(&mut history, &mut |_| Ok("unexpected-append".into())) + .await + .unwrap(); + assert_eq!(observed.lock().unwrap().len(), 1); + assert!(unanswered_tool_call_entries(&history).is_empty()); + } +} + +#[tokio::test] +async fn started_restore_with_reused_provider_id_closes_unknown_without_reexecution() { + let (mut engine, observed) = engine(false); + let mut history = History::from_entries(vec![ + entry(call("reused"), "host-answered"), + entry( + Item::tool_result("reused", "already completed"), + "answered-result", + ), + entry( + call("reused") + .with_tool_execution_metadata(0, Some("old-attempt".into())) + .with_status(ItemStatus::InProgress), + "host-started", + ), + ]); + let (pending, started) = engine.get_unanswered_tool_calls(&history); + assert!(pending.is_empty()); + assert_eq!(started, ["reused"]); + engine + .run_turn_loop(&mut history, &mut |_| Ok("unknown-result".into())) + .await + .unwrap(); + assert!(observed.lock().unwrap().is_empty()); + assert!(matches!( + &history.entries().last().unwrap().item, + Item::ToolResult { call_id, disposition: ToolResultDisposition::OutcomeUnknown, .. } + if call_id == "reused" + )); + let len = history.len(); + engine + .run_turn_loop(&mut history, &mut |_| Ok("unexpected-append".into())) + .await + .unwrap(); + assert_eq!(history.len(), len); + assert!(observed.lock().unwrap().is_empty()); +} diff --git a/crates/agen/src/tool.rs b/crates/agen/src/tool.rs index 29fbb026a..f621e0fcd 100644 --- a/crates/agen/src/tool.rs +++ b/crates/agen/src/tool.rs @@ -336,15 +336,16 @@ pub type ToolDefinition = Arc (ToolMeta, Arc) + Send + Syn /// Per-call context supplied by the engine when executing a tool call. /// -/// The context identifies a tool call within one assistant response's tool-call -/// batch without imposing any scheduling policy on the engine. Tool -/// implementations may use it for response-local ordering, diagnostics, or -/// correlation, but it is intentionally not a handle to engine state, history, -/// or session mutation. +/// Live batch identity and ordering are response-local. An optional opaque host +/// identity comes from the committed ToolCall annotation and survives restore. +/// Neither is a handle to engine state, history, or session mutation. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ToolExecutionContext { /// Provider/tool-call id for the call being executed. pub call_id: String, + /// Opaque host identity from the committed ToolCall annotation, stable across + /// restoration. Provider call IDs and live batch IDs are not replay authority. + pub logical_call_id: Option, /// Engine-local identity shared by all tool calls from one execution batch. pub batch_id: String, /// Zero-based order of this call in the model-returned tool-call list. @@ -355,11 +356,18 @@ impl ToolExecutionContext { pub fn new(call_id: impl Into, batch_id: impl Into, call_index: usize) -> Self { Self { call_id: call_id.into(), + logical_call_id: None, batch_id: batch_id.into(), call_index, } } + /// Bind an identity supplied by the host's persisted-call annotation. + pub fn with_logical_call_id(mut self, logical_call_id: impl Into) -> Self { + self.logical_call_id = Some(logical_call_id.into()); + self + } + /// Identifies one live execution attempt without making the batch id a durable /// replay or idempotency authority. pub fn execution_id(&self) -> String { @@ -572,7 +580,7 @@ pub trait Tool: Send + Sync { /// /// # Arguments /// * `input_json` - JSON-formatted arguments generated by LLM - /// * `ctx` - response-local call identity and ordering context + /// * `ctx` - per-call identity and live execution ordering context /// /// # Returns /// A [`ToolOutput`] with summary and optional detailed content. diff --git a/crates/worker/src/interrupt_prep.rs b/crates/worker/src/interrupt_prep.rs index b02a7284f..d3eb2fd1d 100644 --- a/crates/worker/src/interrupt_prep.rs +++ b/crates/worker/src/interrupt_prep.rs @@ -18,18 +18,25 @@ use agen::{Item, ToolResultDisposition}; /// Build synthetic `Item::ToolResult` items for every unanswered /// `Item::ToolCall` in `history`, preserving order. pub(crate) fn orphan_tool_result_closures(history: &[Item], summary: &str) -> Vec { - let mut answered: std::collections::HashSet<&str> = std::collections::HashSet::new(); + let mut pending: Vec> = Vec::new(); for item in history { - if let Item::ToolResult { call_id, .. } = item { - answered.insert(call_id.as_str()); + match item { + Item::ToolCall { call_id, .. } => pending.push(Some(call_id)), + Item::ToolResult { call_id, .. } => { + if let Some(index) = pending + .iter() + .position(|entry| *entry == Some(call_id.as_str())) + { + pending[index] = None; + } + } + _ => {} } } let mut out = Vec::new(); - for item in history { - if let Item::ToolCall { call_id, .. } = item { - if !answered.contains(call_id.as_str()) { - out.push(Item::tool_result_item_with_disposition_and_attachments( - call_id.clone(), + for call_id in pending.into_iter().flatten() { + out.push(Item::tool_result_item_with_disposition_and_attachments( + call_id, summary, Some( "Execution ended before completion could be confirmed. Completion and side effects are unknown." @@ -38,8 +45,6 @@ pub(crate) fn orphan_tool_result_closures(history: &[Item], summary: &str) -> Ve ToolResultDisposition::OutcomeUnknown, Vec::new(), )); - } - } } out } @@ -59,6 +64,20 @@ fn interrupt_tool_result_summary() -> String { mod tests { use super::*; + #[test] + fn answered_provider_id_does_not_hide_later_unanswered_logical_call() { + let history = vec![ + Item::tool_call("reused", "plugin", "{}"), + Item::tool_result("reused", "first finished"), + Item::tool_call("reused", "plugin", "{}"), + ]; + let closed = orphan_tool_result_closures(&history, "Interrupted"); + assert_eq!(closed.len(), 1); + assert!( + matches!(&closed[0], Item::ToolResult { call_id, disposition: ToolResultDisposition::OutcomeUnknown, .. } if call_id == "reused") + ); + } + #[test] fn no_orphans_returns_empty() { let history = vec![Item::user_message("hi"), Item::assistant_message("hello")]; diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 3f2e89209..f916725ff 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -437,6 +437,10 @@ impl WorkerInterceptor { #[async_trait] impl Interceptor for WorkerInterceptor { + fn logical_tool_call_id(&self, annotation: &SessionHistoryMetadata) -> Option { + Some(annotation.entry_id.0.clone()) + } + fn supports_early_tool_dispatch(&self) -> bool { self.registry.on_turn_end.is_empty() } diff --git a/crates/worker/src/plugin.rs b/crates/worker/src/plugin.rs index 70757e38b..aaf20d792 100644 --- a/crates/worker/src/plugin.rs +++ b/crates/worker/src/plugin.rs @@ -291,18 +291,20 @@ impl Tool for PluginTool { input_json: &str, context: ToolExecutionContext, ) -> Result { - if context.call_id.is_empty() || context.call_id == "direct" { - return Err(failure("Plugin Tool requires committed Tool-call identity")); - } + let logical_call_id = context + .logical_call_id + .as_deref() + .filter(|id| !id.is_empty()) + .ok_or_else(|| failure("Plugin Tool requires committed logical Tool-call identity"))?; let _capacity = self .capacity .clone() .try_acquire_owned() .map_err(|_| failure("Plugin invocation capacity exceeded"))?; let lease = self.start(&context)?; - // The persisted Session and provider call ID survive Runtime restoration; - // the engine-local batch ID identifies only cancellation of a live attempt. - let id_bytes = serde_json::to_vec(&(&self.session_id, &context.call_id)) + // Only the host-owned persisted call identity is durable. Provider IDs may + // repeat in a Session; batch/provider IDs identify live attempts only. + let id_bytes = serde_json::to_vec(&(&self.session_id, logical_call_id)) .expect("identity serialization"); let hex: String = Sha256::digest(&id_bytes) .iter() diff --git a/crates/worker/src/plugin/tests.rs b/crates/worker/src/plugin/tests.rs index a3088e56b..b97fea88e 100644 --- a/crates/worker/src/plugin/tests.rs +++ b/crates/worker/src/plugin/tests.rs @@ -320,6 +320,7 @@ fn install(feature: PluginToolsFeature) -> (Engine, HookRegi fn context(call: &str, batch: &str) -> ToolExecutionContext { ToolExecutionContext::new(call, batch, 0) + .with_logical_call_id(format!("persisted-entry-{call}")) } fn stopping() -> WorkerStoppingContext { @@ -502,15 +503,41 @@ async fn adapter_late_publish_denial_withholds_successful_guest_output() { } #[tokio::test] -async fn consumer_identity_survives_batch_changes_but_isolates_sessions_and_calls() { +async fn consumer_identity_uses_host_entry_not_provider_call_or_live_batch() { let fixture = Fixture::new(); let owner = Owner::new(Behavior::Allow); // Re-preparation simulates a restored adapter, not just reuse of a live Tool. - for (session, call, batch) in [ - ("session-a", "provider-call", "batch-before-restore"), - ("session-a", "provider-call", "batch-after-restore"), - ("session-b", "provider-call", "batch-after-restore"), - ("session-a", "different-call", "batch-after-restore"), + for (session, call, logical, batch) in [ + ( + "session-a", + "provider-call", + "host-entry-a", + "batch-before-restore", + ), + ( + "session-a", + "provider-call", + "host-entry-a", + "batch-after-restore", + ), + ( + "session-b", + "provider-call", + "host-entry-a", + "batch-after-restore", + ), + ( + "session-a", + "provider-call", + "host-entry-b", + "batch-after-restore", + ), + ( + "session-a", + "different-provider-call", + "host-entry-a", + "different-batch", + ), ] { let feature = fixture .prepare(&["installation_a"], owner.clone(), session) @@ -518,7 +545,11 @@ async fn consumer_identity_survives_batch_changes_but_isolates_sessions_and_call let (engine, _) = install(feature); engine .tool_server_handle() - .call_tool(TOOL_A, INPUT, context(call, batch)) + .call_tool( + TOOL_A, + INPUT, + ToolExecutionContext::new(call, batch, 0).with_logical_call_id(logical), + ) .await .unwrap(); } @@ -526,6 +557,34 @@ async fn consumer_identity_survives_batch_changes_but_isolates_sessions_and_call assert_eq!(consumers[0], consumers[1]); assert_ne!(consumers[0], consumers[2]); assert_ne!(consumers[0], consumers[3]); + assert_eq!(consumers[0], consumers[4]); +} + +#[tokio::test] +async fn missing_or_empty_logical_identity_refuses_owner_admission() { + let fixture = Fixture::new(); + let owner = Owner::new(Behavior::Allow); + let feature = fixture + .prepare(&["installation_a"], owner.clone(), "session-a") + .await; + let (engine, _) = install(feature); + let tools = engine.tool_server_handle(); + for context in [ + ToolExecutionContext::direct(), + ToolExecutionContext::new("provider-call", "batch", 0), + ToolExecutionContext::new("provider-call", "batch", 0).with_logical_call_id(""), + ] { + let error = tools.call_tool(TOOL_A, INPUT, context).await.unwrap_err(); + assert!( + error + .to_string() + .contains("committed logical Tool-call identity") + ); + } + assert!( + owner.requests().is_empty(), + "identity rejection must precede Begin" + ); } struct FailInstall; diff --git a/crates/workspace-server/src/server_plugin_controller_tests.rs b/crates/workspace-server/src/server_plugin_controller_tests.rs index 416afb246..bcb181a1a 100644 --- a/crates/workspace-server/src/server_plugin_controller_tests.rs +++ b/crates/workspace-server/src/server_plugin_controller_tests.rs @@ -13,13 +13,14 @@ // Run: cargo test -p yoi-workspace-server --lib plugin_controller_tests // Prerequisites: openssl and the normal connected-restore test environment. mod plugin_controller_tests { - use super::plugin_host_tests::plugin_runtime_tests::{Model, assert_https_history}; + use super::plugin_host_tests::plugin_runtime_tests::{ + Model, assert_https_history, persisted_call_consumers, + }; use super::plugin_tests::request; use super::*; use crate::plugin_https::test_support::{HOST, Peer, response}; use crate::plugin_plan::PlanPinState; use session_store::Store; - use sha2::Digest; use std::sync::{ Mutex, atomic::{AtomicUsize, Ordering}, @@ -41,6 +42,10 @@ mod plugin_controller_tests { impl Fixture { async fn new(peer: &Peer) -> Self { + Self::with_request_count(peer, 1).await + } + + async fn with_request_count(peer: &Peer, request_count: usize) -> Self { // Use the established short-path/environment guard for real restore // (not a mock controller or a replacement execution backend). let directory = super::restored_parent_control_tests::connected_restore_tempdir(); @@ -57,7 +62,7 @@ mod plugin_controller_tests { let clients = model_clients.clone(); let observed = requests.clone(); let tool = model_tool.clone(); - let transport = Arc::new(peer.transport()); + let transport = Arc::new(peer.transport_for_requests(request_count)); let api = WorkspaceApi::new_with_embedded_configurators( config, store, @@ -345,6 +350,14 @@ mod plugin_controller_tests { &self, worker: &RuntimeWorkerRef, ) -> (session_store::SessionId, Vec) { + let (session, saved) = self.saved_state(worker); + (session, saved.history) + } + + fn saved_state( + &self, + worker: &RuntimeWorkerRef, + ) -> (session_store::SessionId, session_store::RestoredState) { let reopened = session_store::WorkerSessionStore::new(self.aggregate(worker).join("session")) .unwrap(); @@ -352,7 +365,7 @@ mod plugin_controller_tests { let segments = reopened.list_segments(session).unwrap(); assert_eq!(segments.len(), 1); let saved = session_store::restore(&reopened, session, segments[0]).unwrap(); - (session, saved.history) + (session, saved) } fn count(&self, table: &str) -> i64 { @@ -444,10 +457,14 @@ mod plugin_controller_tests { } async fn submit(connection: &mut WorkspaceWorkerProtocolConnection) { + submit_id(connection, "controller-https-submit").await; + } + + async fn submit_id(connection: &mut WorkspaceWorkerProtocolConnection, request_id: &str) { connection .methods .send(protocol::Method::Submit { - submission_request_id: "controller-https-submit".into(), + submission_request_id: request_id.into(), input: vec![protocol::Segment::Text { content: "Check the service status.".into(), }], @@ -464,6 +481,14 @@ mod plugin_controller_tests { async fn finish( connection: &mut WorkspaceWorkerProtocolConnection, expected: protocol::RunResult, + ) { + finish_id(connection, expected, "controller-https-submit").await; + } + + async fn finish_id( + connection: &mut WorkspaceWorkerProtocolConnection, + expected: protocol::RunResult, + request_id: &str, ) { tokio::time::timeout(WATCHDOG, async { let mut accepted = false; @@ -474,7 +499,7 @@ mod plugin_controller_tests { submission_request_id, .. } => { - assert_eq!(submission_request_id, "controller-https-submit"); + assert_eq!(submission_request_id, request_id); accepted = true; } protocol::Event::SubmissionRejected { message, .. } => { @@ -489,7 +514,7 @@ mod plugin_controller_tests { if finished && snapshot.state == protocol::WorkerState::Idle && snapshot.last_finished_submission_request_id.as_deref() - == Some("controller-https-submit") => + == Some(request_id) => { return; } @@ -501,6 +526,116 @@ mod plugin_controller_tests { .expect("real Controller HTTPS submission did not finish"); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn repeated_provider_call_id_has_distinct_durable_receipts_and_completed_restore_never_reruns() + { + // Gemini also reuses a response-local id for the same Tool. Deliberately + // emit the identical provider id in two explicit user submissions. + let peer = Peer::new( + HOST, + vec![response(503, &[], b"old"), response(200, &[], b"fresh")], + ) + .await; + let f = Fixture::with_request_count(&peer, 2).await; + f.install(&peer).await; + let worker = f.create(PROFILE).await; + let public_name = f.model_tool.lock().unwrap().clone(); + let mut connection = connect_workspace_worker_protocol(&f.api, &worker, None) + .await + .unwrap(); + snapshot(&mut connection).await; + for request in ["first-explicit-call", "second-explicit-call"] { + submit_id(&mut connection, request).await; + finish_id(&mut connection, protocol::RunResult::Finished, request).await; + } + assert_eq!( + peer.connection_count(), + 2, + "a new logical call must not replay the first receipt; model requests: {:?}", + f.requests.lock().unwrap().iter().map(|r| &r.items).collect::>() + ); + assert_eq!(f.count("plugin_invocations"), 2); + assert_eq!(f.count("plugin_effect_receipts"), 2); + let statuses = |history: &[agen::Item]| -> Vec { + history + .iter() + .filter_map(|item| match item { + agen::Item::ToolResult { + call_id, + is_error: false, + content: Some(content), + .. + } if call_id == "https-call" => Some( + serde_json::from_str::(content).unwrap()["structured"]["status"] + .as_u64() + .unwrap() as u16, + ), + _ => None, + }) + .collect() + }; + { + let requests = f.requests.lock().unwrap(); + assert_eq!(requests.len(), 4); + assert_eq!(statuses(&requests[3].items), [503, 200]); + } + drop(connection); + f.stop(&worker).await; + let (session, saved) = f.saved_state(&worker); + assert_eq!(statuses(&saved.history), [503, 200]); + let calls: Vec<_> = saved + .annotated_history + .iter() + .filter_map(|entry| match &entry.item { + session_store::LoggedItem::ToolCall { call_id, name, .. } + if name == &public_name => + { + Some((call_id, &entry.metadata.entry_id)) + } + _ => None, + }) + .collect(); + assert_eq!(calls.len(), 2); + assert_eq!( + calls[0].0, calls[1].0, + "provider id must actually be reused in the fixture" + ); + assert_ne!(calls[0].1, calls[1].1); + let mut expected = persisted_call_consumers(session, &saved, &public_name); + expected.sort(); + let consumers: Vec = f.api.config_store.with_conn(|conn| { + let mut statement = conn.prepare("SELECT json_extract(consumer_json,'$.tool_call_id') FROM plugin_invocations WHERE workspace_id=?1 AND worker_id=?2 ORDER BY 1")?; + Ok(statement.query_map(rusqlite::params![TEST_WORKSPACE_ID, worker.worker_id], |row| row.get(0))?.collect::,_>>()?) + }).unwrap(); + assert_eq!( + consumers, expected, + "receipts must bind committed Host annotation ids, not response-local ids" + ); + let restored = f.restore(&worker, "completed-reused-id-restore").await; + assert_eq!( + restored.state, + server_api::WorkerRestoreState::Accepted, + "{restored:?}" + ); + assert_eq!(f.requests.lock().unwrap().len(), 4); + let mut connection = connect_workspace_worker_protocol(&f.api, &worker, None) + .await + .unwrap(); + assert!(snapshot(&mut connection).await.tools.contains(&public_name)); + assert_eq!(peer.connection_count(), 2); + assert_eq!(f.count("plugin_invocations"), 2); + assert_eq!(f.count("plugin_effect_receipts"), 2); + drop(connection); + f.stop(&worker).await; + let (restored_session, saved) = f.saved_state(&worker); + assert_eq!(restored_session, session); + let mut restored = persisted_call_consumers(session, &saved, &public_name); + restored.sort(); + assert_eq!(restored, expected); + assert_eq!(statuses(&saved.history), [503, 200]); + f.close().await; + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn owner_diagnostics_observes_saved_pin_and_live_refusal_without_execution_or_paths() { let peer = Peer::new(HOST, vec![]).await; @@ -735,14 +870,9 @@ mod plugin_controller_tests { rusqlite::params![TEST_WORKSPACE_ID, saved.plan_id, worker.worker_id], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)))?) }).unwrap(); - let identity = serde_json::to_vec(&(session.to_string(), "https-call")).unwrap(); - let consumer_id = format!( - "tool-{}", - sha2::Sha256::digest(identity) - .iter() - .map(|byte| format!("{byte:02x}")) - .collect::() - ); + let consumers = persisted_call_consumers(session, &f.saved_state(&worker).1, &public_name); + assert_eq!(consumers.len(), 1); + let consumer_id = consumers[0].clone(); assert_eq!( serde_json::from_str::(&consumer).unwrap(), server_api::PluginConsumer::ToolCall { diff --git a/crates/workspace-server/src/server_plugin_remote_controller_tests.rs b/crates/workspace-server/src/server_plugin_remote_controller_tests.rs index ea7bf4a1f..09768a865 100644 --- a/crates/workspace-server/src/server_plugin_remote_controller_tests.rs +++ b/crates/workspace-server/src/server_plugin_remote_controller_tests.rs @@ -7,7 +7,9 @@ // Run: cargo test -p yoi-workspace-server --lib plugin_remote_controller_tests // Prerequisites: openssl and the existing connected Runtime test environment. mod plugin_remote_controller_tests { - use super::plugin_host_tests::plugin_runtime_tests::{Model, assert_https_history}; + use super::plugin_host_tests::plugin_runtime_tests::{ + Model, assert_https_history, persisted_call_consumers, + }; use super::plugin_tests::request; use super::subject_authorization_tests::{CONNECTED_RUNTIME, ConnectedWorkerHarness}; use super::*; @@ -15,7 +17,6 @@ mod plugin_remote_controller_tests { use crate::plugin_plan::PlanPinState; use agen::llm_client::{ClientError, LlmClient, Request as ModelRequest}; use session_store::Store; - use sha2::Digest; use std::sync::Mutex; use std::time::Duration; @@ -380,14 +381,9 @@ mod plugin_remote_controller_tests { rusqlite::params![TEST_WORKSPACE_ID, saved.plan_id, worker.worker_id], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)))?) }).unwrap(); - let identity = serde_json::to_vec(&(session.to_string(), "https-call")).unwrap(); - let consumer_id = format!( - "tool-{}", - sha2::Sha256::digest(identity) - .iter() - .map(|byte| format!("{byte:02x}")) - .collect::() - ); + let consumers = persisted_call_consumers(session, &history, &public_name); + assert_eq!(consumers.len(), 1); + let consumer_id = consumers[0].clone(); assert_eq!( serde_json::from_str::(&consumer).unwrap(), server_api::PluginConsumer::ToolCall { diff --git a/crates/workspace-server/src/server_plugin_runtime_tests.rs b/crates/workspace-server/src/server_plugin_runtime_tests.rs index 0069b21b5..5d5bc117b 100644 --- a/crates/workspace-server/src/server_plugin_runtime_tests.rs +++ b/crates/workspace-server/src/server_plugin_runtime_tests.rs @@ -147,7 +147,7 @@ pub(super) mod plugin_runtime_tests { turn }; let events = match turn { - 0 => vec![ + 0 | 2 => vec![ ModelEvent::tool_use_start(0, "https-call", &self.tool_name), ModelEvent::tool_input_delta( 0, @@ -159,7 +159,7 @@ pub(super) mod plugin_runtime_tests { ), ModelEvent::tool_use_stop(0), ], - 1 => vec![ + 1 | 3 => vec![ ModelEvent::text_block_start(0), ModelEvent::text_delta(0, "Observed the HTTPS tool result."), ModelEvent::text_block_stop(0, None), @@ -177,6 +177,32 @@ pub(super) mod plugin_runtime_tests { } } + pub(in crate::server::tests) fn persisted_call_consumers( + session: session_store::SessionId, + saved: &session_store::RestoredState, + public_name: &str, + ) -> Vec { + saved + .annotated_history + .iter() + .filter_map(|entry| match &entry.item { + session_store::LoggedItem::ToolCall { name, .. } if name == public_name => { + let bytes = + serde_json::to_vec(&(session.to_string(), &entry.metadata.entry_id.0)) + .unwrap(); + Some(format!( + "tool-{}", + Sha256::digest(bytes) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + )) + } + _ => None, + }) + .collect() + } + pub(in crate::server::tests) fn assert_https_history(items: &[agen::Item], public_name: &str) { let calls = items .iter() @@ -234,7 +260,19 @@ pub(super) mod plugin_runtime_tests { assert_eq!(output["structured"], json!({"status":503,"body_bytes":9})); } - async fn https_through_runtime_and_history(remote: bool) { + struct PauseBeforeFirstCall; + + #[async_trait::async_trait] + impl worker::hook::Hook for PauseBeforeFirstCall { + async fn call( + &self, + _: &worker::hook::ToolCallSummary, + ) -> std::result::Result { + Ok(worker::hook::HookPreToolAction::Pause) + } + } + + async fn https_through_runtime_and_history(remote: bool, restore_pending: bool) { let mut peer = Peer::new(HOST, vec![response(503, &[], b"peer data")]).await; let mut f = Fixture::new(Some(peer_binding(&peer))).await; f.api.plugin_https_transport = Some(Arc::new(peer.transport())); @@ -263,17 +301,34 @@ pub(super) mod plugin_runtime_tests { tool_name: public_name.clone(), requests: Arc::default(), }; - let manifest = worker::WorkerManifest::from_toml(&format!( + let mut manifest = worker::WorkerManifest::from_toml(&format!( "[worker]\nname=\"{}\"\n[model]\nscheme=\"anthropic\"\nmodel_id=\"scripted\"\n[engine]\n[scope]\nallow=[]\n", f.source.worker_id )) .unwrap(); + if restore_pending { + manifest + .plugins + .selections + .push(manifest::PluginProfileSelection { + installation_id: f.installed.installation.installation_id.clone(), + tools: vec![f.tool.into()], + logging: false, + https_slots: vec!["api".into()], + config: Default::default(), + }); + } let store_root = f._directory.path().join("tool-history"); - let mut host = worker::Worker::new( - manifest, - agen::Engine::new_annotated(model.clone()) - .system_prompt("Use the selected HTTPS tool."), + let metadata_root = f._directory.path().join("worker-metadata"); + let store = session_store::CombinedStore::new( session_store::FsStore::new(&store_root).unwrap(), + session_store::FsWorkerStore::new(&metadata_root).unwrap(), + ); + let mut host = worker::Worker::new( + manifest.clone(), + agen::Engine::new_annotated(Box::new(model.clone()) as Box) + .system_prompt("Use the selected HTTPS tool."), + store, worker::WorkerWorkspaceContext::unavailable( Some(worker::WorkspaceId::new(TEST_WORKSPACE_ID).unwrap()), "unrelated Workspace operations are not bound by this fixture", @@ -330,6 +385,10 @@ pub(super) mod plugin_runtime_tests { tools.get_tool(f.tool).is_none(), "bare guest names are not public tools" ); + if restore_pending { + host.enable_worker_metadata_write_through().unwrap(); + host.add_pre_tool_call_hook(PauseBeforeFirstCall); + } host.wire_history_persistence(); let result = tokio::time::timeout( std::time::Duration::from_secs(30), @@ -340,7 +399,63 @@ pub(super) mod plugin_runtime_tests { .await .expect("HTTPS Tool/Engine must finish within the fixture deadline") .unwrap(); - assert!(matches!(result, worker::WorkerRunResult::Finished)); + if restore_pending { + assert!(matches!(result, worker::WorkerRunResult::Paused)); + assert_eq!(peer.connection_count(), 0); + assert_eq!(f.count("plugin_invocations"), 0); + let reopened = session_store::FsStore::new(&store_root).unwrap(); + let segments = reopened.list_segments(session_id).unwrap(); + let pending = session_store::restore(&reopened, session_id, segments[0]).unwrap(); + let expected = persisted_call_consumers(session_id, &pending, &public_name); + assert_eq!( + expected.len(), + 1, + "paused call must already have a durable Host identity" + ); + assert!( + !pending + .history + .iter() + .any(|item| matches!(item, agen::Item::ToolResult { .. })) + ); + drop(host); + let mut restored = worker::Worker::prepare_restored_from_worker_metadata_with_context_and_model_client( + &manifest.worker.name, manifest.clone(), session_store::CombinedStore::new( + reopened, session_store::FsWorkerStore::new(&metadata_root).unwrap(), + ), + worker::PromptCatalogSource::builtins_only(), + worker::WorkerWorkspaceContext::unavailable(Some(worker::WorkspaceId::new(TEST_WORKSPACE_ID).unwrap()), "unrelated Workspace operations unavailable"), + worker::WorkerFilesystemAuthority::None, Some(Box::new(model.clone())), + ).await.unwrap().activate().await.unwrap(); + assert_eq!(restored.session_id(), session_id); + let reloaded = session_store::restore( + &session_store::FsStore::new(&store_root).unwrap(), + session_id, + segments[0], + ) + .unwrap(); + assert_eq!( + persisted_call_consumers(session_id, &reloaded, &public_name), + expected + ); + let feature = bridge + .prepare(&worker_ref, &snapshot, &session_id.to_string()) + .await + .unwrap(); + let report = + restored.install_features(FeatureRegistryBuilder::new().with_module(feature)); + assert!(!report.has_errors(), "{}", report.error_message()); + restored.wire_history_persistence(); + let result = + tokio::time::timeout(std::time::Duration::from_secs(30), restored.resume()) + .await + .unwrap() + .unwrap(); + assert!(matches!(result, worker::WorkerRunResult::Finished)); + host = restored; + } else { + assert!(matches!(result, worker::WorkerRunResult::Finished)); + } assert_https_history(&host.history(), &public_name); { let requests = model.requests.lock().unwrap(); @@ -371,14 +486,9 @@ pub(super) mod plugin_runtime_tests { assert_eq!(segments.len(), 1); let saved = session_store::restore(&reopened, session_id, segments[0]).unwrap(); assert_https_history(&saved.history, &public_name); - let identity_bytes = serde_json::to_vec(&(session_id.to_string(), "https-call")).unwrap(); - let consumer_id = format!( - "tool-{}", - Sha256::digest(identity_bytes) - .iter() - .map(|byte| format!("{byte:02x}")) - .collect::() - ); + let consumers = persisted_call_consumers(session_id, &saved, &public_name); + assert_eq!(consumers.len(), 1); + let consumer_id = consumers[0].clone(); let (consumer, state, status, body): (String, String, i64, Vec) = f .api .config_store @@ -418,7 +528,8 @@ pub(super) mod plugin_runtime_tests { .call_tool( &public_name, "{}", - ToolExecutionContext::new("denied-call", "denied-batch", 0), + ToolExecutionContext::new("denied-call", "denied-batch", 0) + .with_logical_call_id("denied-persisted-call"), ) .await .is_err() @@ -444,12 +555,17 @@ pub(super) mod plugin_runtime_tests { #[tokio::test] async fn remote_pinned_https_tool_commits_normal_history_and_receipt_then_fences_live_denial() { - https_through_runtime_and_history(true).await; + https_through_runtime_and_history(true, false).await; } #[tokio::test] async fn embedded_pinned_https_tool_commits_normal_history_and_receipt_then_fences_live_denial() { - https_through_runtime_and_history(false).await; + https_through_runtime_and_history(false, false).await; + } + + #[tokio::test] + async fn pending_tool_call_restores_its_committed_host_identity_before_first_effect() { + https_through_runtime_and_history(false, true).await; } }