Fix Warping indicator stuck after Bedrock LLM finishes responding
Bedrock was calling suggest_next_prompt tool which created a SuggestPrompt action that waited forever on a oneshot channel for UI interaction that never fires in the Bedrock path, keeping the conversation permanently InProgress. Fixed by filtering the tool from the Bedrock tool list and skipping it at the stream level when the LLM calls it from context history. Also includes: Bedrock cache token tracking, cost estimation, LSP improvements, conversation usage view updates, and external config support. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
95b4708e44
commit
f37a744692
@@ -7,6 +7,7 @@ use aws_sdk_bedrockruntime::Client as BedrockRuntimeClient;
|
||||
|
||||
use crate::settings::ai::BedrockAuthMethod;
|
||||
|
||||
use super::external_config::ExternalBedrockConfig;
|
||||
use super::convert::{build_converse_request, ConversationMessage, ToolDefinition};
|
||||
use super::diagnostic::BedrockDiagnosticLogger;
|
||||
use super::models::apply_cross_region_prefix;
|
||||
@@ -29,6 +30,33 @@ pub struct BedrockClientConfig {
|
||||
pub fallback_to_warp: bool,
|
||||
}
|
||||
|
||||
impl BedrockClientConfig {
|
||||
/// Applies external config (from Claude Code / OpenCode) as fallback values
|
||||
/// when Galaxy's own settings are at their defaults.
|
||||
pub fn with_external_fallbacks(mut self) -> Self {
|
||||
let external = ExternalBedrockConfig::load();
|
||||
if external.is_empty() {
|
||||
return self;
|
||||
}
|
||||
|
||||
if self.profile == "default" {
|
||||
if let Some(profile) = external.profile {
|
||||
log::info!("[bedrock] Using profile from external config: {profile}");
|
||||
self.profile = profile;
|
||||
}
|
||||
}
|
||||
|
||||
if self.region.is_empty() {
|
||||
if let Some(region) = external.region {
|
||||
log::info!("[bedrock] Using region from external config: {region}");
|
||||
self.region = region;
|
||||
}
|
||||
}
|
||||
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum BedrockError {
|
||||
#[error("Bedrock credentials not configured")]
|
||||
|
||||
@@ -872,6 +872,11 @@ pub fn extract_tools(request: &api::Request) -> Vec<ToolDefinition> {
|
||||
tools = default_tool_definitions();
|
||||
}
|
||||
|
||||
// Filter out suggest_next_prompt — its action executor waits on a oneshot
|
||||
// channel for UI interaction that never fires in the Bedrock path, causing
|
||||
// the conversation to stay InProgress forever.
|
||||
tools.retain(|t| t.name != "suggest_next_prompt");
|
||||
|
||||
tools
|
||||
}
|
||||
|
||||
@@ -963,7 +968,6 @@ fn default_tool_definitions() -> Vec<ToolDefinition> {
|
||||
tool_definition_for_name("apply_file_diffs"),
|
||||
tool_definition_for_name("grep"),
|
||||
tool_definition_for_name("file_glob"),
|
||||
tool_definition_for_name("suggest_next_prompt"),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,229 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use crate::settings::ai::BedrockModelConfig;
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ExternalBedrockConfig {
|
||||
pub profile: Option<String>,
|
||||
pub region: Option<String>,
|
||||
pub models: Vec<BedrockModelConfig>,
|
||||
}
|
||||
|
||||
impl ExternalBedrockConfig {
|
||||
pub fn load() -> Self {
|
||||
let claude_config = load_claude_code_config();
|
||||
let opencode_config = load_opencode_config();
|
||||
|
||||
// Claude Code takes priority over OpenCode since it has richer model mappings
|
||||
Self {
|
||||
profile: claude_config
|
||||
.profile
|
||||
.or(opencode_config.profile),
|
||||
region: claude_config
|
||||
.region
|
||||
.or(opencode_config.region),
|
||||
models: if claude_config.models.is_empty() {
|
||||
opencode_config.models
|
||||
} else {
|
||||
claude_config.models
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.profile.is_none() && self.region.is_none() && self.models.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
fn load_claude_code_config() -> ExternalBedrockConfig {
|
||||
let path = dirs::home_dir()
|
||||
.map(|h| h.join(".claude").join("settings.json"))
|
||||
.unwrap_or_default();
|
||||
|
||||
parse_claude_code_config(path)
|
||||
}
|
||||
|
||||
fn parse_claude_code_config(path: PathBuf) -> ExternalBedrockConfig {
|
||||
let contents = match std::fs::read_to_string(&path) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return ExternalBedrockConfig::default(),
|
||||
};
|
||||
|
||||
let json: serde_json::Value = match serde_json::from_str(&contents) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return ExternalBedrockConfig::default(),
|
||||
};
|
||||
|
||||
let env = match json.get("env").and_then(|v| v.as_object()) {
|
||||
Some(e) => e,
|
||||
None => return ExternalBedrockConfig::default(),
|
||||
};
|
||||
|
||||
let profile = env
|
||||
.get("AWS_PROFILE")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
let region = env
|
||||
.get("AWS_REGION")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
let models = parse_claude_code_model_map(env);
|
||||
|
||||
ExternalBedrockConfig {
|
||||
profile,
|
||||
region,
|
||||
models,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_claude_code_model_map(
|
||||
env: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Vec<BedrockModelConfig> {
|
||||
let model_map_str = match env.get("DCP_MODEL_MAP").and_then(|v| v.as_str()) {
|
||||
Some(s) => s,
|
||||
None => return Vec::new(),
|
||||
};
|
||||
|
||||
let model_map: HashMap<String, String> = match serde_json::from_str(model_map_str) {
|
||||
Ok(m) => m,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
|
||||
model_map
|
||||
.into_iter()
|
||||
.map(|(source_id, arn)| {
|
||||
let display_name = derive_display_name(&source_id);
|
||||
BedrockModelConfig {
|
||||
model_id: arn,
|
||||
display_name,
|
||||
vision_supported: true,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn derive_display_name(model_id: &str) -> String {
|
||||
// Strip region prefix (e.g. "us." from "us.anthropic.claude-opus-4-6-v1")
|
||||
let id = model_id
|
||||
.strip_prefix("us.")
|
||||
.or_else(|| model_id.strip_prefix("eu."))
|
||||
.or_else(|| model_id.strip_prefix("ap."))
|
||||
.or_else(|| model_id.strip_prefix("global."))
|
||||
.unwrap_or(model_id);
|
||||
|
||||
// Strip vendor prefix
|
||||
let id = id
|
||||
.strip_prefix("anthropic.")
|
||||
.or_else(|| id.strip_prefix("amazon."))
|
||||
.or_else(|| id.strip_prefix("meta."))
|
||||
.unwrap_or(id);
|
||||
|
||||
// Convert model slug to display name
|
||||
// e.g. "claude-opus-4-6-v1" -> "Claude Opus 4.6"
|
||||
// e.g. "claude-sonnet-4-5-20250929-v1:0" -> "Claude Sonnet 4.5"
|
||||
prettify_model_slug(id)
|
||||
}
|
||||
|
||||
fn prettify_model_slug(slug: &str) -> String {
|
||||
// Remove version suffixes like "-v1:0", "-v1", "-v2:0"
|
||||
let slug = slug
|
||||
.split("-v")
|
||||
.next()
|
||||
.unwrap_or(slug);
|
||||
|
||||
// Remove date suffixes like "-20250514" or "-20251001"
|
||||
let parts: Vec<&str> = slug.split('-').collect();
|
||||
let mut cleaned: Vec<&str> = Vec::new();
|
||||
let mut i = 0;
|
||||
while i < parts.len() {
|
||||
// Skip parts that look like dates (8 digits)
|
||||
if parts[i].len() == 8 && parts[i].chars().all(|c| c.is_ascii_digit()) {
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
cleaned.push(parts[i]);
|
||||
i += 1;
|
||||
}
|
||||
|
||||
// Try to detect version numbers: consecutive single-digit parts become "X.Y"
|
||||
let mut result = String::new();
|
||||
let mut i = 0;
|
||||
while i < cleaned.len() {
|
||||
if i > 0 {
|
||||
result.push(' ');
|
||||
}
|
||||
|
||||
// Check if this and the next part form a version number (e.g. "4" "6" -> "4.6")
|
||||
if cleaned[i].len() == 1
|
||||
&& cleaned[i].chars().all(|c| c.is_ascii_digit())
|
||||
&& i + 1 < cleaned.len()
|
||||
&& cleaned[i + 1].len() == 1
|
||||
&& cleaned[i + 1].chars().all(|c| c.is_ascii_digit())
|
||||
{
|
||||
result.push_str(cleaned[i]);
|
||||
result.push('.');
|
||||
result.push_str(cleaned[i + 1]);
|
||||
i += 2;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Capitalize first letter of each word
|
||||
let word = cleaned[i];
|
||||
let mut chars = word.chars();
|
||||
if let Some(first) = chars.next() {
|
||||
result.push(first.to_ascii_uppercase());
|
||||
result.extend(chars);
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn load_opencode_config() -> ExternalBedrockConfig {
|
||||
let path = dirs::home_dir()
|
||||
.map(|h| h.join(".config").join("opencode").join("opencode.json"))
|
||||
.unwrap_or_default();
|
||||
|
||||
parse_opencode_config(path)
|
||||
}
|
||||
|
||||
fn parse_opencode_config(path: PathBuf) -> ExternalBedrockConfig {
|
||||
let contents = match std::fs::read_to_string(&path) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return ExternalBedrockConfig::default(),
|
||||
};
|
||||
|
||||
let json: serde_json::Value = match serde_json::from_str(&contents) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return ExternalBedrockConfig::default(),
|
||||
};
|
||||
|
||||
let bedrock_opts = json
|
||||
.get("provider")
|
||||
.and_then(|p| p.get("amazon-bedrock"))
|
||||
.and_then(|b| b.get("options"));
|
||||
|
||||
let profile = bedrock_opts
|
||||
.and_then(|o| o.get("profile"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
let region = bedrock_opts
|
||||
.and_then(|o| o.get("region"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
ExternalBedrockConfig {
|
||||
profile,
|
||||
region,
|
||||
models: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "external_config_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,130 @@
|
||||
use std::io::Write;
|
||||
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_claude_code_config_extracts_profile_and_region() {
|
||||
let mut file = NamedTempFile::new().unwrap();
|
||||
write!(
|
||||
file,
|
||||
r#"{{
|
||||
"env": {{
|
||||
"AWS_PROFILE": "coding-assistant",
|
||||
"AWS_REGION": "us-west-2",
|
||||
"CLAUDE_CODE_USE_BEDROCK": "1"
|
||||
}}
|
||||
}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = parse_claude_code_config(file.path().to_path_buf());
|
||||
assert_eq!(config.profile, Some("coding-assistant".to_string()));
|
||||
assert_eq!(config.region, Some("us-west-2".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_claude_code_config_extracts_model_map() {
|
||||
let mut file = NamedTempFile::new().unwrap();
|
||||
write!(
|
||||
file,
|
||||
r#"{{
|
||||
"env": {{
|
||||
"AWS_PROFILE": "test",
|
||||
"AWS_REGION": "us-east-1",
|
||||
"DCP_MODEL_MAP": "{{\"us.anthropic.claude-opus-4-6-v1\": \"arn:aws:bedrock:us-east-1:123456:application-inference-profile/abc123\"}}"
|
||||
}}
|
||||
}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = parse_claude_code_config(file.path().to_path_buf());
|
||||
assert_eq!(config.models.len(), 1);
|
||||
assert_eq!(
|
||||
config.models[0].model_id,
|
||||
"arn:aws:bedrock:us-east-1:123456:application-inference-profile/abc123"
|
||||
);
|
||||
assert!(config.models[0].vision_supported);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_claude_code_config_missing_file() {
|
||||
let config = parse_claude_code_config(PathBuf::from("/nonexistent/path/settings.json"));
|
||||
assert!(config.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_opencode_config_extracts_bedrock_options() {
|
||||
let mut file = NamedTempFile::new().unwrap();
|
||||
write!(
|
||||
file,
|
||||
r#"{{
|
||||
"provider": {{
|
||||
"amazon-bedrock": {{
|
||||
"options": {{
|
||||
"region": "eu-west-1",
|
||||
"profile": "my-profile"
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = parse_opencode_config(file.path().to_path_buf());
|
||||
assert_eq!(config.profile, Some("my-profile".to_string()));
|
||||
assert_eq!(config.region, Some("eu-west-1".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_opencode_config_missing_file() {
|
||||
let config = parse_opencode_config(PathBuf::from("/nonexistent/opencode.json"));
|
||||
assert!(config.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_claude_code_takes_priority_over_opencode() {
|
||||
// When both configs provide a profile, Claude Code wins
|
||||
let claude = ExternalBedrockConfig {
|
||||
profile: Some("claude-profile".to_string()),
|
||||
region: Some("us-east-1".to_string()),
|
||||
models: Vec::new(),
|
||||
};
|
||||
let opencode = ExternalBedrockConfig {
|
||||
profile: Some("opencode-profile".to_string()),
|
||||
region: Some("eu-west-1".to_string()),
|
||||
models: Vec::new(),
|
||||
};
|
||||
|
||||
let merged = ExternalBedrockConfig {
|
||||
profile: claude.profile.or(opencode.profile),
|
||||
region: claude.region.or(opencode.region),
|
||||
models: Vec::new(),
|
||||
};
|
||||
|
||||
assert_eq!(merged.profile, Some("claude-profile".to_string()));
|
||||
assert_eq!(merged.region, Some("us-east-1".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prettify_model_slug() {
|
||||
assert_eq!(prettify_model_slug("claude-opus-4-6"), "Claude Opus 4.6");
|
||||
assert_eq!(
|
||||
prettify_model_slug("claude-sonnet-4-5-20250929"),
|
||||
"Claude Sonnet 4.5"
|
||||
);
|
||||
assert_eq!(prettify_model_slug("claude-haiku-4-5"), "Claude Haiku 4.5");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_derive_display_name_strips_prefixes() {
|
||||
assert_eq!(
|
||||
derive_display_name("us.anthropic.claude-opus-4-6-v1"),
|
||||
"Claude Opus 4.6"
|
||||
);
|
||||
assert_eq!(
|
||||
derive_display_name("global.anthropic.claude-haiku-4-5-20251001-v1:0"),
|
||||
"Claude Haiku 4.5"
|
||||
);
|
||||
}
|
||||
@@ -3,6 +3,7 @@ pub mod convert;
|
||||
pub mod convert_request;
|
||||
pub mod diagnostic;
|
||||
pub mod discovery;
|
||||
pub mod external_config;
|
||||
pub mod models;
|
||||
pub mod stream;
|
||||
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use crate::settings::ai::BedrockModelConfig;
|
||||
|
||||
use super::external_config::ExternalBedrockConfig;
|
||||
|
||||
pub struct DefaultModel {
|
||||
pub model_id: &'static str,
|
||||
pub display_name: &'static str,
|
||||
@@ -55,18 +57,43 @@ pub const DEFAULT_BEDROCK_MODELS: &[DefaultModel] = &[
|
||||
];
|
||||
|
||||
pub fn get_effective_models(user_models: &[BedrockModelConfig]) -> Vec<BedrockModelConfig> {
|
||||
if user_models.is_empty() {
|
||||
DEFAULT_BEDROCK_MODELS
|
||||
if !user_models.is_empty() {
|
||||
return user_models.to_vec();
|
||||
}
|
||||
|
||||
// Fall back to models from external configs (Claude Code / OpenCode)
|
||||
let external = ExternalBedrockConfig::load();
|
||||
if !external.models.is_empty() {
|
||||
log::info!(
|
||||
"[bedrock] Using {} model(s) from external config",
|
||||
external.models.len()
|
||||
);
|
||||
// Merge external models with defaults so the user still sees all defaults
|
||||
let mut models = external.models;
|
||||
let defaults: Vec<BedrockModelConfig> = DEFAULT_BEDROCK_MODELS
|
||||
.iter()
|
||||
.map(|m| BedrockModelConfig {
|
||||
model_id: m.model_id.to_string(),
|
||||
display_name: m.display_name.to_string(),
|
||||
vision_supported: m.vision_supported,
|
||||
})
|
||||
.collect()
|
||||
} else {
|
||||
user_models.to_vec()
|
||||
.collect();
|
||||
for default in defaults {
|
||||
if !models.iter().any(|m| m.model_id == default.model_id) {
|
||||
models.push(default);
|
||||
}
|
||||
}
|
||||
return models;
|
||||
}
|
||||
|
||||
DEFAULT_BEDROCK_MODELS
|
||||
.iter()
|
||||
.map(|m| BedrockModelConfig {
|
||||
model_id: m.model_id.to_string(),
|
||||
display_name: m.display_name.to_string(),
|
||||
vision_supported: m.vision_supported,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn apply_cross_region_prefix(model_id: &str, region: &str) -> String {
|
||||
|
||||
@@ -55,6 +55,8 @@ pub fn bedrock_stream_to_response_events(
|
||||
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
|
||||
@@ -155,31 +157,41 @@ 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 { ¤t_tool_use_id });
|
||||
if !current_tool_use_id.is_empty() {
|
||||
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
|
||||
));
|
||||
// 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 {
|
||||
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();
|
||||
}
|
||||
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) => {
|
||||
@@ -204,6 +216,8 @@ pub fn bedrock_stream_to_response_events(
|
||||
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);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
@@ -265,7 +279,9 @@ pub fn bedrock_stream_to_response_events(
|
||||
}
|
||||
}
|
||||
|
||||
log::info!("[bedrock] Stream finished: input_tokens={input_tokens}, output_tokens={output_tokens}");
|
||||
log::info!(
|
||||
"[bedrock] Stream finished: input_tokens={input_tokens}, output_tokens={output_tokens}, cache_read={cache_read_input_tokens}, cache_write={cache_write_input_tokens}"
|
||||
);
|
||||
|
||||
// Build and store the assistant message into bedrock_messages_sent
|
||||
// so the controller can persist it as part of conversation history.
|
||||
@@ -317,7 +333,13 @@ pub fn bedrock_stream_to_response_events(
|
||||
};
|
||||
logger.log_result_success(input_tokens, output_tokens, stop_reason_str);
|
||||
}
|
||||
let finished_event = build_stream_finished(stop_reason, input_tokens, output_tokens);
|
||||
let finished_event = build_stream_finished(
|
||||
stop_reason,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cache_read_input_tokens,
|
||||
cache_write_input_tokens,
|
||||
);
|
||||
yield Ok(finished_event);
|
||||
};
|
||||
|
||||
@@ -365,8 +387,11 @@ 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,
|
||||
) -> ResponseEvent {
|
||||
let total_tokens = (input_tokens + output_tokens) as u32;
|
||||
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 {
|
||||
@@ -381,6 +406,20 @@ pub(super) fn build_stream_finished(
|
||||
);
|
||||
}
|
||||
|
||||
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,
|
||||
),
|
||||
}];
|
||||
|
||||
#[allow(deprecated)]
|
||||
let conversation_usage_metadata = Some(stream_finished::ConversationUsageMetadata {
|
||||
context_window_usage: 0.0,
|
||||
@@ -396,7 +435,7 @@ pub(super) fn build_stream_finished(
|
||||
r#type: Some(api::response_event::Type::Finished(
|
||||
api::response_event::StreamFinished {
|
||||
reason: Some(reason),
|
||||
token_usage: vec![],
|
||||
token_usage,
|
||||
should_refresh_model_config: false,
|
||||
request_cost: None,
|
||||
conversation_usage_metadata,
|
||||
@@ -405,6 +444,23 @@ pub(super) fn build_stream_finished(
|
||||
}
|
||||
}
|
||||
|
||||
/// Estimates cost in cents based on Anthropic Claude Bedrock pricing.
|
||||
/// Uses Sonnet-tier pricing as a conservative default since we don't
|
||||
/// know the exact model at this layer.
|
||||
/// Pricing (per 1M tokens): input $3, output $15, cache_read $0.30, cache_write $3.75
|
||||
fn estimate_cost_cents(
|
||||
input_tokens: u32,
|
||||
output_tokens: u32,
|
||||
cache_read_tokens: u32,
|
||||
cache_write_tokens: u32,
|
||||
) -> f32 {
|
||||
let input_cost = input_tokens as f64 * 0.3 / 100_000.0;
|
||||
let output_cost = output_tokens as f64 * 1.5 / 100_000.0;
|
||||
let cache_read_cost = cache_read_tokens as f64 * 0.03 / 100_000.0;
|
||||
let cache_write_cost = cache_write_tokens as f64 * 0.375 / 100_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,
|
||||
|
||||
@@ -19,7 +19,7 @@ fn test_build_stream_init_has_valid_ids() {
|
||||
#[test]
|
||||
fn test_build_stream_finished_done_reason() {
|
||||
let reason = stream_finished::Reason::Done(stream_finished::Done {});
|
||||
let event = build_stream_finished(reason, 100, 50);
|
||||
let event = build_stream_finished(reason, 100, 50, 20, 10);
|
||||
|
||||
match event.r#type {
|
||||
Some(api::response_event::Type::Finished(finished)) => {
|
||||
@@ -35,8 +35,16 @@ fn test_build_stream_finished_done_reason() {
|
||||
.get("bedrock")
|
||||
.unwrap()
|
||||
.total_tokens,
|
||||
150
|
||||
180
|
||||
);
|
||||
// Verify token_usage includes cache breakdown
|
||||
assert_eq!(finished.token_usage.len(), 1);
|
||||
let usage = &finished.token_usage[0];
|
||||
assert_eq!(usage.total_input, 100);
|
||||
assert_eq!(usage.output, 50);
|
||||
assert_eq!(usage.input_cache_read, 20);
|
||||
assert_eq!(usage.input_cache_write, 10);
|
||||
assert!(usage.cost_in_cents > 0.0);
|
||||
}
|
||||
other => panic!("Expected Finished event, got {:?}", other),
|
||||
}
|
||||
@@ -45,7 +53,7 @@ fn test_build_stream_finished_done_reason() {
|
||||
#[test]
|
||||
fn test_build_stream_finished_max_token_limit() {
|
||||
let reason = stream_finished::Reason::MaxTokenLimit(stream_finished::ReachedMaxTokenLimit {});
|
||||
let event = build_stream_finished(reason, 200, 100);
|
||||
let event = build_stream_finished(reason, 200, 100, 0, 0);
|
||||
|
||||
match event.r#type {
|
||||
Some(api::response_event::Type::Finished(finished)) => {
|
||||
@@ -61,7 +69,7 @@ fn test_build_stream_finished_max_token_limit() {
|
||||
#[test]
|
||||
fn test_build_stream_finished_other_reason() {
|
||||
let reason = stream_finished::Reason::Other(stream_finished::Other {});
|
||||
let event = build_stream_finished(reason, 0, 0);
|
||||
let event = build_stream_finished(reason, 0, 0, 0, 0);
|
||||
|
||||
match event.r#type {
|
||||
Some(api::response_event::Type::Finished(finished)) => {
|
||||
|
||||
Reference in New Issue
Block a user