Fix Rig tool continuation and plan creation
This commit is contained in:
@@ -79,6 +79,16 @@ fn supported_tools_expose_local_subagents_without_hosted_orchestration_tools() {
|
||||
assert!(!supported_tools.contains(&api::ToolType::StartAgentV2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supported_tools_include_plan_document_capabilities() {
|
||||
let params = request_params_with_ask_user_question_enabled(false);
|
||||
let supported_tools = get_supported_tools(¶ms);
|
||||
|
||||
assert!(supported_tools.contains(&api::ToolType::ReadDocuments));
|
||||
assert!(supported_tools.contains(&api::ToolType::CreateDocuments));
|
||||
assert!(supported_tools.contains(&api::ToolType::EditDocuments));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supported_tools_omit_subagents_when_orchestration_is_disabled() {
|
||||
let params = request_params_with_ask_user_question_enabled(false);
|
||||
|
||||
@@ -1895,7 +1895,7 @@ pub fn default_tool_definitions() -> Vec<ToolDefinition> {
|
||||
},
|
||||
ToolDefinition {
|
||||
name: "create_plan".to_string(),
|
||||
description: "Create a new plan document in Galaxy Drive's Plans folder. Plans are rich-text documents for tracking tasks, architecture decisions, and project notes.".to_string(),
|
||||
description: "Create a new plan document in Galaxy Drive's Plans folder. When the user asks to create a plan for review, use this tool after completing the necessary research instead of only returning plan prose. Plans are rich-text documents for tracking tasks, architecture decisions, and project notes.".to_string(),
|
||||
input_schema: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
|
||||
@@ -147,6 +147,11 @@ impl RuntimeResponseTranslator {
|
||||
events
|
||||
}
|
||||
|
||||
pub(crate) fn begin_followup_turn(&mut self) {
|
||||
self.text_message_id = None;
|
||||
self.reasoning_message_id = None;
|
||||
}
|
||||
|
||||
fn initialize(&mut self, events: &mut Vec<ResponseEvent>) {
|
||||
if self.initialized {
|
||||
return;
|
||||
|
||||
+260
-146
@@ -4,8 +4,8 @@ use std::sync::Arc;
|
||||
use futures::channel::oneshot;
|
||||
use futures::{FutureExt, StreamExt};
|
||||
use galaxy_agent_core::{
|
||||
turn_control, AgentError, AgentEvent, AgentRuntime, MessageContent, MessageRole, ToolCall,
|
||||
ToolCallDecision, ToolEvent, ToolPolicy, ToolResult, TurnCommand,
|
||||
turn_control, AgentError, AgentEvent, AgentRuntime, MessageContent, MessageRole, StopReason,
|
||||
ToolCall, ToolCallDecision, ToolEvent, ToolPolicy, ToolResult, TurnCommand, Usage,
|
||||
};
|
||||
use galaxy_agent_rig::{
|
||||
AnthropicRuntime, AnthropicRuntimeConfig, ChatGPTSubscriptionRuntime,
|
||||
@@ -30,6 +30,10 @@ use crate::ai::runtime::{RuntimeResponseConfig, RuntimeResponseTranslator};
|
||||
use crate::server::server_api::AIApiError;
|
||||
use crate::settings::OpenAIProviderKind;
|
||||
|
||||
const MAX_INLINE_TOOL_CONTINUATIONS: usize = 3;
|
||||
const INLINE_TOOL_LOOP_MESSAGE: &str =
|
||||
"I couldn't continue because the model repeatedly searched prior tool history without making progress. Please retry with a more specific instruction.";
|
||||
|
||||
pub(crate) fn rig_openai_response_stream(
|
||||
config: OpenAIClientConfig,
|
||||
params: RequestParams,
|
||||
@@ -190,30 +194,8 @@ where
|
||||
let model_id = turn_request.model.as_str().to_string();
|
||||
let tool_policy = ToolPolicy::new(&turn_request.tools);
|
||||
let stream = async_stream::stream! {
|
||||
let (control_sender, control) = turn_control();
|
||||
let start_future = runtime.start_turn(turn_request, control).fuse();
|
||||
let cancel_future = cancellation_rx.fuse();
|
||||
futures::pin_mut!(start_future, cancel_future);
|
||||
|
||||
let mut agent_events = futures::select_biased! {
|
||||
_ = cancel_future => {
|
||||
let _ = control_sender.try_send(TurnCommand::Cancel);
|
||||
match start_future.await {
|
||||
Ok(stream) => stream,
|
||||
Err(error) => {
|
||||
yield Err(agent_error(error, stream_type));
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
result = start_future => match result {
|
||||
Ok(stream) => stream,
|
||||
Err(error) => {
|
||||
yield Err(agent_error(error, stream_type));
|
||||
return;
|
||||
}
|
||||
},
|
||||
};
|
||||
futures::pin_mut!(cancel_future);
|
||||
|
||||
let conversation_id = conversation_id.unwrap_or_else(|| Uuid::new_v4().to_string());
|
||||
let mut translator = RuntimeResponseTranslator::new(RuntimeResponseConfig {
|
||||
@@ -226,58 +208,185 @@ where
|
||||
capabilities: runtime_capabilities,
|
||||
empty_output_message: None,
|
||||
});
|
||||
let mut full_text = String::new();
|
||||
let mut full_reasoning = String::new();
|
||||
let mut reasoning_signature = None;
|
||||
let mut proposed_tools = Vec::new();
|
||||
let mut assistant_history_index = None;
|
||||
let mut turn_request = turn_request;
|
||||
let mut cumulative_usage = Usage::default();
|
||||
let mut inline_continuation_count = 0;
|
||||
|
||||
loop {
|
||||
let next_event = agent_events.next().fuse();
|
||||
futures::pin_mut!(next_event);
|
||||
futures::select_biased! {
|
||||
'provider_turns: loop {
|
||||
let (control_sender, control) = turn_control();
|
||||
let start_future = runtime.start_turn(turn_request.clone(), control).fuse();
|
||||
futures::pin_mut!(start_future);
|
||||
|
||||
let mut agent_events = futures::select_biased! {
|
||||
_ = cancel_future => {
|
||||
let _ = control_sender.try_send(TurnCommand::Cancel);
|
||||
}
|
||||
event = next_event => {
|
||||
let Some(event) = event else {
|
||||
yield Err(Arc::new(AIApiError::UnexpectedEof));
|
||||
return;
|
||||
};
|
||||
let event = match event {
|
||||
Ok(event) => event,
|
||||
match start_future.await {
|
||||
Ok(stream) => stream,
|
||||
Err(error) => {
|
||||
yield Err(agent_error(error, stream_type));
|
||||
return;
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
result = start_future => match result {
|
||||
Ok(stream) => stream,
|
||||
Err(error) => {
|
||||
yield Err(agent_error(error, stream_type));
|
||||
return;
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
match event {
|
||||
AgentEvent::Tool {
|
||||
event: ToolEvent::Proposed { call },
|
||||
} => {
|
||||
proposed_tools.push(call.clone());
|
||||
sync_assistant_turn(
|
||||
&messages_sent,
|
||||
&full_reasoning,
|
||||
reasoning_signature.as_deref(),
|
||||
&full_text,
|
||||
&proposed_tools,
|
||||
&mut assistant_history_index,
|
||||
);
|
||||
let history = messages_sent
|
||||
.lock()
|
||||
.map(|sent| sent.clone())
|
||||
.unwrap_or_default();
|
||||
match tool_policy.decide(&call, &history, &tool_result_archive) {
|
||||
ToolCallDecision::Execute => {
|
||||
match build_tool_proposed(
|
||||
&task_id,
|
||||
&call,
|
||||
&skill_path_origin,
|
||||
&mcp_tool_aliases,
|
||||
let mut full_text = String::new();
|
||||
let mut full_reasoning = String::new();
|
||||
let mut reasoning_signature = None;
|
||||
let mut proposed_tools = Vec::new();
|
||||
let mut assistant_history_index = None;
|
||||
let mut handled_inline_tool = false;
|
||||
let mut proposed_client_tool = false;
|
||||
|
||||
loop {
|
||||
let next_event = agent_events.next().fuse();
|
||||
futures::pin_mut!(next_event);
|
||||
futures::select_biased! {
|
||||
_ = cancel_future => {
|
||||
let _ = control_sender.try_send(TurnCommand::Cancel);
|
||||
}
|
||||
event = next_event => {
|
||||
let Some(event) = event else {
|
||||
yield Err(Arc::new(AIApiError::UnexpectedEof));
|
||||
return;
|
||||
};
|
||||
let event = match event {
|
||||
Ok(event) => event,
|
||||
Err(error) => {
|
||||
yield Err(agent_error(error, stream_type));
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
match event {
|
||||
AgentEvent::Tool {
|
||||
event: ToolEvent::Proposed { call },
|
||||
} => {
|
||||
proposed_tools.push(call.clone());
|
||||
sync_assistant_turn(
|
||||
&messages_sent,
|
||||
&full_reasoning,
|
||||
reasoning_signature.as_deref(),
|
||||
&full_text,
|
||||
&proposed_tools,
|
||||
&mut assistant_history_index,
|
||||
);
|
||||
let history = messages_sent
|
||||
.lock()
|
||||
.map(|sent| sent.clone())
|
||||
.unwrap_or_default();
|
||||
match tool_policy.decide(&call, &history, &tool_result_archive) {
|
||||
ToolCallDecision::Execute => {
|
||||
proposed_client_tool = true;
|
||||
match build_tool_proposed(
|
||||
&task_id,
|
||||
&call,
|
||||
&skill_path_origin,
|
||||
&mcp_tool_aliases,
|
||||
) {
|
||||
Ok(action) => yield Ok(StreamEvent::ToolProposed(action)),
|
||||
Err(message) => {
|
||||
yield Err(agent_error(AgentError::new(
|
||||
galaxy_agent_core::AgentErrorKind::Protocol,
|
||||
message,
|
||||
), stream_type));
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
ToolCallDecision::Inline(result) => {
|
||||
handled_inline_tool = true;
|
||||
append_tool_result(&messages_sent, result);
|
||||
}
|
||||
ToolCallDecision::Reject(result) => {
|
||||
log::warn!(
|
||||
"Rig model called unavailable tool '{}' (id={})",
|
||||
call.name,
|
||||
call.id
|
||||
);
|
||||
let error_display = format!(
|
||||
"Failed tool call: `{}`\n\n{}",
|
||||
call.name, result.content
|
||||
);
|
||||
append_tool_result(&messages_sent, result);
|
||||
let message_id = Uuid::new_v4().to_string();
|
||||
yield Ok(StreamEvent::Response(build_add_agent_output_message(
|
||||
&task_id,
|
||||
&message_id,
|
||||
&error_display,
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
AgentEvent::UsageUpdated { usage } => {
|
||||
accumulate_usage(&mut cumulative_usage, &usage);
|
||||
let response_events = match translator.translate(
|
||||
AgentEvent::UsageUpdated {
|
||||
usage: cumulative_usage.clone(),
|
||||
},
|
||||
) {
|
||||
Ok(response_events) => response_events,
|
||||
Err(message) => {
|
||||
yield Err(agent_error(AgentError::new(
|
||||
galaxy_agent_core::AgentErrorKind::Protocol,
|
||||
message,
|
||||
), stream_type));
|
||||
return;
|
||||
}
|
||||
};
|
||||
for response_event in response_events {
|
||||
yield Ok(StreamEvent::Response(response_event));
|
||||
}
|
||||
}
|
||||
AgentEvent::TurnStopped { mut reason } => {
|
||||
sync_assistant_turn(
|
||||
&messages_sent,
|
||||
&full_reasoning,
|
||||
reasoning_signature.as_deref(),
|
||||
&full_text,
|
||||
&proposed_tools,
|
||||
&mut assistant_history_index,
|
||||
);
|
||||
if reason == StopReason::Completed
|
||||
&& handled_inline_tool
|
||||
&& !proposed_client_tool
|
||||
{
|
||||
if inline_continuation_count < MAX_INLINE_TOOL_CONTINUATIONS {
|
||||
inline_continuation_count += 1;
|
||||
turn_request.messages = match copy_messages(&messages_sent) {
|
||||
Ok(messages) => messages,
|
||||
Err(()) => {
|
||||
yield Err(agent_error(AgentError::new(
|
||||
galaxy_agent_core::AgentErrorKind::Protocol,
|
||||
"could not access Rig conversation history for inline tool continuation",
|
||||
), stream_type));
|
||||
return;
|
||||
}
|
||||
};
|
||||
translator.begin_followup_turn();
|
||||
log::info!(
|
||||
"Continuing Rig provider turn after inline tool result ({inline_continuation_count}/{MAX_INLINE_TOOL_CONTINUATIONS})"
|
||||
);
|
||||
continue 'provider_turns;
|
||||
}
|
||||
|
||||
log::warn!(
|
||||
"Rig provider exceeded {MAX_INLINE_TOOL_CONTINUATIONS} inline tool continuations"
|
||||
);
|
||||
append_assistant_text(&messages_sent, INLINE_TOOL_LOOP_MESSAGE);
|
||||
let response_events = match translator.translate(
|
||||
AgentEvent::RuntimeNotice {
|
||||
message: INLINE_TOOL_LOOP_MESSAGE.to_string(),
|
||||
},
|
||||
) {
|
||||
Ok(action) => yield Ok(StreamEvent::ToolProposed(action)),
|
||||
Ok(response_events) => response_events,
|
||||
Err(message) => {
|
||||
yield Err(agent_error(AgentError::new(
|
||||
galaxy_agent_core::AgentErrorKind::Protocol,
|
||||
@@ -285,90 +394,63 @@ where
|
||||
), stream_type));
|
||||
return;
|
||||
}
|
||||
};
|
||||
for response_event in response_events {
|
||||
yield Ok(StreamEvent::Response(response_event));
|
||||
}
|
||||
reason = StopReason::ToolLoopLimit;
|
||||
}
|
||||
ToolCallDecision::Inline(result) => {
|
||||
append_tool_result(&messages_sent, result);
|
||||
}
|
||||
ToolCallDecision::Reject(result) => {
|
||||
log::warn!(
|
||||
"Rig model called unavailable tool '{}' (id={})",
|
||||
call.name,
|
||||
call.id
|
||||
);
|
||||
let error_display = format!(
|
||||
"Failed tool call: `{}`\n\n{}",
|
||||
call.name, result.content
|
||||
);
|
||||
append_tool_result(&messages_sent, result);
|
||||
let message_id = Uuid::new_v4().to_string();
|
||||
yield Ok(StreamEvent::Response(build_add_agent_output_message(
|
||||
&task_id,
|
||||
&message_id,
|
||||
&error_display,
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
AgentEvent::TurnStopped { reason } => {
|
||||
sync_assistant_turn(
|
||||
&messages_sent,
|
||||
&full_reasoning,
|
||||
reasoning_signature.as_deref(),
|
||||
&full_text,
|
||||
&proposed_tools,
|
||||
&mut assistant_history_index,
|
||||
);
|
||||
let response_events = match translator
|
||||
.translate(AgentEvent::TurnStopped { reason })
|
||||
{
|
||||
Ok(response_events) => response_events,
|
||||
Err(message) => {
|
||||
yield Err(agent_error(AgentError::new(
|
||||
galaxy_agent_core::AgentErrorKind::Protocol,
|
||||
message,
|
||||
), stream_type));
|
||||
return;
|
||||
}
|
||||
};
|
||||
for response_event in response_events {
|
||||
yield Ok(StreamEvent::Response(response_event));
|
||||
}
|
||||
return;
|
||||
}
|
||||
event => {
|
||||
match &event {
|
||||
AgentEvent::TextDelta { text } => full_text.push_str(text),
|
||||
AgentEvent::ReasoningDelta { text } => {
|
||||
full_reasoning.push_str(text);
|
||||
}
|
||||
AgentEvent::ReasoningCompleted { text, signature } => {
|
||||
if !text.is_empty() {
|
||||
full_reasoning.clone_from(text);
|
||||
let response_events = match translator
|
||||
.translate(AgentEvent::TurnStopped { reason })
|
||||
{
|
||||
Ok(response_events) => response_events,
|
||||
Err(message) => {
|
||||
yield Err(agent_error(AgentError::new(
|
||||
galaxy_agent_core::AgentErrorKind::Protocol,
|
||||
message,
|
||||
), stream_type));
|
||||
return;
|
||||
}
|
||||
reasoning_signature.clone_from(signature);
|
||||
};
|
||||
for response_event in response_events {
|
||||
yield Ok(StreamEvent::Response(response_event));
|
||||
}
|
||||
AgentEvent::TurnStarted { .. }
|
||||
| AgentEvent::Tool { .. }
|
||||
| AgentEvent::UsageUpdated { .. }
|
||||
| AgentEvent::RuntimeActivityUpdated { .. }
|
||||
| AgentEvent::ContextUsageUpdated { .. }
|
||||
| AgentEvent::UserInputAccepted { .. }
|
||||
| AgentEvent::RuntimeNotice { .. }
|
||||
| AgentEvent::TurnStopped { .. } => {}
|
||||
return;
|
||||
}
|
||||
let response_events = match translator.translate(event) {
|
||||
Ok(response_events) => response_events,
|
||||
Err(message) => {
|
||||
yield Err(agent_error(AgentError::new(
|
||||
galaxy_agent_core::AgentErrorKind::Protocol,
|
||||
message,
|
||||
), stream_type));
|
||||
return;
|
||||
event => {
|
||||
match &event {
|
||||
AgentEvent::TextDelta { text } => full_text.push_str(text),
|
||||
AgentEvent::ReasoningDelta { text } => {
|
||||
full_reasoning.push_str(text);
|
||||
}
|
||||
AgentEvent::ReasoningCompleted { text, signature } => {
|
||||
if !text.is_empty() {
|
||||
full_reasoning.clone_from(text);
|
||||
}
|
||||
reasoning_signature.clone_from(signature);
|
||||
}
|
||||
AgentEvent::TurnStarted { .. }
|
||||
| AgentEvent::Tool { .. }
|
||||
| AgentEvent::UsageUpdated { .. }
|
||||
| AgentEvent::RuntimeActivityUpdated { .. }
|
||||
| AgentEvent::ContextUsageUpdated { .. }
|
||||
| AgentEvent::UserInputAccepted { .. }
|
||||
| AgentEvent::RuntimeNotice { .. }
|
||||
| AgentEvent::TurnStopped { .. } => {}
|
||||
}
|
||||
let response_events = match translator.translate(event) {
|
||||
Ok(response_events) => response_events,
|
||||
Err(message) => {
|
||||
yield Err(agent_error(AgentError::new(
|
||||
galaxy_agent_core::AgentErrorKind::Protocol,
|
||||
message,
|
||||
), stream_type));
|
||||
return;
|
||||
}
|
||||
};
|
||||
for response_event in response_events {
|
||||
yield Ok(StreamEvent::Response(response_event));
|
||||
}
|
||||
};
|
||||
for response_event in response_events {
|
||||
yield Ok(StreamEvent::Response(response_event));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -390,6 +472,15 @@ fn store_messages_sent(
|
||||
*sent = messages.to_vec();
|
||||
}
|
||||
|
||||
fn copy_messages(
|
||||
messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>,
|
||||
) -> Result<Vec<ConversationMessage>, ()> {
|
||||
messages_sent
|
||||
.lock()
|
||||
.map(|sent| sent.clone())
|
||||
.map_err(|_| ())
|
||||
}
|
||||
|
||||
fn append_tool_result(
|
||||
messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>,
|
||||
result: ToolResult,
|
||||
@@ -408,6 +499,29 @@ fn append_tool_result(
|
||||
}
|
||||
}
|
||||
|
||||
fn append_assistant_text(
|
||||
messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>,
|
||||
text: &str,
|
||||
) {
|
||||
if let Ok(mut sent) = messages_sent.lock() {
|
||||
sent.push(ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content: MessageContent::Text(text.to_string()),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn accumulate_usage(total: &mut Usage, usage: &Usage) {
|
||||
total.input_tokens = total.input_tokens.saturating_add(usage.input_tokens);
|
||||
total.output_tokens = total.output_tokens.saturating_add(usage.output_tokens);
|
||||
total.cached_input_tokens = total
|
||||
.cached_input_tokens
|
||||
.saturating_add(usage.cached_input_tokens);
|
||||
total.cache_creation_input_tokens = total
|
||||
.cache_creation_input_tokens
|
||||
.saturating_add(usage.cache_creation_input_tokens);
|
||||
}
|
||||
|
||||
fn sync_assistant_turn(
|
||||
messages_sent: &std::sync::Arc<std::sync::Mutex<Vec<ConversationMessage>>>,
|
||||
reasoning_text: &str,
|
||||
|
||||
@@ -798,6 +798,11 @@ fn build_system_prompt(
|
||||
.join(", "),
|
||||
);
|
||||
prompt.push_str(".\nNever invent tool names or parameters. Check command exit codes and tool error results before claiming success.\n");
|
||||
if tools.iter().any(|tool| tool.name == "create_plan") {
|
||||
prompt.push_str(
|
||||
"Plan document creation is available through `create_plan`. When the user asks you to create a plan for review, research first as needed, then call `create_plan`; do not merely return the plan as prose or claim that no plan-creation tool is available. If the user asks to review the plan before implementation, creating the document and presenting it for review is the requested outcome; do not implement it until they approve.\n",
|
||||
);
|
||||
}
|
||||
}
|
||||
prompt
|
||||
}
|
||||
|
||||
@@ -124,6 +124,31 @@ fn builds_a_rig_turn_directly_from_galaxy_request_state() {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normal_turn_advertises_plan_creation_and_corrects_false_unavailability_claims() {
|
||||
let mut params = RequestParams::new_for_test();
|
||||
params.planning_enabled = true;
|
||||
params.input = vec![user_query("Please create a plan, and let's review.")];
|
||||
|
||||
let prepared = prepare_rig_turn(
|
||||
&config(),
|
||||
params,
|
||||
vec![ToolType::CreateDocuments],
|
||||
Vec::new(),
|
||||
);
|
||||
let prompt = prepared.request.system_prompt.expect("system prompt");
|
||||
|
||||
assert!(prepared
|
||||
.request
|
||||
.tools
|
||||
.iter()
|
||||
.any(|tool| tool.name == "create_plan"));
|
||||
assert!(prompt.contains("Plan document creation is available through `create_plan`"));
|
||||
assert!(prompt.contains("do not merely return the plan as prose"));
|
||||
assert!(prompt.contains("no plan-creation tool is available"));
|
||||
assert!(prompt.contains("do not implement it until they approve"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_tools_turn_flattens_historical_tool_protocol_messages() {
|
||||
let mut params = RequestParams::new_for_test();
|
||||
|
||||
@@ -1,12 +1,23 @@
|
||||
use std::collections::HashMap;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use ai::skills::SkillPathOrigin;
|
||||
use async_trait::async_trait;
|
||||
use futures::channel::oneshot;
|
||||
use futures::StreamExt;
|
||||
use galaxy_agent_core::{
|
||||
ContentPart, MessageContent, MessageRole, ToolCall, ToolResult, ToolResultStatus,
|
||||
AgentError, AgentErrorKind, AgentEvent, AgentEventStream, AgentRuntime, ContentPart,
|
||||
ConversationMessage, MessageContent, MessageRole, RuntimeCapabilities, RuntimeDescriptor,
|
||||
RuntimeKind, StopReason, ToolCall, ToolDefinition, ToolEvent, ToolResult, ToolResultStatus,
|
||||
TurnControl, TurnRequest, Usage, RECALL_TOOL_HISTORY_NAME,
|
||||
};
|
||||
use warp_multi_agent_api::{client_action, message, response_event};
|
||||
|
||||
use super::{append_tool_result, build_tool_proposed, sync_assistant_turn};
|
||||
use super::{
|
||||
append_tool_result, build_tool_proposed, rig_response_stream, sync_assistant_turn,
|
||||
PreparedRigTurn, INLINE_TOOL_LOOP_MESSAGE, MAX_INLINE_TOOL_CONTINUATIONS,
|
||||
};
|
||||
use crate::ai::agent::api::StreamEvent;
|
||||
|
||||
#[test]
|
||||
fn tool_proposal_matches_the_domain_permission_contract() {
|
||||
@@ -194,3 +205,264 @@ fn rejected_tool_result_is_paired_with_the_assistant_call_in_history() {
|
||||
} if tool_use_id == "call-unknown" && content == "tool is unavailable"
|
||||
));
|
||||
}
|
||||
|
||||
struct ScriptedRuntime {
|
||||
descriptor: RuntimeDescriptor,
|
||||
turns: Mutex<VecDeque<Vec<AgentEvent>>>,
|
||||
requests: Arc<Mutex<Vec<TurnRequest>>>,
|
||||
}
|
||||
|
||||
impl ScriptedRuntime {
|
||||
fn new(turns: Vec<Vec<AgentEvent>>, requests: Arc<Mutex<Vec<TurnRequest>>>) -> Self {
|
||||
Self {
|
||||
descriptor: RuntimeDescriptor {
|
||||
id: "scripted-provider".to_string(),
|
||||
display_name: "Scripted provider".to_string(),
|
||||
kind: RuntimeKind::Provider,
|
||||
capabilities: RuntimeCapabilities::provider(),
|
||||
},
|
||||
turns: Mutex::new(turns.into()),
|
||||
requests,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AgentRuntime for ScriptedRuntime {
|
||||
fn descriptor(&self) -> &RuntimeDescriptor {
|
||||
&self.descriptor
|
||||
}
|
||||
|
||||
async fn start_turn(
|
||||
&self,
|
||||
request: TurnRequest,
|
||||
_control: TurnControl,
|
||||
) -> Result<AgentEventStream, AgentError> {
|
||||
self.requests.lock().unwrap().push(request);
|
||||
let events = self.turns.lock().unwrap().pop_front().ok_or_else(|| {
|
||||
AgentError::new(
|
||||
AgentErrorKind::Protocol,
|
||||
"scripted provider ran out of turns",
|
||||
)
|
||||
})?;
|
||||
Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok))))
|
||||
}
|
||||
}
|
||||
|
||||
fn recall_turn(index: usize) -> Vec<AgentEvent> {
|
||||
vec![
|
||||
AgentEvent::TurnStarted {
|
||||
runtime_request_id: format!("request-{index}"),
|
||||
},
|
||||
AgentEvent::Tool {
|
||||
event: ToolEvent::Proposed {
|
||||
call: ToolCall {
|
||||
id: format!("recall-{index}"),
|
||||
name: RECALL_TOOL_HISTORY_NAME.to_string(),
|
||||
arguments: serde_json::json!({"search_query": "missing"}),
|
||||
},
|
||||
},
|
||||
},
|
||||
AgentEvent::UsageUpdated {
|
||||
usage: Usage {
|
||||
input_tokens: 10,
|
||||
output_tokens: 1,
|
||||
..Usage::default()
|
||||
},
|
||||
},
|
||||
AgentEvent::TurnStopped {
|
||||
reason: StopReason::Completed,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
fn answer_turn() -> Vec<AgentEvent> {
|
||||
vec![
|
||||
AgentEvent::TurnStarted {
|
||||
runtime_request_id: "request-answer".to_string(),
|
||||
},
|
||||
AgentEvent::TextDelta {
|
||||
text: "Continuing with the answer.".to_string(),
|
||||
},
|
||||
AgentEvent::UsageUpdated {
|
||||
usage: Usage {
|
||||
input_tokens: 20,
|
||||
output_tokens: 3,
|
||||
..Usage::default()
|
||||
},
|
||||
},
|
||||
AgentEvent::TurnStopped {
|
||||
reason: StopReason::Completed,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
fn prepared_turn(messages_sent: Arc<Mutex<Vec<ConversationMessage>>>) -> PreparedRigTurn {
|
||||
let initial_messages = vec![ConversationMessage {
|
||||
role: MessageRole::User,
|
||||
content: MessageContent::Text("Inspect the issue.".to_string()),
|
||||
}];
|
||||
let mut request = TurnRequest::new("test-model", initial_messages.clone());
|
||||
request.conversation_id = Some("conversation".to_string());
|
||||
request.tools = vec![ToolDefinition {
|
||||
name: RECALL_TOOL_HISTORY_NAME.to_string(),
|
||||
description: "Recall prior tool output".to_string(),
|
||||
input_schema: serde_json::json!({"type": "object"}),
|
||||
}];
|
||||
PreparedRigTurn {
|
||||
task_id: "task".to_string(),
|
||||
needs_create_task: false,
|
||||
user_query: None,
|
||||
request,
|
||||
persistent_messages: initial_messages,
|
||||
tool_result_archive: Vec::new(),
|
||||
messages_sent,
|
||||
mcp_tool_aliases: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_scripted_turn(
|
||||
turns: Vec<Vec<AgentEvent>>,
|
||||
) -> (Vec<StreamEvent>, Vec<TurnRequest>, Vec<ConversationMessage>) {
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let messages_sent = Arc::new(Mutex::new(Vec::new()));
|
||||
let runtime = ScriptedRuntime::new(turns, requests.clone());
|
||||
let (cancel_tx, cancellation_rx) = oneshot::channel();
|
||||
let events = rig_response_stream(
|
||||
runtime,
|
||||
prepared_turn(messages_sent.clone()),
|
||||
SkillPathOrigin::Local,
|
||||
Some(100_000),
|
||||
"scripted",
|
||||
cancellation_rx,
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.expect("scripted response should succeed");
|
||||
drop(cancel_tx);
|
||||
let requests = requests.lock().unwrap().clone();
|
||||
let messages_sent = messages_sent.lock().unwrap().clone();
|
||||
(events, requests, messages_sent)
|
||||
}
|
||||
|
||||
fn agent_output_texts(events: &[StreamEvent]) -> Vec<&str> {
|
||||
let mut texts = Vec::new();
|
||||
for event in events {
|
||||
let StreamEvent::Response(response) = event else {
|
||||
continue;
|
||||
};
|
||||
let Some(response_event::Type::ClientActions(actions)) = &response.r#type else {
|
||||
continue;
|
||||
};
|
||||
for action in &actions.actions {
|
||||
let Some(client_action::Action::AddMessagesToTask(add)) = &action.action else {
|
||||
continue;
|
||||
};
|
||||
for message in &add.messages {
|
||||
if let Some(message::Message::AgentOutput(output)) = &message.message {
|
||||
texts.push(output.text.as_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
texts
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inline_recall_starts_a_followup_provider_turn_with_the_paired_result() {
|
||||
let (events, requests, messages_sent) =
|
||||
run_scripted_turn(vec![recall_turn(1), answer_turn()]).await;
|
||||
|
||||
assert_eq!(requests.len(), 2);
|
||||
assert_eq!(requests[1].messages.len(), 3);
|
||||
assert!(matches!(
|
||||
&requests[1].messages[1].content,
|
||||
MessageContent::ToolUse {
|
||||
tool_use_id,
|
||||
name,
|
||||
..
|
||||
} if tool_use_id == "recall-1" && name == RECALL_TOOL_HISTORY_NAME
|
||||
));
|
||||
assert!(matches!(
|
||||
&requests[1].messages[2].content,
|
||||
MessageContent::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error: false,
|
||||
} if tool_use_id == "recall-1"
|
||||
&& content == "No matching tool calls found in conversation history."
|
||||
));
|
||||
assert_eq!(
|
||||
events
|
||||
.iter()
|
||||
.filter(|event| matches!(
|
||||
event,
|
||||
StreamEvent::Response(response)
|
||||
if matches!(response.r#type, Some(response_event::Type::Init(_)))
|
||||
))
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
events
|
||||
.iter()
|
||||
.filter(|event| matches!(
|
||||
event,
|
||||
StreamEvent::Response(response)
|
||||
if matches!(response.r#type, Some(response_event::Type::Finished(_)))
|
||||
))
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
assert!(events
|
||||
.iter()
|
||||
.all(|event| !matches!(event, StreamEvent::ToolProposed(_))));
|
||||
assert_eq!(agent_output_texts(&events), ["Continuing with the answer."]);
|
||||
|
||||
let finished = events.iter().find_map(|event| {
|
||||
let StreamEvent::Response(response) = event else {
|
||||
return None;
|
||||
};
|
||||
let Some(response_event::Type::Finished(finished)) = &response.r#type else {
|
||||
return None;
|
||||
};
|
||||
Some(finished)
|
||||
});
|
||||
let finished = finished.expect("stream should finish");
|
||||
assert_eq!(finished.token_usage[0].total_input, 30);
|
||||
assert_eq!(finished.token_usage[0].output, 4);
|
||||
assert!(matches!(
|
||||
messages_sent.last().map(|message| &message.content),
|
||||
Some(MessageContent::Text(text)) if text == "Continuing with the answer."
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn repeated_inline_recall_stops_with_a_visible_loop_limit_message() {
|
||||
let turns = (0..=MAX_INLINE_TOOL_CONTINUATIONS)
|
||||
.map(recall_turn)
|
||||
.collect();
|
||||
let (events, requests, messages_sent) = run_scripted_turn(turns).await;
|
||||
|
||||
assert_eq!(requests.len(), MAX_INLINE_TOOL_CONTINUATIONS + 1);
|
||||
assert!(agent_output_texts(&events).contains(&INLINE_TOOL_LOOP_MESSAGE));
|
||||
assert!(matches!(
|
||||
messages_sent.last().map(|message| &message.content),
|
||||
Some(MessageContent::Text(text)) if text == INLINE_TOOL_LOOP_MESSAGE
|
||||
));
|
||||
let finished = events.iter().find_map(|event| {
|
||||
let StreamEvent::Response(response) = event else {
|
||||
return None;
|
||||
};
|
||||
let Some(response_event::Type::Finished(finished)) = &response.r#type else {
|
||||
return None;
|
||||
};
|
||||
Some(finished)
|
||||
});
|
||||
assert!(matches!(
|
||||
finished.and_then(|finished| finished.reason.as_ref()),
|
||||
Some(response_event::stream_finished::Reason::Other(_))
|
||||
));
|
||||
}
|
||||
|
||||
@@ -46,6 +46,32 @@ fn shell_calls_become_domain_actions_without_a_proto_round_trip() {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_plan_calls_become_document_actions() {
|
||||
let action = action_from_tool_call(
|
||||
"task-1",
|
||||
&call(
|
||||
"create_plan",
|
||||
serde_json::json!({
|
||||
"documents": [{
|
||||
"title": "Duplicate content items",
|
||||
"content": "# Implementation plan"
|
||||
}]
|
||||
}),
|
||||
),
|
||||
&SkillPathOrigin::Local,
|
||||
&HashMap::new(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let AIAgentActionType::CreateDocuments(request) = action.action else {
|
||||
panic!("expected create-documents action");
|
||||
};
|
||||
assert_eq!(request.documents.len(), 1);
|
||||
assert_eq!(request.documents[0].title, "Duplicate content items");
|
||||
assert_eq!(request.documents[0].content, "# Implementation plan");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn edit_calls_preserve_file_edits_in_the_domain_model() {
|
||||
let action = action_from_tool_call(
|
||||
|
||||
Reference in New Issue
Block a user