Skip to content

fix(shell): preserve task origin in job snapshots #5869

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
1 commit merged into from
Sep 5, 2026
Merged

fix(shell): preserve task origin in job snapshots #5869

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 37 additions & 4 deletions crates/tui/src/core/engine/turn_loop.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2436,6 +2436,7 @@ impl Engine {
let outcomes = self
.execute_planned_tools(
plans,
&turn.id,
&current_text_visible,
&tool_catalog,
&mut active_tool_names,
Expand Down Expand Up @@ -3165,6 +3166,7 @@ impl Engine {
async fn execute_planned_tools(
&mut self,
plans: Vec<ToolExecutionPlan>,
origin_turn_id: &str,
current_text_visible: &str,
tool_catalog: &[crate::models::Tool],
active_tool_names: &mut std::collections::HashSet<String>,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)),
Expand Down Expand Up @@ -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),
}
};
Expand Down Expand Up @@ -4988,6 +4996,13 @@ impl Engine {
}
}

fn tool_context_for_call(
context: Option<crate::tools::ToolContext>,
tool_call_id: &str,
) -> Option<crate::tools::ToolContext> {
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,
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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(),
}],
"",
Expand All @@ -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(),
},
]);
Expand All @@ -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]
Expand Down
2 changes: 2 additions & 0 deletions crates/tui/src/runtime_handoff.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
})
Expand Down
40 changes: 40 additions & 0 deletions crates/tui/src/tools/shell.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,10 @@ pub struct ShellJobSnapshot {
pub owner_agent_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub owner_agent_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub origin_tool_call_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub origin_turn_id: Option<String>,
/// 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")]
Expand Down Expand Up @@ -247,6 +251,10 @@ pub struct ShellCompletionEvent {
pub owner_agent_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub owner_agent_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub origin_tool_call_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub origin_turn_id: Option<String>,
#[serde(default, skip_serializing_if = "String::is_empty")]
pub owner_session_id: String,
}
Expand Down Expand Up @@ -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),
})
Expand Down Expand Up @@ -936,6 +946,8 @@ pub struct BackgroundShell {
pub linked_task_id: Option<String>,
pub owner_agent: Option<ShellJobOwner>,
owner_session_id: String,
origin_tool_call_id: Option<String>,
origin_turn_id: Option<String>,
ownership: ShellOwnership,
stdout_buffer: SharedRawOutput,
stderr_buffer: Option<SharedRawOutput>,
Expand Down Expand Up @@ -1023,6 +1035,8 @@ struct ShellSpawnIntentGuard {
struct ShellSpawnContext {
owner_agent: Option<ShellJobOwner>,
owner_session_id: String,
origin_tool_call_id: Option<String>,
origin_turn_id: Option<String>,
work_lifecycle: Option<ShellWorkLifecycle>,
}

Expand Down Expand Up @@ -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(),
}
}
Expand Down Expand Up @@ -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,
}
}
Expand Down Expand Up @@ -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()),
Expand Down Expand Up @@ -1922,6 +1942,8 @@ impl ShellManager {
owner_session_id.to_string(),
None,
None,
None,
None,
false,
(1_000, 600_000),
)
Expand Down Expand Up @@ -1956,6 +1978,8 @@ impl ShellManager {
String::new(),
None,
None,
None,
None,
false,
(1_000, 600_000),
)
Expand Down Expand Up @@ -1990,6 +2014,8 @@ impl ShellManager {
owner_session_id.to_string(),
None,
None,
None,
None,
false,
(1_000, 600_000),
)
Expand All @@ -2009,6 +2035,8 @@ impl ShellManager {
extra_env: HashMap<String, String>,
owner_agent: Option<ShellJobOwner>,
owner_session_id: String,
origin_tool_call_id: Option<String>,
origin_turn_id: Option<String>,
work_lifecycle: Option<ShellWorkLifecycle>,
readonly_workspace: Option<&std::path::Path>,
persist_pending: bool,
Expand Down Expand Up @@ -2065,6 +2093,8 @@ impl ShellManager {
ShellSpawnContext {
owner_agent,
owner_session_id,
origin_tool_call_id,
origin_turn_id,
work_lifecycle,
},
persist_pending,
Expand Down Expand Up @@ -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]);
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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(),
},
);
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
37 changes: 34 additions & 3 deletions crates/tui/src/tools/shell/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}),
Expand Down Expand Up @@ -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()]);
}
Expand All @@ -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}),
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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(),
Expand All @@ -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");
Expand Down Expand Up @@ -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,
Expand Down
Loading