Significant progress. Performing cleanup now

This commit is contained in:
Ryan Ward
2026-05-11 14:31:02 -05:00
parent fe105e0369
commit 140da74f99
37 changed files with 843 additions and 1220 deletions
+9
View File
@@ -131,6 +131,13 @@ pub struct RequestParams {
pub parent_agent_id: Option<String>,
/// The display name for this agent (e.g. "Agent 1"), assigned by the orchestrator.
pub agent_name: Option<String>,
/// Full Bedrock conversation history for direct Bedrock calls.
/// When present, the Bedrock path uses this instead of extracting from task_context.
pub bedrock_message_history: Vec<crate::ai::bedrock::convert::ConversationMessage>,
/// Populated by the Bedrock path after building the message list.
/// Contains the full messages sent (old history + new input) so the controller
/// can store them back into the conversation for the next request cycle.
pub bedrock_messages_sent: std::sync::Arc<std::sync::Mutex<Vec<crate::ai::bedrock::convert::ConversationMessage>>>,
}
pub type Event = Result<warp_multi_agent_api::ResponseEvent, Arc<AIApiError>>;
@@ -317,6 +324,8 @@ impl RequestParams {
.map(|id| id.to_string()),
parent_agent_id: None,
agent_name: None,
bedrock_message_history: Vec::new(),
bedrock_messages_sent: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())),
}
}
}
+24 -4
View File
@@ -59,7 +59,7 @@ pub async fn generate_multi_agent_output(
api_keys.allow_use_of_warp_credits = params.allow_use_of_warp_credits_with_byok;
}
let request = api::Request {
let mut request = api::Request {
task_context: Some(api::request::TaskContext {
tasks: params.tasks,
}),
@@ -183,8 +183,22 @@ pub async fn generate_multi_agent_output(
logger.log_protobuf_input(&request);
}
let messages =
crate::ai::bedrock::convert_request::extract_messages_from_request(&request);
// Build message list from bedrock_message_history + new input messages.
// The history contains all prior messages. We extract only NEW messages
// from the current request input and append them.
let new_input_messages =
crate::ai::bedrock::convert_request::extract_new_input_messages(&request);
let mut messages = params.bedrock_message_history.clone();
if !new_input_messages.is_empty() {
log::info!(
"[bedrock] Appending {} new input messages to history of {}",
new_input_messages.len(),
messages.len()
);
messages.extend(new_input_messages);
}
let system_prompt =
crate::ai::bedrock::convert_request::extract_system_prompt(&request);
let tools = crate::ai::bedrock::convert_request::extract_tools(&request);
@@ -249,17 +263,23 @@ pub async fn generate_multi_agent_output(
&model_id,
&task_id,
needs_create_task,
messages,
messages.clone(),
system_prompt,
tools,
64000,
None,
true,
diagnostic_logger,
params.bedrock_messages_sent.clone(),
)
.await
{
Ok(stream) => {
// Store the input messages we sent so the controller can
// persist them. The stream will append the assistant response.
if let Ok(mut sent) = params.bedrock_messages_sent.lock() {
*sent = messages;
}
let output_stream = stream.take_until(cancellation_rx);
return Ok(Box::pin(output_stream));
}
+2
View File
@@ -40,6 +40,8 @@ fn request_params_with_ask_user_question_enabled(ask_user_question_enabled: bool
root_task_id: None,
parent_agent_id: None,
agent_name: None,
bedrock_message_history: Vec::new(),
bedrock_messages_sent: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())),
}
}
+20
View File
@@ -228,6 +228,12 @@ pub struct AIConversation {
/// event log. Used on restore to resume event delivery without
/// re-delivering already-processed events.
last_event_sequence: Option<i64>,
/// Accumulated message history for direct Bedrock conversations.
/// Contains the full ordered sequence of user messages, assistant responses,
/// tool calls, and tool results sent to/received from Bedrock across all
/// request cycles. This is the source of truth for what Bedrock sees.
bedrock_message_history: Vec<crate::ai::bedrock::convert::ConversationMessage>,
}
pub(crate) fn artifact_from_fork_proto(
@@ -278,6 +284,7 @@ impl AIConversation {
parent_conversation_id: None,
is_remote_child: false,
last_event_sequence: None,
bedrock_message_history: Vec::new(),
}
}
@@ -459,6 +466,7 @@ impl AIConversation {
parent_conversation_id,
is_remote_child: false,
last_event_sequence,
bedrock_message_history: Vec::new(),
})
}
@@ -466,6 +474,18 @@ impl AIConversation {
self.id
}
pub fn bedrock_message_history(&self) -> &[crate::ai::bedrock::convert::ConversationMessage] {
&self.bedrock_message_history
}
pub fn bedrock_message_history_mut(&mut self) -> &mut Vec<crate::ai::bedrock::convert::ConversationMessage> {
&mut self.bedrock_message_history
}
pub fn append_to_bedrock_history(&mut self, messages: Vec<crate::ai::bedrock::convert::ConversationMessage>) {
self.bedrock_message_history.extend(messages);
}
/// Assigns fresh exchange IDs to all exchanges in this conversation.
/// Used when forking conversations to avoid ID collisions with persisted blocks.
pub fn reassign_exchange_ids(&mut self) {
+1 -1
View File
@@ -290,7 +290,7 @@ static DEFAULT_TIPS: LazyLock<Vec<AgentTip>> = LazyLock::new(|| {
kind: AgentTipKind::Context,
},
AgentTip {
description: "Warpify a remote SSH session to enable Oz inside that environment.".to_string(),
description: "Galaxify a remote SSH session to enable the agent inside that environment.".to_string(),
link: Some("https://docs.warp.dev/terminal/warpify".to_string()),
binding_name: None,
action: None,
+3 -1
View File
@@ -1,4 +1,4 @@
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use anyhow::Result;
use aws_config::BehaviorVersion;
@@ -113,6 +113,7 @@ impl BedrockClient {
temperature: Option<f32>,
cross_region_inference: bool,
diagnostic_logger: Option<Arc<BedrockDiagnosticLogger>>,
messages_sent: Arc<Mutex<Vec<ConversationMessage>>>,
) -> Result<ResponseStream, BedrockError> {
let effective_model_id = if cross_region_inference {
apply_cross_region_prefix(model_id, &self.region)
@@ -192,6 +193,7 @@ impl BedrockClient {
task_id.to_string(),
needs_create_task,
diagnostic_logger,
messages_sent,
)))
}
+45 -7
View File
@@ -1,9 +1,9 @@
use std::collections::HashMap;
use aws_sdk_bedrockruntime::types::{
ContentBlock, ConversationRole, InferenceConfiguration, Message as BedrockMessage,
SystemContentBlock, Tool, ToolConfiguration, ToolInputSchema, ToolResultBlock,
ToolResultContentBlock, ToolResultStatus, ToolSpecification, ToolUseBlock,
CachePointBlock, CachePointType, ContentBlock, ConversationRole, InferenceConfiguration,
Message as BedrockMessage, SystemContentBlock, Tool, ToolConfiguration, ToolInputSchema,
ToolResultBlock, ToolResultContentBlock, ToolResultStatus, ToolSpecification, ToolUseBlock,
};
use aws_smithy_types::Document;
use serde_json::Value as JsonValue;
@@ -15,7 +15,7 @@ pub struct ConvertedRequest {
pub tool_config: Option<ToolConfiguration>,
}
#[derive(Clone)]
#[derive(Clone, Debug)]
pub struct ConversationMessage {
pub role: MessageRole,
pub content: MessageContent,
@@ -208,7 +208,30 @@ fn convert_messages(messages: Vec<ConversationMessage>) -> Vec<BedrockMessage> {
result.push(message);
}
coalesce_consecutive_roles(result)
let mut messages = coalesce_consecutive_roles(result);
// Add a cache point to the second-to-last message (the conversation prefix
// that is stable between requests). This allows Bedrock to cache all prior
// context and only process the latest message as new input tokens.
if messages.len() >= 2 {
let cache_idx = messages.len() - 2;
let msg = messages.remove(cache_idx);
let mut content = msg.content().to_vec();
content.push(ContentBlock::CachePoint(
CachePointBlock::builder()
.r#type(CachePointType::Default)
.build()
.expect("valid cache point"),
));
let cached_msg = BedrockMessage::builder()
.role(msg.role().clone())
.set_content(Some(content))
.build()
.expect("valid message with cache point");
messages.insert(cache_idx, cached_msg);
}
messages
}
fn coalesce_consecutive_roles(messages: Vec<BedrockMessage>) -> Vec<BedrockMessage> {
@@ -245,7 +268,15 @@ fn coalesce_consecutive_roles(messages: Vec<BedrockMessage>) -> Vec<BedrockMessa
fn convert_system_prompt(system_prompt: Option<String>) -> Vec<SystemContentBlock> {
match system_prompt {
Some(prompt) if !prompt.is_empty() => {
vec![SystemContentBlock::Text(prompt)]
vec![
SystemContentBlock::Text(prompt),
SystemContentBlock::CachePoint(
CachePointBlock::builder()
.r#type(CachePointType::Default)
.build()
.expect("valid cache point"),
),
]
}
_ => vec![],
}
@@ -277,7 +308,7 @@ fn build_tool_config(tools: Vec<ToolDefinition>) -> Option<ToolConfiguration> {
return None;
}
let tool_specs: Vec<Tool> = tools
let mut tool_specs: Vec<Tool> = tools
.into_iter()
.map(|tool| {
let input_schema_doc = json_to_document(tool.input_schema);
@@ -292,6 +323,13 @@ fn build_tool_config(tools: Vec<ToolDefinition>) -> Option<ToolConfiguration> {
})
.collect();
tool_specs.push(Tool::CachePoint(
CachePointBlock::builder()
.r#type(CachePointType::Default)
.build()
.expect("valid cache point"),
));
Some(
ToolConfiguration::builder()
.set_tools(Some(tool_specs))
+359 -61
View File
@@ -1,79 +1,366 @@
use warp_multi_agent_api as api;
use super::convert::{ConversationMessage, MessageContent, MessageRole, ToolDefinition};
use super::convert::{ConversationMessage, ContentPart, MessageContent, MessageRole, ToolDefinition};
/// Extract new input messages from the current request and convert them directly
/// to ConversationMessage format for the Bedrock message history.
/// This extracts UserQuery and ToolCallResult from request.input only.
pub fn extract_new_input_messages(request: &api::Request) -> Vec<ConversationMessage> {
let mut results = Vec::new();
let Some(input) = &request.input else {
return results;
};
let Some(input_type) = &input.r#type else {
return results;
};
match input_type {
api::request::input::Type::UserInputs(user_inputs) => {
let mut tool_results: Vec<ConversationMessage> = Vec::new();
for user_input in &user_inputs.inputs {
match &user_input.input {
Some(api::request::input::user_inputs::user_input::Input::ToolCallResult(
result,
)) => {
if !result.tool_call_id.is_empty() {
let content = extract_tool_result_content(result);
tool_results.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: result.tool_call_id.clone(),
content,
is_error: false,
},
});
}
}
Some(api::request::input::user_inputs::user_input::Input::UserQuery(
query,
)) => {
if !query.query.is_empty() {
results.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(query.query.clone()),
});
}
}
_ => {}
}
}
if !tool_results.is_empty() {
if tool_results.len() == 1 {
results.extend(tool_results);
} else {
let parts: Vec<ContentPart> = tool_results
.into_iter()
.map(|tr| match tr.content {
MessageContent::ToolResult { tool_use_id, content, is_error } => {
ContentPart::ToolResult { tool_use_id, content, is_error }
}
_ => unreachable!(),
})
.collect();
results.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::MultiPart(parts),
});
}
}
}
#[allow(deprecated)]
api::request::input::Type::UserQuery(query) => {
if !query.query.is_empty() {
results.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(query.query.clone()),
});
}
}
api::request::input::Type::InitProjectRules(_) => {
results.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(", "))
};
results.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) => {
results.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(format!(
"Create a new project: {}",
project.query,
)),
});
}
api::request::input::Type::CloneRepository(repo) => {
results.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::ResumeConversation(_) => {
results.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::CodeReview(_) => {
results.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(),
),
});
}
_ => {}
}
for msg in &results {
let desc = match &msg.content {
MessageContent::Text(t) => format!("Text({}chars)", t.len()),
MessageContent::ToolResult { tool_use_id, .. } => format!("ToolResult({})", tool_use_id),
MessageContent::MultiPart(parts) => format!("MultiPart({} parts)", parts.len()),
_ => "Other".to_string(),
};
log::info!("[bedrock] New input message: role={:?}, content={}", msg.role, desc);
}
results
}
/// For the Bedrock direct path: inject all input messages (user queries and tool call
/// results) into the task's messages so they persist in conversation history for future
/// requests. Without this, inputs are lost after the current request cycle because they
/// only exist in `request.input` and are never stored in `task_context.tasks[].messages`.
pub fn inject_input_messages_into_task(request: &mut api::Request) {
let input_messages: Vec<api::Message> = extract_input_messages(request);
if input_messages.is_empty() {
return;
}
log::info!(
"[bedrock] Injecting {} input messages into task history",
input_messages.len()
);
if let Some(task_context) = &mut request.task_context {
if let Some(task) = task_context.tasks.first_mut() {
task.messages.extend(input_messages);
} else {
// No task exists yet — create one to hold the messages
let task_id = uuid::Uuid::new_v4().to_string();
task_context.tasks.push(api::Task {
id: task_id,
messages: input_messages,
..Default::default()
});
}
} else {
let task_id = uuid::Uuid::new_v4().to_string();
request.task_context = Some(api::request::TaskContext {
tasks: vec![api::Task {
id: task_id,
messages: input_messages,
..Default::default()
}],
});
}
}
fn extract_input_messages(request: &api::Request) -> Vec<api::Message> {
let mut results = Vec::new();
let Some(input) = &request.input else {
return results;
};
let Some(input_type) = &input.r#type else {
return results;
};
let task_id = request
.task_context
.as_ref()
.and_then(|tc| tc.tasks.first())
.map(|t| t.id.clone())
.unwrap_or_default();
match input_type {
api::request::input::Type::UserInputs(user_inputs) => {
for user_input in &user_inputs.inputs {
match &user_input.input {
Some(api::request::input::user_inputs::user_input::Input::ToolCallResult(
result,
)) => {
if !result.tool_call_id.is_empty() {
let content = extract_tool_result_content(result);
results.push(api::Message {
id: uuid::Uuid::new_v4().to_string(),
task_id: task_id.clone(),
request_id: String::new(),
timestamp: None,
server_message_data: String::new(),
citations: vec![],
message: Some(api::message::Message::ToolCallResult(
api::message::ToolCallResult {
tool_call_id: result.tool_call_id.clone(),
context: None,
result: Some(
api::message::tool_call_result::Result::Server(
api::message::tool_call_result::ServerResult {
serialized_result: content,
},
),
),
},
)),
});
}
}
Some(api::request::input::user_inputs::user_input::Input::UserQuery(
query,
)) => {
if !query.query.is_empty() {
results.push(api::Message {
id: uuid::Uuid::new_v4().to_string(),
task_id: task_id.clone(),
request_id: String::new(),
timestamp: None,
server_message_data: String::new(),
citations: vec![],
message: Some(api::message::Message::UserQuery(
api::message::UserQuery {
query: query.query.clone(),
..Default::default()
},
)),
});
}
}
_ => {}
}
}
}
api::request::input::Type::UserQuery(query) => {
if !query.query.is_empty() {
results.push(api::Message {
id: uuid::Uuid::new_v4().to_string(),
task_id: task_id.clone(),
request_id: String::new(),
timestamp: None,
server_message_data: String::new(),
citations: vec![],
message: Some(api::message::Message::UserQuery(
api::message::UserQuery {
query: query.query.clone(),
..Default::default()
},
)),
});
}
}
api::request::input::Type::ToolCallResult(result) => {
if !result.tool_call_id.is_empty() {
let content = extract_tool_result_content(result);
results.push(api::Message {
id: uuid::Uuid::new_v4().to_string(),
task_id: task_id.clone(),
request_id: String::new(),
timestamp: None,
server_message_data: String::new(),
citations: vec![],
message: Some(api::message::Message::ToolCallResult(
api::message::ToolCallResult {
tool_call_id: result.tool_call_id.clone(),
context: None,
result: Some(
api::message::tool_call_result::Result::Server(
api::message::tool_call_result::ServerResult {
serialized_result: content,
},
),
),
},
)),
});
}
}
_ => {}
}
results
}
pub fn extract_messages_from_request(request: &api::Request) -> Vec<ConversationMessage> {
let mut messages = Vec::new();
if let Some(task_context) = &request.task_context {
log::info!(
"[bedrock-debug] extract_messages: {} tasks in task_context",
task_context.tasks.len()
);
for task in &task_context.tasks {
log::info!(
"[bedrock-debug] extract_messages: task '{}' has {} messages",
task.id,
task.messages.len()
);
for msg in &task.messages {
let msg_type = msg.message.as_ref().map(|m| match m {
api::message::Message::UserQuery(_) => "UserQuery",
api::message::Message::AgentOutput(_) => "AgentOutput",
api::message::Message::ToolCall(_) => "ToolCall",
api::message::Message::ToolCallResult(_) => "ToolCallResult",
api::message::Message::AgentReasoning(_) => "AgentReasoning",
_ => "Other",
}).unwrap_or("None");
log::info!(
"[bedrock-debug] extract_messages: msg id='{}' type={}",
msg.id,
msg_type
);
if let Some(converted) = convert_proto_message(msg) {
messages.push(converted);
}
}
}
} else {
log::warn!("[bedrock-debug] extract_messages: NO task_context in request!");
}
if let Some(input) = &request.input {
if let Some(input_type) = &input.r#type {
#[allow(deprecated)]
match input_type {
api::request::input::Type::UserInputs(user_inputs) => {
for user_input in &user_inputs.inputs {
if let Some(input_variant) = &user_input.input {
match input_variant {
api::request::input::user_inputs::user_input::Input::UserQuery(
query,
) => {
if !query.query.is_empty() {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(query.query.clone()),
});
}
}
api::request::input::user_inputs::user_input::Input::ToolCallResult(
result,
) => {
let content = extract_tool_result_content(result);
if !result.tool_call_id.is_empty() {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: result.tool_call_id.clone(),
content,
is_error: false,
},
});
}
}
_ => {}
}
}
}
}
api::request::input::Type::UserQuery(query) => {
if !query.query.is_empty() {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text(query.query.clone()),
});
}
}
api::request::input::Type::ToolCallResult(result) => {
let content = extract_tool_result_content(result);
if !result.tool_call_id.is_empty() {
messages.push(ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: result.tool_call_id.clone(),
content,
is_error: false,
},
});
}
}
// UserInputs (UserQuery + ToolCallResult) are already injected
// into task messages by inject_input_messages_into_task().
api::request::input::Type::UserInputs(_) => {}
api::request::input::Type::UserQuery(_) => {}
api::request::input::Type::ToolCallResult(_) => {}
api::request::input::Type::InitProjectRules(_) => {
messages.push(ConversationMessage {
role: MessageRole::User,
@@ -159,6 +446,14 @@ pub fn extract_messages_from_request(request: &api::Request) -> Vec<Conversation
ensure_starts_with_user_message(&mut messages);
ensure_tool_results_paired(&mut messages);
log::info!(
"[bedrock-debug] extract_messages: TOTAL {} messages to send to Bedrock (User={}, Assistant={}, ToolResult={}, ToolUse={})",
messages.len(),
messages.iter().filter(|m| m.role == MessageRole::User && matches!(&m.content, MessageContent::Text(_))).count(),
messages.iter().filter(|m| m.role == MessageRole::Assistant && matches!(&m.content, MessageContent::Text(_))).count(),
messages.iter().filter(|m| matches!(&m.content, MessageContent::ToolResult { .. })).count(),
messages.iter().filter(|m| matches!(&m.content, MessageContent::ToolUse { .. })).count(),
);
messages
}
@@ -331,11 +626,11 @@ fn tool_definition_for_name(name: &str) -> ToolDefinition {
},
"read_files" => ToolDefinition {
name: "read_files".to_string(),
description: "Read the contents of one or more files.".to_string(),
description: "Read the contents of one or more files. ALWAYS pass all files you need in a single call rather than making multiple separate calls.".to_string(),
input_schema: serde_json::json!({
"type": "object",
"properties": {
"files": { "type": "array", "items": { "type": "string" }, "description": "File paths to read" }
"files": { "type": "array", "items": { "type": "string" }, "description": "File paths to read. Include ALL files you need in one call for efficiency." }
},
"required": ["files"]
}),
@@ -353,11 +648,11 @@ fn tool_definition_for_name(name: &str) -> ToolDefinition {
},
"grep" => ToolDefinition {
name: "grep".to_string(),
description: "Search for patterns in files using grep.".to_string(),
description: "Search for patterns in files. Pass all search patterns in one call.".to_string(),
input_schema: serde_json::json!({
"type": "object",
"properties": {
"queries": { "type": "array", "items": { "type": "string" }, "description": "Search patterns" },
"queries": { "type": "array", "items": { "type": "string" }, "description": "Search patterns. Include ALL patterns you need in one call." },
"path": { "type": "string", "description": "Directory to search in" }
},
"required": ["queries"]
@@ -365,7 +660,7 @@ fn tool_definition_for_name(name: &str) -> ToolDefinition {
},
"file_glob" => ToolDefinition {
name: "file_glob".to_string(),
description: "Find files matching glob patterns.".to_string(),
description: "Find files matching glob patterns. Pass all patterns in one call.".to_string(),
input_schema: serde_json::json!({
"type": "object",
"properties": {
@@ -544,6 +839,9 @@ fn format_tool_call_result(result: &api::message::ToolCallResult) -> String {
_ => "Read files completed.".to_string(),
}
}
api::message::tool_call_result::Result::Server(server_result) => {
server_result.serialized_result.clone()
}
_ => "Tool completed successfully.".to_string(),
}
} else {
+10
View File
@@ -1,3 +1,5 @@
use std::sync::{Arc, Mutex};
use futures::StreamExt;
use serde_json::json;
use std::path::PathBuf;
@@ -553,6 +555,7 @@ impl AgentSimulation {
8192,
None,
false,
Arc::new(Mutex::new(Vec::new())),
)
.await
.expect("converse_stream should succeed");
@@ -1137,6 +1140,7 @@ async fn test_reasoning_model_produces_substantial_output() {
8192,
None,
false,
Arc::new(Mutex::new(Vec::new())),
)
.await
.expect("converse_stream should succeed");
@@ -1257,6 +1261,7 @@ async fn test_event_sequence_matches_controller_expectations() {
100,
None,
false,
Arc::new(Mutex::new(Vec::new())),
)
.await
.expect("should connect");
@@ -1369,6 +1374,7 @@ async fn test_followup_turn_does_not_send_create_task() {
100,
None,
false,
Arc::new(Mutex::new(Vec::new())),
)
.await
.expect("should connect");
@@ -1446,6 +1452,7 @@ async fn run_slash_command_test(
4096,
None,
true,
Arc::new(Mutex::new(Vec::new())),
)
.await
.expect("stream should connect");
@@ -1671,6 +1678,7 @@ async fn test_slash_resume_conversation() {
256,
None,
true,
Arc::new(Mutex::new(Vec::new())),
)
.await
.expect("stream should connect");
@@ -1780,6 +1788,7 @@ async fn test_empty_messages_safety_check() {
100,
None,
true,
Arc::new(Mutex::new(Vec::new())),
)
.await
.expect("safety fallback message should work");
@@ -1929,6 +1938,7 @@ async fn test_full_proto_round_trip_with_tool_history() {
1024,
None,
true,
Arc::new(Mutex::new(Vec::new())),
)
.await;
+3
View File
@@ -1,3 +1,5 @@
use std::sync::{Arc, Mutex};
use futures::StreamExt;
use serde_json::json;
@@ -62,6 +64,7 @@ async fn collect_stream_output(
8192,
None,
false,
Arc::new(Mutex::new(Vec::new())),
)
.await
.expect("converse_stream should succeed");
+59 -1
View File
@@ -1,4 +1,4 @@
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use aws_sdk_bedrockruntime::operation::converse_stream::ConverseStreamOutput;
use aws_sdk_bedrockruntime::types::{
@@ -13,6 +13,7 @@ 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, MessageContent, MessageRole, ConversationMessage};
use super::diagnostic::BedrockDiagnosticLogger;
pub fn bedrock_stream_to_response_events(
@@ -20,6 +21,7 @@ pub fn bedrock_stream_to_response_events(
task_id: String,
needs_create_task: bool,
diagnostic_logger: Option<Arc<BedrockDiagnosticLogger>>,
messages_sent: Arc<Mutex<Vec<ConversationMessage>>>,
) -> BoxStream<'static, Event> {
let request_id = Uuid::new_v4().to_string();
let conversation_id = Uuid::new_v4().to_string();
@@ -55,6 +57,10 @@ pub fn bedrock_stream_to_response_events(
let mut output_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();
let mut event_count: u32 = 0;
loop {
match output.stream.recv().await {
@@ -109,6 +115,7 @@ pub fn bedrock_stream_to_response_events(
match d {
ContentBlockDelta::Text(text) => {
log::info!("[bedrock-debug] Event #{event_count}: TextDelta ({} chars): {:?}", text.len(), &text[..text.len().min(80)]);
history_text.push_str(text);
if text_flushed {
let msg_id = current_text_message_id.as_ref().unwrap();
let append = build_append_text(
@@ -149,6 +156,14 @@ pub fn bedrock_stream_to_response_events(
log::info!("[bedrock-debug] Event #{event_count}: ContentBlockStop (tool_use_id={:?})", if current_tool_use_id.is_empty() { "none" } else { &current_tool_use_id });
if !current_tool_use_id.is_empty() {
log::debug!("[bedrock] Tool call complete: {} ({})", current_tool_name, current_tool_use_id);
// Track for bedrock_message_history
let input_json: serde_json::Value = serde_json::from_str(&current_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={}",
@@ -242,6 +257,49 @@ pub fn bedrock_stream_to_response_events(
}
log::info!("[bedrock] Stream finished: input_tokens={input_tokens}, output_tokens={output_tokens}");
// 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);
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",
@@ -829,6 +829,9 @@ impl BlocklistAIActionExecutor {
}
fn should_autoexecute(&self, input: ExecuteActionInput, ctx: &mut ModelContext<Self>) -> bool {
if cfg!(feature = "bedrock_smoke_test") {
return true;
}
match input.action.action {
AIAgentActionType::RequestCommandOutput { .. }
| AIAgentActionType::WriteToLongRunningShellCommand { .. }
-21
View File
@@ -48,27 +48,6 @@ pub static ENTER_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE: LazyLock<Keystroke> = La
}
});
pub static ENTER_CLOUD_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE: LazyLock<Keystroke> =
LazyLock::new(|| {
cfg_if::cfg_if! {
if #[cfg(target_os = "macos")] {
Keystroke {
cmd: true,
alt: true,
key: "enter".to_owned(),
..Default::default()
}
} else {
Keystroke {
ctrl: true,
alt: true,
key: "enter".to_owned(),
..Default::default()
}
}
}
});
pub fn agent_view_bg_fill(app: &AppContext) -> Fill {
let appearance = Appearance::as_ref(app);
appearance.theme().surface_overlay_1()
@@ -13,7 +13,6 @@ use galaxyui::{
AppContext, Element, SingletonEntity,
};
use crate::ai::blocklist::agent_view::ENTER_CLOUD_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE;
use crate::{
ai::blocklist::agent_view::ENTER_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE,
cmd_or_ctrl_shift,
@@ -193,16 +192,9 @@ pub fn render_agent_shortcuts_view(
app,
));
// Use cloud keystroke (cmd+opt+enter) for cloud mode, regular keystroke (cmd+enter) otherwise.
let new_conversation_keystroke = if context.is_cloud_agent {
ENTER_CLOUD_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE.clone()
} else {
ENTER_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE.clone()
};
shortcuts.push(render_shortcut(
ShortcutProps {
keystroke: new_conversation_keystroke.clone(),
keystroke: ENTER_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE.clone(),
text: "start a new conversation".into(),
..Default::default()
},
@@ -25,7 +25,6 @@ use crate::{
agent_view::{
agent_view_bg_color, AgentViewController, AgentViewEntryOrigin,
ENTER_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE,
ENTER_CLOUD_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE,
},
history_model::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel},
},
@@ -723,21 +722,6 @@ fn render_body(props: ZeroStateBodyProps<'_>, app: &AppContext) -> Vec<Box<dyn E
)]),
app,
),
render_standard_message(
Message::new(vec![MessageItem::clickable(
vec![
MessageItem::keystroke(
ENTER_CLOUD_AGENT_VIEW_NEW_CONVERSATION_KEYSTROKE.clone(),
),
MessageItem::text("start a new cloud agent conversation"),
],
|ctx| {
ctx.dispatch_typed_action(TerminalAction::EnterCloudAgentView);
},
state_handles.start_cloud_conversation.clone(),
)]),
app,
),
render_standard_message(
Message::new(vec![MessageItem::clickable(
vec![
+23
View File
@@ -1904,6 +1904,7 @@ impl BlocklistAIController {
active_tasks,
parent_agent_id,
agent_name,
bedrock_history,
) = {
let Some(conversation) = history_model
.as_ref(ctx)
@@ -1926,6 +1927,7 @@ impl BlocklistAIController {
active_tasks,
conversation.parent_agent_id().map(str::to_string),
conversation.agent_name().map(str::to_string),
conversation.bedrock_message_history().to_vec(),
)
};
@@ -1992,6 +1994,7 @@ impl BlocklistAIController {
);
request_params.parent_agent_id = parent_agent_id;
request_params.agent_name = agent_name;
request_params.bedrock_message_history = bedrock_history;
let server_conversation_token_for_identifiers =
conversation_data.server_conversation_token.clone();
@@ -2299,6 +2302,26 @@ impl BlocklistAIController {
did_input_contain_user_query,
ctx,
);
// After the stream finishes, persist the full message
// history (input + assistant response) from the Arc back
// into the conversation for the next request cycle.
let messages_sent_arc = response_stream.as_ref(ctx).bedrock_messages_sent().clone();
let new_history = messages_sent_arc.lock().ok().and_then(|sent| {
if sent.is_empty() { None } else { Some(sent.clone()) }
});
if let Some(new_history) = new_history {
let history_model = BlocklistAIHistoryModel::handle(ctx);
history_model.update(ctx, |history_model, _| {
if let Some(conversation) = history_model.conversation_mut(&conversation_id) {
*conversation.bedrock_message_history_mut() = new_history;
log::info!(
"[bedrock] Updated conversation bedrock history: {} messages",
conversation.bedrock_message_history().len()
);
}
});
}
}
warp_multi_agent_api::response_event::Type::ClientActions(actions) => {
let client_actions = actions.actions;
@@ -155,6 +155,10 @@ impl ResponseStream {
&self.id
}
pub fn bedrock_messages_sent(&self) -> &std::sync::Arc<std::sync::Mutex<Vec<crate::ai::bedrock::convert::ConversationMessage>>> {
&self.params.bedrock_messages_sent
}
/// Returns true if we should attempt to resume the conversation after the stream finishes.
pub fn should_resume_conversation_after_stream_finished(&self) -> bool {
self.should_resume_conversation_after_stream_finished