diff --git a/lib/crates/fabro-agent/src/session.rs b/lib/crates/fabro-agent/src/session.rs index 2d7dd1212..41de33cc0 100644 --- a/lib/crates/fabro-agent/src/session.rs +++ b/lib/crates/fabro-agent/src/session.rs @@ -15,6 +15,7 @@ use crate::types::{AgentEvent, SessionState, Turn}; use fabro_llm::client::Client; use fabro_llm::error::{ProviderErrorKind, SdkError}; use fabro_llm::generate::StreamAccumulator; +use fabro_llm::provider::StreamEventStream; use fabro_llm::types::{Message, Request, StreamEvent, ToolChoice}; use fabro_mcp::config::McpServerConfig; use futures::StreamExt; @@ -389,6 +390,38 @@ impl Session { AgentError::Aborted(reason) } + fn emit_llm_error(&mut self, err: SdkError) -> AgentError { + self.event_emitter.emit( + self.id.clone(), + AgentEvent::Error { + error: AgentError::Llm(err.clone()), + }, + ); + if is_auth_error(&err) { + self.state = SessionState::Closed; + } + AgentError::Llm(err) + } + + async fn open_stream_with_retry( + &mut self, + client: &Client, + request: &Request, + retry_policy: &fabro_llm::types::RetryPolicy, + ) -> Result { + let stream_result = fabro_llm::retry::retry(retry_policy, || { + let client = client.clone(); + let request = request.clone(); + async move { client.stream(&request).await } + }) + .await; + + match stream_result { + Ok(stream) => Ok(stream), + Err(err) => Err(self.emit_llm_error(err)), + } + } + #[must_use] pub fn followup_queue_handle(&self) -> Arc>> { self.followup_queue.clone() @@ -596,42 +629,27 @@ impl Session { ..Default::default() }; let client = self.llm_client.clone(); - let stream_result = fabro_llm::retry::retry(&retry_policy, || { - let c = client.clone(); - let r = request.clone(); - async move { c.stream(&r).await } - }) - .await; - let event_stream = match stream_result { - Ok(stream) => stream, - Err(err) => { - self.event_emitter.emit( - self.id.clone(), - AgentEvent::Error { - error: AgentError::Llm(err.clone()), - }, - ); - if is_auth_error(&err) { - self.state = SessionState::Closed; - } - return Err(AgentError::Llm(err)); - } - }; + let mut event_stream = self + .open_stream_with_retry(&client, &request, &retry_policy) + .await?; - // Consume the stream, retrying up to 3 times on transient - // "stream ended without Finish" errors (e.g. dropped connections). + // Consume the stream, retrying up to 3 times if the provider + // closes the stream without sending a Finish event. If visible + // output was already emitted, clear it before replaying the turn. const STREAM_CONSUME_RETRIES: usize = 3; let mut response = None; - let mut last_stream = event_stream; for stream_attempt in 0..=STREAM_CONSUME_RETRIES { let mut accumulator = StreamAccumulator::new(); + let mut emitted_text = String::new(); + let mut emitted_reasoning = String::new(); - while let Some(event_result) = last_stream.next().await { + while let Some(event_result) = event_stream.next().await { match event_result { Ok(event) => { match &event { StreamEvent::TextDelta { ref delta, .. } => { + emitted_text.push_str(delta); self.event_emitter.emit( self.id.clone(), AgentEvent::TextDelta { @@ -640,6 +658,7 @@ impl Session { ); } StreamEvent::ReasoningDelta { ref delta } => { + emitted_reasoning.push_str(delta); self.event_emitter.emit( self.id.clone(), AgentEvent::ReasoningDelta { @@ -652,13 +671,7 @@ impl Session { accumulator.process(&event); } Err(err) => { - self.event_emitter.emit( - self.id.clone(), - AgentEvent::Error { - error: AgentError::Llm(err.clone()), - }, - ); - return Err(AgentError::Llm(err)); + return Err(self.emit_llm_error(err)); } } @@ -671,7 +684,7 @@ impl Session { // If aborted during streaming, drop the stream to cancel the HTTP // connection, then close the session before returning. if self.cancel_token.is_cancelled() { - drop(last_stream); + drop(event_stream); self.close(); return Err(self.aborted_error()); } @@ -688,23 +701,29 @@ impl Session { max = STREAM_CONSUME_RETRIES, "Stream ended without Finish event, retrying turn" ); - // Re-open the stream with the same request - match client.stream(&request).await { - Ok(new_stream) => { - last_stream = new_stream; - } - Err(err) => { - return Err(AgentError::Llm(err)); - } + if !emitted_text.is_empty() || !emitted_reasoning.is_empty() { + self.event_emitter.emit( + self.id.clone(), + AgentEvent::AssistantOutputReplace { + text: String::new(), + reasoning: None, + }, + ); } + event_stream = self + .open_stream_with_retry(&client, &request, &retry_policy) + .await?; } } - let response = response.ok_or_else(|| { - AgentError::Llm(SdkError::Stream { - message: "Stream ended without a Finish event (after retries)".into(), - }) - })?; + let response = match response { + Some(response) => response, + None => { + return Err(self.emit_llm_error(SdkError::Stream { + message: "Stream ended without a Finish event (after retries)".into(), + })) + } + }; // Record assistant turn let text = response.text(); @@ -905,10 +924,96 @@ mod tests { use super::*; use crate::test_support::*; use crate::tool_registry::{RegisteredTool, ToolRegistry}; - use fabro_llm::error::ProviderErrorDetail; + use fabro_llm::error::{ProviderErrorDetail, ProviderErrorKind}; use fabro_llm::provider::{ProviderAdapter, StreamEventStream}; - use fabro_llm::types::{Response, Role, ToolDefinition}; + use fabro_llm::types::{Request, Response, Role, StreamEvent, ToolDefinition}; use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + #[derive(Clone)] + enum ScriptedStreamCall { + Response(Response), + Events(Vec>), + Error(SdkError), + } + + struct ScriptedStreamProvider { + calls: Vec, + call_index: AtomicUsize, + } + + impl ScriptedStreamProvider { + fn new(calls: Vec) -> Self { + assert!( + !calls.is_empty(), + "scripted stream provider needs at least one call" + ); + Self { + calls, + call_index: AtomicUsize::new(0), + } + } + + fn events_for_response(response: Response) -> Vec> { + let mut events = Vec::new(); + let text = response.text(); + if !text.is_empty() { + events.push(Ok(StreamEvent::text_delta(text, None))); + } + + for part in &response.message.content { + if let fabro_llm::types::ContentPart::ToolCall(tool_call) = part { + events.push(Ok(StreamEvent::ToolCallEnd { + tool_call: tool_call.clone(), + })); + } + } + + events.push(Ok(StreamEvent::finish( + response.finish_reason.clone(), + response.usage.clone(), + response, + ))); + events + } + } + + #[async_trait::async_trait] + impl ProviderAdapter for ScriptedStreamProvider { + fn name(&self) -> &'static str { + "mock" + } + + async fn complete(&self, _request: &Request) -> Result { + Err(SdkError::Configuration { + message: "ScriptedStreamProvider does not implement complete()".into(), + }) + } + + async fn stream(&self, _request: &Request) -> Result { + let idx = self.call_index.fetch_add(1, Ordering::SeqCst); + let scripted = if idx < self.calls.len() { + self.calls[idx].clone() + } else { + self.calls[self.calls.len() - 1].clone() + }; + + match scripted { + ScriptedStreamCall::Response(response) => Ok(Box::pin(futures::stream::iter( + Self::events_for_response(response), + ))), + ScriptedStreamCall::Events(events) => Ok(Box::pin(futures::stream::iter(events))), + ScriptedStreamCall::Error(err) => Err(err), + } + } + } + + async fn make_session_with_provider(provider: Arc) -> Session { + let client = make_client(provider).await; + let profile = Arc::new(TestProfile::new()); + let env = Arc::new(MockSandbox::default()); + Session::new(client, profile, env, SessionConfig::default()) + } // --- Tests --- @@ -1878,6 +1983,153 @@ mod tests { )); } + #[tokio::test] + async fn stream_retries_when_stream_ends_without_finish_before_any_deltas() { + let provider = Arc::new(ScriptedStreamProvider::new(vec![ + ScriptedStreamCall::Events(vec![]), + ScriptedStreamCall::Response(text_response("Recovered")), + ])); + let mut session = make_session_with_provider(provider.clone()).await; + let mut rx = session.subscribe(); + + session.process_input("Hello").await.unwrap(); + + assert_eq!(provider.call_index.load(Ordering::SeqCst), 2); + let turns = session.history().turns(); + assert!(matches!( + turns.last(), + Some(Turn::Assistant { content, .. }) if content == "Recovered" + )); + + let mut assistant_text_start_count = 0; + let mut replace_count = 0; + let mut deltas = Vec::new(); + let mut assistant_messages = Vec::new(); + while let Ok(event) = rx.try_recv() { + match event.event { + AgentEvent::AssistantTextStart => assistant_text_start_count += 1, + AgentEvent::AssistantOutputReplace { .. } => replace_count += 1, + AgentEvent::TextDelta { delta } => deltas.push(delta), + AgentEvent::AssistantMessage { text, .. } => assistant_messages.push(text), + _ => {} + } + } + + assert_eq!(assistant_text_start_count, 1); + assert_eq!(replace_count, 0); + assert_eq!(deltas, vec!["Recovered".to_string()]); + assert_eq!(assistant_messages, vec!["Recovered".to_string()]); + } + + #[tokio::test] + async fn stream_retries_with_output_replace_after_partial_text() { + let provider = Arc::new(ScriptedStreamProvider::new(vec![ + ScriptedStreamCall::Events(vec![Ok(StreamEvent::text_delta("Hel", None))]), + ScriptedStreamCall::Response(text_response("Hello")), + ])); + let mut session = make_session_with_provider(provider.clone()).await; + let mut rx = session.subscribe(); + + session.process_input("Hello").await.unwrap(); + + assert_eq!(provider.call_index.load(Ordering::SeqCst), 2); + let turns = session.history().turns(); + assert!(matches!( + turns.last(), + Some(Turn::Assistant { content, .. }) if content == "Hello" + )); + + let mut observed = Vec::new(); + while let Ok(event) = rx.try_recv() { + match event.event { + AgentEvent::AssistantTextStart => observed.push("start".to_string()), + AgentEvent::TextDelta { delta } => observed.push(format!("delta:{delta}")), + AgentEvent::AssistantOutputReplace { text, reasoning } => { + observed.push(format!("replace:{text}:{reasoning:?}")); + } + AgentEvent::AssistantMessage { text, .. } => { + observed.push(format!("message:{text}")); + } + _ => {} + } + } + + assert_eq!( + observed, + vec![ + "start".to_string(), + "delta:Hel".to_string(), + "replace::None".to_string(), + "delta:Hello".to_string(), + "message:Hello".to_string(), + ] + ); + } + + #[tokio::test] + async fn retry_open_auth_error_emits_error_and_closes_session() { + let auth_error = SdkError::Provider { + kind: ProviderErrorKind::Authentication, + detail: Box::new(ProviderErrorDetail { + status_code: Some(401), + ..ProviderErrorDetail::new("bad key", "mock") + }), + }; + let provider = Arc::new(ScriptedStreamProvider::new(vec![ + ScriptedStreamCall::Events(vec![Ok(StreamEvent::text_delta("Hel", None))]), + ScriptedStreamCall::Error(auth_error.clone()), + ])); + let mut session = make_session_with_provider(provider.clone()).await; + let mut rx = session.subscribe(); + + let result = session.process_input("Hello").await; + assert!(matches!( + result, + Err(AgentError::Llm(SdkError::Provider { + kind: ProviderErrorKind::Authentication, + .. + })) + )); + + assert_eq!(provider.call_index.load(Ordering::SeqCst), 2); + assert_eq!(session.state(), SessionState::Closed); + + let mut observed = Vec::new(); + let mut found_auth_error_event = false; + while let Ok(event) = rx.try_recv() { + match event.event { + AgentEvent::AssistantTextStart => observed.push("start".to_string()), + AgentEvent::TextDelta { delta } => observed.push(format!("delta:{delta}")), + AgentEvent::AssistantOutputReplace { text, reasoning } => { + observed.push(format!("replace:{text}:{reasoning:?}")); + } + AgentEvent::Error { error } => { + observed.push("error".to_string()); + found_auth_error_event = matches!( + error, + AgentError::Llm(SdkError::Provider { + kind: ProviderErrorKind::Authentication, + .. + }) + ); + } + AgentEvent::AssistantMessage { .. } => observed.push("message".to_string()), + _ => {} + } + } + + assert_eq!( + observed, + vec![ + "start".to_string(), + "delta:Hel".to_string(), + "replace::None".to_string(), + "error".to_string(), + ] + ); + assert!(found_auth_error_event, "expected auth error event"); + } + #[tokio::test] async fn compaction_triggered_when_over_threshold() { // Tiny context window to trigger compaction diff --git a/lib/crates/fabro-agent/src/subagent.rs b/lib/crates/fabro-agent/src/subagent.rs index 92fb0a587..b29124be4 100644 --- a/lib/crates/fabro-agent/src/subagent.rs +++ b/lib/crates/fabro-agent/src/subagent.rs @@ -80,6 +80,7 @@ impl SubAgentManager { if matches!( &event.event, AgentEvent::TextDelta { .. } + | AgentEvent::AssistantOutputReplace { .. } | AgentEvent::ReasoningDelta { .. } | AgentEvent::ToolCallOutputDelta { .. } | AgentEvent::AssistantTextStart diff --git a/lib/crates/fabro-agent/src/types.rs b/lib/crates/fabro-agent/src/types.rs index 3423b78f9..8fd178b80 100644 --- a/lib/crates/fabro-agent/src/types.rs +++ b/lib/crates/fabro-agent/src/types.rs @@ -94,6 +94,12 @@ pub enum AgentEvent { text: String, }, AssistantTextStart, + /// Replaces the current in-progress assistant output buffers. + AssistantOutputReplace { + text: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + reasoning: Option, + }, AssistantMessage { text: String, model: String, @@ -206,6 +212,7 @@ impl AgentEvent { Self::AssistantTextStart => { debug!(session_id, "Assistant response started"); } + Self::AssistantOutputReplace { .. } => {} Self::AssistantMessage { model, usage, @@ -640,6 +647,23 @@ mod tests { } } + #[test] + fn agent_event_assistant_output_replace_roundtrip() { + let event = AgentEvent::AssistantOutputReplace { + text: "Hello again".into(), + reasoning: Some("Retrying from scratch".into()), + }; + let json = serde_json::to_string(&event).unwrap(); + let deserialized: AgentEvent = serde_json::from_str(&json).unwrap(); + match deserialized { + AgentEvent::AssistantOutputReplace { text, reasoning } => { + assert_eq!(text, "Hello again"); + assert_eq!(reasoning.as_deref(), Some("Retrying from scratch")); + } + _ => panic!("expected AssistantOutputReplace"), + } + } + // --- Phase 4: Typed error event tests --- #[test] diff --git a/lib/crates/fabro-workflows/src/cli/backend.rs b/lib/crates/fabro-workflows/src/cli/backend.rs index d33bf5a3d..b86c92c58 100644 --- a/lib/crates/fabro-workflows/src/cli/backend.rs +++ b/lib/crates/fabro-workflows/src/cli/backend.rs @@ -100,6 +100,7 @@ fn spawn_event_forwarder( AgentEvent::SessionStarted | AgentEvent::SessionEnded | AgentEvent::AssistantTextStart + | AgentEvent::AssistantOutputReplace { .. } | AgentEvent::TextDelta { .. } | AgentEvent::ReasoningDelta { .. } | AgentEvent::ToolCallOutputDelta { .. } diff --git a/lib/crates/fabro-workflows/src/cli/logs.rs b/lib/crates/fabro-workflows/src/cli/logs.rs index 3f6bd6324..d84313912 100644 --- a/lib/crates/fabro-workflows/src/cli/logs.rs +++ b/lib/crates/fabro-workflows/src/cli/logs.rs @@ -573,6 +573,7 @@ pub fn format_event_pretty(line: &str, styles: &fabro_util::terminal::Styles) -> "Agent.SessionStarted" | "Agent.SessionEnded" | "Agent.AssistantTextStart" + | "Agent.AssistantOutputReplace" | "Agent.TextDelta" | "Agent.ReasoningDelta" | "Agent.ToolCallOutputDelta" @@ -815,6 +816,14 @@ mod tests { assert!(format_event_pretty(line, &styles).is_none()); } + #[test] + fn pretty_skips_assistant_output_replace_noise_event() { + let styles = no_color_styles(); + let line = + r#"{"ts":"2026-01-01T14:23:12Z","event":"Agent.AssistantOutputReplace","text":""}"#; + assert!(format_event_pretty(line, &styles).is_none()); + } + #[test] fn pretty_unknown_events_return_none() { let styles = no_color_styles(); diff --git a/lib/crates/fabro-workflows/src/event.rs b/lib/crates/fabro-workflows/src/event.rs index 85a579792..66e2d9849 100644 --- a/lib/crates/fabro-workflows/src/event.rs +++ b/lib/crates/fabro-workflows/src/event.rs @@ -1219,6 +1219,32 @@ mod tests { assert!(!json.contains("reasoning_tokens")); } + #[test] + fn agent_assistant_output_replace_serialization() { + let event = WorkflowRunEvent::Agent { + stage: "code".to_string(), + event: AgentEvent::AssistantOutputReplace { + text: String::new(), + reasoning: None, + }, + }; + let json = serde_json::to_string(&event).unwrap(); + assert!(json.contains("AssistantOutputReplace")); + + let deserialized: WorkflowRunEvent = serde_json::from_str(&json).unwrap(); + match deserialized { + WorkflowRunEvent::Agent { + stage, + event: AgentEvent::AssistantOutputReplace { text, reasoning }, + } => { + assert_eq!(stage, "code"); + assert!(text.is_empty()); + assert_eq!(reasoning, None); + } + _ => panic!("expected Agent(AssistantOutputReplace)"), + } + } + #[test] fn stage_completed_event_serialization_with_new_fields() { use crate::error::FailureClass; @@ -1753,6 +1779,23 @@ mod tests { assert!(!fields.contains_key("stage")); } + #[test] + fn flatten_event_agent_assistant_output_replace() { + let event = WorkflowRunEvent::Agent { + stage: "plan".to_string(), + event: AgentEvent::AssistantOutputReplace { + text: String::new(), + reasoning: None, + }, + }; + let (name, fields) = flatten_event(&event); + assert_eq!(name, "Agent.AssistantOutputReplace"); + assert_eq!(fields["node_id"], "plan"); + assert_eq!(fields["node_label"], "plan"); + assert_eq!(fields["text"], ""); + assert!(!fields.contains_key("stage")); + } + #[test] fn rename_fields_workflow_run_started() { let event = WorkflowRunEvent::WorkflowRunStarted { diff --git a/lib/crates/fabro-workflows/src/retro_agent.rs b/lib/crates/fabro-workflows/src/retro_agent.rs index 99e9d218c..b829d037a 100644 --- a/lib/crates/fabro-workflows/src/retro_agent.rs +++ b/lib/crates/fabro-workflows/src/retro_agent.rs @@ -330,6 +330,7 @@ fn spawn_retro_event_forwarder(session: &Session, emitter: Arc) -> AgentEvent::SessionStarted | AgentEvent::SessionEnded | AgentEvent::AssistantTextStart + | AgentEvent::AssistantOutputReplace { .. } | AgentEvent::TextDelta { .. } | AgentEvent::ReasoningDelta { .. } | AgentEvent::ToolCallOutputDelta { .. }