Bedrock improvements. Getting back conversations, but losing context after tool calls

Cargo.lock

@@ -4686,6 +4686,7 @@ name = "field_mask"
 version = "0.0.0"
 dependencies = [
  "itertools 0.14.0",
+ "log",
  "prost 0.14.3",
  "prost-reflect",
  "prost-types",

app/src/ai/agent/conversation.rs

@@ -2477,20 +2477,38 @@ impl AIConversation {
                 mask: Some(mask),
             }) => {
                 let task_id = TaskId::new(task_id);
-                let exchange_id = self
+                log::info!("[bedrock-debug] AppendToMessageContent: task_id={:?}, message_id={:?}", task_id, message.id);
+                let exchange_id = match self
                     .added_exchanges_by_response
                     .get(response_stream_id)
-                    .ok_or(UpdateConversationError::NoPendingRequest)?
+                {
-                    .iter()
+                    Some(exchanges) => {
-                    .find_map(|new_exchange| {
+                        log::info!("[bedrock-debug] AppendToMessageContent: found {} exchanges for stream", exchanges.len());
-                        (new_exchange.task_id == task_id).then_some(new_exchange.exchange_id)
+                        for ex in exchanges.iter() {
-                    })
+                            log::info!("[bedrock-debug]   exchange: task_id={:?}, exchange_id={:?}", ex.task_id, ex.exchange_id);
-                    .ok_or(UpdateConversationError::ExchangeNotFound)?;
+                        }
+                        match exchanges.iter().find_map(|new_exchange| {
+                            (new_exchange.task_id == task_id).then_some(new_exchange.exchange_id)
+                        }) {
+                            Some(id) => id,
+                            None => {
+                                log::error!("[bedrock-debug] AppendToMessageContent: ExchangeNotFound - no exchange with matching task_id");
+                                return Err(UpdateConversationError::ExchangeNotFound);
+                            }
+                        }
+                    }
+                    None => {
+                        log::error!("[bedrock-debug] AppendToMessageContent: NoPendingRequest - no exchanges for this stream_id");
+                        return Err(UpdateConversationError::NoPendingRequest);
+                    }
+                };
+
+                log::info!("[bedrock-debug] AppendToMessageContent: found exchange_id={:?}", exchange_id);
 
                 let current_todo_list = self.todo_lists.last().cloned();
                 let current_comment_state = self.code_review.as_ref().cloned();
                 // Update the message and get the updated todos op, if any.
-                let todos_op = self
+                let todos_op = match self
                     .task_store
                     .modify_task(&task_id, |task| {
                         task.append_to_message_content(
@@ -2501,8 +2519,24 @@ impl AIConversation {
                             mask,
                         )
                         .map(|msg| msg.todos_op().cloned())
-                    })
+                    }) {
-                    .ok_or(UpdateConversationError::TaskNotFound)??;
+                    Some(result) => {
+                        match result {
+                            Ok(todos_op) => {
+                                log::info!("[bedrock-debug] AppendToMessageContent: append succeeded");
+                                todos_op
+                            }
+                            Err(e) => {
+                                log::error!("[bedrock-debug] AppendToMessageContent: append_to_message_content failed: {e:?}");
+                                return Err(e.into());
+                            }
+                        }
+                    }
+                    None => {
+                        log::error!("[bedrock-debug] AppendToMessageContent: TaskNotFound in task_store");
+                        return Err(UpdateConversationError::TaskNotFound);
+                    }
+                };
                 // Update todo list if needed
                 if let Some(todos_op) = todos_op {
                     update_todo_list_from_todo_op(&mut self.todo_lists, todos_op);

app/src/ai/agent/task.rs

@@ -766,6 +766,12 @@ impl Task {
                 .apply()
                 .map_err(UpdateTaskError::from)?;
 
+        let text_len = updated_message.message.as_ref().map(|m| match m {
+            api::message::Message::AgentOutput(o) => o.text.len(),
+            _ => 0,
+        }).unwrap_or(0);
+        log::info!("[bedrock-debug] append_to_message_content: accumulated text_len={}", text_len);
+
         let id = self.id.clone();
         let exchange_to_update = self
             .exchange_mut(exchange_id)

app/src/ai/bedrock/stream.rs

@@ -55,11 +55,17 @@ pub fn bedrock_stream_to_response_events(
         let mut output_tokens: i32 = 0;
         let mut stop_reason = stream_finished::Reason::Done(stream_finished::Done {});
 
+        let mut event_count: u32 = 0;
         loop {
             match output.stream.recv().await {
-                Ok(Some(event)) => match event {
+                Ok(Some(event)) => {
-                    StreamEvent::MessageStart(_) => {}
+                    event_count += 1;
+                    match event {
+                    StreamEvent::MessageStart(_) => {
+                        log::info!("[bedrock-debug] Event #{event_count}: MessageStart");
+                    }
                     StreamEvent::ContentBlockStart(block_start) => {
+                        log::info!("[bedrock-debug] Event #{event_count}: ContentBlockStart");
                         if let Some(start) = block_start.start() {
                             match start {
                                 ContentBlockStart::ToolUse(tool_start) => {
@@ -98,10 +104,11 @@ pub fn bedrock_stream_to_response_events(
                         }
                     }
                     StreamEvent::ContentBlockDelta(delta) => {
+                        log::info!("[bedrock-debug] Event #{event_count}: ContentBlockDelta");
                         if let Some(d) = delta.delta() {
                             match d {
                                 ContentBlockDelta::Text(text) => {
-                                    log::trace!("[bedrock] Text delta ({} chars)", text.len());
+                                    log::info!("[bedrock-debug] Event #{event_count}: TextDelta ({} chars): {:?}", text.len(), &text[..text.len().min(80)]);
                                     if text_flushed {
                                         let msg_id = current_text_message_id.as_ref().unwrap();
                                         let append = build_append_text(
@@ -139,6 +146,7 @@ pub fn bedrock_stream_to_response_events(
                         }
                     }
                     StreamEvent::ContentBlockStop(_) => {
+                        log::info!("[bedrock-debug] Event #{event_count}: ContentBlockStop (tool_use_id={:?})", if current_tool_use_id.is_empty() { "none" } else { &current_tool_use_id });
                         if !current_tool_use_id.is_empty() {
                             log::debug!("[bedrock] Tool call complete: {} ({})", current_tool_name, current_tool_use_id);
                             if let Some(ref logger) = diagnostic_logger {
@@ -160,6 +168,7 @@ pub fn bedrock_stream_to_response_events(
                         }
                     }
                     StreamEvent::MessageStop(stop) => {
+                        log::info!("[bedrock-debug] Event #{event_count}: MessageStop (reason={:?})", stop.stop_reason());
                         stop_reason = match stop.stop_reason() {
                             StopReason::EndTurn => {
                                 stream_finished::Reason::Done(stream_finished::Done {})
@@ -176,19 +185,23 @@ pub fn bedrock_stream_to_response_events(
                         };
                     }
                     StreamEvent::Metadata(metadata) => {
+                        log::info!("[bedrock-debug] Event #{event_count}: Metadata");
                         if let Some(usage) = metadata.usage() {
                             input_tokens = usage.input_tokens();
                             output_tokens = usage.output_tokens();
                         }
                     }
-                    _ => {}
+                    _ => {
-                },
+                        log::info!("[bedrock-debug] Event #{event_count}: Unknown/Other event");
+                    }
+                }
+                }
                 Ok(None) => {
-                    log::info!("[bedrock] Stream ended normally");
+                    log::info!("[bedrock-debug] Stream ended normally after {event_count} events");
                     break;
                 }
                 Err(e) => {
-                    log::error!("[bedrock] Stream error: {e}");
+                    log::error!("[bedrock-debug] Stream error after {event_count} events: {e}");
                     if let Some(ref logger) = diagnostic_logger {
                         logger.log_stream_error(&format!("{e}"));
                     }
@@ -378,7 +391,7 @@ fn build_append_text(task_id: &str, message_id: &str, text_delta: &str) -> Respo
     };
 
     let mask = prost_types::FieldMask {
-        paths: vec!["message.agent_output.text".to_string()],
+        paths: vec!["agent_output.text".to_string()],
     };
 
     let action = ClientAction {

app/src/ai/blocklist/block.rs

@@ -1688,6 +1688,14 @@ impl AIBlock {
         let status = self.model.status(ctx);
         let is_udi_enabled = InputSettings::as_ref(ctx).is_universal_developer_input_enabled(ctx);
 
+        log::info!("[bedrock-debug] on_output_status_update: status={}", match &status {
+            AIBlockOutputStatus::Pending => "Pending".to_string(),
+            AIBlockOutputStatus::PartiallyReceived { output } => format!("PartiallyReceived(messages={})", output.get().messages.len()),
+            AIBlockOutputStatus::Complete { output } => format!("Complete(messages={})", output.get().messages.len()),
+            AIBlockOutputStatus::Cancelled { .. } => "Cancelled".to_string(),
+            AIBlockOutputStatus::Failed { .. } => "Failed".to_string(),
+        });
+
         match status {
             AIBlockOutputStatus::Pending => {
                 self.requested_action_ids.clear();

app/src/ai/blocklist/block/model/model_impl.rs

@@ -203,6 +203,7 @@ where
         let exchange_id = self.exchange_id;
         let conversation_id = self.conversation_id;
         let history_model = BlocklistAIHistoryModel::handle(ctx);
+        log::info!("[bedrock-debug] on_updated_output: subscribing for exchange_id={:?}", exchange_id);
         ctx.subscribe_to_model(&history_model, move |me, _, event, ctx| {
             let BlocklistAIHistoryEvent::UpdatedStreamingExchange {
                 exchange_id: event_exchange_id,
@@ -213,6 +214,7 @@ where
                 return;
             };
             if *event_exchange_id == exchange_id {
+                log::info!("[bedrock-debug] on_updated_output: callback fired for matching exchange_id={:?}", exchange_id);
                 callback(me, ctx);
             } else if *event_conversation_id == conversation_id {
                 ctx.notify();

app/src/ai/blocklist/block/view_impl.rs

@@ -826,6 +826,7 @@ impl View for AIBlock {
     }
 
     fn render(&self, app: &AppContext) -> Box<dyn Element> {
+        log::info!("[bedrock-debug] AIBlock::render() called");
         // When the AI block is hidden, we don't need to render anything.
         if self.is_hidden(app) {
             return ConstrainedBox::new(Empty::new().finish())

app/src/ai/blocklist/block/view_impl/output.rs

@@ -220,8 +220,17 @@ pub(super) fn render(props: Props, app: &AppContext) -> Box<dyn Element> {
         | AIBlockOutputStatus::Complete { .. }
         | AIBlockOutputStatus::Cancelled { .. }
         | AIBlockOutputStatus::Failed { .. } => {
-            if let Some(output) = status.output_to_render() {
+             if let Some(output) = status.output_to_render() {
                 let output = output.get();
+                let total_text_len: usize = output.messages.iter().map(|m| match &m.message {
+                    AIAgentOutputMessageType::Text(t) => t.sections.iter().map(|s| match s {
+                        AIAgentTextSection::PlainText { text } => text.text().len(),
+                        AIAgentTextSection::Code { code, .. } => code.len(),
+                        _ => 0,
+                    }).sum::<usize>(),
+                    _ => 0,
+                }).sum();
+                log::info!("[bedrock-debug] render output: messages={}, total_text_len={}", output.messages.len(), total_text_len);
                 let is_complete = matches!(status, AIBlockOutputStatus::Complete { .. });
                 let is_output_for_static_prompt_suggestions =
                     props.model.contains_static_prompt_suggestion_input(app);

app/src/ai/blocklist/controller/response_stream.rs

@@ -255,12 +255,26 @@ impl ResponseStream {
         ctx: &mut ModelContext<Self>,
     ) {
         if self.current_request_id.is_none_or(|id| id != request_id) {
+            log::info!("[bedrock-debug] handle_response_stream_event: stale request_id, dropping event");
             return;
         }
         self.time_to_latest_event = Local::now().signed_duration_since(self.start_time);
 
         match &event {
             Ok(response_event) => {
+                let event_type_name = match &response_event.r#type {
+                    Some(warp_multi_agent_api::response_event::Type::Init(_)) => "Init",
+                    Some(warp_multi_agent_api::response_event::Type::ClientActions(a)) => {
+                        log::info!("[bedrock-debug] ResponseStream received ClientActions with {} actions", a.actions.len());
+                        "ClientActions"
+                    }
+                    Some(warp_multi_agent_api::response_event::Type::Finished(f)) => {
+                        log::info!("[bedrock-debug] ResponseStream received Finished (reason={:?})", f.reason.as_ref().map(|r| format!("{r:?}")).unwrap_or("None".into()));
+                        "Finished"
+                    }
+                    None => "None",
+                };
+                log::info!("[bedrock-debug] ResponseStream emitting event type={event_type_name}");
                 if let Some(event_type) = &response_event.r#type {
                     match event_type {
                         warp_multi_agent_api::response_event::Type::Init(init_event) => {
@@ -300,6 +314,7 @@ impl ResponseStream {
                 ctx.emit(ResponseStreamEvent::ReceivedEvent(Consumable::new(event)));
             }
             Err(e) => {
+                log::info!("[bedrock-debug] ResponseStream received ERROR: {e:?}");
                 // Store original error if this is the first error
                 if self.retry_count == 0 {
                     self.original_error = Some(format!("{e:?}"));
@@ -378,7 +393,9 @@ impl ResponseStream {
     }
 
     fn on_response_stream_complete(&mut self, request_id: Uuid, ctx: &mut ModelContext<Self>) {
+        log::info!("[bedrock-debug] on_response_stream_complete called (request_id={request_id})");
         if self.current_request_id.is_none_or(|id| id != request_id) {
+            log::info!("[bedrock-debug] on_response_stream_complete: stale request_id, ignoring");
             return;
         }
         ctx.emit(ResponseStreamEvent::AfterStreamFinished { cancellation: None });

crates/field_mask/Cargo.toml

@@ -8,6 +8,7 @@ license.workspace = true
 
 [dependencies]
 itertools.workspace = true
+log.workspace = true
 prost.workspace = true
 prost-reflect.workspace = true
 prost-types.workspace = true

crates/field_mask/src/lib.rs

@@ -101,11 +101,12 @@ fn apply_path(
     let field_desc = match target.descriptor().get_field_by_name(field_name) {
         Some(f) => f,
         None => {
-            // Applying a field mask on unknown fields are a no-op.
+            log::warn!("[field_mask] Unknown field '{}' in message '{}' (path: {:?}). Available fields: {:?}",
-            //
+                field_name,
-            // This implies the client's API version is outdated with respect
+                target.descriptor().full_name(),
-            // to the server response. Adding fields is backwards-compatible
+                path_segments,
-            // in protobuf, where expected behavior is to no-op.
+                target.descriptor().fields().map(|f| f.name().to_string()).collect::<Vec<_>>()
+            );
             return Ok(());
         }
     };