refactor(cli): stream run attach over SSE

Replace the attach polling loop with the existing run attach SSE endpoint.
Seed from stored history once, fetch interview questions only when needed,
and keep completed-run replay behavior intact.
This commit is contained in:
Bryan Helmkamp 2026-04-07 12:58:37 -04:00
parent 80dfbd5309
commit 6a804e9808
No known key found for this signature in database
6 changed files with 451 additions and 234 deletions

View file

@ -4,8 +4,6 @@ use std::path::Path;
#[cfg(test)]
use std::path::PathBuf;
use std::process::ExitCode;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
use anyhow::Result;
@ -19,7 +17,7 @@ use fabro_util::terminal::Styles;
use fabro_workflow::outcome::StageStatus;
use fabro_workflow::run_status::RunStatus;
use tokio::signal::ctrl_c;
use tokio::time::sleep;
use tokio::time::{sleep, timeout};
use super::run_progress;
use crate::server_client;
@ -72,157 +70,52 @@ pub(crate) async fn attach_run_with_client(
.as_ref()
.is_some_and(|record| record.settings.verbose_enabled());
let events = client.list_run_events(run_id, None, None).await?;
let event_lines = events
.iter()
.map(event_payload_line)
.collect::<Result<Vec<_>>>()?;
let replay_events = events.clone();
let next_seq = events.last().map_or(1, |event| event.seq.saturating_add(1));
let initial_exit_code = events.iter().rev().find_map(event_exit_code);
attach_run_server(
client,
run_id,
verbose,
event_lines,
events.last().map_or(0, |event| event.seq),
initial_exit_code,
kill_on_detach,
styles,
json_output,
)
.await
if state_is_terminal(&state) || initial_exit_code.is_some() {
return replay_run_with_client(client, run_id, verbose, events, json_output).await;
}
match client.attach_run_events(run_id, Some(next_seq)).await {
Ok(stream) => {
attach_live_run_with_client(
client,
run_id,
verbose,
events,
stream,
kill_on_detach,
styles,
json_output,
)
.await
}
Err(server_client::RunAttachStreamError::Gone) => {
replay_run_with_client(client, run_id, verbose, replay_events, json_output).await
}
Err(server_client::RunAttachStreamError::Other(err)) => Err(err),
}
}
async fn attach_run_server(
async fn replay_run_with_client(
client: &server_client::ServerStoreClient,
run_id: &RunId,
verbose: bool,
existing_events: Vec<String>,
last_seq: u32,
initial_exit_code: Option<ExitCode>,
kill_on_detach: bool,
styles: &'static Styles,
events: Vec<EventEnvelope>,
json_output: bool,
) -> Result<ExitCode> {
let is_tty = std::io::stderr().is_terminal();
let mut progress_ui = run_progress::ProgressUI::new(is_tty, verbose);
let mut terminal_exit_code = None;
// Install Ctrl+C handler
let cancelled = Arc::new(AtomicBool::new(false));
{
let cancelled = Arc::clone(&cancelled);
tokio::spawn(async move {
let _ = ctrl_c().await;
cancelled.store(true, Ordering::Relaxed);
});
}
for line in &existing_events {
emit_progress_line(&mut progress_ui, line, json_output)?;
}
if json_output && !client.list_run_questions(run_id).await?.is_empty() {
eprintln!("{JSON_INTERVIEW_MESSAGE}");
return Ok(ExitCode::from(1));
}
let mut next_seq = if last_seq == 0 { 1 } else { last_seq + 1 };
let mut terminal_exit_code = initial_exit_code;
let mut terminal_event_seen_at = initial_exit_code.map(|_| Instant::now());
loop {
if cancelled.load(Ordering::Relaxed) {
if kill_on_detach {
let _ = client.cancel_run(run_id).await;
// Wait briefly for a terminal status or conclusion
for _ in 0..20 {
if client
.get_run_state(run_id)
.await
.ok()
.is_some_and(|state| {
state.conclusion.is_some()
|| state
.status
.is_some_and(|record| record.status.is_terminal())
})
{
break;
}
sleep(Duration::from_millis(100)).await;
}
} else {
eprintln!("Detached from run (engine continues in background)");
}
break;
}
let mut saw_event = false;
let events = match client.list_run_events(run_id, Some(next_seq), None).await {
Ok(events) => events,
Err(err) if terminal_event_seen_at.is_some() && is_run_not_found_error(&err) => break,
Err(err) => return Err(err),
};
for event in events {
if let Some(exit_code) = event_exit_code(&event) {
terminal_exit_code = Some(exit_code);
terminal_event_seen_at = Some(Instant::now());
}
let line = event_payload_line(&event)?;
emit_progress_line(&mut progress_ui, &line, json_output)?;
next_seq = event.seq.saturating_add(1);
saw_event = true;
}
if let Some(seen_at) = terminal_event_seen_at {
if !saw_event && seen_at.elapsed() >= ATTACH_FINAL_STATUS_GRACE {
break;
}
if !saw_event {
sleep(Duration::from_millis(50)).await;
}
continue;
}
// Check for server-backed interview request
if let Some(question) = client.list_run_questions(run_id).await?.into_iter().next() {
if json_output {
eprintln!("{JSON_INTERVIEW_MESSAGE}");
return Ok(ExitCode::from(1));
}
hide_progress(&mut progress_ui, json_output);
let interviewer = ConsoleInterviewer::new(styles);
let answer = fabro_interview::Interviewer::ask(
&interviewer,
api_question_to_question(&question),
)
.await;
show_progress(&mut progress_ui, json_output);
if answer_requires_reattach(&answer) {
eprintln!("{INTERVIEW_UNANSWERED_MESSAGE}");
return Ok(ExitCode::from(1));
}
submit_server_interview_answer(client, run_id, &question.id, &answer).await?;
continue;
}
let terminal_status = client
.get_run_state(run_id)
.await
.ok()
.and_then(|state| state.status.map(|record| record.status))
.filter(|status| status.is_terminal());
if terminal_status.is_some() && !saw_event {
flush_remaining_server_events(client, run_id, next_seq, &mut progress_ui, json_output)
.await?;
break;
}
if !saw_event {
sleep(Duration::from_millis(100)).await;
for event in events {
if let Some(exit_code) = event_exit_code(&event) {
terminal_exit_code = Some(exit_code);
}
let line = event_payload_line(&event)?;
emit_progress_line(&mut progress_ui, &line, json_output)?;
}
finish_progress(&mut progress_ui, json_output);
@ -233,6 +126,194 @@ async fn attach_run_server(
})
}
async fn attach_live_run_with_client(
client: &server_client::ServerStoreClient,
run_id: &RunId,
verbose: bool,
existing_events: Vec<EventEnvelope>,
mut stream: server_client::RunAttachEventStream,
kill_on_detach: bool,
styles: &'static Styles,
json_output: bool,
) -> Result<ExitCode> {
let is_tty = std::io::stderr().is_terminal();
let mut progress_ui = run_progress::ProgressUI::new(is_tty, verbose);
let ctrl_c_signal = ctrl_c();
tokio::pin!(ctrl_c_signal);
let mut next_seq = 1;
let mut terminal_exit_code = None;
let mut terminal_event_seen_at: Option<Instant> = None;
for event in existing_events {
next_seq = event.seq.saturating_add(1);
if let Some(exit_code) = event_exit_code(&event) {
terminal_exit_code = Some(exit_code);
terminal_event_seen_at = Some(Instant::now());
}
let line = event_payload_line(&event)?;
emit_progress_line(&mut progress_ui, &line, json_output)?;
}
if let Some(exit_code) =
handle_pending_server_interview(client, run_id, &mut progress_ui, styles, json_output)
.await?
{
return Ok(exit_code);
}
loop {
let next_event = if let Some(seen_at) = terminal_event_seen_at {
let remaining = ATTACH_FINAL_STATUS_GRACE.saturating_sub(seen_at.elapsed());
if remaining.is_zero() {
break;
}
tokio::select! {
_ = &mut ctrl_c_signal => {
handle_detach_signal(client, run_id, kill_on_detach).await;
break;
}
result = timeout(remaining, stream.next_event()) => {
match result {
Ok(result) => result?,
Err(_) => break,
}
}
}
} else {
tokio::select! {
_ = &mut ctrl_c_signal => {
handle_detach_signal(client, run_id, kill_on_detach).await;
break;
}
result = stream.next_event() => result?,
}
};
let Some(event) = next_event else {
break;
};
next_seq = event.seq.saturating_add(1);
if let Some(exit_code) = event_exit_code(&event) {
terminal_exit_code = Some(exit_code);
terminal_event_seen_at = Some(Instant::now());
}
let line = event_payload_line(&event)?;
emit_progress_line(&mut progress_ui, &line, json_output)?;
if event_starts_interview(&event) {
if let Some(exit_code) = handle_pending_server_interview(
client,
run_id,
&mut progress_ui,
styles,
json_output,
)
.await?
{
return Ok(exit_code);
}
}
}
if terminal_exit_code.is_none() {
let (_, trailing_exit_code) =
emit_server_events_from(client, run_id, next_seq, &mut progress_ui, json_output)
.await?;
terminal_exit_code = trailing_exit_code;
}
finish_progress(&mut progress_ui, json_output);
Ok(match terminal_exit_code {
Some(exit_code) => exit_code,
None => determine_exit_code_with_server(client, run_id).await,
})
}
async fn handle_pending_server_interview(
client: &server_client::ServerStoreClient,
run_id: &RunId,
progress_ui: &mut run_progress::ProgressUI,
styles: &'static Styles,
json_output: bool,
) -> Result<Option<ExitCode>> {
let Some(question) = client.list_run_questions(run_id).await?.into_iter().next() else {
return Ok(None);
};
if json_output {
eprintln!("{JSON_INTERVIEW_MESSAGE}");
return Ok(Some(ExitCode::from(1)));
}
hide_progress(progress_ui, json_output);
let interviewer = ConsoleInterviewer::new(styles);
let answer =
fabro_interview::Interviewer::ask(&interviewer, api_question_to_question(&question)).await;
show_progress(progress_ui, json_output);
if answer_requires_reattach(&answer) {
eprintln!("{INTERVIEW_UNANSWERED_MESSAGE}");
return Ok(Some(ExitCode::from(1)));
}
submit_server_interview_answer(client, run_id, &question.id, &answer).await?;
Ok(None)
}
async fn handle_detach_signal(
client: &server_client::ServerStoreClient,
run_id: &RunId,
kill_on_detach: bool,
) {
if kill_on_detach {
let _ = client.cancel_run(run_id).await;
for _ in 0..20 {
if client
.get_run_state(run_id)
.await
.ok()
.is_some_and(|state| state_is_terminal(&state))
{
break;
}
sleep(Duration::from_millis(100)).await;
}
} else {
eprintln!("Detached from run (engine continues in background)");
}
}
async fn emit_server_events_from(
client: &server_client::ServerStoreClient,
run_id: &RunId,
next_seq: u32,
progress_ui: &mut run_progress::ProgressUI,
json_output: bool,
) -> Result<(u32, Option<ExitCode>)> {
let events = match client.list_run_events(run_id, Some(next_seq), None).await {
Ok(events) => events,
Err(err) if is_run_not_found_error(&err) => Vec::new(),
Err(err) => return Err(err),
};
let mut current_seq = next_seq;
let mut terminal_exit_code = None;
for event in events {
if let Some(exit_code) = event_exit_code(&event) {
terminal_exit_code = Some(exit_code);
}
let line = event_payload_line(&event)?;
emit_progress_line(progress_ui, &line, json_output)?;
current_seq = event.seq.saturating_add(1);
}
Ok((current_seq, terminal_exit_code))
}
fn api_question_to_question(question: &types::ApiQuestion) -> Question {
let question_type = match question.question_type {
types::QuestionType::YesNo => QuestionType::YesNo,
@ -282,45 +363,19 @@ async fn submit_server_interview_answer(
Ok(true)
}
async fn flush_remaining_server_events(
client: &server_client::ServerStoreClient,
run_id: &RunId,
mut next_seq: u32,
progress_ui: &mut run_progress::ProgressUI,
json_output: bool,
) -> Result<()> {
let deadline = Instant::now() + ATTACH_FINAL_STATUS_GRACE;
loop {
let mut saw_new_event = false;
let events = match client.list_run_events(run_id, Some(next_seq), None).await {
Ok(events) => events,
Err(err) if is_run_not_found_error(&err) => break,
Err(err) => return Err(err),
};
for event in events {
let line = event_payload_line(&event)?;
emit_progress_line(progress_ui, &line, json_output)?;
next_seq = event.seq.saturating_add(1);
saw_new_event = true;
}
if Instant::now() >= deadline {
break;
}
if !saw_new_event {
sleep(Duration::from_millis(100)).await;
}
}
Ok(())
}
fn is_run_not_found_error(err: &anyhow::Error) -> bool {
err.chain()
.any(|cause| cause.to_string() == "Run not found.")
}
fn state_is_terminal(state: &server_client::RunProjection) -> bool {
state.conclusion.is_some()
|| state
.status
.as_ref()
.is_some_and(|record| record.status.is_terminal())
}
fn emit_progress_line(
progress_ui: &mut run_progress::ProgressUI,
line: &str,
@ -450,6 +505,13 @@ fn event_exit_code(event: &EventEnvelope) -> Option<ExitCode> {
}
}
fn event_starts_interview(event: &EventEnvelope) -> bool {
let Ok(run_event) = RunEvent::try_from(&event.payload) else {
return false;
};
matches!(run_event.body, EventBody::InterviewStarted(_))
}
#[cfg(test)]
mod tests {
#![allow(clippy::absolute_paths)]

View file

@ -3,6 +3,7 @@ use futures::StreamExt;
use crate::args::{GlobalArgs, SystemEventsArgs};
use crate::server_client;
use crate::sse;
pub(super) async fn events_command(args: &SystemEventsArgs, globals: &GlobalArgs) -> Result<()> {
let client = server_client::connect_server_backed_api_client_with_storage_dir(
@ -23,25 +24,15 @@ pub(super) async fn events_command(args: &SystemEventsArgs, globals: &GlobalArgs
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|err| anyhow::anyhow!("{err}"))?;
pending.extend_from_slice(&chunk);
drain_sse_lines(&mut pending, globals.json)?;
}
if !pending.is_empty() {
drain_sse_lines(&mut pending, globals.json)?;
}
Ok(())
}
fn drain_sse_lines(buffer: &mut Vec<u8>, json_output: bool) -> Result<()> {
while let Some(pos) = buffer.iter().position(|byte| *byte == b'\n') {
let line = buffer.drain(..=pos).collect::<Vec<_>>();
let line = String::from_utf8_lossy(&line);
let line = line.trim_end_matches(['\r', '\n']);
if let Some(data) = line.strip_prefix("data:") {
render_sse_payload(data.trim(), json_output)?;
for payload in sse::drain_sse_payloads(&mut pending, false) {
render_sse_payload(&payload, globals.json)?;
}
}
for payload in sse::drain_sse_payloads(&mut pending, true) {
render_sse_payload(&payload, globals.json)?;
}
Ok(())
}

View file

@ -9,6 +9,7 @@ mod server_runs;
mod shared;
#[cfg(feature = "sleep_inhibitor")]
mod sleep_inhibitor;
mod sse;
mod user_config;
use anyhow::Result;

View file

@ -1,3 +1,4 @@
use std::collections::VecDeque;
use std::num::NonZeroU64;
use std::path::{Path, PathBuf};
use std::time::Duration;
@ -13,6 +14,7 @@ use tokio::time::sleep;
use crate::args::ServerTargetArgs;
use crate::commands::server::start;
use crate::sse;
use crate::user_config;
#[derive(Clone)]
@ -26,6 +28,55 @@ struct LocalServerRuntime {
storage_dir: PathBuf,
}
pub(crate) struct RunAttachEventStream {
stream: progenitor_client::ByteStream,
pending_bytes: Vec<u8>,
buffered_events: VecDeque<EventEnvelope>,
}
pub(crate) enum RunAttachStreamError {
Gone,
Other(anyhow::Error),
}
impl RunAttachEventStream {
fn new(stream: progenitor_client::ByteStream) -> Self {
Self {
stream,
pending_bytes: Vec::new(),
buffered_events: VecDeque::new(),
}
}
pub(crate) async fn next_event(&mut self) -> Result<Option<EventEnvelope>> {
loop {
if let Some(event) = self.buffered_events.pop_front() {
return Ok(Some(event));
}
match self.stream.next().await {
Some(chunk) => {
let chunk = chunk.map_err(|err| anyhow!("{err}"))?;
self.pending_bytes.extend_from_slice(&chunk);
self.buffer_sse_events(false)?;
}
None => {
self.buffer_sse_events(true)?;
return Ok(self.buffered_events.pop_front());
}
}
}
}
fn buffer_sse_events(&mut self, finalize: bool) -> Result<()> {
for payload in sse::drain_sse_payloads(&mut self.pending_bytes, finalize) {
let event: types::EventEnvelope = serde_json::from_str(&payload)?;
self.buffered_events.push_back(convert_type(event)?);
}
Ok(())
}
}
pub(crate) use fabro_store::RunProjection;
pub(crate) async fn connect_server(storage_dir: &Path) -> Result<ServerStoreClient> {
@ -300,6 +351,19 @@ impl ServerStoreClient {
.collect::<Result<Vec<_>>>()
}
pub(crate) async fn attach_run_events(
&self,
run_id: &RunId,
since_seq: Option<u32>,
) -> std::result::Result<RunAttachEventStream, RunAttachStreamError> {
let mut request = self.client.attach_run_events().id(run_id.to_string());
if let Some(seq) = since_seq.and_then(non_zero_u64_from_u32) {
request = request.since_seq(seq);
}
let response = request.send().await.map_err(map_attach_run_stream_error)?;
Ok(RunAttachEventStream::new(response.into_inner()))
}
pub(crate) async fn list_run_questions(
&self,
run_id: &RunId,
@ -520,6 +584,19 @@ where
}
}
fn map_attach_run_stream_error(
err: progenitor_client::Error<types::ErrorResponse>,
) -> RunAttachStreamError {
match &err {
progenitor_client::Error::ErrorResponse(response)
if response.status() == reqwest::StatusCode::GONE =>
{
RunAttachStreamError::Gone
}
_ => RunAttachStreamError::Other(map_api_error(err)),
}
}
fn convert_type<TInput, TOutput>(value: TInput) -> Result<TOutput>
where
TInput: serde::Serialize,

View file

@ -0,0 +1,51 @@
pub(crate) fn drain_sse_payloads(buffer: &mut Vec<u8>, finalize: bool) -> Vec<String> {
let mut payloads = Vec::new();
while let Some(pos) = buffer.iter().position(|byte| *byte == b'\n') {
let line = buffer.drain(..=pos).collect::<Vec<_>>();
if let Some(payload) = sse_data_line(&line) {
payloads.push(payload);
}
}
if finalize && !buffer.is_empty() {
let line = std::mem::take(buffer);
if let Some(payload) = sse_data_line(&line) {
payloads.push(payload);
}
}
payloads
}
fn sse_data_line(line: &[u8]) -> Option<String> {
let line = String::from_utf8_lossy(line);
let line = line.trim_end_matches(['\r', '\n']);
line.strip_prefix("data:")
.map(|data| data.trim().to_string())
}
#[cfg(test)]
mod tests {
use super::drain_sse_payloads;
#[test]
fn drain_sse_payloads_handles_chunk_boundaries() {
let mut buffer = b"data: {\"a\":1}\n\ndat".to_vec();
assert_eq!(drain_sse_payloads(&mut buffer, false), vec![r#"{"a":1}"#]);
assert_eq!(buffer, b"dat");
buffer.extend_from_slice(b"a: {\"b\":2}\n\n");
assert_eq!(drain_sse_payloads(&mut buffer, false), vec![r#"{"b":2}"#]);
assert!(buffer.is_empty());
}
#[test]
fn drain_sse_payloads_ignores_keepalives_and_finalizes_trailing_line() {
let mut buffer = b": keep-alive\n\n\ndata: {\"a\":1}".to_vec();
assert_eq!(drain_sse_payloads(&mut buffer, true), vec![r#"{"a":1}"#]);
assert!(buffer.is_empty());
}
}

View file

@ -10,6 +10,48 @@ use super::support::{output_stdout, resolve_run, wait_for_status, write_gated_wo
const SHARED_DAEMON_TIMEOUT: Duration = Duration::from_secs(30);
fn live_run_state_response() -> serde_json::Value {
serde_json::json!({
"run": null,
"graph_source": null,
"start": null,
"status": {
"status": "running",
"reason": null,
"updated_at": "2026-04-05T12:00:01Z"
},
"checkpoint": null,
"checkpoints": [],
"conclusion": null,
"retro": null,
"retro_prompt": null,
"retro_response": null,
"sandbox": null,
"final_patch": null,
"pull_request": null,
"nodes": {}
})
}
fn run_sse_body(run_id: &str) -> String {
let completed = serde_json::json!({
"seq": 2,
"payload": {
"event": "run.completed",
"id": "evt-run-completed",
"run_id": run_id,
"ts": "2026-04-05T12:00:01Z",
"properties": {
"duration_ms": 12,
"artifact_count": 0,
"status": "success"
}
}
});
format!("data: {completed}\n\n")
}
#[test]
fn help() {
let context = test_context!();
@ -76,7 +118,7 @@ fn attach_uses_configured_server_target_without_server_flag() {
"labels": {},
"host_repo_path": null,
"start_time": "2026-04-05T12:00:00Z",
"status": "succeeded",
"status": "running",
"status_reason": null,
"duration_ms": 12,
"total_cost": null
@ -88,59 +130,49 @@ fn attach_uses_configured_server_target_without_server_flag() {
server.mock(|when, then| {
when.method("GET")
.path(format!("/api/v1/runs/{run_id}/events"));
then.status(200)
.header("Content-Type", "application/json")
.body(r#"{"data":[],"meta":{"has_more":false}}"#);
});
server.mock(|when, then| {
when.method("GET")
.path(format!("/api/v1/runs/{run_id}/state"));
then.status(200)
.header("Content-Type", "application/json")
.body(
serde_json::json!({
"run": null,
"graph_source": null,
"start": null,
"status": {
"status": "succeeded",
"reason": null,
"updated_at": "2026-04-05T12:00:01Z"
},
"checkpoint": null,
"checkpoints": [],
"conclusion": {
"timestamp": "2026-04-05T12:00:01Z",
"status": "success",
"duration_ms": 12,
"stages": [],
"total_cost": null,
"total_retries": 0,
"total_input_tokens": 0,
"total_output_tokens": 0,
"total_cache_read_tokens": 0,
"total_cache_write_tokens": 0,
"total_reasoning_tokens": 0,
"has_pricing": false
},
"retro": null,
"retro_prompt": null,
"retro_response": null,
"sandbox": null,
"final_patch": null,
"pull_request": null,
"nodes": {}
"data": [{
"seq": 1,
"payload": {
"event": "run.running",
"id": "evt-run-running",
"run_id": run_id,
"ts": "2026-04-05T12:00:00Z",
"properties": {}
}
}],
"meta": { "has_more": false }
})
.to_string(),
);
});
server.mock(|when, then| {
when.method("GET")
.path(format!("/api/v1/runs/{run_id}/questions"));
.path(format!("/api/v1/runs/{run_id}/state"));
then.status(200)
.header("Content-Type", "application/json")
.body(live_run_state_response().to_string());
});
server.mock(|when, then| {
when.method("GET")
.path(format!("/api/v1/runs/{run_id}/questions"))
.query_param("page[limit]", "100")
.query_param("page[offset]", "0");
then.status(200)
.header("Content-Type", "application/json")
.body(r#"{"data":[],"meta":{"has_more":false}}"#);
});
let attach_mock = server.mock(|when, then| {
when.method("GET")
.path(format!("/api/v1/runs/{run_id}/attach"))
.query_param("since_seq", "2");
then.status(200)
.header("Content-Type", "text/event-stream")
.body(run_sse_body(run_id.as_str()));
});
context.write_home(
".fabro/settings.toml",
format!("[server]\ntarget = \"{}/api/v1\"\n", server.base_url()),
@ -148,7 +180,7 @@ fn attach_uses_configured_server_target_without_server_flag() {
let output = context
.command()
.args(["attach", &run_id])
.args(["--json", "attach", &run_id])
.output()
.expect("attach should execute");
@ -159,6 +191,9 @@ fn attach_uses_configured_server_target_without_server_flag() {
String::from_utf8_lossy(&output.stderr)
);
list_mock.assert();
attach_mock.assert();
let stdout = String::from_utf8(output.stdout).expect("stdout should be UTF-8");
assert!(stdout.contains("\"event\":\"run.completed\""), "{stdout}");
}
#[test]