public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Ryan Ward
2026-05-08 15:24:14 -05:00
parent 9e0af95953
commit fe105e0369
11 changed files with 117 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1
View File
@@ -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
+44 -10
View File
@@ -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()
.find_map(|new_exchange| {
(new_exchange.task_id == task_id).then_some(new_exchange.exchange_id)
})
.ok_or(UpdateConversationError::ExchangeNotFound)?;
{
Some(exchanges) => {
log::info!("[bedrock-debug] AppendToMessageContent: found {} exchanges for stream", exchanges.len());
for ex in exchanges.iter() {
log::info!("[bedrock-debug] exchange: task_id={:?}, exchange_id={:?}", ex.task_id, ex.exchange_id);
}
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
+6
View File
@@ -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
+21 -8
View File
@@ -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 {
StreamEvent::MessageStart(_) => {}
Ok(Some(event)) => {
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
+8
View File
@@ -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
+2
View File
@@ -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
+1
View File
@@ -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
+10 -1
View File
@@ -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
+17
View File
@@ -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
+1
View File
@@ -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
+6 -5
View File
@@ -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.
//
// This implies the client's API version is outdated with respect
// to the server response. Adding fields is backwards-compatible
// in protobuf, where expected behavior is to no-op.
log::warn!("[field_mask] Unknown field '{}' in message '{}' (path: {:?}). Available fields: {:?}",
field_name,
target.descriptor().full_name(),
path_segments,
target.descriptor().fields().map(|f| f.name().to_string()).collect::<Vec<_>>()
);
return Ok(());
}
};