diff --git a/crates/worker/src/in_flight.rs b/crates/worker/src/in_flight.rs index a11fcf3e3..dec0a435b 100644 --- a/crates/worker/src/in_flight.rs +++ b/crates/worker/src/in_flight.rs @@ -1,4 +1,7 @@ -use std::sync::{Arc, Mutex, MutexGuard}; +use std::{ + collections::HashSet, + sync::{Arc, Mutex, MutexGuard}, +}; use protocol::{ CommandEvent, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, Event, @@ -22,6 +25,9 @@ pub struct InFlightEvents { pub(crate) struct InFlightInner { next_block_id: u64, blocks: Vec, + // The physical Workdir provider can be shared with scoped Internal Workers. + // Keep command telemetry only for tool calls observed in this Worker stream. + known_tool_call_ids: HashSet, commands: Vec, compaction: Option, } @@ -53,6 +59,7 @@ impl InFlightEvents { inner: Arc::new(Mutex::new(InFlightInner { next_block_id: 1, blocks: Vec::new(), + known_tool_call_ids: HashSet::new(), commands: Vec::new(), compaction: None, })), @@ -150,6 +157,7 @@ impl InFlightEvents { pub(crate) fn tool_call_start(&self, id: String, name: String) -> InFlightBlockId { let mut inner = self.lock(); let block_id = inner.next_id(); + inner.known_tool_call_ids.insert(id.clone()); inner.blocks.push(TrackedBlock::ToolCall { block_id, id: id.clone(), @@ -213,12 +221,26 @@ impl InFlightEvents { } pub(crate) fn publish_command_event(&self, event: CommandEvent) { - self.lock().apply_command_event(&event); + let mut inner = self.lock(); + if !inner.accepts_command_event(&event) { + return; + } + inner.apply_command_event(&event); + drop(inner); let _ = self.working_event_tx.send(Event::Command { event }); } pub(crate) fn replace_command_snapshot(&self, commands: Vec) { - self.lock().commands = commands; + let mut inner = self.lock(); + inner.commands = commands + .into_iter() + .filter(|command| { + command + .tool_call_id + .as_ref() + .is_some_and(|tool_call_id| inner.known_tool_call_ids.contains(tool_call_id)) + }) + .collect(); } /// Atomically update reconnect state and publish the matching live progress event. @@ -253,6 +275,23 @@ impl InFlightInner { .find(|block| block.block_id() == block_id) } + fn accepts_command_event(&self, event: &CommandEvent) -> bool { + match event { + CommandEvent::Started { + tool_call_id: Some(tool_call_id), + .. + } => self.known_tool_call_ids.contains(tool_call_id), + CommandEvent::Started { + tool_call_id: None, .. + } => false, + CommandEvent::Output { command_id, .. } + | CommandEvent::Terminal { command_id, .. } => self + .commands + .iter() + .any(|command| command.command_id == *command_id), + } + } + fn apply_command_event(&mut self, event: &CommandEvent) { match event { CommandEvent::Started { @@ -394,10 +433,14 @@ impl InFlightInner { } fn clear(&mut self) -> bool { - if self.blocks.is_empty() && self.commands.is_empty() { + if self.blocks.is_empty() + && self.known_tool_call_ids.is_empty() + && self.commands.is_empty() + { false } else { self.blocks.clear(); + self.known_tool_call_ids.clear(); self.commands.clear(); true } @@ -504,6 +547,20 @@ pub(crate) fn snapshot_from_guard(guard: &MutexGuard<'_, InFlightInner>) -> InFl mod tests { use super::*; + fn command_snapshot(command_id: &str, tool_call_id: Option<&str>) -> CommandSnapshot { + CommandSnapshot { + command_id: command_id.into(), + tool_call_id: tool_call_id.map(str::to_owned), + status: CommandStatus::Running, + started_at_ms: 100, + observed_at_ms: 100, + last_output_at_ms: None, + stdout: CommandStreamSlice::default(), + stderr: CommandStreamSlice::default(), + exit_code: None, + } + } + #[test] fn snapshot_boundary_does_not_duplicate_or_gap_delta_sent_after_subscribe() { let (working_event_tx, _) = broadcast::channel(16); @@ -706,6 +763,8 @@ mod tests { let (working_event_tx, _) = broadcast::channel(16); let mut rx = working_event_tx.subscribe(); let in_flight = InFlightEvents::new(working_event_tx); + in_flight.tool_call_start("tool-1".into(), "Bash".into()); + assert!(matches!(rx.try_recv().unwrap(), Event::ToolCallStart { .. })); in_flight.publish_command_event(CommandEvent::Started { command_id: "command-1".into(), tool_call_id: Some("tool-1".into()), @@ -752,6 +811,49 @@ mod tests { assert!(snapshot_from_guard(&guard).commands.is_empty()); } + #[test] + fn foreign_command_telemetry_is_not_published_or_restored() { + let (working_event_tx, _) = broadcast::channel(16); + let mut rx = working_event_tx.subscribe(); + let in_flight = InFlightEvents::new(working_event_tx); + in_flight.tool_call_start("own-tool".into(), "Bash".into()); + assert!(matches!(rx.try_recv().unwrap(), Event::ToolCallStart { .. })); + + in_flight.publish_command_event(CommandEvent::Started { + command_id: "foreign-command".into(), + tool_call_id: Some("foreign-tool".into()), + observed_at_ms: 100, + }); + in_flight.publish_command_event(CommandEvent::Output { + command_id: "foreign-command".into(), + stream: CommandStream::Stdout, + start_offset: 0, + end_offset: 6, + content: "leaked".into(), + observed_at_ms: 110, + }); + in_flight.publish_command_event(CommandEvent::Terminal { + command_id: "foreign-command".into(), + status: CommandStatus::Completed, + exit_code: Some(0), + stdout_end_offset: 6, + stderr_end_offset: 0, + observed_at_ms: 120, + }); + + assert!(rx.try_recv().is_err()); + assert!(snapshot_from_guard(&in_flight.snapshot_guard()).commands.is_empty()); + + in_flight.replace_command_snapshot(vec![ + command_snapshot("own-command", Some("own-tool")), + command_snapshot("foreign-command", Some("foreign-tool")), + command_snapshot("unattributed-command", None), + ]); + let snapshot = snapshot_from_guard(&in_flight.snapshot_guard()); + assert_eq!(snapshot.commands.len(), 1); + assert_eq!(snapshot.commands[0].command_id, "own-command"); + } + #[test] fn compaction_progress_updates_snapshot_and_live_event_atomically() { let (working_event_tx, _) = broadcast::channel(16); diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index fe32f2842..0872affcc 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -162,6 +162,24 @@ fn simple_text_events() -> Vec { ] } +fn bash_tool_events(call_id: &str, command: &str) -> Vec { + vec![ + LlmEvent::tool_use_start(0, call_id, "Bash"), + LlmEvent::tool_input_delta( + 0, + serde_json::json!({ + "command": command, + "cwd": ".", + }) + .to_string(), + ), + LlmEvent::tool_use_stop(0), + LlmEvent::Status(StatusEvent { + status: ResponseStatus::Completed, + }), + ] +} + const MANIFEST_TOML: &str = r#" [worker] name = "test-worker" @@ -377,10 +395,10 @@ async fn shutdown_closes_bound_workdir_session() { } #[tokio::test] -async fn controller_projects_workdir_command_events_and_snapshot_state() { +async fn controller_filters_workdir_command_not_owned_by_worker() { let (mut worker, pwd) = make_worker_with_pwd(MockClient::new(simple_text_events())).await; let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::materialized_bound( - Workdir::new("controller-command-observation-workdir"), + Workdir::new("controller-foreign-command-workdir"), pwd.clone(), pwd, worker.scope().clone(), @@ -392,76 +410,26 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() { let command = session .start_command(CommandRequest { - command: "printf ready; sleep 0.3; printf done".to_owned(), + command: "sleep 0.2".to_owned(), timeout_secs: 5, output_limit: 1024, cwd: workdir::WorkdirPath::root(), spill_dir: None, - tool_call_id: Some("tool-command-1".into()), + tool_call_id: Some("foreign-tool-call".into()), }) .await .unwrap(); - let mut saw_started = false; - let mut saw_output = false; - while !saw_output { - let event = tokio::time::timeout(std::time::Duration::from_secs(2), events.recv()) + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), events.recv()) .await - .expect("command event should arrive") - .unwrap(); - match event { - Event::Command { - event: - protocol::CommandEvent::Started { - command_id, - tool_call_id, - .. - }, - } => { - assert_eq!(command_id, command.0); - assert_eq!(tool_call_id.as_deref(), Some("tool-command-1")); - saw_started = true; - } - Event::Command { - event: - protocol::CommandEvent::Output { - command_id, - stream: protocol::CommandStream::Stdout, - content, - .. - }, - } if command_id == command.0 && content.contains("ready") => saw_output = true, - _ => {} - } - } - assert!(saw_started); - + .is_err(), + "a command without a Worker-owned tool call must not enter its protocol stream" + ); let Event::Snapshot { in_flight, .. } = handle.snapshot_event() else { panic!("worker snapshot expected"); }; - assert_eq!(in_flight.commands.len(), 1); - assert_eq!(in_flight.commands[0].command_id, command.0); - assert_eq!(in_flight.commands[0].stdout.content, "ready"); - assert_eq!( - in_flight.commands[0].status, - protocol::CommandStatus::Running - ); - - let saw_terminal = drain_until(&mut events, std::time::Duration::from_secs(2), |event| { - matches!( - event, - Event::Command { - event: protocol::CommandEvent::Terminal { - command_id, - status: protocol::CommandStatus::Completed, - exit_code: Some(0), - .. - } - } if command_id == &command.0 - ) - }) - .await; - assert!(saw_terminal, "completed command event should arrive"); + assert!(in_flight.commands.is_empty()); let output = session .command_output(CommandOutputRequest { @@ -473,12 +441,104 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() { .await .unwrap(); assert_eq!(output.status, workdir::CommandStatus::Completed); - let (entries, _) = handle.sink.subscribe_with_snapshot(); - let durable_history = serde_json::to_string(&entries).unwrap(); - assert!( - !durable_history.contains("ready") && !durable_history.contains("done"), - "operational command chunks must not be appended to Worker history: {durable_history}" + handle + .send(Method::Shutdown { + command: worker_command(&handle), + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn controller_projects_workdir_command_events_and_snapshot_state() { + let client = MockClient::sequential(vec![ + MockResponse::Complete(bash_tool_events( + "tool-command-1", + "printf ready; sleep 0.3; printf done", + )), + MockResponse::Complete(simple_text_events()), + ]); + let (mut worker, pwd) = make_worker_with_pwd(client).await; + let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::materialized_bound( + Workdir::new("controller-command-observation-workdir"), + pwd.clone(), + pwd, + worker.scope().clone(), + WorkdirSessionCapabilities::ALL, + )); + worker.bind_single_workdir_session(Some(Arc::clone(&session))); + let handle = spawn_controller(worker).await; + let mut events = handle.subscribe(); + handle + .send(Method::submit_text( + protocol::new_submission_request_id(), + "run the command", + )) + .await + .unwrap(); + + let mut command_id = None; + let mut saw_output = false; + while !saw_output { + let event = tokio::time::timeout(std::time::Duration::from_secs(2), events.recv()) + .await + .expect("command event should arrive") + .unwrap(); + match event { + Event::Command { + event: + protocol::CommandEvent::Started { + command_id: started_id, + tool_call_id, + .. + }, + } => { + assert_eq!(tool_call_id.as_deref(), Some("tool-command-1")); + command_id = Some(started_id); + } + Event::Command { + event: + protocol::CommandEvent::Output { + command_id: output_id, + stream: protocol::CommandStream::Stdout, + content, + .. + }, + } if command_id.as_deref() == Some(output_id.as_str()) && content.contains("ready") => { + saw_output = true; + } + _ => {} + } + } + let command_id = command_id.expect("owned command start should be projected"); + + let Event::Snapshot { in_flight, .. } = handle.snapshot_event() else { + panic!("worker snapshot expected"); + }; + assert_eq!(in_flight.commands.len(), 1); + assert_eq!(in_flight.commands[0].command_id, command_id); + assert_eq!(in_flight.commands[0].stdout.content, "ready"); + assert_eq!( + in_flight.commands[0].status, + protocol::CommandStatus::Running ); + + let saw_terminal = drain_until(&mut events, std::time::Duration::from_secs(2), |event| { + matches!( + event, + Event::Command { + event: protocol::CommandEvent::Terminal { + command_id: terminal_id, + status: protocol::CommandStatus::Completed, + exit_code: Some(0), + .. + } + } if terminal_id == &command_id + ) + }) + .await; + assert!(saw_terminal, "completed command event should arrive"); + handle .send(Method::Shutdown { command: worker_command(&handle), @@ -489,7 +549,12 @@ async fn controller_projects_workdir_command_events_and_snapshot_state() { #[tokio::test] async fn controller_refreshes_command_snapshot_after_high_output_provider_lag() { - let (mut worker, pwd) = make_worker_with_pwd(MockClient::new(simple_text_events())).await; + let command_text = "dd if=/dev/zero bs=8192 count=300 2>/dev/null | tr '\\0' x; sleep 5"; + let client = MockClient::sequential(vec![ + MockResponse::Complete(bash_tool_events("tool-high-output", command_text)), + MockResponse::Complete(simple_text_events()), + ]); + let (mut worker, pwd) = make_worker_with_pwd(client).await; let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::materialized_bound( Workdir::new("controller-command-lag-recovery-workdir"), pwd.clone(), @@ -499,23 +564,37 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag() )); worker.bind_single_workdir_session(Some(Arc::clone(&session))); let handle = spawn_controller(worker).await; + let mut events = handle.subscribe(); + handle + .send(Method::submit_text( + protocol::new_submission_request_id(), + "run the high-output command", + )) + .await + .unwrap(); + + let command_id = loop { + let event = tokio::time::timeout(std::time::Duration::from_secs(2), events.recv()) + .await + .expect("owned command start should arrive") + .unwrap(); + if let Event::Command { + event: + protocol::CommandEvent::Started { + command_id, + tool_call_id: Some(tool_call_id), + .. + }, + } = event + && tool_call_id == "tool-high-output" + { + break command_id; + } + }; // Local command telemetry uses 8 KiB chunks and a 256-event channel. One // synchronous file-poll burst with 300 chunks deterministically makes the // worker-side receiver observe `Lagged` before this command terminates. - let command = session - .start_command(CommandRequest { - command: "dd if=/dev/zero bs=8192 count=300 2>/dev/null | tr '\\0' x; sleep 5" - .to_owned(), - timeout_secs: 10, - output_limit: 1024, - cwd: workdir::WorkdirPath::root(), - spill_dir: None, - tool_call_id: Some("tool-high-output".into()), - }) - .await - .unwrap(); - let expected_end_offset = 300_u64 * 8192; let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(3); let recovered = loop { @@ -525,7 +604,7 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag() if let Some(snapshot) = in_flight .commands .iter() - .find(|snapshot| snapshot.command_id == command.0) + .find(|snapshot| snapshot.command_id == command_id) && snapshot.stdout.end_offset >= expected_end_offset { break snapshot.clone(); @@ -545,17 +624,8 @@ async fn controller_refreshes_command_snapshot_after_high_output_provider_lag() assert!(recovered.stdout.content.len() <= 32 * 1024); assert!(recovered.stdout.content.bytes().all(|byte| byte == b'x')); - session.cancel_command(command.clone()).await.unwrap(); - let output = session - .command_output(CommandOutputRequest { - handle: command, - cursor: 0, - limit: 1024, - wait: true, - }) - .await - .unwrap(); - assert_eq!(output.status, workdir::CommandStatus::Cancelled); + let command = workdir::CommandHandle(command_id); + session.cancel_command(command).await.unwrap(); handle .send(Method::Shutdown { command: worker_command(&handle),