diff --git a/crates/tui/src/core/engine/turn_loop.rs b/crates/tui/src/core/engine/turn_loop.rs index 956f9aaa8d..347c7b613a 100644 --- a/crates/tui/src/core/engine/turn_loop.rs +++ b/crates/tui/src/core/engine/turn_loop.rs @@ -2436,6 +2436,7 @@ impl Engine { let outcomes = self .execute_planned_tools( plans, + &turn.id, ¤t_text_visible, &tool_catalog, &mut active_tool_names, @@ -3165,6 +3166,7 @@ impl Engine { async fn execute_planned_tools( &mut self, plans: Vec, + origin_turn_id: &str, current_text_visible: &str, tool_catalog: &[crate::models::Tool], active_tool_names: &mut std::collections::HashSet, @@ -3313,7 +3315,9 @@ impl Engine { continue; } - let batch_tool_context = self.live_tool_context(tool_registry); + let batch_tool_context = self + .live_tool_context(tool_registry) + .map(|context| context.with_origin_turn_id(origin_turn_id)); if parallel_allowed { let parallel_plan_receipts: Vec<_> = plans @@ -3371,7 +3375,8 @@ impl Engine { let started_at = Instant::now(); let shell_permits = shell_permits.clone(); let workspace = self.session.workspace.clone(); - let context_override = batch_tool_context.clone(); + let context_override = + tool_context_for_call(batch_tool_context.clone(), &plan.id); let cancel_token = self.cancel_token.clone(); tool_tasks.push(async move { @@ -3553,7 +3558,7 @@ impl Engine { tool_input.clone(), tool_registry, tool_exec_lock.clone(), - batch_tool_context.clone(), + tool_context_for_call(batch_tool_context.clone(), &tool_id), ) => match result { Ok(rich) => ( ToolExecutionOutcome::from_legacy(Ok(rich.result)), @@ -3871,7 +3876,10 @@ impl Engine { self.session.workspace.clone(), tool_registry, mcp_pool.clone(), - context_override.or_else(|| batch_tool_context.clone()), + tool_context_for_call( + context_override.or_else(|| batch_tool_context.clone()), + &tool_id, + ), ) => (result, false), } }; @@ -4988,6 +4996,13 @@ impl Engine { } } +fn tool_context_for_call( + context: Option, + tool_call_id: &str, +) -> Option { + context.map(|context| context.with_origin_tool_call_id(tool_call_id)) +} + pub(super) fn shell_completion_status_text( events: &[crate::tools::shell::ShellCompletionEvent], timing: &str, @@ -5834,6 +5849,18 @@ mod tests { use std::time::Duration; use tempfile::tempdir; + #[test] + fn tool_context_for_call_preserves_turn_and_sets_call_origin() { + let context = crate::tools::ToolContext::new(".").with_origin_turn_id("turn-origin"); + + let context = tool_context_for_call(Some(context), "tool-origin") + .expect("tool context remains available"); + + assert_eq!(context.origin_turn_id.as_deref(), Some("turn-origin")); + assert_eq!(context.origin_tool_call_id.as_deref(), Some("tool-origin")); + assert!(tool_context_for_call(None, "tool-origin").is_none()); + } + #[tokio::test] async fn child_owned_background_completion_is_not_delivered_to_parent() { let tmp = tempdir().expect("tempdir"); @@ -6048,6 +6075,8 @@ mod tests { linked_task_id: Some("task_1".to_string()), owner_agent_id: Some("agent_verifier".to_string()), owner_agent_name: Some("verifier".to_string()), + origin_tool_call_id: Some("tool_abc".to_string()), + origin_turn_id: Some("turn_abc".to_string()), owner_session_id: "session-test".to_string(), }], "", @@ -6072,6 +6101,8 @@ mod tests { linked_task_id: Some("task_1".to_string()), owner_agent_id: Some("agent_verifier".to_string()), owner_agent_name: Some("verifier".to_string()), + origin_tool_call_id: Some("tool_abc".to_string()), + origin_turn_id: Some("turn_abc".to_string()), owner_session_id: "session-test".to_string(), }, ]); @@ -6089,6 +6120,8 @@ mod tests { assert!(text.contains("art_shell_abc")); assert!(text.contains("cargo test -p codewhale-tui")); assert!(text.contains("test failed")); + assert!(text.contains(r#""origin_tool_call_id":"tool_abc""#)); + assert!(text.contains(r#""origin_turn_id":"turn_abc""#)); } #[test] diff --git a/crates/tui/src/runtime_handoff.rs b/crates/tui/src/runtime_handoff.rs index bdd8b66a74..31d7f240c7 100644 --- a/crates/tui/src/runtime_handoff.rs +++ b/crates/tui/src/runtime_handoff.rs @@ -210,6 +210,8 @@ pub(crate) fn shell_completion_runtime_message( "linked_task_id": event.linked_task_id, "owner_agent_id": event.owner_agent_id, "owner_agent_name": event.owner_agent_name, + "origin_tool_call_id": event.origin_tool_call_id, + "origin_turn_id": event.origin_turn_id, }) .to_string() }) diff --git a/crates/tui/src/tools/shell.rs b/crates/tui/src/tools/shell.rs index 976789c954..87ea13543c 100644 --- a/crates/tui/src/tools/shell.rs +++ b/crates/tui/src/tools/shell.rs @@ -220,6 +220,10 @@ pub struct ShellJobSnapshot { pub owner_agent_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub owner_agent_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub origin_tool_call_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub origin_turn_id: Option, /// Immutable root session that launched the job. Empty legacy records are /// intentionally hidden from session-scoped completion drains. #[serde(default, skip_serializing_if = "String::is_empty")] @@ -247,6 +251,10 @@ pub struct ShellCompletionEvent { pub owner_agent_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub owner_agent_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub origin_tool_call_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub origin_turn_id: Option, #[serde(default, skip_serializing_if = "String::is_empty")] pub owner_session_id: String, } @@ -302,6 +310,8 @@ impl ShellCompletionEvidence { "status": format!("{:?}", self.event.status), "exit_code": self.event.exit_code, "duration_ms": self.event.duration_ms, + "origin_tool_call_id": self.event.origin_tool_call_id, + "origin_turn_id": self.event.origin_turn_id, "stdout": stream(&self.stdout, self.stdout_omitted), "stderr": stream(&self.stderr, self.stderr_omitted), }) @@ -936,6 +946,8 @@ pub struct BackgroundShell { pub linked_task_id: Option, pub owner_agent: Option, owner_session_id: String, + origin_tool_call_id: Option, + origin_turn_id: Option, ownership: ShellOwnership, stdout_buffer: SharedRawOutput, stderr_buffer: Option, @@ -1023,6 +1035,8 @@ struct ShellSpawnIntentGuard { struct ShellSpawnContext { owner_agent: Option, owner_session_id: String, + origin_tool_call_id: Option, + origin_turn_id: Option, work_lifecycle: Option, } @@ -1524,6 +1538,8 @@ impl BackgroundShell { .owner_agent .as_ref() .map(|owner| owner.agent_name.clone()), + origin_tool_call_id: self.origin_tool_call_id.clone(), + origin_turn_id: self.origin_turn_id.clone(), owner_session_id: self.owner_session_id.clone(), } } @@ -1558,6 +1574,8 @@ impl BackgroundShell { linked_task_id: snapshot.linked_task_id, owner_agent_id: snapshot.owner_agent_id, owner_agent_name: snapshot.owner_agent_name, + origin_tool_call_id: snapshot.origin_tool_call_id, + origin_turn_id: snapshot.origin_turn_id, owner_session_id: snapshot.owner_session_id, } } @@ -1791,6 +1809,8 @@ impl ShellManager { linked_task_id: None, owner_agent: None, owner_session_id: String::new(), + origin_tool_call_id: None, + origin_turn_id: None, ownership: ShellOwnership::Managed, stdout_buffer: new_shared_raw_output(), stderr_buffer: Some(new_shared_raw_output()), @@ -1922,6 +1942,8 @@ impl ShellManager { owner_session_id.to_string(), None, None, + None, + None, false, (1_000, 600_000), ) @@ -1956,6 +1978,8 @@ impl ShellManager { String::new(), None, None, + None, + None, false, (1_000, 600_000), ) @@ -1990,6 +2014,8 @@ impl ShellManager { owner_session_id.to_string(), None, None, + None, + None, false, (1_000, 600_000), ) @@ -2009,6 +2035,8 @@ impl ShellManager { extra_env: HashMap, owner_agent: Option, owner_session_id: String, + origin_tool_call_id: Option, + origin_turn_id: Option, work_lifecycle: Option, readonly_workspace: Option<&std::path::Path>, persist_pending: bool, @@ -2065,6 +2093,8 @@ impl ShellManager { ShellSpawnContext { owner_agent, owner_session_id, + origin_tool_call_id, + origin_turn_id, work_lifecycle, }, persist_pending, @@ -2421,6 +2451,8 @@ impl ShellManager { let ShellSpawnContext { owner_agent, owner_session_id, + origin_tool_call_id, + origin_turn_id, work_lifecycle, } = spawn_context; let task_id = format!("shell_{}", &Uuid::new_v4().to_string()[..8]); @@ -2619,6 +2651,8 @@ impl ShellManager { linked_task_id: None, owner_agent, owner_session_id, + origin_tool_call_id, + origin_turn_id, ownership: if persist_pending { ShellOwnership::PersistPending } else { @@ -3253,6 +3287,8 @@ impl ShellManager { linked_task_id, owner_agent_id: None, owner_agent_name: None, + origin_tool_call_id: None, + origin_turn_id: None, owner_session_id: String::new(), }, ); @@ -4111,6 +4147,8 @@ async fn execute_foreground_via_background( extra_env, owner, context.state_namespace.clone(), + context.origin_tool_call_id.clone(), + context.origin_turn_id.clone(), lifecycle, direct_argv.then_some(context.workspace.as_path()), false, @@ -5129,6 +5167,8 @@ impl ToolSpec for BashTool { extra_env, shell_job_owner_from_context(context), context.state_namespace.clone(), + context.origin_tool_call_id.clone(), + context.origin_turn_id.clone(), shell_work_lifecycle_from_context(context), None, persist, diff --git a/crates/tui/src/tools/shell/tests.rs b/crates/tui/src/tools/shell/tests.rs index 6a04dd7a59..4c78aeee76 100644 --- a/crates/tui/src/tools/shell/tests.rs +++ b/crates/tui/src/tools/shell/tests.rs @@ -1568,9 +1568,12 @@ async fn background_start_advertises_task_status_completion() { } #[tokio::test] -async fn background_shell_job_carries_subagent_owner() { +async fn background_shell_job_preserves_origin_identity() { let tmp = tempdir().expect("tempdir"); - let ctx = ToolContext::new(tmp.path()).with_owner_agent("agent_owner", "verifier"); + let ctx = ToolContext::new(tmp.path()) + .with_origin_turn_id("turn-origin") + .with_origin_tool_call_id("tool-origin") + .with_owner_agent("agent_owner", "verifier"); let result = BashTool::new("Bash") .execute( json!({"command": sleep_command(2), "background": true}), @@ -1621,6 +1624,16 @@ async fn background_shell_job_carries_subagent_owner() { .expect("owned shell job snapshot"); assert_eq!(snapshot.owner_agent_id.as_deref(), Some("agent_owner")); assert_eq!(snapshot.owner_agent_name.as_deref(), Some("verifier")); + assert_eq!(snapshot.origin_tool_call_id.as_deref(), Some("tool-origin")); + assert_eq!(snapshot.origin_turn_id.as_deref(), Some("turn-origin")); + let mut legacy_json = serde_json::to_value(&snapshot).expect("serialize snapshot"); + let legacy_object = legacy_json.as_object_mut().expect("snapshot object"); + legacy_object.remove("origin_tool_call_id"); + legacy_object.remove("origin_turn_id"); + let legacy_snapshot: ShellJobSnapshot = + serde_json::from_value(legacy_json).expect("deserialize legacy snapshot"); + assert_eq!(legacy_snapshot.origin_tool_call_id, None); + assert_eq!(legacy_snapshot.origin_turn_id, None); let owners = manager.running_owner_agent_ids(); assert_eq!(owners, vec!["agent_owner".to_string()]); } @@ -1634,7 +1647,9 @@ async fn background_shell_job_carries_subagent_owner() { #[tokio::test] async fn drain_finished_jobs_reports_once() { let tmp = tempdir().expect("tempdir"); - let ctx = ToolContext::new(tmp.path()); + let ctx = ToolContext::new(tmp.path()) + .with_origin_turn_id("turn-origin") + .with_origin_tool_call_id("tool-origin"); let result = BashTool::new("Bash") .execute( json!({"command": echo_command("drain-finished-once"), "background": true}), @@ -1669,6 +1684,16 @@ async fn drain_finished_jobs_reports_once() { assert_eq!(first[0].task_id, task_id); assert_eq!(first[0].status, ShellStatus::Completed); assert!(first[0].stdout_tail.contains("drain-finished-once")); + assert_eq!(first[0].origin_tool_call_id.as_deref(), Some("tool-origin")); + assert_eq!(first[0].origin_turn_id.as_deref(), Some("turn-origin")); + let mut legacy_json = serde_json::to_value(&first[0]).expect("serialize completion"); + let legacy_object = legacy_json.as_object_mut().expect("completion object"); + legacy_object.remove("origin_tool_call_id"); + legacy_object.remove("origin_turn_id"); + let legacy_completion: ShellCompletionEvent = + serde_json::from_value(legacy_json).expect("deserialize legacy completion"); + assert_eq!(legacy_completion.origin_tool_call_id, None); + assert_eq!(legacy_completion.origin_turn_id, None); let second = manager.drain_finished_jobs_with_evidence(); assert!(second.is_empty(), "completion should be reported only once"); @@ -1757,6 +1782,8 @@ fn completion_evidence_preserves_arbitrary_stream_bytes() { linked_task_id: None, owner_agent_id: None, owner_agent_name: None, + origin_tool_call_id: Some("tool-origin".to_string()), + origin_turn_id: Some("turn-origin".to_string()), owner_session_id: "session-test".to_string(), }, stdout: stdout.clone(), @@ -1769,6 +1796,8 @@ fn completion_evidence_preserves_arbitrary_stream_bytes() { serde_json::from_slice(&evidence.artifact_bytes()).expect("evidence JSON"); assert_eq!(payload["stdout"]["encoding"], "base64"); assert_eq!(payload["stderr"]["encoding"], "base64"); + assert_eq!(payload["origin_tool_call_id"], "tool-origin"); + assert_eq!(payload["origin_turn_id"], "turn-origin"); let decoded_stdout = base64::engine::general_purpose::STANDARD .decode(payload["stdout"]["content"].as_str().expect("stdout data")) .expect("decode stdout"); @@ -3352,6 +3381,8 @@ fn killed_shell_does_not_wait_for_blocked_reader_threads() { ownership: ShellOwnership::Managed, linked_task_id: None, owner_agent: None, + origin_tool_call_id: None, + origin_turn_id: None, stdout_buffer: super::new_shared_raw_output(), stderr_buffer: None, heavy_permit: None, diff --git a/crates/tui/src/tools/spec.rs b/crates/tui/src/tools/spec.rs index ab2840b202..c71dc91045 100644 --- a/crates/tui/src/tools/spec.rs +++ b/crates/tui/src/tools/spec.rs @@ -593,6 +593,11 @@ pub struct ToolExecutionState { /// jobs can be attributed in UI surfaces. pub owner_agent_id: Option, pub owner_agent_name: Option, + /// Tool call and engine turn that created long-running work through this + /// context. Hosts use these stable identities to reconcile later updates + /// with the originating transcript position. + pub(crate) origin_tool_call_id: Option, + pub(crate) origin_turn_id: Option, /// Outer process authority cap installed by Fleet/headless dispatch. /// `None` for ordinary interactive/root sessions. pub(crate) tool_authority: Option>, @@ -768,6 +773,8 @@ impl ToolContext { file_read_tracker: new_shared_file_read_tracker(), owner_agent_id: None, owner_agent_name: None, + origin_tool_call_id: None, + origin_turn_id: None, tool_authority, trust_mode, sandbox_policy: SandboxPolicy::None, @@ -845,6 +852,20 @@ impl ToolContext { self } + /// Bind long-running work to the engine turn that created it. + #[must_use] + pub(crate) fn with_origin_turn_id(mut self, turn_id: impl Into) -> Self { + self.origin_turn_id = Some(turn_id.into()); + self + } + + /// Bind long-running work to the tool call that created it. + #[must_use] + pub(crate) fn with_origin_tool_call_id(mut self, tool_call_id: impl Into) -> Self { + self.origin_tool_call_id = Some(tool_call_id.into()); + self + } + #[cfg(test)] pub(crate) fn with_tool_authority( mut self, diff --git a/crates/tui/src/tui/shell_job_routing.rs b/crates/tui/src/tui/shell_job_routing.rs index fc2e83a27f..3d553b2029 100644 --- a/crates/tui/src/tui/shell_job_routing.rs +++ b/crates/tui/src/tui/shell_job_routing.rs @@ -169,6 +169,8 @@ mod tests { linked_task_id: Some("task_1".to_string()), owner_agent_id: None, owner_agent_name: None, + origin_tool_call_id: None, + origin_turn_id: None, owner_session_id: "session-test".to_string(), }]; let formatted = format_shell_job_list(&jobs); diff --git a/crates/tui/src/tui/ui/tests.rs b/crates/tui/src/tui/ui/tests.rs index aa30b8be76..61e9e123f3 100644 --- a/crates/tui/src/tui/ui/tests.rs +++ b/crates/tui/src/tui/ui/tests.rs @@ -7275,6 +7275,8 @@ fn shell_live_output_update_matches_exact_task_id_only() { linked_task_id: None, owner_agent_id: None, owner_agent_name: None, + origin_tool_call_id: None, + origin_turn_id: None, owner_session_id: "session-test".to_string(), }, ); @@ -7298,6 +7300,8 @@ fn shell_live_output_update_matches_exact_task_id_only() { linked_task_id: None, owner_agent_id: None, owner_agent_name: None, + origin_tool_call_id: None, + origin_turn_id: None, owner_session_id: "session-test".to_string(), }, ); @@ -7362,6 +7366,8 @@ fn shell_live_output_update_marks_stale_running_job_static() { linked_task_id: None, owner_agent_id: None, owner_agent_name: None, + origin_tool_call_id: None, + origin_turn_id: None, owner_session_id: "session-test".to_string(), }, ); @@ -7440,6 +7446,8 @@ fn shell_live_output_update_finalizes_background_exec_output() { linked_task_id: None, owner_agent_id: None, owner_agent_name: None, + origin_tool_call_id: None, + origin_turn_id: None, owner_session_id: "session-test".to_string(), }, ); @@ -7525,6 +7533,8 @@ fn shell_live_output_update_skips_finalized_exec_cell() { linked_task_id: None, owner_agent_id: None, owner_agent_name: None, + origin_tool_call_id: None, + origin_turn_id: None, owner_session_id: "session-test".to_string(), }, );