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

Rebrand to Galaxy, major improvements to Bedrock support, still needs some TLC though

Browse Source
This commit is contained in:
Ryan Ward
2026-05-07 11:29:34 -05:00
parent f4e2475c60
commit a41cbd8cc7
2433 changed files with 14208 additions and 9409 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/bedrock/client.rs
+34 -3
View File
@@ -1,3 +1,5 @@
use std::sync::Arc;
use anyhow::Result;
use aws_config::BehaviorVersion;
use aws_sdk_bedrockruntime::config::Region;
@@ -6,6 +8,7 @@ use aws_sdk_bedrockruntime::Client as BedrockRuntimeClient;
use crate::settings::ai::BedrockAuthMethod;
use super::convert::{build_converse_request, ConversationMessage, ToolDefinition};
use super::diagnostic::BedrockDiagnosticLogger;
use super::models::apply_cross_region_prefix;
use super::stream::bedrock_stream_to_response_events;
use crate::ai::agent::api::ResponseStream;
@@ -102,12 +105,14 @@ impl BedrockClient {
&self,
model_id: &str,
task_id: &str,
needs_create_task: bool,
messages: Vec<ConversationMessage>,
system_prompt: Option<String>,
tools: Vec<ToolDefinition>,
max_tokens: i32,
temperature: Option<f32>,
cross_region_inference: bool,
diagnostic_logger: Option<Arc<BedrockDiagnosticLogger>>,
) -> Result<ResponseStream, BedrockError> {
let effective_model_id = if cross_region_inference {
apply_cross_region_prefix(model_id, &self.region)
@@ -122,8 +127,26 @@ impl BedrockClient {
tools.len()
);
let converted =
build_converse_request(messages, system_prompt, tools, max_tokens, temperature, None, None);
let converted = build_converse_request(
messages.clone(),
system_prompt.clone(),
tools.clone(),
max_tokens,
temperature,
None,
None,
);
if let Some(ref logger) = diagnostic_logger {
logger.log_bedrock_input(
&messages,
&system_prompt,
&tools,
max_tokens,
temperature,
cross_region_inference,
);
}
let mut request = self
.runtime_client
@@ -147,6 +170,9 @@ impl BedrockClient {
} else {
display_msg
};
if let Some(ref logger) = diagnostic_logger {
logger.log_result_fail(&msg);
}
if msg.contains("AccessDenied") || msg.contains("access denied") {
BedrockError::AccessDenied(msg)
} else if msg.contains("ThrottlingException") || msg.contains("throttl") {
@@ -161,7 +187,12 @@ impl BedrockClient {
})?;
log::info!("[bedrock] Stream connected successfully");
Ok(Box::pin(bedrock_stream_to_response_events(output, task_id.to_string())))
Ok(Box::pin(bedrock_stream_to_response_events(
output,
task_id.to_string(),
needs_create_task,
diagnostic_logger,
)))
}
pub fn runtime_client(&self) -> &BedrockRuntimeClient {
app/src/ai/bedrock/convert.rs
+6 -3
View File
@@ -15,16 +15,19 @@ pub struct ConvertedRequest {
pub tool_config: Option<ToolConfiguration>,
}
#[derive(Clone)]
pub struct ConversationMessage {
pub role: MessageRole,
pub content: MessageContent,
}
#[derive(Clone, Debug, PartialEq)]
pub enum MessageRole {
User,
Assistant,
}
#[derive(Clone, Debug)]
pub enum MessageContent {
Text(String),
ToolUse {
@@ -40,6 +43,7 @@ pub enum MessageContent {
MultiPart(Vec<ContentPart>),
}
#[derive(Clone, Debug)]
pub enum ContentPart {
Text(String),
ToolUse {
@@ -54,6 +58,7 @@ pub enum ContentPart {
},
}
#[derive(Clone)]
pub struct ToolDefinition {
pub name: String,
pub description: String,
@@ -96,9 +101,7 @@ fn json_to_document(value: JsonValue) -> Document {
}
}
JsonValue::String(s) => Document::String(s),
JsonValue::Array(arr) => {
Document::Array(arr.into_iter().map(json_to_document).collect())
}
JsonValue::Array(arr) => Document::Array(arr.into_iter().map(json_to_document).collect()),
JsonValue::Object(obj) => {
let map: HashMap<String, Document> = obj
.into_iter()
app/src/ai/bedrock/convert_request.rs
+192 -18
View File
@@ -74,14 +74,187 @@ pub fn extract_messages_from_request(request: &api::Request) -> Vec<Conversation
});
}
}
api::request::input::Type::InitProjectRules(_) => {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(
"Initialize this project. Analyze the codebase structure and files, \
generate an AGENTS.md file documenting project conventions and setup \
instructions, and offer to create a development environment configuration. \
Use the available tools to inspect the project before responding."
.to_string(),
),
});
}
api::request::input::Type::CreateEnvironment(env) => {
let repo_info = if env.repo_paths.is_empty() {
String::new()
} else {
format!(" Repositories: {}", env.repo_paths.join(", "))
};
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(format!(
"Create a development environment for this project. \
Set up necessary dependencies, configuration files, and tooling.{}",
repo_info
)),
});
}
api::request::input::Type::CreateNewProject(project) => {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(format!(
"Create a new project: {}",
project.query
)),
});
}
api::request::input::Type::CloneRepository(repo) => {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(format!(
"Clone the repository at {} and set it up for development.",
repo.url
)),
});
}
api::request::input::Type::AutoCodeDiffQuery(diff) => {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(format!(
"Apply code changes: {}",
diff.query
)),
});
}
api::request::input::Type::ResumeConversation(_) => {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(
"Continue where we left off. Review the conversation history and proceed with the next steps."
.to_string(),
),
});
}
api::request::input::Type::QueryWithCannedResponse(canned) => {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(canned.query.clone()),
});
}
api::request::input::Type::CodeReview(_) => {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(
"Review the following code changes and provide detailed feedback on correctness, style, and potential issues."
.to_string(),
),
});
}
_ => {}
}
}
}
ensure_starts_with_user_message(&mut messages);
ensure_tool_results_paired(&mut messages);
messages
}
fn ensure_starts_with_user_message(messages: &mut Vec<ConversationMessage>) {
if messages.is_empty() {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("Please proceed with the requested task.".to_string()),
});
return;
}
if messages[0].role != MessageRole::User {
messages.insert(
0,
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(
"Please proceed with the requested task.".to_string(),
),
},
);
}
}
fn ensure_tool_results_paired(messages: &mut Vec<ConversationMessage>) {
let mut tool_use_ids: Vec<String> = Vec::new();
let mut tool_result_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
for msg in messages.iter() {
match &msg.content {
MessageContent::ToolUse { tool_use_id, .. } => {
tool_use_ids.push(tool_use_id.clone());
}
MessageContent::ToolResult { tool_use_id, .. } => {
tool_result_ids.insert(tool_use_id.clone());
}
MessageContent::MultiPart(parts) => {
for part in parts {
match part {
super::convert::ContentPart::ToolUse { tool_use_id, .. } => {
tool_use_ids.push(tool_use_id.clone());
}
super::convert::ContentPart::ToolResult { tool_use_id, .. } => {
tool_result_ids.insert(tool_use_id.clone());
}
_ => {}
}
}
}
_ => {}
}
}
let orphaned: Vec<String> = tool_use_ids
.into_iter()
.filter(|id| !tool_result_ids.contains(id))
.collect();
if orphaned.is_empty() {
return;
}
log::debug!(
"[bedrock] Synthesizing {} missing toolResult messages for orphaned tool calls",
orphaned.len()
);
for orphaned_id in &orphaned {
let insert_idx = messages
.iter()
.rposition(|m| match &m.content {
MessageContent::ToolUse { tool_use_id, .. } => tool_use_id == orphaned_id,
MessageContent::MultiPart(parts) => parts.iter().any(|p| matches!(
p,
super::convert::ContentPart::ToolUse { tool_use_id, .. } if tool_use_id == orphaned_id
)),
_ => false,
})
.map(|i| i + 1)
.unwrap_or(messages.len());
messages.insert(
insert_idx,
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: orphaned_id.clone(),
content: "Tool executed successfully.".to_string(),
is_error: false,
},
},
);
}
}
pub fn extract_system_prompt(_request: &api::Request) -> Option<String> {
Some("You are a helpful AI coding assistant. You help users with software engineering tasks including writing code, debugging, and explaining concepts.".to_string())
}
@@ -109,7 +282,10 @@ pub fn extract_tools(request: &api::Request) -> Vec<ToolDefinition> {
match input_type {
api::request::input::Type::UserInputs(user_inputs) => {
for user_input in &user_inputs.inputs {
if let Some(api::request::input::user_inputs::user_input::Input::ToolCallResult(_)) = &user_input.input {
if let Some(
api::request::input::user_inputs::user_input::Input::ToolCallResult(_),
) = &user_input.input
{
break;
}
}
@@ -133,6 +309,10 @@ pub fn extract_tools(request: &api::Request) -> Vec<ToolDefinition> {
}
}
if tools.is_empty() {
tools = default_tool_definitions();
}
tools
}
@@ -232,7 +412,7 @@ fn convert_proto_message(msg: &api::Message) -> Option<ConversationMessage> {
Some(ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: msg.id.clone(),
tool_use_id: tool_call.tool_call_id.clone(),
name,
input,
},
@@ -296,14 +476,12 @@ fn extract_tool_result_content(result: &api::request::input::ToolCallResult) ->
match result_type {
api::request::input::tool_call_result::Result::RunShellCommand(cmd_result) => {
match &cmd_result.result {
Some(
api::run_shell_command_result::Result::CommandFinished(finished),
) => finished.output.clone(),
Some(
api::run_shell_command_result::Result::LongRunningCommandSnapshot(
snapshot,
),
) => snapshot.output.clone(),
Some(api::run_shell_command_result::Result::CommandFinished(finished)) => {
finished.output.clone()
}
Some(api::run_shell_command_result::Result::LongRunningCommandSnapshot(
snapshot,
)) => snapshot.output.clone(),
_ => "Command completed.".to_string(),
}
}
@@ -341,19 +519,15 @@ fn format_tool_call_result(result: &api::message::ToolCallResult) -> String {
match result_type {
api::message::tool_call_result::Result::RunShellCommand(cmd_result) => {
match &cmd_result.result {
Some(
api::run_shell_command_result::Result::CommandFinished(finished),
) => {
Some(api::run_shell_command_result::Result::CommandFinished(finished)) => {
format!(
"Exit code: {}\nOutput: {}",
finished.exit_code, finished.output
)
}
Some(
api::run_shell_command_result::Result::LongRunningCommandSnapshot(
snapshot,
),
) => {
Some(api::run_shell_command_result::Result::LongRunningCommandSnapshot(
snapshot,
)) => {
format!("Output (running): {}", snapshot.output)
}
_ => "Command completed.".to_string(),
app/src/ai/bedrock/convert_tests.rs
+6 -2
View File
@@ -158,7 +158,8 @@ fn test_system_prompt_separated_from_messages() {
#[test]
fn test_empty_system_prompt_produces_empty_vec() {
let result = build_converse_request(vec![], Some("".to_string()), vec![], 4096, None, None, None);
let result =
build_converse_request(vec![], Some("".to_string()), vec![], 4096, None, None, None);
assert!(result.system.is_empty());
let result2 = build_converse_request(vec![], None, vec![], 4096, None, None, None);
@@ -235,7 +236,10 @@ fn test_multipart_content_produces_multiple_blocks() {
let result = build_converse_request(messages, None, vec![], 4096, None, None, None);
assert_eq!(result.messages[0].content().len(), 2);
assert!(matches!(&result.messages[0].content()[0], ContentBlock::Text(_)));
assert!(matches!(
&result.messages[0].content()[0],
ContentBlock::Text(_)
));
assert!(matches!(
&result.messages[0].content()[1],
ContentBlock::ToolUse(_)
app/src/ai/bedrock/diagnostic.rs
+363
View File
@@ -0,0 +1,363 @@
use std::fs::{self, File, OpenOptions};
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use std::sync::Mutex;
use chrono::Utc;
use serde_json::Value as JsonValue;
use super::convert::{ConversationMessage, MessageContent, MessageRole, ToolDefinition};
const ENV_VAR: &str = "GALAXY_BEDROCK_DIAGNOSTICS";
const LOG_FILENAME: &str = "bedrock-diagnostics.log";
const MAX_ROTATIONS: usize = 5;
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Layer {
Protobuf,
Bedrock,
}
impl std::fmt::Display for Layer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Layer::Protobuf => write!(f, "PROTOBUF"),
Layer::Bedrock => write!(f, "BEDROCK"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Direction {
Input,
Stream,
Result,
}
impl std::fmt::Display for Direction {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Direction::Input => write!(f, "INPUT"),
Direction::Stream => write!(f, "STREAM"),
Direction::Result => write!(f, "RESULT"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Status {
Pending,
Success,
Fail,
}
impl std::fmt::Display for Status {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Status::Pending => write!(f, "PENDING"),
Status::Success => write!(f, "SUCCESS"),
Status::Fail => write!(f, "FAIL"),
}
}
}
pub struct BedrockDiagnosticLogger {
writer: Mutex<BufWriter<File>>,
model_id: String,
conversation_id: Mutex<String>,
request_id: Mutex<String>,
task_id: String,
}
impl BedrockDiagnosticLogger {
pub fn try_new(
model_id: &str,
conversation_id: &str,
request_id: &str,
task_id: &str,
) -> Option<Self> {
if !is_enabled() {
return None;
}
let log_path = match log_file_path() {
Some(path) => path,
None => {
log::warn!("[bedrock-diag] Could not determine log directory");
return None;
}
};
if let Some(parent) = log_path.parent() {
let _ = fs::create_dir_all(parent);
}
rotate_if_needed(&log_path);
let file = match OpenOptions::new().create(true).append(true).open(&log_path) {
Ok(f) => f,
Err(e) => {
log::warn!(
"[bedrock-diag] Failed to open log file {:?}: {}",
log_path,
e
);
return None;
}
};
log::info!(
"[bedrock-diag] Diagnostic logging enabled -> {:?}",
log_path
);
Some(Self {
writer: Mutex::new(BufWriter::new(file)),
model_id: model_id.to_string(),
conversation_id: Mutex::new(conversation_id.to_string()),
request_id: Mutex::new(request_id.to_string()),
task_id: task_id.to_string(),
})
}
pub fn set_ids(&self, conversation_id: &str, request_id: &str) {
if let Ok(mut cid) = self.conversation_id.lock() {
*cid = conversation_id.to_string();
}
if let Ok(mut rid) = self.request_id.lock() {
*rid = request_id.to_string();
}
}
pub fn log_protobuf_input(&self, request: &warp_multi_agent_api::Request) {
let payload = format!("{:?}", request);
self.write_line(Layer::Protobuf, Direction::Input, Status::Pending, &payload);
}
pub fn log_bedrock_input(
&self,
messages: &[ConversationMessage],
system_prompt: &Option<String>,
tools: &[ToolDefinition],
max_tokens: i32,
temperature: Option<f32>,
cross_region_inference: bool,
) {
let messages_json = serialize_messages(messages);
let tools_json = serialize_tools(tools);
let payload = serde_json::json!({
"model_id": self.model_id,
"cross_region_inference": cross_region_inference,
"max_tokens": max_tokens,
"temperature": temperature,
"system_prompt": system_prompt,
"messages": messages_json,
"tools": tools_json,
});
self.write_line(
Layer::Bedrock,
Direction::Input,
Status::Pending,
&payload.to_string(),
);
}
pub fn log_stream_event(&self, event_description: &str) {
self.write_line(
Layer::Bedrock,
Direction::Stream,
Status::Success,
event_description,
);
}
pub fn log_stream_error(&self, error: &str) {
self.write_line(Layer::Bedrock, Direction::Stream, Status::Fail, error);
}
pub fn log_result_success(&self, input_tokens: i32, output_tokens: i32, stop_reason: &str) {
let payload = serde_json::json!({
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"stop_reason": stop_reason,
});
self.write_line(
Layer::Bedrock,
Direction::Result,
Status::Success,
&payload.to_string(),
);
}
pub fn log_result_fail(&self, error: &str) {
let payload = serde_json::json!({
"error": error,
});
self.write_line(
Layer::Bedrock,
Direction::Result,
Status::Fail,
&payload.to_string(),
);
}
fn write_line(&self, layer: Layer, direction: Direction, status: Status, payload: &str) {
let timestamp = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
let conversation_id = self
.conversation_id
.lock()
.map(|c| c.clone())
.unwrap_or_default();
let request_id = self
.request_id
.lock()
.map(|r| r.clone())
.unwrap_or_default();
let line = format!(
"[{}][{}][{}][{}][{}][{}][{}][{}] {}\n",
timestamp,
layer,
direction,
self.model_id,
status,
conversation_id,
request_id,
self.task_id,
payload,
);
if let Ok(mut writer) = self.writer.lock() {
let _ = writer.write_all(line.as_bytes());
let _ = writer.flush();
}
}
}
pub fn is_enabled() -> bool {
std::env::var(ENV_VAR)
.map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
.unwrap_or(false)
}
fn log_file_path() -> Option<PathBuf> {
galaxy_logging::log_directory()
.ok()
.map(|dir| dir.join(LOG_FILENAME))
}
fn rotate_if_needed(path: &PathBuf) {
let metadata = match fs::metadata(path) {
Ok(m) => m,
Err(_) => return,
};
const TEN_MB: u64 = 10 * 1024 * 1024;
if metadata.len() < TEN_MB {
return;
}
for i in (0..MAX_ROTATIONS - 1).rev() {
let from = if i == 0 {
path.clone()
} else {
path.with_extension(format!("log.{}", i))
};
let to = path.with_extension(format!("log.{}", i + 1));
let _ = fs::rename(&from, &to);
}
let first_rotation = path.with_extension("log.1");
let _ = fs::rename(path, &first_rotation);
}
fn serialize_messages(messages: &[ConversationMessage]) -> JsonValue {
let entries: Vec<JsonValue> = messages
.iter()
.map(|msg| {
let role = match msg.role {
MessageRole::User => "user",
MessageRole::Assistant => "assistant",
};
let content = match &msg.content {
MessageContent::Text(t) => serde_json::json!({"type": "text", "text": t}),
MessageContent::ToolUse {
tool_use_id,
name,
input,
} => {
serde_json::json!({
"type": "tool_use",
"tool_use_id": tool_use_id,
"name": name,
"input": input,
})
}
MessageContent::ToolResult {
tool_use_id,
content,
is_error,
} => {
serde_json::json!({
"type": "tool_result",
"tool_use_id": tool_use_id,
"content": content,
"is_error": is_error,
})
}
MessageContent::MultiPart(parts) => {
let part_values: Vec<JsonValue> = parts
.iter()
.map(|p| match p {
super::convert::ContentPart::Text(t) => {
serde_json::json!({"type": "text", "text": t})
}
super::convert::ContentPart::ToolUse {
tool_use_id,
name,
input,
} => {
serde_json::json!({
"type": "tool_use",
"tool_use_id": tool_use_id,
"name": name,
"input": input,
})
}
super::convert::ContentPart::ToolResult {
tool_use_id,
content,
is_error,
} => {
serde_json::json!({
"type": "tool_result",
"tool_use_id": tool_use_id,
"content": content,
"is_error": is_error,
})
}
})
.collect();
serde_json::json!({"type": "multi_part", "parts": part_values})
}
};
serde_json::json!({"role": role, "content": content})
})
.collect();
JsonValue::Array(entries)
}
fn serialize_tools(tools: &[ToolDefinition]) -> JsonValue {
let entries: Vec<JsonValue> = tools
.iter()
.map(|t| {
serde_json::json!({
"name": t.name,
"description": t.description,
"input_schema": t.input_schema,
})
})
.collect();
JsonValue::Array(entries)
}
app/src/ai/bedrock/e2e_tests.rs
+2126
View File
File diff suppressed because it is too large Load Diff
app/src/ai/bedrock/integration_tests.rs
+604
View File
@@ -0,0 +1,604 @@
use futures::StreamExt;
use serde_json::json;
use super::client::{BedrockClient, BedrockClientConfig};
use super::convert::{ConversationMessage, MessageContent, MessageRole, ToolDefinition};
use crate::settings::ai::BedrockAuthMethod;
fn get_test_config() -> Option<BedrockClientConfig> {
if std::env::var("BEDROCK_INTEGRATION_TEST").is_err() {
return None;
}
let profile =
std::env::var("BEDROCK_TEST_PROFILE").unwrap_or_else(|_| "coding-assistant".into());
let region = std::env::var("BEDROCK_TEST_REGION").unwrap_or_else(|_| "us-east-1".into());
Some(BedrockClientConfig {
auth_method: BedrockAuthMethod::Profile,
profile,
region,
access_key_id: String::new(),
secret_access_key: String::new(),
cross_region_inference: false,
fallback_to_warp: false,
})
}
fn get_test_model() -> String {
std::env::var("BEDROCK_TEST_MODEL").unwrap_or_else(|_| {
"arn:aws:bedrock:us-east-1:156729649053:application-inference-profile/fma5bsw4oxhy".into()
})
}
struct StreamOutput {
text: String,
tool_calls: Vec<ToolCallInfo>,
finished_reason: Option<String>,
total_tokens: u32,
}
#[derive(Debug)]
struct ToolCallInfo {
name: String,
input_json: String,
}
async fn collect_stream_output(
client: &BedrockClient,
model: &str,
messages: Vec<ConversationMessage>,
system_prompt: Option<String>,
tools: Vec<ToolDefinition>,
) -> StreamOutput {
let stream = client
.converse_stream(
model,
"test-task-id",
true,
messages,
system_prompt,
tools,
8192,
None,
false,
)
.await
.expect("converse_stream should succeed");
let mut text = String::new();
let tool_calls = Vec::new();
let mut finished_reason = None;
let mut total_tokens = 0u32;
let mut stream = stream;
while let Some(event) = stream.next().await {
let event = event.expect("stream event should be Ok");
if let Some(event_type) = event.r#type {
use warp_multi_agent_api::response_event::Type;
match event_type {
Type::ClientActions(actions) => {
for action in actions.actions {
if let Some(action_type) = action.action {
use warp_multi_agent_api::client_action::Action;
match action_type {
Action::AddMessagesToTask(add) => {
for msg in add.messages {
if let Some(msg_content) = msg.message {
use warp_multi_agent_api::message::Message;
match msg_content {
Message::AgentOutput(output) => {
text.push_str(&output.text);
}
_ => {}
}
}
}
}
Action::AppendToMessageContent(append) => {
if let Some(msg) = append.message {
if let Some(msg_content) = msg.message {
use warp_multi_agent_api::message::Message;
if let Message::AgentOutput(output) = msg_content {
text.push_str(&output.text);
}
}
}
}
_ => {}
}
}
}
}
Type::Finished(finished) => {
finished_reason = Some(format!("{:?}", finished.reason));
if let Some(meta) = finished.conversation_usage_metadata {
if let Some(usage) = meta.byok_token_usage.get("bedrock") {
total_tokens = usage.total_tokens;
}
}
}
_ => {}
}
}
}
StreamOutput {
text,
tool_calls,
finished_reason,
total_tokens,
}
}
#[tokio::test]
async fn test_simple_text_response() {
let Some(config) = get_test_config() else {
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
return;
};
let client = BedrockClient::from_config(config)
.await
.expect("client creation");
let model = get_test_model();
let messages = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("Say exactly: Hello there, how are you?".into()),
}];
let output = collect_stream_output(&client, &model, messages, None, vec![]).await;
println!("[test] Text output: {:?}", output.text);
println!("[test] Finished reason: {:?}", output.finished_reason);
println!("[test] Total tokens: {}", output.total_tokens);
assert!(!output.text.is_empty(), "Expected non-empty text response");
assert!(
output.finished_reason.is_some(),
"Expected stream to finish"
);
}
#[tokio::test]
async fn test_simple_with_system_prompt() {
let Some(config) = get_test_config() else {
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
return;
};
let client = BedrockClient::from_config(config)
.await
.expect("client creation");
let model = get_test_model();
let messages = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("What is your name?".into()),
}];
let output = collect_stream_output(
&client,
&model,
messages,
Some("You are a helpful assistant named Warp.".into()),
vec![],
)
.await;
println!("[test] Text output: {:?}", output.text);
assert!(!output.text.is_empty());
assert!(
output.text.to_lowercase().contains("warp"),
"Expected response to mention 'Warp', got: {}",
&output.text[..output.text.len().min(200)]
);
}
#[tokio::test]
async fn test_tool_call_round_trip() {
let Some(config) = get_test_config() else {
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
return;
};
let client = BedrockClient::from_config(config)
.await
.expect("client creation");
let model = get_test_model();
let tools = vec![ToolDefinition {
name: "list_files".into(),
description: "List files in a directory".into(),
input_schema: json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Directory path" }
},
"required": ["path"]
}),
}];
let messages = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("List the files in the /project directory.".into()),
}];
let output = collect_stream_output(&client, &model, messages, None, tools).await;
println!(
"[test] Text: {:?}",
&output.text[..output.text.len().min(200)]
);
println!("[test] Tool calls: {:?}", output.tool_calls);
println!("[test] Finished: {:?}", output.finished_reason);
assert!(
!output.text.is_empty() || !output.tool_calls.is_empty(),
"Expected either text or a tool call"
);
}
#[tokio::test]
async fn test_multi_turn_with_tool_result() {
let Some(config) = get_test_config() else {
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
return;
};
let client = BedrockClient::from_config(config)
.await
.expect("client creation");
let model = get_test_model();
let tools = vec![ToolDefinition {
name: "list_files".into(),
description: "List files in a directory".into(),
input_schema: json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Directory path" }
},
"required": ["path"]
}),
}];
let messages = vec![
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(
"List files in /project and tell me what you see.".into(),
),
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "tool_1".into(),
name: "list_files".into(),
input: json!({"path": "/project"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "tool_1".into(),
content: "README.md\nsrc/\nCargo.toml\n.gitignore".into(),
is_error: false,
},
},
];
let output = collect_stream_output(&client, &model, messages, None, tools).await;
println!(
"[test] Text after tool result: {:?}",
&output.text[..output.text.len().min(300)]
);
assert!(
!output.text.is_empty(),
"Expected text response after tool result"
);
}
#[tokio::test]
async fn test_multi_turn_conversation_with_tools() {
let Some(config) = get_test_config() else {
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
return;
};
let client = BedrockClient::from_config(config)
.await
.expect("client creation");
let model = get_test_model();
let tools = vec![
ToolDefinition {
name: "list_files".into(),
description: "List files in a directory".into(),
input_schema: json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "Directory path" }
},
"required": ["path"]
}),
},
ToolDefinition {
name: "read_file".into(),
description: "Read contents of a file".into(),
input_schema: json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "File path" }
},
"required": ["path"]
}),
},
ToolDefinition {
name: "run_command".into(),
description: "Run a shell command".into(),
input_schema: json!({
"type": "object",
"properties": {
"command": { "type": "string", "description": "Shell command" }
},
"required": ["command"]
}),
},
];
let system = Some("You are a helpful coding assistant.".into());
// Turn 1: Ask model to inspect project
let turn1_messages = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(
"Inspect this project and tell me what it does. Start by listing files in /project."
.into(),
),
}];
let turn1 = collect_stream_output(
&client,
&model,
turn1_messages,
system.clone(),
tools.clone(),
)
.await;
println!(
"[test] Turn 1 text: {:?}",
&turn1.text[..turn1.text.len().min(200)]
);
println!("[test] Turn 1 tool_calls: {:?}", turn1.tool_calls);
// Turn 2: Provide tool result, continue
let turn2_messages = vec![
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(
"Inspect this project and tell me what it does. Start by listing files in /project."
.into(),
),
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "tool_turn1".into(),
name: "list_files".into(),
input: json!({"path": "/project"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "tool_turn1".into(),
content: "README.md\nsrc/main.rs\nsrc/lib.rs\nCargo.toml\ntests/\n.gitignore".into(),
is_error: false,
},
},
];
let turn2 = collect_stream_output(
&client,
&model,
turn2_messages,
system.clone(),
tools.clone(),
)
.await;
println!(
"[test] Turn 2 text: {:?}",
&turn2.text[..turn2.text.len().min(200)]
);
println!("[test] Turn 2 tool_calls: {:?}", turn2.tool_calls);
// Turn 3: Provide README content and ask for summary
let turn3_messages = vec![
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(
"Inspect this project and tell me what it does. Start by listing files in /project."
.into(),
),
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "tool_turn1".into(),
name: "list_files".into(),
input: json!({"path": "/project"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "tool_turn1".into(),
content: "README.md\nsrc/main.rs\nsrc/lib.rs\nCargo.toml\ntests/\n.gitignore".into(),
is_error: false,
},
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "tool_turn2".into(),
name: "read_file".into(),
input: json!({"path": "/project/README.md"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "tool_turn2".into(),
content: "# My CLI Tool\n\nA Rust command-line tool for managing developer workflows.\n\n## Features\n- Task tracking\n- Git integration\n- Custom scripts\n\n## Usage\n```\ncargo run -- <command>\n```".into(),
is_error: false,
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(
"Based on what you've seen, give me a brief summary of this project. Do not use any tools."
.into(),
),
},
];
let turn3 = collect_stream_output(
&client,
&model,
turn3_messages,
system.clone(),
tools.clone(),
)
.await;
println!(
"[test] Turn 3 text: {:?}",
&turn3.text[..turn3.text.len().min(500)]
);
assert!(
!turn3.text.is_empty() || !turn3.tool_calls.is_empty(),
"Expected final summary or tool use after multi-turn conversation"
);
}
#[tokio::test]
async fn test_tool_error_recovery() {
let Some(config) = get_test_config() else {
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
return;
};
let client = BedrockClient::from_config(config)
.await
.expect("client creation");
let model = get_test_model();
let tools = vec![ToolDefinition {
name: "read_file".into(),
description: "Read contents of a file".into(),
input_schema: json!({
"type": "object",
"properties": {
"path": { "type": "string", "description": "File path" }
},
"required": ["path"]
}),
}];
let messages = vec![
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("Read the file /project/config.yaml".into()),
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "tool_err".into(),
name: "read_file".into(),
input: json!({"path": "/project/config.yaml"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "tool_err".into(),
content: "Error: File not found: /project/config.yaml".into(),
is_error: true,
},
},
];
let output = collect_stream_output(&client, &model, messages, None, tools).await;
println!(
"[test] Error recovery text: {:?}",
&output.text[..output.text.len().min(300)]
);
assert!(
!output.text.is_empty(),
"Expected model to respond to tool error"
);
}
#[tokio::test]
async fn test_arn_based_model() {
let Some(config) = get_test_config() else {
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
return;
};
let arn = std::env::var("BEDROCK_TEST_ARN").unwrap_or_else(|_| {
"arn:aws:bedrock:us-east-1:156729649053:application-inference-profile/1tim45pgo320".into()
});
let client = BedrockClient::from_config(config)
.await
.expect("client creation");
let messages = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("Reply with the single word: confirmed".into()),
}];
let output = collect_stream_output(&client, &arn, messages, None, vec![]).await;
println!("[test] ARN model text: {:?}", output.text);
assert!(
!output.text.is_empty(),
"Expected response from ARN-based model"
);
}
#[tokio::test]
async fn test_reasoning_model_output() {
let Some(config) = get_test_config() else {
eprintln!("Skipping: BEDROCK_INTEGRATION_TEST not set");
return;
};
let client = BedrockClient::from_config(config)
.await
.expect("client creation");
let model = get_test_model();
let messages = vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("What is 15 * 37? Show your reasoning step by step.".into()),
}];
let output = collect_stream_output(&client, &model, messages, None, vec![]).await;
println!(
"[test] Reasoning model text ({} chars): {:?}",
output.text.len(),
&output.text[..output.text.len().min(500)]
);
println!("[test] Finished: {:?}", output.finished_reason);
assert!(!output.text.is_empty(), "Expected reasoning output");
assert!(
output.text.contains("555"),
"Expected correct answer (555) in output, got: {}",
&output.text[..output.text.len().min(300)]
);
}
app/src/ai/bedrock/mod.rs
+5
View File
@@ -1,6 +1,7 @@
pub mod client;
pub mod convert;
pub mod convert_request;
pub mod diagnostic;
pub mod discovery;
pub mod models;
pub mod stream;
@@ -8,6 +9,10 @@ pub mod stream;
#[cfg(test)]
mod convert_tests;
#[cfg(test)]
mod e2e_tests;
#[cfg(test)]
mod integration_tests;
#[cfg(test)]
mod models_tests;
#[cfg(test)]
mod stream_tests;
app/src/ai/bedrock/models_tests.rs
+13 -5
View File
@@ -20,7 +20,10 @@ fn test_cross_region_prefix_eu_west() {
#[test]
fn test_cross_region_prefix_ap_northeast_1() {
assert_eq!(
apply_cross_region_prefix("anthropic.claude-3-5-sonnet-20241022-v1:0", "ap-northeast-1"),
apply_cross_region_prefix(
"anthropic.claude-3-5-sonnet-20241022-v1:0",
"ap-northeast-1"
),
"jp.anthropic.claude-3-5-sonnet-20241022-v1:0"
);
}
@@ -28,7 +31,10 @@ fn test_cross_region_prefix_ap_northeast_1() {
#[test]
fn test_cross_region_prefix_ap_southeast_2() {
assert_eq!(
apply_cross_region_prefix("anthropic.claude-3-5-sonnet-20241022-v1:0", "ap-southeast-2"),
apply_cross_region_prefix(
"anthropic.claude-3-5-sonnet-20241022-v1:0",
"ap-southeast-2"
),
"au.anthropic.claude-3-5-sonnet-20241022-v1:0"
);
}
@@ -36,7 +42,10 @@ fn test_cross_region_prefix_ap_southeast_2() {
#[test]
fn test_cross_region_prefix_ap_southeast_1() {
assert_eq!(
apply_cross_region_prefix("anthropic.claude-3-5-sonnet-20241022-v1:0", "ap-southeast-1"),
apply_cross_region_prefix(
"anthropic.claude-3-5-sonnet-20241022-v1:0",
"ap-southeast-1"
),
"apac.anthropic.claude-3-5-sonnet-20241022-v1:0"
);
}
@@ -131,8 +140,7 @@ fn test_is_bedrock_model_arn() {
#[test]
fn test_cross_region_prefix_skips_arn() {
let arn =
"arn:aws:bedrock:us-east-1:156729649053:application-inference-profile/fma5bsw4oxhy";
let arn = "arn:aws:bedrock:us-east-1:156729649053:application-inference-profile/fma5bsw4oxhy";
assert_eq!(apply_cross_region_prefix(arn, "us-east-1"), arn);
}
app/src/ai/bedrock/stream.rs
+257 -49
View File
@@ -13,26 +13,50 @@ use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent};
use crate::ai::agent::api::Event;
use crate::server::server_api::AIApiError;
use super::diagnostic::BedrockDiagnosticLogger;
pub fn bedrock_stream_to_response_events(
mut output: ConverseStreamOutput,
task_id: String,
needs_create_task: bool,
diagnostic_logger: Option<Arc<BedrockDiagnosticLogger>>,
) -> 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}");
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);
}
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 stop_reason = stream_finished::Reason::Done(stream_finished::Done {});
const TEXT_FLUSH_THRESHOLD: usize = 20;
loop {
match output.stream.recv().await {
Ok(Some(event)) => match event {
@@ -41,6 +65,24 @@ pub fn bedrock_stream_to_response_events(
if let Some(start) = block_start.start() {
match start {
ContentBlockStart::ToolUse(tool_start) => {
has_tool_calls = true;
if !text_flushed && !buffered_text.is_empty() {
if buffered_text.len() >= TEXT_FLUSH_THRESHOLD {
let msg_id = Uuid::new_v4().to_string();
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] Discarding short text fragment ({} chars) before tool call: {:?}", buffered_text.len(), &buffered_text);
}
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();
@@ -54,17 +96,7 @@ pub fn bedrock_stream_to_response_events(
match d {
ContentBlockDelta::Text(text) => {
log::trace!("[bedrock] Text delta ({} chars): {:?}", text.len(), &text[..text.len().min(100)]);
if current_text_message_id.is_none() {
let msg_id = Uuid::new_v4().to_string();
current_text_message_id = Some(msg_id.clone());
log::debug!("[bedrock] First text chunk, creating message msg_id={msg_id}");
let add_msg = build_add_agent_output_message(
&task_id,
&msg_id,
text,
);
yield Ok(add_msg);
} else {
if text_flushed {
let msg_id = current_text_message_id.as_ref().unwrap();
let append = build_append_text(
&task_id,
@@ -72,30 +104,26 @@ pub fn bedrock_stream_to_response_events(
text,
);
yield Ok(append);
} else {
buffered_text.push_str(text);
if buffered_text.len() >= TEXT_FLUSH_THRESHOLD {
let msg_id = Uuid::new_v4().to_string();
current_text_message_id = Some(msg_id.clone());
text_flushed = true;
log::debug!("[bedrock] Text reached flush threshold, creating message msg_id={msg_id}");
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)", text.len());
if current_text_message_id.is_none() {
let msg_id = Uuid::new_v4().to_string();
current_text_message_id = Some(msg_id.clone());
log::debug!("[bedrock] First reasoning chunk, creating message msg_id={msg_id}");
let add_msg = build_add_agent_output_message(
&task_id,
&msg_id,
text,
);
yield Ok(add_msg);
} else {
let msg_id = current_text_message_id.as_ref().unwrap();
let append = build_append_text(
&task_id,
msg_id,
text,
);
yield Ok(append);
}
log::trace!("[bedrock] Reasoning delta ({} chars) - not displayed to user", text.len());
}
}
ContentBlockDelta::ToolUse(tool_delta) => {
@@ -108,6 +136,12 @@ pub fn bedrock_stream_to_response_events(
StreamEvent::ContentBlockStop(_) => {
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 {
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,
&current_tool_use_id,
@@ -150,6 +184,9 @@ pub fn bedrock_stream_to_response_events(
}
Err(e) => {
log::error!("[bedrock] Stream error: {e}");
if let Some(ref logger) = diagnostic_logger {
logger.log_stream_error(&format!("{e}"));
}
yield Err(Arc::new(AIApiError::Stream {
stream_type: "bedrock_converse",
source: anyhow::anyhow!("Bedrock stream error: {}", e),
@@ -159,7 +196,24 @@ pub fn bedrock_stream_to_response_events(
}
}
if !text_flushed && !buffered_text.is_empty() && !has_tool_calls {
let msg_id = Uuid::new_v4().to_string();
log::debug!("[bedrock] Flushing remaining buffered text ({} chars) at stream end", buffered_text.len());
let add_msg = build_add_agent_output_message(&task_id, &msg_id, &buffered_text);
yield Ok(add_msg);
} else if !text_flushed && !buffered_text.is_empty() && has_tool_calls {
log::debug!("[bedrock] Discarding short unflushed text ({} chars) - stream ended with tool calls", buffered_text.len());
}
log::info!("[bedrock] Stream finished: input_tokens={input_tokens}, output_tokens={output_tokens}");
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);
yield Ok(finished_event);
};
@@ -167,6 +221,31 @@ pub fn bedrock_stream_to_response_events(
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],
},
)),
}
}
pub(super) fn build_stream_init(request_id: &str, conversation_id: &str) -> ResponseEvent {
ResponseEvent {
r#type: Some(api::response_event::Type::Init(
@@ -304,22 +383,153 @@ fn build_tool_call_message(
tool_name: &str,
tool_input_json: &str,
) -> ResponseEvent {
let _tool_use_id = tool_use_id.to_string();
let _tool_name = tool_name.to_string();
let _tool_input_json = tool_input_json.to_string();
let input: serde_json::Value =
serde_json::from_str(tool_input_json).unwrap_or(serde_json::json!({}));
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::AgentOutput(
api::message::AgentOutput {
text: format!("[Tool call: {} ({})]", _tool_name, _tool_use_id),
},
)),
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: false,
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(),
},
))
}
_ => {
log::warn!("[bedrock] Unknown tool name: {tool_name}, emitting as text");
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 {
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!(
"[Tool call: {} ({})]\nInput: {}",
tool_name, tool_use_id, tool_input_json
),
},
)),
}
};
let action = ClientAction {
@@ -339,5 +549,3 @@ fn build_tool_call_message(
)),
}
}
app/src/ai/bedrock/stream_tests.rs
+46 -3
View File
@@ -29,7 +29,14 @@ fn test_build_stream_finished_done_reason() {
));
assert!(!finished.should_refresh_model_config);
let metadata = finished.conversation_usage_metadata.unwrap();
assert_eq!(metadata.byok_token_usage.get("bedrock").unwrap().total_tokens, 150);
assert_eq!(
metadata
.byok_token_usage
.get("bedrock")
.unwrap()
.total_tokens,
150
);
}
other => panic!("Expected Finished event, got {:?}", other),
}
@@ -37,8 +44,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 reason = stream_finished::Reason::MaxTokenLimit(stream_finished::ReachedMaxTokenLimit {});
let event = build_stream_finished(reason, 200, 100);
match event.r#type {
@@ -67,3 +73,40 @@ fn test_build_stream_finished_other_reason() {
other => panic!("Expected Finished event, got {:?}", other),
}
}
#[test]
fn test_build_create_task_event() {
let event = super::stream::build_create_task("task-abc-123");
match event.r#type {
Some(api::response_event::Type::ClientActions(actions)) => {
assert_eq!(actions.actions.len(), 1);
match &actions.actions[0].action {
Some(api::client_action::Action::CreateTask(create)) => {
let task = create.task.as_ref().unwrap();
assert_eq!(task.id, "task-abc-123");
assert!(task.messages.is_empty());
assert!(task.description.is_empty());
assert!(task.dependencies.is_none());
}
other => panic!("Expected CreateTask action, got {:?}", other),
}
}
other => panic!("Expected ClientActions event, got {:?}", other),
}
}
#[test]
fn test_build_create_task_has_no_parent() {
let event = super::stream::build_create_task("root-task-id");
if let Some(api::response_event::Type::ClientActions(actions)) = event.r#type {
if let Some(api::client_action::Action::CreateTask(create)) = &actions.actions[0].action {
let task = create.task.as_ref().unwrap();
assert!(
task.dependencies.is_none(),
"Root task CreateTask must have no dependencies (no parent_id)"
);
}
}
}
app/src/ai/bedrock/test_fixtures/sample_project/Cargo.toml
+9
View File
@@ -0,0 +1,9 @@
[package]
name = "sample-project"
version = "0.1.0"
edition = "2021"
[dependencies]
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tokio = { version = "1", features = ["full"] }
app/src/ai/bedrock/test_fixtures/sample_project/README.md
+19
View File
@@ -0,0 +1,19 @@
# Sample Project
A simple Rust project used for testing the Warp AI Bedrock integration.
## Features
- Configuration from environment variables
- Basic math utilities
- File reading helpers
## Usage
```bash
cargo run
```
Set environment variables:
- `APP_NAME` - Application name (default: "sample-app")
- `PORT` - Server port (default: 8080)
- `DEBUG` - Enable debug mode (set to "1")
app/src/ai/bedrock/test_fixtures/sample_project/src/lib.rs
+24
View File
@@ -0,0 +1,24 @@
pub mod utils;
pub fn add(a: i32, b: i32) -> i32 {
a + b
}
pub fn multiply(a: i32, b: i32) -> i32 {
a * b
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_add() {
assert_eq!(add(2, 3), 5);
}
#[test]
fn test_multiply() {
assert_eq!(multiply(3, 4), 12);
}
}
app/src/ai/bedrock/test_fixtures/sample_project/src/main.rs
+30
View File
@@ -0,0 +1,30 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Serialize, Deserialize)]
struct Config {
name: String,
port: u16,
debug: bool,
}
impl Config {
fn from_env() -> Self {
Self {
name: std::env::var("APP_NAME").unwrap_or_else(|_| "sample-app".into()),
port: std::env::var("PORT")
.ok()
.and_then(|p| p.parse().ok())
.unwrap_or(8080),
debug: std::env::var("DEBUG").map(|v| v == "1").unwrap_or(false),
}
}
}
#[tokio::main]
async fn main() {
let config = Config::from_env();
println!("Starting {} on port {}", config.name, config.port);
if config.debug {
println!("Debug mode enabled");
}
}
app/src/ai/bedrock/test_fixtures/sample_project/src/utils.rs
+12
View File
@@ -0,0 +1,12 @@
use std::path::Path;
pub fn file_exists(path: &str) -> bool {
Path::new(path).exists()
}
pub fn read_lines(path: &str) -> Result<Vec<String>, std::io::Error> {
use std::io::BufRead;
let file = std::fs::File::open(path)?;
let reader = std::io::BufReader::new(file);
reader.lines().collect()
}