Ryan WardandClaude Opus 4.6
f278e53b7e
New features: - Inline subagent panels with expand/collapse and click-to-toggle - /context slash command to inspect bedrock_message_history - Child-to-parent question routing with auto-answer for subagents - Subagent token usage and cost merging into parent conversation - Randomized session-colored user avatar silhouettes Bug fixes: - Bedrock: remove orphaned tool_results after compaction - Bedrock: self-healing exchange lookup for out-of-order streaming - Bedrock: append continuation prompt when conversation ends with assistant - Duration sanity check rejects epoch-time artifacts from session restore - Cache hit rate calculation uses actual total_input_tokens - Hide "Time to first token" when value is zero Improvements: - Demote verbose bedrock-debug logs to debug/trace levels - Bedrock tool usage counting falls back to action counting - Remove logout menu item from workspace menu Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1174 lines
49 KiB
Rust
1174 lines
49 KiB
Rust
use std::sync::{Arc, Mutex};
|
|
|
|
use aws_sdk_bedrockruntime::operation::converse_stream::ConverseStreamOutput;
|
|
use aws_sdk_bedrockruntime::types::{
|
|
ContentBlockDelta, ContentBlockStart, ConverseStreamOutput as StreamEvent,
|
|
ReasoningContentBlockDelta, StopReason,
|
|
};
|
|
use futures::stream::BoxStream;
|
|
use uuid::Uuid;
|
|
use warp_multi_agent_api::response_event::stream_finished;
|
|
use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent};
|
|
|
|
use crate::ai::agent::api::Event;
|
|
use crate::server::server_api::AIApiError;
|
|
|
|
use super::convert::{ContentPart, ConversationMessage, MessageContent, MessageRole};
|
|
use super::diagnostic::BedrockDiagnosticLogger;
|
|
|
|
fn json_to_prost_struct(value: &serde_json::Value) -> prost_types::Struct {
|
|
let fields = match value.as_object() {
|
|
Some(map) => map
|
|
.iter()
|
|
.map(|(k, v)| (k.clone(), json_to_prost_value(v)))
|
|
.collect(),
|
|
None => std::collections::BTreeMap::new(),
|
|
};
|
|
prost_types::Struct { fields }
|
|
}
|
|
|
|
fn json_to_prost_value(value: &serde_json::Value) -> prost_types::Value {
|
|
use prost_types::value::Kind;
|
|
let kind = match value {
|
|
serde_json::Value::Null => Kind::NullValue(0),
|
|
serde_json::Value::Bool(b) => Kind::BoolValue(*b),
|
|
serde_json::Value::Number(n) => Kind::NumberValue(n.as_f64().unwrap_or(0.0)),
|
|
serde_json::Value::String(s) => Kind::StringValue(s.clone()),
|
|
serde_json::Value::Array(arr) => Kind::ListValue(prost_types::ListValue {
|
|
values: arr.iter().map(json_to_prost_value).collect(),
|
|
}),
|
|
serde_json::Value::Object(_) => Kind::StructValue(json_to_prost_struct(value)),
|
|
};
|
|
prost_types::Value { kind: Some(kind) }
|
|
}
|
|
|
|
/// Returns the context window size (in tokens) for a given model ID.
|
|
/// Models with "[1m]" in their identifier support 1M token context.
|
|
pub fn context_window_for_model(model_id: &str) -> u32 {
|
|
let lower = model_id.to_lowercase();
|
|
if lower.contains("[1m]") {
|
|
1_000_000
|
|
} else if lower.contains("nova") {
|
|
300_000
|
|
} else if lower.contains("deepseek") {
|
|
128_000
|
|
} else {
|
|
200_000
|
|
}
|
|
}
|
|
|
|
pub fn bedrock_stream_to_response_events(
|
|
mut output: ConverseStreamOutput,
|
|
task_id: String,
|
|
needs_create_task: bool,
|
|
user_query: Option<String>,
|
|
diagnostic_logger: Option<Arc<BedrockDiagnosticLogger>>,
|
|
messages_sent: Arc<Mutex<Vec<ConversationMessage>>>,
|
|
model_id: String,
|
|
is_summarization: bool,
|
|
) -> BoxStream<'static, Event> {
|
|
let request_id = Uuid::new_v4().to_string();
|
|
let conversation_id = Uuid::new_v4().to_string();
|
|
|
|
if let Some(ref logger) = diagnostic_logger {
|
|
logger.set_ids(&conversation_id, &request_id);
|
|
}
|
|
|
|
let stream = async_stream::stream! {
|
|
log::info!("[bedrock] Stream started: task_id={task_id}, request_id={request_id}, needs_create_task={needs_create_task}");
|
|
|
|
if let Some(ref logger) = diagnostic_logger {
|
|
logger.log_stream_event(&format!("StreamInit: request_id={request_id}, conversation_id={conversation_id}"));
|
|
}
|
|
|
|
let init_event = build_stream_init(&request_id, &conversation_id);
|
|
yield Ok(init_event);
|
|
|
|
if needs_create_task {
|
|
log::debug!("[bedrock] Emitting CreateTask to upgrade optimistic root task");
|
|
let create_task_event = build_create_task(&task_id);
|
|
yield Ok(create_task_event);
|
|
}
|
|
|
|
// Emit the user's query as a proto message in the task so it persists
|
|
// across sessions and can be used for the conversation title.
|
|
if let Some(ref query_text) = user_query {
|
|
let user_query_msg = build_user_query_message(&task_id, query_text);
|
|
yield Ok(user_query_msg);
|
|
}
|
|
|
|
let mut current_text_message_id: Option<String> = None;
|
|
let mut buffered_text = String::new();
|
|
let mut text_flushed = false;
|
|
let mut current_tool_use_id = String::new();
|
|
let mut current_tool_name = String::new();
|
|
let mut current_tool_input_json = String::new();
|
|
let mut _has_tool_calls = false;
|
|
let mut input_tokens: i32 = 0;
|
|
let mut output_tokens: i32 = 0;
|
|
let mut cache_read_input_tokens: i32 = 0;
|
|
let mut cache_write_input_tokens: i32 = 0;
|
|
let mut stop_reason = stream_finished::Reason::Done(stream_finished::Done {});
|
|
|
|
// Track full assistant text and tool calls for bedrock_message_history
|
|
let mut history_text = String::new();
|
|
let mut history_tool_calls: Vec<ContentPart> = Vec::new();
|
|
// Synthetic tool_results for unknown/hallucinated tools — these get
|
|
// paired with their tool_use in history so the next request is valid.
|
|
let mut synthetic_tool_results: Vec<ContentPart> = Vec::new();
|
|
|
|
let mut event_count: u32 = 0;
|
|
loop {
|
|
match output.stream.recv().await {
|
|
Ok(Some(event)) => {
|
|
event_count += 1;
|
|
match event {
|
|
StreamEvent::MessageStart(_) => {
|
|
log::debug!("[bedrock] Event #{event_count}: MessageStart");
|
|
}
|
|
StreamEvent::ContentBlockStart(block_start) => {
|
|
log::debug!("[bedrock] Event #{event_count}: ContentBlockStart");
|
|
if let Some(start) = block_start.start() {
|
|
match start {
|
|
ContentBlockStart::ToolUse(tool_start) => {
|
|
_has_tool_calls = true;
|
|
if !buffered_text.is_empty() {
|
|
let msg_id = current_text_message_id
|
|
.clone()
|
|
.unwrap_or_else(|| Uuid::new_v4().to_string());
|
|
if !text_flushed {
|
|
current_text_message_id = Some(msg_id.clone());
|
|
text_flushed = true;
|
|
log::debug!("[bedrock] Flushing buffered text ({} chars) before tool call", buffered_text.len());
|
|
let add_msg = build_add_agent_output_message(
|
|
&task_id,
|
|
&msg_id,
|
|
&buffered_text,
|
|
);
|
|
yield Ok(add_msg);
|
|
} else {
|
|
log::debug!("[bedrock] Flushing remaining buffered text ({} chars) as append before tool call", buffered_text.len());
|
|
let append = build_append_text(
|
|
&task_id,
|
|
&msg_id,
|
|
&buffered_text,
|
|
);
|
|
yield Ok(append);
|
|
}
|
|
buffered_text.clear();
|
|
}
|
|
current_tool_use_id = tool_start.tool_use_id().to_string();
|
|
current_tool_name = tool_start.name().to_string();
|
|
current_tool_input_json.clear();
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
StreamEvent::ContentBlockDelta(delta) => {
|
|
log::trace!("[bedrock] Event #{event_count}: ContentBlockDelta");
|
|
if let Some(d) = delta.delta() {
|
|
match d {
|
|
ContentBlockDelta::Text(text) => {
|
|
log::debug!("[bedrock] TextDelta ({} chars)", text.len());
|
|
history_text.push_str(text);
|
|
if text_flushed {
|
|
let msg_id = current_text_message_id.as_ref().unwrap();
|
|
let append = build_append_text(
|
|
&task_id,
|
|
msg_id,
|
|
text,
|
|
);
|
|
yield Ok(append);
|
|
} else {
|
|
buffered_text.push_str(text);
|
|
// Buffer a few initial deltas so the first
|
|
// AddMessagesToTask carries enough content for
|
|
// the exchange to be fully registered before
|
|
// subsequent AppendToMessageContent events arrive.
|
|
if buffered_text.len() >= 1 {
|
|
let msg_id = Uuid::new_v4().to_string();
|
|
current_text_message_id = Some(msg_id.clone());
|
|
text_flushed = true;
|
|
let add_msg = build_add_agent_output_message(
|
|
&task_id,
|
|
&msg_id,
|
|
&buffered_text,
|
|
);
|
|
yield Ok(add_msg);
|
|
buffered_text.clear();
|
|
}
|
|
}
|
|
}
|
|
ContentBlockDelta::ReasoningContent(reasoning) => {
|
|
if let ReasoningContentBlockDelta::Text(text) = reasoning {
|
|
log::trace!("[bedrock] Reasoning delta ({} chars) - not displayed to user", text.len());
|
|
}
|
|
}
|
|
ContentBlockDelta::ToolUse(tool_delta) => {
|
|
current_tool_input_json.push_str(tool_delta.input());
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
StreamEvent::ContentBlockStop(_) => {
|
|
log::debug!("[bedrock] Event #{event_count}: ContentBlockStop (tool_use_id={:?})", if current_tool_use_id.is_empty() { "none" } else { ¤t_tool_use_id });
|
|
if !current_tool_use_id.is_empty() {
|
|
// Skip suggest_next_prompt — its executor hangs forever
|
|
// waiting for UI interaction that doesn't exist in the
|
|
// Bedrock path.
|
|
if current_tool_name == "suggest_next_prompt" {
|
|
log::info!("[bedrock] Skipping suggest_next_prompt tool call");
|
|
current_tool_use_id.clear();
|
|
current_tool_name.clear();
|
|
current_tool_input_json.clear();
|
|
} else if !is_known_tool(¤t_tool_name) {
|
|
// Unknown/hallucinated tool: record it in history
|
|
// with a paired error result so the conversation
|
|
// doesn't deadlock waiting for a tool_result that
|
|
// will never come.
|
|
log::warn!(
|
|
"[bedrock] Model called unknown tool '{}' (id={}), synthesizing error result",
|
|
current_tool_name, current_tool_use_id
|
|
);
|
|
let input_json: serde_json::Value = serde_json::from_str(¤t_tool_input_json)
|
|
.unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
|
|
history_tool_calls.push(ContentPart::ToolUse {
|
|
tool_use_id: current_tool_use_id.clone(),
|
|
name: current_tool_name.clone(),
|
|
input: input_json,
|
|
});
|
|
// Immediately pair with a synthetic error result
|
|
// so ensure_tool_results_paired doesn't need to
|
|
// fix it up later (and so the executor doesn't hang).
|
|
synthetic_tool_results.push(ContentPart::ToolResult {
|
|
tool_use_id: current_tool_use_id.clone(),
|
|
content: format!(
|
|
"Error: '{}' is not a valid tool. Available tools are: run_shell_command, read_files, apply_file_diffs, grep, file_glob. Please use one of these tools instead.",
|
|
current_tool_name
|
|
),
|
|
is_error: true,
|
|
});
|
|
if let Some(ref logger) = diagnostic_logger {
|
|
logger.log_stream_event(&format!(
|
|
"UnknownToolCall: name={}, id={} — synthesized error result",
|
|
current_tool_name, current_tool_use_id
|
|
));
|
|
}
|
|
current_tool_use_id.clear();
|
|
current_tool_name.clear();
|
|
current_tool_input_json.clear();
|
|
} else {
|
|
log::debug!("[bedrock] Tool call complete: {} ({})", current_tool_name, current_tool_use_id);
|
|
// Track for bedrock_message_history
|
|
let input_json: serde_json::Value = serde_json::from_str(¤t_tool_input_json)
|
|
.unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
|
|
history_tool_calls.push(ContentPart::ToolUse {
|
|
tool_use_id: current_tool_use_id.clone(),
|
|
name: current_tool_name.clone(),
|
|
input: input_json,
|
|
});
|
|
if let Some(ref logger) = diagnostic_logger {
|
|
logger.log_stream_event(&format!(
|
|
"ToolCall: name={}, id={}, input={}",
|
|
current_tool_name, current_tool_use_id, current_tool_input_json
|
|
));
|
|
}
|
|
let tool_msg = build_tool_call_message(
|
|
&task_id,
|
|
¤t_tool_use_id,
|
|
¤t_tool_name,
|
|
¤t_tool_input_json,
|
|
);
|
|
yield Ok(tool_msg);
|
|
current_tool_use_id.clear();
|
|
current_tool_name.clear();
|
|
current_tool_input_json.clear();
|
|
}
|
|
}
|
|
}
|
|
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 {})
|
|
}
|
|
StopReason::MaxTokens => {
|
|
stream_finished::Reason::MaxTokenLimit(
|
|
stream_finished::ReachedMaxTokenLimit {},
|
|
)
|
|
}
|
|
StopReason::ToolUse => {
|
|
stream_finished::Reason::Done(stream_finished::Done {})
|
|
}
|
|
_ => stream_finished::Reason::Other(stream_finished::Other {}),
|
|
};
|
|
}
|
|
StreamEvent::Metadata(metadata) => {
|
|
if let Some(usage) = metadata.usage() {
|
|
input_tokens = usage.input_tokens();
|
|
output_tokens = usage.output_tokens();
|
|
cache_read_input_tokens = usage.cache_read_input_tokens().unwrap_or(0);
|
|
cache_write_input_tokens = usage.cache_write_input_tokens().unwrap_or(0);
|
|
log::info!(
|
|
"[bedrock-debug] Event #{event_count}: Metadata (input={}, output={}, cache_read={}, cache_write={})",
|
|
input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens
|
|
);
|
|
} else {
|
|
log::warn!("[bedrock-debug] Event #{event_count}: Metadata with NO usage data");
|
|
}
|
|
}
|
|
_ => {
|
|
log::info!("[bedrock-debug] Event #{event_count}: Unknown/Other event");
|
|
}
|
|
}
|
|
}
|
|
Ok(None) => {
|
|
log::info!("[bedrock-debug] Stream ended normally after {event_count} events");
|
|
break;
|
|
}
|
|
Err(e) => {
|
|
log::error!("[bedrock-debug] Stream error after {event_count} events: {e}");
|
|
if let Some(ref logger) = diagnostic_logger {
|
|
let error_msg = format!("{e}");
|
|
let debug_error = format!("{e:?}");
|
|
logger.log_stream_error(&error_msg);
|
|
logger.log_result_fail(&error_msg);
|
|
if let Some(path) = logger.dump_error_snapshot(&error_msg, &debug_error) {
|
|
log::error!(
|
|
"[bedrock] Wrote Bedrock failure snapshot to {}",
|
|
path.display()
|
|
);
|
|
}
|
|
}
|
|
if !buffered_text.is_empty() {
|
|
let msg_id = current_text_message_id
|
|
.clone()
|
|
.unwrap_or_else(|| Uuid::new_v4().to_string());
|
|
if !text_flushed {
|
|
let add_msg = build_add_agent_output_message(&task_id, &msg_id, &buffered_text);
|
|
yield Ok(add_msg);
|
|
} else {
|
|
let append = build_append_text(&task_id, &msg_id, &buffered_text);
|
|
yield Ok(append);
|
|
}
|
|
buffered_text.clear();
|
|
}
|
|
yield Err(Arc::new(AIApiError::Stream {
|
|
stream_type: "bedrock_converse",
|
|
source: anyhow::anyhow!("Bedrock stream error: {}", e),
|
|
}));
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
if !buffered_text.is_empty() {
|
|
let msg_id = current_text_message_id
|
|
.clone()
|
|
.unwrap_or_else(|| Uuid::new_v4().to_string());
|
|
log::debug!("[bedrock] Flushing remaining buffered text ({} chars) at stream end", buffered_text.len());
|
|
if !text_flushed {
|
|
let add_msg = build_add_agent_output_message(&task_id, &msg_id, &buffered_text);
|
|
yield Ok(add_msg);
|
|
} else {
|
|
let append = build_append_text(&task_id, &msg_id, &buffered_text);
|
|
yield Ok(append);
|
|
}
|
|
}
|
|
|
|
let cost = estimate_cost_cents(
|
|
input_tokens as u32,
|
|
output_tokens as u32,
|
|
cache_read_input_tokens as u32,
|
|
cache_write_input_tokens as u32,
|
|
&model_id,
|
|
);
|
|
log::info!(
|
|
"[bedrock] Stream finished: {event_count} events, model={model_id}, input_tokens={input_tokens}, output_tokens={output_tokens}, cache_read={cache_read_input_tokens}, cache_write={cache_write_input_tokens}, cost_cents={cost:.4}"
|
|
);
|
|
|
|
// Build and store the assistant message into bedrock_messages_sent
|
|
// so the controller can persist it as part of conversation history.
|
|
{
|
|
let mut parts: Vec<ContentPart> = Vec::new();
|
|
if !history_text.is_empty() {
|
|
parts.push(ContentPart::Text(history_text));
|
|
}
|
|
parts.extend(history_tool_calls);
|
|
|
|
if !parts.is_empty() {
|
|
let assistant_msg = if parts.len() == 1 {
|
|
match parts.remove(0) {
|
|
ContentPart::Text(t) => ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::Text(t),
|
|
},
|
|
ContentPart::ToolUse { tool_use_id, name, input } => ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::ToolUse { tool_use_id, name, input },
|
|
},
|
|
other => ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::MultiPart(vec![other]),
|
|
},
|
|
}
|
|
} else {
|
|
ConversationMessage {
|
|
role: MessageRole::Assistant,
|
|
content: MessageContent::MultiPart(parts),
|
|
}
|
|
};
|
|
|
|
if let Ok(mut sent) = messages_sent.lock() {
|
|
sent.push(assistant_msg);
|
|
|
|
// If the model hallucinated unknown tools, append a user
|
|
// message with synthetic error results so the history is
|
|
// valid for the next Bedrock request (every tool_use must
|
|
// be followed by a tool_result).
|
|
if !synthetic_tool_results.is_empty() {
|
|
let result_msg = if synthetic_tool_results.len() == 1 {
|
|
match synthetic_tool_results.remove(0) {
|
|
ContentPart::ToolResult { tool_use_id, content, is_error } => {
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::ToolResult { tool_use_id, content, is_error },
|
|
}
|
|
}
|
|
_ => unreachable!(),
|
|
}
|
|
} else {
|
|
ConversationMessage {
|
|
role: MessageRole::User,
|
|
content: MessageContent::MultiPart(synthetic_tool_results),
|
|
}
|
|
};
|
|
sent.push(result_msg);
|
|
log::info!(
|
|
"[bedrock] Stored synthetic tool error results in history. Total messages: {}",
|
|
sent.len()
|
|
);
|
|
}
|
|
|
|
log::info!(
|
|
"[bedrock] Stored assistant message in history. Total messages: {}",
|
|
sent.len()
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Some(ref logger) = diagnostic_logger {
|
|
let stop_reason_str = match &stop_reason {
|
|
stream_finished::Reason::Done(_) => "EndTurn",
|
|
stream_finished::Reason::MaxTokenLimit(_) => "MaxTokens",
|
|
_ => "Other",
|
|
};
|
|
logger.log_result_success(input_tokens, output_tokens, stop_reason_str);
|
|
}
|
|
let finished_event = build_stream_finished(
|
|
stop_reason,
|
|
input_tokens,
|
|
output_tokens,
|
|
cache_read_input_tokens,
|
|
cache_write_input_tokens,
|
|
&model_id,
|
|
is_summarization,
|
|
);
|
|
yield Ok(finished_event);
|
|
};
|
|
|
|
Box::pin(stream)
|
|
}
|
|
|
|
pub(crate) fn build_create_task(task_id: &str) -> ResponseEvent {
|
|
let task = api::Task {
|
|
id: task_id.to_string(),
|
|
description: String::new(),
|
|
dependencies: None,
|
|
messages: vec![],
|
|
summary: String::new(),
|
|
server_data: String::new(),
|
|
};
|
|
|
|
let action = ClientAction {
|
|
action: Some(api::client_action::Action::CreateTask(
|
|
api::client_action::CreateTask { task: Some(task) },
|
|
)),
|
|
};
|
|
|
|
ResponseEvent {
|
|
r#type: Some(api::response_event::Type::ClientActions(
|
|
api::response_event::ClientActions {
|
|
actions: vec![action],
|
|
},
|
|
)),
|
|
}
|
|
}
|
|
|
|
fn build_user_query_message(task_id: &str, query_text: &str) -> ResponseEvent {
|
|
let message = api::Message {
|
|
id: Uuid::new_v4().to_string(),
|
|
task_id: task_id.to_string(),
|
|
request_id: String::new(),
|
|
timestamp: None,
|
|
server_message_data: String::new(),
|
|
citations: vec![],
|
|
message: Some(api::message::Message::UserQuery(
|
|
api::message::UserQuery {
|
|
query: query_text.to_string(),
|
|
..Default::default()
|
|
},
|
|
)),
|
|
};
|
|
|
|
let action = ClientAction {
|
|
action: Some(api::client_action::Action::AddMessagesToTask(
|
|
api::client_action::AddMessagesToTask {
|
|
task_id: task_id.to_string(),
|
|
messages: vec![message],
|
|
},
|
|
)),
|
|
};
|
|
|
|
ResponseEvent {
|
|
r#type: Some(api::response_event::Type::ClientActions(
|
|
api::response_event::ClientActions {
|
|
actions: vec![action],
|
|
},
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(super) fn build_stream_init(request_id: &str, conversation_id: &str) -> ResponseEvent {
|
|
ResponseEvent {
|
|
r#type: Some(api::response_event::Type::Init(
|
|
api::response_event::StreamInit {
|
|
conversation_id: conversation_id.to_string(),
|
|
request_id: request_id.to_string(),
|
|
run_id: String::new(),
|
|
},
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(super) fn build_stream_finished(
|
|
reason: stream_finished::Reason,
|
|
input_tokens: i32,
|
|
output_tokens: i32,
|
|
cache_read_input_tokens: i32,
|
|
cache_write_input_tokens: i32,
|
|
model_id: &str,
|
|
is_summarization: bool,
|
|
) -> ResponseEvent {
|
|
let total_tokens =
|
|
(input_tokens + output_tokens + cache_read_input_tokens + cache_write_input_tokens) as u32;
|
|
|
|
let mut byok_token_usage = std::collections::HashMap::new();
|
|
if total_tokens > 0 {
|
|
#[allow(deprecated)]
|
|
byok_token_usage.insert(
|
|
"bedrock".to_string(),
|
|
stream_finished::ModelTokenUsage {
|
|
model_id: String::new(),
|
|
total_tokens,
|
|
token_usage_by_category: std::collections::HashMap::new(),
|
|
},
|
|
);
|
|
}
|
|
|
|
let token_usage = vec![stream_finished::TokenUsage {
|
|
model_id: "bedrock".to_string(),
|
|
total_input: input_tokens as u32,
|
|
output: output_tokens as u32,
|
|
input_cache_read: cache_read_input_tokens as u32,
|
|
input_cache_write: cache_write_input_tokens as u32,
|
|
cost_in_cents: estimate_cost_cents(
|
|
input_tokens as u32,
|
|
output_tokens as u32,
|
|
cache_read_input_tokens as u32,
|
|
cache_write_input_tokens as u32,
|
|
model_id,
|
|
),
|
|
}];
|
|
|
|
let max_context_tokens = context_window_for_model(model_id);
|
|
let context_usage = if max_context_tokens > 0 {
|
|
(input_tokens as f32 + cache_read_input_tokens as f32 + cache_write_input_tokens as f32)
|
|
/ max_context_tokens as f32
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
#[allow(deprecated)]
|
|
let conversation_usage_metadata = Some(stream_finished::ConversationUsageMetadata {
|
|
context_window_usage: context_usage,
|
|
summarized: is_summarization,
|
|
credits_spent: 0.0,
|
|
token_usage: vec![],
|
|
tool_usage_metadata: None,
|
|
warp_token_usage: std::collections::HashMap::new(),
|
|
byok_token_usage,
|
|
});
|
|
|
|
ResponseEvent {
|
|
r#type: Some(api::response_event::Type::Finished(
|
|
api::response_event::StreamFinished {
|
|
reason: Some(reason),
|
|
token_usage,
|
|
should_refresh_model_config: false,
|
|
request_cost: None,
|
|
conversation_usage_metadata,
|
|
},
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Estimates cost in cents based on Bedrock model pricing.
|
|
/// Pricing varies by model (per 1M tokens):
|
|
/// Opus 4.6/4.7: input $15, output $75, cache_read $1.50, cache_write $18.75
|
|
/// Sonnet 4/4.6: input $3, output $15, cache_read $0.30, cache_write $3.75
|
|
/// Haiku 4.5: input $0.80, output $4, cache_read $0.08, cache_write $1.00
|
|
/// Nova Pro: input $0.80, output $3.20
|
|
/// Nova Lite: input $0.06, output $0.24
|
|
/// Nova Micro: input $0.035, output $0.14
|
|
/// DeepSeek R1: input $1.35, output $5.40
|
|
fn estimate_cost_cents(
|
|
input_tokens: u32,
|
|
output_tokens: u32,
|
|
cache_read_tokens: u32,
|
|
cache_write_tokens: u32,
|
|
model_id: &str,
|
|
) -> f32 {
|
|
let lower = model_id.to_lowercase();
|
|
|
|
// (input_per_1m, output_per_1m, cache_read_per_1m, cache_write_per_1m) in dollars
|
|
let (input_rate, output_rate, cache_read_rate, cache_write_rate) =
|
|
if lower.contains("opus") {
|
|
(15.0, 75.0, 1.50, 18.75)
|
|
} else if lower.contains("haiku") {
|
|
(0.80, 4.0, 0.08, 1.0)
|
|
} else if lower.contains("nova-pro") {
|
|
(0.80, 3.20, 0.0, 0.0)
|
|
} else if lower.contains("nova-lite") {
|
|
(0.06, 0.24, 0.0, 0.0)
|
|
} else if lower.contains("nova-micro") {
|
|
(0.035, 0.14, 0.0, 0.0)
|
|
} else if lower.contains("deepseek") {
|
|
(1.35, 5.40, 0.0, 0.0)
|
|
} else {
|
|
// Default to Sonnet pricing
|
|
(3.0, 15.0, 0.30, 3.75)
|
|
};
|
|
|
|
// Convert from dollars per 1M tokens to cents per token
|
|
let input_cost = input_tokens as f64 * input_rate * 100.0 / 1_000_000.0;
|
|
let output_cost = output_tokens as f64 * output_rate * 100.0 / 1_000_000.0;
|
|
let cache_read_cost = cache_read_tokens as f64 * cache_read_rate * 100.0 / 1_000_000.0;
|
|
let cache_write_cost = cache_write_tokens as f64 * cache_write_rate * 100.0 / 1_000_000.0;
|
|
(input_cost + output_cost + cache_read_cost + cache_write_cost) as f32
|
|
}
|
|
|
|
fn build_add_agent_output_message(
|
|
task_id: &str,
|
|
message_id: &str,
|
|
initial_text: &str,
|
|
) -> ResponseEvent {
|
|
let message = api::Message {
|
|
id: message_id.to_string(),
|
|
task_id: task_id.to_string(),
|
|
request_id: String::new(),
|
|
timestamp: None,
|
|
server_message_data: String::new(),
|
|
citations: vec![],
|
|
message: Some(api::message::Message::AgentOutput(
|
|
api::message::AgentOutput {
|
|
text: initial_text.to_string(),
|
|
},
|
|
)),
|
|
};
|
|
|
|
let action = ClientAction {
|
|
action: Some(api::client_action::Action::AddMessagesToTask(
|
|
api::client_action::AddMessagesToTask {
|
|
task_id: task_id.to_string(),
|
|
messages: vec![message],
|
|
},
|
|
)),
|
|
};
|
|
|
|
ResponseEvent {
|
|
r#type: Some(api::response_event::Type::ClientActions(
|
|
api::response_event::ClientActions {
|
|
actions: vec![action],
|
|
},
|
|
)),
|
|
}
|
|
}
|
|
|
|
fn build_append_text(task_id: &str, message_id: &str, text_delta: &str) -> ResponseEvent {
|
|
let message = api::Message {
|
|
id: message_id.to_string(),
|
|
task_id: task_id.to_string(),
|
|
request_id: String::new(),
|
|
timestamp: None,
|
|
server_message_data: String::new(),
|
|
citations: vec![],
|
|
message: Some(api::message::Message::AgentOutput(
|
|
api::message::AgentOutput {
|
|
text: text_delta.to_string(),
|
|
},
|
|
)),
|
|
};
|
|
|
|
let mask = prost_types::FieldMask {
|
|
paths: vec!["agent_output.text".to_string()],
|
|
};
|
|
|
|
let action = ClientAction {
|
|
action: Some(api::client_action::Action::AppendToMessageContent(
|
|
api::client_action::AppendToMessageContent {
|
|
task_id: task_id.to_string(),
|
|
message: Some(message),
|
|
mask: Some(mask),
|
|
},
|
|
)),
|
|
};
|
|
|
|
ResponseEvent {
|
|
r#type: Some(api::response_event::Type::ClientActions(
|
|
api::response_event::ClientActions {
|
|
actions: vec![action],
|
|
},
|
|
)),
|
|
}
|
|
}
|
|
|
|
fn build_tool_call_message(
|
|
task_id: &str,
|
|
tool_use_id: &str,
|
|
tool_name: &str,
|
|
tool_input_json: &str,
|
|
) -> ResponseEvent {
|
|
let input: serde_json::Value =
|
|
serde_json::from_str(tool_input_json).unwrap_or(serde_json::json!({}));
|
|
|
|
let tool = match tool_name {
|
|
"run_shell_command" => {
|
|
let command = input
|
|
.get("command")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
Some(api::message::tool_call::Tool::RunShellCommand(
|
|
api::message::tool_call::RunShellCommand {
|
|
command,
|
|
is_read_only: false,
|
|
uses_pager: true,
|
|
citations: vec![],
|
|
is_risky: false,
|
|
risk_category: 0,
|
|
wait_until_complete_value: None,
|
|
},
|
|
))
|
|
}
|
|
"read_files" => {
|
|
let files = input
|
|
.get("files")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|f| f.as_str())
|
|
.map(|name| api::message::tool_call::read_files::File {
|
|
name: name.to_string(),
|
|
line_ranges: vec![],
|
|
})
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
Some(api::message::tool_call::Tool::ReadFiles(
|
|
api::message::tool_call::ReadFiles { files },
|
|
))
|
|
}
|
|
"apply_file_diffs" => {
|
|
let diffs = input
|
|
.get("diffs")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|d| {
|
|
Some(api::message::tool_call::apply_file_diffs::FileDiff {
|
|
file_path: d.get("file_path")?.as_str()?.to_string(),
|
|
search: d
|
|
.get("search")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
replace: d
|
|
.get("replace")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
})
|
|
})
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
Some(api::message::tool_call::Tool::ApplyFileDiffs(
|
|
api::message::tool_call::ApplyFileDiffs {
|
|
summary: String::new(),
|
|
diffs,
|
|
new_files: vec![],
|
|
deleted_files: vec![],
|
|
v4a_updates: vec![],
|
|
},
|
|
))
|
|
}
|
|
"grep" => {
|
|
let queries = input
|
|
.get("queries")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|q| q.as_str().map(String::from))
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
let path = input
|
|
.get("path")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
Some(api::message::tool_call::Tool::Grep(
|
|
api::message::tool_call::Grep { queries, path },
|
|
))
|
|
}
|
|
"file_glob" => {
|
|
let patterns = input
|
|
.get("patterns")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|p| p.as_str().map(String::from))
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
#[allow(deprecated)]
|
|
Some(api::message::tool_call::Tool::FileGlob(
|
|
api::message::tool_call::FileGlob {
|
|
patterns,
|
|
path: String::new(),
|
|
},
|
|
))
|
|
}
|
|
"search_codebase" => {
|
|
let query = input
|
|
.get("query")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
let codebase_path = input
|
|
.get("path")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
Some(api::message::tool_call::Tool::SearchCodebase(
|
|
api::message::tool_call::SearchCodebase {
|
|
query,
|
|
path_filters: vec![],
|
|
codebase_path,
|
|
},
|
|
))
|
|
}
|
|
"write_to_long_running_shell_command" => {
|
|
let text_input = input
|
|
.get("input")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
Some(api::message::tool_call::Tool::WriteToLongRunningShellCommand(
|
|
api::message::tool_call::WriteToLongRunningShellCommand {
|
|
input: text_input.into_bytes(),
|
|
mode: None,
|
|
command_id: String::new(),
|
|
},
|
|
))
|
|
}
|
|
"read_shell_command_output" => {
|
|
Some(api::message::tool_call::Tool::ReadShellCommandOutput(
|
|
api::message::tool_call::ReadShellCommandOutput {
|
|
command_id: String::new(),
|
|
delay: None,
|
|
},
|
|
))
|
|
}
|
|
"read_mcp_resource" => {
|
|
let server_id = input
|
|
.get("server_id")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
let uri = input
|
|
.get("uri")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
Some(api::message::tool_call::Tool::ReadMcpResource(
|
|
api::message::tool_call::ReadMcpResource {
|
|
uri,
|
|
server_id,
|
|
},
|
|
))
|
|
}
|
|
"read_documents" => {
|
|
let documents = input
|
|
.get("document_ids")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|d| d.as_str())
|
|
.map(|id| api::message::tool_call::read_documents::Document {
|
|
document_id: id.to_string(),
|
|
line_ranges: vec![],
|
|
})
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
Some(api::message::tool_call::Tool::ReadDocuments(
|
|
api::message::tool_call::ReadDocuments { documents },
|
|
))
|
|
}
|
|
"create_documents" => {
|
|
let new_documents = input
|
|
.get("documents")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|d| {
|
|
Some(api::message::tool_call::create_documents::NewDocument {
|
|
title: d.get("title")?.as_str()?.to_string(),
|
|
content: d.get("content")?.as_str()?.to_string(),
|
|
})
|
|
})
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
Some(api::message::tool_call::Tool::CreateDocuments(
|
|
api::message::tool_call::CreateDocuments { new_documents },
|
|
))
|
|
}
|
|
"edit_documents" => {
|
|
let diffs = input
|
|
.get("diffs")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|d| {
|
|
Some(api::message::tool_call::edit_documents::DocumentDiff {
|
|
document_id: d.get("document_id")?.as_str()?.to_string(),
|
|
search: d.get("search").and_then(|v| v.as_str()).unwrap_or("").to_string(),
|
|
replace: d.get("replace").and_then(|v| v.as_str()).unwrap_or("").to_string(),
|
|
})
|
|
})
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
Some(api::message::tool_call::Tool::EditDocuments(
|
|
api::message::tool_call::EditDocuments { diffs },
|
|
))
|
|
}
|
|
"start_agent" => {
|
|
let _name = input.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
|
let prompt = input.get("prompt").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
|
Some(api::message::tool_call::Tool::Subagent(
|
|
api::message::tool_call::Subagent {
|
|
task_id: String::new(),
|
|
payload: prompt,
|
|
metadata: None,
|
|
},
|
|
))
|
|
}
|
|
"send_message_to_agent" => {
|
|
let agent_id = input.get("agent_id").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
|
let message = input.get("message").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
|
Some(api::message::tool_call::Tool::SendMessageToAgent(
|
|
api::SendMessageToAgent {
|
|
addresses: vec![agent_id],
|
|
subject: String::new(),
|
|
message,
|
|
},
|
|
))
|
|
}
|
|
"ask_user_question" => {
|
|
let question_text = input.get("question").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
|
let options: Vec<api::ask_user_question::Option> = input
|
|
.get("options")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|o| o.as_str())
|
|
.map(|label| api::ask_user_question::Option { label: label.to_string() })
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
let question = api::ask_user_question::Question {
|
|
question_id: Uuid::new_v4().to_string(),
|
|
question: question_text,
|
|
question_type: Some(api::ask_user_question::question::QuestionType::MultipleChoice(
|
|
api::ask_user_question::MultipleChoice {
|
|
options,
|
|
is_multiselect: false,
|
|
supports_other: true,
|
|
recommended_option_index: 0,
|
|
},
|
|
)),
|
|
};
|
|
Some(api::message::tool_call::Tool::AskUserQuestion(
|
|
api::AskUserQuestion {
|
|
questions: vec![question],
|
|
},
|
|
))
|
|
}
|
|
"read_skill" => {
|
|
let skill = input.get("skill").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
|
Some(api::message::tool_call::Tool::ReadSkill(
|
|
api::message::tool_call::ReadSkill {
|
|
name: skill.clone(),
|
|
skill_reference: Some(
|
|
api::message::tool_call::read_skill::SkillReference::SkillPath(skill),
|
|
),
|
|
},
|
|
))
|
|
}
|
|
"fetch_conversation" => {
|
|
let conversation_id = input.get("conversation_id").and_then(|v| v.as_str()).unwrap_or("").to_string();
|
|
Some(api::message::tool_call::Tool::FetchConversation(
|
|
api::message::tool_call::FetchConversation { conversation_id },
|
|
))
|
|
}
|
|
"suggest_next_prompt" => {
|
|
let prompt = input
|
|
.get("prompt")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
let label = input
|
|
.get("label")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string();
|
|
Some(api::message::tool_call::Tool::SuggestPrompt(
|
|
api::message::tool_call::SuggestPrompt {
|
|
is_trigger_irrelevant: false,
|
|
display_mode: Some(
|
|
api::message::tool_call::suggest_prompt::DisplayMode::PromptChip(
|
|
api::message::tool_call::suggest_prompt::PromptChip { prompt, label },
|
|
),
|
|
),
|
|
},
|
|
))
|
|
}
|
|
name if name.starts_with("mcp__") => {
|
|
// MCP tool call: parse server and tool name from "mcp__{server}__{tool}"
|
|
let parts: Vec<&str> = name.splitn(3, "__").collect();
|
|
let (server_name, mcp_tool_name) = if parts.len() == 3 {
|
|
(parts[1].to_string(), parts[2].to_string())
|
|
} else {
|
|
(String::new(), name.strip_prefix("mcp__").unwrap_or(name).to_string())
|
|
};
|
|
let args = json_to_prost_struct(&input);
|
|
Some(api::message::tool_call::Tool::CallMcpTool(
|
|
api::message::tool_call::CallMcpTool {
|
|
name: mcp_tool_name,
|
|
args: Some(args),
|
|
server_id: server_name,
|
|
},
|
|
))
|
|
}
|
|
_ => {
|
|
log::error!("[bedrock] build_tool_call_message called with unknown tool: {tool_name}");
|
|
None
|
|
}
|
|
};
|
|
|
|
let message = if let Some(tool_variant) = tool {
|
|
api::Message {
|
|
id: tool_use_id.to_string(),
|
|
task_id: task_id.to_string(),
|
|
request_id: String::new(),
|
|
timestamp: None,
|
|
server_message_data: String::new(),
|
|
citations: vec![],
|
|
message: Some(api::message::Message::ToolCall(api::message::ToolCall {
|
|
tool_call_id: tool_use_id.to_string(),
|
|
tool: Some(tool_variant),
|
|
})),
|
|
}
|
|
} else {
|
|
// Fallback: emit as agent output text so the stream doesn't break,
|
|
// but this should not happen in normal operation.
|
|
log::error!("[bedrock] Emitting unknown tool as text (should have been caught earlier): {tool_name}");
|
|
api::Message {
|
|
id: Uuid::new_v4().to_string(),
|
|
task_id: task_id.to_string(),
|
|
request_id: String::new(),
|
|
timestamp: None,
|
|
server_message_data: String::new(),
|
|
citations: vec![],
|
|
message: Some(api::message::Message::AgentOutput(
|
|
api::message::AgentOutput {
|
|
text: format!(
|
|
"Error: Model attempted to use unknown tool '{}'. This tool does not exist.",
|
|
tool_name
|
|
),
|
|
},
|
|
)),
|
|
}
|
|
};
|
|
|
|
let action = ClientAction {
|
|
action: Some(api::client_action::Action::AddMessagesToTask(
|
|
api::client_action::AddMessagesToTask {
|
|
task_id: task_id.to_string(),
|
|
messages: vec![message],
|
|
},
|
|
)),
|
|
};
|
|
|
|
ResponseEvent {
|
|
r#type: Some(api::response_event::Type::ClientActions(
|
|
api::response_event::ClientActions {
|
|
actions: vec![action],
|
|
},
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Built-in tools that Galaxy knows how to execute directly.
|
|
const KNOWN_TOOLS: &[&str] = &[
|
|
"run_shell_command",
|
|
"read_files",
|
|
"apply_file_diffs",
|
|
"grep",
|
|
"file_glob",
|
|
"search_codebase",
|
|
"write_to_long_running_shell_command",
|
|
"read_shell_command_output",
|
|
"read_mcp_resource",
|
|
"read_documents",
|
|
"create_documents",
|
|
"edit_documents",
|
|
"start_agent",
|
|
"send_message_to_agent",
|
|
"ask_user_question",
|
|
"suggest_next_prompt",
|
|
"read_skill",
|
|
"fetch_conversation",
|
|
];
|
|
|
|
fn is_known_tool(name: &str) -> bool {
|
|
KNOWN_TOOLS.contains(&name) || name.starts_with("mcp__")
|
|
}
|