diff --git a/crates/aether-cli/src/headless/run.rs b/crates/aether-cli/src/headless/run.rs index 5b55b4103..c2ac71922 100644 --- a/crates/aether-cli/src/headless/run.rs +++ b/crates/aether-cli/src/headless/run.rs @@ -254,31 +254,34 @@ fn setup_tracing(verbose: bool) { #[cfg(test)] mod tests { - use aether_core::events::ContextUsage; + use aether_core::events::{ContextUsage, StreamState}; use super::*; #[test] fn format_text_formats_complete_text() { - assert_eq!(format_text(&AgentEvent::text("id", "hello world", true)), Some("hello world".to_string())); + assert_eq!( + format_text(&AgentEvent::text("id", "hello world", StreamState::Complete)), + Some("hello world".to_string()) + ); } #[test] fn format_text_skips_incomplete_text() { - assert_eq!(format_text(&AgentEvent::text("id", "partial", false)), None); + assert_eq!(format_text(&AgentEvent::text("id", "partial", StreamState::Partial)), None); } #[test] fn format_text_formats_complete_thought() { assert_eq!( - format_text(&AgentEvent::thought("id", "reasoning here", true)), + format_text(&AgentEvent::thought("id", "reasoning here", StreamState::Complete)), Some("Thought: reasoning here".to_string()) ); } #[test] fn format_text_skips_incomplete_thought() { - assert_eq!(format_text(&AgentEvent::thought("id", "partial", false)), None); + assert_eq!(format_text(&AgentEvent::thought("id", "partial", StreamState::Partial)), None); } #[test] @@ -419,8 +422,8 @@ mod tests { #[test] fn event_kind_none_for_non_output_fragments() { - assert_eq!(event_kind(&AgentEvent::text("id", "x", false)), None); - assert_eq!(event_kind(&AgentEvent::thought("id", "x", false)), None); + assert_eq!(event_kind(&AgentEvent::text("id", "x", StreamState::Partial)), None); + assert_eq!(event_kind(&AgentEvent::thought("id", "x", StreamState::Partial)), None); assert_eq!( event_kind(&AgentEvent::Tool(ToolEvent::CallUpdate { tool_call_id: "tc1".to_string(), @@ -443,7 +446,7 @@ mod tests { &[] )); assert!(should_emit(&AgentEvent::turn_ended(TurnOutcome::Completed), &[])); - assert!(!should_emit(&AgentEvent::text("id", "x", false), &[])); + assert!(!should_emit(&AgentEvent::text("id", "x", StreamState::Partial), &[])); assert!(!should_emit( &AgentEvent::Tool(ToolEvent::CallUpdate { tool_call_id: "tc1".to_string(), chunk: "x".to_string() }), &[], @@ -478,8 +481,8 @@ mod tests { use clap::ValueEnum; let samples = vec![ - (AgentEvent::text("id", "x", true), CliEventKind::Text), - (AgentEvent::thought("id", "x", true), CliEventKind::Thought), + (AgentEvent::text("id", "x", StreamState::Complete), CliEventKind::Text), + (AgentEvent::thought("id", "x", StreamState::Complete), CliEventKind::Thought), (tool_call_msg(), CliEventKind::ToolCall), (tool_result_msg(), CliEventKind::ToolResult), ( diff --git a/crates/aether-core/src/context/ext.rs b/crates/aether-core/src/context/ext.rs index 8dd7ebf3c..2d9600d4f 100644 --- a/crates/aether-core/src/context/ext.rs +++ b/crates/aether-core/src/context/ext.rs @@ -84,7 +84,7 @@ fn apply_agent_event(ctx: &mut Context, event: &AgentEvent, acc: &mut TurnAccumu #[cfg(test)] mod tests { - use crate::events::TurnOutcome; + use crate::events::{StreamState, TurnOutcome}; use crate::session::{SessionControlEvent, last_agent_from_events}; use super::*; @@ -103,7 +103,7 @@ mod tests { } fn text_complete(chunk: &str) -> AgentEvent { - AgentEvent::text("msg_1", chunk, true) + AgentEvent::text("msg_1", chunk, StreamState::Complete) } fn tool_result(id: &str, name: &str, result: &str) -> AgentEvent { @@ -253,7 +253,7 @@ mod tests { #[test] fn streaming_chunks_are_ignored() { - let ctx = run_agent_events(&[AgentEvent::text("msg_1", "partial", false)]); + let ctx = run_agent_events(&[AgentEvent::text("msg_1", "partial", StreamState::Partial)]); assert_eq!(ctx.message_count(), 1); } @@ -262,7 +262,7 @@ mod tests { let ctx = run_agent_events(&[ text_complete("Turn 1"), AgentEvent::turn_ended(TurnOutcome::Completed), - AgentEvent::text("msg_2", "Turn 2", true), + AgentEvent::text("msg_2", "Turn 2", StreamState::Complete), AgentEvent::turn_ended(TurnOutcome::Completed), ]); assert_eq!(ctx.message_count(), 3); diff --git a/crates/aether-core/src/core/agent.rs b/crates/aether-core/src/core/agent.rs index 646f5f823..7f57169e5 100644 --- a/crates/aether-core/src/core/agent.rs +++ b/crates/aether-core/src/core/agent.rs @@ -4,7 +4,7 @@ use crate::core::prompt_cache_key::derive_prompt_cache_key; pub use crate::core::retry_config::RetryConfig; use crate::events::{ AgentCommand, AgentEvent, AgentObserver, Command, CompactionOutcome, ContextEvent, ContextUsage, LlmCallOutcome, - LlmCallPurpose, ModelEvent, ToolEvent, TurnEvent, TurnOutcome, UserCommand, + LlmCallPurpose, ModelEvent, StreamState, ToolEvent, TurnEvent, TurnOutcome, UserCommand, }; use crate::mcp::run_mcp_task::{McpCommand, ToolExecutionEvent}; use futures::Stream; @@ -211,10 +211,10 @@ impl Agent { let reasoning = AssistantReasoning::from_parts(reasoning_summary_text.clone(), encrypted_reasoning); self.update_context(&message_content, reasoning, completed_tool_calls); - self.emit(AgentEvent::text(&id, &message_content, true)).await; + self.emit(AgentEvent::text(&id, &message_content, StreamState::Complete)).await; if !reasoning_summary_text.is_empty() { - self.emit(AgentEvent::thought(&id, &reasoning_summary_text, true)).await; + self.emit(AgentEvent::thought(&id, &reasoning_summary_text, StreamState::Complete)).await; } } @@ -451,7 +451,7 @@ impl Agent { Reasoning { chunk } => { state.reasoning_summary_text.push_str(&chunk); if let Some(id) = state.current_message_id.clone() { - self.emit(AgentEvent::thought(&id, &chunk, false)).await; + self.emit(AgentEvent::thought(&id, &chunk, StreamState::Partial)).await; } } @@ -498,7 +498,7 @@ impl Agent { state.message_content.push_str(&chunk); if let Some(id) = state.current_message_id.clone() { - self.emit(AgentEvent::text(&id, &chunk, false)).await; + self.emit(AgentEvent::text(&id, &chunk, StreamState::Partial)).await; } } diff --git a/crates/aether-core/src/events/agent_event.rs b/crates/aether-core/src/events/agent_event.rs index 3dcca614e..5dcf52787 100644 --- a/crates/aether-core/src/events/agent_event.rs +++ b/crates/aether-core/src/events/agent_event.rs @@ -1,7 +1,7 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use super::{ContextEvent, MessageEvent, ModelEvent, ToolEvent, TurnEvent, TurnOutcome}; +use super::{ContextEvent, MessageEvent, ModelEvent, StreamState, ToolEvent, TurnEvent, TurnOutcome}; /// A canonical event on the agent's output stream. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] @@ -15,12 +15,20 @@ pub enum AgentEvent { } impl AgentEvent { - pub fn text(message_id: &str, chunk: &str, is_complete: bool) -> Self { - Self::Message(MessageEvent::Text { message_id: message_id.into(), chunk: chunk.into(), is_complete }) + pub fn text(message_id: &str, chunk: &str, state: StreamState) -> Self { + Self::Message(MessageEvent::Text { + message_id: message_id.into(), + chunk: chunk.into(), + is_complete: state.is_complete(), + }) } - pub fn thought(message_id: &str, chunk: &str, is_complete: bool) -> Self { - Self::Message(MessageEvent::Thought { message_id: message_id.into(), chunk: chunk.into(), is_complete }) + pub fn thought(message_id: &str, chunk: &str, state: StreamState) -> Self { + Self::Message(MessageEvent::Thought { + message_id: message_id.into(), + chunk: chunk.into(), + is_complete: state.is_complete(), + }) } pub fn turn_ended(outcome: TurnOutcome) -> Self { @@ -56,7 +64,7 @@ mod tests { #[test] fn serializes_nested_event_contract() { - let event = AgentEvent::text("m1", "hello", true); + let event = AgentEvent::text("m1", "hello", StreamState::Complete); assert_eq!( serde_json::to_value(event).unwrap(), serde_json::json!({"category":"message","event":{"type":"text","message_id":"m1","chunk":"hello","is_complete":true}}) @@ -66,7 +74,7 @@ mod tests { #[test] fn nested_events_roundtrip() { let events = [ - AgentEvent::text("m", "text", true), + AgentEvent::text("m", "text", StreamState::Complete), AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools: vec![] }), AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose: LlmCallPurpose::Chat, @@ -85,7 +93,7 @@ mod tests { #[test] fn turn_outcome_returns_outcome_only_for_turn_end() { assert_eq!(AgentEvent::turn_ended(TurnOutcome::Completed).turn_outcome(), Some(&TurnOutcome::Completed)); - assert_eq!(AgentEvent::text("m", "text", true).turn_outcome(), None); + assert_eq!(AgentEvent::text("m", "text", StreamState::Complete).turn_outcome(), None); assert_eq!(AgentEvent::Turn(TurnEvent::Started { content: vec![] }).turn_outcome(), None); } } diff --git a/crates/aether-core/src/events/message_event.rs b/crates/aether-core/src/events/message_event.rs index ad4d44ba2..455fff721 100644 --- a/crates/aether-core/src/events/message_event.rs +++ b/crates/aether-core/src/events/message_event.rs @@ -1,6 +1,29 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; +/// Whether a streamed text or thought chunk is the final one for its message. +/// +/// `Partial` chunks stream as they arrive; a single `Complete` event carries the +/// full accumulated text and is emitted when the turn wraps up (which may be after +/// the originating LLM call's +/// [`TurnEvent::LlmCallEnded`](crate::events::TurnEvent::LlmCallEnded)). +/// +/// This stands in for the raw `is_complete: bool` on event constructors so that +/// call sites read as `StreamState::Complete` instead of an opaque `true` literal. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamState { + /// More chunks may follow for this message. + Partial, + /// This is the final chunk for the message. + Complete, +} + +impl StreamState { + pub fn is_complete(self) -> bool { + matches!(self, Self::Complete) + } +} + /// Streaming message content from the agent. /// /// Chunks stream with `is_complete: false`; a final event with `is_complete: true` diff --git a/crates/aether-core/src/events/mod.rs b/crates/aether-core/src/events/mod.rs index 4bf5321f7..6c96399a9 100644 --- a/crates/aether-core/src/events/mod.rs +++ b/crates/aether-core/src/events/mod.rs @@ -19,7 +19,7 @@ mod user_message; pub use acp::{aether_tool_name_meta, humanize_tool_name, mcp_tool_name, parse_tool_call_chunk}; pub use agent_event::AgentEvent; pub use context_event::{CompactionOutcome, ContextEvent, ContextUsage}; -pub use message_event::MessageEvent; +pub use message_event::{MessageEvent, StreamState}; pub use model_event::ModelEvent; pub use observer::{AgentObserver, DynObserverFactory, McpRequestInstrumentation, ObserverFactory}; pub use sub_agent_progress::SubAgentProgressPayload; diff --git a/crates/aether-core/src/session.rs b/crates/aether-core/src/session.rs index 4becca35d..afedb2743 100644 --- a/crates/aether-core/src/session.rs +++ b/crates/aether-core/src/session.rs @@ -202,7 +202,7 @@ pub fn last_agent_from_events(initial: Option, events: &[SessionEvent]) #[cfg(test)] mod tests { use super::*; - use crate::events::{LlmCallPurpose, TurnOutcome}; + use crate::events::{LlmCallPurpose, StreamState, TurnOutcome}; fn agent(event: AgentEvent) -> SessionEvent { SessionEvent::Agent(event) @@ -217,7 +217,7 @@ mod tests { delay_ms: 10, })); let cancelled = agent(AgentEvent::Turn(TurnEvent::Ended { outcome: TurnOutcome::Cancelled })); - let partial = agent(AgentEvent::text("m", "partial", false)); + let partial = agent(AgentEvent::text("m", "partial", StreamState::Partial)); let compaction_ended = agent(AgentEvent::Context(ContextEvent::CompactionEnded { outcome: crate::events::CompactionOutcome::Completed, })); diff --git a/crates/aether-core/src/testing/agent_event_builder.rs b/crates/aether-core/src/testing/agent_event_builder.rs index d64f97f8b..06b2437dc 100644 --- a/crates/aether-core/src/testing/agent_event_builder.rs +++ b/crates/aether-core/src/testing/agent_event_builder.rs @@ -1,5 +1,4 @@ -use crate::events::AgentEvent; -use crate::events::ToolEvent; +use crate::events::{AgentEvent, StreamState, ToolEvent}; use llm::{ToolCallError, ToolCallRequest, ToolCallResult}; use serde::Serialize; @@ -20,7 +19,7 @@ impl AgentEventBuilder { pub fn text(mut self, chunks: &[&str]) -> Self { for chunk in chunks { - self.chunks.push(AgentEvent::text(&self.message_id, chunk, false)); + self.chunks.push(AgentEvent::text(&self.message_id, chunk, StreamState::Partial)); self.full_text.push_str(chunk); } self @@ -80,7 +79,7 @@ impl AgentEventBuilder { } pub fn build(mut self) -> Vec { - self.chunks.push(AgentEvent::text(&self.message_id, &self.full_text, true)); + self.chunks.push(AgentEvent::text(&self.message_id, &self.full_text, StreamState::Complete)); self.chunks } diff --git a/crates/aether-evals/src/agents/fake_agent.rs b/crates/aether-evals/src/agents/fake_agent.rs index d4981c206..bbe361ca6 100644 --- a/crates/aether-evals/src/agents/fake_agent.rs +++ b/crates/aether-evals/src/agents/fake_agent.rs @@ -1,6 +1,6 @@ use super::agent::{Agent, AgentRunResult, RunError}; use crate::Task; -use aether_core::events::{AgentEvent, ToolEvent, TurnOutcome}; +use aether_core::events::{AgentEvent, StreamState, ToolEvent, TurnOutcome}; use async_stream::try_stream; use futures::Stream; use llm::ToolCallResult; @@ -21,7 +21,7 @@ impl FakeAgent { pub fn success() -> Self { Self::new(vec![ - AgentEvent::text("fake_1", "Task completed successfully", true), + AgentEvent::text("fake_1", "Task completed successfully", StreamState::Complete), AgentEvent::turn_ended(TurnOutcome::Completed), ]) } @@ -38,7 +38,7 @@ impl FakeAgent { }, result_meta: None, }), - AgentEvent::text("fake_2", "Task completed using tools", true), + AgentEvent::text("fake_2", "Task completed using tools", StreamState::Complete), AgentEvent::turn_ended(TurnOutcome::Completed), ]) } diff --git a/crates/aether-evals/src/judge.rs b/crates/aether-evals/src/judge.rs index 8a540449a..5a3489a6d 100644 --- a/crates/aether-evals/src/judge.rs +++ b/crates/aether-evals/src/judge.rs @@ -465,7 +465,7 @@ fn truncate_chars(value: &str, max_chars: usize) -> String { #[cfg(test)] mod tests { use super::*; - use aether_core::events::{AgentEvent, TurnOutcome}; + use aether_core::events::{AgentEvent, StreamState, TurnOutcome}; use llm::testing::FakeLlmProvider; use llm::{LlmError, ToolCallRequest, ToolCallResult}; @@ -481,14 +481,17 @@ mod tests { }, }); - assert_eq!(get_transcript_line(&AgentEvent::text("msg_1", "hi", true), 100).unwrap(), "[agent] hi"); + assert_eq!( + get_transcript_line(&AgentEvent::text("msg_1", "hi", StreamState::Complete), 100).unwrap(), + "[agent] hi" + ); assert_eq!(get_transcript_line(&call, 100).unwrap(), "[tool-call] bash arguments={}"); assert_eq!(get_transcript_line(&AgentEvent::turn_ended(TurnOutcome::Completed), 100).unwrap(), "[done]"); } #[test] fn transcript_lines_truncate_long_payloads() { - let line = get_transcript_line(&AgentEvent::text("msg_1", &"a".repeat(50), true), 10).unwrap(); + let line = get_transcript_line(&AgentEvent::text("msg_1", &"a".repeat(50), StreamState::Complete), 10).unwrap(); assert_eq!(line, format!("[agent] {}... [truncated]", "a".repeat(10))); } @@ -545,7 +548,7 @@ mod tests { arguments: "{}".to_string(), }, }), - AgentEvent::text("msg_1", "all done", true), + AgentEvent::text("msg_1", "all done", StreamState::Complete), ]; let judge = judge() diff --git a/crates/aether-telemetry/tests/observer_tests.rs b/crates/aether-telemetry/tests/observer_tests.rs index 1cb5236f6..77998e724 100644 --- a/crates/aether-telemetry/tests/observer_tests.rs +++ b/crates/aether-telemetry/tests/observer_tests.rs @@ -4,7 +4,7 @@ use std::time::Duration; use aether_core::core::RetryConfig; use aether_core::events::{ - AgentEvent, AgentObserver, LlmCallOutcome, LlmCallPurpose, ToolEvent, TurnEvent, TurnOutcome, + AgentEvent, AgentObserver, LlmCallOutcome, LlmCallPurpose, StreamState, ToolEvent, TurnEvent, TurnOutcome, }; use aether_core::testing::{AddNumbersRequest, AgentTrace, DivideNumbersRequest, TestScenario, test_agent}; use aether_telemetry::{ @@ -285,7 +285,7 @@ async fn completed_llm_calls_capture_token_usage_breakdown() -> Result<(), Box Result<(), Box> { let trace = AgentTrace::from_events(vec![ AgentEvent::Turn(TurnEvent::Started { content: vec![] }), - AgentEvent::text("m1", "complete response", true), + AgentEvent::text("m1", "complete response", StreamState::Complete), AgentEvent::turn_ended(TurnOutcome::Completed), ]);