bug fixes in error handling

Author: brandonskiser

crates/agent/src/agent/agent_loop/mod.rs

@@ -41,6 +41,7 @@ use tracing::{
 use types::{
     ContentBlock,
     Message,
+    MessageStartEvent,
     MessageStopEvent,
     MetadataEvent,
     Role,
@@ -403,6 +404,8 @@ struct StreamParseState {
     parsing_tool_use: Option<(String, String, String)>,
     /// Buffered metadata event returned from the response stream
     metadata: Option<MetadataEvent>,
+    /// Buffered message start event returned from the response stream
+    message_start: Option<MessageStartEvent>,
     /// Buffered message stop event returned from the response stream
     message_stop: Option<MessageStopEvent>,
     /// Buffered error event returned from the response stream
@@ -425,6 +428,7 @@ impl StreamParseState {
             user_message,
             message_id: None,
             metadata: None,
+            message_start: None,
             message_stop: None,
             stream_err: None,
             ended_time: None,
@@ -433,15 +437,12 @@ impl StreamParseState {
     }
 
     pub fn next(&mut self, ev: Option<StreamResult>, buf: &mut Vec<AgentLoopEventKind>) {
-        if self.errored {
-            if let Some(ev) = ev {
-                warn!(?ev, "ignoring unexpected event after having received an error");
-            }
-            return;
-        }
-
         let Some(ev) = ev else {
             // No event received means the stream has ended.
+            debug_assert!(
+                self.ended_time.is_none(),
+                "unexpected call to next after stream has already ended"
+            );
             self.ended_time = Some(self.ended_time.unwrap_or(Instant::now()));
             self.errored = self.errored || !self.invalid_tool_uses.is_empty();
             let result = self.make_result();
@@ -453,14 +454,31 @@ impl StreamParseState {
             return;
         };
 
+        if self.errored {
+            warn!(?ev, "ignoring unexpected event after having received an error");
+            return;
+        }
+
+        // Debug assertion that we always start with either a MessageStart, or an error.
+        match &ev {
+            StreamResult::Ok(StreamEvent::MessageStart(_)) | StreamResult::Err(_) => (),
+            other @ StreamResult::Ok(_) => debug_assert!(
+                self.message_start.is_some(),
+                "received an unexpected event at the start of the response stream: {:?}",
+                other
+            ),
+        }
+
         // Pushing low-level stream events in case end users want to consume these directly. Likely
         // not required.
         buf.push(AgentLoopEventKind::Stream(ev.clone()));
 
         match ev {
             StreamResult::Ok(s) => match s {
                 StreamEvent::MessageStart(ev) => {
+                    debug_assert!(self.message_start.is_none());
                     debug_assert!(ev.role == Role::Assistant);
+                    self.message_start = Some(ev);
                 },
                 StreamEvent::MessageStop(ev) => {
                     debug_assert!(self.message_stop.is_none());
@@ -547,7 +565,6 @@ impl StreamParseState {
                 );
                 self.stream_err = Some(err);
                 self.errored = true;
-                self.ended_time = Some(Instant::now());
             },
         }
     }

crates/agent/src/agent/rts/mod.rs

@@ -55,6 +55,7 @@ use crate::agent::agent_loop::types::{
     ToolUseBlockDelta,
     ToolUseBlockStart,
 };
+use crate::agent_loop::types::MessageStartEvent;
 use crate::api_client::error::{
     ApiClientError,
     ConverseStreamError,
@@ -387,7 +388,12 @@ struct ResponseParser {
     /// Whether or not the stream has completed.
     ended: bool,
     /// Buffer to hold the next event in [SendMessageOutput].
+    ///
+    /// Required since the RTS stream needs 1 look-ahead token to ensure we don't emit assistant
+    /// response events that are immediately followed by a code reference event.
     peek: Option<ChatResponseStream>,
+    /// Whether or not we have sent a [MessageStartEvent].
+    message_start_pushed: bool,
     /// Whether or not we are currently receiving tool use delta events. Tuple of
     /// `Some((tool_use_id, name))` if true, [None] otherwise.
     parsing_tool_use: Option<(String, String)>,
@@ -421,6 +427,7 @@ impl ResponseParser {
             cancel_token,
             ended: false,
             peek: None,
+            message_start_pushed: false,
             parsing_tool_use: None,
             tool_use_seen: false,
             buf: vec![],
@@ -601,6 +608,14 @@ impl ResponseParser {
             Ok(ev) => {
                 trace!(?ev, "Received new event");
 
+                if !self.message_start_pushed {
+                    self.buf
+                        .push(StreamResult::Ok(StreamEvent::MessageStart(MessageStartEvent {
+                            role: Role::Assistant,
+                        })));
+                    self.message_start_pushed = true;
+                }
+
                 // Track metadata about the chunk.
                 self.time_to_first_chunk
                     .get_or_insert_with(|| self.request_start_time.elapsed());