Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

feat: integrate command output streaming

Browse Source
This commit is contained in:
Hare
2026-08-21 03:54:47 +09:00
parent 4b8dc302ee cfb173c570
commit 6fe36f3e46
16 changed files with 1614 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+123 -2
View File
@@ -28,8 +28,14 @@ use crate::worker::{
WorkerRunResult,
};
use protocol::{
AlertLevel, AlertSource, ErrorCode, Event, Method, RewindTargetId, RunResult, Segment,
TurnResult, WorkerStatus,
AlertLevel, AlertSource, CommandEvent as ProtocolCommandEvent,
CommandSnapshot as ProtocolCommandSnapshot, CommandStatus as ProtocolCommandStatus,
CommandStream as ProtocolCommandStream, CommandStreamSlice as ProtocolCommandStreamSlice,
ErrorCode, Event, Method, RewindTargetId, RunResult, Segment, TurnResult, WorkerStatus,
};
use workdir::{
CommandEvent as WorkdirCommandEvent, CommandSnapshot as WorkdirCommandSnapshot,
CommandStatus as WorkdirCommandStatus, CommandStream as WorkdirCommandStream, WorkdirSession,
};
// ---------------------------------------------------------------------------
@@ -424,6 +430,9 @@ impl WorkerController {
Some(method_tx.downgrade()),
)
.await?;
if let Some(session) = fs_for_view.as_ref() {
wire_workdir_command_events(session, &in_flight);
}
// Intake role Workers self-terminate only after a successful
// TicketIntakeReady turn has fully settled back to Idle. The request
@@ -498,6 +507,118 @@ impl WorkerController {
}
}
pub(crate) fn wire_workdir_command_events(
session: &Arc<dyn WorkdirSession>,
in_flight: &InFlightEvents,
) {
in_flight.replace_command_snapshot(
session
.command_snapshot()
.into_iter()
.map(protocol_command_snapshot)
.collect(),
);
let Some(mut events) = session.subscribe_command_events() else {
return;
};
let in_flight = in_flight.clone();
tokio::spawn(async move {
loop {
match events.recv().await {
Ok(event) => in_flight.publish_command_event(protocol_command_event(event)),
Err(broadcast::error::RecvError::Lagged(_)) => {
// Never retain stale command output after a provider-local
// observer lag. The next chunk reconstructs a bounded tail
// with its absolute offset and marks the gap truncated.
in_flight.replace_command_snapshot(Vec::new());
}
Err(broadcast::error::RecvError::Closed) => break,
}
}
});
}
fn protocol_command_snapshot(snapshot: WorkdirCommandSnapshot) -> ProtocolCommandSnapshot {
ProtocolCommandSnapshot {
command_id: snapshot.command_id,
tool_call_id: snapshot.tool_call_id,
status: protocol_command_status(snapshot.status),
started_at_ms: snapshot.started_at_ms,
observed_at_ms: snapshot.observed_at_ms,
last_output_at_ms: snapshot.last_output_at_ms,
stdout: ProtocolCommandStreamSlice {
start_offset: snapshot.stdout.start_offset,
end_offset: snapshot.stdout.end_offset,
content: snapshot.stdout.content,
truncated: snapshot.stdout.truncated,
},
stderr: ProtocolCommandStreamSlice {
start_offset: snapshot.stderr.start_offset,
end_offset: snapshot.stderr.end_offset,
content: snapshot.stderr.content,
truncated: snapshot.stderr.truncated,
},
exit_code: snapshot.exit_code,
}
}
fn protocol_command_event(event: WorkdirCommandEvent) -> ProtocolCommandEvent {
match event {
WorkdirCommandEvent::Started {
command_id,
tool_call_id,
observed_at_ms,
} => ProtocolCommandEvent::Started {
command_id,
tool_call_id,
observed_at_ms,
},
WorkdirCommandEvent::Output {
command_id,
stream,
start_offset,
end_offset,
content,
observed_at_ms,
} => ProtocolCommandEvent::Output {
command_id,
stream: match stream {
WorkdirCommandStream::Stdout => ProtocolCommandStream::Stdout,
WorkdirCommandStream::Stderr => ProtocolCommandStream::Stderr,
},
start_offset,
end_offset,
content,
observed_at_ms,
},
WorkdirCommandEvent::Terminal {
command_id,
status,
exit_code,
stdout_end_offset,
stderr_end_offset,
observed_at_ms,
} => ProtocolCommandEvent::Terminal {
command_id,
status: protocol_command_status(status),
exit_code,
stdout_end_offset,
stderr_end_offset,
observed_at_ms,
},
}
}
fn protocol_command_status(status: WorkdirCommandStatus) -> ProtocolCommandStatus {
match status {
WorkdirCommandStatus::Running => ProtocolCommandStatus::Running,
WorkdirCommandStatus::Completed => ProtocolCommandStatus::Completed,
WorkdirCommandStatus::Failed => ProtocolCommandStatus::Failed,
WorkdirCommandStatus::TimedOut => ProtocolCommandStatus::TimedOut,
WorkdirCommandStatus::Cancelled => ProtocolCommandStatus::Cancelled,
}
}
/// Wire the per-event broadcast bridges on the Worker's Engine. Each callback
/// re-publishes a worker-level signal as a `protocol::Event` on `event_tx`
/// so subscribers (TUI, socket clients) get a single typed stream.
crates/worker/src/in_flight.rs
+157 -2
View File
@@ -1,9 +1,14 @@
use std::sync::{Arc, Mutex, MutexGuard};
use protocol::{Event, InFlightBlock, InFlightSnapshot, InFlightToolCallState};
use protocol::{
CommandEvent, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, Event,
InFlightBlock, InFlightSnapshot, InFlightToolCallState,
};
use session_store::{LoggedContentPart, LoggedItem};
use tokio::sync::broadcast;
const COMMAND_SNAPSHOT_STREAM_BYTES: usize = 32 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InFlightBlockId(u64);
@@ -17,6 +22,7 @@ pub struct InFlightEvents {
pub(crate) struct InFlightInner {
next_block_id: u64,
blocks: Vec<TrackedBlock>,
commands: Vec<CommandSnapshot>,
}
#[derive(Debug, Clone)]
@@ -46,6 +52,7 @@ impl InFlightEvents {
inner: Arc::new(Mutex::new(InFlightInner {
next_block_id: 1,
blocks: Vec::new(),
commands: Vec::new(),
})),
event_tx,
}
@@ -201,6 +208,15 @@ impl InFlightEvents {
f()
}
pub(crate) fn publish_command_event(&self, event: CommandEvent) {
self.lock().apply_command_event(&event);
let _ = self.event_tx.send(Event::Command { event });
}
pub(crate) fn replace_command_snapshot(&self, commands: Vec<CommandSnapshot>) {
self.lock().commands = commands;
}
pub(crate) fn clear(&self) {
let mut inner = self.lock();
inner.clear();
@@ -224,6 +240,92 @@ impl InFlightInner {
.find(|block| block.block_id() == block_id)
}
fn apply_command_event(&mut self, event: &CommandEvent) {
match event {
CommandEvent::Started {
command_id,
tool_call_id,
observed_at_ms,
} => {
self.commands
.retain(|command| command.command_id != *command_id);
self.commands.push(CommandSnapshot {
command_id: command_id.clone(),
tool_call_id: tool_call_id.clone(),
status: CommandStatus::Running,
started_at_ms: *observed_at_ms,
observed_at_ms: *observed_at_ms,
last_output_at_ms: None,
stdout: CommandStreamSlice::default(),
stderr: CommandStreamSlice::default(),
exit_code: None,
});
}
CommandEvent::Output {
command_id,
stream,
start_offset,
end_offset,
content,
observed_at_ms,
} => {
let command = match self
.commands
.iter_mut()
.find(|command| command.command_id == *command_id)
{
Some(command) => command,
None => {
self.commands.push(CommandSnapshot {
command_id: command_id.clone(),
tool_call_id: None,
status: CommandStatus::Running,
started_at_ms: *observed_at_ms,
observed_at_ms: *observed_at_ms,
last_output_at_ms: Some(*observed_at_ms),
stdout: CommandStreamSlice::default(),
stderr: CommandStreamSlice::default(),
exit_code: None,
});
self.commands.last_mut().expect("command was inserted")
}
};
command.observed_at_ms = *observed_at_ms;
command.last_output_at_ms = Some(*observed_at_ms);
let target = match stream {
CommandStream::Stdout => &mut command.stdout,
CommandStream::Stderr => &mut command.stderr,
};
if target.end_offset != *start_offset {
target.content.clear();
target.start_offset = *start_offset;
target.truncated = *start_offset > 0;
}
target.content.push_str(content);
target.end_offset = *end_offset;
if target.content.len() > COMMAND_SNAPSHOT_STREAM_BYTES {
let mut cut = target.content.len() - COMMAND_SNAPSHOT_STREAM_BYTES;
while cut < target.content.len() && !target.content.is_char_boundary(cut) {
cut += 1;
}
target.content.drain(..cut);
target.start_offset = target
.end_offset
.saturating_sub(target.content.len() as u64);
target.truncated = true;
}
}
CommandEvent::Terminal { command_id, .. } => {
// Terminal state is delivered as a live protocol event. It is
// no longer in-flight snapshot state, and removing it here
// also prevents queued output from an aborted turn from
// surviving the subsequent terminal event after `clear()`.
self.commands
.retain(|command| command.command_id != *command_id);
}
}
}
fn clear_for_committed_item(&mut self, item: &LoggedItem) {
match item {
LoggedItem::Message { role, content }
@@ -273,14 +375,16 @@ impl InFlightInner {
.iter()
.filter_map(TrackedBlock::to_snapshot_block)
.collect(),
commands: self.commands.clone(),
}
}
fn clear(&mut self) -> bool {
if self.blocks.is_empty() {
if self.blocks.is_empty() && self.commands.is_empty() {
false
} else {
self.blocks.clear();
self.commands.clear();
true
}
}
@@ -583,6 +687,57 @@ mod tests {
);
}
#[test]
fn command_events_are_bounded_and_recoverable_from_snapshot() {
let (event_tx, _) = broadcast::channel(16);
let mut rx = event_tx.subscribe();
let in_flight = InFlightEvents::new(event_tx);
in_flight.publish_command_event(CommandEvent::Started {
command_id: "command-1".into(),
tool_call_id: Some("tool-1".into()),
observed_at_ms: 100,
});
in_flight.publish_command_event(CommandEvent::Output {
command_id: "command-1".into(),
stream: CommandStream::Stdout,
start_offset: 0,
end_offset: 5,
content: "ready".into(),
observed_at_ms: 110,
});
let guard = in_flight.snapshot_guard();
let snapshot = snapshot_from_guard(&guard);
assert_eq!(snapshot.commands.len(), 1);
assert_eq!(snapshot.commands[0].tool_call_id.as_deref(), Some("tool-1"));
assert_eq!(snapshot.commands[0].stdout.content, "ready");
assert_eq!(snapshot.commands[0].status, CommandStatus::Running);
drop(guard);
assert!(matches!(
rx.try_recv().unwrap(),
Event::Command {
event: CommandEvent::Started { .. }
}
));
assert!(matches!(
rx.try_recv().unwrap(),
Event::Command {
event: CommandEvent::Output { .. }
}
));
in_flight.publish_command_event(CommandEvent::Terminal {
command_id: "command-1".into(),
status: CommandStatus::TimedOut,
exit_code: None,
stdout_end_offset: 5,
stderr_end_offset: 0,
observed_at_ms: 200,
});
let guard = in_flight.snapshot_guard();
assert!(snapshot_from_guard(&guard).commands.is_empty());
}
#[test]
fn clear_discards_uncommitted_blocks_without_protocol_event() {
let (event_tx, _) = broadcast::channel(16);
crates/worker/src/internal_worker.rs
+4 -1
View File
@@ -17,7 +17,7 @@ use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntr
use tokio::sync::broadcast;
use uuid::Uuid;
use crate::controller::wire_event_bridges_on_engine;
use crate::controller::{wire_event_bridges_on_engine, wire_workdir_command_events};
use crate::feature::FeatureRegistryBuilder;
use crate::in_flight::{InFlightEvents, snapshot_from_guard};
use crate::ipc::alerter::Alerter;
@@ -576,6 +576,9 @@ pub(crate) async fn prepare_internal_worker_session(
spawn_internal_log_event_bridge(sink.clone(), event_tx.clone());
let alerter = Alerter::new(event_tx.clone());
let in_flight = InFlightEvents::new(event_tx.clone());
if let Some(session) = worker.workdir_session() {
wire_workdir_command_events(session, &in_flight);
}
let actor_in_flight = in_flight.clone();
worker.attach_alerter(alerter.clone());
worker.attach_event_tx(event_tx.clone());
crates/worker/tests/controller_test.rs
+108 -2
View File
@@ -12,8 +12,8 @@ use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use session_store::{CombinedStore, FsWorkerStore};
use session_store::{FsStore, LogEntry};
use workdir::{
CommandRequest, LocalWorkdirSession, Workdir, WorkdirError, WorkdirSessionCapabilities,
WorkdirSessionHandle,
CommandOutputRequest, CommandRequest, LocalWorkdirSession, Workdir, WorkdirError,
WorkdirSessionCapabilities, WorkdirSessionHandle,
};
use worker::{
@@ -232,6 +232,7 @@ async fn shutdown_closes_bound_workdir_session() {
command: "sleep 30".to_owned(),
timeout_secs: 60,
output_limit: 1024,
tool_call_id: None,
})
.await
.unwrap();
@@ -253,6 +254,110 @@ async fn shutdown_closes_bound_workdir_session() {
));
}
#[tokio::test]
async fn controller_projects_workdir_command_events_and_snapshot_state() {
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"),
pwd.clone(),
pwd,
worker.scope().clone(),
WorkdirSessionCapabilities::ALL,
));
worker.bind_workdir_session(Some(Arc::clone(&session)));
let handle = spawn_controller(worker).await;
let mut events = handle.subscribe();
let command = session
.start_command(CommandRequest {
command: "printf ready; sleep 0.3; printf done".to_owned(),
timeout_secs: 5,
output_limit: 1024,
tool_call_id: Some("tool-command-1".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())
.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);
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");
let output = session
.command_output(CommandOutputRequest {
handle: command,
cursor: 0,
limit: 1024,
wait: true,
})
.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).await.unwrap();
}
#[tokio::test]
async fn controller_startup_failure_closes_bound_workdir_session() {
let (mut worker, pwd) = make_worker_with_pwd(MockClient::new(simple_text_events())).await;
@@ -279,6 +384,7 @@ async fn controller_startup_failure_closes_bound_workdir_session() {
command: "printf unreachable".to_owned(),
timeout_secs: 5,
output_limit: 1024,
tool_call_id: None,
})
.await,
Err(WorkdirError::Unavailable(_))