From f4a04d02408168ff50ca36d0fd895f69046b30d6 Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Fri, 14 Aug 2026 09:01:53 -0500 Subject: [PATCH] Fix Rig tool continuation and plan creation --- AGENTS.md | 4 +- app/src/ai/agent/api/impl_tests.rs | 10 + app/src/ai/bedrock/request_translator.rs | 2 +- app/src/ai/runtime/event_translator.rs | 5 + app/src/ai/runtime/rig.rs | 406 +++++++++++------- app/src/ai/runtime/rig_request.rs | 5 + app/src/ai/runtime/rig_request_tests.rs | 25 ++ app/src/ai/runtime/rig_tests.rs | 278 +++++++++++- app/src/ai/runtime/rig_tool_tests.rs | 26 ++ crates/galaxy_agent_core/src/tool_policy.rs | 1 + .../src/tool_policy_tests.rs | 15 + .../src/openai_compatible_tests.rs | 20 +- 12 files changed, 642 insertions(+), 155 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 76872820..91bf291c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -119,8 +119,10 @@ context_size = 128000 Key invariants: - Known tools are in `KNOWN_TOOLS` constant in `response_translator.rs` - Tool definitions are built via `tool_definition_for_name()` in `convert_request.rs`; includes `recall_tool_history` for retrieving past tool results +- Direct-provider normal and plan turns must advertise `read_plan`, `create_plan`, and `edit_plan` when the matching document capabilities are enabled; plan-creation requests should call `create_plan` after research rather than only returning prose or claiming the tool is unavailable - Unknown/hallucinated tool calls are caught in the stream, paired with synthetic error results, and now emit a visible `AgentOutput` text message to the UI -- `recall_tool_history` is handled inline in the response translator (synthetic result from `messages_sent`) +- `recall_tool_history` is handled inline by direct-provider adapters using a synthetic result from `messages_sent`; the Rig adapter must automatically start a bounded follow-up provider turn after pairing that result, without continuing turns that proposed client-executed tools +- `recall_tool_history` must exclude prior calls to itself from candidates so inline continuation cannot recursively recall synthetic recall results - Tool result archive: before progressive summarization drains messages, `ConversationMessage::archive_tool_results()` extracts all tool_use/tool_result pairs into a separate `tool_result_archive` vec. `recall_tool_history` searches both live history + archived results, and supports a `tool_use_id` parameter for exact ID lookup - Prompt caching uses three cache points: system prompt, conversation history (second-to-last message), tool config - `ensure_tool_results_paired()` enforces Bedrock's invariant that every `tool_use` has a matching `tool_result` diff --git a/app/src/ai/agent/api/impl_tests.rs b/app/src/ai/agent/api/impl_tests.rs index d924e4a9..0338120f 100644 --- a/app/src/ai/agent/api/impl_tests.rs +++ b/app/src/ai/agent/api/impl_tests.rs @@ -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); diff --git a/app/src/ai/bedrock/request_translator.rs b/app/src/ai/bedrock/request_translator.rs index 167afd9d..0c5fc1a8 100644 --- a/app/src/ai/bedrock/request_translator.rs +++ b/app/src/ai/bedrock/request_translator.rs @@ -1895,7 +1895,7 @@ pub fn default_tool_definitions() -> Vec { }, 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": { diff --git a/app/src/ai/runtime/event_translator.rs b/app/src/ai/runtime/event_translator.rs index 41cbc62c..0666defe 100644 --- a/app/src/ai/runtime/event_translator.rs +++ b/app/src/ai/runtime/event_translator.rs @@ -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) { if self.initialized { return; diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index 1e125cc1..5bed5df4 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -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>>, +) -> Result, ()> { + messages_sent + .lock() + .map(|sent| sent.clone()) + .map_err(|_| ()) +} + fn append_tool_result( messages_sent: &std::sync::Arc>>, result: ToolResult, @@ -408,6 +499,29 @@ fn append_tool_result( } } +fn append_assistant_text( + messages_sent: &std::sync::Arc>>, + 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>>, reasoning_text: &str, diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index 7b4fcceb..f32bffae 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -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 } diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index c7c16947..c673d8d4 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -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(); diff --git a/app/src/ai/runtime/rig_tests.rs b/app/src/ai/runtime/rig_tests.rs index 7be07d0c..a1a22ebe 100644 --- a/app/src/ai/runtime/rig_tests.rs +++ b/app/src/ai/runtime/rig_tests.rs @@ -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>>, + requests: Arc>>, +} + +impl ScriptedRuntime { + fn new(turns: Vec>, requests: Arc>>) -> 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 { + 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 { + 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 { + 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>>) -> 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, Vec, Vec) { + 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::>() + .await + .into_iter() + .collect::, _>>() + .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(_)) + )); +} diff --git a/app/src/ai/runtime/rig_tool_tests.rs b/app/src/ai/runtime/rig_tool_tests.rs index 6fdd0faf..8bb17192 100644 --- a/app/src/ai/runtime/rig_tool_tests.rs +++ b/app/src/ai/runtime/rig_tool_tests.rs @@ -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( diff --git a/crates/galaxy_agent_core/src/tool_policy.rs b/crates/galaxy_agent_core/src/tool_policy.rs index 678a01b0..76052944 100644 --- a/crates/galaxy_agent_core/src/tool_policy.rs +++ b/crates/galaxy_agent_core/src/tool_policy.rs @@ -188,6 +188,7 @@ pub fn recall_tool_history( let filtered = entries .iter() + .filter(|entry| entry.name != RECALL_TOOL_HISTORY_NAME) .filter(|entry| { (query.tool_use_id.is_empty() || entry.tool_use_id == query.tool_use_id) && (query.tool_name.is_empty() || entry.name == query.tool_name) diff --git a/crates/galaxy_agent_core/src/tool_policy_tests.rs b/crates/galaxy_agent_core/src/tool_policy_tests.rs index 96986bcc..4b4dd597 100644 --- a/crates/galaxy_agent_core/src/tool_policy_tests.rs +++ b/crates/galaxy_agent_core/src/tool_policy_tests.rs @@ -141,6 +141,21 @@ fn recall_supports_exact_call_id_and_case_insensitive_text_search() { ); } +#[test] +fn recall_never_returns_a_prior_recall_call() { + let messages = tool_exchange( + "recall-1", + RECALL_TOOL_HISTORY_NAME, + serde_json::json!({}), + "a prior synthetic recall result", + ); + + assert_eq!( + recall_tool_history(&messages, &[], ToolHistoryQuery::default()), + "No matching tool calls found in conversation history." + ); +} + #[test] fn loop_guard_detects_repeated_failures_and_resets_after_detection() { let mut guard = ToolLoopGuard::new(5, 3); diff --git a/crates/galaxy_agent_rig/src/openai_compatible_tests.rs b/crates/galaxy_agent_rig/src/openai_compatible_tests.rs index b1dfece0..520caf20 100644 --- a/crates/galaxy_agent_rig/src/openai_compatible_tests.rs +++ b/crates/galaxy_agent_rig/src/openai_compatible_tests.rs @@ -226,16 +226,28 @@ fn request_conversion_preserves_history_tools_and_limits() { request.system_prompt = Some("Be useful".to_string()); request.max_output_tokens = Some(123); request.tools.push(galaxy_agent_core::ToolDefinition { - name: "shell".to_string(), - description: "Run a command".to_string(), - input_schema: serde_json::json!({"type": "object"}), + name: "create_plan".to_string(), + description: "Create a plan document".to_string(), + input_schema: serde_json::json!({ + "type": "object", + "properties": {"documents": {"type": "array"}}, + "required": ["documents"] + }), }); let converted = build_completion_request(request, Some(999), true).unwrap(); assert_eq!(converted.max_tokens, Some(123)); assert_eq!(converted.tools.len(), 1); - assert_eq!(converted.tools[0].name, "shell"); + assert_eq!(converted.tools[0].name, "create_plan"); + assert_eq!( + converted.tools[0].parameters, + serde_json::json!({ + "type": "object", + "properties": {"documents": {"type": "array"}}, + "required": ["documents"] + }) + ); assert_eq!(converted.chat_history.len(), 2); assert!(matches!( converted.chat_history.iter().next(),