From 4aae078d7563f54bf88a99f0c87c5e31de3a715e Mon Sep 17 00:00:00 2001 From: Ahmed Hesham Abdelkader <23265119+ahmedhesham6@users.noreply.github.com> Date: Sat, 9 May 2026 00:34:04 +0200 Subject: [PATCH 1/5] refactor(runtime): use stakai message models --- cli/src/commands/acp/fs_handler.rs | 12 +- cli/src/commands/acp/server.rs | 257 +- cli/src/commands/agent/run/checkpoint.rs | 224 +- cli/src/commands/agent/run/helpers.rs | 97 +- cli/src/commands/agent/run/mcp_init.rs | 6 +- cli/src/commands/agent/run/mode_async.rs | 351 ++- .../commands/agent/run/mode_interactive.rs | 236 +- cli/src/commands/agent/run/pause.rs | 42 +- cli/src/commands/agent/run/renderer.rs | 37 +- cli/src/commands/agent/run/stream.rs | 731 ++--- cli/src/commands/agent/run/tooling.rs | 17 +- cli/src/commands/agent/run/tui.rs | 3 +- cli/src/commands/sessions/messages.rs | 37 +- cli/src/commands/sessions/output.rs | 85 +- cli/src/commands/sessions/tests.rs | 112 +- cli/src/utils/agent_context.rs | 47 + libs/ai/docs/PROVIDER_OPTIONS_PLAN.md | 17 +- libs/ai/src/providers/gemini/convert.rs | 3 +- libs/api/src/client/mod.rs | 3 +- libs/api/src/client/provider.rs | 318 +- libs/api/src/client/stakai.rs | 165 ++ libs/api/src/lib.rs | 17 +- libs/api/src/local/context_managers/common.rs | 108 +- .../file_scratchpad_context_manager.rs | 445 +-- libs/api/src/local/context_managers/mod.rs | 4 +- .../scratchpad_context_manager.rs | 15 +- .../simple_context_manager.rs | 22 +- .../task_board_context_manager.rs | 2591 +---------------- .../hooks/file_scratchpad_context/mod.rs | 24 +- .../hooks/inline_scratchpad_context/mod.rs | 24 +- .../src/local/hooks/task_board_context/mod.rs | 35 +- libs/api/src/local/tests.rs | 200 +- libs/api/src/models.rs | 146 +- libs/api/src/stakpak/models.rs | 3 +- libs/api/src/storage.rs | 7 +- libs/mcp/client/src/lib.rs | 2 +- libs/mcp/client/src/local.rs | 2 +- libs/mcp/server/src/local_tools.rs | 6 +- libs/server/README.md | 2 +- libs/server/src/lib.rs | 1 - libs/server/src/message_bridge.rs | 82 - libs/server/src/routes.rs | 3 +- libs/server/src/session_actor.rs | 11 +- libs/shared/src/models/agent_runtime.rs | 91 + libs/shared/src/models/async_manifest.rs | 8 +- libs/shared/src/models/integrations/mcp.rs | 22 +- libs/shared/src/models/integrations/openai.rs | 1065 +------ libs/shared/src/models/llm.rs | 229 -- libs/shared/src/models/mod.rs | 2 +- libs/shared/src/models/stakai_adapter.rs | 1863 ------------ libs/shared/src/remote_connection.rs | 60 +- tui/src/app/events.rs | 20 +- tui/src/app/types.rs | 42 +- tui/src/event_loop.rs | 4 +- tui/src/services/approval_bar.rs | 17 +- tui/src/services/auto_approve.rs | 25 +- tui/src/services/bash_block.rs | 64 +- tui/src/services/board_tasks.rs | 25 +- tui/src/services/changeset.rs | 16 +- tui/src/services/file_diff.rs | 7 +- tui/src/services/handlers/ask_user.rs | 14 +- tui/src/services/handlers/dialog.rs | 28 +- tui/src/services/handlers/input.rs | 6 +- tui/src/services/handlers/mod.rs | 52 +- tui/src/services/handlers/shell.rs | 19 +- tui/src/services/handlers/tool.rs | 89 +- tui/src/services/image_upload.rs | 20 +- tui/src/services/message.rs | 103 +- 68 files changed, 2107 insertions(+), 8334 deletions(-) create mode 100644 libs/api/src/client/stakai.rs delete mode 100644 libs/server/src/message_bridge.rs create mode 100644 libs/shared/src/models/agent_runtime.rs delete mode 100644 libs/shared/src/models/stakai_adapter.rs diff --git a/cli/src/commands/acp/fs_handler.rs b/cli/src/commands/acp/fs_handler.rs index 42e194184..bd0887d48 100644 --- a/cli/src/commands/acp/fs_handler.rs +++ b/cli/src/commands/acp/fs_handler.rs @@ -80,14 +80,13 @@ pub fn spawn_fs_handler( /// Execute filesystem tool using native ACP protocol via channel pub async fn execute_acp_fs_tool( fs_tx: &mpsc::UnboundedSender, - tool_call: &stakpak_shared::models::integrations::openai::ToolCall, + tool_call: &stakai::ToolCall, session_id: &acp::SessionId, ) -> Result, String> { - let args: serde_json::Value = serde_json::from_str(&tool_call.function.arguments) - .map_err(|e| format!("Failed to parse tool arguments: {}", e))?; + let args = tool_call.arguments.clone(); use super::tool_names; - let stripped_name = super::utils::strip_tool_name(&tool_call.function.name); + let stripped_name = super::utils::strip_tool_name(&tool_call.name); match stripped_name { tool_names::VIEW => { let path = args @@ -278,9 +277,6 @@ pub async fn execute_acp_fs_tool( structured_content: None, })) } - _ => Err(format!( - "Unknown filesystem tool: {}", - tool_call.function.name - )), + _ => Err(format!("Unknown filesystem tool: {}", tool_call.name)), } } diff --git a/cli/src/commands/acp/server.rs b/cli/src/commands/acp/server.rs index adfddbd49..890c942b4 100644 --- a/cli/src/commands/acp/server.rs +++ b/cli/src/commands/acp/server.rs @@ -6,18 +6,14 @@ use agent_client_protocol::{ SetSessionModelRequest, SetSessionModelResponse, }; use futures_util::StreamExt; -use stakpak_api::models::ApiStreamError; +use stakai::{ContentPart, Message, MessageContent, Role, StreamEvent, Tool, ToolCall}; +use stakpak_api::models::{AgentStreamEvent, ApiStreamError, CompletionResponse}; use stakpak_api::storage::CreateSessionRequest; use stakpak_api::{AgentClient, AgentClientConfig, AgentProvider, StakpakConfig}; use stakpak_api::{Model, ModelLimit}; use stakpak_mcp_client::McpClient; +use stakpak_shared::models::agent_runtime::{ToolCallResultProgress, ToolCallResultStatus}; use stakpak_shared::models::integrations::mcp::CallToolResultExt; -use stakpak_shared::models::integrations::openai::{ - ChatCompletionChoice, ChatCompletionResponse, ChatCompletionStreamResponse, ChatMessage, - FinishReason, MessageContent, Role, Tool, ToolCall, ToolCallResultProgress, - ToolCallResultStatus, -}; -use stakpak_shared::models::llm::LLMTokenUsage; use std::cell::Cell; use std::path::Path; use std::sync::Arc; @@ -38,7 +34,7 @@ pub struct StakpakAcpAgent { current_session_id: Cell>, progress_tx: Option>, // Add persistent message history for conversation context - messages: Arc>>, + messages: Arc>>, // Add permission request channel permission_request_tx: Option< mpsc::UnboundedSender<( @@ -61,6 +57,28 @@ pub struct StakpakAcpAgent { client_capabilities: Arc>, } +fn stakai_tool_calls_from_message(message: &Message) -> Vec { + message + .parts() + .into_iter() + .filter_map(|part| match part { + ContentPart::ToolCall { + id, + name, + arguments, + metadata, + .. + } => Some(ToolCall { + id, + name, + arguments, + metadata, + }), + _ => None, + }) + .collect() +} + impl StakpakAcpAgent { /// Convert internal Model to ACP ModelInfo fn model_to_acp_model_info(model: &Model) -> ModelInfo { @@ -294,7 +312,7 @@ impl StakpakAcpAgent { ) -> Result { log::info!( "Requesting permission for tool: {} - {}", - tool_call.function.name, + tool_call.name, tool_title ); log::info!("Tool Call ID: {}", tool_call_id); @@ -320,10 +338,7 @@ impl StakpakAcpAgent { acp::ToolCallId::new(tool_call_id.clone()), acp::ToolCallUpdateFields::new() .title(tool_title.to_string()) - .raw_input( - serde_json::from_str(&tool_call.function.arguments) - .unwrap_or(serde_json::Value::Null), - ), + .raw_input(tool_call.arguments.clone()), ), options, ); @@ -759,7 +774,7 @@ impl StakpakAcpAgent { &self, tool_calls: Vec, session_id: &acp::SessionId, - ) -> Result, acp::Error> { + ) -> Result, acp::Error> { log::info!("Processing {} tool calls", tool_calls.len()); let mut tool_calls_queue = tool_calls; @@ -790,7 +805,7 @@ impl StakpakAcpAgent { log::info!( "🔧 DEBUG: Processing tool call: {} (original_id: {}, new_id: {})", - tool_call.function.name, + tool_call.name, tool_call.id, tool_call_id ); @@ -800,10 +815,8 @@ impl StakpakAcpAgent { let mut active_tool_calls = self.active_tool_calls.lock().await; active_tool_calls.push(tool_call.clone()); } - let raw_input = serde_json::from_str(&tool_call.function.arguments) - .unwrap_or(serde_json::Value::Null); - let stripped_name = - crate::commands::acp::utils::strip_tool_name(&tool_call.function.name); + let raw_input = tool_call.arguments.clone(); + let stripped_name = crate::commands::acp::utils::strip_tool_name(&tool_call.name); let tool_title = self.generate_tool_title(stripped_name, &raw_input); let tool_kind = self.get_tool_kind(stripped_name); @@ -938,7 +951,7 @@ impl StakpakAcpAgent { let result = if should_delegate { log::info!( "🔧 DEBUG: Executing filesystem tool via native ACP: {}", - tool_call.function.name + tool_call.name ); // Execute using native ACP filesystem protocol @@ -953,10 +966,7 @@ impl StakpakAcpAgent { acp::Error::internal_error().data(format!("Tool execution failed: {e}")) })? } else if let Some(ref mcp_client) = self.mcp_client { - log::info!( - "Executing tool call: {} with MCP client", - tool_call.function.name - ); + log::info!("Executing tool call: {} with MCP client", tool_call.name); // Create cancellation receiver for this tool call let tool_cancel_rx = self.tool_cancel_tx.as_ref().map(|tx| tx.subscribe()); @@ -976,10 +986,8 @@ impl StakpakAcpAgent { acp::Error::internal_error().data(format!("MCP tool execution failed: {e}")) })? } else { - let error_msg = format!( - "No execution method available for tool: {}", - tool_call.function.name - ); + let error_msg = + format!("No execution method available for tool: {}", tool_call.name); log::error!("{error_msg}"); return Err(acp::Error::internal_error().data(error_msg)); }; @@ -1138,38 +1146,22 @@ impl StakpakAcpAgent { async fn process_acp_streaming_response_with_cancellation( &self, - stream: impl futures_util::Stream>, + stream: impl futures_util::Stream>, session_id: &acp::SessionId, - ) -> Result { + ) -> Result { let mut stream = Box::pin(stream); let current_model = self.model.read().await; - let mut chat_completion_response = ChatCompletionResponse { + let mut completion_response = CompletionResponse { id: "".to_string(), - object: "".to_string(), created: 0, model: current_model.id.clone(), - choices: vec![], - usage: LLMTokenUsage { - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - prompt_tokens_details: None, - }, - system_fingerprint: None, + message: Message::new(Role::Assistant, ""), + usage: stakai::Usage::default(), metadata: None, }; - let mut chat_message = ChatMessage { - role: Role::Assistant, - content: None, - name: None, - tool_calls: None, - tool_call_id: None, - usage: None, - ..Default::default() - }; - + let mut text = String::new(); let mut tool_call_accumulator = ToolCallAccumulator::new(); // Compile regex once outside the loop @@ -1208,35 +1200,18 @@ impl StakpakAcpAgent { }; match &response { - Ok(response) => { - if response.choices.is_empty() { - continue; + Ok(response) => match response { + AgentStreamEvent::Model(model) => { + completion_response.model = model.id.clone(); } - let delta = &response.choices[0].delta; - - chat_completion_response = ChatCompletionResponse { - id: response.id.clone(), - object: response.object.clone(), - created: response.created, - model: response.model.clone(), - choices: vec![], - usage: LLMTokenUsage { - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - prompt_tokens_details: None, - }, - system_fingerprint: None, - metadata: None, - }; - - if let Some(content) = &delta.content { - chat_message.content = - Some(MessageContent::String(match chat_message.content { - Some(MessageContent::String(old_content)) => old_content + content, - _ => content.clone(), - })); - + AgentStreamEvent::Metadata(metadata) => { + completion_response.metadata = Some(metadata.clone()); + } + AgentStreamEvent::Event(StreamEvent::Start { id }) => { + completion_response.id = id.clone(); + } + AgentStreamEvent::Event(StreamEvent::TextDelta { delta: content, .. }) => { + text.push_str(content); // Accumulate the raw content in the current streaming message BEFORE filtering { let mut current_message = self.current_streaming_message.lock().await; @@ -1281,14 +1256,21 @@ impl StakpakAcpAgent { rx.await.map_err(|_| "Failed to await streaming chunk")?; } } - - // Handle tool calls streaming - if let Some(tool_calls) = &delta.tool_calls { - for delta_tool_call in tool_calls { - tool_call_accumulator.process_delta(delta_tool_call); - } + AgentStreamEvent::Event( + event @ (StreamEvent::ToolCallStart { .. } + | StreamEvent::ToolCallDelta { .. } + | StreamEvent::ToolCallEnd { .. }), + ) => { + tool_call_accumulator.process_event(event); } - } + AgentStreamEvent::Event(StreamEvent::Finish { usage, .. }) => { + completion_response.usage = usage.clone(); + } + AgentStreamEvent::Event(StreamEvent::ReasoningDelta { .. }) => {} + AgentStreamEvent::Event(StreamEvent::Error { message }) => { + return Err(format!("Stream error: {}", message)); + } + }, Err(e) => { return Err(format!("Stream error: {:?}", e)); } @@ -1316,20 +1298,20 @@ impl StakpakAcpAgent { // Get accumulated tool calls (already filtered for empty IDs) let final_tool_calls = tool_call_accumulator.into_tool_calls(); - chat_message.tool_calls = if final_tool_calls.is_empty() { - None + completion_response.message = if final_tool_calls.is_empty() { + Message::new(Role::Assistant, text) } else { - Some(final_tool_calls) + let mut parts = Vec::new(); + if !text.is_empty() { + parts.push(ContentPart::text(text)); + } + parts.extend(final_tool_calls.into_iter().map(|tool_call| { + ContentPart::tool_call(tool_call.id, tool_call.name, tool_call.arguments) + })); + Message::new(Role::Assistant, MessageContent::Parts(parts)) }; - chat_completion_response.choices.push(ChatCompletionChoice { - index: 0, - message: chat_message.clone(), - finish_reason: FinishReason::Stop, - logprobs: None, - }); - - Ok(chat_completion_response) + Ok(completion_response) } pub async fn run_stdio(&self) -> Result<(), String> { @@ -1802,7 +1784,7 @@ impl acp::Agent for StakpakAcpAgent { async fn prompt(&self, args: acp::PromptRequest) -> Result { log::info!("Received prompt request {args:?}"); - // Convert prompt to your ChatMessage format + // Convert prompt to StakAI message format let prompt_text = args .prompt .iter() @@ -1878,48 +1860,20 @@ impl acp::Agent for StakpakAcpAgent { } }; log::info!( - "Chat completion successful, response choices: {}", - response.choices.len() + "Chat completion successful, response message: {:?}", + response.message ); - if !response.choices.is_empty() { - log::info!("First choice message: {:?}", response.choices[0].message); - log::info!( - "First choice content: {:?}", - response.choices[0].message.content - ); - } // Add assistant response to conversation history { let mut messages = self.messages.lock().await; - messages.push(response.choices[0].message.clone()); + messages.push(response.message.clone()); } - let content = if let Some(content) = &response.choices[0].message.content { - match content { - MessageContent::String(s) => { - log::info!("Content from chat completion: '{}'", s); - s.clone() - } - MessageContent::Array(parts) => { - let extracted_content = parts - .iter() - .filter_map(|part| part.text.as_ref()) - .map(|text| text.as_str()) - .filter(|text| !text.starts_with("")) - .collect::>() - .join("\n"); - log::info!( - "Content from chat completion array: '{}'", - extracted_content - ); - extracted_content - } - } - } else { + let content = response.message.content.text().unwrap_or_default(); + if content.is_empty() { log::warn!("No content in chat completion response"); - String::new() - }; + } log::info!("Final content to send: '{}'", content); @@ -1936,12 +1890,7 @@ impl acp::Agent for StakpakAcpAgent { }; // Check if the initial response has tool calls - let mut has_tool_calls = response.choices[0] - .message - .tool_calls - .as_ref() - .map(|tc| !tc.is_empty()) - .unwrap_or(false); + let mut has_tool_calls = !stakai_tool_calls_from_message(&response.message).is_empty(); log::info!("Initial response has tool calls: {}", has_tool_calls); @@ -1982,7 +1931,8 @@ impl acp::Agent for StakpakAcpAgent { } }; - if let Some(tool_calls) = latest_message.tool_calls.as_ref() { + let tool_calls = stakai_tool_calls_from_message(latest_message); + { if tool_calls.is_empty() { break; // No more tool calls, exit loop } @@ -1991,7 +1941,7 @@ impl acp::Agent for StakpakAcpAgent { // Process tool calls with cancellation support let tool_results = self - .process_tool_calls_with_cancellation(tool_calls.clone(), &args.session_id) + .process_tool_calls_with_cancellation(tool_calls, &args.session_id) .await .map_err(|e| { log::error!("Tool call processing failed: {}", e); @@ -2000,11 +1950,9 @@ impl acp::Agent for StakpakAcpAgent { // Check if any tool calls were cancelled in the current processing let has_cancelled_tool_calls = tool_results.iter().any(|msg| { - if let Some(MessageContent::String(text)) = &msg.content { - text.contains("TOOL_CALL_CANCELLED") - } else { - false - } + msg.content + .text() + .is_some_and(|text| text.contains("TOOL_CALL_CANCELLED")) }); // Add tool results to conversation history @@ -2073,24 +2021,21 @@ impl acp::Agent for StakpakAcpAgent { // Add follow-up response to conversation history { let mut messages = self.messages.lock().await; - messages.push(follow_up_response.choices[0].message.clone()); + messages.push(follow_up_response.message.clone()); } // Update current_messages for the next iteration - current_messages.push(follow_up_response.choices[0].message.clone()); + current_messages.push(follow_up_response.message.clone()); // Check if the follow-up response has more tool calls - has_tool_calls = follow_up_response.choices[0] - .message - .tool_calls - .as_ref() - .map(|tc| !tc.is_empty()) - .unwrap_or(false); + has_tool_calls = !stakai_tool_calls_from_message( + current_messages + .last() + .unwrap_or(&Message::new(Role::Assistant, "")), + ) + .is_empty(); log::info!("Follow-up response has tool calls: {}", has_tool_calls); - } else { - // No tool calls in the latest message, exit the loop - break; } } @@ -2143,7 +2088,7 @@ impl acp::Agent for StakpakAcpAgent { // Add cancellation messages for each active tool call for tool_call in active_tool_calls { - log::info!("Cancelling tool call: {}", tool_call.function.name); + log::info!("Cancelling tool call: {}", tool_call.name); // Add cancellation message to conversation history (like rejection logic) { diff --git a/cli/src/commands/agent/run/checkpoint.rs b/cli/src/commands/agent/run/checkpoint.rs index 5bfc9c612..526749600 100644 --- a/cli/src/commands/agent/run/checkpoint.rs +++ b/cli/src/commands/agent/run/checkpoint.rs @@ -1,17 +1,15 @@ use crate::commands::agent::run::tui::send_input_event; -use rmcp::model::CallToolResult; +use crate::utils::agent_context::strip_injected_context_blocks; +use stakai::{ContentPart, Message, MessageContent, Role, ToolCall}; use stakpak_api::AgentProvider; -use stakpak_shared::models::integrations::{ - mcp::CallToolResultExt, - openai::{ChatMessage, MessageContent, Role, ToolCall, ToolCallResult}, -}; +use stakpak_shared::models::agent_runtime::{ToolCallResult, ToolCallResultStatus}; use stakpak_tui::{InputEvent, LoadingOperation}; use uuid::Uuid; pub async fn get_checkpoint_messages( client: &dyn AgentProvider, checkpoint_id: &str, -) -> Result<(Vec, Option), String> { +) -> Result<(Vec, Option), String> { let checkpoint_uuid = Uuid::parse_str(checkpoint_id).map_err(|_| { format!( "Invalid checkpoint ID '{}' - must be a valid UUID", @@ -30,8 +28,8 @@ pub async fn get_checkpoint_messages( pub async fn extract_checkpoint_messages_and_tool_calls( checkpoint_id: &str, input_tx: &tokio::sync::mpsc::Sender, - messages: Vec, -) -> Result<(Vec, Vec), String> { + messages: Vec, +) -> Result<(Vec, Vec), String> { let mut checkpoint_messages = messages; // Append checkpoint_id to the last assistant message if present if let Some(last_message) = checkpoint_messages @@ -40,71 +38,65 @@ pub async fn extract_checkpoint_messages_and_tool_calls( .find(|message| message.role != Role::User && message.role != Role::Tool) && last_message.role == Role::Assistant { - last_message.content = Some(MessageContent::String(format!( + last_message.content = MessageContent::Text(format!( "{}\n{}", - last_message - .content - .as_ref() - .unwrap_or(&MessageContent::String(String::new())), + last_message.content.text().unwrap_or_default(), checkpoint_id - ))); + )); } for message in &checkpoint_messages { match message.role { Role::Assistant => { - if let Some(content) = &message.content { + if let Some(content) = message.content.text() { let _ = input_tx - .send(InputEvent::StreamAssistantMessage( - Uuid::new_v4(), - content.to_string(), - )) + .send(InputEvent::StreamAssistantMessage(Uuid::new_v4(), content)) .await; } } Role::User => { - if let Some(content) = &message.content { - let _ = input_tx - .send(InputEvent::AddUserMessage(content.to_string())) - .await; + if let Some(content) = message.content.text() { + let content = strip_injected_context_blocks(&content); + let _ = input_tx.send(InputEvent::AddUserMessage(content)).await; } } Role::Tool => { - let tool_call = checkpoint_messages - .iter() - .find(|checkpoint_message| { - checkpoint_message - .tool_calls - .as_ref() - .is_some_and(|tool_calls| { - message.tool_call_id.as_ref().is_some_and(|tool_call_id| { - tool_calls - .iter() - .any(|tool_call| tool_call.id == *tool_call_id) - }) - }) - }) - .and_then(|chat_message| { - chat_message.tool_calls.as_ref().and_then(|tool_calls| { - message.tool_call_id.as_ref().and_then(|tool_call_id| { - tool_calls - .iter() - .find(|tool_call| tool_call.id == *tool_call_id) - }) - }) - }); - - if let Some(tool_call) = tool_call { + for part in message.parts() { + let ContentPart::ToolResult { + tool_call_id, + content, + .. + } = part + else { + continue; + }; + let tool_call = checkpoint_messages + .iter() + .flat_map(|checkpoint_message| checkpoint_message.parts()) + .find_map(|part| match part { + ContentPart::ToolCall { + id, + name, + arguments, + metadata, + .. + } if id == tool_call_id => Some(ToolCall { + id, + name, + arguments, + metadata, + }), + _ => None, + }); + let Some(tool_call) = tool_call else { + continue; + }; let _ = send_input_event( input_tx, InputEvent::ToolResult(ToolCallResult { - call: tool_call.clone(), - result: message - .content - .as_ref() - .unwrap_or(&MessageContent::String(String::new())) - .to_string(), - status: CallToolResult::get_status_from_chat_message(message), + call: tool_call, + result: content.to_string(), + status: ToolCallResultStatus::Success, }), ) .await; @@ -118,21 +110,49 @@ pub async fn extract_checkpoint_messages_and_tool_calls( let tool_calls = checkpoint_messages .iter() .rev() - .find(|msg| msg.role == Role::Assistant && msg.tool_calls.is_some()) - .and_then(|msg| msg.tool_calls.as_ref()); + .find(|msg| { + msg.role == Role::Assistant + && msg + .parts() + .iter() + .any(|part| matches!(part, ContentPart::ToolCall { .. })) + }) + .map(|msg| { + msg.parts() + .into_iter() + .filter_map(|part| match part { + ContentPart::ToolCall { + id, + name, + arguments, + metadata, + .. + } => Some(ToolCall { + id, + name, + arguments, + metadata, + }), + _ => None, + }) + .collect::>() + }); // Filter out tool calls that already have results (Role::Tool messages) let executed_tool_ids: std::collections::HashSet = checkpoint_messages .iter() .filter(|msg| msg.role == Role::Tool) - .filter_map(|msg| msg.tool_call_id.clone()) + .flat_map(|msg| msg.parts()) + .filter_map(|part| match part { + ContentPart::ToolResult { tool_call_id, .. } => Some(tool_call_id), + _ => None, + }) .collect(); let pending_tool_calls: Vec = tool_calls .map(|tcs| { - tcs.iter() + tcs.into_iter() .filter(|tc| !executed_tool_ids.contains(&tc.id)) - .cloned() .collect() }) .unwrap_or_default(); @@ -140,36 +160,15 @@ pub async fn extract_checkpoint_messages_and_tool_calls( Ok((checkpoint_messages, pending_tool_calls)) } -pub fn extract_checkpoint_id_from_messages(messages: &[ChatMessage]) -> Option { +pub fn extract_checkpoint_id_from_messages(messages: &[Message]) -> Option { messages .last() - .and_then(|msg| msg.content.as_ref()) - .as_ref() - .and_then(|content| match content { - MessageContent::String(text) => { - if let Some(start) = text.find("") { - if let Some(end) = text.find("") { - let start_pos = start + "".len(); - Some(text[start_pos..end].to_string()) - } else { - None - } - } else { - None - } - } - MessageContent::Array(items) => { - for item in items { - if let Some(text) = &item.text - && let Some(start) = text.find("") - && let Some(end) = text.find("") - { - let start_pos = start + "".len(); - return Some(text[start_pos..end].to_string()); - } - } - None - } + .and_then(|msg| msg.content.text()) + .and_then(|text| { + let start = text.find("")?; + let end = text.find("")?; + let start_pos = start + "".len(); + Some(text[start_pos..end].to_string()) }) } @@ -178,15 +177,7 @@ pub async fn resume_session_from_checkpoint( client: &dyn AgentProvider, session_id: &str, input_tx: &tokio::sync::mpsc::Sender, -) -> Result< - ( - Vec, - Vec, - Uuid, - Option, - ), - String, -> { +) -> Result<(Vec, Vec, Uuid, Option), String> { let session_uuid = Uuid::parse_str(session_id).map_err(|e| e.to_string())?; match client.get_active_checkpoint(session_uuid).await { @@ -212,3 +203,42 @@ pub async fn resume_session_from_checkpoint( } } } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::mpsc; + + #[tokio::test] + async fn extract_checkpoint_messages_replays_user_text_without_injected_context_blocks() { + let (input_tx, mut input_rx) = mpsc::channel(8); + let messages = vec![Message::new( + Role::User, + "fix tests\n\n# Available Skills:\n- very long\n\n\nrepo instructions\n" + .to_string(), + )]; + + let (checkpoint_messages, pending_tool_calls) = extract_checkpoint_messages_and_tool_calls( + "checkpoint-id", + &input_tx, + messages.clone(), + ) + .await + .expect("checkpoint extraction"); + + assert_eq!(checkpoint_messages.len(), 1); + assert!( + checkpoint_messages[0] + .content + .text() + .expect("checkpoint text") + .contains("") + ); + assert!(pending_tool_calls.is_empty()); + + match input_rx.recv().await.expect("input event") { + InputEvent::AddUserMessage(content) => assert_eq!(content, "fix tests"), + event => panic!("unexpected input event: {event:?}"), + } + } +} diff --git a/cli/src/commands/agent/run/helpers.rs b/cli/src/commands/agent/run/helpers.rs index fa762795e..4c5e1e214 100644 --- a/cli/src/commands/agent/run/helpers.rs +++ b/cli/src/commands/agent/run/helpers.rs @@ -1,6 +1,5 @@ -use stakpak_shared::models::integrations::openai::{ - ChatMessage, FunctionDefinition, MessageContent, Role, Tool, ToolCallResult, -}; +use stakai::{ContentPart, Message, MessageContent, Role, Tool}; +use stakpak_shared::models::agent_runtime::ToolCallResult; use uuid::Uuid; /// Resolve a short model name against the provider's model catalog. @@ -71,22 +70,25 @@ pub fn build_resume_command( } /// Extract the checkpoint ID from the last assistant message that contains one. -pub fn extract_last_checkpoint_id(messages: &[ChatMessage]) -> Option { +pub fn extract_last_checkpoint_id(messages: &[Message]) -> Option { messages .iter() .rev() .filter(|m| m.role == Role::Assistant) - .find_map(|m| { - m.content - .as_ref() - .and_then(MessageContent::extract_checkpoint_id) + .filter_map(|m| m.content.text()) + .find_map(|content| { + content + .split("") + .nth(1) + .and_then(|s| s.split("").next()) + .and_then(|id| Uuid::parse_str(id.trim()).ok()) }) } /// Returns true when there are no user/assistant/tool turns yet. /// /// System messages (e.g. injected prompts) do not count as conversation turns. -pub fn is_first_non_system_message(messages: &[ChatMessage]) -> bool { +pub fn is_first_non_system_message(messages: &[Message]) -> bool { messages.iter().all(|message| message.role == Role::System) } @@ -107,52 +109,36 @@ pub fn convert_tools_with_filter( return None; } - Some(Tool { - r#type: "function".to_string(), - function: FunctionDefinition { - name: tool_name.to_owned(), - description: tool.description.clone().map(|d| d.to_string()), - parameters: serde_json::Value::Object((*tool.input_schema).clone()), - }, - }) + Some( + Tool::function( + tool_name.to_owned(), + tool.description + .clone() + .map(|d| d.to_string()) + .unwrap_or_default(), + ) + .parameters(serde_json::Value::Object((*tool.input_schema).clone())), + ) }) .collect() } -pub fn user_message(user_input: String) -> ChatMessage { - ChatMessage { - role: Role::User, - content: Some(MessageContent::String(user_input)), - name: None, - tool_calls: None, - tool_call_id: None, - usage: None, - ..Default::default() - } +pub fn user_message(user_input: String) -> Message { + Message::new(Role::User, user_input) } -pub fn system_message(system_prompt: String) -> ChatMessage { - ChatMessage { - role: Role::System, - content: Some(MessageContent::String(system_prompt)), - name: None, - tool_calls: None, - tool_call_id: None, - usage: None, - ..Default::default() - } +pub fn system_message(system_prompt: String) -> Message { + Message::new(Role::System, system_prompt) } -pub fn tool_result(tool_call_id: String, result: String) -> ChatMessage { - ChatMessage { - role: Role::Tool, - content: Some(MessageContent::String(result)), - name: None, - tool_calls: None, - tool_call_id: Some(tool_call_id), - usage: None, - ..Default::default() - } +pub fn tool_result(tool_call_id: String, result: String) -> Message { + Message::new( + Role::Tool, + MessageContent::Parts(vec![ContentPart::tool_result( + tool_call_id, + serde_json::Value::String(result), + )]), + ) } pub fn tool_call_history_string(tool_calls: &[ToolCallResult]) -> Option { @@ -162,16 +148,13 @@ pub fn tool_call_history_string(tool_calls: &[ToolCallResult]) -> Option let history = tool_calls .iter() .map(|tc| { - let command = if let Ok(json) = - serde_json::from_str::(&tc.call.function.arguments) - { - json.get("command") - .and_then(|v| v.as_str()) - .unwrap_or(&tc.call.function.arguments) - .to_string() - } else { - tc.call.function.arguments.clone() - }; + let command = tc + .call + .arguments + .get("command") + .and_then(|v| v.as_str()) + .map(String::from) + .unwrap_or_else(|| tc.call.arguments.to_string()); let output = if tc.result.trim().is_empty() { "No output".to_string() diff --git a/cli/src/commands/agent/run/mcp_init.rs b/cli/src/commands/agent/run/mcp_init.rs index 521648502..f93658f38 100644 --- a/cli/src/commands/agent/run/mcp_init.rs +++ b/cli/src/commands/agent/run/mcp_init.rs @@ -19,7 +19,7 @@ use stakpak_mcp_server::{ EnabledToolsConfig, MCPServerConfig, SubagentConfig, ToolMode, start_server, }; use stakpak_shared::cert_utils::CertificateChain; -use stakpak_shared::models::integrations::openai::ToolCallResultProgress; +use stakpak_shared::models::agent_runtime::ToolCallResultProgress; use std::collections::HashMap; use std::sync::Arc; use tokio::net::TcpListener; @@ -68,8 +68,8 @@ pub struct McpInitResult { pub client: Arc, /// Raw MCP tools from the server pub mcp_tools: Vec, - /// Converted tools for OpenAI format - pub tools: Vec, + /// Tools in StakAI format + pub tools: Vec, /// Shutdown handle for the MCP server pub server_shutdown_tx: broadcast::Sender<()>, /// Shutdown handle for the proxy server diff --git a/cli/src/commands/agent/run/mode_async.rs b/cli/src/commands/agent/run/mode_async.rs index 51bef495e..293313b1c 100644 --- a/cli/src/commands/agent/run/mode_async.rs +++ b/cli/src/commands/agent/run/mode_async.rs @@ -15,7 +15,6 @@ use stakpak_api::{AgentClient, AgentClientConfig, AgentProvider, Model, SessionS use stakpak_mcp_server::EnabledToolsConfig; use stakpak_shared::local_store::LocalStore; use stakpak_shared::models::async_manifest::{AsyncManifest, PauseReason, PendingToolCall}; -use stakpak_shared::models::integrations::openai::{ChatMessage, MessageContent, Role}; use stakpak_shared::models::llm::LLMTokenUsage; use stakpak_shared::secret_manager::SecretManager; use stakpak_shared::utils::{backward_compatibility_mapping, strip_tool_name}; @@ -23,6 +22,37 @@ use std::collections::HashMap; use std::time::Instant; use uuid::Uuid; +fn message_text(message: &stakai::Message) -> Option { + message.content.text().map(|text| { + text.lines() + .filter(|line| !line.starts_with("")) + .collect::>() + .join("\n") + }) +} + +fn message_tool_calls(message: &stakai::Message) -> Vec { + message + .parts() + .into_iter() + .filter_map(|part| match part { + stakai::ContentPart::ToolCall { + id, + name, + arguments, + metadata, + .. + } => Some(stakai::ToolCall { + id, + name, + arguments, + metadata, + }), + _ => None, + }) + .collect() +} + pub struct RunAsyncConfig { pub prompt: String, pub checkpoint_id: Option, @@ -183,7 +213,7 @@ impl AsyncAutoApproveConfig { pub async fn run_async(ctx: AppConfig, mut config: RunAsyncConfig) -> Result { let start_time = Instant::now(); let mut llm_response_time = std::time::Duration::new(0, 0); - let mut chat_messages: Vec = Vec::new(); + let mut chat_messages: Vec = Vec::new(); let mut total_usage = LLMTokenUsage::default(); let renderer = OutputRenderer::new(config.output_format.clone(), config.verbose); let secret_manager = SecretManager::new(config.redact_secrets, config.privacy_mode); @@ -286,7 +316,7 @@ pub async fn run_async(ctx: AppConfig, mut config: RunAsyncConfig) -> Result Result Result result?, Err(_) => { - let error_msg = format!( - "Tool '{}' timed out after 60 minutes", - tool_call.function.name - ); + let error_msg = + format!("Tool '{}' timed out after 60 minutes", tool_call.name); print!("{}", renderer.render_error(&error_msg)); chat_messages.push(tool_result(tool_call.id.clone(), error_msg)); continue; @@ -383,7 +411,7 @@ pub async fn run_async(ctx: AppConfig, mut config: RunAsyncConfig) -> Result Result Result Result Result s.clone(), - MessageContent::Array(parts) => parts - .iter() - .filter_map(|part| part.text.as_ref()) - .map(|text| text.as_str()) - .filter(|text| !text.starts_with("")) - .collect::>() - .join("\n"), - }); + let agent_message = message_text(&response.message); // Show assistant response if let Some(content_str) = &agent_message @@ -586,162 +599,144 @@ pub async fn run_async(ctx: AppConfig, mut config: RunAsyncConfig) -> Result = tool_calls - .iter() - .map(|tc| tc.function.name.as_str()) - .collect(); - if auto_approve_config.any_requires_approval(&tool_names) { - // PAUSE: tools require approval - let pending: Vec = - tool_calls.iter().map(PendingToolCall::from).collect(); - - let checkpoint_id_str = current_checkpoint_id.map(|id| id.to_string()); - let session_id_str = current_session_id.map(|id| id.to_string()); - - let pause_reason = PauseReason::ToolApprovalRequired { - pending_tool_calls: pending, - }; + // Check if pause_on_approval is enabled and any tools require approval + if let Some(ref auto_approve_config) = auto_approve { + let tool_names: Vec<&str> = tool_calls.iter().map(|tc| tc.name.as_str()).collect(); + if auto_approve_config.any_requires_approval(&tool_names) { + // PAUSE: tools require approval + let pending: Vec = + tool_calls.iter().map(PendingToolCall::from).collect(); - let resume_hint = checkpoint_id_str - .as_ref() - .map(|cid| build_resume_hint(cid, &pause_reason)); - - let manifest = AsyncManifest { - outcome: "paused".to_string(), - checkpoint_id: checkpoint_id_str.clone(), - session_id: session_id_str.clone(), - model: config.model.id.clone(), - agent_message: agent_message.clone(), - steps: step, - total_steps: prior_steps + step, - usage: total_usage.clone(), - pause_reason: Some(pause_reason.clone()), - resume_hint, - }; + let checkpoint_id_str = current_checkpoint_id.map(|id| id.to_string()); + let session_id_str = current_session_id.map(|id| id.to_string()); - // Write pause manifest - if let Err(e) = write_pause_manifest(&manifest) { - print!( - "{}", - renderer - .render_warning(&format!("Failed to write pause manifest: {}", e)) - ); - } + let pause_reason = PauseReason::ToolApprovalRequired { + pending_tool_calls: pending, + }; - // Output JSON to stdout if in JSON mode - if config.output_format == OutputFormat::Json - && let Ok(json) = serde_json::to_string_pretty(&manifest) - { - println!("{}", json); - } + let resume_hint = checkpoint_id_str + .as_ref() + .map(|cid| build_resume_hint(cid, &pause_reason)); + + let manifest = AsyncManifest { + outcome: "paused".to_string(), + checkpoint_id: checkpoint_id_str.clone(), + session_id: session_id_str.clone(), + model: config.model.id.clone(), + agent_message: agent_message.clone(), + steps: step, + total_steps: prior_steps + step, + usage: total_usage.clone(), + pause_reason: Some(pause_reason.clone()), + resume_hint, + }; + // Write pause manifest + if let Err(e) = write_pause_manifest(&manifest) { print!( "{}", - renderer.render_info("Agent paused - tools require approval") + renderer.render_warning(&format!("Failed to write pause manifest: {}", e)) ); + } - // Shutdown MCP - let _ = server_shutdown_tx.send(()); - let _ = proxy_shutdown_tx.send(()); - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - - return Ok(AsyncOutcome::Paused { - checkpoint_id: checkpoint_id_str, - session_id: session_id_str, - pause_reason, - agent_message, - }); + // Output JSON to stdout if in JSON mode + if config.output_format == OutputFormat::Json + && let Ok(json) = serde_json::to_string_pretty(&manifest) + { + println!("{}", json); } - } - // Execute all tool calls (either auto-approved or pause_on_approval is disabled) - for (i, tool_call) in tool_calls.iter().enumerate() { - // Print tool start with arguments print!( "{}", - renderer.render_tool_execution( - &tool_call.function.name, - &tool_call.function.arguments, - i, - tool_calls.len(), - ) + renderer.render_info("Agent paused - tools require approval") ); - // Add timeout for tool execution - let tool_execution = async { - run_tool_call( - &mcp_client, - &mcp_tools, - tool_call, - None, - current_session_id, - Some(config.model.id.clone()), - Some(config.model.provider.clone()), - ) - .await - }; + // Shutdown MCP + let _ = server_shutdown_tx.send(()); + let _ = proxy_shutdown_tx.send(()); + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - let result = match tokio::time::timeout( - std::time::Duration::from_secs(60 * 60), // 60 minute timeout - tool_execution, + return Ok(AsyncOutcome::Paused { + checkpoint_id: checkpoint_id_str, + session_id: session_id_str, + pause_reason, + agent_message, + }); + } + } + + // Execute all tool calls (either auto-approved or pause_on_approval is disabled) + for (i, tool_call) in tool_calls.iter().enumerate() { + // Print tool start with arguments + print!( + "{}", + renderer.render_tool_execution( + &tool_call.name, + &tool_call.arguments.to_string(), + i, + tool_calls.len(), + ) + ); + + // Add timeout for tool execution + let tool_execution = async { + run_tool_call( + &mcp_client, + &mcp_tools, + tool_call, + None, + current_session_id, + Some(config.model.id.clone()), + Some(config.model.provider.clone()), ) .await - { - Ok(result) => result?, - Err(_) => { - let error_msg = format!( - "Tool '{}' timed out after 60 minutes", - tool_call.function.name - ); - print!("{}", renderer.render_error(&error_msg)); - chat_messages.push(tool_result(tool_call.id.clone(), error_msg)); - continue; - } - }; + }; - if let Some(result) = result { - let result_content = result - .content - .iter() - .map(|c| match c.raw.as_text() { - Some(text) => text.text.clone(), - None => String::new(), - }) - .collect::>() - .join("\n"); - - // Print tool result - print!("{}", renderer.render_tool_result(&result_content)); - - chat_messages.push(tool_result(tool_call.id.clone(), result_content.clone())); - } else { - print!( - "{}", - renderer.render_warning(&format!( - "Tool '{}' returned no result", - tool_call.function.name - )) - ); + let result = match tokio::time::timeout( + std::time::Duration::from_secs(60 * 60), // 60 minute timeout + tool_execution, + ) + .await + { + Ok(result) => result?, + Err(_) => { + let error_msg = format!("Tool '{}' timed out after 60 minutes", tool_call.name); + print!("{}", renderer.render_error(&error_msg)); + chat_messages.push(tool_result(tool_call.id.clone(), error_msg)); + continue; } + }; + + if let Some(result) = result { + let result_content = result + .content + .iter() + .map(|c| match c.raw.as_text() { + Some(text) => text.text.clone(), + None => String::new(), + }) + .collect::>() + .join("\n"); + + // Print tool result + print!("{}", renderer.render_tool_result(&result_content)); + + chat_messages.push(tool_result(tool_call.id.clone(), result_content.clone())); + } else { + print!( + "{}", + renderer + .render_warning(&format!("Tool '{}' returned no result", tool_call.name)) + ); } - } else { - print!( - "{}", - renderer.render_success("No more tools to execute - agent completed successfully") - ); - break; } // Plan mode: check if plan.md status has changed after tool execution @@ -815,18 +810,8 @@ pub async fn run_async(ctx: AppConfig, mut config: RunAsyncConfig) -> Result s.clone(), - MessageContent::Array(parts) => parts - .iter() - .filter_map(|part| part.text.as_ref()) - .map(|text| text.as_str()) - .filter(|text| !text.starts_with("")) - .collect::>() - .join("\n"), - }); + .find(|m| m.role == stakai::Role::Assistant) + .and_then(message_text); // Use generic renderer functions to build the completion output print!("{}", renderer.render_section_break()); diff --git a/cli/src/commands/agent/run/mode_interactive.rs b/cli/src/commands/agent/run/mode_interactive.rs index d960d8cf5..5d80dec25 100644 --- a/cli/src/commands/agent/run/mode_interactive.rs +++ b/cli/src/commands/agent/run/mode_interactive.rs @@ -23,11 +23,10 @@ use stakpak_api::local::skills::{default_skill_directories, discover_skills}; use stakpak_api::models::{ApiStreamError, Skill}; use stakpak_api::{AgentClient, AgentClientConfig, AgentProvider, Model}; +use stakai::{ContentPart, Message, MessageContent, Role, ToolCall}; use stakpak_mcp_server::EnabledToolsConfig; +use stakpak_shared::models::agent_runtime::ToolCallResultStatus; use stakpak_shared::models::integrations::mcp::CallToolResultExt; -use stakpak_shared::models::integrations::openai::{ - ChatMessage, MessageContent, Role, ToolCall, ToolCallResultStatus, -}; use stakpak_shared::models::llm::{LLMTokenUsage, PromptTokensDetails}; use stakpak_shared::secret_manager::SecretManager; @@ -41,7 +40,7 @@ use uuid::Uuid; type ClientTaskResult = Result< ( - Vec, + Vec, Option, Option, LLMTokenUsage, @@ -97,17 +96,51 @@ async fn set_session_id( /// Returns the IDs of tool_calls from the last assistant message that don't have corresponding tool_results. /// This is used to add cancelled tool_results before inserting a user message. -fn get_unresolved_tool_call_ids(messages: &[ChatMessage]) -> Vec { +fn tool_calls_from_message(message: &Message) -> Vec { + message + .parts() + .into_iter() + .filter_map(|part| match part { + ContentPart::ToolCall { + id, + name, + arguments, + metadata, + .. + } => Some(ToolCall { + id, + name, + arguments, + metadata, + }), + _ => None, + }) + .collect() +} + +fn tool_result_ids_from_message(message: &Message) -> Vec { + message + .parts() + .into_iter() + .filter_map(|part| match part { + ContentPart::ToolResult { tool_call_id, .. } => Some(tool_call_id), + _ => None, + }) + .collect() +} + +fn get_unresolved_tool_call_ids(messages: &[Message]) -> Vec { // Find the last assistant message and check if it has tool_calls - if let Some(last_assistant_msg) = messages.iter().rev().find(|m| m.role == Role::Assistant) - && let Some(tool_calls) = &last_assistant_msg.tool_calls - && !tool_calls.is_empty() - { + if let Some(last_assistant_msg) = messages.iter().rev().find(|m| m.role == Role::Assistant) { + let tool_calls = tool_calls_from_message(last_assistant_msg); + if tool_calls.is_empty() { + return Vec::new(); + } // Collect all tool_result IDs from messages let tool_result_ids: std::collections::HashSet<_> = messages .iter() - .filter(|m| m.role == Role::Tool && m.tool_call_id.is_some()) - .filter_map(|m| m.tool_call_id.as_ref()) + .filter(|m| m.role == Role::Tool) + .flat_map(tool_result_ids_from_message) .collect(); // Return tool_call IDs that don't have corresponding tool_results @@ -124,7 +157,7 @@ fn get_unresolved_tool_call_ids(messages: &[ChatMessage]) -> Vec { /// Checks if there are pending tool calls that don't have corresponding tool_results. /// This is used to prevent sending messages to the API when tool_use blocks would be orphaned, /// which causes Anthropic API 400 errors. -fn has_pending_tool_calls(messages: &[ChatMessage], tools_queue: &[ToolCall]) -> bool { +fn has_pending_tool_calls(messages: &[Message], tools_queue: &[ToolCall]) -> bool { // If there are tools in the queue waiting to be processed, we have pending tool calls if !tools_queue.is_empty() { return true; @@ -136,7 +169,7 @@ fn has_pending_tool_calls(messages: &[ChatMessage], tools_queue: &[ToolCall]) -> /// Find the index in the messages Vec of the nth user message (1-indexed). /// Used for reverting to a specific user message by truncating the messages array. -fn find_nth_user_message_index(messages: &[ChatMessage], n: usize) -> Option { +fn find_nth_user_message_index(messages: &[Message], n: usize) -> Option { let mut count = 0; for (idx, msg) in messages.iter().enumerate() { if msg.role == Role::User { @@ -158,17 +191,16 @@ async fn send_next_tool_from_queue( tool_call: &ToolCall, ) -> Result<(), String> { let tool_name = tool_call - .function .name .strip_prefix("stakpak__") - .unwrap_or(&tool_call.function.name); + .unwrap_or(&tool_call.name); // Auto-approve ask_user — show popup directly, skip the approval bar. // If parsing fails or questions are empty, fall through to normal approval flow. if tool_name == "ask_user" - && let Ok(request) = serde_json::from_str::< - stakpak_shared::models::integrations::openai::AskUserRequest, - >(&tool_call.function.arguments) + && let Ok(request) = serde_json::from_value::< + stakpak_shared::models::agent_runtime::AskUserRequest, + >(tool_call.arguments.clone()) && !request.questions.is_empty() { send_input_event( @@ -216,7 +248,7 @@ pub async fn run_interactive( // Outer loop for profile switching 'profile_switch_loop: loop { let mut model = config.model.clone(); - let mut messages: Vec = Vec::new(); + let mut messages: Vec = Vec::new(); let mut tools_queue: Vec = Vec::new(); // Plan mode tracking — written in PlanModeActivated, read in later phases #[allow(unused_variables, unused_assignments)] @@ -591,7 +623,7 @@ pub async fn run_interactive( ) => { // Handle revert if provided - truncate messages to the specified user message index if let Some(target_user_idx) = revert_index { - // Find the ChatMessage index for the nth user message + // Find the message index for the nth user message let truncate_at = find_nth_user_message_index(&messages, target_user_idx); @@ -651,24 +683,10 @@ pub async fn run_interactive( } else { let mut parts = Vec::new(); if !redacted_user_input.trim().is_empty() { - parts.push( - stakpak_shared::models::integrations::openai::ContentPart { - r#type: "text".to_string(), - text: Some(redacted_user_input), - image_url: None, - }, - ); + parts.push(ContentPart::text(redacted_user_input)); } parts.extend(image_parts); - ChatMessage { - role: Role::User, - content: Some(MessageContent::Array(parts)), - name: None, - tool_calls: None, - tool_call_id: None, - usage: None, - ..Default::default() - } + Message::new(Role::User, MessageContent::Parts(parts)) }; send_input_event(&input_tx, InputEvent::HasUserMessage).await?; @@ -709,15 +727,14 @@ pub async fn run_interactive( OutputEvent::AcceptTool(tool_call) => { // Check if this is the ask_user tool - handle it specially let tool_name = tool_call - .function .name .strip_prefix("stakpak__") - .unwrap_or(&tool_call.function.name); + .unwrap_or(&tool_call.name); if tool_name == "ask_user" { // Parse the questions from the tool call arguments - match serde_json::from_str::< - stakpak_shared::models::integrations::openai::AskUserRequest, - >(&tool_call.function.arguments) + match serde_json::from_value::< + stakpak_shared::models::agent_runtime::AskUserRequest, + >(tool_call.arguments.clone()) { Ok(request) if !request.questions.is_empty() => { // Send the popup event to TUI @@ -741,7 +758,7 @@ pub async fn run_interactive( send_input_event( &input_tx, InputEvent::ToolResult( - stakpak_shared::models::integrations::openai::ToolCallResult { + stakpak_shared::models::agent_runtime::ToolCallResult { call: tool_call.clone(), result: error_msg, status: ToolCallResultStatus::Error, @@ -759,7 +776,7 @@ pub async fn run_interactive( send_input_event( &input_tx, InputEvent::ToolResult( - stakpak_shared::models::integrations::openai::ToolCallResult { + stakpak_shared::models::agent_runtime::ToolCallResult { call: tool_call.clone(), result: error_msg, status: ToolCallResultStatus::Error, @@ -826,7 +843,9 @@ pub async fn run_interactive( // skip adding the real result to avoid duplicate tool_call_ids. let already_resolved = messages.iter().any(|m| { m.role == Role::Tool - && m.tool_call_id.as_deref() == Some(&tool_call.id) + && tool_result_ids_from_message(m) + .iter() + .any(|id| id == &tool_call.id) }); if already_resolved { // Skip — a CANCELLED placeholder was already inserted @@ -860,7 +879,7 @@ pub async fn run_interactive( send_input_event( &input_tx, InputEvent::ToolResult( - stakpak_shared::models::integrations::openai::ToolCallResult { + stakpak_shared::models::agent_runtime::ToolCallResult { call: tool_call.clone(), result: result_content, status, @@ -1490,7 +1509,7 @@ pub async fn run_interactive( match response_result { Ok(response) => { - messages.push(response.choices[0].message.clone()); + messages.push(response.message.clone()); if let Some(session_id) = response .metadata @@ -1512,13 +1531,14 @@ pub async fn run_interactive( current_metadata = Some(state_metadata.clone()); } - // Accumulate usage from response - total_session_usage.prompt_tokens += response.usage.prompt_tokens; - total_session_usage.completion_tokens += response.usage.completion_tokens; - total_session_usage.total_tokens += response.usage.total_tokens; + let response_usage = + stakpak_api::models::stakai_usage_to_llm_usage(&response.usage); + total_session_usage.prompt_tokens += response_usage.prompt_tokens; + total_session_usage.completion_tokens += response_usage.completion_tokens; + total_session_usage.total_tokens += response_usage.total_tokens; // Accumulate prompt token details if available - if let Some(response_details) = &response.usage.prompt_tokens_details { + if let Some(response_details) = &response_usage.prompt_tokens_details { if total_session_usage.prompt_tokens_details.is_none() { total_session_usage.prompt_tokens_details = Some(PromptTokensDetails { @@ -1581,7 +1601,12 @@ pub async fn run_interactive( } // Send tool calls to TUI if present - if let Some(tool_calls) = &response.choices[0].message.tool_calls { + let tool_calls = tool_calls_from_message( + messages + .last() + .unwrap_or(&Message::new(Role::Assistant, "")), + ); + if !tool_calls.is_empty() { // Send MessageToolCalls only once with all new tools from AI send_input_event( &input_tx, @@ -1831,46 +1856,40 @@ mod tests { fn test_tool_call(id: &str) -> ToolCall { ToolCall { id: id.to_string(), - r#type: "function".to_string(), - function: stakpak_shared::models::integrations::openai::FunctionCall { - name: format!("{}_fn", id), - arguments: "{}".to_string(), - }, + name: format!("{}_fn", id), + arguments: serde_json::json!({}), metadata: None, } } - fn assistant_with_tool_calls(ids: &[&str]) -> ChatMessage { - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("assistant".to_string())), - tool_calls: Some(ids.iter().map(|id| test_tool_call(id)).collect()), - ..Default::default() - } + fn assistant_with_tool_calls(ids: &[&str]) -> Message { + let mut parts = vec![ContentPart::text("assistant")]; + parts.extend(ids.iter().map(|id| { + let tool_call = test_tool_call(id); + ContentPart::tool_call(tool_call.id, tool_call.name, tool_call.arguments) + })); + Message::new(Role::Assistant, MessageContent::Parts(parts)) } - fn tool_message(id: &str, content: &str) -> ChatMessage { - ChatMessage { - role: Role::Tool, - content: Some(MessageContent::String(content.to_string())), - tool_call_id: Some(id.to_string()), - ..Default::default() - } + fn tool_message(id: &str, content: &str) -> Message { + Message::new( + Role::Tool, + MessageContent::Parts(vec![ContentPart::tool_result( + id, + serde_json::Value::String(content.to_string()), + )]), + ) } #[test] fn get_unresolved_tool_call_ids_returns_empty_when_no_messages() { - let messages: Vec = vec![]; + let messages: Vec = vec![]; assert!(get_unresolved_tool_call_ids(&messages).is_empty()); } #[test] fn get_unresolved_tool_call_ids_returns_empty_when_no_assistant_message() { - let messages = vec![ChatMessage { - role: Role::User, - content: Some(MessageContent::String("hello".to_string())), - ..Default::default() - }]; + let messages = vec![Message::new(Role::User, "hello".to_string())]; assert!(get_unresolved_tool_call_ids(&messages).is_empty()); } @@ -1905,7 +1924,7 @@ mod tests { #[test] fn has_pending_tool_calls_returns_true_when_queue_not_empty() { - let messages: Vec = vec![]; + let messages: Vec = vec![]; let tools_queue = vec![test_tool_call("tool_1")]; assert!(has_pending_tool_calls(&messages, &tools_queue)); @@ -1913,7 +1932,7 @@ mod tests { #[test] fn has_pending_tool_calls_returns_false_when_empty_queue_and_no_messages() { - let messages: Vec = vec![]; + let messages: Vec = vec![]; let tools_queue: Vec = vec![]; assert!(!has_pending_tool_calls(&messages, &tools_queue)); @@ -1951,12 +1970,10 @@ mod tests { #[test] fn has_pending_tool_calls_returns_false_when_assistant_has_empty_tool_calls() { - let messages = vec![ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("test".to_string())), - tool_calls: Some(vec![]), - ..Default::default() - }]; + let messages = vec![Message::new( + Role::Assistant, + MessageContent::Parts(vec![ContentPart::text("test")]), + )]; let tools_queue: Vec = vec![]; assert!(!has_pending_tool_calls(&messages, &tools_queue)); @@ -1964,12 +1981,7 @@ mod tests { #[test] fn has_pending_tool_calls_returns_false_when_assistant_has_no_tool_calls() { - let messages = vec![ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("test".to_string())), - tool_calls: None, - ..Default::default() - }]; + let messages = vec![Message::new(Role::Assistant, "test".to_string())]; let tools_queue: Vec = vec![]; assert!(!has_pending_tool_calls(&messages, &tools_queue)); @@ -1980,11 +1992,7 @@ mod tests { let messages = vec![ assistant_with_tool_calls(&["tool_old"]), tool_message("tool_old", "old result"), - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("continue".to_string())), - ..Default::default() - }, + Message::new(Role::User, "continue".to_string()), assistant_with_tool_calls(&["tool_new"]), tool_message("tool_new", "new result"), ]; @@ -1999,27 +2007,15 @@ mod tests { let older = Uuid::from_u128(0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa); let newer = Uuid::from_u128(0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb); let messages = vec![ - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "{}", - older - ))), - ..Default::default() - }, - ChatMessage { - role: Role::Tool, - content: Some(MessageContent::String("tool output".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "{}", - newer - ))), - ..Default::default() - }, + Message::new( + Role::Assistant, + format!("{}", older), + ), + Message::new(Role::Tool, "tool output".to_string()), + Message::new( + Role::Assistant, + format!("{}", newer), + ), ]; assert_eq!(extract_last_checkpoint_id(&messages), Some(newer)); @@ -2027,11 +2023,7 @@ mod tests { #[test] fn extract_last_checkpoint_id_returns_none_without_tag() { - let messages = vec![ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("no checkpoint".to_string())), - ..Default::default() - }]; + let messages = vec![Message::new(Role::Assistant, "no checkpoint".to_string())]; assert_eq!(extract_last_checkpoint_id(&messages), None); } diff --git a/cli/src/commands/agent/run/pause.rs b/cli/src/commands/agent/run/pause.rs index 3136e07e3..2b44011fe 100644 --- a/cli/src/commands/agent/run/pause.rs +++ b/cli/src/commands/agent/run/pause.rs @@ -1,6 +1,6 @@ +use stakai::{ContentPart, Message, Role, ToolCall}; use stakpak_shared::local_store::LocalStore; use stakpak_shared::models::async_manifest::{AsyncManifest, PauseReason}; -use stakpak_shared::models::integrations::openai::{ChatMessage, Role, ToolCall}; use std::collections::HashSet; /// Exit code indicating the agent has paused and needs input or approval to resume. @@ -67,13 +67,38 @@ impl ResumeInput { /// Detect pending tool calls from checkpoint messages. /// Returns tool calls from the last assistant message that don't have corresponding tool results. -pub fn detect_pending_tool_calls(messages: &[ChatMessage]) -> Vec { +pub fn detect_pending_tool_calls(messages: &[Message]) -> Vec { // Find the last assistant message with tool_calls let tool_calls = messages .iter() .rev() - .find(|msg| msg.role == Role::Assistant && msg.tool_calls.is_some()) - .and_then(|msg| msg.tool_calls.as_ref()); + .find(|msg| { + msg.role == Role::Assistant + && msg + .parts() + .iter() + .any(|part| matches!(part, ContentPart::ToolCall { .. })) + }) + .map(|msg| { + msg.parts() + .into_iter() + .filter_map(|part| match part { + ContentPart::ToolCall { + id, + name, + arguments, + metadata, + .. + } => Some(ToolCall { + id, + name, + arguments, + metadata, + }), + _ => None, + }) + .collect::>() + }); let Some(tool_calls) = tool_calls else { return Vec::new(); @@ -83,14 +108,17 @@ pub fn detect_pending_tool_calls(messages: &[ChatMessage]) -> Vec { let executed_ids: HashSet = messages .iter() .filter(|msg| msg.role == Role::Tool) - .filter_map(|msg| msg.tool_call_id.clone()) + .flat_map(|msg| msg.parts()) + .filter_map(|part| match part { + ContentPart::ToolResult { tool_call_id, .. } => Some(tool_call_id), + _ => None, + }) .collect(); // Return tool calls without results tool_calls - .iter() + .into_iter() .filter(|tc| !executed_ids.contains(&tc.id)) - .cloned() .collect() } diff --git a/cli/src/commands/agent/run/renderer.rs b/cli/src/commands/agent/run/renderer.rs index 975a3e2ce..812628813 100644 --- a/cli/src/commands/agent/run/renderer.rs +++ b/cli/src/commands/agent/run/renderer.rs @@ -1,7 +1,7 @@ use crossterm::style::Stylize; use serde_json::Value; use stakpak_api::storage::{SessionStats, ToolUsageStats}; -use stakpak_shared::models::{integrations::openai::ChatMessage, llm::LLMTokenUsage}; +use stakpak_shared::models::llm::LLMTokenUsage; use std::fmt; use crate::utils::cli_colors::crossterm_colors; @@ -237,16 +237,17 @@ impl OutputRenderer { } } - pub fn render_final_completion(&self, messages: &[ChatMessage]) -> String { + pub fn render_final_completion(&self, messages: &[stakai::Message]) -> String { match self.format { OutputFormat::Json => { if self.verbose { serde_json::to_string_pretty(messages).unwrap_or_default() } else { // Find the last assistant message - let final_message = messages.iter().rev().find(|m| { - m.role == stakpak_shared::models::integrations::openai::Role::Assistant - }); + let final_message = messages + .iter() + .rev() + .find(|m| m.role == stakai::Role::Assistant); if let Some(message) = final_message { serde_json::to_string_pretty(message).unwrap_or_default() @@ -260,11 +261,12 @@ impl OutputRenderer { let mut output = String::new(); // Show final assistant message - if let Some(final_message) = messages.iter().rev().find(|m| { - m.role == stakpak_shared::models::integrations::openai::Role::Assistant - }) && let Some(content) = &final_message.content + if let Some(final_message) = messages + .iter() + .rev() + .find(|m| m.role == stakai::Role::Assistant) { - let content_str = self.extract_content_string(content); + let content_str = self.extract_content_string(&final_message.content); if !content_str.trim().is_empty() { output.push_str(&self.format_final_assistant_message(&content_str)); } @@ -273,7 +275,7 @@ impl OutputRenderer { // } else { // // Show just the final assistant message content // if let Some(final_message) = messages.iter().rev().find(|m| { - // m.role == stakpak_shared::models::integrations::openai::Role::Assistant + // m.role == stakai::Role::Assistant // }) { // if let Some(content) = &final_message.content { // let content_str = self.extract_content_string(content); @@ -288,16 +290,15 @@ impl OutputRenderer { } } - fn extract_content_string( - &self, - content: &stakpak_shared::models::integrations::openai::MessageContent, - ) -> String { + fn extract_content_string(&self, content: &stakai::MessageContent) -> String { match content { - stakpak_shared::models::integrations::openai::MessageContent::String(s) => s.clone(), - stakpak_shared::models::integrations::openai::MessageContent::Array(parts) => parts + stakai::MessageContent::Text(s) => s.clone(), + stakai::MessageContent::Parts(parts) => parts .iter() - .filter_map(|part| part.text.as_ref()) - .map(|text| text.as_str()) + .filter_map(|part| match part { + stakai::ContentPart::Text { text, .. } => Some(text.as_str()), + _ => None, + }) .filter(|text| !text.starts_with("")) .collect::>() .join("\n"), diff --git a/cli/src/commands/agent/run/stream.rs b/cli/src/commands/agent/run/stream.rs index 10d3405de..dd6b5668b 100644 --- a/cli/src/commands/agent/run/stream.rs +++ b/cli/src/commands/agent/run/stream.rs @@ -1,26 +1,14 @@ use crate::commands::agent::run::tui::send_input_event; use futures_util::{Stream, StreamExt}; -use stakai::Model; -use stakpak_api::models::ApiStreamError; -use stakpak_shared::models::{ - integrations::openai::{ - ChatCompletionChoice, ChatCompletionResponse, ChatCompletionStreamResponse, ChatMessage, - FinishReason, FunctionCall, MessageContent, Role, ToolCall, ToolCallDelta, - ToolCallStreamInfo, - }, - llm::LLMTokenUsage, +use stakai::{ContentPart, Message, MessageContent, Model, StreamEvent, ToolCall}; +use stakpak_api::models::{ + AgentStreamEvent, ApiStreamError, CompletionResponse, stakai_usage_to_llm_usage, }; +use stakpak_shared::models::agent_runtime::ToolCallStreamInfo; use stakpak_tui::{InputEvent, LoadingOperation}; use uuid::Uuid; -/// Accumulates streaming tool call deltas into complete tool calls. -/// -/// Handles two different streaming behaviors: -/// - **ID-based matching**: When delta has an ID, match by ID only (used by Anthropic/StakAI) -/// - **Index-based matching**: When delta has no ID, fall back to index (used by OpenAI) -/// -/// This distinction is important because some providers (like Anthropic via StakAI adapter) -/// send multiple tool calls with the same index but different IDs. +/// Accumulates streaming tool call events into complete tool calls. pub struct ToolCallAccumulator { tool_calls: Vec, } @@ -32,72 +20,70 @@ impl ToolCallAccumulator { } } - /// Process a tool call delta and accumulate it into the appropriate tool call. - pub fn process_delta(&mut self, delta: &ToolCallDelta) { - let delta_id = delta.id.as_deref().filter(|id| !id.is_empty()); - let delta_func = delta.function.as_ref(); + /// Process a StakAI tool call stream event. + pub fn process_event(&mut self, event: &StreamEvent) { + let (id, event_name, event_arguments, event_delta, event_metadata) = match event { + StreamEvent::ToolCallStart { id, name } => (id.as_str(), Some(name), None, None, None), + StreamEvent::ToolCallDelta { id, delta } => { + (id.as_str(), None, None, Some(delta), None) + } + StreamEvent::ToolCallEnd { + id, + name, + arguments, + metadata, + } => ( + id.as_str(), + Some(name), + Some(arguments), + None, + metadata.as_ref(), + ), + _ => return, + }; - match self.find_tool_call(delta_id, delta.index) { + match self.find_tool_call(id) { Some(tool_call) => { - // Update existing tool call - if let Some(func) = delta_func { - if let Some(name) = func.name.as_deref() - && tool_call.function.name.is_empty() - { - tool_call.function.name = name.to_string(); - } - if let Some(args) = &func.arguments { - tool_call.function.arguments.push_str(args); + if let Some(name) = event_name + && tool_call.name.is_empty() + { + tool_call.name = name.to_string(); + } + if let Some(arguments) = event_arguments { + tool_call.arguments = arguments.clone(); + } + if let Some(delta) = event_delta { + let mut raw = tool_call + .arguments + .as_str() + .map(String::from) + .unwrap_or_else(|| tool_call.arguments.to_string()); + if raw == "null" { + raw.clear(); } + raw.push_str(delta); + tool_call.arguments = serde_json::Value::String(raw); } - // Merge metadata (later deltas can carry metadata, e.g. ToolCallEnd) - if delta.metadata.is_some() { - tool_call.metadata = delta.metadata.clone(); + if let Some(metadata) = event_metadata { + tool_call.metadata = Some(metadata.clone()); } } None => { - // Create new tool call - self.create_tool_call(delta); + self.tool_calls.push(ToolCall { + id: id.to_string(), + name: event_name.cloned().unwrap_or_default(), + arguments: event_arguments + .cloned() + .or_else(|| event_delta.cloned().map(serde_json::Value::String)) + .unwrap_or_else(|| serde_json::Value::String(String::new())), + metadata: event_metadata.cloned(), + }); } } } - /// Find an existing tool call by ID or index. - /// Returns None if a new tool call should be created. - fn find_tool_call(&mut self, id: Option<&str>, index: usize) -> Option<&mut ToolCall> { - match id { - // Has ID: only match by ID, never fall back to index - Some(id) => self.tool_calls.iter_mut().find(|tc| tc.id == id), - // No ID: fall back to index-based matching for backwards compatibility - None => self.tool_calls.get_mut(index), - } - } - - /// Create a new tool call from a delta. - fn create_tool_call(&mut self, delta: &ToolCallDelta) { - // Pad with empty tool calls if needed (for sparse indices) - while self.tool_calls.len() < delta.index { - self.tool_calls.push(ToolCall { - id: String::new(), - r#type: "function".to_string(), - function: FunctionCall { - name: String::new(), - arguments: String::new(), - }, - metadata: None, - }); - } - - let func = delta.function.as_ref(); - self.tool_calls.push(ToolCall { - id: delta.id.clone().unwrap_or_default(), - r#type: "function".to_string(), - function: FunctionCall { - name: func.and_then(|f| f.name.clone()).unwrap_or_default(), - arguments: func.and_then(|f| f.arguments.clone()).unwrap_or_default(), - }, - metadata: delta.metadata.clone(), - }); + fn find_tool_call(&mut self, id: &str) -> Option<&mut ToolCall> { + self.tool_calls.iter_mut().find(|tc| tc.id == id) } /// Get the accumulated tool calls, filtering out empty placeholders. @@ -116,13 +102,16 @@ impl ToolCallAccumulator { .filter(|tc| !tc.id.is_empty()) .map(|tc| { // Best-effort: try to extract "description" from partial JSON args - let description = serde_json::from_str::(&tc.function.arguments) - .ok() - .and_then(|v| v.get("description")?.as_str().map(String::from)); + let description = tc + .arguments + .get("description") + .and_then(|v| v.as_str()) + .map(String::from); + let args_len = tc.arguments.to_string().len(); ToolCallStreamInfo { - name: tc.function.name.clone(), - args_tokens: tc.function.arguments.len() / 4, // rough chars-to-tokens estimate + name: tc.name.clone(), + args_tokens: args_len / 4, // rough chars-to-tokens estimate description, } }) @@ -131,39 +120,22 @@ impl ToolCallAccumulator { } pub async fn process_responses_stream( - stream: impl Stream>, + stream: impl Stream>, input_tx: &tokio::sync::mpsc::Sender, -) -> Result { +) -> Result { let mut stream = Box::pin(stream); - let mut chat_completion_response = ChatCompletionResponse { + let mut completion_response = CompletionResponse { id: "".to_string(), - object: "".to_string(), created: 0, model: "".to_string(), - choices: vec![], - usage: LLMTokenUsage { - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - prompt_tokens_details: None, - }, - system_fingerprint: None, + message: Message::new(stakai::Role::Assistant, ""), + usage: stakai::Usage::default(), metadata: None, }; - let mut response_metadata: Option = None; let mut current_model: Option = None; - - let mut chat_message = ChatMessage { - role: Role::Assistant, - content: None, - name: None, - tool_calls: None, - tool_call_id: None, - usage: None, - ..Default::default() - }; + let mut text = String::new(); let message_id = Uuid::new_v4(); let mut tool_call_accumulator = ToolCallAccumulator::new(); @@ -176,506 +148,187 @@ pub async fn process_responses_stream( while let Some(response) = stream.next().await { match &response { - Ok(response) => { - // Handle usage first - it can come in any event, including those with no content - if let Some(usage) = &response.usage { - chat_completion_response.usage = usage.clone(); - - // Send usage to TUI for display immediately when we receive it - send_input_event(input_tx, InputEvent::StreamUsage(usage.clone())).await?; - } - if let Some(metadata) = &response.metadata { - response_metadata = Some(metadata.clone()); - } - - // Skip chunks with no choices (e.g., usage-only events) - if response.choices.is_empty() { - continue; - } - - let delta = &response.choices[0].delta; - if !response.model.is_empty() { - // Look up the model in the catalog to get proper display name - // Use false for use_stakpak since the response already has the resolved model ID - let model = stakpak_api::find_model(&response.model, false) - .unwrap_or_else(|| Model::custom(response.model.clone(), "unknown")); - - // Only send event if model changed + Ok(response) => match response { + AgentStreamEvent::Model(model) => { let should_send = match ¤t_model { Some(existing) => existing.id != model.id, None => true, }; - if should_send { current_model = Some(model.clone()); - send_input_event(input_tx, InputEvent::StreamModel(model)).await?; + completion_response.model = model.id.clone(); + send_input_event(input_tx, InputEvent::StreamModel(model.clone())).await?; } } - - chat_completion_response = ChatCompletionResponse { - id: response.id.clone(), - object: response.object.clone(), - created: response.created, - model: current_model - .as_ref() - .map(|m| m.id.clone()) - .unwrap_or_default(), - choices: vec![], - usage: chat_completion_response.usage.clone(), - system_fingerprint: None, - metadata: None, - }; - - if let Some(content) = &delta.content { - chat_message.content = - Some(MessageContent::String(match chat_message.content { - Some(MessageContent::String(old_content)) => old_content + content, - _ => content.clone(), - })); - - send_input_event( - input_tx, - InputEvent::StreamAssistantMessage(message_id, content.clone()), - ) - .await?; + AgentStreamEvent::Metadata(metadata) => { + completion_response.metadata = Some(metadata.clone()); } - - if let Some(tool_calls) = &delta.tool_calls { - for delta_tool_call in tool_calls { - tool_call_accumulator.process_delta(delta_tool_call); + AgentStreamEvent::Event(event) => match event { + StreamEvent::Start { id } => { + completion_response.id = id.clone(); } - // Send streaming progress to TUI so users see what's being generated - let snapshot = tool_call_accumulator.progress_snapshot(); - if !snapshot.is_empty() { - send_input_event(input_tx, InputEvent::StreamToolCallProgress(snapshot)) + StreamEvent::TextDelta { delta, .. } => { + text.push_str(delta); + send_input_event( + input_tx, + InputEvent::StreamAssistantMessage(message_id, delta.clone()), + ) + .await?; + } + StreamEvent::ToolCallStart { .. } + | StreamEvent::ToolCallDelta { .. } + | StreamEvent::ToolCallEnd { .. } => { + tool_call_accumulator.process_event(event); + let snapshot = tool_call_accumulator.progress_snapshot(); + if !snapshot.is_empty() { + send_input_event( + input_tx, + InputEvent::StreamToolCallProgress(snapshot), + ) .await?; + } } - } - } + StreamEvent::Finish { usage, .. } => { + completion_response.usage = usage.clone(); + send_input_event( + input_tx, + InputEvent::StreamUsage(stakai_usage_to_llm_usage(usage)), + ) + .await?; + } + StreamEvent::ReasoningDelta { .. } => {} + StreamEvent::Error { message } => { + let _ = send_input_event( + input_tx, + InputEvent::EndLoadingOperation(LoadingOperation::StreamProcessing), + ) + .await; + return Err(ApiStreamError::Unknown(message.clone())); + } + }, + }, Err(e) => { - // End stream processing loading when error occurs - let _ = send_input_event( + send_input_event( input_tx, InputEvent::EndLoadingOperation(LoadingOperation::StreamProcessing), ) - .await; + .await?; return Err(e.clone()); } } } - // Get accumulated tool calls (already filtered for empty IDs) let final_tool_calls = tool_call_accumulator.into_tool_calls(); - chat_message.tool_calls = if final_tool_calls.is_empty() { - None + completion_response.message = if final_tool_calls.is_empty() { + Message::new(stakai::Role::Assistant, text) } else { - Some(final_tool_calls) + let mut parts = Vec::new(); + if !text.is_empty() { + parts.push(ContentPart::text(text)); + } + parts.extend(final_tool_calls.into_iter().map(|tool_call| { + ContentPart::tool_call(tool_call.id, tool_call.name, tool_call.arguments) + })); + Message::new(stakai::Role::Assistant, MessageContent::Parts(parts)) }; - chat_completion_response.choices.push(ChatCompletionChoice { - index: 0, - message: chat_message.clone(), - finish_reason: FinishReason::Stop, - logprobs: None, - }); - chat_completion_response.metadata = response_metadata; - - // End stream processing loading when stream completes send_input_event( input_tx, InputEvent::EndLoadingOperation(LoadingOperation::StreamProcessing), ) .await?; - Ok(chat_completion_response) + Ok(completion_response) } #[cfg(test)] mod tests { use super::*; use futures_util::stream; - use stakpak_shared::models::integrations::openai::{ - ChatCompletionStreamChoice, ChatCompletionStreamResponse, ChatMessageDelta, - FunctionCallDelta, - }; - - fn create_stream_response_with_tool_call( - index: usize, - id: Option, - name: Option, - arguments: Option, - ) -> ChatCompletionStreamResponse { - ChatCompletionStreamResponse { - id: "test".to_string(), - object: "chat.completion.chunk".to_string(), - created: 0, - model: "test-model".to_string(), - choices: vec![ChatCompletionStreamChoice { - index: 0, - delta: ChatMessageDelta { - role: None, - content: None, - tool_calls: Some(vec![ToolCallDelta { - index, - id, - r#type: Some("function".to_string()), - function: Some(FunctionCallDelta { name, arguments }), - metadata: None, - }]), - }, - finish_reason: None, - }], - usage: None, - metadata: None, - } - } - - fn create_content_response(content: &str) -> ChatCompletionStreamResponse { - ChatCompletionStreamResponse { - id: "test".to_string(), - object: "chat.completion.chunk".to_string(), - created: 0, - model: "test-model".to_string(), - choices: vec![ChatCompletionStreamChoice { - index: 0, - delta: ChatMessageDelta { - role: None, - content: Some(content.to_string()), - tool_calls: None, - }, - finish_reason: None, - }], - usage: None, - metadata: None, - } - } - fn create_usage_only_response() -> ChatCompletionStreamResponse { - ChatCompletionStreamResponse { - id: "test".to_string(), - object: "chat.completion.chunk".to_string(), - created: 0, - model: "".to_string(), - choices: vec![], // Empty choices — usage-only event - usage: Some(LLMTokenUsage { - prompt_tokens: 100, - completion_tokens: 50, - total_tokens: 150, - prompt_tokens_details: None, - }), - metadata: None, - } + fn event(event: StreamEvent) -> Result { + Ok(AgentStreamEvent::Event(event)) } #[tokio::test] - async fn test_empty_choices_does_not_panic() { - // This is the exact scenario that caused the index-out-of-bounds panic: - // Some providers send a final event with usage data but no choices. + async fn accumulates_text_and_usage_from_stakai_events() { let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - let responses = vec![ - Ok(create_content_response("Hello")), - Ok(create_usage_only_response()), // Empty choices — was panicking + Ok(AgentStreamEvent::Model(Model::custom("test-model", "test"))), + event(StreamEvent::Start { + id: "gen-1".to_string(), + }), + event(StreamEvent::TextDelta { + id: "gen-1".to_string(), + delta: "Hello".to_string(), + }), + event(StreamEvent::TextDelta { + id: "gen-1".to_string(), + delta: " world".to_string(), + }), + event(StreamEvent::Finish { + usage: stakai::Usage::new(100, 50), + reason: stakai::FinishReason::stop(), + }), ]; let test_stream = stream::iter(responses); tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - // Content should still be accumulated from the first chunk - let content = response.choices[0] - .message - .content - .as_ref() - .unwrap() - .to_string(); - assert_eq!(content, "Hello"); - // Usage from the empty-choices event should still be captured + let response = process_responses_stream(test_stream, &input_tx) + .await + .unwrap(); + assert_eq!(response.id, "gen-1"); + assert_eq!(response.model, "test-model"); + assert_eq!(response.message.content.text().unwrap(), "Hello world"); assert_eq!(response.usage.prompt_tokens, 100); assert_eq!(response.usage.completion_tokens, 50); - assert_eq!(response.usage.total_tokens, 150); - } - - #[tokio::test] - async fn test_only_usage_events_no_content() { - // Stream with only usage events and no content at all - let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - - let responses = vec![Ok(create_usage_only_response())]; - - let test_stream = stream::iter(responses); - tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - assert_eq!(response.choices.len(), 1); - assert!(response.choices[0].message.content.is_none()); - assert!(response.choices[0].message.tool_calls.is_none()); - assert_eq!(response.usage.total_tokens, 150); - } - - #[tokio::test] - async fn test_multiple_empty_choices_interspersed() { - // Multiple empty-choices events interspersed with content - let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - - let responses = vec![ - Ok(create_usage_only_response()), - Ok(create_content_response("Hello")), - Ok(create_usage_only_response()), - Ok(create_content_response(" World")), - Ok(create_usage_only_response()), - ]; - - let test_stream = stream::iter(responses); - tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - let content = response.choices[0] - .message - .content - .as_ref() - .unwrap() - .to_string(); - assert_eq!(content, "Hello World"); - } - - #[tokio::test] - async fn test_empty_stream() { - // Completely empty stream — no events at all - let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - - let responses: Vec> = vec![]; - - let test_stream = stream::iter(responses); - tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - // Should have one choice with empty content - assert_eq!(response.choices.len(), 1); - assert!(response.choices[0].message.content.is_none()); - assert!(response.choices[0].message.tool_calls.is_none()); } #[tokio::test] - async fn test_content_accumulation_across_chunks() { + async fn accumulates_tool_call_events_by_id() { let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - let responses = vec![ - Ok(create_content_response("Hello")), - Ok(create_content_response(", ")), - Ok(create_content_response("world")), - Ok(create_content_response("!")), + event(StreamEvent::ToolCallStart { + id: "tool-1".to_string(), + name: "my_function".to_string(), + }), + event(StreamEvent::ToolCallDelta { + id: "tool-1".to_string(), + delta: "{\"key\":".to_string(), + }), + event(StreamEvent::ToolCallDelta { + id: "tool-1".to_string(), + delta: "\"value\"}".to_string(), + }), ]; let test_stream = stream::iter(responses); tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - let content = response.choices[0] + let response = process_responses_stream(test_stream, &input_tx) + .await + .unwrap(); + let tool_calls: Vec<_> = response .message - .content - .as_ref() - .unwrap() - .to_string(); - assert_eq!(content, "Hello, world!"); - } - - #[tokio::test] - async fn test_stream_error_propagated() { - let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - - let responses: Vec> = vec![ - Ok(create_content_response("start")), - Err(ApiStreamError::Unknown("connection lost".to_string())), - ]; - - let test_stream = stream::iter(responses); - tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_err()); - } - - #[tokio::test] - async fn test_consecutive_tool_calls_with_same_index_different_ids() { - // This test verifies the bug fix: consecutive tool calls with same index - // but different IDs should create separate tool calls, not merge arguments - let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - - // Simulate two tool calls with index 0 but different IDs (Stakai behavior) - let responses = vec![ - Ok(create_stream_response_with_tool_call( - 0, - Some("tool-1".to_string()), - Some("function_a".to_string()), - Some("{\"arg\":".to_string()), - )), - Ok(create_stream_response_with_tool_call( - 0, - Some("tool-1".to_string()), - None, - Some("\"value1\"}".to_string()), - )), - Ok(create_stream_response_with_tool_call( - 0, - Some("tool-2".to_string()), - Some("function_b".to_string()), - Some("{\"arg\":".to_string()), - )), - Ok(create_stream_response_with_tool_call( - 0, - Some("tool-2".to_string()), - None, - Some("\"value2\"}".to_string()), - )), - ]; - - let test_stream = stream::iter(responses); - - // Spawn a task to drain the receiver so it doesn't block - tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - let tool_calls = response.choices[0].message.tool_calls.as_ref().unwrap(); - - // Should have 2 separate tool calls, not 1 with merged arguments - assert_eq!(tool_calls.len(), 2, "Should have 2 separate tool calls"); - assert_eq!(tool_calls[0].id, "tool-1"); - assert_eq!(tool_calls[0].function.name, "function_a"); - assert_eq!(tool_calls[0].function.arguments, "{\"arg\":\"value1\"}"); - assert_eq!(tool_calls[1].id, "tool-2"); - assert_eq!(tool_calls[1].function.name, "function_b"); - assert_eq!(tool_calls[1].function.arguments, "{\"arg\":\"value2\"}"); - } - - #[tokio::test] - async fn test_tool_calls_fallback_to_index_when_no_id() { - // This test verifies backwards compatibility: when no ID is provided, - // tool calls should be matched by index - let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - - // Simulate tool calls using index-based matching (OpenAI behavior) - let responses = vec![ - Ok(create_stream_response_with_tool_call( - 0, - Some("tool-1".to_string()), - Some("function_a".to_string()), - Some("{\"arg\":".to_string()), - )), - Ok(create_stream_response_with_tool_call( - 0, - None, // No ID, should fall back to index - None, - Some("\"value1\"}".to_string()), - )), - ]; - - let test_stream = stream::iter(responses); - - tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - let tool_calls = response.choices[0].message.tool_calls.as_ref().unwrap(); - + .parts() + .into_iter() + .filter_map(|part| match part { + ContentPart::ToolCall { + id, + name, + arguments, + .. + } => Some((id, name, arguments)), + _ => None, + }) + .collect(); assert_eq!(tool_calls.len(), 1); - assert_eq!(tool_calls[0].id, "tool-1"); - assert_eq!(tool_calls[0].function.arguments, "{\"arg\":\"value1\"}"); - } - - #[tokio::test] - async fn test_tool_calls_with_incrementing_indices() { - // Test standard behavior with incrementing indices - let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - - let responses = vec![ - Ok(create_stream_response_with_tool_call( - 0, - Some("tool-1".to_string()), - Some("func_a".to_string()), - Some("{\"a\":1}".to_string()), - )), - Ok(create_stream_response_with_tool_call( - 1, - Some("tool-2".to_string()), - Some("func_b".to_string()), - Some("{\"b\":2}".to_string()), - )), - ]; - - let test_stream = stream::iter(responses); - - tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - let tool_calls = response.choices[0].message.tool_calls.as_ref().unwrap(); - - assert_eq!(tool_calls.len(), 2); - assert_eq!(tool_calls[0].id, "tool-1"); - assert_eq!(tool_calls[0].function.name, "func_a"); - assert_eq!(tool_calls[1].id, "tool-2"); - assert_eq!(tool_calls[1].function.name, "func_b"); - } - - #[tokio::test] - async fn test_tool_call_name_and_arguments_in_separate_chunks() { - // Test that name and arguments coming in separate chunks are merged correctly - let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); - - let responses = vec![ - // First chunk: ID and name only - Ok(create_stream_response_with_tool_call( - 0, - Some("tool-1".to_string()), - Some("my_function".to_string()), - None, - )), - // Second chunk: same ID, arguments only - Ok(create_stream_response_with_tool_call( - 0, - Some("tool-1".to_string()), - None, - Some("{\"key\":\"value\"}".to_string()), - )), - ]; - - let test_stream = stream::iter(responses); - - tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); - - let result = process_responses_stream(test_stream, &input_tx).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - let tool_calls = response.choices[0].message.tool_calls.as_ref().unwrap(); - - // Should have 1 tool call with both name and arguments - assert_eq!(tool_calls.len(), 1, "Should have 1 merged tool call"); - assert_eq!(tool_calls[0].id, "tool-1"); - assert_eq!(tool_calls[0].function.name, "my_function"); - assert_eq!(tool_calls[0].function.arguments, "{\"key\":\"value\"}"); + assert_eq!(tool_calls[0].0, "tool-1"); + assert_eq!(tool_calls[0].1, "my_function"); + assert_eq!( + tool_calls[0].2, + serde_json::Value::String("{\"key\":\"value\"}".to_string()) + ); } } diff --git a/cli/src/commands/agent/run/tooling.rs b/cli/src/commands/agent/run/tooling.rs index 84802b278..a5abbad2e 100644 --- a/cli/src/commands/agent/run/tooling.rs +++ b/cli/src/commands/agent/run/tooling.rs @@ -6,7 +6,6 @@ use stakpak_api::AgentProvider; use stakpak_api::storage::ListSessionsQuery; use stakpak_mcp_client::McpClient; use stakpak_shared::models::integrations::mcp::CallToolResultExt; -use stakpak_shared::models::integrations::openai::ToolCall; use stakpak_tui::SessionInfo; use uuid::Uuid; @@ -36,28 +35,18 @@ pub async fn list_sessions(client: &dyn AgentProvider) -> Result>, session_id: Option, model_id: Option, model_provider: Option, ) -> Result, String> { - let tool_name = &tool_call.function.name; + let tool_name = &tool_call.name; let tool_exists = tools.iter().any(|tool| tool.name == *tool_name); if tool_exists { // Parse arguments safely - let arguments = match serde_json::from_str(&tool_call.function.arguments) { - Ok(args) => Some(args), - Err(e) => { - let error_msg = format!("Failed to parse tool arguments as JSON: {}", e); - log::error!("{}", error_msg); - return Ok(Some(CallToolResult::error(vec![ - rmcp::model::Content::text("INVALID_ARGUMENTS"), - rmcp::model::Content::text(error_msg), - ]))); - } - }; + let arguments = tool_call.arguments.as_object().cloned(); // Call tool and handle errors gracefully let metadata = Some({ diff --git a/cli/src/commands/agent/run/tui.rs b/cli/src/commands/agent/run/tui.rs index 9169671ef..c2ec5c0fb 100644 --- a/cli/src/commands/agent/run/tui.rs +++ b/cli/src/commands/agent/run/tui.rs @@ -1,4 +1,3 @@ -use stakpak_shared::models::integrations::openai::ToolCall; use stakpak_tui::InputEvent; pub async fn send_input_event( @@ -10,7 +9,7 @@ pub async fn send_input_event( pub async fn send_tool_call( input_tx: &tokio::sync::mpsc::Sender, - tool_call: &ToolCall, + tool_call: &stakai::ToolCall, ) -> Result<(), String> { send_input_event(input_tx, InputEvent::RunToolCall(tool_call.clone())).await?; Ok(()) diff --git a/cli/src/commands/sessions/messages.rs b/cli/src/commands/sessions/messages.rs index b2636d13e..0b956686d 100644 --- a/cli/src/commands/sessions/messages.rs +++ b/cli/src/commands/sessions/messages.rs @@ -2,7 +2,7 @@ //! //! Applies `--role`, `--limit`, and `--offset` flags to a checkpoint's messages. -use stakpak_shared::models::integrations::openai::{ChatMessage, Role}; +use stakai::{Message, Role}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RoleFilter { @@ -18,7 +18,7 @@ impl RoleFilter { RoleFilter::User => matches!(role, Role::User), RoleFilter::Assistant => matches!(role, Role::Assistant), RoleFilter::Tool => matches!(role, Role::Tool), - RoleFilter::System => matches!(role, Role::System | Role::Developer), + RoleFilter::System => matches!(role, Role::System), } } } @@ -45,12 +45,12 @@ impl std::str::FromStr for RoleFilter { /// Ordering: filter by role first, then compute a chronological window anchored at /// the newest end of the filtered list. pub fn filter_messages( - messages: Vec, + messages: Vec, role: Option, limit: Option, offset: u32, -) -> (Vec, u32) { - let filtered: Vec = match role { +) -> (Vec, u32) { + let filtered: Vec = match role { Some(filter) => messages .into_iter() .filter(|m| filter.matches(&m.role)) @@ -81,17 +81,12 @@ pub fn filter_messages( #[cfg(test)] mod tests { use super::*; - use stakpak_shared::models::integrations::openai::MessageContent; - fn msg(role: Role, text: &str) -> ChatMessage { - ChatMessage { - role, - content: Some(MessageContent::String(text.to_string())), - ..Default::default() - } + fn msg(role: Role, text: &str) -> Message { + Message::new(role, text.to_string()) } - fn sample() -> Vec { + fn sample() -> Vec { vec![ msg(Role::System, "sys"), msg(Role::User, "u1"), @@ -102,13 +97,10 @@ mod tests { ] } - fn contents(messages: &[ChatMessage]) -> Vec<&str> { + fn contents(messages: &[Message]) -> Vec { messages .iter() - .map(|message| match &message.content { - Some(MessageContent::String(content)) => content.as_str(), - other => panic!("expected string content, got {other:?}"), - }) + .map(|message| message.content.text().expect("expected text content")) .collect() } @@ -127,12 +119,11 @@ mod tests { } #[test] - fn role_filter_system_includes_developer_and_reports_total() { - let mut msgs = sample(); - msgs.push(msg(Role::Developer, "dev")); + fn role_filter_system_keeps_only_system_messages_and_reports_total() { + let msgs = sample(); let (out, total) = filter_messages(msgs, Some(RoleFilter::System), None, 0); - assert_eq!(contents(&out), vec!["sys", "dev"]); - assert_eq!(total, 2); + assert_eq!(contents(&out), vec!["sys"]); + assert_eq!(total, 1); } #[test] diff --git a/cli/src/commands/sessions/output.rs b/cli/src/commands/sessions/output.rs index d7a5f5ca4..aebcd50f2 100644 --- a/cli/src/commands/sessions/output.rs +++ b/cli/src/commands/sessions/output.rs @@ -4,7 +4,6 @@ use serde::Serialize; use stakpak_api::{ BackendInfo, BackendKind, Session, SessionStatus, SessionSummary, SessionVisibility, }; -use stakpak_shared::models::integrations::openai::ChatMessage; use stakpak_shared::utils::sanitize_text_output; /// Output format for session commands. @@ -131,13 +130,13 @@ pub struct ShowOutput<'a> { pub updated_at: chrono::DateTime, pub active_checkpoint_id: Option, pub message_count: u32, - pub messages: &'a [ChatMessage], + pub messages: &'a [stakai::Message], pub backend: &'a BackendInfo, } pub fn render_show( session: &Session, - messages: &[ChatMessage], + messages: &[stakai::Message], backend: &BackendInfo, options: ShowRenderOptions<'_>, mode: OutputMode, @@ -158,7 +157,7 @@ pub fn render_show( fn render_show_json( session: &Session, - messages: &[ChatMessage], + messages: &[stakai::Message], message_count: u32, backend: &BackendInfo, ) -> String { @@ -180,7 +179,7 @@ fn render_show_json( fn render_show_human( session: &Session, - messages: &[ChatMessage], + messages: &[stakai::Message], message_count: u32, limit: Option, offset: u32, @@ -228,36 +227,53 @@ fn render_show_human( } for (i, m) in messages.iter().enumerate() { - out.push_str(&format!("\n[{}] {}\n", i + 1, m.role)); - if let Some(content) = &m.content { - let text = sanitize_text_output(&content.to_string()); + out.push_str(&format!("\n[{}] {}\n", i + 1, format_role(&m.role))); + if let Some(content) = m.content.text() { + let text = sanitize_text_output(&content); for line in text.lines() { out.push_str(" "); out.push_str(line); out.push('\n'); } } - if let Some(tool_calls) = &m.tool_calls { - for tc in tool_calls { - out.push_str(&format!( - " [tool_call] {} ({})\n", - sanitize_text_output(&tc.function.name), - sanitize_text_output(&tc.id), - )); - out.push_str(" arguments:\n"); - for line in format_tool_call_arguments(&tc.function.arguments).lines() { - out.push_str(" "); - out.push_str(line); - out.push('\n'); + for part in m.parts() { + match part { + stakai::ContentPart::ToolCall { + id, + name, + arguments, + .. + } => { + out.push_str(&format!( + " [tool_call] {} ({})\n", + sanitize_text_output(&name), + sanitize_text_output(&id), + )); + out.push_str(" arguments:\n"); + for line in format_tool_call_arguments(&arguments.to_string()).lines() { + out.push_str(" "); + out.push_str(line); + out.push('\n'); + } } + stakai::ContentPart::ToolResult { + tool_call_id, + content, + .. + } => { + out.push_str(&format!( + " [tool_call_id] {}\n", + sanitize_text_output(&tool_call_id) + )); + for line in format_tool_result_content(&content).lines() { + out.push_str(" "); + out.push_str(line); + out.push('\n'); + } + } + _ => {} } } - if let Some(tc_id) = &m.tool_call_id { - out.push_str(&format!( - " [tool_call_id] {}\n", - sanitize_text_output(tc_id) - )); - } } if let Some(limit) = limit @@ -331,6 +347,23 @@ fn format_tool_call_arguments(arguments: &str) -> String { sanitize_text_output(&rendered) } +fn format_tool_result_content(content: &serde_json::Value) -> String { + let rendered = match content { + serde_json::Value::String(text) => text.clone(), + value => serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string()), + }; + sanitize_text_output(&rendered) +} + +fn format_role(role: &stakai::Role) -> &'static str { + match role { + stakai::Role::System => "system", + stakai::Role::User => "user", + stakai::Role::Assistant => "assistant", + stakai::Role::Tool => "tool", + } +} + fn collapse_whitespace(s: &str) -> String { let mut out = String::with_capacity(s.len()); let mut last_was_space = false; diff --git a/cli/src/commands/sessions/tests.rs b/cli/src/commands/sessions/tests.rs index de988319a..e6db31ca0 100644 --- a/cli/src/commands/sessions/tests.rs +++ b/cli/src/commands/sessions/tests.rs @@ -4,12 +4,12 @@ use std::sync::Arc; +use stakai::{ContentPart, Message, MessageContent, Role}; use stakpak_api::{ BackendInfo, Checkpoint, CheckpointState, ListSessionsQuery, LocalStorage, Session, SessionStatus, SessionStorage, SessionVisibility, StorageCreateSessionRequest as CreateSessionRequest, StorageError, }; -use stakpak_shared::models::integrations::openai::{ChatMessage, MessageContent, Role}; use uuid::Uuid; use super::classify_storage_error; @@ -22,47 +22,36 @@ async fn in_memory_storage() -> LocalStorage { .expect("in-memory storage") } -fn msg(role: Role, text: &str) -> ChatMessage { - ChatMessage { - role, - content: Some(MessageContent::String(text.to_string())), - ..Default::default() - } +fn msg(role: Role, text: &str) -> Message { + Message::new(role, text.to_string()) } -fn assistant_tool_call_msg(name: &str, tool_call_id: &str) -> ChatMessage { +fn assistant_tool_call_msg(name: &str, tool_call_id: &str) -> Message { assistant_tool_call_msg_with_args(name, tool_call_id, "{}") } -fn assistant_tool_call_msg_with_args( - name: &str, - tool_call_id: &str, - arguments: &str, -) -> ChatMessage { - use stakpak_shared::models::integrations::openai::{FunctionCall, ToolCall}; - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("calling a tool".to_string())), - tool_calls: Some(vec![ToolCall { - id: tool_call_id.to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: name.to_string(), - arguments: arguments.to_string(), - }, - metadata: None, - }]), - ..Default::default() - } +fn assistant_tool_call_msg_with_args(name: &str, tool_call_id: &str, arguments: &str) -> Message { + Message::new( + Role::Assistant, + MessageContent::Parts(vec![ + ContentPart::text("calling a tool"), + ContentPart::tool_call( + tool_call_id, + name, + serde_json::from_str(arguments) + .unwrap_or_else(|_| serde_json::Value::String(arguments.to_string())), + ), + ]), + ) } -fn numbered_messages(count: u32) -> Vec { +fn numbered_messages(count: u32) -> Vec { (1..=count) .map(|index| msg(Role::User, &format!("m{index}"))) .collect() } -fn alternating_assistant_user_messages(assistant_count: u32) -> Vec { +fn alternating_assistant_user_messages(assistant_count: u32) -> Vec { let mut messages = Vec::with_capacity((assistant_count * 2) as usize); for index in 1..=assistant_count { messages.push(msg(Role::User, &format!("u{index}"))); @@ -85,7 +74,7 @@ fn render_list(sessions: &[stakpak_api::SessionSummary], mode: OutputMode) -> St fn render_show( session: &Session, - messages: &[ChatMessage], + messages: &[Message], message_count: u32, limit: Option, offset: u32, @@ -438,13 +427,13 @@ async fn sessions_show_json_roundtrip_preserves_tool_calls() { let value: serde_json::Value = serde_json::from_str(&rendered).unwrap(); let messages = value["messages"].as_array().unwrap(); let assistant = messages.iter().find(|m| m["role"] == "assistant").unwrap(); - let tool_calls = assistant["tool_calls"].as_array().unwrap(); - assert_eq!(tool_calls.len(), 1); - assert_eq!(tool_calls[0]["id"].as_str().unwrap(), "tc_1"); - assert_eq!( - tool_calls[0]["function"]["name"].as_str().unwrap(), - "run_command" - ); + let content = assistant["content"].as_array().unwrap(); + let tool_call = content + .iter() + .find(|part| part["type"] == "tool_call") + .unwrap(); + assert_eq!(tool_call["id"].as_str().unwrap(), "tc_1"); + assert_eq!(tool_call["name"].as_str().unwrap(), "run_command"); } #[tokio::test] @@ -869,12 +858,13 @@ fn render_show_human_with_empty_messages_and_checkpoint_shows_no_messages_marker #[test] fn render_show_human_renders_tool_call_id_line_for_tool_messages() { let session = fake_session(Uuid::new_v4(), true); - let tool_msg = ChatMessage { - role: Role::Tool, - content: Some(MessageContent::String("tool output".to_string())), - tool_call_id: Some("tc_42".to_string()), - ..Default::default() - }; + let tool_msg = Message::new( + Role::Tool, + MessageContent::Parts(vec![ContentPart::tool_result( + "tc_42", + serde_json::Value::String("tool output".to_string()), + )]), + ); let rendered = render_show(&session, &[tool_msg], 1, None, 0, None, OutputMode::Human); assert!(rendered.contains("[tool_call_id] tc_42")); assert!(rendered.contains("tool output")); @@ -912,7 +902,7 @@ fn render_show_human_renders_raw_non_json_tool_call_arguments() { let rendered = render_show(&session, &[asst], 1, Some(1), 0, None, OutputMode::Human); assert!(rendered.contains("[tool_call] run_cmd (tc_7)")); assert!(rendered.contains(" arguments:")); - assert!(rendered.contains(" not-json --flag raw")); + assert!(rendered.contains(" \"not-json --flag raw\"")); } // ============================================================================= @@ -921,29 +911,17 @@ fn render_show_human_renders_raw_non_json_tool_call_arguments() { #[test] fn render_show_human_strips_terminal_escape_sequences() { - use stakpak_shared::models::integrations::openai::{FunctionCall, ToolCall}; - let mut session = fake_session(Uuid::new_v4(), true); session.title = "evil\x1b[2Jtitle\x1b]52;c;YmFkYm95\x07".to_string(); session.cwd = Some("/tmp/\x1b[31mred\x1b[0m".to_string()); - let tool_msg = ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String( - "hi\x1b[2Jhidden\x1b]8;;https://evil.example/\x07click\x1b]8;;\x07".to_string(), - )), - tool_calls: Some(vec![ToolCall { - id: "tc_\x1b[Kx".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "run_\x1b[31mcmd".to_string(), - arguments: "{}".to_string(), - }, - metadata: None, - }]), - tool_call_id: None, - ..Default::default() - }; + let tool_msg = Message::new( + Role::Assistant, + MessageContent::Parts(vec![ + ContentPart::text("hi\x1b[2Jhidden\x1b]8;;https://evil.example/\x07click\x1b]8;;\x07"), + ContentPart::tool_call("tc_\x1b[Kx", "run_\x1b[31mcmd", serde_json::json!({})), + ]), + ); let rendered = render_show(&session, &[tool_msg], 1, None, 0, None, OutputMode::Human); assert!( @@ -984,11 +962,7 @@ async fn render_list_human_strips_terminal_escape_sequences_from_title() { #[test] fn render_show_json_preserves_raw_content_including_escapes() { let session = fake_session(Uuid::new_v4(), true); - let m = ChatMessage { - role: Role::User, - content: Some(MessageContent::String("hi\x1b[2Jhidden".to_string())), - ..Default::default() - }; + let m = Message::new(Role::User, "hi\x1b[2Jhidden".to_string()); let rendered = render_show(&session, &[m], 1, None, 0, None, OutputMode::Json); let v: serde_json::Value = serde_json::from_str(&rendered).unwrap(); assert_eq!( diff --git a/cli/src/utils/agent_context.rs b/cli/src/utils/agent_context.rs index 6ffc45d23..b31d97580 100644 --- a/cli/src/utils/agent_context.rs +++ b/cli/src/utils/agent_context.rs @@ -3,6 +3,9 @@ use crate::utils::apps_md::{AppsMdInfo, format_apps_md_for_context}; use crate::utils::local_context::LocalContext; use stakpak_api::models::Skill; +const INJECTED_CONTEXT_TAGS: &[&str] = + &["local_context", "available_skills", "agents_md", "apps_md"]; + #[derive(Debug, Clone)] pub struct AgentContext { /// Pre-formatted local context string. Snapshotted once at construction; @@ -95,6 +98,36 @@ fn format_skills(skills: &[Skill]) -> String { ) } +pub fn strip_injected_context_blocks(input: &str) -> String { + let mut stripped = input.to_string(); + for tag in INJECTED_CONTEXT_TAGS { + stripped = strip_xml_like_blocks(&stripped, tag); + } + + stripped.trim_end().to_string() +} + +fn strip_xml_like_blocks(input: &str, tag: &str) -> String { + let open = format!("<{tag}>"); + let close = format!(""); + let mut remaining = input; + let mut output = String::with_capacity(input.len()); + + while let Some(start) = remaining.find(&open) { + output.push_str(&remaining[..start]); + let after_open = start + open.len(); + let Some(relative_end) = remaining[after_open..].find(&close) else { + output.push_str(&remaining[start..]); + return output; + }; + let end = after_open + relative_end + close.len(); + remaining = &remaining[end..]; + } + + output.push_str(remaining); + output +} + #[cfg(test)] mod tests { use super::*; @@ -219,4 +252,18 @@ mod tests { ctx.update_skills(Some(make_skills())); assert!(ctx.skills.is_some()); } + + #[test] + fn strip_injected_context_blocks_removes_context_from_display_text() { + let input = "deploy this\n\nsecret cwd\n\n\nlarge skill list\n"; + + assert_eq!(strip_injected_context_blocks(input), "deploy this"); + } + + #[test] + fn strip_injected_context_blocks_preserves_malformed_user_text() { + let input = "show me literally"; + + assert_eq!(strip_injected_context_blocks(input), input); + } } diff --git a/libs/ai/docs/PROVIDER_OPTIONS_PLAN.md b/libs/ai/docs/PROVIDER_OPTIONS_PLAN.md index 79345ae50..35aa11bdc 100644 --- a/libs/ai/docs/PROVIDER_OPTIONS_PLAN.md +++ b/libs/ai/docs/PROVIDER_OPTIONS_PLAN.md @@ -138,30 +138,21 @@ pub enum HarmBlockThreshold { | File | Changes | |------|---------| -| `openai/types.rs` | Add new fields to ChatCompletionRequest | +| `openai/types.rs` | Add new fields to the provider request payload | | `openai/convert.rs` | Use new options in request conversion | | `gemini/types.rs` | Add safety_settings and cached_content to GeminiRequest | | `gemini/convert.rs` | Use safety_settings in conversion | -### Adapter Layer (`libs/shared/src/models/`) - -| File | Changes | -|------|---------| -| `stakai_adapter.rs` | Update option conversion (set new fields to None initially) | - ---- - ## Implementation Order 1. **Add types to `request.rs`** - Define all new structs and enums 2. **Export from `mod.rs`** - Make types publicly available -3. **Update OpenAI types** - Add fields to ChatCompletionRequest +3. **Update OpenAI types** - Add fields to the provider request payload 4. **Update OpenAI convert** - Use options in to_openai_request() 5. **Update Gemini types** - Add fields to GeminiRequest 6. **Update Gemini convert** - Use options in to_gemini_request() -7. **Update stakai_adapter** - Handle new options in conversion -8. **Run tests** - Ensure all existing tests pass -9. **Add new tests** - Test new option handling +7. **Run tests** - Ensure all existing tests pass +8. **Add new tests** - Test new option handling --- diff --git a/libs/ai/src/providers/gemini/convert.rs b/libs/ai/src/providers/gemini/convert.rs index f27199dde..1f750058d 100644 --- a/libs/ai/src/providers/gemini/convert.rs +++ b/libs/ai/src/providers/gemini/convert.rs @@ -465,8 +465,7 @@ mod tests { #[test] fn test_convert_messages_resolves_tool_result_names() { // Simulate the real-world flow: model proposes a tool call, then the - // tool result comes back as a separate message with string content - // (the way stakai_adapter.rs actually produces tool results). + // tool result comes back as a separate message with string content. let messages = vec![ Message::new(Role::User, "List files"), Message { diff --git a/libs/api/src/client/mod.rs b/libs/api/src/client/mod.rs index 8683ea732..e03de9941 100644 --- a/libs/api/src/client/mod.rs +++ b/libs/api/src/client/mod.rs @@ -7,6 +7,7 @@ //! - Integrates with hooks for lifecycle events mod provider; +mod stakai; use crate::local::hooks::task_board_context::{TaskBoardContextHook, TaskBoardContextHookOptions}; use crate::local::storage::LocalStorage; @@ -15,9 +16,9 @@ use crate::stakpak::storage::StakpakStorage; use crate::stakpak::{StakpakApiClient, StakpakApiConfig}; use crate::storage::SessionStorage; +use self::stakai::StakAIClient; use stakpak_shared::hooks::{HookRegistry, LifecycleEvent}; use stakpak_shared::models::llm::{LLMProviderConfig, ProviderConfig}; -use stakpak_shared::models::stakai_adapter::StakAIClient; use std::sync::Arc; // ============================================================================= diff --git a/libs/api/src/client/provider.rs b/libs/api/src/client/provider.rs index e93d78065..45eb4e4a3 100644 --- a/libs/api/src/client/provider.rs +++ b/libs/api/src/client/provider.rs @@ -13,18 +13,11 @@ use crate::storage::{ UpdateSessionRequest as StorageUpdateSessionRequest, }; use async_trait::async_trait; -use futures_util::Stream; +use futures_util::{Stream, StreamExt}; use reqwest::header::HeaderMap; use rmcp::model::Content; -use stakai::Model; +use stakai::{ContentPart, Message, Model, ResponseContent, StreamEvent}; use stakpak_shared::hooks::{HookContext, LifecycleEvent}; -use stakpak_shared::models::integrations::openai::{ - ChatCompletionChoice, ChatCompletionResponse, ChatCompletionStreamChoice, - ChatCompletionStreamResponse, ChatMessage, FinishReason, MessageContent, Role, Tool, -}; -use stakpak_shared::models::llm::{ - GenerationDelta, LLMInput, LLMMessage, LLMMessageContent, LLMStreamInput, -}; use std::pin::Pin; use tokio::sync::mpsc; use uuid::Uuid; @@ -45,7 +38,7 @@ use super::AgentClient; #[derive(Debug)] pub(crate) enum StreamMessage { - Delta(GenerationDelta), + Event(StreamEvent), Ctx(Box>), } @@ -178,11 +171,11 @@ impl AgentProvider for AgentClient { async fn chat_completion( &self, model: Model, - messages: Vec, - tools: Option>, + messages: Vec, + tools: Option>, session_id: Option, metadata: Option, - ) -> Result { + ) -> Result { let mut ctx = HookContext::new( session_id, AgentState::new(model, messages, tools, metadata), @@ -238,9 +231,8 @@ impl AgentProvider for AgentClient { meta.insert("state_metadata".to_string(), state_metadata.clone()); } - Ok(ChatCompletionResponse { + Ok(CompletionResponse { id: ctx.new_checkpoint_id.unwrap().to_string(), - object: "chat.completion".to_string(), created: checkpoint_created_at, model: ctx .state @@ -248,19 +240,18 @@ impl AgentProvider for AgentClient { .as_ref() .map(|llm_input| llm_input.model.id.clone()) .unwrap_or_default(), - choices: vec![ChatCompletionChoice { - index: 0, - message: ctx.state.messages.last().cloned().unwrap(), - logprobs: None, - finish_reason: FinishReason::Stop, - }], + message: ctx + .state + .messages + .last() + .cloned() + .unwrap_or_else(|| Message::new(stakai::Role::Assistant, "")), usage: ctx .state .llm_output .as_ref() .map(|u| u.usage.clone()) .unwrap_or_default(), - system_fingerprint: None, metadata: if meta.is_empty() { None } else { @@ -272,16 +263,14 @@ impl AgentProvider for AgentClient { async fn chat_completion_stream( &self, model: Model, - messages: Vec, - tools: Option>, + messages: Vec, + tools: Option>, _headers: Option, session_id: Option, metadata: Option, ) -> Result< ( - Pin< - Box> + Send>, - >, + Pin> + Send>>, Option, ), String, @@ -369,6 +358,8 @@ impl AgentProvider for AgentClient { let hook_registry = self.hook_registry.clone(); let stream = async_stream::stream! { + yield Ok(AgentStreamEvent::Model(ctx.state.active_model.clone())); + while let Some(delta_result) = rx.recv().await { match delta_result { Ok(delta) => match delta { @@ -384,38 +375,11 @@ impl AgentProvider for AgentClient { if let Some(state_metadata) = &ctx.state.metadata { meta.insert("state_metadata".to_string(), state_metadata.clone()); } - yield Ok(ChatCompletionStreamResponse { - id: ctx.request_id.to_string(), - object: "chat.completion.chunk".to_string(), - created: chrono::Utc::now().timestamp() as u64, - model: String::new(), - choices: vec![], - usage: None, - metadata: Some(serde_json::Value::Object(meta)), - }); + yield Ok(AgentStreamEvent::Metadata(serde_json::Value::Object(meta))); } } - StreamMessage::Delta(delta) => { - // Extract usage from Usage delta variant - let usage = if let GenerationDelta::Usage { usage } = &delta { - Some(usage.clone()) - } else { - None - }; - - yield Ok(ChatCompletionStreamResponse { - id: ctx.request_id.to_string(), - object: "chat.completion.chunk".to_string(), - created: chrono::Utc::now().timestamp() as u64, - model: ctx.state.llm_input.as_ref().map(|llm_input| llm_input.model.clone().to_string()).unwrap_or_default(), - choices: vec![ChatCompletionStreamChoice { - index: 0, - delta: delta.into(), - finish_reason: None, - }], - usage, - metadata: None, - }) + StreamMessage::Event(event) => { + yield Ok(AgentStreamEvent::Event(event)) } } Err(e) => yield Err(ApiStreamError::Unknown(e)), @@ -686,6 +650,99 @@ impl crate::storage::SessionStorage for super::AgentClient { // Helper Methods // ============================================================================= +fn response_to_message(response: &stakai::GenerateResponse) -> Message { + let mut parts = Vec::new(); + + for content in &response.content { + match content { + ResponseContent::Text { text } => parts.push(ContentPart::text(text.clone())), + ResponseContent::Reasoning { reasoning } => { + parts.push(ContentPart::text(format!("[Reasoning: {reasoning}]"))); + } + ResponseContent::ToolCall(tool_call) => { + let mut part = ContentPart::tool_call( + tool_call.id.clone(), + tool_call.name.clone(), + tool_call.arguments.clone(), + ); + if let ContentPart::ToolCall { metadata, .. } = &mut part { + *metadata = tool_call.metadata.clone(); + } + parts.push(part); + } + } + } + + if parts.len() == 1 + && let ContentPart::Text { text, .. } = &parts[0] + { + return Message::new(stakai::Role::Assistant, text.clone()); + } + + Message::new(stakai::Role::Assistant, parts) +} + +fn accumulate_stream_tool_event(parts: &mut Vec, event: &StreamEvent) { + let (id, event_name, event_arguments, event_delta, event_metadata) = match event { + StreamEvent::ToolCallStart { id, name } => (id, Some(name), None, None, None), + StreamEvent::ToolCallDelta { id, delta } => (id, None, None, Some(delta), None), + StreamEvent::ToolCallEnd { + id, + name, + arguments, + metadata, + } => (id, Some(name), Some(arguments), None, metadata.as_ref()), + _ => return, + }; + + let existing = parts + .iter_mut() + .find(|part| matches!(part, ContentPart::ToolCall { id: part_id, .. } if part_id == id)); + + match existing { + Some(ContentPart::ToolCall { + name, + arguments, + metadata, + .. + }) => { + if let Some(new_name) = event_name + && name.is_empty() + { + *name = new_name.clone(); + } + if let Some(arguments_value) = event_arguments { + *arguments = arguments_value.clone(); + } + if let Some(delta) = event_delta { + if let serde_json::Value::String(current) = arguments { + current.push_str(delta); + } else { + *arguments = serde_json::Value::String(delta.clone()); + } + } + if let Some(new_metadata) = event_metadata { + *metadata = Some(new_metadata.clone()); + } + } + _ => { + parts.push(ContentPart::tool_call( + id.clone(), + event_name.cloned().unwrap_or_default(), + event_arguments + .cloned() + .or_else(|| event_delta.cloned().map(serde_json::Value::String)) + .unwrap_or_else(|| serde_json::Value::String(String::new())), + )); + if let Some(new_metadata) = event_metadata + && let Some(ContentPart::ToolCall { metadata, .. }) = parts.last_mut() + { + *metadata = Some(new_metadata.clone()); + } + } + } +} + const TITLE_GENERATOR_PROMPT: &str = include_str!("../prompts/session_title_generator.v1.txt"); impl AgentClient { @@ -789,13 +846,14 @@ impl AgentClient { }) } - fn fallback_session_title(messages: &[ChatMessage]) -> String { + fn fallback_session_title(messages: &[Message]) -> String { messages .iter() - .find(|m| m.role == Role::User) - .and_then(|m| m.content.as_ref()) - .map(|c| { - let text = c.to_string(); + .find(|m| m.role == stakai::Role::User) + .into_iter() + .filter_map(Message::text) + .next() + .map(|text| { text.split_whitespace() .take(5) .collect::>() @@ -808,7 +866,7 @@ impl AgentClient { pub(crate) async fn save_checkpoint( &self, current: &SessionInfo, - messages: Vec, + messages: Vec, metadata: Option, ) -> Result { let mut checkpoint_request = @@ -836,7 +894,7 @@ impl AgentClient { &self, ctx: &mut HookContext, stream_channel_tx: Option>>, - ) -> Result { + ) -> Result { // Execute before inference hooks self.hook_registry .execute_hooks(ctx, &LifecycleEvent::BeforeInference) @@ -855,48 +913,63 @@ impl AgentClient { // Inject session_id header if available if let Some(session_id) = ctx.session_id { - let headers = input + input + .options .headers - .get_or_insert_with(std::collections::HashMap::new); - headers.insert("X-Session-Id".to_string(), session_id.to_string()); + .get_or_insert_with(stakai::Headers::new) + .insert("X-Session-Id", session_id.to_string()); } let (response_message, usage) = if let Some(tx) = stream_channel_tx { // Streaming mode - let (internal_tx, mut internal_rx) = mpsc::channel::(100); - let stream_input = LLMStreamInput { - model: input.model, - messages: input.messages, - max_tokens: input.max_tokens, - tools: input.tools, - stream_channel_tx: internal_tx, - provider_options: input.provider_options, - headers: input.headers, - }; - - let stakai = self.stakai.clone(); - let chat_future = async move { - stakai - .chat_stream(stream_input) - .await - .map_err(|e| e.to_string()) - }; - - let receive_future = async move { - while let Some(delta) = internal_rx.recv().await { - if tx.send(Ok(StreamMessage::Delta(delta))).await.is_err() { - break; + let mut stream = self + .stakai + .stream(&input) + .await + .map_err(|e| e.to_string())?; + let mut text = String::new(); + let mut tool_calls: Vec = Vec::new(); + let mut usage = stakai::Usage::default(); + + while let Some(event_result) = stream.next().await { + let event = event_result.map_err(|e| e.to_string())?; + match &event { + StreamEvent::TextDelta { delta, .. } => text.push_str(delta), + StreamEvent::ToolCallStart { .. } + | StreamEvent::ToolCallDelta { .. } + | StreamEvent::ToolCallEnd { .. } => { + accumulate_stream_tool_event(&mut tool_calls, &event); } + StreamEvent::Finish { + usage: final_usage, .. + } => { + usage = final_usage.clone(); + } + StreamEvent::Start { .. } + | StreamEvent::ReasoningDelta { .. } + | StreamEvent::Error { .. } => {} } - }; + if tx.send(Ok(StreamMessage::Event(event))).await.is_err() { + break; + } + } - let (chat_result, _) = tokio::join!(chat_future, receive_future); - let response = chat_result?; - (response.choices[0].message.clone(), response.usage) + let message = if tool_calls.is_empty() { + Message::new(stakai::Role::Assistant, text) + } else { + let mut parts = vec![ContentPart::text(text)]; + parts.extend(tool_calls); + Message::new(stakai::Role::Assistant, parts) + }; + (message, usage) } else { // Non-streaming mode - let response = self.stakai.chat(input).await.map_err(|e| e.to_string())?; - (response.choices[0].message.clone(), response.usage) + let response = self + .stakai + .generate(&input) + .await + .map_err(|e| e.to_string())?; + (response_to_message(&response), response.usage) }; ctx.state.set_llm_output(response_message, usage); @@ -914,11 +987,11 @@ impl AgentClient { .as_ref() .ok_or_else(|| "LLM output is missing from state".to_string())?; - Ok(ChatMessage::from(llm_output)) + Ok(llm_output.new_message.clone()) } /// Generate a title for a new session - async fn generate_session_title(&self, messages: &[ChatMessage]) -> Result { + async fn generate_session_title(&self, messages: &[Message]) -> Result { // Pick a cheap model from the user's configured providers let use_stakpak = self.stakpak.is_some(); let providers = self.stakai.registry().list_providers(); @@ -940,38 +1013,35 @@ impl AgentClient { }) .ok_or_else(|| "No model available for title generation".to_string())?; - let llm_messages = vec![ - LLMMessage { - role: Role::System.to_string(), - content: LLMMessageContent::String(TITLE_GENERATOR_PROMPT.to_string()), - }, - LLMMessage { - role: Role::User.to_string(), - content: LLMMessageContent::String( - messages - .iter() - .map(|msg| { - msg.content - .as_ref() - .unwrap_or(&MessageContent::String("".to_string())) - .to_string() - }) - .collect(), - ), - }, + let title_messages = vec![ + Message::new(stakai::Role::System, TITLE_GENERATOR_PROMPT), + Message::new( + stakai::Role::User, + messages + .iter() + .filter_map(Message::text) + .collect::>() + .join("\n"), + ), ]; - let input = LLMInput { + let input = stakai::GenerateRequest { model, - messages: llm_messages, - max_tokens: 100, - tools: None, + messages: title_messages, + options: stakai::GenerateOptions { + max_tokens: Some(100), + ..Default::default() + }, provider_options: None, - headers: None, + telemetry_metadata: None, }; - let response = self.stakai.chat(input).await.map_err(|e| e.to_string())?; + let response = self + .stakai + .generate(&input) + .await + .map_err(|e| e.to_string())?; - Ok(response.choices[0].message.content.to_string()) + Ok(response.text()) } } diff --git a/libs/api/src/client/stakai.rs b/libs/api/src/client/stakai.rs new file mode 100644 index 000000000..69fcd2513 --- /dev/null +++ b/libs/api/src/client/stakai.rs @@ -0,0 +1,165 @@ +//! StakAI inference client for the Agent API. + +use stakai::{ + GenerateRequest, GenerateResponse, GenerateStream, Inference, + providers::anthropic::{AnthropicConfig, AnthropicProvider}, + providers::copilot::{CopilotConfig, CopilotProvider}, + providers::gemini::{GeminiConfig, GeminiProvider}, + providers::openai::{OpenAIConfig, OpenAIProvider}, + providers::stakpak::{StakpakProvider, StakpakProviderConfig}, + registry::ProviderRegistry, +}; +use stakpak_shared::models::openai_runtime::{ + OpenAIBackendResolutionInput, resolve_openai_runtime, +}; +use stakpak_shared::models::{ + error::{AgentError, BadRequestErrorMessage}, + llm::{LLMProviderConfig, ProviderConfig}, +}; + +#[derive(Clone)] +pub struct StakAIClient { + inference: Inference, +} + +impl StakAIClient { + pub fn new(config: &LLMProviderConfig) -> Result { + let registry = build_provider_registry(config) + .map_err(|error| invalid_agent_input(format!("Failed to build providers: {error}")))?; + Self::with_registry(registry) + } + + pub fn with_registry(registry: ProviderRegistry) -> Result { + let inference = Inference::builder() + .with_registry(registry) + .build() + .map_err(|error| invalid_agent_input(error.to_string()))?; + + Ok(Self { inference }) + } + + pub async fn generate( + &self, + request: &GenerateRequest, + ) -> Result { + self.inference + .generate(request) + .await + .map_err(|error| invalid_agent_input(error.to_string())) + } + + pub async fn stream(&self, request: &GenerateRequest) -> Result { + self.inference + .stream(request) + .await + .map_err(|error| invalid_agent_input(error.to_string())) + } + + pub fn registry(&self) -> &ProviderRegistry { + self.inference.registry() + } +} + +fn build_provider_registry(config: &LLMProviderConfig) -> Result { + let mut registry = ProviderRegistry::new(); + + for (name, provider_config) in &config.providers { + match provider_config { + ProviderConfig::OpenAI { .. } => { + if let Some(openai_config) = resolve_stakai_openai_config(provider_config)? { + let provider = OpenAIProvider::new(openai_config) + .map_err(|error| format!("Failed to create OpenAI provider: {error}"))?; + registry = registry.register("openai", provider); + } + } + ProviderConfig::Anthropic { api_endpoint, .. } => { + if let Some(config) = anthropic_config(provider_config, api_endpoint.as_deref()) { + let provider = AnthropicProvider::new(config) + .map_err(|error| format!("Failed to create Anthropic provider: {error}"))?; + registry = registry.register("anthropic", provider); + } + } + ProviderConfig::Gemini { api_endpoint, .. } => { + if let Some(api_key) = provider_config.api_key() { + let mut config = GeminiConfig::new(api_key.to_string()); + if let Some(endpoint) = api_endpoint { + config = config.with_base_url(endpoint.clone()); + } + let provider = GeminiProvider::new(config) + .map_err(|error| format!("Failed to create Gemini provider: {error}"))?; + registry = registry.register("google", provider); + } + } + ProviderConfig::Stakpak { api_endpoint, .. } => { + let Some(api_key) = provider_config.api_key() else { + continue; + }; + let mut config = StakpakProviderConfig::new(api_key.to_string()) + .with_user_agent(format!("Stakpak/{}", env!("CARGO_PKG_VERSION"))); + if let Some(endpoint) = api_endpoint { + config = config.with_base_url(endpoint.clone()); + } + let provider = StakpakProvider::new(config) + .map_err(|error| format!("Failed to create Stakpak provider: {error}"))?; + registry = registry.register("stakpak", provider); + } + ProviderConfig::GitHubCopilot { api_endpoint, .. } => { + if let Some(access_token) = provider_config.access_token() { + let mut config = CopilotConfig::new(access_token.to_string()); + if let Some(endpoint) = api_endpoint { + config = config.with_base_url(endpoint.clone()); + } + let provider = CopilotProvider::new(config) + .map_err(|error| format!("Failed to create Copilot provider: {error}"))?; + registry = registry.register("github-copilot", provider); + } + } + ProviderConfig::Custom { api_endpoint, .. } => { + let key = provider_config.api_key().unwrap_or_default().to_string(); + let config = OpenAIConfig::new(key).with_base_url(api_endpoint.clone()); + let provider = OpenAIProvider::new(config).map_err(|error| { + format!("Failed to create custom provider '{name}': {error}") + })?; + registry = registry.register(name, provider); + } + ProviderConfig::Bedrock { .. } => {} + } + } + + Ok(registry) +} + +fn resolve_stakai_openai_config( + provider_config: &ProviderConfig, +) -> Result, String> { + let resolved = resolve_openai_runtime(OpenAIBackendResolutionInput::new( + Some(provider_config.clone()), + provider_config.get_auth(), + )) + .map_err(|error| format!("Failed to resolve OpenAI runtime config: {error}"))?; + + Ok(resolved.map(|config| config.to_stakai_config())) +} + +fn anthropic_config( + provider_config: &ProviderConfig, + api_endpoint: Option<&str>, +) -> Option { + let mut config = if let Some(token) = provider_config.access_token() { + AnthropicConfig::with_oauth(token) + } else if let Some(key) = provider_config.api_key() { + AnthropicConfig::new(key) + } else { + return None; + }; + + if let Some(endpoint) = api_endpoint { + config = config.with_base_url(endpoint); + } + + Some(config) +} + +fn invalid_agent_input(message: String) -> AgentError { + AgentError::BadRequest(BadRequestErrorMessage::InvalidAgentInput(message)) +} diff --git a/libs/api/src/lib.rs b/libs/api/src/lib.rs index eb1306a9e..e20b2a4ad 100644 --- a/libs/api/src/lib.rs +++ b/libs/api/src/lib.rs @@ -3,9 +3,6 @@ use futures_util::Stream; use models::*; use reqwest::header::HeaderMap; use rmcp::model::Content; -use stakpak_shared::models::integrations::openai::{ - ChatCompletionResponse, ChatCompletionStreamResponse, ChatMessage, Tool, -}; use uuid::Uuid; pub mod client; @@ -147,24 +144,22 @@ pub trait AgentProvider: SessionStorage + Send + Sync { async fn chat_completion( &self, model: Model, - messages: Vec, - tools: Option>, + messages: Vec, + tools: Option>, session_id: Option, metadata: Option, - ) -> Result; + ) -> Result; async fn chat_completion_stream( &self, model: Model, - messages: Vec, - tools: Option>, + messages: Vec, + tools: Option>, headers: Option, session_id: Option, metadata: Option, ) -> Result< ( - std::pin::Pin< - Box> + Send>, - >, + std::pin::Pin> + Send>>, Option, ), String, diff --git a/libs/api/src/local/context_managers/common.rs b/libs/api/src/local/context_managers/common.rs index 5a6f4c575..72ce745fb 100644 --- a/libs/api/src/local/context_managers/common.rs +++ b/libs/api/src/local/context_managers/common.rs @@ -1,5 +1,4 @@ use regex::Regex; -use stakpak_shared::models::integrations::openai::{ChatMessage, Role}; use std::fmt::Display; /// Options for history processing shared by context managers @@ -13,80 +12,88 @@ pub struct HistoryProcessingOptions { /// Convert chat messages to history items with the given processing options pub fn messages_to_history( - messages: &[ChatMessage], + messages: &[stakai::Message], options: &HistoryProcessingOptions, ) -> Vec { let mut history_items: Vec = Vec::new(); let mut index = 0; for message in messages.iter() { - match &message.role { - Role::Assistant | Role::User if message.tool_calls.is_none() => { - // clean content from checkpoint_id tag - let content = remove_xml_tag( - "checkpoint_id", - &message.content.clone().unwrap_or_default().to_string(), - ); + let parts = message.parts(); + let tool_calls: Vec<_> = parts + .iter() + .filter_map(|part| match part { + stakai::ContentPart::ToolCall { + id, + name, + arguments, + .. + } => Some((id.clone(), name.clone(), arguments.clone())), + _ => None, + }) + .collect(); + + match message.role { + stakai::Role::Assistant | stakai::Role::User if tool_calls.is_empty() => { + let content = remove_xml_tag("checkpoint_id", &message.text().unwrap_or_default()); history_items.push(HistoryItem { index, content: HistoryItemContent::Message { - role: message.role.clone(), + role: role_label(message.role).to_string(), content, }, }); index += 1; } - Role::Assistant | Role::User if message.tool_calls.is_some() => { - // clean content from checkpoint_id tag + stakai::Role::Assistant | stakai::Role::User => { let content = message - .content - .clone() - .map(|c| remove_xml_tag("checkpoint_id", &c.to_string())); - for tool_call in message.tool_calls.clone().unwrap_or_default() { + .text() + .map(|text| remove_xml_tag("checkpoint_id", &text)); + for (id, name, arguments) in tool_calls { history_items.push(HistoryItem { index, content: HistoryItemContent::Action { - role: message.role.clone(), - id: tool_call.id.clone(), - name: tool_call.function.name.clone(), + role: role_label(message.role).to_string(), + id, + name, status: HistoryItemActionStatus::Pending, message: content.clone(), - arguments: serde_json::from_str(&tool_call.function.arguments) - .unwrap_or_default(), + arguments, result: None, }, }); index += 1; } } - Role::Tool => { - // Find the corresponding tool call item and update it with the result - if let Some(tool_call_id) = &message.tool_call_id { - // Look for the matching tool call in history items + stakai::Role::Tool => { + for part in parts { + let stakai::ContentPart::ToolResult { + tool_call_id, + content, + .. + } = part + else { + continue; + }; + if let Some(history_item) = history_items.iter_mut().find(|item| { if let HistoryItemContent::Action { id, .. } = &item.content { - *id == *tool_call_id + *id == tool_call_id } else { false } - }) { - // Update the tool call with the result - if let HistoryItemContent::Action { status, result, .. } = - &mut history_item.content + }) && let HistoryItemContent::Action { status, result, .. } = + &mut history_item.content + { + *result = Some(content.clone()); + + if let Some(result) = result + && result.as_str().unwrap_or_default() == "TOOL_CALL_CANCELLED" { - let result_content = - message.content.clone().unwrap_or_default().to_string(); - *result = serde_json::from_str(&result_content) - .unwrap_or(Some(serde_json::Value::String(result_content))); - - if let Some(result) = result - && result.as_str().unwrap_or_default() == "TOOL_CALL_CANCELLED" - { - *status = HistoryItemActionStatus::Aborted; - continue; - } - *status = HistoryItemActionStatus::Completed; + *status = HistoryItemActionStatus::Aborted; + continue; } + *status = HistoryItemActionStatus::Completed; } } } @@ -250,11 +257,11 @@ impl Display for HistoryItem { pub enum HistoryItemContent { Message { - role: Role, + role: String, content: String, }, Action { - role: Role, + role: String, id: String, name: String, status: HistoryItemActionStatus, @@ -264,6 +271,15 @@ pub enum HistoryItemContent { }, } +pub fn role_label(role: stakai::Role) -> &'static str { + match role { + stakai::Role::System => "system", + stakai::Role::User => "user", + stakai::Role::Assistant => "assistant", + stakai::Role::Tool => "tool", + } +} + impl HistoryItemContent { pub fn is_action(&self) -> bool { matches!(self, HistoryItemContent::Action { .. }) @@ -294,7 +310,7 @@ mod tests { HistoryItem { index, content: HistoryItemContent::Action { - role: Role::Assistant, + role: "assistant".to_string(), id: format!("action_{}", index), name: "test_action".to_string(), status: HistoryItemActionStatus::Completed, @@ -309,7 +325,7 @@ mod tests { HistoryItem { index, content: HistoryItemContent::Message { - role: Role::User, + role: "user".to_string(), content: "test message".to_string(), }, } diff --git a/libs/api/src/local/context_managers/file_scratchpad_context_manager.rs b/libs/api/src/local/context_managers/file_scratchpad_context_manager.rs index 923c5ce62..e6b237f48 100644 --- a/libs/api/src/local/context_managers/file_scratchpad_context_manager.rs +++ b/libs/api/src/local/context_managers/file_scratchpad_context_manager.rs @@ -1,8 +1,4 @@ use super::common::{HistoryItemActionStatus, HistoryItemContent, remove_xml_tag}; -use stakpak_shared::models::{ - integrations::openai::{ChatMessage, Role}, - llm::{LLMMessage, LLMMessageContent}, -}; use std::{ collections::HashSet, fmt::Display, @@ -39,7 +35,7 @@ pub struct FileScratchpadContextManagerOptions { } impl super::ContextManager for FileScratchpadContextManager { - fn reduce_context(&self, messages: Vec) -> Vec { + fn reduce_context(&self, messages: Vec) -> Vec { self.reduce_context_with_session(messages, None) } } @@ -90,9 +86,9 @@ impl FileScratchpadContextManager { pub fn reduce_context_with_session( &self, - messages: Vec, + messages: Vec, session_id: Option, - ) -> Vec { + ) -> Vec { let session_key = session_id.map(|u| u.to_string()).unwrap_or_default(); let should_recover = { let recovered = self.recovered_sessions.read().unwrap(); @@ -111,18 +107,15 @@ impl FileScratchpadContextManager { let history = self.messages_to_history(&messages, session_id); let context_content = self.history_to_text(&history, &scratchpad_content, &todo_content); - vec![LLMMessage { - role: Role::User.to_string(), - content: LLMMessageContent::String(context_content), - }] + vec![stakai::Message::new(stakai::Role::User, context_content)] } - fn recover_from_history(&self, messages: &[ChatMessage], session_id: Option) { + fn recover_from_history(&self, messages: &[stakai::Message], session_id: Option) { self.recover_file(messages, &self.get_scratchpad_path(session_id)); self.recover_file(messages, &self.get_todo_path(session_id)); } - fn recover_file(&self, messages: &[ChatMessage], path: &Path) { + fn recover_file(&self, messages: &[stakai::Message], path: &Path) { // Try to reconstruct from tool calls let reconstructed_content = self.reconstruct_file_from_history(messages, path); @@ -155,16 +148,19 @@ impl FileScratchpadContextManager { fn reconstruct_file_from_history( &self, - messages: &[ChatMessage], + messages: &[stakai::Message], path: &Path, ) -> Option { let mut current_content: Option = None; let path_str = path.to_string_lossy(); for message in messages { - if let Some(tool_calls) = &message.tool_calls { - for tool_call in tool_calls { - self.apply_tool_call(tool_call, &path_str, &mut current_content); + for part in message.parts() { + if let stakai::ContentPart::ToolCall { + name, arguments, .. + } = part + { + self.apply_tool_call(&name, &arguments, &path_str, &mut current_content); } } } @@ -174,14 +170,11 @@ impl FileScratchpadContextManager { fn apply_tool_call( &self, - tool_call: &stakpak_shared::models::integrations::openai::ToolCall, + tool_name: &str, + args: &serde_json::Value, target_path: &str, current_content: &mut Option, ) { - let args: serde_json::Value = - serde_json::from_str(&tool_call.function.arguments).unwrap_or_default(); - let tool_name = tool_call.function.name.as_str(); - if let Some(path) = args.get("path").and_then(|p| p.as_str()) { if !path.contains(target_path) { return; @@ -217,82 +210,90 @@ impl FileScratchpadContextManager { fn messages_to_history( &self, - messages: &[ChatMessage], + messages: &[stakai::Message], session_id: Option, ) -> Vec { let mut history_items: Vec = Vec::new(); let mut index = 0; for (message_index, message) in messages.iter().enumerate() { - match &message.role { - Role::Assistant | Role::User if message.tool_calls.is_none() => { - // clean content from checkpoint_id tag - let content = remove_xml_tag( - "checkpoint_id", - &message.content.clone().unwrap_or_default().to_string(), - ); + let parts = message.parts(); + let tool_calls: Vec<_> = parts + .iter() + .filter_map(|part| match part { + stakai::ContentPart::ToolCall { + id, + name, + arguments, + .. + } => Some((id.clone(), name.clone(), arguments.clone())), + _ => None, + }) + .collect(); + + match message.role { + stakai::Role::Assistant | stakai::Role::User if tool_calls.is_empty() => { + let content = + remove_xml_tag("checkpoint_id", &message.text().unwrap_or_default()); history_items.push(HistoryItem { index, message_index, content: HistoryItemContent::Message { - role: message.role.clone(), + role: super::common::role_label(message.role).to_string(), content, }, }); index += 1; } - Role::Assistant | Role::User if message.tool_calls.is_some() => { - // clean content from checkpoint_id tag + stakai::Role::Assistant | stakai::Role::User => { let content = message - .content - .clone() - .map(|c| remove_xml_tag("checkpoint_id", &c.to_string())); - for tool_call in message.tool_calls.clone().unwrap_or_default() { + .text() + .map(|text| remove_xml_tag("checkpoint_id", &text)); + for (id, name, arguments) in tool_calls { history_items.push(HistoryItem { index, message_index, content: HistoryItemContent::Action { - role: message.role.clone(), - id: tool_call.id.clone(), - name: tool_call.function.name.clone(), + role: super::common::role_label(message.role).to_string(), + id, + name, status: HistoryItemActionStatus::Pending, message: content.clone(), - arguments: serde_json::from_str(&tool_call.function.arguments) - .unwrap_or_default(), + arguments, result: None, }, }); index += 1; } } - Role::Tool => { - // Find the corresponding tool call item and update it with the result - if let Some(tool_call_id) = &message.tool_call_id { - // Look for the matching tool call in history items + stakai::Role::Tool => { + for part in parts { + let stakai::ContentPart::ToolResult { + tool_call_id, + content, + .. + } = part + else { + continue; + }; if let Some(history_item) = history_items.iter_mut().find(|item| { if let HistoryItemContent::Action { id, .. } = &item.content { - *id == *tool_call_id + *id == tool_call_id } else { false } - }) { - // Update the tool call with the result - if let HistoryItemContent::Action { status, result, .. } = - &mut history_item.content + }) && let HistoryItemContent::Action { status, result, .. } = + &mut history_item.content + { + *result = Some(content.clone()); + + if let Some(result) = result + && result.as_str().unwrap_or_default() == "TOOL_CALL_CANCELLED" { - let result_content = - message.content.clone().unwrap_or_default().to_string(); - *result = serde_json::from_str(&result_content) - .unwrap_or(Some(serde_json::Value::String(result_content))); - - if let Some(result) = result - && result.as_str().unwrap_or_default() == "TOOL_CALL_CANCELLED" - { - *status = HistoryItemActionStatus::Aborted; - continue; - } - *status = HistoryItemActionStatus::Completed; + *status = HistoryItemActionStatus::Aborted; + continue; } + *status = HistoryItemActionStatus::Completed; } } } @@ -513,321 +514,3 @@ impl Display for HistoryItem { Ok(()) } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::local::context_managers::ContextManager; - use stakpak_shared::models::integrations::openai::{FunctionCall, ToolCall}; - use tempfile::TempDir; - - fn create_test_manager() -> (FileScratchpadContextManager, TempDir) { - let temp_dir = TempDir::new().unwrap(); - let scratchpad_path = temp_dir.path().join("scratchpad.md"); - let todo_path = temp_dir.path().join("todo.md"); - - let manager = FileScratchpadContextManager::new(FileScratchpadContextManagerOptions { - history_action_message_size_limit: 100, - history_action_message_keep_last_n: 1, - history_action_result_keep_last_n: 50, - scratchpad_file_path: scratchpad_path, - todo_file_path: todo_path, - overwrite_if_different: false, - }); - - (manager, temp_dir) - } - - #[test] - fn test_load_empty_files() { - let (manager, _temp_dir) = create_test_manager(); - - // Without creating files, load should return empty strings - let scratchpad = manager.load_file(&manager.get_scratchpad_path(None)); - let todo = manager.load_file(&manager.get_todo_path(None)); - - assert!(scratchpad.is_empty()); - assert!(todo.is_empty()); - } - - #[test] - fn test_load_existing_files() { - let (manager, _temp_dir) = create_test_manager(); - - // Write content to files (simulating agent file edits) - fs::write( - &manager.scratchpad_file_path, - "# My Notes\nSome important info", - ) - .unwrap(); - fs::write(&manager.todo_file_path, "- [x] Task 1\n- [ ] Task 2").unwrap(); - - // Load and verify - let scratchpad = manager.load_file(&manager.get_scratchpad_path(None)); - let todo = manager.load_file(&manager.get_todo_path(None)); - - assert!(scratchpad.contains("My Notes")); - assert!(scratchpad.contains("important info")); - assert!(todo.contains("Task 1")); - assert!(todo.contains("Task 2")); - } - - #[test] - fn test_history_to_text_with_content() { - let (manager, _temp_dir) = create_test_manager(); - - let history = vec![]; - let scratchpad = "# Notes\nKey info here"; - let todo = "- [ ] Do something"; - - let result = manager.history_to_text(&history, scratchpad, todo); - - assert!(result.contains("")); - assert!(result.contains("Key info here")); - assert!(result.contains("")); - assert!(result.contains("Do something")); - assert!(result.contains("")); - } - - #[test] - fn test_history_to_text_empty_scratchpad() { - let (manager, _temp_dir) = create_test_manager(); - - let history = vec![]; - let scratchpad = ""; - let todo = ""; - - let result = manager.history_to_text(&history, scratchpad, todo); - - // Should not include empty scratchpad/todo sections - assert!(!result.contains("")); - assert!(!result.contains("")); - assert!(result.contains("")); - } - - #[test] - fn test_paths_accessible() { - let (manager, temp_dir) = create_test_manager(); - - // Verify paths are accessible for system prompt generation - assert_eq!( - &manager.scratchpad_file_path, - &temp_dir.path().join("scratchpad.md") - ); - assert_eq!(&manager.todo_file_path, &temp_dir.path().join("todo.md")); - } - - fn create_tool_call_history( - scratchpad_path: &str, - todo_path: &str, - scratchpad_content: &str, - todo_content: &str, - ) -> Vec { - vec![ - ChatMessage { - role: Role::Assistant, - content: None, - tool_calls: Some(vec![ToolCall { - id: "call_1".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "create".to_string(), - arguments: serde_json::json!({ - "path": scratchpad_path, - "file_text": scratchpad_content - }) - .to_string(), - }, - metadata: None, - }]), - tool_call_id: None, - name: None, - usage: None, - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: None, - tool_calls: Some(vec![ToolCall { - id: "call_2".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "create".to_string(), - arguments: serde_json::json!({ - "path": todo_path, - "file_text": todo_content - }) - .to_string(), - }, - metadata: None, - }]), - tool_call_id: None, - name: None, - usage: None, - ..Default::default() - }, - ] - } - - #[test] - fn test_recover_from_history() { - let (manager, _temp_dir) = create_test_manager(); - let scratchpad_path = manager - .get_scratchpad_path(None) - .to_string_lossy() - .to_string(); - let todo_path = manager.get_todo_path(None).to_string_lossy().to_string(); - - let history = create_tool_call_history( - &scratchpad_path, - &todo_path, - "Recovered Content", - "- [ ] Recovered Task", - ); - - // Initially empty - assert!( - manager - .load_file(&manager.get_scratchpad_path(None)) - .is_empty() - ); - assert!(manager.load_file(&manager.get_todo_path(None)).is_empty()); - - // Trigger recovery - manager.reduce_context(history); - - // Verify recovered content - let scratchpad = manager.load_file(&manager.get_scratchpad_path(None)); - let todo = manager.load_file(&manager.get_todo_path(None)); - - assert!(scratchpad.contains("Recovered Content")); - assert!(todo.contains("Recovered Task")); - } - - #[test] - fn test_recover_no_overwrite_existing() { - let (manager, _temp_dir) = create_test_manager(); - let scratchpad_path = manager - .get_scratchpad_path(None) - .to_string_lossy() - .to_string(); - let todo_path = manager.get_todo_path(None).to_string_lossy().to_string(); - - // Create existing files - fs::write(&manager.scratchpad_file_path, "Existing Scratchpad").unwrap(); - fs::write(&manager.todo_file_path, "Existing Todo").unwrap(); - - let history = - create_tool_call_history(&scratchpad_path, &todo_path, "New Scratchpad", "New Todo"); - - // Trigger recovery - manager.reduce_context(history); - - // Verify content retrieved from disk matches EXISTING, not new - let scratchpad = manager.load_file(&manager.get_scratchpad_path(None)); - let todo = manager.load_file(&manager.get_todo_path(None)); - - assert_eq!(scratchpad, "Existing Scratchpad"); - assert_eq!(todo, "Existing Todo"); - } - - #[test] - fn test_recover_overwrite_if_different() { - let temp_dir = TempDir::new().unwrap(); - let scratchpad_path = temp_dir.path().join("scratchpad.md"); - let todo_path = temp_dir.path().join("todo.md"); - - // Enable overwrite_if_different - let manager = FileScratchpadContextManager::new(FileScratchpadContextManagerOptions { - history_action_message_size_limit: 100, - history_action_message_keep_last_n: 1, - history_action_result_keep_last_n: 50, - scratchpad_file_path: scratchpad_path.clone(), - todo_file_path: todo_path.clone(), - overwrite_if_different: true, - }); - - // Create existing files - fs::write(&scratchpad_path, "Existing Scratchpad").unwrap(); - fs::write(&todo_path, "Existing Todo").unwrap(); - - let history = create_tool_call_history( - &scratchpad_path.to_string_lossy(), - &todo_path.to_string_lossy(), - "New Scratchpad", - "New Todo", - ); - - // Trigger recovery - manager.reduce_context(history); - - // Verify content overwritten - let scratchpad = fs::read_to_string(&scratchpad_path).unwrap(); - let todo = fs::read_to_string(&todo_path).unwrap(); - - assert_eq!(scratchpad.trim(), "New Scratchpad"); - assert_eq!(todo.trim(), "New Todo"); - } - #[test] - fn test_recover_from_tool_calls() { - let (manager, _temp_dir) = create_test_manager(); - let scratchpad_path = manager - .get_scratchpad_path(None) - .to_string_lossy() - .to_string(); - - let history = vec![ - ChatMessage { - role: Role::Assistant, - content: None, - tool_calls: Some(vec![ToolCall { - id: "call_1".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "create".to_string(), - arguments: serde_json::json!({ - "path": scratchpad_path, - "file_text": "Initial Content" - }) - .to_string(), - }, - metadata: None, - }]), - tool_call_id: None, - name: None, - usage: None, - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - // ... (rest of test content clipped, assume context match is enough) - content: None, - tool_calls: Some(vec![ToolCall { - id: "call_2".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "str_replace".to_string(), - arguments: serde_json::json!({ - "path": scratchpad_path, - "old_str": "Initial", - "new_str": "Updated" - }) - .to_string(), - }, - metadata: None, - }]), - tool_call_id: None, - name: None, - usage: None, - ..Default::default() - }, - ]; - - // Trigger recovery - manager.reduce_context(history); - - // Verify recovered content - let scratchpad = manager.load_file(&manager.get_scratchpad_path(None)); - assert_eq!(scratchpad, "Updated Content"); - } -} diff --git a/libs/api/src/local/context_managers/mod.rs b/libs/api/src/local/context_managers/mod.rs index 2e2189ca9..24023bc7b 100644 --- a/libs/api/src/local/context_managers/mod.rs +++ b/libs/api/src/local/context_managers/mod.rs @@ -1,7 +1,5 @@ -use stakpak_shared::models::{integrations::openai::ChatMessage, llm::LLMMessage}; - pub trait ContextManager: Send + Sync { - fn reduce_context(&self, messages: Vec) -> Vec; + fn reduce_context(&self, messages: Vec) -> Vec; } pub mod common; diff --git a/libs/api/src/local/context_managers/scratchpad_context_manager.rs b/libs/api/src/local/context_managers/scratchpad_context_manager.rs index 3760579fd..0209e82ed 100644 --- a/libs/api/src/local/context_managers/scratchpad_context_manager.rs +++ b/libs/api/src/local/context_managers/scratchpad_context_manager.rs @@ -1,10 +1,6 @@ use super::common::{HistoryItem, HistoryProcessingOptions, messages_to_history}; use once_cell::sync::Lazy; use regex::Regex; -use stakpak_shared::models::{ - integrations::openai::{ChatMessage, Role}, - llm::{LLMMessage, LLMMessageContent}, -}; use std::collections::HashMap; pub struct ScratchpadContextManager { @@ -12,12 +8,12 @@ pub struct ScratchpadContextManager { } impl super::ContextManager for ScratchpadContextManager { - fn reduce_context(&self, messages: Vec) -> Vec { + fn reduce_context(&self, messages: Vec) -> Vec { let mut scratchpad = HashMap::new(); for message in messages.iter() { - if let Some(content) = message.content.clone() - && let Some(scratchpad_content) = extract_scratchpad(&content.to_string()) + if let Some(content) = message.text() + && let Some(scratchpad_content) = extract_scratchpad(&content) { scratchpad.extend(scratchpad_content); } @@ -26,10 +22,7 @@ impl super::ContextManager for ScratchpadContextManager { let history = messages_to_history(&messages, &self.options); let context_content = self.history_to_text(&history, &scratchpad); - vec![LLMMessage { - role: Role::User.to_string(), - content: LLMMessageContent::String(context_content), - }] + vec![stakai::Message::new(stakai::Role::User, context_content)] } } diff --git a/libs/api/src/local/context_managers/simple_context_manager.rs b/libs/api/src/local/context_managers/simple_context_manager.rs index d53d473d6..a580c3dea 100644 --- a/libs/api/src/local/context_managers/simple_context_manager.rs +++ b/libs/api/src/local/context_managers/simple_context_manager.rs @@ -1,13 +1,8 @@ -use stakpak_shared::models::{ - integrations::openai::{ChatMessage, Role}, - llm::{LLMMessage, LLMMessageContent}, -}; - #[allow(dead_code)] pub struct SimpleContextManager; impl super::ContextManager for SimpleContextManager { - fn reduce_context(&self, messages: Vec) -> Vec { + fn reduce_context(&self, messages: Vec) -> Vec { if messages.is_empty() { return vec![]; } @@ -18,19 +13,22 @@ impl super::ContextManager for SimpleContextManager { if messages.len() > 1 { let history_content = messages[..messages.len() - 1] .iter() - .map(|m| format!("{}: {}", m.role, m.content.clone().unwrap_or_default())) + .map(|m| { + format!( + "{}: {}", + super::common::role_label(m.role), + m.text().unwrap_or_default() + ) + }) .collect::>() .join("\n"); - context.push(LLMMessage { - role: Role::User.to_string(), - content: LLMMessageContent::String(history_content), - }); + context.push(stakai::Message::new(stakai::Role::User, history_content)); } // 2. Preserve the last message (with images) if let Some(last_message) = messages.last() { - context.push(LLMMessage::from(last_message.clone())); + context.push(last_message.clone()); } context diff --git a/libs/api/src/local/context_managers/task_board_context_manager.rs b/libs/api/src/local/context_managers/task_board_context_manager.rs index 80a9b96ef..aa979efd2 100644 --- a/libs/api/src/local/context_managers/task_board_context_manager.rs +++ b/libs/api/src/local/context_managers/task_board_context_manager.rs @@ -1,29 +1,18 @@ use std::collections::HashMap; -use stakpak_shared::models::{ - integrations::openai::{ChatMessage, MessageContent}, - llm::{LLMMessage, LLMMessageContent, LLMMessageTypedContent, LLMTool}, -}; - pub struct TaskBoardContextManager { keep_last_n_assistant_messages: usize, context_budget_threshold: f32, } impl super::ContextManager for TaskBoardContextManager { - fn reduce_context(&self, messages: Vec) -> Vec { - let llm_messages: Vec<_> = messages + fn reduce_context(&self, messages: Vec) -> Vec { + let messages = messages .into_iter() - .map(|mut message| { - // Remove checkpoint_id XML tags from message content - Self::clean_checkpoint_tags(&mut message); - message - }) - .map(LLMMessage::from) + .map(Self::clean_checkpoint_tags) .collect(); - - let llm_messages = merge_consecutive_same_role(llm_messages); - dedup_tool_results(llm_messages) + let messages = merge_consecutive_same_role(messages); + dedup_tool_results(messages) } } @@ -31,26 +20,20 @@ const TRIMMED_CONTENT_PLACEHOLDER: &str = "[trimmed]"; impl TaskBoardContextManager { /// Remove `...` XML tags from message content. - fn clean_checkpoint_tags(message: &mut ChatMessage) { - if let Some(content) = message.content.take() { - message.content = Some(match content { - MessageContent::String(s) => { - MessageContent::String(super::common::remove_xml_tag("checkpoint_id", &s)) + fn clean_checkpoint_tags(mut message: stakai::Message) -> stakai::Message { + match &mut message.content { + stakai::MessageContent::Text(text) => { + *text = super::common::remove_xml_tag("checkpoint_id", text); + } + stakai::MessageContent::Parts(parts) => { + for part in parts { + if let stakai::ContentPart::Text { text, .. } = part { + *text = super::common::remove_xml_tag("checkpoint_id", text); + } } - MessageContent::Array(parts) => MessageContent::Array( - parts - .into_iter() - .map(|mut part| { - if let Some(text) = part.text { - part.text = - Some(super::common::remove_xml_tag("checkpoint_id", &text)); - } - part - }) - .collect(), - ), - }); + } } + message } } @@ -72,21 +55,23 @@ impl TaskBoardContextManager { /// /// Each content part carries API-level structural tokens (type discriminators, /// JSON keys, tool_use_id fields) that aren't captured by content length alone. - fn estimate_content_part_tokens(part: &LLMMessageTypedContent) -> u64 { + fn estimate_content_part_tokens(part: &stakai::ContentPart) -> u64 { match part { - LLMMessageTypedContent::Text { text } => Self::bytes_to_tokens(text.len()), - LLMMessageTypedContent::ToolCall { name, args, .. } => { - let content_bytes = name.len() + args.to_string().len(); + stakai::ContentPart::Text { text, .. } => Self::bytes_to_tokens(text.len()), + stakai::ContentPart::ToolCall { + name, arguments, .. + } => { + let content_bytes = name.len() + arguments.to_string().len(); // +30 bytes for tool_use_id (~26 chars), "type":"tool_use", JSON structure Self::bytes_to_tokens(content_bytes + 30) } - LLMMessageTypedContent::ToolResult { content, .. } => { + stakai::ContentPart::ToolResult { content, .. } => { // +30 bytes for tool_use_id, "type":"tool_result", JSON structure - Self::bytes_to_tokens(content.len() + 30) + Self::bytes_to_tokens(content.to_string().len() + 30) } // Images: Anthropic charges 1600-6400+ tokens depending on resolution. // Use 2000 as a conservative default for typical images. - LLMMessageTypedContent::Image { .. } => 2000, + stakai::ContentPart::Image { .. } => 2000, } } @@ -103,13 +88,13 @@ impl TaskBoardContextManager { /// - 30 bytes structural overhead per ToolCall/ToolResult (IDs, type fields) /// - 2000 tokens per image (vs previous 1000) /// - 5% safety buffer on the final total - pub fn estimate_tokens(messages: &[LLMMessage]) -> u64 { + pub fn estimate_tokens(messages: &[stakai::Message]) -> u64 { let raw_estimate: u64 = messages .iter() .map(|msg| { let content_tokens = match &msg.content { - LLMMessageContent::String(s) => Self::bytes_to_tokens(s.len()), - LLMMessageContent::List(parts) => { + stakai::MessageContent::Text(s) => Self::bytes_to_tokens(s.len()), + stakai::MessageContent::Parts(parts) => { let part_tokens: u64 = parts.iter().map(Self::estimate_content_part_tokens).sum(); // Per-part structural overhead: each content block in a List @@ -130,23 +115,24 @@ impl TaskBoardContextManager { /// Trim a single message's content, replacing it with a placeholder. /// Preserves message structure (role, tool_call_ids) for API validity. - fn trim_message(msg: &mut LLMMessage) { + fn trim_message(msg: &mut stakai::Message) { match &mut msg.content { - LLMMessageContent::String(s) => { + stakai::MessageContent::Text(s) => { *s = TRIMMED_CONTENT_PLACEHOLDER.to_string(); } - LLMMessageContent::List(parts) => { + stakai::MessageContent::Parts(parts) => { for part in parts.iter_mut() { match part { - LLMMessageTypedContent::Text { text } => { + stakai::ContentPart::Text { text, .. } => { *text = TRIMMED_CONTENT_PLACEHOLDER.to_string(); } - LLMMessageTypedContent::ToolResult { content, .. } => { - *content = TRIMMED_CONTENT_PLACEHOLDER.to_string(); + stakai::ContentPart::ToolResult { content, .. } => { + *content = + serde_json::Value::String(TRIMMED_CONTENT_PLACEHOLDER.to_string()); } // Preserve ToolCall structure - needed for API to match tool_use/tool_result - LLMMessageTypedContent::ToolCall { .. } => {} - LLMMessageTypedContent::Image { .. } => {} + stakai::ContentPart::ToolCall { .. } => {} + stakai::ContentPart::Image { .. } => {} } } } @@ -159,13 +145,14 @@ impl TaskBoardContextManager { /// included in the chat message list, so the trimmer needs to account /// for them separately. The system prompt is already part of the message /// list and is covered by `estimate_tokens`. - pub fn estimate_tool_overhead(tools: Option<&[LLMTool]>) -> u64 { + pub fn estimate_tool_overhead(tools: Option<&[stakai::Tool]>) -> u64 { tools .map(|t| { t.iter() .map(|tool| { - let schema_len = tool.input_schema.to_string().len(); - let tool_bytes = tool.name.len() + tool.description.len() + schema_len; + let schema_len = tool.function.parameters.to_string().len(); + let tool_bytes = + tool.function.name.len() + tool.function.description.len() + schema_len; // 1.2x multiplier accounts for JSON structural overhead not // captured by content length: property names, "type":"function" // wrappers, "required" arrays, enum values, nested refs, etc. @@ -197,19 +184,15 @@ impl TaskBoardContextManager { /// returned as-is with no metadata changes. pub fn reduce_context_with_budget( &self, - messages: Vec, + messages: Vec, context_window: u64, metadata: Option, - tools: Option<&[LLMTool]>, - ) -> (Vec, Option) { + tools: Option<&[stakai::Tool]>, + ) -> (Vec, Option) { // Standard processing: clean, convert, merge, dedup let llm_messages: Vec<_> = messages .into_iter() - .map(|mut message| { - Self::clean_checkpoint_tags(&mut message); - message - }) - .map(LLMMessage::from) + .map(Self::clean_checkpoint_tags) .collect(); let llm_messages = merge_consecutive_same_role(llm_messages); @@ -253,7 +236,7 @@ impl TaskBoardContextManager { if self.keep_last_n_assistant_messages > 0 { let mut assistant_count = 0usize; for i in (0..len).rev() { - if llm_messages[i].role == "assistant" { + if llm_messages[i].role == stakai::Role::Assistant { assistant_count += 1; if assistant_count >= self.keep_last_n_assistant_messages { keep_n_trim_end = i; @@ -268,7 +251,7 @@ impl TaskBoardContextManager { // stable for prompt caching. let prev_clamped = prev_trimmed_up_to.min(len); for msg in &mut llm_messages[..prev_clamped] { - if msg.role == "system" || msg.role == "user" { + if msg.role == stakai::Role::System || msg.role == stakai::Role::User { continue; } Self::trim_message(msg); @@ -296,8 +279,8 @@ impl TaskBoardContextManager { .take(keep_n_trim_end.min(len)) .skip(prev_clamped) { - let role = msg.role.as_str(); - if role == "assistant" || role == "tool" { + let role = msg.role; + if role == stakai::Role::Assistant || role == stakai::Role::Tool { Self::trim_message(msg); } } @@ -315,8 +298,8 @@ impl TaskBoardContextManager { if current_tokens > threshold { let mut scan_idx = candidate; while scan_idx < len { - let role = llm_messages[scan_idx].role.as_str(); - if role == "assistant" || role == "tool" { + let role = llm_messages[scan_idx].role; + if role == stakai::Role::Assistant || role == stakai::Role::Tool { Self::trim_message(&mut llm_messages[scan_idx]); candidate = scan_idx + 1; @@ -342,7 +325,7 @@ impl TaskBoardContextManager { let clamped_end = effective_trim_end.min(len); if clamped_end > prev_clamped { for msg in &mut llm_messages[prev_clamped..clamped_end] { - if msg.role == "system" || msg.role == "user" { + if msg.role == stakai::Role::System || msg.role == stakai::Role::User { continue; } Self::trim_message(msg); @@ -363,37 +346,42 @@ impl TaskBoardContextManager { } } -/// Merge consecutive LLMMessages that share the same role into a single message. +/// Merge consecutive messages that share the same role into a single message. /// /// When the assistant returns N tool_calls, the chat history contains N separate /// `role=tool` messages. Provider conversion layers map `tool` → `user`, which /// creates N consecutive `user` messages — invalid for Anthropic. By merging -/// them here into a single `role=tool` LLMMessage with multiple ToolResult +/// them here into a single `role=tool` StakAI message with multiple tool results. /// content parts, the downstream conversion produces one `user` message with /// all the tool_result blocks. -fn merge_consecutive_same_role(messages: Vec) -> Vec { +fn merge_consecutive_same_role(messages: Vec) -> Vec { if messages.is_empty() { return messages; } - let mut result: Vec = Vec::with_capacity(messages.len()); + let mut result: Vec = Vec::with_capacity(messages.len()); for msg in messages { let should_merge = result.last().is_some_and(|prev| prev.role == msg.role); if should_merge { let prev = result.last_mut().expect("checked above"); - let new_parts = msg.content.into_parts(); - prev.content = match std::mem::take(&mut prev.content) { - LLMMessageContent::String(s) if s.is_empty() => LLMMessageContent::List(new_parts), - LLMMessageContent::String(s) => { - let mut parts = vec![LLMMessageTypedContent::Text { text: s }]; + let new_parts = msg.content.parts(); + prev.content = match std::mem::replace( + &mut prev.content, + stakai::MessageContent::Text(String::new()), + ) { + stakai::MessageContent::Text(s) if s.is_empty() => { + stakai::MessageContent::Parts(new_parts) + } + stakai::MessageContent::Text(s) => { + let mut parts = vec![stakai::ContentPart::text(s)]; parts.extend(new_parts); - LLMMessageContent::List(parts) + stakai::MessageContent::Parts(parts) } - LLMMessageContent::List(mut existing) => { + stakai::MessageContent::Parts(mut existing) => { existing.extend(new_parts); - LLMMessageContent::List(existing) + stakai::MessageContent::Parts(existing) } }; } else { @@ -406,13 +394,13 @@ fn merge_consecutive_same_role(messages: Vec) -> Vec { /// Remove duplicate ToolResult entries that share the same tool_use_id. /// Keeps only the **last** occurrence (the most recent / retried result). -fn dedup_tool_results(mut messages: Vec) -> Vec { +fn dedup_tool_results(mut messages: Vec) -> Vec { for msg in &mut messages { - if msg.role != "tool" { + if msg.role != stakai::Role::Tool { continue; } let parts = match &mut msg.content { - LLMMessageContent::List(p) => p, + stakai::MessageContent::Parts(p) => p, _ => continue, }; @@ -420,9 +408,9 @@ fn dedup_tool_results(mut messages: Vec) -> Vec { let mut last_index: HashMap = HashMap::new(); let mut counts: HashMap = HashMap::new(); for (i, part) in parts.iter().enumerate() { - if let LLMMessageTypedContent::ToolResult { tool_use_id, .. } = part { - last_index.insert(tool_use_id.clone(), i); - *counts.entry(tool_use_id.clone()).or_insert(0) += 1; + if let stakai::ContentPart::ToolResult { tool_call_id, .. } = part { + last_index.insert(tool_call_id.clone(), i); + *counts.entry(tool_call_id.clone()).or_insert(0) += 1; } } @@ -434,10 +422,10 @@ fn dedup_tool_results(mut messages: Vec) -> Vec { let mut idx = 0; parts.retain(|part| { - let keep = if let LLMMessageTypedContent::ToolResult { tool_use_id, .. } = part { - if counts.get(tool_use_id).copied().unwrap_or(0) > 1 { + let keep = if let stakai::ContentPart::ToolResult { tool_call_id, .. } = part { + if counts.get(tool_call_id).copied().unwrap_or(0) > 1 { // Duplicate — keep only the last one - last_index.get(tool_use_id).copied() == Some(idx) + last_index.get(tool_call_id).copied() == Some(idx) } else { true } @@ -452,2427 +440,6 @@ fn dedup_tool_results(mut messages: Vec) -> Vec { messages } -#[cfg(test)] -#[allow(clippy::items_after_test_module)] -mod tests { - use super::super::ContextManager; - use super::*; - use stakpak_shared::models::integrations::openai::{ - FunctionCall, MessageContent, Role, ToolCall, - }; - use stakpak_shared::models::llm::{LLMMessageContent, LLMMessageTypedContent}; - - fn create_context_manager() -> TaskBoardContextManager { - TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 2, // Only keep last 2 assistant messages untrimmed - context_budget_threshold: 0.8, - }) - } - - fn create_tool_call_msg(id: &str) -> ChatMessage { - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("Thinking...".to_string())), - tool_calls: Some(vec![ToolCall { - id: id.to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "test_tool".to_string(), - arguments: "{}".to_string(), - }, - metadata: None, - }]), - ..Default::default() - } - } - - fn create_tool_call_msg_with_ids(ids: &[&str]) -> ChatMessage { - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("Thinking...".to_string())), - tool_calls: Some( - ids.iter() - .map(|id| ToolCall { - id: (*id).to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "test_tool".to_string(), - arguments: "{}".to_string(), - }, - metadata: None, - }) - .collect(), - ), - ..Default::default() - } - } - - fn create_tool_result_msg(id: &str, content: &str) -> ChatMessage { - ChatMessage { - role: Role::Tool, - tool_call_id: Some(id.to_string()), - content: Some(MessageContent::String(content.to_string())), - ..Default::default() - } - } - - #[test] - fn test_reduce_context_preserves_messages() { - let cm = create_context_manager(); - let messages = vec![ChatMessage { - role: Role::User, - content: Some(MessageContent::String("Hello".to_string())), - ..Default::default() - }]; - - let reduced = cm.reduce_context(messages); - assert_eq!(reduced.len(), 1); - assert_eq!(reduced[0].role, "user"); - } - - #[test] - fn test_reduce_context_preserves_old_tool_results() { - let cm = create_context_manager(); - let messages = vec![ - create_tool_call_msg("call_1"), - create_tool_result_msg("call_1", "Result 1 - detailed output from first tool"), - create_tool_call_msg("call_2"), - create_tool_result_msg("call_2", "Result 2"), - create_tool_call_msg("call_3"), - create_tool_result_msg("call_3", "Result 3"), - ]; - - let reduced = cm.reduce_context(messages); - - // call_1 result should still have its original content (not truncated) - let result_1 = &reduced[1]; // tool result for call_1 - match &result_1.content { - LLMMessageContent::List(parts) => { - if let LLMMessageTypedContent::ToolResult { content, .. } = &parts[0] { - assert!( - !content.contains("truncated"), - "Result should NOT be truncated, got: {}", - content - ); - assert!( - content.contains("Result 1"), - "Original content should be preserved" - ); - } - } - LLMMessageContent::String(s) => { - assert!(!s.contains("truncated"), "Result should NOT be truncated"); - } - } - } - - #[test] - fn test_reduce_context_preserves_large_assistant_thoughts() { - let cm = create_context_manager(); // size_limit = 10 - let long_thought = "This is a very long thought that exceeds the size limit"; - - let messages = vec![ - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(long_thought.to_string())), - tool_calls: Some(vec![ToolCall { - id: "call_1".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "t".to_string(), - arguments: "{}".to_string(), - }, - metadata: None, - }]), - ..Default::default() - }, - create_tool_result_msg("call_1", "res"), - create_tool_call_msg("call_2"), - create_tool_call_msg("call_3"), - ]; - - let reduced = cm.reduce_context(messages); - if let LLMMessageContent::List(parts) = &reduced[0].content { - let has_text = parts - .iter() - .any(|p| matches!(p, LLMMessageTypedContent::Text { .. })); - assert!(has_text, "Large thought should be preserved"); - } else { - panic!("Expected list content for assistant message with tool calls"); - } - } - - #[test] - fn test_reduce_context_removes_checkpoint_id_tags() { - let cm = create_context_manager(); - let messages = vec![ChatMessage { - role: Role::User, - content: Some(MessageContent::String( - "abc-123\nHello, world!".to_string(), - )), - ..Default::default() - }]; - - let reduced = cm.reduce_context(messages); - - match &reduced[0].content { - LLMMessageContent::String(s) => { - assert!( - !s.contains("checkpoint_id"), - "checkpoint_id tag should be removed" - ); - assert!( - s.contains("Hello, world!"), - "actual content should be preserved" - ); - } - _ => panic!("Expected string content"), - } - } - - #[test] - fn test_reduce_context_merges_consecutive_tool_results_for_same_assistant_turn() { - let cm = create_context_manager(); - let messages = vec![ - create_tool_call_msg_with_ids(&["call_1", "call_2"]), - create_tool_result_msg("call_1", "Result 1"), - create_tool_result_msg("call_2", "Result 2"), - ]; - - let reduced = cm.reduce_context(messages); - - assert_eq!(reduced.len(), 2); - assert_eq!(reduced[1].role, "tool"); - match &reduced[1].content { - LLMMessageContent::List(parts) => { - let tool_results: Vec<_> = parts - .iter() - .filter_map(|part| { - if let LLMMessageTypedContent::ToolResult { - tool_use_id, - content, - } = part - { - Some((tool_use_id.as_str(), content.as_str())) - } else { - None - } - }) - .collect(); - - assert_eq!( - tool_results, - vec![("call_1", "Result 1"), ("call_2", "Result 2")] - ); - } - other => panic!( - "Expected list content for merged tool results, got {:?}", - other - ), - } - } - - #[test] - fn test_reduce_context_deduplicates_tool_results_keeping_last() { - let cm = create_context_manager(); - let messages = vec![ - create_tool_call_msg("call_1"), - create_tool_result_msg("call_1", "old_result"), - create_tool_result_msg("call_1", "new_result"), - ]; - - let reduced = cm.reduce_context(messages); - - assert_eq!(reduced.len(), 2); - assert_eq!(reduced[1].role, "tool"); - match &reduced[1].content { - LLMMessageContent::List(parts) => { - let tool_results: Vec<_> = parts - .iter() - .filter_map(|part| { - if let LLMMessageTypedContent::ToolResult { - tool_use_id, - content, - } = part - { - Some((tool_use_id.as_str(), content.as_str())) - } else { - None - } - }) - .collect(); - assert_eq!(tool_results, vec![("call_1", "new_result")]); - } - other => panic!( - "Expected list content for deduplicated tool results, got {:?}", - other - ), - } - } - - #[test] - fn test_reduce_context_preserves_tool_role() { - let cm = create_context_manager(); - let messages = vec![ - create_tool_call_msg("call_1"), - create_tool_result_msg("call_1", "Result content"), - ]; - - let reduced = cm.reduce_context(messages); - - let result_msg = &reduced[1]; - assert_eq!( - result_msg.role, "tool", - "Tool results should preserve tool role (provider layer handles format conversion)" - ); - } - - // ========================================================================= - // Token estimation tests - // ========================================================================= - - #[test] - fn test_estimate_tokens_simple_message() { - let messages = vec![LLMMessage { - role: "user".to_string(), - content: LLMMessageContent::String("Hello world".to_string()), - }]; - let tokens = TaskBoardContextManager::estimate_tokens(&messages); - // "Hello world" = 11 bytes / BYTES_PER_TOKEN + 8 msg overhead, then 5% buffer - let content_tokens = (11.0 / BYTES_PER_TOKEN).ceil() as u64; - let raw = content_tokens + 8; - let expected = (raw as f64 * 1.05).ceil() as u64; - assert_eq!(tokens, expected); - } - - #[test] - fn test_estimate_tokens_multiple_messages() { - let messages = vec![ - LLMMessage { - role: "user".to_string(), - content: LLMMessageContent::String("Hello".to_string()), - }, - LLMMessage { - role: "assistant".to_string(), - content: LLMMessageContent::String("World".to_string()), - }, - ]; - let tokens = TaskBoardContextManager::estimate_tokens(&messages); - // Each: 5 bytes / BYTES_PER_TOKEN + 8 msg overhead, total raw, then 5% buffer - let per_msg = (5.0 / BYTES_PER_TOKEN).ceil() as u64 + 8; - let raw = per_msg * 2; - let expected = (raw as f64 * 1.05).ceil() as u64; - assert_eq!(tokens, expected); - } - - #[test] - fn test_estimate_tokens_with_tool_results() { - let messages = vec![LLMMessage { - role: "tool".to_string(), - content: LLMMessageContent::List(vec![LLMMessageTypedContent::ToolResult { - tool_use_id: "tc_1".to_string(), - content: "A".repeat(400), // 400 bytes - }]), - }]; - let tokens = TaskBoardContextManager::estimate_tokens(&messages); - // 400 content + 30 structural = 430 bytes / BYTES_PER_TOKEN - // + 3 per-part overhead (1 part) + 8 msg overhead, then 5% buffer - let part_tokens = (430.0 / BYTES_PER_TOKEN).ceil() as u64; - let raw = part_tokens + 3 + 8; - let expected = (raw as f64 * 1.05).ceil() as u64; - assert_eq!(tokens, expected); - } - - // ========================================================================= - // Budget-aware trimming tests - // ========================================================================= - - #[test] - fn test_no_trimming_when_under_threshold() { - let cm = create_context_manager(); - let messages = vec![ - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("Hello".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("Hi there".to_string())), - ..Default::default() - }, - ]; - - let (result, metadata) = cm.reduce_context_with_budget(messages, 200_000, None, None); - assert_eq!(result.len(), 2); - // No trimming should have occurred - assert!(metadata.is_none()); - // Content should be preserved - match &result[0].content { - LLMMessageContent::String(s) => assert_eq!(s, "Hello"), - _ => panic!("Expected string content"), - } - } - - #[test] - fn test_trimming_triggers_at_threshold() { - let cm = create_context_manager(); - // Create messages that exceed 80% of a small context window - let mut messages = Vec::new(); - for i in 0..20 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "Message {}: {}", - i, - "x".repeat(100) - ))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Response {}: {}", - i, - "y".repeat(100) - ))), - ..Default::default() - }); - } - - // Use a small context window so trimming triggers - let (result, metadata) = cm.reduce_context_with_budget(messages, 100, None, None); - - // Should have metadata with trimmed index - assert!(metadata.is_some()); - let meta = metadata.unwrap(); - assert!(meta.get("trimmed_up_to_message_index").is_some()); - - // With keep_last_n_assistant_messages = 2, the trim boundary is placed - // just before the 2nd-from-last assistant message. In a 40-message - // alternating sequence (user/assistant × 20), the last 2 assistants are - // at indices 37 and 39, so trim_end = 37 and messages 37-39 are untrimmed. - let trimmed_idx = meta["trimmed_up_to_message_index"].as_u64().unwrap() as usize; - for msg in &result[trimmed_idx..] { - if let LLMMessageContent::String(s) = &msg.content { - assert_ne!( - s, TRIMMED_CONTENT_PLACEHOLDER, - "Messages after trim boundary should not be trimmed" - ); - } - } - - // Earlier assistant messages should be trimmed, user messages preserved - let first_assistant = &result[1]; // index 1 is assistant - if let LLMMessageContent::String(s) = &first_assistant.content { - assert_eq!( - s, TRIMMED_CONTENT_PLACEHOLDER, - "Early assistant messages should be trimmed" - ); - } - - // User messages should NOT be trimmed - let first_user = &result[0]; // index 0 is user - if let LLMMessageContent::String(s) = &first_user.content { - assert_ne!( - s, TRIMMED_CONTENT_PLACEHOLDER, - "User messages should never be trimmed" - ); - } - } - - #[test] - fn test_already_trimmed_messages_not_re_trimmed() { - let cm = create_context_manager(); - let mut messages = Vec::new(); - for i in 0..10 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!("Msg {}", i))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!("Resp {}", i))), - ..Default::default() - }); - } - - // First trim - let (_, metadata) = cm.reduce_context_with_budget(messages.clone(), 100, None, None); - let trimmed_index = metadata - .as_ref() - .unwrap() - .get("trimmed_up_to_message_index") - .unwrap() - .as_u64() - .unwrap(); - - // Add more messages and trim again - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String("New message".to_string())), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("New response".to_string())), - ..Default::default() - }); - - let (_, metadata2) = cm.reduce_context_with_budget(messages, 100, metadata, None); - let new_trimmed_index = metadata2 - .as_ref() - .unwrap() - .get("trimmed_up_to_message_index") - .unwrap() - .as_u64() - .unwrap(); - - // New trimmed index should be >= old one (never goes backward) - assert!( - new_trimmed_index >= trimmed_index, - "Trimmed index should not decrease: {} < {}", - new_trimmed_index, - trimmed_index - ); - } - - #[test] - fn test_trimming_preserves_message_structure() { - let cm = create_context_manager(); - let messages = vec![ - create_tool_call_msg("call_1"), - create_tool_result_msg("call_1", "x".repeat(200).as_str()), - create_tool_call_msg("call_2"), - create_tool_result_msg("call_2", "y".repeat(200).as_str()), - // These should be kept - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("Recent".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("Response".to_string())), - ..Default::default() - }, - ]; - - let (result, _) = cm.reduce_context_with_budget(messages, 50, None, None); - - // All messages should still be present (structure preserved) - assert!(!result.is_empty()); - - // Verify roles alternate correctly - for msg in &result { - assert!( - msg.role == "user" - || msg.role == "assistant" - || msg.role == "tool" - || msg.role == "system", - "Invalid role: {}", - msg.role - ); - } - } - - // ========================================================================= - // Full lifecycle integration tests - // ========================================================================= - - /// Helper: check if a message has been trimmed - fn is_trimmed(msg: &LLMMessage) -> bool { - match &msg.content { - LLMMessageContent::String(s) => s == TRIMMED_CONTENT_PLACEHOLDER, - LLMMessageContent::List(parts) => parts.iter().all(|p| match p { - LLMMessageTypedContent::Text { text } => text == TRIMMED_CONTENT_PLACEHOLDER, - LLMMessageTypedContent::ToolResult { content, .. } => { - content == TRIMMED_CONTENT_PLACEHOLDER - } - LLMMessageTypedContent::ToolCall { .. } => true, // tool calls are never trimmed - LLMMessageTypedContent::Image { .. } => true, - }), - } - } - - #[test] - fn test_full_lifecycle_trim_save_resume_trim() { - let cm = create_context_manager(); // threshold = 0.8 - let context_window = 200u64; // small window to force trimming - - // === Turn 1-5: Build up conversation === - let mut messages = Vec::new(); - for i in 0..5 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "User turn {}: {}", - i, - "a".repeat(80) - ))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Assistant turn {}: {}", - i, - "b".repeat(80) - ))), - ..Default::default() - }); - } - - // === First inference: should trigger trimming === - let (result1, metadata1) = - cm.reduce_context_with_budget(messages.clone(), context_window, None, None); - - assert!( - metadata1.is_some(), - "Trimming should have produced metadata" - ); - let trimmed_idx_1 = metadata1.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!(trimmed_idx_1 > 0, "Should have trimmed some messages"); - - // Verify: non-user messages before trimmed_idx are trimmed - for (i, msg) in result1.iter().enumerate() { - if i < trimmed_idx_1 && msg.role != "user" { - assert!( - is_trimmed(msg), - "Non-user message {} should be trimmed (trimmed_idx={}), content: {:?}", - i, - trimmed_idx_1, - msg.content - ); - } - if i < trimmed_idx_1 && msg.role == "user" { - assert!( - !is_trimmed(msg), - "User message {} should NOT be trimmed, content: {:?}", - i, - msg.content - ); - } - } - - // Verify: messages after trim boundary are NOT trimmed - for (i, msg) in result1.iter().enumerate().skip(trimmed_idx_1) { - assert!( - !is_trimmed(msg), - "Message {} (after trim boundary {}) should NOT be trimmed", - i, - trimmed_idx_1 - ); - } - - // === Simulate checkpoint save/load: metadata round-trips through JSON === - let saved_metadata_json = serde_json::to_string(&metadata1).unwrap(); - let loaded_metadata: Option = - serde_json::from_str(&saved_metadata_json).unwrap(); - - // === Turn 6-7: Add more messages (simulating continued conversation) === - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "User turn 5: {}", - "c".repeat(80) - ))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Assistant turn 5: {}", - "d".repeat(80) - ))), - ..Default::default() - }); - - // === Second inference with loaded metadata === - let (result2, metadata2) = - cm.reduce_context_with_budget(messages.clone(), context_window, loaded_metadata, None); - - let trimmed_idx_2 = metadata2.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - // Trimmed index should advance (more messages to trim) - assert!( - trimmed_idx_2 >= trimmed_idx_1, - "Trimmed index should not decrease: {} < {}", - trimmed_idx_2, - trimmed_idx_1 - ); - - // Verify: the previously-trimmed prefix is still trimmed (stable for caching) - // User messages are never trimmed, only assistant/tool messages - for (i, msg) in result2.iter().enumerate().take(trimmed_idx_1) { - if msg.role != "user" { - assert!( - is_trimmed(msg), - "Previously trimmed non-user message {} should still be trimmed", - i - ); - } - } - - // Verify: messages after trim boundary are NOT trimmed - for (i, msg) in result2.iter().enumerate().skip(trimmed_idx_2) { - assert!( - !is_trimmed(msg), - "Message {} (after trim boundary {}) should NOT be trimmed after second trim", - i, - trimmed_idx_2 - ); - } - - // Verify: total message count increased - assert!( - result2.len() > result1.len(), - "Should have more messages after adding turns" - ); - } - - #[test] - fn test_tool_call_ids_preserved_after_trimming() { - let cm = create_context_manager(); - - let messages = vec![ - // Old tool interaction (should be trimmed) - create_tool_call_msg_with_ids(&["tc_old_1", "tc_old_2"]), - create_tool_result_msg("tc_old_1", &"x".repeat(200)), - create_tool_result_msg("tc_old_2", &"y".repeat(200)), - // Recent tool interaction (should be kept) - create_tool_call_msg_with_ids(&["tc_new_1"]), - create_tool_result_msg("tc_new_1", "recent result"), - // Recent conversation - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("What happened?".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("Here's what I found".to_string())), - ..Default::default() - }, - ]; - - let (result, _) = cm.reduce_context_with_budget(messages, 80, None, None); - - // Find all ToolCall parts and verify they still have their IDs - for msg in &result { - if let LLMMessageContent::List(parts) = &msg.content { - for part in parts { - if let LLMMessageTypedContent::ToolCall { id, name, .. } = part { - assert!(!id.is_empty(), "ToolCall ID should be preserved"); - assert!(!name.is_empty(), "ToolCall name should be preserved"); - } - if let LLMMessageTypedContent::ToolResult { tool_use_id, .. } = part { - assert!( - !tool_use_id.is_empty(), - "ToolResult tool_use_id should be preserved" - ); - } - } - } - } - } - - #[test] - fn test_system_messages_never_trimmed() { - let cm = create_context_manager(); - - // Simulate a system message at the start (though normally added by hook, - // test that the trimmer skips system role) - let mut messages = vec![ChatMessage { - role: Role::User, // system messages are added by hook, not in ChatMessage - content: Some(MessageContent::String("System-like setup".to_string())), - ..Default::default() - }]; - - for i in 0..10 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "Msg {}: {}", - i, - "z".repeat(100) - ))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Resp {}: {}", - i, - "w".repeat(100) - ))), - ..Default::default() - }); - } - - let (result, metadata) = cm.reduce_context_with_budget(messages, 100, None, None); - - assert!(metadata.is_some(), "Should have trimmed"); - - // Verify no system-role or user-role messages were trimmed - for msg in &result { - if msg.role == "system" { - assert!(!is_trimmed(msg), "System messages must never be trimmed"); - } - if msg.role == "user" { - assert!(!is_trimmed(msg), "User messages must never be trimmed"); - } - } - } - - #[test] - fn test_trimming_with_empty_metadata_object() { - let cm = create_context_manager(); - // Use large assistant messages and small user messages so that trimming - // the first assistant (per keep_last_n=2) is sufficient to get under budget. - let messages = vec![ - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u1".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("y".repeat(500))), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u2".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("w".repeat(500))), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("recent".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("response".to_string())), - ..Default::default() - }, - ]; - - // Context window sized so total (~380 tokens) exceeds 80% threshold (~320) - // but after trimming 1 assistant per keep_last_n=2, remaining fits. - let (result, metadata) = - cm.reduce_context_with_budget(messages, 400, Some(serde_json::json!({})), None); - - assert!(metadata.is_some()); - let meta = metadata.unwrap(); - assert!( - meta.get("trimmed_up_to_message_index").is_some(), - "Should have set trimmed_up_to_message_index" - ); - - // Verify trimming happened on assistant messages (user messages are preserved) - assert!( - !is_trimmed(&result[0]), - "First user message should NOT be trimmed" - ); - assert!( - is_trimmed(&result[1]), - "First assistant message should be trimmed" - ); - - // Verify last messages preserved - let len = result.len(); - assert!( - !is_trimmed(&result[len - 1]), - "Last message should not be trimmed" - ); - } - - /// Verify that keep_last_n_assistant_messages counts only assistant-role - /// messages, not all messages. With interleaved user/tool messages between - /// assistants, the trim boundary should be placed based on assistant count. - #[test] - fn test_keep_last_n_counts_only_assistant_messages() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 2, - context_budget_threshold: 0.8, - }); - - // Build: user, assistant, user, assistant, user, user, assistant - // After merge_consecutive_same_role, the two consecutive user messages - // get merged, so the sequence is: - // [0] user, [1] asst(a1), [2] user, [3] asst(a2), [4] user(merged), [5] asst(a3) - // Assistants at indices: 1, 3, 5 - // keep_last_n_assistant_messages = 2 → trim_end = index of 2nd-from-last asst = 3 - // So messages 0..3 are in the trim zone, messages 3..6 are untrimmed. - // - // Use large assistant messages so trimming them actually reduces token count - // below the budget threshold. User messages are small so the remaining - // content fits within budget after trimming the prefix. - let messages = vec![ - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u1".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a1".to_string() + &"y".repeat(300))), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u2".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a2".to_string() + &"y".repeat(300))), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u3".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u4".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a3".to_string() + &"y".repeat(300))), - ..Default::default() - }, - ]; - - // Context window sized so that: - // - Total tokens (~380) exceed 80% threshold (~320) → triggers trimming - // - After trimming 1 assistant per keep_last_n=2 (~100 tokens saved), - // remaining (~280) fits under threshold - let (result, metadata) = cm.reduce_context_with_budget(messages, 400, None, None); - assert!(metadata.is_some(), "Should trigger trimming"); - - let trimmed_idx = metadata.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - // The 2nd-from-last assistant is "a2". After merge_consecutive_same_role, - // the two consecutive user messages get merged, so the sequence is: - // [0] user, [1] asst(a1), [2] user, [3] asst(a2), [4] user(merged), [5] asst(a3) - // Assistants at 1, 3, 5. 2nd-from-last = index 3. trim_end = 3. - assert_eq!( - trimmed_idx, 3, - "Trim boundary should be at the 2nd-from-last assistant (index 3)" - ); - - // a1 (index 1) should be trimmed - assert!( - is_trimmed(&result[1]), - "First assistant (a1) should be trimmed" - ); - - // a2 (index 3) should NOT be trimmed (it's the boundary — kept) - assert!( - !is_trimmed(&result[3]), - "Second assistant (a2) should NOT be trimmed (within keep window)" - ); - - // a3 (index 5) should NOT be trimmed - assert!( - !is_trimmed(&result[5]), - "Third assistant (a3) should NOT be trimmed" - ); - - // All user messages should never be trimmed - for msg in &result { - if msg.role == "user" { - assert!(!is_trimmed(msg), "User messages should never be trimmed"); - } - } - } - - /// Verify lazy trimming: when under threshold but with existing - /// trimmed_up_to_message_index, re-apply old trim without advancing. - #[test] - fn test_lazy_trimming_reapplies_without_advancing() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 2, - context_budget_threshold: 0.8, - }); - - // Build 10 turns of user/assistant - let mut messages = Vec::new(); - for i in 0..10 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "Msg {}: {}", - i, - "x".repeat(100) - ))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Resp {}: {}", - i, - "y".repeat(100) - ))), - ..Default::default() - }); - } - - // First call: small window → triggers trimming, establishes trim index - let (_, metadata1) = cm.reduce_context_with_budget(messages.clone(), 100, None, None); - assert!(metadata1.is_some(), "First call should trigger trimming"); - let trimmed_idx_1 = metadata1.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!(trimmed_idx_1 > 0); - - // Second call: LARGE window (under threshold) but with existing metadata. - // Should re-apply trimming up to trimmed_idx_1 but NOT advance it. - let (result2, metadata2) = - cm.reduce_context_with_budget(messages.clone(), 200_000, metadata1.clone(), None); - - assert!(metadata2.is_some(), "Should still return metadata"); - let trimmed_idx_2 = metadata2.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - // Index should NOT have advanced (we're under threshold) - assert_eq!( - trimmed_idx_2, trimmed_idx_1, - "Trim index should not advance when under threshold" - ); - - // But the old prefix should still be trimmed (re-applied) - for (i, msg) in result2.iter().enumerate() { - if i < trimmed_idx_1 && msg.role != "user" { - assert!( - is_trimmed(msg), - "Non-user message {} should still be trimmed (re-applied from metadata)", - i - ); - } - } - - // Messages after the trim boundary should NOT be trimmed - for (i, msg) in result2.iter().enumerate() { - if i >= trimmed_idx_1 { - assert!( - !is_trimmed(msg), - "Message {} (after trim boundary) should NOT be trimmed", - i - ); - } - } - - // Third call: small window again → should advance the index - // Add more messages first to ensure we're over budget - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String("x".repeat(200))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("y".repeat(200))), - ..Default::default() - }); - - let (_, metadata3) = cm.reduce_context_with_budget(messages, 100, metadata2, None); - let trimmed_idx_3 = metadata3.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - assert!( - trimmed_idx_3 >= trimmed_idx_1, - "Trim index should advance when over threshold again: {} < {}", - trimmed_idx_3, - trimmed_idx_1 - ); - } - - /// Tool messages between assistants should be trimmed in the prefix, - /// and the trim boundary should still be based on assistant count only. - #[test] - fn test_trim_boundary_with_tool_messages() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 1, - context_budget_threshold: 0.8, - }); - - // Realistic agent flow: - // user → assistant(tool_call) → tool(result) → assistant → user → assistant(tool_call) → tool(result) → assistant - // - // After merge_consecutive_same_role the LLM sequence is: - // [0] user [1] assistant(tc_1) [2] tool(result_1) [3] assistant - // [4] user [5] assistant(tc_2) [6] tool(result_2) [7] assistant - // - // Assistants at: 1, 3, 5, 7. keep_last_n = 1 → trim_end = index 7. - // So messages 0..7 are in the trim zone, only message 7 is untrimmed. - // Within the trim zone: user messages (0, 4) are preserved; assistant (1, 3, 5) - // and tool (2, 6) messages are trimmed. - let messages = vec![ - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u1".to_string())), - ..Default::default() - }, - create_tool_call_msg("tc_1"), - create_tool_result_msg("tc_1", &"r".repeat(200)), - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("follow-up 1".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u2".to_string())), - ..Default::default() - }, - create_tool_call_msg("tc_2"), - create_tool_result_msg("tc_2", &"r".repeat(200)), - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("follow-up 2".to_string())), - ..Default::default() - }, - ]; - - // Compute a context window that triggers trimming: use 90% of the - // estimated total so the 80% threshold is exceeded. - let pre_messages: Vec<_> = messages.iter().cloned().map(LLMMessage::from).collect(); - let total_estimate = TaskBoardContextManager::estimate_tokens(&pre_messages); - let context_window = (total_estimate as f64 / 0.9).ceil() as u64; - - let (result, metadata) = - cm.reduce_context_with_budget(messages, context_window, None, None); - assert!(metadata.is_some()); - let trimmed_idx = metadata.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - // Verify every message's trim state - for (i, msg) in result.iter().enumerate() { - if i >= trimmed_idx { - // After boundary: nothing trimmed - assert!( - !is_trimmed(msg), - "Message {} (role={}) after trim boundary should NOT be trimmed", - i, - msg.role - ); - } else if msg.role == "user" { - // Before boundary but user: never trimmed - assert!( - !is_trimmed(msg), - "User message {} before trim boundary should NOT be trimmed", - i - ); - } else { - // Before boundary, assistant or tool: trimmed - assert!( - is_trimmed(msg), - "Non-user message {} (role={}) before trim boundary should be trimmed", - i, - msg.role - ); - } - } - - // The last assistant should be the only untrimmed assistant - let untrimmed_assistants: Vec<_> = result - .iter() - .enumerate() - .filter(|(_, m)| m.role == "assistant" && !is_trimmed(m)) - .map(|(i, _)| i) - .collect(); - assert_eq!( - untrimmed_assistants.len(), - 1, - "Exactly 1 assistant should be untrimmed, got {:?}", - untrimmed_assistants - ); - } - - /// When there are fewer assistant messages than keep_last_n but we're - /// over budget, the budget constraint takes priority — the oldest - /// assistant gets trimmed to make progress. keep_last_n is best-effort. - #[test] - fn test_fewer_assistants_than_keep_last_n_budget_overrides() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 10, - context_budget_threshold: 0.8, - }); - - // Only 3 assistant messages but keep_last_n = 10 - let messages = vec![ - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("x".repeat(300))), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a1".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("x".repeat(300))), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a2".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("x".repeat(300))), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a3".to_string())), - ..Default::default() - }, - ]; - - // Small window — over budget despite fewer assistants than keep_last_n - let (result, metadata) = cm.reduce_context_with_budget(messages, 50, None, None); - - // Budget overrides keep_last_n: at least the oldest assistant is trimmed - assert!(metadata.is_some(), "Should produce metadata when trimming"); - let trimmed_idx = metadata.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!( - trimmed_idx > 0, - "Budget should force trimming even with fewer assistants than keep_last_n" - ); - - // All user messages should be preserved - for msg in &result { - if msg.role == "user" { - assert!(!is_trimmed(msg), "User messages should never be trimmed"); - } - } - } - - /// With keep_last_n_assistant_messages = 0, every assistant/tool message - /// in the entire history should be trimmed. - #[test] - fn test_keep_zero_trims_all_assistant_messages() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 0, - context_budget_threshold: 0.8, - }); - - let messages = vec![ - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("x".repeat(200))), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a1".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("x".repeat(200))), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a2".to_string())), - ..Default::default() - }, - ]; - - let (result, metadata) = cm.reduce_context_with_budget(messages, 50, None, None); - assert!(metadata.is_some()); - - // Every assistant message should be trimmed - for msg in &result { - if msg.role == "assistant" { - assert!( - is_trimmed(msg), - "All assistant messages should be trimmed when keep_last_n = 0" - ); - } - } - - // Every user message should be preserved - for msg in &result { - if msg.role == "user" { - assert!(!is_trimmed(msg), "User messages should never be trimmed"); - } - } - } - - /// With keep_last_n_assistant_messages = 1, only the very last assistant - /// message should be kept untrimmed. - #[test] - fn test_keep_one_preserves_only_last_assistant() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 1, - context_budget_threshold: 0.8, - }); - - // Use large assistant messages and small user messages so that trimming - // the first 2 assistants (per keep_last_n=1) is sufficient to get under budget. - let messages = vec![ - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u1".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a1".to_string() + &"y".repeat(300))), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u2".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("a2".to_string() + &"y".repeat(300))), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("u3".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String( - "a3_last".to_string() + &"y".repeat(300), - )), - ..Default::default() - }, - ]; - - // Context window sized so total (~380 tokens) exceeds 80% threshold (~280) - // but after trimming 2 assistants per keep_last_n=1, remaining fits. - let (result, metadata) = cm.reduce_context_with_budget(messages, 350, None, None); - assert!(metadata.is_some()); - - // Collect assistant messages with their trim state - let assistants: Vec<_> = result - .iter() - .enumerate() - .filter(|(_, m)| m.role == "assistant") - .map(|(i, m)| (i, is_trimmed(m))) - .collect(); - - // All but the last assistant should be trimmed - for &(idx, trimmed) in &assistants[..assistants.len() - 1] { - assert!( - trimmed, - "Assistant at index {} should be trimmed (not the last)", - idx - ); - } - - // The last assistant should NOT be trimmed - let (last_idx, last_trimmed) = assistants.last().unwrap(); - assert!( - !last_trimmed, - "Last assistant at index {} should NOT be trimmed", - last_idx - ); - } - - /// The trim index should never go backward, even if a new call computes - /// a smaller trim_end than the stored prev_trimmed_up_to. - #[test] - fn test_trim_index_never_goes_backward() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 2, - context_budget_threshold: 0.8, - }); - - // 10 turns → 20 messages, small window → establishes a trim index - let mut messages = Vec::new(); - for i in 0..10 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "u{}: {}", - i, - "x".repeat(100) - ))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "a{}: {}", - i, - "y".repeat(100) - ))), - ..Default::default() - }); - } - - let (_, metadata1) = cm.reduce_context_with_budget(messages.clone(), 100, None, None); - let idx1 = metadata1.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!(idx1 > 0); - - // Now use a LARGER keep_last_n (would compute a smaller trim_end) - // but pass the existing metadata with the higher index. - let cm_generous = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 8, - context_budget_threshold: 0.8, - }); - - let (_, metadata2) = cm_generous.reduce_context_with_budget(messages, 100, metadata1, None); - let idx2 = metadata2.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - assert!( - idx2 >= idx1, - "Trim index must never decrease: got {} < {}", - idx2, - idx1 - ); - } - - /// Tool definitions should be included in the budget check, pushing - /// an otherwise-under-threshold conversation over the limit. - #[test] - fn test_tool_overhead_pushes_over_threshold() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 2, - context_budget_threshold: 0.8, - }); - - // Build a conversation that's just under threshold without tools - let messages = vec![ - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("x".repeat(100))), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("y".repeat(100))), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("x".repeat(100))), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("y".repeat(100))), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("x".repeat(100))), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("y".repeat(100))), - ..Default::default() - }, - ]; - - // Estimate tokens for these messages - let llm_msgs: Vec<_> = messages.iter().cloned().map(LLMMessage::from).collect(); - let msg_tokens = TaskBoardContextManager::estimate_tokens(&llm_msgs); - - // Pick a context window where messages alone are under 80% threshold - // but messages + tool overhead exceed it. - let context_window = (msg_tokens as f64 / 0.75) as u64; // 80% of this ≈ msg_tokens * 1.067 - let threshold = (context_window as f32 * 0.8) as u64; - - // Without tools: under threshold → no trimming - let (_, meta_no_tools) = - cm.reduce_context_with_budget(messages.clone(), context_window, None, None); - assert!( - meta_no_tools.is_none(), - "Without tools, should be under threshold (tokens={}, threshold={})", - msg_tokens, - threshold - ); - - // With many large tool definitions: over threshold → trimming triggers - let big_tools: Vec = (0..50) - .map(|i| LLMTool { - name: format!("tool_{}", i), - description: "x".repeat(200), - input_schema: serde_json::json!({ - "type": "object", - "properties": { - "arg1": { "type": "string", "description": "x".repeat(200) }, - "arg2": { "type": "string", "description": "x".repeat(200) }, - } - }), - }) - .collect(); - - let tool_overhead = TaskBoardContextManager::estimate_tool_overhead(Some(&big_tools)); - assert!( - msg_tokens + tool_overhead > threshold, - "Tools should push over threshold: msg={} + tools={} > threshold={}", - msg_tokens, - tool_overhead, - threshold - ); - - let (result, meta_with_tools) = - cm.reduce_context_with_budget(messages, context_window, None, Some(&big_tools)); - assert!( - meta_with_tools.is_some(), - "With tool overhead, should exceed threshold" - ); - - // Verify some assistant message got trimmed - let has_trimmed_assistant = result - .iter() - .any(|m| m.role == "assistant" && is_trimmed(m)); - assert!( - has_trimmed_assistant, - "At least one assistant message should be trimmed" - ); - } - - /// Realistic multi-turn tool-call flow verifying that the trim boundary - /// is based on assistant count and that interleaved tool messages are - /// correctly trimmed/preserved relative to the boundary. - #[test] - fn test_realistic_tool_call_flow_trim_boundary() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 3, - context_budget_threshold: 0.8, - }); - - // 5 turns of: user → assistant(tool_call) → tool(result) → assistant(follow-up) - // Each turn produces 4 ChatMessages. After merge, each turn is: - // user, assistant(tc), tool(result), assistant(follow-up) - // So 5 turns = 20 LLM messages. - // Assistants at indices: 1, 3, 5, 7, 9, 11, 13, 15, 17, 19 - // (every assistant(tc) and assistant(follow-up)) - // keep_last_n = 3 → 3rd-from-last assistant = index 15 → trim_end = 15 - let mut messages = Vec::new(); - for turn in 0..5 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!("user turn {}", turn))), - ..Default::default() - }); - messages.push(create_tool_call_msg(&format!("tc_{}", turn))); - messages.push(create_tool_result_msg( - &format!("tc_{}", turn), - &format!("result {}: {}", turn, "r".repeat(200)), - )); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!("follow-up {}", turn))), - ..Default::default() - }); - } - - let (result, metadata) = cm.reduce_context_with_budget(messages, 700, None, None); - assert!(metadata.is_some()); - let trimmed_idx = metadata.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - // Count untrimmed assistants — should be exactly 3 - let untrimmed_assistants: Vec = result - .iter() - .enumerate() - .filter(|(_, m)| m.role == "assistant" && !is_trimmed(m)) - .map(|(i, _)| i) - .collect(); - assert_eq!( - untrimmed_assistants.len(), - 3, - "Should have exactly 3 untrimmed assistants, got {:?}", - untrimmed_assistants - ); - - // All untrimmed assistants should be at or after the trim boundary - for &idx in &untrimmed_assistants { - assert!( - idx >= trimmed_idx, - "Untrimmed assistant at {} should be >= trim boundary {}", - idx, - trimmed_idx - ); - } - - // Tool messages before the boundary should be trimmed - let trimmed_tools: Vec = result - .iter() - .enumerate() - .filter(|(i, m)| *i < trimmed_idx && m.role == "tool" && is_trimmed(m)) - .map(|(i, _)| i) - .collect(); - let total_tools_before_boundary = result - .iter() - .enumerate() - .filter(|(i, m)| *i < trimmed_idx && m.role == "tool") - .count(); - assert_eq!( - trimmed_tools.len(), - total_tools_before_boundary, - "All tool messages before trim boundary should be trimmed" - ); - - // Tool messages at or after the boundary should NOT be trimmed - for (i, msg) in result.iter().enumerate() { - if i >= trimmed_idx && msg.role == "tool" { - assert!( - !is_trimmed(msg), - "Tool message {} at/after trim boundary should NOT be trimmed", - i - ); - } - } - - // User messages should never be trimmed regardless of position - for (i, msg) in result.iter().enumerate() { - if msg.role == "user" { - assert!( - !is_trimmed(msg), - "User message {} should never be trimmed", - i - ); - } - } - } - - /// Verbose simulation of a realistic multi-turn agent session. - /// Run with `cargo test -p stakpak-api test_verbose_session_simulation -- --nocapture` - #[test] - fn test_verbose_session_simulation() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 50, - context_budget_threshold: 0.8, - }); - - // Simulate a 200k context window model (like Claude) - // With ~1000 chars per tool result, each turn is ~1200 chars = ~300 tokens - // 30 turns = ~9000 tokens. To trigger trimming at 80%, we need context_window - // where 80% < 9000, so context_window < 11250. Use 10000 to simulate a - // session approaching the limit. - let context_window = 10_000u64; - let mut messages: Vec = Vec::new(); - let metadata: Option = None; - - // === Simulate 30 turns of agent conversation with tool calls === - // Each turn: user message + assistant with tool_call + tool_result + assistant follow-up - // This generates ~4 messages per turn, with tool results being large (like file contents) - for turn in 0..30 { - // User message - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "Turn {}: Please check the status of the deployment and fix any issues you find.", - turn - ))), - ..Default::default() - }); - - // Assistant with tool call - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "I'll check the deployment status for turn {}.", - turn - ))), - tool_calls: Some(vec![ - stakpak_shared::models::integrations::openai::ToolCall { - id: format!("tc_{}", turn), - r#type: "function".to_string(), - function: stakpak_shared::models::integrations::openai::FunctionCall { - name: "run_command".to_string(), - arguments: format!( - r#"{{"command":"kubectl get pods -n app-{} -o wide"}}"#, - turn - ), - }, - metadata: None, - }, - ]), - ..Default::default() - }); - - // Tool result (simulating large kubectl output) - messages.push(ChatMessage { - role: Role::Tool, - tool_call_id: Some(format!("tc_{}", turn)), - content: Some(MessageContent::String(format!( - "NAME READY STATUS RESTARTS AGE IP NODE\n\ - app-{t}-deploy-abc123-xyz 1/1 Running 0 2d 10.0.{t}.1 node-1\n\ - app-{t}-deploy-abc123-def 1/1 Running 0 2d 10.0.{t}.2 node-2\n\ - app-{t}-deploy-abc123-ghi 0/1 CrashLoop 5 2d 10.0.{t}.3 node-3\n\ - app-{t}-worker-jkl456-mno 1/1 Running 0 5d 10.0.{t}.4 node-1\n\ - app-{t}-worker-jkl456-pqr 1/1 Running 0 5d 10.0.{t}.5 node-2\n\ - {extra}", - t = turn, - extra = "x".repeat(800) // simulate verbose output - ))), - ..Default::default() - }); - - // Assistant follow-up - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "I found an issue in turn {}. Pod app-{}-deploy-abc123-ghi is in CrashLoopBackOff. \ - Let me check the logs and fix the configuration.", - turn, turn - ))), - ..Default::default() - }); - } - - // === Run the context manager (simulating what the hook does) === - let (result, new_metadata) = - cm.reduce_context_with_budget(messages.clone(), context_window, metadata.clone(), None); - - // === Print what the agent sees === - println!("\n{}", "=".repeat(80)); - println!("CONTEXT TRIMMING SIMULATION RESULTS"); - println!("{}", "=".repeat(80)); - println!( - "Input: {} ChatMessages ({} turns)", - messages.len(), - messages.len() / 4 - ); - println!("Output: {} LLMMessages (after merge/dedup)", result.len()); - println!("Context window: {} tokens", context_window); - println!( - "Threshold: 80% = {} tokens", - (context_window as f32 * 0.8) as u64 - ); - - let _estimated_before = TaskBoardContextManager::estimate_tokens( - &result - .iter() - .map(|m| { - // Create untrimmed version for comparison - LLMMessage { - role: m.role.clone(), - content: LLMMessageContent::String("x".repeat(200)), - } - }) - .collect::>(), - ); - - let estimated_after = TaskBoardContextManager::estimate_tokens(&result); - println!("Estimated tokens after trimming: {}", estimated_after); - - if let Some(ref meta) = new_metadata { - println!( - "Metadata: {}", - serde_json::to_string_pretty(meta).unwrap_or_default() - ); - } else { - println!("Metadata: None (no trimming occurred)"); - } - - println!("\n--- Messages the agent sees ---"); - for (i, msg) in result.iter().enumerate() { - let content_preview = match &msg.content { - LLMMessageContent::String(s) => { - if s.chars().count() > 80 { - let truncated: String = s.chars().take(80).collect(); - format!("{}...", truncated) - } else { - s.clone() - } - } - LLMMessageContent::List(parts) => { - let mut preview = String::new(); - for p in parts { - match p { - LLMMessageTypedContent::Text { text } => { - let t = if text.chars().count() > 60 { - let truncated: String = text.chars().take(60).collect(); - format!("{}...", truncated) - } else { - text.clone() - }; - preview.push_str(&format!("[text:{}] ", t)); - } - LLMMessageTypedContent::ToolCall { id, name, .. } => { - preview.push_str(&format!("[tool_call:{}:{}] ", id, name)); - } - LLMMessageTypedContent::ToolResult { - tool_use_id, - content, - .. - } => { - let c = if content.chars().count() > 40 { - let truncated: String = content.chars().take(40).collect(); - format!("{}...", truncated) - } else { - content.clone() - }; - preview.push_str(&format!("[tool_result:{}:{}] ", tool_use_id, c)); - } - LLMMessageTypedContent::Image { .. } => { - preview.push_str("[image] "); - } - } - } - preview - } - }; - - let trimmed_marker = if is_trimmed(msg) { " [TRIMMED]" } else { "" }; - println!( - " [{:2}] {:>10}: {}{}", - i, msg.role, content_preview, trimmed_marker - ); - } - - // === Simulate session resume: add 2 more turns and trim again === - println!("\n--- After adding 2 more turns (simulating session resume) ---"); - for turn in 30..32 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "Turn {}: Check again please.", - turn - ))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Checking turn {}... {}", - turn, - "z".repeat(200) - ))), - ..Default::default() - }); - } - - let (result2, metadata2) = - cm.reduce_context_with_budget(messages.clone(), context_window, new_metadata, None); - - println!("After resume: {} LLMMessages", result2.len()); - let estimated_after2 = TaskBoardContextManager::estimate_tokens(&result2); - println!("Estimated tokens after second trim: {}", estimated_after2); - if let Some(ref meta) = metadata2 { - println!( - "Metadata: {}", - serde_json::to_string_pretty(meta).unwrap_or_default() - ); - } - - let trimmed_count = result2.iter().filter(|m| is_trimmed(m)).count(); - let untrimmed_count = result2.iter().filter(|m| !is_trimmed(m)).count(); - println!("Trimmed: {}, Untrimmed: {}", trimmed_count, untrimmed_count); - - // Print last 6 messages (what the agent actually works with) - println!("\n--- Last 6 messages (agent's working context) ---"); - let len2 = result2.len(); - let start = len2.saturating_sub(6); - for (i, msg) in result2[start..].iter().enumerate() { - let content_preview = match &msg.content { - LLMMessageContent::String(s) => { - if s.chars().count() > 100 { - let truncated: String = s.chars().take(100).collect(); - format!("{}...", truncated) - } else { - s.clone() - } - } - LLMMessageContent::List(parts) => format!("[{} parts]", parts.len()), - }; - println!(" [{:2}] {:>10}: {}", start + i, msg.role, content_preview); - } - - // Assertions - assert!(trimmed_count > 0, "Should have trimmed some messages"); - assert!( - untrimmed_count >= 4, - "Should keep at least 4 untrimmed messages" - ); - } - - // ========================================================================= - // Regression tests - // - // These tests cover bugs found in production. Each documents: - // 1. The production config that triggered the bug - // 2. Why the old code was wrong - // 3. The correct behavior - // - // Key design invariants: - // - Budget (context_budget_threshold) is the HARD constraint — context - // must not grow past it unchecked. - // - keep_last_n_assistant_messages is BEST-EFFORT — respected when budget - // allows, overridden when budget demands trimming. - // - Lazy trimming: the trim boundary only advances when the *effective* - // token count (after re-applying previous trims) exceeds the threshold. - // This keeps the message prefix stable across turns, which is critical - // for Anthropic prompt caching — the less the prefix changes, the - // higher the cache hit rate. - // ========================================================================= - - /// Production config: keep_last_n=50, threshold=0.3, context_window=200k. - /// - /// Bug: with 30% threshold you hit the budget after ~20 assistant messages, - /// well before accumulating 50. The old code's keep_last_n loop never found - /// 50 assistants, so trim_end stayed at 0 → nothing got trimmed → context - /// grew past 50% unchecked. - /// - /// Fix: budget overrides keep_last_n. When over budget and keep_last_n - /// can't produce a boundary, the fallback trims the oldest assistant to - /// make progress. - #[test] - fn test_production_config_budget_overrides_keep_last_n() { - // Mirror the actual production config from client/mod.rs - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 50, - context_budget_threshold: 0.3, - }); - - // Simulate a session with 10 turns (20 messages) — well under 50 - // assistants but enough content to exceed 30% of a small window. - let mut messages = Vec::new(); - for i in 0..10 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "User turn {}: {}", - i, - "x".repeat(200) - ))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Assistant turn {}: {}", - i, - "y".repeat(200) - ))), - ..Default::default() - }); - } - - // Context window where 30% threshold is easily exceeded by 10 turns - let context_window = 500; - let (result, metadata) = - cm.reduce_context_with_budget(messages, context_window, None, None); - - // Must trim despite having only 10 assistants (< keep_last_n of 50) - assert!(metadata.is_some(), "Should trigger trimming"); - let trimmed_idx = metadata.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!( - trimmed_idx > 0, - "Budget must override keep_last_n when over threshold \ - (had 10 assistants < keep_last_n of 50, but was over 30% budget)" - ); - - // At least the first assistant should be trimmed - let first_assistant = result.iter().find(|m| m.role == "assistant").unwrap(); - assert!( - is_trimmed(first_assistant), - "Oldest assistant should be trimmed to make progress" - ); - - // User messages are never trimmed - for msg in &result { - if msg.role == "user" { - assert!(!is_trimmed(msg), "User messages should never be trimmed"); - } - } - } - - /// Regression test: keep_last_n produces a trim boundary but the remaining - /// N assistant messages (with large tool results) still exceed the budget. - /// - /// Bug: the old code used keep_last_n as a hard boundary — when - /// keep_n_trim_end > 0, it skipped the budget-driven scan entirely. - /// With keep_last_n=20 and 200K context window, 20 assistant messages - /// with large tool results could easily be 160K+ tokens, blowing past - /// the 80% threshold. - /// - /// Fix: keep_last_n is best-effort. After applying keep_last_n trimming, - /// if still over budget, continue scanning forward to trim more messages. - #[test] - fn test_keep_last_n_boundary_insufficient_budget_continues_trimming() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 3, - context_budget_threshold: 0.8, - }); - - // 6 turns: user + assistant with large content. - // keep_last_n=3 would normally keep the last 3 assistants untrimmed, - // but each assistant is so large that 3 of them alone exceed the budget. - let mut messages = Vec::new(); - for i in 0..6 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!("user {}", i))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "assistant {}: {}", - i, - "x".repeat(500) // ~170 tokens each - ))), - ..Default::default() - }); - } - - // Context window where 80% ≈ 400 tokens. - // 6 assistants × ~170 tokens = ~1020 tokens total. - // keep_last_n=3 would leave ~510 tokens — still over 400. - // Budget must override and trim past the keep_last_n boundary. - let context_window = 500; - let (result, metadata) = - cm.reduce_context_with_budget(messages, context_window, None, None); - - assert!(metadata.is_some(), "Should trigger trimming"); - let trimmed_idx = metadata.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - // The trim index must go PAST the keep_last_n boundary (index of - // 3rd-from-last assistant). With 6 assistants at indices 1,3,5,7,9,11, - // keep_last_n=3 boundary = index 7. Budget must push past that. - let keep_n_boundary = 7; // 3rd-from-last assistant in merged sequence - assert!( - trimmed_idx > keep_n_boundary, - "Budget should override keep_last_n: trimmed_idx {} should be > keep_n boundary {}", - trimmed_idx, - keep_n_boundary - ); - - // Effective tokens after trimming should be under threshold - let effective_tokens = TaskBoardContextManager::estimate_tokens(&result); - let threshold = (context_window as f32 * 0.8) as u64; - assert!( - effective_tokens <= threshold, - "Effective tokens {} should be <= threshold {} after budget-driven trimming", - effective_tokens, - threshold - ); - - // User messages are never trimmed - for msg in &result { - if msg.role == "user" { - assert!(!is_trimmed(msg), "User messages should never be trimmed"); - } - } - } - - /// Simulates the exact production scenario from the bug report: - /// Claude Opus 4.5 with 200K context, keep_last_n=20, threshold=0.8. - /// A session grows to 196K/200K tokens because the last 20 assistant - /// messages (with large tool results) themselves exceed the budget. - /// - /// The OLD code would stop trimming at the keep_last_n boundary and - /// never get under budget. The fix continues trimming past keep_last_n - /// when budget demands it. - /// - /// We also verify the hook-level fix: the system prompt (~8K tokens) - /// and max_output_tokens (16K) are subtracted from the context window - /// before the trimmer sees it, so the effective budget is ~176K × 0.8 - /// = ~140K, not 200K × 0.8 = 160K. - #[test] - fn test_production_scenario_200k_context_keep_20() { - // Mirror production config from client/mod.rs - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 20, - context_budget_threshold: 0.8, - }); - - // Simulate a session with 30 turns of tool-heavy interaction. - // Each turn: user (small) → assistant with tool_call → tool result (large) → assistant follow-up - // This produces 60 assistant messages and 30 tool results. - // After merge: 120 messages (30 × 4). - let mut messages = Vec::new(); - for turn in 0..30 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "User turn {}: {}", - turn, - "q".repeat(100) - ))), - ..Default::default() - }); - messages.push(create_tool_call_msg(&format!("tc_{}", turn))); - messages.push(create_tool_result_msg( - &format!("tc_{}", turn), - // Large tool results (~2000 chars each ≈ 700 tokens) - &format!("result {}: {}", turn, "r".repeat(2000)), - )); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Follow-up {}: {}", - turn, - "a".repeat(500) - ))), - ..Default::default() - }); - } - - // Use a context window that simulates the effective budget after - // subtracting system prompt and max_output_tokens in the hook. - // Real: 200K - ~8K (system prompt) - 16K (max_output) ≈ 176K - // Scale down proportionally for test: use 10000 as context_window. - // 80% threshold = 8000 tokens. - let context_window = 10_000u64; - let threshold = (context_window as f32 * 0.8) as u64; - - let (result, metadata) = - cm.reduce_context_with_budget(messages.clone(), context_window, None, None); - - assert!(metadata.is_some(), "Should trigger trimming"); - let trimmed_idx = metadata.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - // KEY ASSERTION: effective tokens must be under threshold. - // This is what the old code failed to guarantee — it would stop at - // keep_last_n boundary even if remaining tokens exceeded the budget. - let effective_tokens = TaskBoardContextManager::estimate_tokens(&result); - assert!( - effective_tokens <= threshold, - "PRODUCTION BUG: effective tokens {} exceed threshold {} \ - (this is the exact scenario where context hit 196K/200K). \ - trimmed_idx={}, total_messages={}", - effective_tokens, - threshold, - trimmed_idx, - result.len() - ); - - // The trim index should go well past the keep_last_n boundary. - // With 60 assistants and keep_last_n=20, the keep_n boundary is at - // the 20th-from-last assistant. But 20 assistants with ~700-token - // tool results = ~14000 tokens > 8000 threshold, so budget must - // push further. - assert!(trimmed_idx > 0, "Should have trimmed some messages"); - - // User messages are never trimmed - for msg in &result { - if msg.role == "user" { - assert!(!is_trimmed(msg), "User messages should never be trimmed"); - } - } - - // Simulate a second call (next turn) — verify trimming state persists - // and effective tokens stay under budget. - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String("Next question".to_string())), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Response: {}", - "a".repeat(500) - ))), - ..Default::default() - }); - - let (result2, metadata2) = - cm.reduce_context_with_budget(messages, context_window, metadata, None); - - let effective_tokens2 = TaskBoardContextManager::estimate_tokens(&result2); - assert!( - effective_tokens2 <= threshold, - "Second call: effective tokens {} exceed threshold {} — trimming not keeping up", - effective_tokens2, - threshold - ); - - let trimmed_idx2 = metadata2.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!( - trimmed_idx2 >= trimmed_idx, - "Trim index should never go backward: {} < {}", - trimmed_idx2, - trimmed_idx - ); - } - - /// Lazy trimming preserves prompt cache: after trimming brings effective - /// tokens under threshold, the trim boundary must NOT advance on the next - /// call — even though raw (untrimmed) history keeps growing. - /// - /// Bug: the old code estimated tokens on raw messages, so once raw content - /// exceeded the threshold it was permanently true. The trim boundary - /// advanced every turn, invalidating the prompt cache each time. - /// - /// Fix: re-apply previous trims first, then re-estimate on the *effective* - /// (trimmed) messages. Only advance when effective tokens exceed threshold. - #[test] - fn test_lazy_trimming_preserves_cache_prefix() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 2, - context_budget_threshold: 0.8, - }); - - // Build conversation that exceeds threshold - let mut messages = Vec::new(); - for i in 0..5 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "User {}: {}", - i, - "x".repeat(100) - ))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Asst {}: {}", - i, - "y".repeat(100) - ))), - ..Default::default() - }); - } - - // First call: triggers trimming, establishes trim boundary - let context_window = 500; - let (result1, metadata1) = - cm.reduce_context_with_budget(messages.clone(), context_window, None, None); - assert!(metadata1.is_some(), "First call should trigger trimming"); - let trimmed_idx_1 = metadata1.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!(trimmed_idx_1 > 0); - - // Add small messages — raw total grows but effective total stays under - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String("short".to_string())), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("brief".to_string())), - ..Default::default() - }); - - // Second call: effective tokens under threshold → index must NOT advance - let (result2, metadata2) = - cm.reduce_context_with_budget(messages.clone(), context_window, metadata1, None); - let trimmed_idx_2 = metadata2.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - - assert_eq!( - trimmed_idx_2, trimmed_idx_1, - "Trim index must not advance when effective tokens are under threshold — \ - advancing would invalidate the prompt cache (was {}, now {})", - trimmed_idx_1, trimmed_idx_2 - ); - - // The trimmed prefix must be identical across calls (cache stability) - for i in 0..trimmed_idx_1.min(result1.len()).min(result2.len()) { - let t1 = is_trimmed(&result1[i]); - let t2 = is_trimmed(&result2[i]); - assert_eq!( - t1, t2, - "Message {} trim state changed between calls (was trimmed={}, now trimmed={}) — \ - this would invalidate the prompt cache", - i, t1, t2 - ); - } - - // New messages at the end should NOT be trimmed - let len = result2.len(); - assert!( - !is_trimmed(&result2[len - 1]), - "Last message should not be trimmed" - ); - assert!( - !is_trimmed(&result2[len - 2]), - "Second-to-last message should not be trimmed" - ); - } - - /// Under budget with fewer assistants than keep_last_n: no trimming at all. - #[test] - fn test_under_budget_fewer_assistants_no_trimming() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 10, - context_budget_threshold: 0.8, - }); - - let messages = vec![ - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("hello".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("hi".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("how are you".to_string())), - ..Default::default() - }, - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("good".to_string())), - ..Default::default() - }, - ]; - - // Large window — well under budget - let (result, metadata) = cm.reduce_context_with_budget(messages, 200_000, None, None); - - assert!( - metadata.is_none(), - "Should not produce metadata when under budget" - ); - for msg in &result { - assert!( - !is_trimmed(msg), - "No messages should be trimmed when under budget" - ); - } - } - - /// Multi-turn progressive trimming: each call that's still over budget - /// trims one more assistant, eventually bringing context under threshold. - /// Once under threshold, the boundary freezes (cache stability). - #[test] - fn test_progressive_trimming_then_freeze() { - let cm = TaskBoardContextManager::new(TaskBoardContextManagerOptions { - keep_last_n_assistant_messages: 50, // high keep_last_n, like production - context_budget_threshold: 0.3, - }); - - // 5 turns with large assistant responses - let mut messages = Vec::new(); - for i in 0..5 { - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!("q{}", i))), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("y".repeat(300))), - ..Default::default() - }); - } - - let context_window = 800; // 30% = 240 tokens budget - - // Turn 1: over budget → trims oldest assistant - let (_, meta1) = - cm.reduce_context_with_budget(messages.clone(), context_window, None, None); - assert!(meta1.is_some()); - let idx1 = meta1.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!(idx1 > 0, "Should trim on first call"); - - // Turn 2: add small message, call again — if still over budget, - // index advances further; if under, it freezes. - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String("small".to_string())), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("ok".to_string())), - ..Default::default() - }); - - let (_, meta2) = - cm.reduce_context_with_budget(messages.clone(), context_window, meta1, None); - let idx2 = meta2.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!( - idx2 >= idx1, - "Trim index must never go backward: {} < {}", - idx2, - idx1 - ); - - // Turn 3: another small message - messages.push(ChatMessage { - role: Role::User, - content: Some(MessageContent::String("tiny".to_string())), - ..Default::default() - }); - messages.push(ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String("k".to_string())), - ..Default::default() - }); - - let (_, meta3) = - cm.reduce_context_with_budget(messages.clone(), context_window, meta2.clone(), None); - let idx3 = meta3.as_ref().unwrap()["trimmed_up_to_message_index"] - .as_u64() - .unwrap() as usize; - assert!( - idx3 >= idx2, - "Trim index must never go backward: {} < {}", - idx3, - idx2 - ); - - // If idx2 == idx3, trimming has frozen (effective tokens under budget). - // If idx3 > idx2, we're still over budget and making progress. - // Either way, the invariant holds: monotonically non-decreasing. - } -} - pub struct TaskBoardContextManagerOptions { /// How many recent **assistant** messages to keep untrimmed when context /// trimming is triggered. The trim boundary is placed just before the diff --git a/libs/api/src/local/hooks/file_scratchpad_context/mod.rs b/libs/api/src/local/hooks/file_scratchpad_context/mod.rs index 7f19e992b..c16f209ac 100644 --- a/libs/api/src/local/hooks/file_scratchpad_context/mod.rs +++ b/libs/api/src/local/hooks/file_scratchpad_context/mod.rs @@ -4,8 +4,6 @@ use crate::local::context_managers::file_scratchpad_context_manager::{ use crate::models::AgentState; use stakpak_shared::define_hook; use stakpak_shared::hooks::{Hook, HookAction, HookContext, HookError, LifecycleEvent}; -use stakpak_shared::models::integrations::openai::Role; -use stakpak_shared::models::llm::{LLMInput, LLMMessage, LLMMessageContent}; const SYSTEM_PROMPT: &str = include_str!("./system_prompt.txt"); const SCRATCHPAD_FILE: &str = ".stakpak/session/scratchpad.md"; @@ -59,12 +57,6 @@ define_hook!( let model = ctx.state.active_model.clone(); - let tools = ctx - .state - .tools - .clone() - .map(|t| t.into_iter().map(Into::into).collect()); - let cwd = std::env::current_dir().unwrap_or_default(); let scratchpad_file_path = cwd.join(self.context_manager.get_scratchpad_path(ctx.session_id)); @@ -77,22 +69,22 @@ define_hook!( .replace("{{TODO_PATH}}", &todo_file_path.display().to_string()); let mut messages = Vec::new(); - messages.push(LLMMessage { - role: Role::System.to_string(), - content: LLMMessageContent::String(system_prompt), - }); + messages.push(stakai::Message::new(stakai::Role::System, system_prompt)); messages.extend( self.context_manager .reduce_context_with_session(ctx.state.messages.clone(), ctx.session_id), ); - ctx.state.llm_input = Some(LLMInput { + ctx.state.llm_input = Some(stakai::GenerateRequest { model, messages, - max_tokens: 16000, - tools, + options: stakai::GenerateOptions { + max_tokens: Some(16000), + tools: ctx.state.tools.clone(), + ..Default::default() + }, provider_options: None, - headers: None, + telemetry_metadata: None, }); Ok(HookAction::Continue) diff --git a/libs/api/src/local/hooks/inline_scratchpad_context/mod.rs b/libs/api/src/local/hooks/inline_scratchpad_context/mod.rs index 11e24ab59..8d66c107a 100644 --- a/libs/api/src/local/hooks/inline_scratchpad_context/mod.rs +++ b/libs/api/src/local/hooks/inline_scratchpad_context/mod.rs @@ -1,7 +1,5 @@ use stakpak_shared::define_hook; use stakpak_shared::hooks::{Hook, HookAction, HookContext, HookError, LifecycleEvent}; -use stakpak_shared::models::integrations::openai::Role; -use stakpak_shared::models::llm::{LLMInput, LLMMessage, LLMMessageContent}; use crate::local::context_managers::ContextManager; use crate::local::context_managers::scratchpad_context_manager::{ @@ -48,29 +46,23 @@ define_hook!( let model = ctx.state.active_model.clone(); - let tools = ctx - .state - .tools - .clone() - .map(|t| t.into_iter().map(Into::into).collect()); - let mut messages = Vec::new(); - messages.push(LLMMessage { - role: Role::System.to_string(), - content: LLMMessageContent::String(SYSTEM_PROMPT.to_string()), - }); + messages.push(stakai::Message::new(stakai::Role::System, SYSTEM_PROMPT)); messages.extend( self.context_manager .reduce_context(ctx.state.messages.clone()), ); - ctx.state.llm_input = Some(LLMInput { + ctx.state.llm_input = Some(stakai::GenerateRequest { model, messages, - max_tokens: 16000, - tools, + options: stakai::GenerateOptions { + max_tokens: Some(16000), + tools: ctx.state.tools.clone(), + ..Default::default() + }, provider_options: None, - headers: None, + telemetry_metadata: None, }); Ok(HookAction::Continue) diff --git a/libs/api/src/local/hooks/task_board_context/mod.rs b/libs/api/src/local/hooks/task_board_context/mod.rs index 797a4dae4..9d53f346d 100644 --- a/libs/api/src/local/hooks/task_board_context/mod.rs +++ b/libs/api/src/local/hooks/task_board_context/mod.rs @@ -1,7 +1,5 @@ use stakpak_shared::define_hook; use stakpak_shared::hooks::{Hook, HookAction, HookContext, HookError, LifecycleEvent}; -use stakpak_shared::models::integrations::openai::Role; -use stakpak_shared::models::llm::{LLMInput, LLMMessage, LLMMessageContent}; use crate::local::context_managers::task_board_context_manager::{ TaskBoardContextManager, TaskBoardContextManagerOptions, @@ -49,21 +47,16 @@ define_hook!( // only the space actually available for chat messages. // - System prompt: added after trimming (line 67+), not in message list // - max_output_tokens: reserved for the model's response - let system_prompt_tokens = TaskBoardContextManager::estimate_tokens(&[LLMMessage { - role: Role::System.to_string(), - content: LLMMessageContent::String(SYSTEM_PROMPT.to_string()), - }]); + let system_prompt_tokens = + TaskBoardContextManager::estimate_tokens(&[stakai::Message::new( + stakai::Role::System, + SYSTEM_PROMPT, + )]); let context_window = model .limit .context .saturating_sub(system_prompt_tokens + max_output_tokens); - let llm_tools: Option> = ctx - .state - .tools - .clone() - .map(|t| t.into_iter().map(Into::into).collect()); - // Use budget-aware trimming with metadata from checkpoint. // Tool definitions are passed in so the context manager can account // for their token overhead internally. @@ -71,26 +64,26 @@ define_hook!( ctx.state.messages.clone(), context_window, ctx.state.metadata.clone(), - llm_tools.as_deref(), + ctx.state.tools.as_deref(), ); // Write updated metadata back to state for checkpoint persistence ctx.state.metadata = updated_metadata; let mut messages = Vec::new(); - messages.push(LLMMessage { - role: Role::System.to_string(), - content: LLMMessageContent::String(SYSTEM_PROMPT.to_string()), - }); + messages.push(stakai::Message::new(stakai::Role::System, SYSTEM_PROMPT)); messages.extend(reduced_messages); - ctx.state.llm_input = Some(LLMInput { + ctx.state.llm_input = Some(stakai::GenerateRequest { model, messages, - max_tokens: max_output_tokens as u32, - tools: llm_tools, + options: stakai::GenerateOptions { + max_tokens: Some(max_output_tokens as u32), + tools: ctx.state.tools.clone(), + ..Default::default() + }, provider_options: None, - headers: None, + telemetry_metadata: None, }); Ok(HookAction::Continue) diff --git a/libs/api/src/local/tests.rs b/libs/api/src/local/tests.rs index aea1968e3..165a5bd78 100644 --- a/libs/api/src/local/tests.rs +++ b/libs/api/src/local/tests.rs @@ -1,7 +1,7 @@ #[cfg(test)] mod local_storage_tests { use crate::storage::*; - use stakpak_shared::models::integrations::openai::{ChatMessage, MessageContent, Role}; + use stakai::{ContentPart, Message, MessageContent, Role}; use uuid::Uuid; /// Helper: create an in-memory LocalStorage @@ -12,25 +12,32 @@ mod local_storage_tests { } /// Helper: build a simple CreateSessionRequest - fn session_request(title: &str, messages: Vec) -> CreateSessionRequest { + fn session_request(title: &str, messages: Vec) -> CreateSessionRequest { CreateSessionRequest::new(title, messages) } - /// Helper: build a user ChatMessage - fn user_msg(text: &str) -> ChatMessage { - ChatMessage { - role: Role::User, - content: Some(MessageContent::String(text.to_string())), - ..Default::default() - } + /// Helper: build a user message + fn user_msg(text: &str) -> Message { + Message::new(Role::User, text.to_string()) + } + + /// Helper: build an assistant message + fn assistant_msg(text: &str) -> Message { + Message::new(Role::Assistant, text.to_string()) + } + + fn message_text(message: &Message) -> Option { + message.content.text() } - /// Helper: build an assistant ChatMessage - fn assistant_msg(text: &str) -> ChatMessage { - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(text.to_string())), - ..Default::default() + fn is_trimmed_message(message: &Message) -> bool { + match &message.content { + MessageContent::Text(text) => text == "[trimmed]", + MessageContent::Parts(parts) => parts.iter().all(|part| match part { + ContentPart::Text { text, .. } => text == "[trimmed]", + ContentPart::ToolResult { content, .. } => content.as_str() == Some("[trimmed]"), + ContentPart::ToolCall { .. } | ContentPart::Image { .. } => true, + }), } } @@ -53,11 +60,7 @@ mod local_storage_tests { assert!(result.checkpoint.parent_id.is_none()); assert_eq!(result.checkpoint.state.messages.len(), 1); assert_eq!( - result.checkpoint.state.messages[0] - .content - .as_ref() - .unwrap() - .to_string(), + result.checkpoint.state.messages[0].content.text().unwrap(), "hello" ); } @@ -562,12 +565,13 @@ mod local_storage_tests { let msgs = vec![ user_msg("question"), assistant_msg("answer"), - ChatMessage { - role: Role::Tool, - content: Some(MessageContent::String("tool result".to_string())), - tool_call_id: Some("tc_123".to_string()), - ..Default::default() - }, + Message::new( + Role::Tool, + MessageContent::Parts(vec![ContentPart::tool_result( + "tc_123", + serde_json::Value::String("tool result".to_string()), + )]), + ), ]; let session = storage .create_session(&session_request("Test", msgs)) @@ -579,10 +583,12 @@ mod local_storage_tests { assert_eq!(fetched.state.messages[0].role, Role::User); assert_eq!(fetched.state.messages[1].role, Role::Assistant); assert_eq!(fetched.state.messages[2].role, Role::Tool); - assert_eq!( - fetched.state.messages[2].tool_call_id, - Some("tc_123".to_string()) - ); + assert!(fetched.state.messages[2].parts().iter().any(|part| { + matches!( + part, + ContentPart::ToolResult { tool_call_id, .. } if tool_call_id == "tc_123" + ) + })); } // ========================================================================= @@ -1068,29 +1074,16 @@ mod local_storage_tests { // ========================================================================= /// Helper: build a large user message to inflate token count - fn large_user_msg(turn: usize) -> ChatMessage { - ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!( - "Turn {}: {}", - turn, - "x".repeat(200) - ))), - ..Default::default() - } + fn large_user_msg(turn: usize) -> Message { + Message::new(Role::User, format!("Turn {}: {}", turn, "x".repeat(200))) } /// Helper: build a large assistant message to inflate token count - fn large_assistant_msg(turn: usize) -> ChatMessage { - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(format!( - "Response {}: {}", - turn, - "y".repeat(200) - ))), - ..Default::default() - } + fn large_assistant_msg(turn: usize) -> Message { + Message::new( + Role::Assistant, + format!("Response {}: {}", turn, "y".repeat(200)), + ) } /// Simulates the async mode flow: trimming should NOT trigger when @@ -1154,7 +1147,7 @@ mod local_storage_tests { ); // All content should be preserved for msg in &result { - if let stakpak_shared::models::llm::LLMMessageContent::String(s) = &msg.content { + if let Some(s) = message_text(msg) { assert_ne!(s, "[trimmed]", "No messages should be trimmed"); } } @@ -1240,18 +1233,16 @@ mod local_storage_tests { // Verify early assistant messages are trimmed (user messages preserved) // result[0] is user (NOT trimmed), result[1] is assistant (trimmed) - match &result[0].content { - stakpak_shared::models::llm::LLMMessageContent::String(s) => { - assert_ne!(s, "[trimmed]", "First user message should NOT be trimmed"); - } - _ => panic!("Expected string content"), - } - match &result[1].content { - stakpak_shared::models::llm::LLMMessageContent::String(s) => { - assert_eq!(s, "[trimmed]", "First assistant message should be trimmed"); - } - _ => panic!("Expected string content"), - } + assert_ne!( + message_text(&result[0]).expect("Expected text content"), + "[trimmed]", + "First user message should NOT be trimmed" + ); + assert_eq!( + message_text(&result[1]).expect("Expected text content"), + "[trimmed]", + "First assistant message should be trimmed" + ); // Budget is the hard constraint: the trimmer trims all assistant/tool // messages it can. With a 200-token window, user messages alone (~700 @@ -1260,8 +1251,8 @@ mod local_storage_tests { // Verify all assistant messages before the trim boundary are trimmed. for (i, msg) in result.iter().enumerate() { if i < trimmed_idx - && msg.role == "assistant" - && let stakpak_shared::models::llm::LLMMessageContent::String(s) = &msg.content + && msg.role == Role::Assistant + && let Some(s) = message_text(msg) { assert_eq!( s, "[trimmed]", @@ -1273,8 +1264,8 @@ mod local_storage_tests { // User messages are never trimmed regardless of budget pressure for msg in &result { - if msg.role == "user" - && let stakpak_shared::models::llm::LLMMessageContent::String(s) = &msg.content + if msg.role == Role::User + && let Some(s) = message_text(msg) { assert_ne!(s, "[trimmed]", "User messages should never be trimmed"); } @@ -1352,29 +1343,15 @@ mod local_storage_tests { // Budget is the hard constraint — keep_last_n is best-effort, so the // last 4 messages may also be trimmed if budget demands it. for (i, msg) in result1.iter().enumerate() { - let is_trimmed = match &msg.content { - stakpak_shared::models::llm::LLMMessageContent::String(s) => s == "[trimmed]", - stakpak_shared::models::llm::LLMMessageContent::List(parts) => { - parts.iter().all(|p| match p { - stakpak_shared::models::llm::LLMMessageTypedContent::Text { text } => { - text == "[trimmed]" - } - stakpak_shared::models::llm::LLMMessageTypedContent::ToolResult { - content, - .. - } => content == "[trimmed]", - _ => true, - }) - } - }; - if i < trimmed_idx_1 && msg.role != "user" { + let is_trimmed = is_trimmed_message(msg); + if i < trimmed_idx_1 && msg.role != Role::User { assert!( is_trimmed, "Phase 1: non-user message {} should be trimmed", i ); } - if msg.role == "user" { + if msg.role == Role::User { assert!( !is_trimmed, "Phase 1: user message {} should NOT be trimmed", @@ -1387,11 +1364,8 @@ mod local_storage_tests { // so the trimmer can't get fully under budget — but all trimmable // messages before the boundary must be trimmed. for (i, msg) in result1.iter().enumerate() { - if i < trimmed_idx_1 && msg.role == "assistant" { - let is_trimmed = match &msg.content { - stakpak_shared::models::llm::LLMMessageContent::String(s) => s == "[trimmed]", - _ => true, - }; + if i < trimmed_idx_1 && msg.role == Role::Assistant { + let is_trimmed = is_trimmed_message(msg); assert!( is_trimmed, "Phase 1: assistant at {} before boundary should be trimmed", @@ -1467,19 +1441,17 @@ mod local_storage_tests { // Verify: budget overrides keep_last_n — all trimmable messages before // the boundary are trimmed, user messages are always preserved. for (i, msg) in result2.iter().enumerate() { - if msg.role == "user" - && let stakpak_shared::models::llm::LLMMessageContent::String(s) = &msg.content + if msg.role == Role::User + && let Some(s) = message_text(msg) { assert_ne!( s, "[trimmed]", "Phase 2: user message {} should NOT be trimmed", i ); - } else if i < trimmed_idx_2 - && let stakpak_shared::models::llm::LLMMessageContent::String(s) = &msg.content - { - assert_eq!( - s, "[trimmed]", + } else if i < trimmed_idx_2 { + assert!( + is_trimmed_message(msg), "Phase 2: non-user message {} before boundary should be trimmed", i ); @@ -1578,18 +1550,16 @@ mod local_storage_tests { assert!(trimmed_idx > 0, "Should have trimmed messages"); // Verify trimming is correct — user messages preserved, assistant messages trimmed - match &result[0].content { - stakpak_shared::models::llm::LLMMessageContent::String(s) => { - assert_ne!(s, "[trimmed]", "First user message should NOT be trimmed"); - } - _ => panic!("Expected string content"), - } - match &result[1].content { - stakpak_shared::models::llm::LLMMessageContent::String(s) => { - assert_eq!(s, "[trimmed]", "First assistant message should be trimmed"); - } - _ => panic!("Expected string content"), - } + assert_ne!( + message_text(&result[0]).expect("Expected text content"), + "[trimmed]", + "First user message should NOT be trimmed" + ); + assert_eq!( + message_text(&result[1]).expect("Expected text content"), + "[trimmed]", + "First assistant message should be trimmed" + ); } /// Verify that metadata is correctly persisted and loaded through multiple @@ -1745,7 +1715,7 @@ mod local_storage_tests { // Should return the empty metadata as-is (no trimming triggered) // The function returns metadata unchanged when under threshold and prev_trimmed_up_to == 0 for msg in &result { - if let stakpak_shared::models::llm::LLMMessageContent::String(s) = &msg.content { + if let Some(s) = message_text(msg) { assert_ne!(s, "[trimmed]", "No messages should be trimmed"); } } @@ -1828,11 +1798,7 @@ mod local_storage_tests { barrier.wait().await; let request = CreateSessionRequest::new( format!("concurrent-session-{}", i), - vec![ChatMessage { - role: Role::User, - content: Some(MessageContent::String(format!("msg {}", i))), - ..Default::default() - }], + vec![Message::new(Role::User, format!("msg {}", i))], ); storage.create_session(&request).await })); @@ -1883,11 +1849,7 @@ mod local_storage_tests { let writer = tokio::spawn(async move { let request = CreateSessionRequest::new( "contended-session", - vec![ChatMessage { - role: Role::User, - content: Some(MessageContent::String("contended".to_string())), - ..Default::default() - }], + vec![Message::new(Role::User, "contended".to_string())], ); storage2.create_session(&request).await }); diff --git a/libs/api/src/models.rs b/libs/api/src/models.rs index cc051673d..509a53cc2 100644 --- a/libs/api/src/models.rs +++ b/libs/api/src/models.rs @@ -3,10 +3,7 @@ use rmcp::model::Content; use serde::{Deserialize, Serialize}; use serde_json::Value; use stakai::Model; -use stakpak_shared::models::{ - integrations::openai::{ChatMessage, FunctionCall, MessageContent, Role, Tool, ToolCall}, - llm::{LLMInput, LLMMessage, LLMMessageContent, LLMMessageTypedContent, LLMTokenUsage}, -}; +use stakpak_shared::models::llm::LLMTokenUsage; use std::collections::HashMap; use uuid::Uuid; @@ -498,29 +495,6 @@ impl ListRuleBook { } } -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] -pub struct SimpleLLMMessage { - #[serde(rename = "role")] - pub role: SimpleLLMRole, - pub content: String, -} - -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] -#[serde(rename_all = "lowercase")] -pub enum SimpleLLMRole { - User, - Assistant, -} - -impl std::fmt::Display for SimpleLLMRole { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - SimpleLLMRole::User => write!(f, "user"), - SimpleLLMRole::Assistant => write!(f, "assistant"), - } - } -} - #[derive(Debug, Deserialize, Serialize)] pub struct SearchDocsRequest { pub keywords: String, @@ -558,10 +532,10 @@ pub struct SlackSendMessageRequest { pub struct AgentState { /// The active model to use for inference pub active_model: Model, - pub messages: Vec, - pub tools: Option>, + pub messages: Vec, + pub tools: Option>, - pub llm_input: Option, + pub llm_input: Option, pub llm_output: Option, /// Metadata for checkpoint persistence (context trimming state, etc.) @@ -569,74 +543,52 @@ pub struct AgentState { pub metadata: Option, } -#[derive(Debug, Clone, Default, Serialize)] +#[derive(Debug, Clone, Serialize)] pub struct LLMOutput { - pub new_message: LLMMessage, - pub usage: LLMTokenUsage, -} - -impl From<&LLMOutput> for ChatMessage { - fn from(value: &LLMOutput) -> Self { - let message_content = match &value.new_message.content { - LLMMessageContent::String(s) => s.clone(), - LLMMessageContent::List(l) => l - .iter() - .map(|c| match c { - LLMMessageTypedContent::Text { text } => text.clone(), - LLMMessageTypedContent::ToolCall { .. } => String::new(), - LLMMessageTypedContent::ToolResult { content, .. } => content.clone(), - LLMMessageTypedContent::Image { .. } => String::new(), - }) - .collect::>() - .join("\n"), - }; - let tool_calls = if let LLMMessageContent::List(items) = &value.new_message.content { - let calls: Vec = items - .iter() - .filter_map(|item| { - if let LLMMessageTypedContent::ToolCall { - id, - name, - args, - metadata, - } = item - { - Some(ToolCall { - id: id.clone(), - r#type: "function".to_string(), - function: FunctionCall { - name: name.clone(), - arguments: args.to_string(), - }, - metadata: metadata.clone(), - }) - } else { - None - } - }) - .collect(); - - if calls.is_empty() { None } else { Some(calls) } - } else { - None - }; - ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String(message_content)), - name: None, - tool_calls, - tool_call_id: None, - usage: Some(value.usage.clone()), - ..Default::default() - } + pub new_message: stakai::Message, + pub usage: stakai::Usage, +} + +#[derive(Debug, Clone, Serialize)] +pub struct CompletionResponse { + pub id: String, + pub created: u64, + pub model: String, + pub message: stakai::Message, + pub usage: stakai::Usage, + pub metadata: Option, +} + +#[derive(Debug, Clone, Serialize)] +pub enum AgentStreamEvent { + Model(Model), + Event(stakai::StreamEvent), + Metadata(Value), +} + +pub fn stakai_usage_to_llm_usage(usage: &stakai::Usage) -> LLMTokenUsage { + use stakpak_shared::models::llm::PromptTokensDetails; + + LLMTokenUsage { + prompt_tokens: usage.prompt_tokens, + completion_tokens: usage.completion_tokens, + total_tokens: usage.total_tokens, + prompt_tokens_details: usage.input_token_details.as_ref().map(|details| { + PromptTokensDetails { + input_tokens: details.no_cache, + output_tokens: None, + cache_read_input_tokens: details.cache_read, + cache_write_input_tokens: details.cache_write, + } + }), } } impl AgentState { pub fn new( active_model: Model, - messages: Vec, - tools: Option>, + messages: Vec, + tools: Option>, metadata: Option, ) -> Self { Self { @@ -649,11 +601,11 @@ impl AgentState { } } - pub fn set_messages(&mut self, messages: Vec) { + pub fn set_messages(&mut self, messages: Vec) { self.messages = messages; } - pub fn set_tools(&mut self, tools: Option>) { + pub fn set_tools(&mut self, tools: Option>) { self.tools = tools; } @@ -661,18 +613,18 @@ impl AgentState { self.active_model = model; } - pub fn set_llm_input(&mut self, llm_input: Option) { + pub fn set_llm_input(&mut self, llm_input: Option) { self.llm_input = llm_input; } - pub fn set_llm_output(&mut self, new_message: LLMMessage, new_usage: Option) { + pub fn set_llm_output(&mut self, new_message: stakai::Message, new_usage: stakai::Usage) { self.llm_output = Some(LLMOutput { new_message, - usage: new_usage.unwrap_or_default(), + usage: new_usage, }); } - pub fn append_new_message(&mut self, new_message: ChatMessage) { + pub fn append_new_message(&mut self, new_message: stakai::Message) { self.messages.push(new_message); } } diff --git a/libs/api/src/stakpak/models.rs b/libs/api/src/stakpak/models.rs index f1987166c..54ee541f2 100644 --- a/libs/api/src/stakpak/models.rs +++ b/libs/api/src/stakpak/models.rs @@ -4,7 +4,6 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; -use stakpak_shared::models::integrations::openai::ChatMessage; use uuid::Uuid; // ============================================================================= @@ -90,7 +89,7 @@ pub struct CheckpointSummary { #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct CheckpointState { #[serde(default)] - pub messages: Vec, + pub messages: Vec, /// Optional metadata for context trimming state, etc. #[serde(default, skip_serializing_if = "Option::is_none")] pub metadata: Option, diff --git a/libs/api/src/storage.rs b/libs/api/src/storage.rs index e71ddbc5a..dd5a043ef 100644 --- a/libs/api/src/storage.rs +++ b/libs/api/src/storage.rs @@ -6,7 +6,6 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; -use stakpak_shared::models::integrations::openai::ChatMessage; use uuid::Uuid; // Re-export implementations @@ -285,7 +284,7 @@ pub struct CheckpointSummary { #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct CheckpointState { #[serde(default)] - pub messages: Vec, + pub messages: Vec, /// Optional metadata for context trimming state, etc. #[serde(default, skip_serializing_if = "Option::is_none")] pub metadata: Option, @@ -305,7 +304,7 @@ pub struct CreateSessionRequest { } impl CreateSessionRequest { - pub fn new(title: impl Into, messages: Vec) -> Self { + pub fn new(title: impl Into, messages: Vec) -> Self { Self { title: title.into(), visibility: SessionVisibility::Private, @@ -359,7 +358,7 @@ pub struct CreateCheckpointRequest { } impl CreateCheckpointRequest { - pub fn new(messages: Vec) -> Self { + pub fn new(messages: Vec) -> Self { Self { state: CheckpointState { messages, diff --git a/libs/mcp/client/src/lib.rs b/libs/mcp/client/src/lib.rs index 90bdbb387..f5856b611 100644 --- a/libs/mcp/client/src/lib.rs +++ b/libs/mcp/client/src/lib.rs @@ -7,7 +7,7 @@ use rmcp::{ transport::streamable_http_client::StreamableHttpClientTransportConfig, }; use stakpak_shared::cert_utils::CertificateChain; -use stakpak_shared::models::integrations::openai::ToolCallResultProgress; +use stakpak_shared::models::agent_runtime::ToolCallResultProgress; use std::sync::Arc; use tokio::sync::mpsc::Sender; diff --git a/libs/mcp/client/src/local.rs b/libs/mcp/client/src/local.rs index c3fd8681d..739ad6bc2 100644 --- a/libs/mcp/client/src/local.rs +++ b/libs/mcp/client/src/local.rs @@ -5,7 +5,7 @@ use rmcp::{ service::RunningService, transport::TokioChildProcess, }; -use stakpak_shared::models::integrations::openai::ToolCallResultProgress; +use stakpak_shared::models::agent_runtime::ToolCallResultProgress; use tokio::process::Command; use tokio::sync::mpsc::Sender; diff --git a/libs/mcp/server/src/local_tools.rs b/libs/mcp/server/src/local_tools.rs index 3e8ca5d1b..5317374cb 100644 --- a/libs/mcp/server/src/local_tools.rs +++ b/libs/mcp/server/src/local_tools.rs @@ -17,11 +17,11 @@ use ignore::WalkBuilder; use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT}; use serde_json::json; use similar::TextDiff; -use stakpak_shared::models::async_manifest::{AsyncManifest, PendingToolCall}; -use stakpak_shared::models::integrations::mcp::CallToolResultExt; -use stakpak_shared::models::integrations::openai::{ +use stakpak_shared::models::agent_runtime::{ ProgressType, TaskPauseInfo, TaskUpdate, ToolCallResultProgress, }; +use stakpak_shared::models::async_manifest::{AsyncManifest, PendingToolCall}; +use stakpak_shared::models::integrations::mcp::CallToolResultExt; use stakpak_shared::task_manager::{StartTaskOptions, TaskInfo}; use stakpak_shared::tls_client::{TlsClientConfig, create_tls_client}; use stakpak_shared::utils::{ diff --git a/libs/server/README.md b/libs/server/README.md index 2a0766b49..259747846 100644 --- a/libs/server/README.md +++ b/libs/server/README.md @@ -58,7 +58,7 @@ EventLog - `event_log.rs` — durable replay + live pub/sub - `checkpoint_store.rs` — latest checkpoint envelope storage - `idempotency.rs` — `Idempotency-Key` lookup/replay/conflict logic -- `message_bridge.rs` — temporary storage edge adapters (`ChatMessage` <-> `stakai::Message`) +- Routes and session actors use `stakai::Message` directly at the storage/API edge. - `auth.rs` — bearer auth middleware - `openapi.rs` — generated OpenAPI document definitions - `state.rs` — shared `AppState` and MCP tool cache helpers diff --git a/libs/server/src/lib.rs b/libs/server/src/lib.rs index d1f7b2fd6..5670d167b 100644 --- a/libs/server/src/lib.rs +++ b/libs/server/src/lib.rs @@ -4,7 +4,6 @@ pub mod context; pub mod error; pub mod event_log; pub mod idempotency; -pub mod message_bridge; pub mod openapi; pub mod routes; pub mod sandbox; diff --git a/libs/server/src/message_bridge.rs b/libs/server/src/message_bridge.rs deleted file mode 100644 index 14d59636e..000000000 --- a/libs/server/src/message_bridge.rs +++ /dev/null @@ -1,82 +0,0 @@ -use stakpak_shared::models::{ - integrations::openai::ChatMessage, - llm::LLMMessage, - stakai_adapter::{from_stakai_message, to_stakai_message}, -}; - -/// Storage adapter: local/session storage currently persists ChatMessage payloads. -/// -/// Runtime/API boundaries should stay StakAI-native. These adapters exist only at -/// the persistence edge until storage is migrated to `Vec`. -pub fn chat_to_stakai(messages: Vec) -> Vec { - messages - .into_iter() - .map(LLMMessage::from) - .map(|message| to_stakai_message(&message)) - .collect() -} - -pub fn stakai_to_chat(messages: &[stakai::Message]) -> Vec { - messages - .iter() - .map(from_stakai_message) - .map(ChatMessage::from) - .collect() -} - -#[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - use stakai::{ContentPart, Message, MessageContent, Role}; - - #[test] - fn converts_stakai_to_chat_and_back() { - let messages = vec![ - Message::new(Role::System, "system"), - Message::new(Role::User, "hello"), - Message::new(Role::Assistant, "hi"), - ]; - - let chat = stakai_to_chat(&messages); - let back = chat_to_stakai(chat); - - assert_eq!(back.len(), 3); - assert_eq!(back[0].role, Role::System); - assert_eq!(back[1].role, Role::User); - assert_eq!(back[2].role, Role::Assistant); - assert_eq!(back[1].text(), Some("hello".to_string())); - } - - #[test] - fn preserves_tool_result_tool_call_id_through_storage_adapter() { - let messages = vec![Message::new( - Role::Tool, - MessageContent::Parts(vec![ContentPart::tool_result( - "toolu_01Abc123".to_string(), - json!("result payload"), - )]), - )]; - - let chat = stakai_to_chat(&messages); - assert_eq!(chat.len(), 1); - assert_eq!(chat[0].tool_call_id.as_deref(), Some("toolu_01Abc123")); - - let back = chat_to_stakai(chat); - assert_eq!(back.len(), 1); - assert_eq!(back[0].role, Role::Tool); - - match &back[0].content { - MessageContent::Parts(parts) => { - assert_eq!(parts.len(), 1); - match &parts[0] { - ContentPart::ToolResult { tool_call_id, .. } => { - assert_eq!(tool_call_id, "toolu_01Abc123"); - } - other => panic!("expected ToolResult part, got {other:?}"), - } - } - other => panic!("expected parts content, got {other:?}"), - } - } -} diff --git a/libs/server/src/routes.rs b/libs/server/src/routes.rs index 329c2b7fe..62cc55e97 100644 --- a/libs/server/src/routes.rs +++ b/libs/server/src/routes.rs @@ -2,7 +2,6 @@ use crate::{ auth::{AuthConfig, require_bearer}, context::{ContextFile, ContextPriority}, idempotency::{IdempotencyRequest, LookupResult, StoredResponse}, - message_bridge, session_actor::{ACTIVE_MODEL_METADATA_KEY, spawn_session_actor}, state::AppState, types::{AutoApproveOverride, RunConfig, RunOverrides, SessionRuntimeState}, @@ -663,7 +662,7 @@ async fn get_session_messages_handler( .get_active_checkpoint(session_id) .await .map_err(storage_error)?; - message_bridge::chat_to_stakai(checkpoint.state.messages) + checkpoint.state.messages } Err(error) => { return Err(api_error( diff --git a/libs/server/src/session_actor.rs b/libs/server/src/session_actor.rs index 5a4e4b929..6efe38fb9 100644 --- a/libs/server/src/session_actor.rs +++ b/libs/server/src/session_actor.rs @@ -1,6 +1,5 @@ use crate::{ context::{ContextFile, EnvironmentContext, ProjectContext, SessionContextBuilder}, - message_bridge, sandbox::{SandboxConfig, SandboxMode, SandboxedMcpServer}, state::AppState, types::{RunConfig, SessionHandle}, @@ -104,9 +103,7 @@ async fn run_session_actor( Ok(None) => { let messages = active_checkpoint .as_ref() - .map(|checkpoint| { - message_bridge::chat_to_stakai(checkpoint.state.messages.clone()) - }) + .map(|checkpoint| checkpoint.state.messages.clone()) .unwrap_or_default(); let metadata = active_checkpoint .as_ref() @@ -718,10 +715,8 @@ async fn persist_checkpoint( messages: &[Message], metadata: &serde_json::Value, ) -> Result { - // TODO(ahmed): Migrate server/session checkpoint storage to `Vec` directly - // and remove the ChatMessage adapter conversion (`message_bridge::stakai_to_chat`). - let mut request = CreateCheckpointRequest::new(message_bridge::stakai_to_chat(messages)) - .with_metadata(metadata.clone()); + let mut request = + CreateCheckpointRequest::new(messages.to_vec()).with_metadata(metadata.clone()); if let Some(parent_id) = parent_id { request = request.with_parent(parent_id); diff --git a/libs/shared/src/models/agent_runtime.rs b/libs/shared/src/models/agent_runtime.rs new file mode 100644 index 000000000..bcd7555bb --- /dev/null +++ b/libs/shared/src/models/agent_runtime.rs @@ -0,0 +1,91 @@ +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +/// Tool call result status. +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] +pub enum ToolCallResultStatus { + Success, + Error, + Cancelled, +} + +/// Tool call result. +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct ToolCallResult { + pub call: stakai::ToolCall, + pub result: String, + pub status: ToolCallResultStatus, +} + +/// Streaming progress info for a tool call being generated by the LLM. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct ToolCallStreamInfo { + /// Tool name, which may be empty if it has not streamed yet. + pub name: String, + /// Estimated token count of arguments streamed so far. + pub args_tokens: usize, + /// Optional description extracted from arguments. + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, +} + +/// Tool call result progress update. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct ToolCallResultProgress { + pub id: Uuid, + pub message: String, + /// Type of progress update for specialized handling. + #[serde(skip_serializing_if = "Option::is_none")] + pub progress_type: Option, + /// Structured task updates for task wait progress. + #[serde(skip_serializing_if = "Option::is_none")] + pub task_updates: Option>, + /// Overall progress percentage from 0.0 to 100.0. + #[serde(skip_serializing_if = "Option::is_none")] + pub progress: Option, +} + +/// Type of progress update. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum ProgressType { + /// Command output streaming. + CommandOutput, + /// Task wait progress with structured updates. + TaskWait, + /// Generic progress. + Generic, +} + +/// Structured task status update. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct TaskUpdate { + pub task_id: String, + pub status: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub duration_secs: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_preview: Option, + /// Whether this is a target task being waited on. + #[serde(default)] + pub is_target: bool, + /// Pause information for paused subagent tasks. + #[serde(skip_serializing_if = "Option::is_none")] + pub pause_info: Option, +} + +/// Pause information for subagent tasks awaiting approval. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct TaskPauseInfo { + /// The agent's message before pausing. + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_message: Option, + /// Pending tool calls awaiting approval. + #[serde(skip_serializing_if = "Option::is_none")] + pub pending_tool_calls: Option>, +} + +pub use crate::models::tools::ask_user::{ + AskUserAnswer, AskUserOption, AskUserQuestion, AskUserRequest, AskUserResult, +}; diff --git a/libs/shared/src/models/async_manifest.rs b/libs/shared/src/models/async_manifest.rs index baa67344d..772d10c25 100644 --- a/libs/shared/src/models/async_manifest.rs +++ b/libs/shared/src/models/async_manifest.rs @@ -3,9 +3,9 @@ //! These types represent the JSON output produced by async agent runs //! and provide formatting for human/LLM consumption. -use crate::models::integrations::openai::ToolCall; use crate::models::llm::LLMTokenUsage; use serde::{Deserialize, Serialize}; +use stakai::ToolCall; /// Why an async agent paused execution. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -32,12 +32,10 @@ pub struct PendingToolCall { impl From<&ToolCall> for PendingToolCall { fn from(tc: &ToolCall) -> Self { - let arguments = serde_json::from_str(&tc.function.arguments) - .unwrap_or(serde_json::Value::String(tc.function.arguments.clone())); PendingToolCall { id: tc.id.clone(), - name: tc.function.name.clone(), - arguments, + name: tc.name.clone(), + arguments: tc.arguments.clone(), } } } diff --git a/libs/shared/src/models/integrations/mcp.rs b/libs/shared/src/models/integrations/mcp.rs index a3a3e9755..189d2ee40 100644 --- a/libs/shared/src/models/integrations/mcp.rs +++ b/libs/shared/src/models/integrations/mcp.rs @@ -1,12 +1,11 @@ use rmcp::model::{Annotated, CallToolResult, Content, RawContent}; -use crate::models::integrations::openai::{ChatMessage, MessageContent, ToolCallResultStatus}; +use crate::models::agent_runtime::ToolCallResultStatus; pub trait CallToolResultExt { /// Create a success result with a simple text message fn cancel(content: Option<&Vec>>) -> Self; fn get_status(&self) -> ToolCallResultStatus; - fn get_status_from_chat_message(message: &ChatMessage) -> ToolCallResultStatus; } impl CallToolResultExt for CallToolResult { @@ -45,23 +44,4 @@ impl CallToolResultExt for CallToolResult { ToolCallResultStatus::Success } } - - fn get_status_from_chat_message(message: &ChatMessage) -> ToolCallResultStatus { - match message - .content - .as_ref() - .unwrap_or(&MessageContent::String(String::new())) - .to_string() - .contains("TOOL_CALL_CANCELLED") - || message - .content - .as_ref() - .unwrap_or(&MessageContent::String(String::new())) - .to_string() - .contains("cancelled") - { - true => ToolCallResultStatus::Cancelled, - false => ToolCallResultStatus::Success, - } - } } diff --git a/libs/shared/src/models/integrations/openai.rs b/libs/shared/src/models/integrations/openai.rs index 1c0eeede4..fc0b93faa 100644 --- a/libs/shared/src/models/integrations/openai.rs +++ b/libs/shared/src/models/integrations/openai.rs @@ -1,21 +1,10 @@ -//! OpenAI provider configuration and chat message types +//! OpenAI provider configuration and model metadata //! -//! This module contains: -//! - Configuration types for OpenAI provider -//! - OpenAI model enums with pricing info -//! - Chat message types used throughout the TUI -//! - Tool call types for agent interactions -//! -//! Note: Low-level API request/response types are in `libs/ai/src/providers/openai/`. +//! Agent runtime messages and tool calls live in `stakai`. -use crate::models::llm::{ - GenerationDelta, LLMMessage, LLMMessageContent, LLMMessageImageSource, LLMMessageTypedContent, - LLMTokenUsage, LLMTool, -}; use crate::models::model_pricing::{ContextAware, ContextPricingTier, ModelContextInfo}; use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; -use uuid::Uuid; +use serde_json::Value; // ============================================================================= // Provider Configuration @@ -201,1056 +190,10 @@ impl std::fmt::Display for OpenAIModel { } } -// ============================================================================= -// Message Types (used by TUI) -// ============================================================================= - -/// Message role -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] -#[serde(rename_all = "lowercase")] -pub enum Role { - System, - Developer, - User, - #[default] - Assistant, - Tool, -} - -impl std::fmt::Display for Role { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Role::System => write!(f, "system"), - Role::Developer => write!(f, "developer"), - Role::User => write!(f, "user"), - Role::Assistant => write!(f, "assistant"), - Role::Tool => write!(f, "tool"), - } - } -} - -/// Model info for tracking which model generated a message -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ModelInfo { - /// Provider name (e.g., "anthropic", "openai") - pub provider: String, - /// Model identifier (e.g., "claude-sonnet-4-20250514", "gpt-4") - pub id: String, -} - -/// Chat message -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)] -pub struct ChatMessage { - pub role: Role, - pub content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub name: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_calls: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_call_id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub usage: Option, - - // === Extended fields for session tracking === - /// Unique message identifier - #[serde(skip_serializing_if = "Option::is_none")] - pub id: Option, - /// Model that generated this message (for assistant messages) - #[serde(skip_serializing_if = "Option::is_none")] - pub model: Option, - /// Cost in dollars for this message - #[serde(skip_serializing_if = "Option::is_none")] - pub cost: Option, - /// Why the model stopped: "stop", "tool_calls", "length", "error" - #[serde(skip_serializing_if = "Option::is_none")] - pub finish_reason: Option, - /// Unix timestamp (ms) when message was created/sent - #[serde(skip_serializing_if = "Option::is_none")] - pub created_at: Option, - /// Unix timestamp (ms) when assistant finished generating - #[serde(skip_serializing_if = "Option::is_none")] - pub completed_at: Option, - /// Plugin extensibility - unstructured metadata - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - -impl ChatMessage { - pub fn last_server_message(messages: &[ChatMessage]) -> Option<&ChatMessage> { - messages - .iter() - .rev() - .find(|message| message.role != Role::User && message.role != Role::Tool) - } - - pub fn to_xml(&self) -> String { - match &self.content { - Some(MessageContent::String(s)) => { - format!("{}", self.role, s) - } - Some(MessageContent::Array(parts)) => parts - .iter() - .map(|part| { - format!( - "{}", - self.role, - part.r#type, - part.text.clone().unwrap_or_default() - ) - }) - .collect::>() - .join("\n"), - None => String::new(), - } - } -} - -/// Message content (string or array of parts) -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -#[serde(untagged)] -pub enum MessageContent { - String(String), - Array(Vec), -} - -impl MessageContent { - pub fn inject_checkpoint_id(&self, checkpoint_id: Uuid) -> Self { - match self { - MessageContent::String(s) => MessageContent::String(format!( - "{checkpoint_id}\n{s}" - )), - MessageContent::Array(parts) => MessageContent::Array( - std::iter::once(ContentPart { - r#type: "text".to_string(), - text: Some(format!("{checkpoint_id}")), - image_url: None, - }) - .chain(parts.iter().cloned()) - .collect(), - ), - } - } - - /// All indices from rfind()/find() of ASCII XML tags on same string - #[allow(clippy::string_slice)] - pub fn extract_checkpoint_id(&self) -> Option { - match self { - MessageContent::String(s) => s - .rfind("") - .and_then(|start| { - s[start..] - .find("") - .map(|end| (start + "".len(), start + end)) - }) - .and_then(|(start, end)| Uuid::parse_str(&s[start..end]).ok()), - MessageContent::Array(parts) => parts.iter().rev().find_map(|part| { - part.text.as_deref().and_then(|text| { - text.rfind("") - .and_then(|start| { - text[start..] - .find("") - .map(|end| (start + "".len(), start + end)) - }) - .and_then(|(start, end)| Uuid::parse_str(&text[start..end]).ok()) - }) - }), - } - } -} - -impl std::fmt::Display for MessageContent { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - MessageContent::String(s) => write!(f, "{s}"), - MessageContent::Array(parts) => { - let text_parts: Vec = - parts.iter().filter_map(|part| part.text.clone()).collect(); - write!(f, "{}", text_parts.join("\n")) - } - } - } -} - -impl Default for MessageContent { - fn default() -> Self { - MessageContent::String(String::new()) - } -} - -/// Content part (text or image) -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ContentPart { - pub r#type: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub text: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub image_url: Option, -} - -/// Image URL with optional detail level -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ImageUrl { - pub url: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub detail: Option, -} - -// ============================================================================= -// Tool Types (used by TUI) -// ============================================================================= - -/// Tool definition -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct Tool { - pub r#type: String, - pub function: FunctionDefinition, -} - -/// Function definition for tools -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct FunctionDefinition { - pub name: String, - pub description: Option, - pub parameters: serde_json::Value, -} - -impl From for LLMTool { - fn from(tool: Tool) -> Self { - LLMTool { - name: tool.function.name, - description: tool.function.description.unwrap_or_default(), - input_schema: tool.function.parameters, - } - } -} - -/// Tool choice configuration -#[derive(Debug, Clone, PartialEq)] -pub enum ToolChoice { - Auto, - Required, - Object(ToolChoiceObject), -} - -impl Serialize for ToolChoice { - fn serialize(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - match self { - ToolChoice::Auto => serializer.serialize_str("auto"), - ToolChoice::Required => serializer.serialize_str("required"), - ToolChoice::Object(obj) => obj.serialize(serializer), - } - } -} - -impl<'de> Deserialize<'de> for ToolChoice { - fn deserialize(deserializer: D) -> Result - where - D: serde::Deserializer<'de>, - { - struct ToolChoiceVisitor; - - impl<'de> serde::de::Visitor<'de> for ToolChoiceVisitor { - type Value = ToolChoice; - - fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { - formatter.write_str("string or object") - } - - fn visit_str(self, value: &str) -> Result - where - E: serde::de::Error, - { - match value { - "auto" => Ok(ToolChoice::Auto), - "required" => Ok(ToolChoice::Required), - _ => Err(serde::de::Error::unknown_variant( - value, - &["auto", "required"], - )), - } - } - - fn visit_map(self, map: M) -> Result - where - M: serde::de::MapAccess<'de>, - { - let obj = ToolChoiceObject::deserialize( - serde::de::value::MapAccessDeserializer::new(map), - )?; - Ok(ToolChoice::Object(obj)) - } - } - - deserializer.deserialize_any(ToolChoiceVisitor) - } -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ToolChoiceObject { - pub r#type: String, - pub function: FunctionChoice, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct FunctionChoice { - pub name: String, -} - -/// Tool call from assistant -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ToolCall { - pub id: String, - pub r#type: String, - pub function: FunctionCall, - /// Opaque provider-specific metadata (e.g., Gemini thought_signature) - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - -/// Function call details -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct FunctionCall { - pub name: String, - pub arguments: String, -} - -/// Tool call result status -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub enum ToolCallResultStatus { - Success, - Error, - Cancelled, -} - -/// Tool call result -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ToolCallResult { - pub call: ToolCall, - pub result: String, - pub status: ToolCallResultStatus, -} - -/// Streaming progress info for a tool call being generated by the LLM -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct ToolCallStreamInfo { - /// Tool name (may be empty if not yet streamed) - pub name: String, - /// Estimated token count of arguments streamed so far (~chars/4) - pub args_tokens: usize, - /// Optional description extracted from arguments (best-effort, may be None if JSON incomplete) - #[serde(skip_serializing_if = "Option::is_none")] - pub description: Option, -} - -/// Tool call result progress update -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct ToolCallResultProgress { - pub id: Uuid, - pub message: String, - /// Type of progress update for specialized handling - #[serde(skip_serializing_if = "Option::is_none")] - pub progress_type: Option, - /// Structured task updates for task wait progress - #[serde(skip_serializing_if = "Option::is_none")] - pub task_updates: Option>, - /// Overall progress percentage (0.0 - 100.0) - #[serde(skip_serializing_if = "Option::is_none")] - pub progress: Option, -} - -/// Type of progress update -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub enum ProgressType { - /// Command output streaming - CommandOutput, - /// Task wait progress with structured updates - TaskWait, - /// Generic progress - Generic, -} - -/// Structured task status update -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct TaskUpdate { - pub task_id: String, - pub status: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub description: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub duration_secs: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub output_preview: Option, - /// Whether this is a target task being waited on - #[serde(default)] - pub is_target: bool, - /// Pause information for paused subagent tasks - #[serde(skip_serializing_if = "Option::is_none")] - pub pause_info: Option, -} - -/// Pause information for subagent tasks awaiting approval -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct TaskPauseInfo { - /// The agent's message before pausing - #[serde(skip_serializing_if = "Option::is_none")] - pub agent_message: Option, - /// Pending tool calls awaiting approval - #[serde(skip_serializing_if = "Option::is_none")] - pub pending_tool_calls: Option>, -} - -pub use crate::models::tools::ask_user::{ - AskUserAnswer, AskUserOption, AskUserQuestion, AskUserRequest, AskUserResult, -}; - -/// Chat completion request -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ChatCompletionRequest { - pub model: String, - pub messages: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub frequency_penalty: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub logit_bias: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub logprobs: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub max_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub n: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub presence_penalty: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub response_format: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub seed: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub stop: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub stream: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub temperature: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub top_p: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub tools: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_choice: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub user: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub context: Option, -} - -impl ChatCompletionRequest { - pub fn new( - model: String, - messages: Vec, - tools: Option>, - stream: Option, - ) -> Self { - Self { - model, - messages, - frequency_penalty: None, - logit_bias: None, - logprobs: None, - max_tokens: None, - n: None, - presence_penalty: None, - response_format: None, - seed: None, - stop: None, - stream, - temperature: None, - top_p: None, - tools, - tool_choice: None, - user: None, - context: None, - } - } -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ChatCompletionContext { - pub scratchpad: Option, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ResponseFormat { - pub r#type: String, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -#[serde(untagged)] -pub enum StopSequence { - String(String), - Array(Vec), -} - -/// Chat completion response -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ChatCompletionResponse { - pub id: String, - pub object: String, - pub created: u64, - pub model: String, - pub choices: Vec, - pub usage: LLMTokenUsage, - #[serde(skip_serializing_if = "Option::is_none")] - pub system_fingerprint: Option, - pub metadata: Option, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ChatCompletionChoice { - pub index: usize, - pub message: ChatMessage, - pub logprobs: Option, - pub finish_reason: FinishReason, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -#[serde(rename_all = "snake_case")] -pub enum FinishReason { - Stop, - Length, - ContentFilter, - ToolCalls, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct LogProbs { - pub content: Option>, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct LogProbContent { - pub token: String, - pub logprob: f32, - pub bytes: Option>, - pub top_logprobs: Option>, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct TokenLogprob { - pub token: String, - pub logprob: f32, - pub bytes: Option>, -} - -// ============================================================================= -// Streaming Types -// ============================================================================= - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ChatCompletionStreamResponse { - pub id: String, - pub object: String, - pub created: u64, - pub model: String, - pub choices: Vec, - pub usage: Option, - pub metadata: Option, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ChatCompletionStreamChoice { - pub index: usize, - pub delta: ChatMessageDelta, - pub finish_reason: Option, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ChatMessageDelta { - #[serde(skip_serializing_if = "Option::is_none")] - pub role: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_calls: Option>, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct ToolCallDelta { - pub index: usize, - pub id: Option, - pub r#type: Option, - pub function: Option, - /// Opaque provider-specific metadata (e.g., Gemini thought_signature) - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - -#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -pub struct FunctionCallDelta { - pub name: Option, - pub arguments: Option, -} - -// ============================================================================= -// Conversions -// ============================================================================= - -impl From for ChatMessage { - fn from(llm_message: LLMMessage) -> Self { - let role = match llm_message.role.as_str() { - "system" => Role::System, - "user" => Role::User, - "assistant" => Role::Assistant, - "tool" => Role::Tool, - "developer" => Role::Developer, - _ => Role::User, - }; - - let (content, tool_calls, tool_call_id) = match llm_message.content { - LLMMessageContent::String(text) => (Some(MessageContent::String(text)), None, None), - LLMMessageContent::List(items) => { - let mut text_parts = Vec::new(); - let mut tool_call_parts = Vec::new(); - let mut tool_result_id: Option = None; - - for item in items { - match item { - LLMMessageTypedContent::Text { text } => { - text_parts.push(ContentPart { - r#type: "text".to_string(), - text: Some(text), - image_url: None, - }); - } - LLMMessageTypedContent::ToolCall { - id, - name, - args, - metadata, - } => { - tool_call_parts.push(ToolCall { - id, - r#type: "function".to_string(), - function: FunctionCall { - name, - arguments: args.to_string(), - }, - metadata, - }); - } - LLMMessageTypedContent::ToolResult { - tool_use_id, - content, - } => { - if tool_result_id.is_none() { - tool_result_id = Some(tool_use_id); - } - text_parts.push(ContentPart { - r#type: "text".to_string(), - text: Some(content), - image_url: None, - }); - } - LLMMessageTypedContent::Image { source } => { - text_parts.push(ContentPart { - r#type: "image_url".to_string(), - text: None, - image_url: Some(ImageUrl { - url: format!( - "data:{};base64,{}", - source.media_type, source.data - ), - detail: None, - }), - }); - } - } - } - - let content = if !text_parts.is_empty() { - Some(MessageContent::Array(text_parts)) - } else { - None - }; - - let tool_calls = if !tool_call_parts.is_empty() { - Some(tool_call_parts) - } else { - None - }; - - (content, tool_calls, tool_result_id) - } - }; - - ChatMessage { - role, - content, - name: None, - tool_calls, - tool_call_id, - usage: None, - ..Default::default() - } - } -} - -impl From for LLMMessage { - fn from(chat_message: ChatMessage) -> Self { - let mut content_parts = Vec::new(); - - match chat_message.content { - Some(MessageContent::String(s)) => { - if !s.is_empty() { - content_parts.push(LLMMessageTypedContent::Text { text: s }); - } - } - Some(MessageContent::Array(parts)) => { - for part in parts { - if let Some(text) = part.text { - content_parts.push(LLMMessageTypedContent::Text { text }); - } else if let Some(image_url) = part.image_url { - let (media_type, data) = if image_url.url.starts_with("data:") { - let parts: Vec<&str> = image_url.url.splitn(2, ',').collect(); - if parts.len() == 2 { - let meta = parts[0]; - let data = parts[1]; - let media_type = meta - .trim_start_matches("data:") - .trim_end_matches(";base64") - .to_string(); - (media_type, data.to_string()) - } else { - ("image/jpeg".to_string(), image_url.url) - } - } else { - ("image/jpeg".to_string(), image_url.url) - }; - - content_parts.push(LLMMessageTypedContent::Image { - source: LLMMessageImageSource { - r#type: "base64".to_string(), - media_type, - data, - }, - }); - } - } - } - None => {} - } - - if let Some(tool_calls) = chat_message.tool_calls { - for tool_call in tool_calls { - let args = serde_json::from_str(&tool_call.function.arguments).unwrap_or(json!({})); - content_parts.push(LLMMessageTypedContent::ToolCall { - id: tool_call.id, - name: tool_call.function.name, - args, - metadata: tool_call.metadata, - }); - } - } - - // Handle tool result messages: when role is Tool and tool_call_id is present, - // convert the content to a ToolResult content part. This is the generic - // intermediate representation - each provider's conversion layer handles - // the specifics (e.g., Anthropic converts to user role with tool_result blocks) - if chat_message.role == Role::Tool - && let Some(tool_call_id) = chat_message.tool_call_id - { - // Extract content as string for the tool result - let content_str = content_parts - .iter() - .filter_map(|p| match p { - LLMMessageTypedContent::Text { text } => Some(text.clone()), - _ => None, - }) - .collect::>() - .join("\n"); - - // Replace content with a single ToolResult - content_parts = vec![LLMMessageTypedContent::ToolResult { - tool_use_id: tool_call_id, - content: content_str, - }]; - } - - LLMMessage { - role: chat_message.role.to_string(), - content: if content_parts.is_empty() { - LLMMessageContent::String(String::new()) - } else if content_parts.len() == 1 { - match &content_parts[0] { - LLMMessageTypedContent::Text { text } => { - LLMMessageContent::String(text.clone()) - } - _ => LLMMessageContent::List(content_parts), - } - } else { - LLMMessageContent::List(content_parts) - }, - } - } -} - -impl From for ChatMessageDelta { - fn from(delta: GenerationDelta) -> Self { - match delta { - GenerationDelta::Content { content } => ChatMessageDelta { - role: Some(Role::Assistant), - content: Some(content), - tool_calls: None, - }, - GenerationDelta::Thinking { thinking: _ } => ChatMessageDelta { - role: Some(Role::Assistant), - content: None, - tool_calls: None, - }, - GenerationDelta::ToolUse { tool_use } => ChatMessageDelta { - role: Some(Role::Assistant), - content: None, - tool_calls: Some(vec![ToolCallDelta { - index: tool_use.index, - id: tool_use.id, - r#type: Some("function".to_string()), - function: Some(FunctionCallDelta { - name: tool_use.name, - arguments: tool_use.input, - }), - metadata: tool_use.metadata, - }]), - }, - _ => ChatMessageDelta { - role: Some(Role::Assistant), - content: None, - tool_calls: None, - }, - } - } -} - #[cfg(test)] mod tests { use super::*; - - #[test] - fn test_serialize_basic_request() { - let request = ChatCompletionRequest { - model: "gpt-4".to_string(), - messages: vec![ - ChatMessage { - role: Role::System, - content: Some(MessageContent::String( - "You are a helpful assistant.".to_string(), - )), - name: None, - tool_calls: None, - tool_call_id: None, - usage: None, - ..Default::default() - }, - ChatMessage { - role: Role::User, - content: Some(MessageContent::String("Hello!".to_string())), - name: None, - tool_calls: None, - tool_call_id: None, - usage: None, - ..Default::default() - }, - ], - frequency_penalty: None, - logit_bias: None, - logprobs: None, - max_tokens: Some(100), - n: None, - presence_penalty: None, - response_format: None, - seed: None, - stop: None, - stream: None, - temperature: Some(0.7), - top_p: None, - tools: None, - tool_choice: None, - user: None, - context: None, - }; - - let json = serde_json::to_string(&request).unwrap(); - assert!(json.contains("\"model\":\"gpt-4\"")); - assert!(json.contains("\"messages\":[")); - assert!(json.contains("\"role\":\"system\"")); - } - - #[test] - fn test_llm_message_to_chat_message() { - let llm_message = LLMMessage { - role: "user".to_string(), - content: LLMMessageContent::String("Hello, world!".to_string()), - }; - - let chat_message = ChatMessage::from(llm_message); - assert_eq!(chat_message.role, Role::User); - match &chat_message.content { - Some(MessageContent::String(text)) => assert_eq!(text, "Hello, world!"), - _ => panic!("Expected string content"), - } - } - - #[test] - fn test_llm_message_to_chat_message_tool_result_preserves_tool_call_id() { - let llm_message = LLMMessage { - role: "tool".to_string(), - content: LLMMessageContent::List(vec![LLMMessageTypedContent::ToolResult { - tool_use_id: "toolu_01Abc123".to_string(), - content: "Tool execution result".to_string(), - }]), - }; - - let chat_message = ChatMessage::from(llm_message); - assert_eq!(chat_message.role, Role::Tool); - assert_eq!(chat_message.tool_call_id.as_deref(), Some("toolu_01Abc123")); - assert_eq!( - chat_message.content, - Some(MessageContent::Array(vec![ContentPart { - r#type: "text".to_string(), - text: Some("Tool execution result".to_string()), - image_url: None, - }])) - ); - } - - #[test] - fn test_chat_message_to_llm_message_tool_result() { - // Test that Tool role messages with tool_call_id are converted to ToolResult content - // This is critical for Anthropic compatibility - the provider layer converts - // role="tool" to role="user" with tool_result content blocks - let chat_message = ChatMessage { - role: Role::Tool, - content: Some(MessageContent::String("Tool execution result".to_string())), - name: None, - tool_calls: None, - tool_call_id: Some("toolu_01Abc123".to_string()), - usage: None, - ..Default::default() - }; - - let llm_message: LLMMessage = chat_message.into(); - - // Role should be preserved as "tool" - provider layer handles conversion - assert_eq!(llm_message.role, "tool"); - - // Content should be a ToolResult with the tool_call_id - match &llm_message.content { - LLMMessageContent::List(parts) => { - assert_eq!(parts.len(), 1, "Should have exactly one content part"); - match &parts[0] { - LLMMessageTypedContent::ToolResult { - tool_use_id, - content, - } => { - assert_eq!(tool_use_id, "toolu_01Abc123"); - assert_eq!(content, "Tool execution result"); - } - _ => panic!("Expected ToolResult content part, got {:?}", parts[0]), - } - } - _ => panic!( - "Expected List content with ToolResult, got {:?}", - llm_message.content - ), - } - } - - #[test] - fn test_chat_message_to_llm_message_tool_result_empty_content() { - // Test tool result with empty content - let chat_message = ChatMessage { - role: Role::Tool, - content: None, - name: None, - tool_calls: None, - tool_call_id: Some("toolu_02Xyz789".to_string()), - usage: None, - ..Default::default() - }; - - let llm_message: LLMMessage = chat_message.into(); - - assert_eq!(llm_message.role, "tool"); - match &llm_message.content { - LLMMessageContent::List(parts) => { - assert_eq!(parts.len(), 1); - match &parts[0] { - LLMMessageTypedContent::ToolResult { - tool_use_id, - content, - } => { - assert_eq!(tool_use_id, "toolu_02Xyz789"); - assert_eq!(content, ""); // Empty content - } - _ => panic!("Expected ToolResult content part"), - } - } - _ => panic!("Expected List content with ToolResult"), - } - } - - #[test] - fn test_chat_message_to_llm_message_assistant_with_tool_calls() { - // Test that assistant messages with tool_calls are converted correctly - let chat_message = ChatMessage { - role: Role::Assistant, - content: Some(MessageContent::String( - "I'll help you with that.".to_string(), - )), - name: None, - tool_calls: Some(vec![ToolCall { - id: "call_abc123".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "get_weather".to_string(), - arguments: r#"{"location": "Paris"}"#.to_string(), - }, - metadata: None, - }]), - tool_call_id: None, - usage: None, - ..Default::default() - }; - - let llm_message: LLMMessage = chat_message.into(); - - assert_eq!(llm_message.role, "assistant"); - match &llm_message.content { - LLMMessageContent::List(parts) => { - assert_eq!(parts.len(), 2, "Should have text and tool call"); - - // First part should be text - match &parts[0] { - LLMMessageTypedContent::Text { text } => { - assert_eq!(text, "I'll help you with that."); - } - _ => panic!("Expected Text content part first"), - } - - // Second part should be tool call - match &parts[1] { - LLMMessageTypedContent::ToolCall { id, name, args, .. } => { - assert_eq!(id, "call_abc123"); - assert_eq!(name, "get_weather"); - assert_eq!(args["location"], "Paris"); - } - _ => panic!("Expected ToolCall content part second"), - } - } - _ => panic!("Expected List content"), - } - } + use serde_json::json; #[test] fn test_extract_chatgpt_account_id_from_access_token() { diff --git a/libs/shared/src/models/llm.rs b/libs/shared/src/models/llm.rs index bed7e2b9c..7d7ff1a7b 100644 --- a/libs/shared/src/models/llm.rs +++ b/libs/shared/src/models/llm.rs @@ -52,7 +52,6 @@ //! ``` use serde::{Deserialize, Serialize}; -use stakai::Model; use std::collections::HashMap; use super::auth::ProviderAuth; @@ -740,213 +739,6 @@ pub struct LLMGoogleOptions { pub thinking_budget: Option, } -#[derive(Clone, Debug, Serialize)] -pub struct LLMInput { - pub model: Model, - pub messages: Vec, - pub max_tokens: u32, - pub tools: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub provider_options: Option, - /// Custom headers to pass to the inference provider - #[serde(skip_serializing_if = "Option::is_none")] - pub headers: Option>, -} - -#[derive(Debug)] -pub struct LLMStreamInput { - pub model: Model, - pub messages: Vec, - pub max_tokens: u32, - pub stream_channel_tx: tokio::sync::mpsc::Sender, - pub tools: Option>, - pub provider_options: Option, - /// Custom headers to pass to the inference provider - pub headers: Option>, -} - -impl From<&LLMStreamInput> for LLMInput { - fn from(value: &LLMStreamInput) -> Self { - LLMInput { - model: value.model.clone(), - messages: value.messages.clone(), - max_tokens: value.max_tokens, - tools: value.tools.clone(), - provider_options: value.provider_options.clone(), - headers: value.headers.clone(), - } - } -} - -#[derive(Serialize, Deserialize, Debug, Clone, Default)] -pub struct LLMMessage { - pub role: String, - pub content: LLMMessageContent, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct SimpleLLMMessage { - #[serde(rename = "role")] - pub role: SimpleLLMRole, - pub content: String, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -#[serde(rename_all = "lowercase")] -pub enum SimpleLLMRole { - User, - Assistant, -} - -impl std::fmt::Display for SimpleLLMRole { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - SimpleLLMRole::User => write!(f, "user"), - SimpleLLMRole::Assistant => write!(f, "assistant"), - } - } -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -#[serde(untagged)] -pub enum LLMMessageContent { - String(String), - List(Vec), -} - -#[allow(clippy::to_string_trait_impl)] -impl ToString for LLMMessageContent { - fn to_string(&self) -> String { - match self { - LLMMessageContent::String(s) => s.clone(), - LLMMessageContent::List(l) => l - .iter() - .map(|c| match c { - LLMMessageTypedContent::Text { text } => text.clone(), - LLMMessageTypedContent::ToolCall { .. } => String::new(), - LLMMessageTypedContent::ToolResult { content, .. } => content.clone(), - LLMMessageTypedContent::Image { .. } => String::new(), - }) - .collect::>() - .join("\n"), - } - } -} - -impl From for LLMMessageContent { - fn from(value: String) -> Self { - LLMMessageContent::String(value) - } -} - -impl Default for LLMMessageContent { - fn default() -> Self { - LLMMessageContent::String(String::new()) - } -} - -impl LLMMessageContent { - /// Convert into a Vec of typed content parts. - /// A `String` variant is returned as a single `Text` part (empty strings yield an empty vec). - pub fn into_parts(self) -> Vec { - match self { - LLMMessageContent::List(parts) => parts, - LLMMessageContent::String(s) if s.is_empty() => vec![], - LLMMessageContent::String(s) => vec![LLMMessageTypedContent::Text { text: s }], - } - } -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -#[serde(tag = "type")] -pub enum LLMMessageTypedContent { - #[serde(rename = "text")] - Text { text: String }, - #[serde(rename = "tool_use")] - ToolCall { - id: String, - name: String, - #[serde(alias = "input")] - args: serde_json::Value, - /// Opaque provider-specific metadata (e.g., Gemini thought_signature). - #[serde(skip_serializing_if = "Option::is_none")] - metadata: Option, - }, - #[serde(rename = "tool_result")] - ToolResult { - tool_use_id: String, - content: String, - }, - #[serde(rename = "image")] - Image { source: LLMMessageImageSource }, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct LLMMessageImageSource { - #[serde(rename = "type")] - pub r#type: String, - pub media_type: String, - pub data: String, -} - -impl Default for LLMMessageTypedContent { - fn default() -> Self { - LLMMessageTypedContent::Text { - text: String::new(), - } - } -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct LLMChoice { - pub finish_reason: Option, - pub index: u32, - pub message: LLMMessage, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct LLMCompletionResponse { - pub model: String, - pub object: String, - pub choices: Vec, - pub created: u64, - pub usage: Option, - pub id: String, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct LLMStreamDelta { - #[serde(skip_serializing_if = "Option::is_none")] - pub content: Option, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct LLMStreamChoice { - pub finish_reason: Option, - pub index: u32, - pub message: Option, - pub delta: LLMStreamDelta, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct LLMCompletionStreamResponse { - pub model: String, - pub object: String, - pub choices: Vec, - pub created: u64, - #[serde(skip_serializing_if = "Option::is_none")] - pub usage: Option, - pub id: String, - pub citations: Option>, -} - -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] -pub struct LLMTool { - pub name: String, - pub description: String, - pub input_schema: serde_json::Value, -} - #[derive(Default, Debug, Serialize, Deserialize, Clone, PartialEq)] pub struct LLMTokenUsage { pub prompt_tokens: u32, @@ -1029,27 +821,6 @@ impl std::ops::AddAssign for PromptTokensDetails { } } -#[derive(Serialize, Deserialize, Debug, Clone)] -#[serde(tag = "type")] -pub enum GenerationDelta { - Content { content: String }, - Thinking { thinking: String }, - ToolUse { tool_use: GenerationDeltaToolUse }, - Usage { usage: LLMTokenUsage }, - Metadata { metadata: serde_json::Value }, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct GenerationDeltaToolUse { - pub id: Option, - pub name: Option, - pub input: Option, - pub index: usize, - /// Opaque provider-specific metadata (e.g., Gemini thought_signature) - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, -} - #[cfg(test)] mod tests { use super::*; diff --git a/libs/shared/src/models/mod.rs b/libs/shared/src/models/mod.rs index be677a32f..d1cc721d7 100644 --- a/libs/shared/src/models/mod.rs +++ b/libs/shared/src/models/mod.rs @@ -1,3 +1,4 @@ +pub mod agent_runtime; pub mod async_manifest; pub mod auth; pub mod billing; @@ -9,5 +10,4 @@ pub mod llm; pub mod model_pricing; pub mod openai_runtime; pub mod overrides; -pub mod stakai_adapter; pub mod tools; diff --git a/libs/shared/src/models/stakai_adapter.rs b/libs/shared/src/models/stakai_adapter.rs deleted file mode 100644 index cd9f0c317..000000000 --- a/libs/shared/src/models/stakai_adapter.rs +++ /dev/null @@ -1,1863 +0,0 @@ -//! Adapter layer between CLI LLM types and StakAI SDK -//! -//! This module provides conversion functions and a wrapper client to use the StakAI SDK -//! with the CLI's existing LLM types, enabling BYOM (Bring Your Own Model) functionality. - -use crate::models::error::{AgentError, BadRequestErrorMessage}; -use crate::models::llm::{ - GenerationDelta, GenerationDeltaToolUse, LLMChoice, LLMCompletionResponse, LLMInput, - LLMMessage, LLMMessageContent, LLMMessageImageSource, LLMMessageTypedContent, - LLMProviderConfig, LLMProviderOptions, LLMStreamInput, LLMTokenUsage, LLMTool, ProviderConfig, -}; -use crate::models::openai_runtime::{OpenAIBackendResolutionInput, resolve_openai_runtime}; -use futures::StreamExt; -use stakai::{ - AnthropicOptions, ContentPart, FinishReason, GenerateOptions, GenerateRequest, - GenerateResponse, GoogleOptions, Headers, Inference, InferenceConfig, Message, MessageContent, - Model, OpenAIApiConfig, OpenAIOptions, ProviderOptions, ReasoningEffort, ResponsesConfig, Role, - StreamEvent, ThinkingOptions, Tool, ToolFunction, Usage, - providers::anthropic::AnthropicConfig as StakaiAnthropicConfig, - providers::openai::OpenAIConfig as StakaiOpenAIConfig, registry::ProviderRegistry, -}; - -/// Convert CLI LLMMessage to StakAI Message -pub fn to_stakai_message(msg: &LLMMessage) -> Message { - let role = match msg.role.as_str() { - "system" => Role::System, - "user" => Role::User, - "assistant" => Role::Assistant, - "tool" => Role::Tool, - _ => Role::User, - }; - - let content: MessageContent = match &msg.content { - LLMMessageContent::String(s) => s.clone().into(), - LLMMessageContent::List(parts) => { - let content_parts: Vec = - parts.iter().map(to_stakai_content_part).collect(); - content_parts.into() - } - }; - - Message::new(role, content) -} - -/// Convert CLI LLMMessageTypedContent to StakAI ContentPart -fn to_stakai_content_part(part: &LLMMessageTypedContent) -> ContentPart { - match part { - LLMMessageTypedContent::Text { text } => ContentPart::text(text), - LLMMessageTypedContent::Image { source } => { - // Convert base64 image to data URI - ContentPart::image(format!("data:{};base64,{}", source.media_type, source.data)) - } - LLMMessageTypedContent::ToolCall { - id, - name, - args, - metadata, - } => { - let mut part = ContentPart::tool_call(id, name, args.clone()); - if let ContentPart::ToolCall { - metadata: ref mut m, - .. - } = part - { - *m = metadata.clone(); - } - part - } - LLMMessageTypedContent::ToolResult { - tool_use_id, - content, - } => ContentPart::tool_result(tool_use_id, serde_json::Value::String(content.clone())), - } -} - -/// Convert StakAI Message to CLI LLMMessage -pub fn from_stakai_message(msg: &Message) -> LLMMessage { - let role = match msg.role { - Role::System => "system", - Role::User => "user", - Role::Assistant => "assistant", - Role::Tool => "tool", - } - .to_string(); - - let content = match &msg.content { - MessageContent::Text(s) => LLMMessageContent::String(s.clone()), - MessageContent::Parts(parts) => { - LLMMessageContent::List(parts.iter().map(from_stakai_content_part).collect()) - } - }; - - LLMMessage { role, content } -} - -/// Convert StakAI ContentPart to CLI LLMMessageTypedContent -fn from_stakai_content_part(part: &ContentPart) -> LLMMessageTypedContent { - match part { - ContentPart::Text { text, .. } => LLMMessageTypedContent::Text { text: text.clone() }, - ContentPart::Image { url, .. } => { - // Parse data URI back to base64 - let (media_type, data) = if url.starts_with("data:") { - let parts: Vec<&str> = url.splitn(2, ',').collect(); - if parts.len() == 2 { - let media = parts[0] - .strip_prefix("data:") - .unwrap_or("image/png") - .strip_suffix(";base64") - .unwrap_or("image/png"); - (media.to_string(), parts[1].to_string()) - } else { - ("image/png".to_string(), url.clone()) - } - } else { - ("image/png".to_string(), url.clone()) - }; - - LLMMessageTypedContent::Image { - source: LLMMessageImageSource { - r#type: "base64".to_string(), - media_type, - data, - }, - } - } - ContentPart::ToolCall { - id, - name, - arguments, - metadata, - .. - } => LLMMessageTypedContent::ToolCall { - id: id.clone(), - name: name.clone(), - args: arguments.clone(), - metadata: metadata.clone(), - }, - ContentPart::ToolResult { - tool_call_id, - content, - .. - } => LLMMessageTypedContent::ToolResult { - tool_use_id: tool_call_id.clone(), - content: match content { - serde_json::Value::String(s) => s.clone(), - other => other.to_string(), - }, - }, - } -} - -/// Convert CLI LLMTool to StakAI Tool -pub fn to_stakai_tool(tool: &LLMTool) -> Tool { - Tool { - tool_type: "function".to_string(), - function: ToolFunction { - name: tool.name.clone(), - description: tool.description.clone(), - parameters: tool.input_schema.clone(), - }, - provider_options: None, - } -} - -/// Convert StakAI Tool to CLI LLMTool -pub fn from_stakai_tool(tool: &Tool) -> LLMTool { - LLMTool { - name: tool.function.name.clone(), - description: tool.function.description.clone(), - input_schema: tool.function.parameters.clone(), - } -} - -/// Convert StakAI StreamEvent to CLI GenerationDelta -pub fn from_stakai_stream_event(event: &StreamEvent) -> Option { - match event { - StreamEvent::TextDelta { delta, .. } => Some(GenerationDelta::Content { - content: delta.clone(), - }), - StreamEvent::ReasoningDelta { delta, .. } => Some(GenerationDelta::Thinking { - thinking: delta.clone(), - }), - StreamEvent::ToolCallStart { id, name } => Some(GenerationDelta::ToolUse { - tool_use: GenerationDeltaToolUse { - id: Some(id.clone()), - name: Some(name.clone()), - input: None, - index: 0, - metadata: None, - }, - }), - StreamEvent::ToolCallDelta { id, delta } => Some(GenerationDelta::ToolUse { - tool_use: GenerationDeltaToolUse { - id: Some(id.clone()), - name: None, - input: Some(delta.clone()), - index: 0, - metadata: None, - }, - }), - StreamEvent::ToolCallEnd { - id, name, metadata, .. - } => { - // ToolCallEnd signals completion - don't emit arguments here as they - // were already accumulated via ToolCallDelta events. Including them - // would cause doubling for providers like Anthropic that stream deltas. - // We emit name and metadata (e.g. Gemini thought_signature). - Some(GenerationDelta::ToolUse { - tool_use: GenerationDeltaToolUse { - id: Some(id.clone()), - name: Some(name.clone()), - input: None, - index: 0, - metadata: metadata.clone(), - }, - }) - } - StreamEvent::Finish { usage, .. } => { - let llm_usage = from_stakai_usage(usage); - Some(GenerationDelta::Usage { usage: llm_usage }) - } - StreamEvent::Start { .. } | StreamEvent::Error { .. } => None, - } -} - -/// Convert StakAI Usage to CLI LLMTokenUsage -pub fn from_stakai_usage(usage: &Usage) -> LLMTokenUsage { - use crate::models::llm::PromptTokensDetails; - - // Convert StakAI input_token_details to CLI PromptTokensDetails - let prompt_tokens_details = usage.input_token_details.as_ref().map(|details| { - PromptTokensDetails { - input_tokens: details.no_cache, - output_tokens: None, // Output tokens are tracked separately - cache_read_input_tokens: details.cache_read, - cache_write_input_tokens: details.cache_write, - } - }); - - LLMTokenUsage { - prompt_tokens: usage.prompt_tokens, - completion_tokens: usage.completion_tokens, - total_tokens: usage.total_tokens, - prompt_tokens_details, - } -} - -/// Convert StakAI FinishReason to string -pub fn finish_reason_to_string(reason: &FinishReason) -> String { - // If raw value is available, use it; otherwise use the unified reason - if let Some(raw) = &reason.raw { - raw.clone() - } else { - match reason.unified { - stakai::FinishReasonKind::Stop => "stop".to_string(), - stakai::FinishReasonKind::Length => "length".to_string(), - stakai::FinishReasonKind::ContentFilter => "content_filter".to_string(), - stakai::FinishReasonKind::ToolCalls => "tool_calls".to_string(), - stakai::FinishReasonKind::Error => "error".to_string(), - stakai::FinishReasonKind::Other => "other".to_string(), - } - } -} - -/// Convert CLI LLMProviderOptions to StakAI ProviderOptions -pub fn to_stakai_provider_options( - opts: &LLMProviderOptions, - model: &Model, -) -> Option { - // Determine provider from model's provider field - match model.provider.as_str() { - "anthropic" => { - if let Some(anthropic) = &opts.anthropic { - let thinking = anthropic - .thinking - .as_ref() - .map(|t| ThinkingOptions::new(t.budget_tokens)); - - Some(ProviderOptions::Anthropic(AnthropicOptions { - thinking, - effort: None, - })) - } else { - None - } - } - "openai" => { - opts.openai.as_ref().map(|openai| { - let reasoning_effort = openai.reasoning_effort.as_ref().and_then(|e| { - match e.to_lowercase().as_str() { - "low" => Some(ReasoningEffort::Low), - "medium" => Some(ReasoningEffort::Medium), - "high" => Some(ReasoningEffort::High), - _ => None, - } - }); - - ProviderOptions::OpenAI(OpenAIOptions { - api_config: Some(OpenAIApiConfig::Responses(ResponsesConfig { - reasoning_effort, - reasoning_summary: None, - session_id: None, - service_tier: None, - cache_retention: None, - })), - system_message_mode: None, - store: None, - user: None, - }) - }) - } - "google" | "gemini" => opts.google.as_ref().map(|google| { - ProviderOptions::Google(GoogleOptions { - thinking_budget: google.thinking_budget, - cached_content: None, - }) - }), - _ => { - // For custom/unknown providers, try to infer from which options are set - if let Some(anthropic) = &opts.anthropic { - let thinking = anthropic - .thinking - .as_ref() - .map(|t| ThinkingOptions::new(t.budget_tokens)); - Some(ProviderOptions::Anthropic(AnthropicOptions { - thinking, - effort: None, - })) - } else if let Some(openai) = &opts.openai { - let reasoning_effort = openai.reasoning_effort.as_ref().and_then(|e| { - match e.to_lowercase().as_str() { - "low" => Some(ReasoningEffort::Low), - "medium" => Some(ReasoningEffort::Medium), - "high" => Some(ReasoningEffort::High), - _ => None, - } - }); - Some(ProviderOptions::OpenAI(OpenAIOptions { - api_config: Some(OpenAIApiConfig::Responses(ResponsesConfig { - reasoning_effort, - reasoning_summary: None, - session_id: None, - service_tier: None, - cache_retention: None, - })), - system_message_mode: None, - store: None, - user: None, - })) - } else { - opts.google.as_ref().map(|google| { - ProviderOptions::Google(GoogleOptions { - thinking_budget: google.thinking_budget, - cached_content: None, - }) - }) - } - } - } -} - -/// Convert StakAI GenerateResponse to CLI LLMCompletionResponse -pub fn from_stakai_response(response: GenerateResponse, model: &str) -> LLMCompletionResponse { - let mut content_parts: Vec = Vec::new(); - - for content in &response.content { - match content { - stakai::ResponseContent::Text { text } => { - content_parts.push(LLMMessageTypedContent::Text { text: text.clone() }); - } - stakai::ResponseContent::Reasoning { reasoning } => { - // Include reasoning as a text block with a prefix for visibility - // This matches how Anthropic's thinking is typically displayed - content_parts.push(LLMMessageTypedContent::Text { - text: format!("[Reasoning: {}]", reasoning), - }); - } - stakai::ResponseContent::ToolCall(tool_call) => { - content_parts.push(LLMMessageTypedContent::ToolCall { - id: tool_call.id.clone(), - name: tool_call.name.clone(), - args: tool_call.arguments.clone(), - metadata: tool_call.metadata.clone(), - }); - } - } - } - - let message_content = if content_parts.len() == 1 { - if let LLMMessageTypedContent::Text { text } = &content_parts[0] { - LLMMessageContent::String(text.clone()) - } else { - LLMMessageContent::List(content_parts) - } - } else { - LLMMessageContent::List(content_parts) - }; - - LLMCompletionResponse { - id: uuid::Uuid::new_v4().to_string(), - model: model.to_string(), - object: "chat.completion".to_string(), - choices: vec![LLMChoice { - finish_reason: Some(finish_reason_to_string(&response.finish_reason)), - index: 0, - message: LLMMessage { - role: "assistant".to_string(), - content: message_content, - }, - }], - created: chrono::Utc::now().timestamp_millis() as u64, - usage: Some(from_stakai_usage(&response.usage)), - } -} - -fn resolve_stakai_openai_config( - provider_config: &ProviderConfig, -) -> Result, String> { - let resolved = resolve_openai_runtime(OpenAIBackendResolutionInput::new( - Some(provider_config.clone()), - provider_config.get_auth(), - )) - .map_err(|error| format!("Failed to resolve OpenAI runtime config: {}", error))?; - - Ok(resolved.map(|config| config.to_stakai_config())) -} - -/// Build StakAI InferenceConfig from CLI LLMProviderConfig -pub fn build_inference_config(config: &LLMProviderConfig) -> Result { - let mut inference_config = InferenceConfig::new(); - - for (name, provider_config) in &config.providers { - match provider_config { - ProviderConfig::OpenAI { .. } => { - if let Some(openai_config) = resolve_stakai_openai_config(provider_config)? { - inference_config = inference_config.openai_config(openai_config); - } - } - ProviderConfig::Anthropic { api_endpoint, .. } => { - // Use get_auth() to resolve credentials - // Check for OAuth access token first, then API key - let stakai_config = if let Some(token) = provider_config.access_token() { - // OAuth authentication - uses Bearer token header - let mut cfg = StakaiAnthropicConfig::with_oauth(token); - if let Some(endpoint) = api_endpoint { - cfg = cfg.with_base_url(endpoint); - } - Some(cfg) - } else if let Some(key) = provider_config.api_key() { - // API key authentication - uses x-api-key header - let mut cfg = StakaiAnthropicConfig::new(key); - if let Some(endpoint) = api_endpoint { - cfg = cfg.with_base_url(endpoint); - } - Some(cfg) - } else { - None - }; - - if let Some(cfg) = stakai_config { - inference_config = inference_config.anthropic_config(cfg); - } - } - ProviderConfig::Gemini { api_endpoint, .. } => { - if let Some(api_key) = provider_config.api_key() { - inference_config = - inference_config.gemini(api_key.to_string(), api_endpoint.clone()); - } - } - ProviderConfig::Stakpak { api_endpoint, .. } => { - // Skip if no api_key - stakpak is optional - if let Some(api_key) = provider_config.api_key() { - inference_config = - inference_config.stakpak(api_key.to_string(), api_endpoint.clone()); - } - } - ProviderConfig::GitHubCopilot { .. } => { - tracing::debug!( - provider = %name, - "GitHub Copilot provider is not configured via InferenceConfig; \ - it will be routed through the provider registry instead" - ); - } - ProviderConfig::Custom { .. } => { - // Custom providers are handled by build_provider_registry_direct - // InferenceConfig doesn't support custom providers directly - let _ = name; // Suppress unused warning - } - ProviderConfig::Bedrock { - region, - profile_name, - } => { - #[cfg(feature = "bedrock")] - { - use stakai::providers::bedrock::BedrockConfig; - let mut bedrock_config = BedrockConfig::new(region.clone()); - if let Some(profile) = profile_name { - bedrock_config = bedrock_config.with_profile_name(profile.clone()); - } - inference_config = inference_config.bedrock_config(bedrock_config); - } - #[cfg(not(feature = "bedrock"))] - { - let _ = (name, region, profile_name); - tracing::warn!( - "Bedrock provider configured but bedrock feature is not enabled" - ); - } - } - } - } - - Ok(inference_config) -} - -/// Build a ProviderRegistry directly with all providers including custom ones -fn build_provider_registry_direct(config: &LLMProviderConfig) -> Result { - use stakai::providers::anthropic::{ - AnthropicConfig as StakaiAnthropicConfig, AnthropicProvider, - }; - use stakai::providers::copilot::{CopilotConfig, CopilotProvider}; - use stakai::providers::gemini::{GeminiConfig as StakaiGeminiConfig, GeminiProvider}; - use stakai::providers::openai::{OpenAIConfig as StakaiOpenAIConfig, OpenAIProvider}; - use stakai::providers::stakpak::{StakpakProvider, StakpakProviderConfig}; - - let mut registry = ProviderRegistry::new(); - - for (name, provider_config) in &config.providers { - match provider_config { - ProviderConfig::OpenAI { .. } => { - if let Some(openai_config) = resolve_stakai_openai_config(provider_config)? { - let provider = OpenAIProvider::new(openai_config) - .map_err(|e| format!("Failed to create OpenAI provider: {}", e))?; - registry = registry.register("openai", provider); - } - } - ProviderConfig::Anthropic { api_endpoint, .. } => { - let stakai_config = if let Some(token) = provider_config.access_token() { - let mut cfg = StakaiAnthropicConfig::with_oauth(token); - if let Some(endpoint) = api_endpoint { - cfg = cfg.with_base_url(endpoint); - } - Some(cfg) - } else if let Some(key) = provider_config.api_key() { - let mut cfg = StakaiAnthropicConfig::new(key); - if let Some(endpoint) = api_endpoint { - cfg = cfg.with_base_url(endpoint); - } - Some(cfg) - } else { - None - }; - - if let Some(cfg) = stakai_config { - let provider = AnthropicProvider::new(cfg) - .map_err(|e| format!("Failed to create Anthropic provider: {}", e))?; - registry = registry.register("anthropic", provider); - } - } - ProviderConfig::Gemini { api_endpoint, .. } => { - if let Some(api_key) = provider_config.api_key() { - let mut gemini_config = StakaiGeminiConfig::new(api_key.to_string()); - if let Some(endpoint) = api_endpoint { - gemini_config = gemini_config.with_base_url(endpoint.clone()); - } - let provider = GeminiProvider::new(gemini_config) - .map_err(|e| format!("Failed to create Gemini provider: {}", e))?; - registry = registry.register("google", provider); - } - } - ProviderConfig::Stakpak { api_endpoint, .. } => { - // Skip if no api_key - stakpak is optional - let Some(api_key) = provider_config.api_key() else { - continue; - }; - let mut stakpak_config = StakpakProviderConfig::new(api_key.to_string()) - .with_user_agent(format!("Stakpak/{}", env!("CARGO_PKG_VERSION"))); - if let Some(endpoint) = api_endpoint { - stakpak_config = stakpak_config.with_base_url(endpoint.clone()); - } - let provider = StakpakProvider::new(stakpak_config) - .map_err(|e| format!("Failed to create Stakpak provider: {}", e))?; - registry = registry.register("stakpak", provider); - } - ProviderConfig::GitHubCopilot { api_endpoint, .. } => { - if let Some(access_token) = provider_config.access_token() { - let mut copilot_config = CopilotConfig::new(access_token.to_string()); - if let Some(endpoint) = api_endpoint { - copilot_config = copilot_config.with_base_url(endpoint.clone()); - } - let provider = CopilotProvider::new(copilot_config) - .map_err(|e| format!("Failed to create GitHub Copilot provider: {}", e))?; - registry = registry.register("github-copilot", provider); - } - } - ProviderConfig::Custom { api_endpoint, .. } => { - // Custom providers are registered as OpenAI-compatible providers. - // The provider is registered with the config key (e.g., "litellm") as its ID. - // When a model like "litellm/anthropic/claude-opus" is used: - // 1. "litellm" is matched to this provider - // 2. "anthropic/claude-opus" is sent as the model name to the API - let key = provider_config.api_key().unwrap_or_default().to_string(); - let openai_config = - StakaiOpenAIConfig::new(key).with_base_url(api_endpoint.clone()); - - let provider = OpenAIProvider::new(openai_config) - .map_err(|e| format!("Failed to create custom provider '{}': {}", name, e))?; - - // Register with the config key as provider ID (e.g., "litellm", "ollama") - registry = registry.register(name, provider); - } - ProviderConfig::Bedrock { - region, - profile_name, - } => { - #[cfg(feature = "bedrock")] - { - use stakai::providers::bedrock::{BedrockConfig, BedrockProvider}; - let mut bedrock_config = BedrockConfig::new(region.clone()); - if let Some(profile) = profile_name { - bedrock_config = bedrock_config.with_profile_name(profile.clone()); - } - let provider = BedrockProvider::new(bedrock_config); - registry = registry.register("amazon-bedrock", provider); - } - #[cfg(not(feature = "bedrock"))] - { - let _ = (name, region, profile_name); - tracing::warn!( - "Bedrock provider configured but bedrock feature is not enabled" - ); - } - } - } - } - - Ok(registry) -} - -/// Get model string for StakAI -pub fn get_stakai_model_string(model: &Model) -> String { - model.id.clone() -} - -/// Wrapper around StakAI Inference for CLI usage -#[derive(Clone)] -pub struct StakAIClient { - inference: Inference, -} - -impl StakAIClient { - /// Create a new StakAI client from CLI provider config - pub fn new(config: &LLMProviderConfig) -> Result { - // Build registry with all providers including custom ones - let registry = build_provider_registry_direct(config) - .map_err(|e| AgentError::BadRequest(BadRequestErrorMessage::InvalidAgentInput(e)))?; - - let inference = Inference::builder() - .with_registry(registry) - .build() - .map_err(|e| { - AgentError::BadRequest(BadRequestErrorMessage::InvalidAgentInput(e.to_string())) - })?; - - Ok(Self { inference }) - } - - /// Create a new StakAI client with custom provider registry - pub fn with_registry(registry: ProviderRegistry) -> Result { - let inference = Inference::builder() - .with_registry(registry) - .build() - .map_err(|e| { - AgentError::BadRequest(BadRequestErrorMessage::InvalidAgentInput(e.to_string())) - })?; - - Ok(Self { inference }) - } - - /// Non-streaming chat completion - pub async fn chat(&self, input: LLMInput) -> Result { - let messages: Vec = input.messages.iter().map(to_stakai_message).collect(); - - let mut options = GenerateOptions::new().max_tokens(input.max_tokens); - - if let Some(tools) = &input.tools { - for tool in tools { - options = options.add_tool(to_stakai_tool(tool)); - } - } - - // Add custom headers if present - if let Some(headers) = &input.headers { - let mut stakai_headers = Headers::new(); - for (key, value) in headers { - stakai_headers.insert(key, value); - } - options = options.headers(stakai_headers); - } - - // Convert provider options if present - let provider_options = input - .provider_options - .as_ref() - .and_then(|opts| to_stakai_provider_options(opts, &input.model)); - let request = GenerateRequest { - model: input.model.clone(), - messages, - options, - provider_options, - telemetry_metadata: None, - }; - - let response = self.inference.generate(&request).await.map_err(|e| { - AgentError::BadRequest(BadRequestErrorMessage::InvalidAgentInput(e.to_string())) - })?; - - Ok(from_stakai_response(response, &input.model.id)) - } - - /// Streaming chat completion - pub async fn chat_stream( - &self, - input: LLMStreamInput, - ) -> Result { - let messages: Vec = input.messages.iter().map(to_stakai_message).collect(); - - let mut options = GenerateOptions::new().max_tokens(input.max_tokens); - - if let Some(tools) = &input.tools { - for tool in tools { - options = options.add_tool(to_stakai_tool(tool)); - } - } - - // Add custom headers if present - if let Some(headers) = &input.headers { - let mut stakai_headers = Headers::new(); - for (key, value) in headers { - stakai_headers.insert(key, value); - } - options = options.headers(stakai_headers); - } - - // Convert provider options if present - let provider_options = input - .provider_options - .as_ref() - .and_then(|opts| to_stakai_provider_options(opts, &input.model)); - let model_id = input.model.id.clone(); - let request = GenerateRequest { - model: input.model.clone(), - messages, - options, - provider_options, - telemetry_metadata: None, - }; - - let mut stream = self.inference.stream(&request).await.map_err(|e| { - AgentError::BadRequest(BadRequestErrorMessage::InvalidAgentInput(e.to_string())) - })?; - - let tx = input.stream_channel_tx; - let mut accumulated_text = String::new(); - let mut accumulated_tool_calls: Vec = Vec::new(); - let mut final_usage = LLMTokenUsage::default(); - let mut finish_reason = "stop".to_string(); - - while let Some(event_result) = stream.next().await { - match event_result { - Ok(event) => { - // Forward event to channel - if let Some(delta) = from_stakai_stream_event(&event) { - // Accumulate content for final response - match &delta { - GenerationDelta::Content { content } => { - accumulated_text.push_str(content); - } - GenerationDelta::ToolUse { tool_use } => { - if let Some(id) = &tool_use.id { - // Find existing tool call by id - let existing = accumulated_tool_calls.iter_mut().find(|tc| { - matches!(tc, LLMMessageTypedContent::ToolCall { id: tc_id, .. } if tc_id == id) - }); - - match existing { - Some(LLMMessageTypedContent::ToolCall { - args, - name: existing_name, - metadata: existing_metadata, - .. - }) => { - // Update existing tool call - // Update name if provided and current is empty - if let Some(new_name) = &tool_use.name - && existing_name.is_empty() - { - *existing_name = new_name.clone(); - } - // Append arguments if provided - if let Some(input) = &tool_use.input { - // Accumulate as string first, parse later - if let serde_json::Value::String(s) = args { - s.push_str(input); - } else { - *args = - serde_json::Value::String(input.clone()); - } - } - // Set metadata if provided (e.g. Gemini thought_signature from ToolCallEnd) - if tool_use.metadata.is_some() { - *existing_metadata = tool_use.metadata.clone(); - } - } - _ => { - // Create new tool call - let name = tool_use.name.clone().unwrap_or_default(); - let args = tool_use - .input - .clone() - .map(serde_json::Value::String) - .unwrap_or_else(|| { - serde_json::Value::String(String::new()) - }); - accumulated_tool_calls.push( - LLMMessageTypedContent::ToolCall { - id: id.clone(), - name, - args, - metadata: None, - }, - ); - } - } - } - } - GenerationDelta::Usage { usage } => { - final_usage = usage.clone(); - } - _ => {} - } - - // Send to channel (ignore errors if receiver dropped) - let _ = tx.send(delta).await; - } - - // Check for finish - if let StreamEvent::Finish { reason, usage } = event { - finish_reason = finish_reason_to_string(&reason); - final_usage = from_stakai_usage(&usage); - } - } - Err(e) => { - return Err(AgentError::BadRequest( - BadRequestErrorMessage::InvalidAgentInput(e.to_string()), - )); - } - } - } - - // Build final response - // Parse accumulated JSON string arguments into proper JSON values - let parsed_tool_calls: Vec = accumulated_tool_calls - .into_iter() - .map(|tc| { - if let LLMMessageTypedContent::ToolCall { - id, - name, - args, - metadata, - } = tc - { - let parsed_args = match args { - serde_json::Value::String(s) if !s.is_empty() => { - serde_json::from_str(&s).unwrap_or(serde_json::Value::String(s)) - } - other => other, - }; - LLMMessageTypedContent::ToolCall { - id, - name, - args: parsed_args, - metadata, - } - } else { - tc - } - }) - .collect(); - - let message_content = if parsed_tool_calls.is_empty() { - LLMMessageContent::String(accumulated_text) - } else { - let mut parts = vec![LLMMessageTypedContent::Text { - text: accumulated_text, - }]; - parts.extend(parsed_tool_calls); - LLMMessageContent::List(parts) - }; - - Ok(LLMCompletionResponse { - id: uuid::Uuid::new_v4().to_string(), - model: model_id, - object: "chat.completion".to_string(), - choices: vec![LLMChoice { - finish_reason: Some(finish_reason), - index: 0, - message: LLMMessage { - role: "assistant".to_string(), - content: message_content, - }, - }], - created: chrono::Utc::now().timestamp_millis() as u64, - usage: Some(final_usage), - }) - } - - /// Get the provider registry for model listing - pub fn registry(&self) -> &ProviderRegistry { - self.inference.registry() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - // ==================== Role Conversion Tests ==================== - - #[test] - fn test_role_conversion_user() { - let msg = LLMMessage { - role: "user".to_string(), - content: LLMMessageContent::String("Hello".to_string()), - }; - - let stakai_msg = to_stakai_message(&msg); - assert!(matches!(stakai_msg.role, Role::User)); - - let back = from_stakai_message(&stakai_msg); - assert_eq!(back.role, "user"); - } - - #[test] - fn test_role_conversion_assistant() { - let msg = LLMMessage { - role: "assistant".to_string(), - content: LLMMessageContent::String("Hi there!".to_string()), - }; - - let stakai_msg = to_stakai_message(&msg); - assert!(matches!(stakai_msg.role, Role::Assistant)); - - let back = from_stakai_message(&stakai_msg); - assert_eq!(back.role, "assistant"); - } - - #[test] - fn test_role_conversion_system() { - let msg = LLMMessage { - role: "system".to_string(), - content: LLMMessageContent::String("You are a helpful assistant.".to_string()), - }; - - let stakai_msg = to_stakai_message(&msg); - assert!(matches!(stakai_msg.role, Role::System)); - - let back = from_stakai_message(&stakai_msg); - assert_eq!(back.role, "system"); - } - - #[test] - fn test_role_conversion_tool() { - let msg = LLMMessage { - role: "tool".to_string(), - content: LLMMessageContent::String("Tool result".to_string()), - }; - - let stakai_msg = to_stakai_message(&msg); - assert!(matches!(stakai_msg.role, Role::Tool)); - - let back = from_stakai_message(&stakai_msg); - assert_eq!(back.role, "tool"); - } - - #[test] - fn test_role_conversion_unknown_defaults_to_user() { - let msg = LLMMessage { - role: "unknown_role".to_string(), - content: LLMMessageContent::String("Test".to_string()), - }; - - let stakai_msg = to_stakai_message(&msg); - assert!(matches!(stakai_msg.role, Role::User)); - } - - // ==================== Content Conversion Tests ==================== - - #[test] - fn test_string_content_conversion() { - let msg = LLMMessage { - role: "user".to_string(), - content: LLMMessageContent::String("Simple text message".to_string()), - }; - - let stakai_msg = to_stakai_message(&msg); - assert_eq!(stakai_msg.text(), Some("Simple text message".to_string())); - - let back = from_stakai_message(&stakai_msg); - if let LLMMessageContent::String(text) = back.content { - assert_eq!(text, "Simple text message"); - } else { - panic!("Expected String content"); - } - } - - #[test] - fn test_list_content_with_text() { - let msg = LLMMessage { - role: "assistant".to_string(), - content: LLMMessageContent::List(vec![LLMMessageTypedContent::Text { - text: "Hello world".to_string(), - }]), - }; - - let stakai_msg = to_stakai_message(&msg); - let parts = stakai_msg.parts(); - assert_eq!(parts.len(), 1); - - let back = from_stakai_message(&stakai_msg); - if let LLMMessageContent::List(parts) = back.content { - assert_eq!(parts.len(), 1); - assert!( - matches!(&parts[0], LLMMessageTypedContent::Text { text } if text == "Hello world") - ); - } else { - panic!("Expected List content"); - } - } - - #[test] - fn test_list_content_with_tool_call() { - let msg = LLMMessage { - role: "assistant".to_string(), - content: LLMMessageContent::List(vec![ - LLMMessageTypedContent::Text { - text: "Let me check the weather.".to_string(), - }, - LLMMessageTypedContent::ToolCall { - id: "call_abc123".to_string(), - name: "get_weather".to_string(), - args: serde_json::json!({"location": "New York", "unit": "celsius"}), - metadata: None, - }, - ]), - }; - - let stakai_msg = to_stakai_message(&msg); - let parts = stakai_msg.parts(); - assert_eq!(parts.len(), 2); - - let back = from_stakai_message(&stakai_msg); - if let LLMMessageContent::List(parts) = back.content { - assert_eq!(parts.len(), 2); - - // Check text part - assert!( - matches!(&parts[0], LLMMessageTypedContent::Text { text } if text == "Let me check the weather.") - ); - - // Check tool call part - if let LLMMessageTypedContent::ToolCall { id, name, args, .. } = &parts[1] { - assert_eq!(id, "call_abc123"); - assert_eq!(name, "get_weather"); - assert_eq!(args["location"], "New York"); - assert_eq!(args["unit"], "celsius"); - } else { - panic!("Expected ToolCall content"); - } - } else { - panic!("Expected List content"); - } - } - - #[test] - fn test_list_content_with_tool_result() { - let msg = LLMMessage { - role: "tool".to_string(), - content: LLMMessageContent::List(vec![LLMMessageTypedContent::ToolResult { - tool_use_id: "call_abc123".to_string(), - content: "Temperature: 22°C, Sunny".to_string(), - }]), - }; - - let stakai_msg = to_stakai_message(&msg); - let parts = stakai_msg.parts(); - assert_eq!(parts.len(), 1); - - let back = from_stakai_message(&stakai_msg); - if let LLMMessageContent::List(parts) = back.content { - assert_eq!(parts.len(), 1); - if let LLMMessageTypedContent::ToolResult { - tool_use_id, - content, - } = &parts[0] - { - assert_eq!(tool_use_id, "call_abc123"); - assert_eq!(content, "Temperature: 22°C, Sunny"); - } else { - panic!("Expected ToolResult content"); - } - } else { - panic!("Expected List content"); - } - } - - #[test] - fn test_image_content_conversion() { - let msg = LLMMessage { - role: "user".to_string(), - content: LLMMessageContent::List(vec![LLMMessageTypedContent::Image { - source: LLMMessageImageSource { - r#type: "base64".to_string(), - media_type: "image/png".to_string(), - data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==".to_string(), - }, - }]), - }; - - let stakai_msg = to_stakai_message(&msg); - let parts = stakai_msg.parts(); - assert_eq!(parts.len(), 1); - - // Verify it's converted to a data URI - if let ContentPart::Image { url, .. } = &parts[0] { - assert!(url.starts_with("data:image/png;base64,")); - } else { - panic!("Expected Image content part"); - } - - let back = from_stakai_message(&stakai_msg); - if let LLMMessageContent::List(parts) = back.content { - if let LLMMessageTypedContent::Image { source } = &parts[0] { - assert_eq!(source.media_type, "image/png"); - assert!(!source.data.is_empty()); - } else { - panic!("Expected Image content"); - } - } else { - panic!("Expected List content"); - } - } - - // ==================== Tool Conversion Tests ==================== - - #[test] - fn test_tool_conversion_basic() { - let tool = LLMTool { - name: "get_weather".to_string(), - description: "Get weather for a location".to_string(), - input_schema: serde_json::json!({ - "type": "object", - "properties": { - "location": {"type": "string"} - } - }), - }; - - let stakai_tool = to_stakai_tool(&tool); - assert_eq!(stakai_tool.tool_type, "function"); - assert_eq!(stakai_tool.function.name, "get_weather"); - assert_eq!( - stakai_tool.function.description, - "Get weather for a location" - ); - assert_eq!(stakai_tool.function.parameters["type"], "object"); - - let back = from_stakai_tool(&stakai_tool); - assert_eq!(back.name, "get_weather"); - assert_eq!(back.description, "Get weather for a location"); - assert_eq!(back.input_schema["type"], "object"); - } - - #[test] - fn test_tool_conversion_complex_schema() { - let tool = LLMTool { - name: "search_database".to_string(), - description: "Search a database with filters".to_string(), - input_schema: serde_json::json!({ - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query" - }, - "filters": { - "type": "array", - "items": { - "type": "object", - "properties": { - "field": {"type": "string"}, - "value": {"type": "string"} - } - } - }, - "limit": { - "type": "integer", - "default": 10 - } - }, - "required": ["query"] - }), - }; - - let stakai_tool = to_stakai_tool(&tool); - let back = from_stakai_tool(&stakai_tool); - - assert_eq!(back.name, "search_database"); - assert_eq!(back.input_schema["properties"]["query"]["type"], "string"); - assert_eq!(back.input_schema["properties"]["filters"]["type"], "array"); - assert_eq!(back.input_schema["required"][0], "query"); - } - - // ==================== Stream Event Conversion Tests ==================== - - #[test] - fn test_stream_event_text_delta() { - let event = StreamEvent::TextDelta { - id: "gen_123".to_string(), - delta: "Hello ".to_string(), - }; - - let delta = from_stakai_stream_event(&event); - assert!(delta.is_some()); - - if let Some(GenerationDelta::Content { content }) = delta { - assert_eq!(content, "Hello "); - } else { - panic!("Expected Content delta"); - } - } - - #[test] - fn test_stream_event_reasoning_delta() { - let event = StreamEvent::ReasoningDelta { - id: "gen_123".to_string(), - delta: "Let me think about this...".to_string(), - }; - - let delta = from_stakai_stream_event(&event); - assert!(delta.is_some()); - - if let Some(GenerationDelta::Thinking { thinking }) = delta { - assert_eq!(thinking, "Let me think about this..."); - } else { - panic!("Expected Thinking delta"); - } - } - - #[test] - fn test_stream_event_tool_call_start() { - let event = StreamEvent::ToolCallStart { - id: "call_xyz".to_string(), - name: "run_command".to_string(), - }; - - let delta = from_stakai_stream_event(&event); - assert!(delta.is_some()); - - if let Some(GenerationDelta::ToolUse { tool_use }) = delta { - assert_eq!(tool_use.id, Some("call_xyz".to_string())); - assert_eq!(tool_use.name, Some("run_command".to_string())); - assert!(tool_use.input.is_none()); - } else { - panic!("Expected ToolUse delta"); - } - } - - #[test] - fn test_stream_event_tool_call_delta() { - let event = StreamEvent::ToolCallDelta { - id: "call_xyz".to_string(), - delta: r#"{"command": "ls"#.to_string(), - }; - - let delta = from_stakai_stream_event(&event); - assert!(delta.is_some()); - - if let Some(GenerationDelta::ToolUse { tool_use }) = delta { - assert_eq!(tool_use.id, Some("call_xyz".to_string())); - assert!(tool_use.name.is_none()); - assert_eq!(tool_use.input, Some(r#"{"command": "ls"#.to_string())); - } else { - panic!("Expected ToolUse delta"); - } - } - - #[test] - fn test_stream_event_tool_call_end() { - let event = StreamEvent::ToolCallEnd { - id: "call_xyz".to_string(), - name: "run_command".to_string(), - arguments: serde_json::json!({"command": "ls -la"}), - metadata: None, - }; - - let delta = from_stakai_stream_event(&event); - assert!(delta.is_some()); - - if let Some(GenerationDelta::ToolUse { tool_use }) = delta { - assert_eq!(tool_use.id, Some("call_xyz".to_string())); - // ToolCallEnd emits name (for providers like Gemini) but NOT input - // to avoid doubling arguments that were already accumulated via ToolCallDelta - assert_eq!(tool_use.name, Some("run_command".to_string())); - assert!(tool_use.input.is_none()); - } else { - panic!("Expected ToolUse delta"); - } - } - - #[test] - fn test_stream_event_finish() { - let event = StreamEvent::Finish { - usage: Usage::new(100, 50), - reason: FinishReason::stop(), - }; - - let delta = from_stakai_stream_event(&event); - assert!(delta.is_some()); - - if let Some(GenerationDelta::Usage { usage }) = delta { - assert_eq!(usage.prompt_tokens, 100); - assert_eq!(usage.completion_tokens, 50); - assert_eq!(usage.total_tokens, 150); - } else { - panic!("Expected Usage delta"); - } - } - - #[test] - fn test_stream_event_start_returns_none() { - let event = StreamEvent::Start { - id: "gen_123".to_string(), - }; - - let delta = from_stakai_stream_event(&event); - assert!(delta.is_none()); - } - - #[test] - fn test_stream_event_error_returns_none() { - let event = StreamEvent::Error { - message: "Something went wrong".to_string(), - }; - - let delta = from_stakai_stream_event(&event); - assert!(delta.is_none()); - } - - // ==================== Usage Conversion Tests ==================== - - #[test] - fn test_usage_conversion() { - let usage = Usage::new(500, 200); - - let llm_usage = from_stakai_usage(&usage); - assert_eq!(llm_usage.prompt_tokens, 500); - assert_eq!(llm_usage.completion_tokens, 200); - assert_eq!(llm_usage.total_tokens, 700); - assert!(llm_usage.prompt_tokens_details.is_none()); - } - - // ==================== Finish Reason Tests ==================== - - #[test] - fn test_finish_reason_conversion() { - assert_eq!(finish_reason_to_string(&FinishReason::stop()), "stop"); - assert_eq!(finish_reason_to_string(&FinishReason::length()), "length"); - assert_eq!( - finish_reason_to_string(&FinishReason::content_filter()), - "content_filter" - ); - assert_eq!( - finish_reason_to_string(&FinishReason::tool_calls()), - "tool_calls" - ); - assert_eq!(finish_reason_to_string(&FinishReason::other()), "other"); - assert_eq!(finish_reason_to_string(&FinishReason::error()), "error"); - - // Test with raw values - should return the raw value - use stakai::FinishReasonKind; - let reason = FinishReason::with_raw(FinishReasonKind::Stop, "end_turn"); - assert_eq!(finish_reason_to_string(&reason), "end_turn"); - } - - // ==================== Model String Tests ==================== - - #[test] - fn test_model_string_anthropic() { - let model = Model::custom("claude-sonnet-4-5-20250929", "anthropic"); - let model_str = get_stakai_model_string(&model); - assert_eq!(model_str, "claude-sonnet-4-5-20250929"); - } - - #[test] - fn test_model_string_openai() { - let model = Model::custom("gpt-5", "openai"); - let model_str = get_stakai_model_string(&model); - assert_eq!(model_str, "gpt-5"); - } - - #[test] - fn test_model_string_gemini() { - let model = Model::custom("gemini-2.5-flash", "google"); - let model_str = get_stakai_model_string(&model); - assert_eq!(model_str, "gemini-2.5-flash"); - } - - #[test] - fn test_model_string_custom() { - let model = Model::custom("claude-opus-4-5", "litellm"); - let model_str = get_stakai_model_string(&model); - assert_eq!(model_str, "claude-opus-4-5"); - } - - // ==================== Response Conversion Tests ==================== - - #[test] - fn test_response_conversion_text_only() { - let response = GenerateResponse { - content: vec![stakai::ResponseContent::Text { - text: "Hello, how can I help?".to_string(), - }], - usage: Usage::new(10, 5), - finish_reason: FinishReason::stop(), - metadata: None, - warnings: None, - }; - - let llm_response = from_stakai_response(response, "gpt-4"); - - assert_eq!(llm_response.model, "gpt-4"); - assert_eq!(llm_response.object, "chat.completion"); - assert_eq!(llm_response.choices.len(), 1); - assert_eq!( - llm_response.choices[0].finish_reason, - Some("stop".to_string()) - ); - assert_eq!(llm_response.choices[0].message.role, "assistant"); - - if let LLMMessageContent::String(text) = &llm_response.choices[0].message.content { - assert_eq!(text, "Hello, how can I help?"); - } else { - panic!("Expected String content"); - } - - let usage = llm_response.usage.unwrap(); - assert_eq!(usage.prompt_tokens, 10); - assert_eq!(usage.completion_tokens, 5); - } - - #[test] - fn test_response_conversion_with_tool_calls() { - let response = GenerateResponse { - content: vec![ - stakai::ResponseContent::Text { - text: "I'll check the weather for you.".to_string(), - }, - stakai::ResponseContent::ToolCall(stakai::ToolCall { - id: "call_123".to_string(), - name: "get_weather".to_string(), - arguments: serde_json::json!({"location": "NYC"}), - metadata: None, - }), - ], - usage: Usage::new(20, 15), - finish_reason: FinishReason::tool_calls(), - metadata: None, - warnings: None, - }; - - let llm_response = from_stakai_response(response, "claude-3"); - - assert_eq!( - llm_response.choices[0].finish_reason, - Some("tool_calls".to_string()) - ); - - if let LLMMessageContent::List(parts) = &llm_response.choices[0].message.content { - assert_eq!(parts.len(), 2); - - // Check text part - assert!( - matches!(&parts[0], LLMMessageTypedContent::Text { text } if text == "I'll check the weather for you.") - ); - - // Check tool call part - if let LLMMessageTypedContent::ToolCall { id, name, args, .. } = &parts[1] { - assert_eq!(id, "call_123"); - assert_eq!(name, "get_weather"); - assert_eq!(args["location"], "NYC"); - } else { - panic!("Expected ToolCall"); - } - } else { - panic!("Expected List content"); - } - } - - // ==================== Provider Options Conversion Tests ==================== - - #[test] - fn test_provider_options_anthropic_thinking() { - use crate::models::llm::{LLMAnthropicOptions, LLMProviderOptions, LLMThinkingOptions}; - - let opts = LLMProviderOptions { - anthropic: Some(LLMAnthropicOptions { - thinking: Some(LLMThinkingOptions::new(8000)), - }), - openai: None, - google: None, - }; - - let model = Model::custom("claude-sonnet-4-5-20250929", "anthropic"); - let result = to_stakai_provider_options(&opts, &model); - - assert!(result.is_some()); - if let Some(ProviderOptions::Anthropic(anthropic)) = result { - assert!(anthropic.thinking.is_some()); - assert_eq!(anthropic.thinking.unwrap().budget_tokens, 8000); - } else { - panic!("Expected Anthropic provider options"); - } - } - - #[test] - fn test_provider_options_openai_reasoning() { - use crate::models::llm::{LLMOpenAIOptions, LLMProviderOptions}; - - let opts = LLMProviderOptions { - anthropic: None, - openai: Some(LLMOpenAIOptions { - reasoning_effort: Some("high".to_string()), - }), - google: None, - }; - - let model = Model::custom("gpt-5", "openai"); - let result = to_stakai_provider_options(&opts, &model); - - assert!(result.is_some()); - if let Some(ProviderOptions::OpenAI(openai)) = result { - if let Some(OpenAIApiConfig::Responses(config)) = openai.api_config { - assert_eq!(config.reasoning_effort, Some(ReasoningEffort::High)); - } else { - panic!("Expected Responses API config"); - } - } else { - panic!("Expected OpenAI provider options"); - } - } - - #[test] - fn test_provider_options_openai_none_when_empty() { - use crate::models::llm::LLMProviderOptions; - - let opts = LLMProviderOptions::default(); - let model = Model::custom("gpt-4.1-mini", "openai"); - let result = to_stakai_provider_options(&opts, &model); - - assert!(result.is_none()); - } - - #[test] - fn test_provider_options_custom_none_when_empty() { - use crate::models::llm::LLMProviderOptions; - - let opts = LLMProviderOptions::default(); - let model = Model::custom("llama3.2", "ollama"); - let result = to_stakai_provider_options(&opts, &model); - - assert!(result.is_none()); - } - - #[test] - fn test_provider_options_google_thinking() { - use crate::models::llm::{LLMGoogleOptions, LLMProviderOptions}; - - let opts = LLMProviderOptions { - anthropic: None, - openai: None, - google: Some(LLMGoogleOptions { - thinking_budget: Some(5000), - }), - }; - - let model = Model::custom("gemini-2.5-flash", "google"); - let result = to_stakai_provider_options(&opts, &model); - - assert!(result.is_some()); - if let Some(ProviderOptions::Google(google)) = result { - assert_eq!(google.thinking_budget, Some(5000)); - } else { - panic!("Expected Google provider options"); - } - } - - #[test] - fn test_provider_options_none_when_empty() { - use crate::models::llm::LLMProviderOptions; - - let opts = LLMProviderOptions::default(); - - let model = Model::custom("claude-sonnet-4-5-20250929", "anthropic"); - let result = to_stakai_provider_options(&opts, &model); - - assert!(result.is_none()); - } - - // ==================== Config Building Tests ==================== - - #[test] - fn test_build_inference_config_empty() { - let config = LLMProviderConfig::new(); - - let result = build_inference_config(&config); - assert!(result.is_ok()); - } - - #[test] - fn test_build_inference_config_with_openai() { - use crate::models::llm::ProviderConfig; - - let mut config = LLMProviderConfig::new(); - config.add_provider( - "openai", - ProviderConfig::OpenAI { - api_key: Some("sk-test-key".to_string()), - api_endpoint: Some("https://api.openai.com/v1".to_string()), - auth: None, - }, - ); - - let result = build_inference_config(&config); - assert!(result.is_ok()); - } - - #[test] - fn test_build_inference_config_with_anthropic() { - use crate::models::llm::ProviderConfig; - - let mut config = LLMProviderConfig::new(); - config.add_provider( - "anthropic", - ProviderConfig::Anthropic { - api_key: Some("sk-ant-test".to_string()), - api_endpoint: None, - access_token: None, - auth: None, - }, - ); - - let result = build_inference_config(&config); - assert!(result.is_ok()); - } - - #[test] - fn test_build_inference_config_with_gemini() { - use crate::models::llm::ProviderConfig; - - let mut config = LLMProviderConfig::new(); - config.add_provider( - "gemini", - ProviderConfig::Gemini { - api_key: Some("gemini-test-key".to_string()), - api_endpoint: None, - auth: None, - }, - ); - - let result = build_inference_config(&config); - assert!(result.is_ok()); - } - - #[test] - fn test_build_inference_config_all_providers() { - use crate::models::llm::ProviderConfig; - - let mut config = LLMProviderConfig::new(); - config.add_provider( - "anthropic", - ProviderConfig::Anthropic { - api_key: Some("sk-ant-test".to_string()), - api_endpoint: None, - access_token: None, - auth: None, - }, - ); - config.add_provider( - "openai", - ProviderConfig::OpenAI { - api_key: Some("sk-openai-test".to_string()), - api_endpoint: None, - auth: None, - }, - ); - config.add_provider( - "gemini", - ProviderConfig::Gemini { - api_key: Some("gemini-test".to_string()), - api_endpoint: None, - auth: None, - }, - ); - - let result = build_inference_config(&config); - assert!(result.is_ok()); - } - - #[test] - fn test_build_provider_registry_with_custom_providers() { - use crate::models::llm::ProviderConfig; - - let mut config = LLMProviderConfig::new(); - config.add_provider( - "litellm", - ProviderConfig::Custom { - api_endpoint: "http://localhost:4000".to_string(), - api_key: Some("sk-1234".to_string()), - auth: None, - }, - ); - config.add_provider( - "ollama", - ProviderConfig::Custom { - api_endpoint: "http://localhost:11434/v1".to_string(), - api_key: None, - auth: None, - }, - ); - - let result = build_provider_registry_direct(&config); - assert!( - result.is_ok(), - "Failed to build registry: {:?}", - result.err() - ); - - let registry = result.unwrap(); - assert!(registry.has_provider("litellm")); - assert!(registry.has_provider("ollama")); - } - - #[test] - fn test_build_provider_registry_registers_openai_from_oauth_auth() { - use crate::models::auth::ProviderAuth; - use crate::models::llm::ProviderConfig; - use base64::Engine; - - let payload = serde_json::json!({ - "https://api.openai.com/auth": { - "chatgpt_account_id": "acct_test_789" - } - }); - let encoded_payload = - base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(payload.to_string().as_bytes()); - let access_token = format!("header.{}.signature", encoded_payload); - - let mut config = LLMProviderConfig::new(); - config.add_provider( - "openai", - ProviderConfig::OpenAI { - api_key: None, - api_endpoint: None, - auth: Some(ProviderAuth::oauth_with_name( - access_token, - "refresh-token", - i64::MAX, - "ChatGPT Plus/Pro", - )), - }, - ); - - let registry = build_provider_registry_direct(&config).expect("registry should build"); - - assert!(registry.has_provider("openai")); - } - - // ==================== Round-trip Tests ==================== - - #[test] - fn test_message_roundtrip_simple() { - let original = LLMMessage { - role: "user".to_string(), - content: LLMMessageContent::String("What is 2+2?".to_string()), - }; - - let stakai_msg = to_stakai_message(&original); - let back = from_stakai_message(&stakai_msg); - - assert_eq!(back.role, original.role); - if let (LLMMessageContent::String(orig), LLMMessageContent::String(converted)) = - (&original.content, &back.content) - { - assert_eq!(orig, converted); - } else { - panic!("Content type mismatch"); - } - } - - #[test] - fn test_message_roundtrip_complex() { - let original = LLMMessage { - role: "assistant".to_string(), - content: LLMMessageContent::List(vec![ - LLMMessageTypedContent::Text { - text: "Here's the result:".to_string(), - }, - LLMMessageTypedContent::ToolCall { - id: "call_001".to_string(), - name: "calculator".to_string(), - args: serde_json::json!({"expression": "2+2"}), - metadata: None, - }, - ]), - }; - - let stakai_msg = to_stakai_message(&original); - let back = from_stakai_message(&stakai_msg); - - assert_eq!(back.role, original.role); - if let LLMMessageContent::List(parts) = back.content { - assert_eq!(parts.len(), 2); - } else { - panic!("Expected List content"); - } - } - - #[test] - fn test_tool_roundtrip() { - let original = LLMTool { - name: "file_reader".to_string(), - description: "Read contents of a file".to_string(), - input_schema: serde_json::json!({ - "type": "object", - "properties": { - "path": { - "type": "string", - "description": "File path to read" - } - }, - "required": ["path"] - }), - }; - - let stakai_tool = to_stakai_tool(&original); - let back = from_stakai_tool(&stakai_tool); - - assert_eq!(back.name, original.name); - assert_eq!(back.description, original.description); - assert_eq!(back.input_schema, original.input_schema); - } -} diff --git a/libs/shared/src/remote_connection.rs b/libs/shared/src/remote_connection.rs index 400fe1959..79e66b771 100644 --- a/libs/shared/src/remote_connection.rs +++ b/libs/shared/src/remote_connection.rs @@ -507,18 +507,28 @@ impl RemoteConnection { && options.with_progress && !text.trim().is_empty() { - let _ = ctx.peer.notify_progress(rmcp::model::ProgressNotificationParam { - progress_token: rmcp::model::ProgressToken(rmcp::model::NumberOrString::Number(0)), + let _ = ctx + .peer + .notify_progress(rmcp::model::ProgressNotificationParam { + progress_token: rmcp::model::ProgressToken( + rmcp::model::NumberOrString::Number(0), + ), progress: 50.0, total: Some(100.0), - message: Some(serde_json::to_string(&crate::models::integrations::openai::ToolCallResultProgress { - id: progress_id, - message: text, - progress_type: None, - task_updates: None, - progress: None, - }).unwrap_or_default()), - }).await; + message: Some( + serde_json::to_string( + &crate::models::agent_runtime::ToolCallResultProgress { + id: progress_id, + message: text, + progress_type: None, + task_updates: None, + progress: None, + }, + ) + .unwrap_or_default(), + ), + }) + .await; } } russh::ChannelMsg::ExtendedData { data, ext: _ } => { @@ -533,18 +543,28 @@ impl RemoteConnection { && options.with_progress && !text.trim().is_empty() { - let _ = ctx.peer.notify_progress(rmcp::model::ProgressNotificationParam { - progress_token: rmcp::model::ProgressToken(rmcp::model::NumberOrString::Number(0)), + let _ = ctx + .peer + .notify_progress(rmcp::model::ProgressNotificationParam { + progress_token: rmcp::model::ProgressToken( + rmcp::model::NumberOrString::Number(0), + ), progress: 50.0, total: Some(100.0), - message: Some(serde_json::to_string(&crate::models::integrations::openai::ToolCallResultProgress { - id: progress_id, - message: text, - progress_type: None, - task_updates: None, - progress: None, - }).unwrap_or_default()), - }).await; + message: Some( + serde_json::to_string( + &crate::models::agent_runtime::ToolCallResultProgress { + id: progress_id, + message: text, + progress_type: None, + task_updates: None, + progress: None, + }, + ) + .unwrap_or_default(), + ), + }) + .await; } } russh::ChannelMsg::ExitStatus { exit_status } => { diff --git a/tui/src/app/events.rs b/tui/src/app/events.rs index bc5440b51..94ba11a39 100644 --- a/tui/src/app/events.rs +++ b/tui/src/app/events.rs @@ -2,7 +2,7 @@ use ratatui::style::Color; use stakai::Model; use stakpak_api::models::ListRuleBook; use stakpak_shared::models::{ - integrations::openai::{ToolCall, ToolCallResult, ToolCallResultProgress, ToolCallStreamInfo}, + agent_runtime::{ToolCallResult, ToolCallResultProgress, ToolCallStreamInfo}, llm::LLMTokenUsage, }; use uuid::Uuid; @@ -16,7 +16,7 @@ pub enum InputEvent { AssistantMessage(String), AddUserMessage(String), StreamAssistantMessage(Uuid, String), - RunToolCall(ToolCall), + RunToolCall(stakai::ToolCall), ToolResult(ToolCallResult), StreamToolResult(ToolCallResultProgress), /// Progress update while tool calls are being streamed/generated by the LLM @@ -38,7 +38,7 @@ pub enum InputEvent { InputSubmitted, InputSubmittedWith(String), InputSubmittedWithColor(String, Color), - MessageToolCalls(Vec), + MessageToolCalls(Vec), ScrollUp, ScrollDown, PageUp, @@ -54,7 +54,7 @@ pub enum InputEvent { CursorRight, ToggleCursorVisible, Resized(u16, u16), - ShowConfirmationDialog(ToolCall), + ShowConfirmationDialog(stakai::ToolCall), HasUserMessage, Tab, ToggleApprovalStatus, @@ -203,8 +203,8 @@ pub enum InputEvent { // Ask User popup events ShowAskUserPopup( - ToolCall, - Vec, + stakai::ToolCall, + Vec, ), AskUserNextTab, AskUserPrevTab, @@ -286,15 +286,15 @@ pub enum OutputEvent { UserMessage( String, Option>, - Vec, + Vec, Option, // revert_to_user_message_index ), - AcceptTool(ToolCall), - RejectTool(ToolCall, bool), + AcceptTool(stakai::ToolCall), + RejectTool(stakai::ToolCall, bool), ListSessions, SwitchToSession(String), NewSession, - SendToolResult(ToolCallResult, bool, Vec), + SendToolResult(ToolCallResult, bool, Vec), ResumeSession, RequestProfileSwitch(String), RequestRulebookUpdate(Vec), diff --git a/tui/src/app/types.rs b/tui/src/app/types.rs index 1eb1fb14e..6b10986b0 100644 --- a/tui/src/app/types.rs +++ b/tui/src/app/types.rs @@ -13,11 +13,9 @@ use crate::services::shell_mode::ShellCommand; use crate::services::text_selection::SelectionState; use crate::services::textarea::{TextArea, TextAreaState}; use ratatui::text::Line; -use stakai::Model; +use stakai::{ContentPart, Model, ToolCall}; use stakpak_api::models::ListRuleBook; -use stakpak_shared::models::integrations::openai::{ - ContentPart, TaskPauseInfo, ToolCall, ToolCallResult, -}; +use stakpak_shared::models::agent_runtime::{TaskPauseInfo, ToolCallResult}; use stakpak_shared::models::llm::LLMTokenUsage; use stakpak_shared::secret_manager::SecretManager; use std::collections::{HashMap, HashSet, VecDeque}; @@ -152,7 +150,7 @@ pub struct AttachedImage { pub end_pos: usize, } -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub struct PendingUserMessage { pub final_input: String, pub shell_tool_calls: Option>, @@ -684,9 +682,9 @@ pub struct AskUserState { /// Whether the ask user interaction is active pub is_visible: bool, /// Questions to display in the inline block - pub questions: Vec, + pub questions: Vec, /// User's answers (question label -> answer) - pub answers: HashMap, + pub answers: HashMap, /// Currently selected tab index (question index, or questions.len() for Submit) pub current_tab: usize, /// Currently selected option index within the current question @@ -920,19 +918,15 @@ impl LoadingStateManager { #[cfg(test)] mod tests { use super::*; - use stakpak_shared::models::integrations::openai::{ - FunctionCall, ToolCall, ToolCallResultStatus, - }; + use stakai::ToolCall; + use stakpak_shared::models::agent_runtime::ToolCallResultStatus; fn tool_result(id: &str) -> ToolCallResult { ToolCallResult { call: ToolCall { id: id.to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "run_command".to_string(), - arguments: "{}".to_string(), - }, + name: "run_command".to_string(), + arguments: serde_json::json!({}), metadata: None, }, result: format!("result-{id}"), @@ -945,22 +939,14 @@ mod tests { let mut first = PendingUserMessage::new( "first".to_string(), Some(vec![tool_result("t1")]), - vec![ContentPart { - r#type: "text".to_string(), - text: Some("img-1".to_string()), - image_url: None, - }], + vec![ContentPart::text("img-1")], "first".to_string(), ); let second = PendingUserMessage::new( "second".to_string(), Some(vec![tool_result("t2")]), - vec![ContentPart { - r#type: "text".to_string(), - text: Some("img-2".to_string()), - image_url: None, - }], + vec![ContentPart::text("img-2")], "second".to_string(), ); @@ -986,11 +972,7 @@ mod tests { let second = PendingUserMessage::new( "second".to_string(), None, - vec![ContentPart { - r#type: "text".to_string(), - text: Some("img-2".to_string()), - image_url: None, - }], + vec![ContentPart::text("img-2")], "second".to_string(), ); diff --git a/tui/src/event_loop.rs b/tui/src/event_loop.rs index 27114f461..631a77e98 100644 --- a/tui/src/event_loop.rs +++ b/tui/src/event_loop.rs @@ -16,7 +16,7 @@ use crossterm::event::{ use crossterm::{execute, terminal::EnterAlternateScreen}; use ratatui::{Terminal, backend::CrosstermBackend}; use stakai::Model; -use stakpak_shared::models::integrations::openai::ToolCallResultStatus; +use stakpak_shared::models::agent_runtime::ToolCallResultStatus; use stakpak_shared::utils::strip_tool_name; use std::io; use std::sync::Arc; @@ -270,7 +270,7 @@ pub async fn run_tui( state.session_tool_calls_state.session_tool_calls_queue.insert(tool_call_result.call.id.clone(), ToolCallStatus::Executed); update_session_tool_calls_queue(&mut state, tool_call_result); - let tool_name = strip_tool_name(&tool_call_result.call.function.name); + let tool_name = strip_tool_name(&tool_call_result.call.name); let is_fg_cmd = matches!(tool_name, "run_command" | "run_remote_command"); if tool_call_result.status == ToolCallResultStatus::Cancelled && is_fg_cmd { diff --git a/tui/src/services/approval_bar.rs b/tui/src/services/approval_bar.rs index b4e1f3d20..06ce2fe9c 100644 --- a/tui/src/services/approval_bar.rs +++ b/tui/src/services/approval_bar.rs @@ -21,7 +21,7 @@ use ratatui::layout::Rect; use ratatui::style::{Modifier, Style}; use ratatui::text::{Line, Span}; use ratatui::widgets::{Clear, Paragraph}; -use stakpak_shared::models::integrations::openai::ToolCall; +use stakai::ToolCall; use stakpak_shared::utils::strip_tool_name; /// Approval status for a tool call @@ -42,12 +42,11 @@ pub struct ApprovalAction { impl ApprovalAction { pub fn new(tool_call: ToolCall) -> Self { - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); let label = if tool_name == "dynamic_subagent_task" { // Parse description and sandbox flag from arguments for a meaningful label - let args = - serde_json::from_str::(&tool_call.function.arguments).ok(); + let args = Some(&tool_call.arguments); let desc = args .as_ref() .and_then(|a| a.get("description").and_then(|v| v.as_str())) @@ -578,16 +577,12 @@ impl ApprovalBar { #[cfg(test)] mod tests { use super::*; - use stakpak_shared::models::integrations::openai::FunctionCall; - fn make_tool_call(name: &str, args: &str) -> ToolCall { ToolCall { id: format!("call_{}", name), - r#type: "function".to_string(), - function: FunctionCall { - name: name.to_string(), - arguments: args.to_string(), - }, + name: name.to_string(), + arguments: serde_json::from_str(args) + .unwrap_or_else(|_| serde_json::Value::String(args.to_string())), metadata: None, } } diff --git a/tui/src/services/auto_approve.rs b/tui/src/services/auto_approve.rs index 1488f4259..72daee824 100644 --- a/tui/src/services/auto_approve.rs +++ b/tui/src/services/auto_approve.rs @@ -1,7 +1,7 @@ use crate::app::InputEvent; use crate::constants::AUTO_APPROVE_CONFIG_PATH; use serde::{Deserialize, Serialize}; -use stakpak_shared::models::integrations::openai::ToolCall; +use stakai::ToolCall; use stakpak_shared::utils::{backward_compatibility_mapping, strip_tool_name}; use std::collections::HashMap; use std::fs; @@ -253,7 +253,7 @@ impl AutoApproveManager { } pub fn get_policy_for_tool(&self, tool_call: &ToolCall) -> AutoApprovePolicy { - let binding = tool_call.function.name.clone(); + let binding = tool_call.name.clone(); let tool_name = strip_tool_name(&binding); // For shell commands, resolve hierarchical scope keys @@ -469,9 +469,8 @@ fn resolve_shell_scope( rules: &HashMap, default: &AutoApprovePolicy, ) -> Option { - let args: serde_json::Value = serde_json::from_str(&tool_call.function.arguments).ok()?; - let command_str = args.get("command")?.as_str()?; - let tool_name = strip_tool_name(&tool_call.function.name); + let command_str = tool_call.arguments.get("command")?.as_str()?; + let tool_name = strip_tool_name(&tool_call.name); let fallback_scopes = if SHELL_TOOLS.contains(&tool_name) && tool_name != BASE_SHELL_TOOL { vec![BASE_SHELL_TOOL] } else { @@ -654,16 +653,13 @@ mod tests { // --- Tests for hierarchical shell scope resolution --- - use stakpak_shared::models::integrations::openai::{FunctionCall, ToolCall}; + use stakai::ToolCall; fn make_tool_call(tool_name: &str, command: &str) -> ToolCall { ToolCall { id: "tc-1".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: tool_name.to_string(), - arguments: serde_json::json!({"command": command}).to_string(), - }, + name: tool_name.to_string(), + arguments: serde_json::json!({"command": command}), metadata: None, } } @@ -736,11 +732,8 @@ mod tests { let rules = HashMap::new(); let tc = ToolCall { id: "tc-2".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "run_command".to_string(), - arguments: serde_json::json!({"command": ""}).to_string(), - }, + name: "run_command".to_string(), + arguments: serde_json::json!({"command": ""}), metadata: None, }; let result = resolve_shell_scope(&tc, &rules, &AutoApprovePolicy::Prompt); diff --git a/tui/src/services/bash_block.rs b/tui/src/services/bash_block.rs index 4eadb4756..1e3ec606c 100644 --- a/tui/src/services/bash_block.rs +++ b/tui/src/services/bash_block.rs @@ -8,8 +8,9 @@ use console::strip_ansi_codes; use ratatui::style::{Color, Modifier, Style}; use ratatui::text::{Line, Span}; use regex::Regex; -use stakpak_shared::models::integrations::openai::{ - ToolCall, ToolCallResult, ToolCallResultStatus, ToolCallStreamInfo, +use stakai::ToolCall; +use stakpak_shared::models::agent_runtime::{ + ToolCallResult, ToolCallResultStatus, ToolCallStreamInfo, }; use stakpak_shared::utils::strip_tool_name; use std::sync::OnceLock; @@ -587,7 +588,7 @@ pub fn extract_bash_block_info( }; let outside_title = get_command_type_name(tool_call); let bubble_title = extract_command_purpose(&command, &outside_title); - let colors = match strip_tool_name(&tool_call.function.name) { + let colors = match strip_tool_name(&tool_call.name) { "create_file" => BubbleColors { border_color: ThemeColors::success(), title_color: tool_muted_color(), @@ -625,11 +626,11 @@ pub fn extract_bash_block_info( tool_type: "Delete File".to_string(), }, "dynamic_subagent_task" => { - let is_sandbox = - serde_json::from_str::(&tool_call.function.arguments) - .ok() - .and_then(|a| a.get("enable_sandbox").and_then(|v| v.as_bool())) - .unwrap_or(false); + let is_sandbox = tool_call + .arguments + .get("enable_sandbox") + .and_then(|v| v.as_bool()) + .unwrap_or(false); if is_sandbox { BubbleColors { border_color: ThemeColors::success(), @@ -990,7 +991,7 @@ pub fn render_file_diff_full( } pub fn render_file_diff(tool_call: &ToolCall, terminal_width: usize) -> Vec> { - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if tool_name == "str_replace" || tool_name == "create" { // Use full diff (not truncated) for pending approval blocks let (_, mut diff_lines) = render_file_diff_block(tool_call, terminal_width); @@ -1047,7 +1048,7 @@ pub fn render_bash_block( &bubble_title, Some(colors.clone()), terminal_width, - strip_tool_name(&tool_call.function.name), + strip_tool_name(&tool_call.name), None, ) } @@ -1120,7 +1121,7 @@ pub fn render_result_block(tool_call_result: &ToolCallResult, width: usize) -> V // Handle str_replace/create with diff-only content // If render_diff_result_block returns None (no diff), fall through to standard rendering - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if tool_name == "str_replace" || tool_name == "create" { // Check for rejected/cancelled in result text if result.contains("TOOL_CALL_REJECTED") { @@ -1723,7 +1724,7 @@ pub fn render_refreshed_terminal_bubble( } pub fn is_collapsed_tool_call(tool_call: &ToolCall) -> bool { - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if matches!(tool_name, "run_command_task" | "run_remote_command_task") { return false; } @@ -2323,10 +2324,10 @@ pub fn render_run_command_block( /// Render an ask_user tool block inline, similar to render_run_command_block. /// Shows a bordered block with tab bar, question content or review, and help text. pub fn render_ask_user_block( - questions: &[stakpak_shared::models::integrations::openai::AskUserQuestion], + questions: &[stakpak_shared::models::agent_runtime::AskUserQuestion], answers: &std::collections::HashMap< String, - stakpak_shared::models::integrations::openai::AskUserAnswer, + stakpak_shared::models::agent_runtime::AskUserAnswer, >, current_tab: usize, selected_option: usize, @@ -3027,7 +3028,7 @@ pub fn render_ask_user_block( /// Render a task wait block showing progress of background tasks /// Displays a bordered box with task statuses and overall progress pub fn render_task_wait_block( - task_updates: &[stakpak_shared::models::integrations::openai::TaskUpdate], + task_updates: &[stakpak_shared::models::agent_runtime::TaskUpdate], progress: f64, target_task_ids: &[String], terminal_width: usize, @@ -3352,7 +3353,7 @@ pub fn render_task_wait_block( pub fn render_subagent_resume_pending_block<'a>( tool_call: &ToolCall, is_auto_approved: bool, - pause_info: Option<&stakpak_shared::models::integrations::openai::TaskPauseInfo>, + pause_info: Option<&stakpak_shared::models::agent_runtime::TaskPauseInfo>, width: usize, ) -> Vec> { let mut formatted_lines: Vec> = Vec::new(); @@ -3365,22 +3366,25 @@ pub fn render_subagent_resume_pending_block<'a>( let inner_width = width.saturating_sub(4); // Parse arguments to determine resume type - let args = serde_json::from_str::(&tool_call.function.arguments).ok(); - // Extract task_id from arguments - let task_id = args - .as_ref() - .and_then(|a| a.get("task_id").and_then(|v| v.as_str()).map(String::from)) + let task_id = tool_call + .arguments + .get("task_id") + .and_then(|v| v.as_str()) + .map(String::from) .unwrap_or_else(|| "unknown".to_string()); // Check if this is an input-based resume (for completed agents) or tool approval resume - let input_text = args - .as_ref() - .and_then(|a| a.get("input").and_then(|v| v.as_str()).map(String::from)); - - let has_approve_all = args - .as_ref() - .and_then(|a| a.get("approve_all").and_then(|v| v.as_bool())) + let input_text = tool_call + .arguments + .get("input") + .and_then(|v| v.as_str()) + .map(String::from); + + let has_approve_all = tool_call + .arguments + .get("approve_all") + .and_then(|v| v.as_bool()) .unwrap_or(false); // Title @@ -3934,7 +3938,7 @@ mod tests { #[test] fn test_tool_call_stream_block_border_alignment() { - use stakpak_shared::models::integrations::openai::ToolCallStreamInfo; + use stakpak_shared::models::agent_runtime::ToolCallStreamInfo; let infos = vec![ ToolCallStreamInfo { @@ -3974,7 +3978,7 @@ mod tests { #[test] fn test_tool_call_stream_block_overflow_summary() { - use stakpak_shared::models::integrations::openai::ToolCallStreamInfo; + use stakpak_shared::models::agent_runtime::ToolCallStreamInfo; let infos: Vec = (0..8) .map(|i| ToolCallStreamInfo { diff --git a/tui/src/services/board_tasks.rs b/tui/src/services/board_tasks.rs index 155977895..a4d014a2b 100644 --- a/tui/src/services/board_tasks.rs +++ b/tui/src/services/board_tasks.rs @@ -239,16 +239,16 @@ pub fn extract_board_agent_id_from_messages( for message in messages.iter().rev() { let text_to_search = match &message.content { // Tool call results contain the output - MessageContent::RenderResultBorderBlock(result) => Some(result.result.as_str()), - MessageContent::RenderFullContentMessage(result) => Some(result.result.as_str()), - MessageContent::RenderCommandCollapsedResult(result) => Some(result.result.as_str()), + MessageContent::RenderResultBorderBlock(result) => Some(result.result.clone()), + MessageContent::RenderFullContentMessage(result) => Some(result.result.clone()), + MessageContent::RenderCommandCollapsedResult(result) => Some(result.result.clone()), // Tool call arguments (for run_command, the command might set the env var) MessageContent::RenderPendingBorderBlock(tool_call, _) => { - Some(tool_call.function.arguments.as_str()) + Some(tool_call.arguments.to_string()) } MessageContent::RenderCollapsedMessage(tool_call) => { - Some(tool_call.function.arguments.as_str()) + Some(tool_call.arguments.to_string()) } // Run command blocks contain command and result @@ -257,18 +257,21 @@ pub fn extract_board_agent_id_from_messages( if let Some(cap) = agent_id_pattern.find(command) { return Some(cap.as_str().to_string()); } - result.as_deref() + result.clone() } // Plain text and assistant messages might contain agent IDs - MessageContent::Plain(text, _) => Some(text.as_str()), - MessageContent::AssistantMD(text, _) => Some(text.as_str()), - MessageContent::PlainText(text) => Some(text.as_str()), + MessageContent::Plain(text, _) => Some(text.clone()), + MessageContent::AssistantMD(text, _) => Some(text.clone()), + MessageContent::PlainText(text) => Some(text.clone()), _ => None, }; - if let Some(cap) = text_to_search.and_then(|text| agent_id_pattern.find(text)) { + if let Some(cap) = text_to_search + .as_deref() + .and_then(|text| agent_id_pattern.find(text)) + { return Some(cap.as_str().to_string()); } } @@ -646,7 +649,7 @@ mod tests { }; // Progress: all 4 should count as completed - let progress = calculate_progress(&[card.clone()]); + let progress = calculate_progress(std::slice::from_ref(&card)); assert_eq!(progress.completed, 4); assert_eq!(progress.total, 4); diff --git a/tui/src/services/changeset.rs b/tui/src/services/changeset.rs index 10874325d..e812c458b 100644 --- a/tui/src/services/changeset.rs +++ b/tui/src/services/changeset.rs @@ -6,7 +6,7 @@ //! - Backup paths for revert functionality use chrono::{DateTime, Utc}; -use stakpak_shared::models::integrations::openai::ToolCall; +use stakai::ToolCall; use std::collections::HashMap; use std::fs; use std::path::Path; @@ -191,7 +191,7 @@ impl Changeset { pub fn track_file(&mut self, path: &str, edit: FileEdit) { // Check if this is a file creation based on tool call let is_creation = edit.tool_call.as_ref().is_some_and(|tc| { - let name = tc.function.name.as_str(); + let name = tc.name.as_str(); name == "stakpak__create" || name == "stakpak__write" || name == "stakpak__create_file" @@ -332,13 +332,12 @@ impl Changeset { /// Apply a single edit in reverse (replace new_str with old_str) fn apply_reverse_edit(content: &str, tool_call: &ToolCall) -> Result { // Only handle str_replace tool calls - if tool_call.function.name != "stakpak__str_replace" { + if tool_call.name != "stakpak__str_replace" { return Ok(content.to_string()); } // Parse the tool call arguments - let args: serde_json::Value = serde_json::from_str(&tool_call.function.arguments) - .map_err(|e| format!("Failed to parse tool call arguments: {}", e))?; + let args = &tool_call.arguments; let old_str = args["old_str"] .as_str() @@ -752,11 +751,8 @@ mod tests { .with_stats(10, 0) .with_tool_call(ToolCall { id: "call_1".to_string(), - r#type: "function".to_string(), - function: stakpak_shared::models::integrations::openai::FunctionCall { - name: "stakpak__create".to_string(), - arguments: "{}".to_string(), - }, + name: "stakpak__create".to_string(), + arguments: serde_json::json!({}), metadata: None, }), ); diff --git a/tui/src/services/file_diff.rs b/tui/src/services/file_diff.rs index 9234d63d9..8fb708a1b 100644 --- a/tui/src/services/file_diff.rs +++ b/tui/src/services/file_diff.rs @@ -2,7 +2,7 @@ use ratatui::style::{Color, Modifier, Style}; use ratatui::text::{Line, Span}; use regex::Regex; use similar::TextDiff; -use stakpak_shared::models::integrations::openai::ToolCall; +use stakai::ToolCall; use stakpak_shared::utils::strip_tool_name; use unicode_width::UnicodeWidthStr; @@ -499,11 +499,10 @@ pub fn render_file_diff_block_from_args( terminal_width: usize, result: Option<&str>, ) -> (Vec>, Vec>) { - let args: serde_json::Value = serde_json::from_str(&tool_call.function.arguments) - .unwrap_or_else(|_| serde_json::json!({})); + let args = &tool_call.arguments; let old_str = args.get("old_str").and_then(|v| v.as_str()).unwrap_or(""); - let new_str = if strip_tool_name(&tool_call.function.name) == "create" { + let new_str = if strip_tool_name(&tool_call.name) == "create" { args.get("file_text").and_then(|v| v.as_str()).unwrap_or("") } else { args.get("new_str").and_then(|v| v.as_str()).unwrap_or("") diff --git a/tui/src/services/handlers/ask_user.rs b/tui/src/services/handlers/ask_user.rs index b9f7b4246..fcfad8763 100644 --- a/tui/src/services/handlers/ask_user.rs +++ b/tui/src/services/handlers/ask_user.rs @@ -4,8 +4,9 @@ //! option selection, custom input, and submission. use crate::app::{AppState, OutputEvent}; -use stakpak_shared::models::integrations::openai::{ - AskUserAnswer, AskUserQuestion, AskUserResult, ToolCall, ToolCallResult, ToolCallResultStatus, +use stakai::ToolCall; +use stakpak_shared::models::agent_runtime::{ + AskUserAnswer, AskUserQuestion, AskUserResult, ToolCallResult, ToolCallResultStatus, }; use tokio::sync::mpsc::Sender; @@ -590,7 +591,7 @@ mod tests { use super::*; use crate::app::AppStateOptions; use stakai::Model; - use stakpak_shared::models::integrations::openai::{AskUserOption, FunctionCall}; + use stakpak_shared::models::agent_runtime::AskUserOption; use tokio::sync::mpsc; /// Helper to create a minimal AppState for testing @@ -662,11 +663,8 @@ mod tests { fn create_test_tool_call() -> ToolCall { ToolCall { id: "call_test123".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "ask_user".to_string(), - arguments: "{}".to_string(), - }, + name: "ask_user".to_string(), + arguments: serde_json::json!({}), metadata: None, } } diff --git a/tui/src/services/handlers/dialog.rs b/tui/src/services/handlers/dialog.rs index 796e856b9..cf1a8491f 100644 --- a/tui/src/services/handlers/dialog.rs +++ b/tui/src/services/handlers/dialog.rs @@ -13,7 +13,7 @@ use crate::services::message::{ use crate::services::text_selection::SelectionState; use ratatui::layout::Size; use ratatui::style::Color; -use stakpak_shared::models::integrations::openai::ToolCall; +use stakai::ToolCall; use stakpak_shared::utils::strip_tool_name; use tokio::sync::mpsc::Sender; use uuid::Uuid; @@ -28,7 +28,7 @@ fn is_foreground_command_tool(tool_name: &str) -> bool { /// Update a run_command block from Pending to Running state /// This should be called when a run_command tool is approved and starts executing pub fn update_run_command_to_running(state: &mut AppState, tool_call: &ToolCall) { - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if !is_foreground_command_tool(tool_name) { return; } @@ -70,7 +70,7 @@ pub fn update_pending_tool_to_first( .retain(|m| m.id != pending_id); } - let tool_name = strip_tool_name(&first_tool.function.name); + let tool_name = strip_tool_name(&first_tool.name); // Create the appropriate pending block based on tool type if is_foreground_command_tool(tool_name) { @@ -137,7 +137,7 @@ pub fn handle_esc_event( .clear(); // Store the latest tool call for potential retry (only for command tools) if let Some(tool_call) = &state.dialog_approval_state.dialog_command - && is_foreground_command_tool(strip_tool_name(&tool_call.function.name)) + && is_foreground_command_tool(strip_tool_name(&tool_call.name)) { state.tool_call_state.latest_tool_call = Some(tool_call.clone()); } @@ -191,7 +191,7 @@ pub fn handle_esc( .output_tx .try_send(OutputEvent::RejectTool(tool_call.clone(), should_stop)); - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if is_foreground_command_tool(tool_name) { // For command tools, remove the pending unified block and add rejected unified block // Remove pending message by tool_call_id @@ -345,7 +345,7 @@ pub fn handle_esc( /// Handle show confirmation dialog event pub fn handle_show_confirmation_dialog( state: &mut AppState, - tool_call: stakpak_shared::models::integrations::openai::ToolCall, + tool_call: stakai::ToolCall, input_tx: &Sender, output_tx: &Sender, _terminal_size: Size, @@ -360,7 +360,7 @@ pub fn handle_show_confirmation_dialog( .map(|status| status == &ToolCallStatus::Executed) .unwrap_or(false) { - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if is_foreground_command_tool(tool_name) { // Use unified block for command tools let command = @@ -396,7 +396,7 @@ pub fn handle_show_confirmation_dialog( } state.dialog_approval_state.dialog_command = Some(tool_call.clone()); - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if is_foreground_command_tool(tool_name) { state.tool_call_state.latest_tool_call = Some(tool_call.clone()); } @@ -435,13 +435,11 @@ pub fn handle_show_confirmation_dialog( } else if tool_name == "resume_subagent_task" { // For resume_subagent_task, use the special subagent pending block // Try to get pause info from cached subagent state - let pause_info = serde_json::from_str::(&tool_call.function.arguments) - .ok() - .and_then(|args| { - args.get("task_id") - .and_then(|v| v.as_str()) - .map(String::from) - }) + let pause_info = tool_call + .arguments + .get("task_id") + .and_then(|v| v.as_str()) + .map(String::from) .and_then(|task_id| { state .tool_call_state diff --git a/tui/src/services/handlers/input.rs b/tui/src/services/handlers/input.rs index 02c5a8044..a44aa538a 100644 --- a/tui/src/services/handlers/input.rs +++ b/tui/src/services/handlers/input.rs @@ -180,7 +180,8 @@ pub fn handle_input_submitted_event( let is_approved = state .dialog_approval_state .message_approved_tools - .contains(tool_call); + .iter() + .any(|approved| approved.id == tool_call.id); let status = if is_approved { ToolCallStatus::Approved } else { @@ -209,7 +210,8 @@ pub fn handle_input_submitted_event( let is_approved = state .dialog_approval_state .message_approved_tools - .contains(first_tool); + .iter() + .any(|approved| approved.id == first_tool.id); // Update the pending display to show the first tool (which is being executed) // This ensures the UI shows the correct tool as "running", not the selected one diff --git a/tui/src/services/handlers/mod.rs b/tui/src/services/handlers/mod.rs index 243ed419a..e541ecfd0 100644 --- a/tui/src/services/handlers/mod.rs +++ b/tui/src/services/handlers/mod.rs @@ -794,7 +794,8 @@ pub fn update( let is_approved = state .dialog_approval_state .message_approved_tools - .contains(tool_call); + .iter() + .any(|approved| approved.id == tool_call.id); let status = if is_approved { crate::app::ToolCallStatus::Approved } else { @@ -822,7 +823,8 @@ pub fn update( let is_approved = state .dialog_approval_state .message_approved_tools - .contains(first_tool); + .iter() + .any(|approved| approved.id == first_tool.id); // Update the pending display to show the first tool (which is being executed) dialog::update_pending_tool_to_first(state, first_tool, is_approved); @@ -1615,10 +1617,8 @@ mod tests { use crate::app::{AppStateOptions, LoadingOperation}; use crate::services::message::MessageContent; use ratatui::layout::Size; - use stakai::Model; - use stakpak_shared::models::integrations::openai::{ - ContentPart, FunctionCall, ToolCall, ToolCallResult, ToolCallResultStatus, - }; + use stakai::{ContentPart, Model, ToolCall}; + use stakpak_shared::models::agent_runtime::{ToolCallResult, ToolCallResultStatus}; use tokio::sync::mpsc; fn build_state() -> AppState { @@ -1643,11 +1643,8 @@ mod tests { ToolCallResult { call: ToolCall { id: id.to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "run_command".to_string(), - arguments: "{}".to_string(), - }, + name: "run_command".to_string(), + arguments: serde_json::json!({}), metadata: None, }, result: format!("result-{id}"), @@ -1656,11 +1653,7 @@ mod tests { } fn make_image_part(label: &str) -> ContentPart { - ContentPart { - r#type: "text".to_string(), - text: Some(label.to_string()), - image_url: None, - } + ContentPart::text(label) } #[tokio::test] @@ -1875,7 +1868,7 @@ mod tests { #[tokio::test] async fn ask_user_arrows_navigate_options_via_update() { - use stakpak_shared::models::integrations::openai::{AskUserOption, AskUserQuestion}; + use stakpak_shared::models::agent_runtime::{AskUserOption, AskUserQuestion}; let mut state = build_state(); let (input_tx, _input_rx) = mpsc::channel(8); @@ -1911,11 +1904,8 @@ mod tests { }]; let tool_call = ToolCall { id: "tc_1".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "ask_user".to_string(), - arguments: "{}".to_string(), - }, + name: "ask_user".to_string(), + arguments: serde_json::json!({}), metadata: None, }; ask_user::handle_show_ask_user_popup(&mut state, tool_call, questions); @@ -1976,7 +1966,7 @@ mod tests { #[tokio::test] async fn ask_user_left_right_navigate_tabs_via_update() { - use stakpak_shared::models::integrations::openai::{AskUserOption, AskUserQuestion}; + use stakpak_shared::models::agent_runtime::{AskUserOption, AskUserQuestion}; let mut state = build_state(); let (input_tx, _input_rx) = mpsc::channel(8); @@ -2012,11 +2002,8 @@ mod tests { ]; let tool_call = ToolCall { id: "tc_2".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "ask_user".to_string(), - arguments: "{}".to_string(), - }, + name: "ask_user".to_string(), + arguments: serde_json::json!({}), metadata: None, }; ask_user::handle_show_ask_user_popup(&mut state, tool_call, questions); @@ -2059,7 +2046,7 @@ mod tests { #[tokio::test] async fn ask_user_arrows_always_navigate_options() { - use stakpak_shared::models::integrations::openai::{AskUserOption, AskUserQuestion}; + use stakpak_shared::models::agent_runtime::{AskUserOption, AskUserQuestion}; let mut state = build_state(); let (input_tx, _input_rx) = mpsc::channel(8); @@ -2088,11 +2075,8 @@ mod tests { }]; let tool_call = ToolCall { id: "tc_3".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "ask_user".to_string(), - arguments: "{}".to_string(), - }, + name: "ask_user".to_string(), + arguments: serde_json::json!({}), metadata: None, }; ask_user::handle_show_ask_user_popup(&mut state, tool_call, questions); diff --git a/tui/src/services/handlers/shell.rs b/tui/src/services/handlers/shell.rs index 2602b2e4d..ec5704296 100644 --- a/tui/src/services/handlers/shell.rs +++ b/tui/src/services/handlers/shell.rs @@ -15,10 +15,9 @@ use crate::services::shell_mode::run_pty_command; use crate::services::shell_mode::{SHELL_PROMPT_PREFIX, ShellEvent}; use ratatui::style::{Color, Modifier, Style}; use ratatui::text::{Line, Span}; +use stakai::ToolCall; use stakpak_shared::helper::truncate_output; -use stakpak_shared::models::integrations::openai::{ - FunctionCall, ToolCall, ToolCallResult, ToolCallResultStatus, -}; +use stakpak_shared::models::agent_runtime::{ToolCallResult, ToolCallResultStatus}; use tokio::sync::mpsc; use tokio::sync::mpsc::Sender; use uuid::Uuid; @@ -163,10 +162,7 @@ pub fn send_shell_input(state: &mut AppState, data: &str) { /// Extract command from tool call pub fn extract_command_from_tool_call(tool_call: &ToolCall) -> Result { // Parse as JSON and extract the command field - let json = serde_json::from_str::(&tool_call.function.arguments) - .map_err(|e| format!("Failed to parse JSON: {}", e))?; - - if let Some(command_value) = json.get("command") { + if let Some(command_value) = tool_call.arguments.get("command") { if let Some(command_str) = command_value.as_str() { return Ok(command_str.to_string()); } else { @@ -1010,7 +1006,7 @@ pub fn shell_command_to_tool_call_result( shell_output: Option, ) -> ToolCallResult { let (id, name) = if let Some(cmd) = &state.dialog_approval_state.dialog_command { - (cmd.id.clone(), cmd.function.name.clone()) + (cmd.id.clone(), cmd.name.clone()) } else { ( format!("tool_{}", Uuid::new_v4()), @@ -1050,11 +1046,8 @@ pub fn shell_command_to_tool_call_result( let call = ToolCall { id, - r#type: "function".to_string(), - function: FunctionCall { - name, - arguments: args, - }, + name, + arguments: serde_json::from_str(&args).unwrap_or(serde_json::Value::String(args)), metadata: None, }; ToolCallResult { diff --git a/tui/src/services/handlers/tool.rs b/tui/src/services/handlers/tool.rs index 75c8295f6..ca75100bd 100644 --- a/tui/src/services/handlers/tool.rs +++ b/tui/src/services/handlers/tool.rs @@ -6,9 +6,9 @@ use crate::app::{AppState, InputEvent, OutputEvent, ToolCallStatus}; use crate::services::commands::{CommandAction, CommandContext, execute_command, filter_commands}; use crate::services::helper_block::push_error_message; use crate::services::message::{Message, invalidate_message_lines_cache}; -use stakpak_shared::models::integrations::openai::{ - ProgressType, ToolCall, ToolCallResult, ToolCallResultProgress, ToolCallResultStatus, - ToolCallStreamInfo, +use stakai::ToolCall; +use stakpak_shared::models::agent_runtime::{ + ProgressType, ToolCallResult, ToolCallResultProgress, ToolCallResultStatus, ToolCallStreamInfo, }; use stakpak_shared::utils::strip_tool_name; use tokio::sync::mpsc::Sender; @@ -113,7 +113,7 @@ pub fn handle_stream_tool_result( .as_ref() .map(|tc| { matches!( - strip_tool_name(&tc.function.name), + strip_tool_name(&tc.name), "run_command" | "run_remote_command" ) }) @@ -600,14 +600,8 @@ pub fn handle_tool_result(state: &mut AppState, result: ToolCallResult) { return; } - let function_name = result.call.function.name.as_str(); - let args_str = &result.call.function.arguments; - - // Parse arguments - let args: serde_json::Value = match serde_json::from_str(args_str) { - Ok(v) => v, - Err(_) => return, // Should not happen if tool call was successful - }; + let function_name = result.call.name.as_str(); + let args = &result.call.arguments; // Normalize/Strip tool name for checking let tool_name_stripped = strip_tool_name(function_name); @@ -902,29 +896,25 @@ fn extract_diff_preview(message: &str) -> Option { pub fn extract_view_params_from_tool_call( tool_call: &ToolCall, ) -> (Option, Option, Option) { - // Try to parse arguments as JSON - if let Ok(args) = serde_json::from_str::(&tool_call.function.arguments) { - let path = args - .get("path") - .or(args.get("filePath")) - .or(args.get("file_path")) - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - let grep = args - .get("grep") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - let glob = args - .get("glob") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - return (path, grep, glob); - } - - (None, None, None) + let args = &tool_call.arguments; + let path = args + .get("path") + .or(args.get("filePath")) + .or(args.get("file_path")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let grep = args + .get("grep") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let glob = args + .get("glob") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + (path, grep, glob) } // ========== Approval Bar Handlers ========== @@ -1019,7 +1009,7 @@ pub fn create_pending_block_for_selected_tool(state: &mut AppState) { // Get the currently selected tool call if let Some(action) = state.dialog_approval_state.approval_bar.selected_action() { let tool_call = &action.tool_call; - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); // Determine the approval state for display let auto_approve = action.status == crate::services::approval_bar::ApprovalStatus::Approved; @@ -1043,21 +1033,18 @@ pub fn create_pending_block_for_selected_tool(state: &mut AppState) { state.messages_scrolling_state.messages.push(msg); } else if tool_name == "resume_subagent_task" { // For resume_subagent_task, use the special subagent pending block - let pause_info = - serde_json::from_str::(&tool_call.function.arguments) - .ok() - .and_then(|args| { - args.get("task_id") - .and_then(|v| v.as_str()) - .map(String::from) - }) - .and_then(|task_id| { - state - .tool_call_state - .subagent_pause_info - .get(&task_id) - .cloned() - }); + let pause_info = tool_call + .arguments + .get("task_id") + .and_then(|v| v.as_str()) + .map(String::from) + .and_then(|task_id| { + state + .tool_call_state + .subagent_pause_info + .get(&task_id) + .cloned() + }); let msg = Message::render_subagent_resume_pending_block( tool_call.clone(), diff --git a/tui/src/services/image_upload.rs b/tui/src/services/image_upload.rs index 734425a0d..bfde37e54 100644 --- a/tui/src/services/image_upload.rs +++ b/tui/src/services/image_upload.rs @@ -7,7 +7,7 @@ //! - Creating ContentParts for API communication use regex::Regex; -use stakpak_shared::models::integrations::openai::{ContentPart, ImageUrl}; +use stakai::ContentPart; use std::path::{Path, PathBuf}; /// Supported image file extensions @@ -56,14 +56,7 @@ pub fn extract_image_urls(text: &str) -> Vec { /// Extract local file paths from text (handles quoted and unquoted paths) /// Create an image content part from a URL (sends URL directly, no download) pub fn create_image_part_from_url(url: &str) -> ContentPart { - ContentPart { - r#type: "input_image".to_string(), - text: None, - image_url: Some(ImageUrl { - url: url.to_string(), - detail: None, - }), - } + ContentPart::image(url) } /// Create an image content part from a file path @@ -98,14 +91,7 @@ pub fn create_image_part_from_path(path: &Path) -> Option { let base64_data = general_purpose::STANDARD.encode(&image_data); let data_url = format!("data:{};base64,{}", mime_type, base64_data); - Some(ContentPart { - r#type: "input_image".to_string(), - text: None, - image_url: Some(ImageUrl { - url: data_url, - detail: None, - }), - }) + Some(ContentPart::image(data_url)) } fn detect_mime_type_from_content(data: &[u8], path: &Path) -> &'static str { diff --git a/tui/src/services/message.rs b/tui/src/services/message.rs index f0b8325e9..f9290154e 100644 --- a/tui/src/services/message.rs +++ b/tui/src/services/message.rs @@ -12,10 +12,9 @@ use ratatui::style::{Modifier, Style}; use ratatui::text::{Line, Span}; use regex::Regex; use serde_json::Value; -#[cfg(test)] -use stakpak_shared::models::integrations::openai::FunctionCall; -use stakpak_shared::models::integrations::openai::{ - ToolCall, ToolCallResult, ToolCallResultStatus, ToolCallStreamInfo, +use stakai::ToolCall; +use stakpak_shared::models::agent_runtime::{ + ToolCallResult, ToolCallResultStatus, ToolCallStreamInfo, }; use stakpak_shared::utils::strip_tool_name; use std::collections::hash_map::DefaultHasher; @@ -74,7 +73,7 @@ pub enum MessageContent { /// Task wait block - shows progress of background tasks being waited on /// (task_updates: Vec, overall_progress: f64, target_task_ids: Vec) RenderTaskWaitBlock( - Vec, + Vec, f64, Vec, ), @@ -83,18 +82,16 @@ pub enum MessageContent { RenderSubagentResumePendingBlock( ToolCall, bool, - Option, + Option, ), /// Tool call streaming preview - shows tools being generated with token counters /// (tool_infos: Vec) RenderToolCallStreamBlock(Vec), /// Ask user inline block - renders questions and options inline in the message area RenderAskUserBlock { - questions: Vec, - answers: std::collections::HashMap< - String, - stakpak_shared::models::integrations::openai::AskUserAnswer, - >, + questions: Vec, + answers: + std::collections::HashMap, current_tab: usize, selected_option: usize, custom_input: String, @@ -149,7 +146,7 @@ pub fn hash_message_content(content: &MessageContent) -> u64 { MessageContent::RenderPendingBorderBlock(tool_call, is_auto) => { 6u8.hash(&mut hasher); tool_call.id.hash(&mut hasher); - tool_call.function.arguments.hash(&mut hasher); + tool_call.arguments.hash(&mut hasher); is_auto.hash(&mut hasher); } MessageContent::RenderPendingBorderBlockWithStallWarning(tool_call, is_auto, msg) => { @@ -594,10 +591,10 @@ impl Message { } pub fn render_ask_user_block( - questions: Vec, + questions: Vec, answers: std::collections::HashMap< String, - stakpak_shared::models::integrations::openai::AskUserAnswer, + stakpak_shared::models::agent_runtime::AskUserAnswer, >, current_tab: usize, selected_option: usize, @@ -653,7 +650,7 @@ impl Message { /// Create a task wait block message /// Shows progress of background tasks being waited on with status indicators pub fn render_task_wait_block( - task_updates: Vec, + task_updates: Vec, progress: f64, target_task_ids: Vec, message_id: Option, @@ -670,7 +667,7 @@ impl Message { pub fn render_subagent_resume_pending_block( tool_call: ToolCall, is_auto_approved: bool, - pause_info: Option, + pause_info: Option, message_id: Option, ) -> Self { Message { @@ -1432,7 +1429,7 @@ fn render_single_message_internal(msg: &Message, width: usize) -> Vec<(Line<'sta } MessageContent::RenderPendingBorderBlock(tool_call, is_auto_approved) => { let full_command = extract_full_command_arguments(tool_call); - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); let rendered = if (tool_name == "str_replace" || tool_name == "create") && !render_file_diff(tool_call, width).is_empty() { @@ -1480,7 +1477,7 @@ fn render_single_message_internal(msg: &Message, width: usize) -> Vec<(Line<'sta lines.extend(convert_to_owned_lines(borrowed)); } MessageContent::RenderCollapsedMessage(tool_call) => { - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if (tool_name == "str_replace" || tool_name == "create") && let Some(rendered) = render_file_diff_full(tool_call, width, Some(true), None) && !rendered.is_empty() @@ -1513,8 +1510,7 @@ fn render_single_message_internal(msg: &Message, width: usize) -> Vec<(Line<'sta lines.extend(convert_to_owned_lines(borrowed)); } MessageContent::RenderFullContentMessage(tool_call_result) => { - let tool_name = - stakpak_shared::utils::strip_tool_name(&tool_call_result.call.function.name); + let tool_name = stakpak_shared::utils::strip_tool_name(&tool_call_result.call.name); // For str_replace/create, use the diff view with proper line numbers if (tool_name == "str_replace" || tool_name == "create") @@ -2138,7 +2134,7 @@ fn get_wrapped_message_lines_internal( } MessageContent::RenderPendingBorderBlock(tool_call, is_auto_approved) => { let full_command = extract_full_command_arguments(tool_call); - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); let rendered_lines = if (tool_name == "str_replace" || tool_name == "create") && !render_file_diff(tool_call, width).is_empty() { @@ -2192,7 +2188,7 @@ fn get_wrapped_message_lines_internal( } MessageContent::RenderCollapsedMessage(tool_call) => { - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if (tool_name == "str_replace" || tool_name == "create") && let Some(rendered_lines) = render_file_diff_full(tool_call, width, Some(true), None) @@ -2234,8 +2230,7 @@ fn get_wrapped_message_lines_internal( all_lines.extend(owned_lines); } MessageContent::RenderFullContentMessage(tool_call_result) => { - let tool_name = - stakpak_shared::utils::strip_tool_name(&tool_call_result.call.function.name); + let tool_name = stakpak_shared::utils::strip_tool_name(&tool_call_result.call.name); // For str_replace/create, use the diff view with proper line numbers if (tool_name == "str_replace" || tool_name == "create") @@ -2445,23 +2440,21 @@ fn get_wrapped_message_lines_internal( } pub fn extract_truncated_command_arguments(tool_call: &ToolCall, sign: Option) -> String { - let arguments = serde_json::from_str::(&tool_call.function.arguments); + let arguments = &tool_call.arguments; // For subagent tasks, show description + tools summary instead of raw args - let tool_name = strip_tool_name(&tool_call.function.name); - if tool_name == "dynamic_subagent_task" - && let Ok(ref args) = arguments - { - let desc = args + let tool_name = strip_tool_name(&tool_call.name); + if tool_name == "dynamic_subagent_task" { + let desc = arguments .get("description") .and_then(|v| v.as_str()) .unwrap_or("subagent"); - let tools_count = args + let tools_count = arguments .get("tools") .and_then(|v| v.as_array()) .map(|a| a.len()) .unwrap_or(0); - let is_sandbox = args + let is_sandbox = arguments .get("enable_sandbox") .and_then(|v| v.as_bool()) .unwrap_or(false); @@ -2471,9 +2464,9 @@ pub fn extract_truncated_command_arguments(tool_call: &ToolCall, sign: Option(&tool_call.function.arguments) + && let Ok(request) = serde_json::from_value::< + stakpak_shared::models::agent_runtime::AskUserRequest, + >(tool_call.arguments.clone()) { let labels: Vec<&str> = request.questions.iter().map(|q| q.label.as_str()).collect(); if !labels.is_empty() { @@ -2484,7 +2477,7 @@ pub fn extract_truncated_command_arguments(tool_call: &ToolCall, sign: Option String { - let tool_name = strip_tool_name(&tool_call.function.name); + let tool_name = strip_tool_name(&tool_call.name); if tool_name == "ask_user" { return "Ask user questions...".to_string(); } - // First try to parse as valid JSON - if let Ok(v) = serde_json::from_str::(&tool_call.function.arguments) { - return format_json_value(&v); + if !tool_call.arguments.is_null() { + return format_json_value(&tool_call.arguments); } - // If JSON parsing fails, try regex patterns for malformed JSON + let raw_arguments = tool_call.arguments.as_str().unwrap_or(""); + + // If arguments are a malformed JSON string, try regex patterns let patterns = vec![ // Pattern for key-value pairs with quotes r#"["']?(\w+)["']?\s*:\s*["']([^"']+)["']"#, @@ -2531,7 +2525,7 @@ pub fn extract_full_command_arguments(tool_call: &ToolCall) -> String { for pattern in patterns { if let Ok(re) = Regex::new(pattern) { let mut results = Vec::new(); - for caps in re.captures_iter(&tool_call.function.arguments) { + for caps in re.captures_iter(raw_arguments) { if caps.len() >= 3 { let key = caps.get(1).map(|m| m.as_str()).unwrap_or(""); let value = caps.get(2).map(|m| m.as_str()).unwrap_or(""); @@ -2545,19 +2539,19 @@ pub fn extract_full_command_arguments(tool_call: &ToolCall) -> String { } // Try to wrap in braces and parse as JSON - let wrapped = format!("{{{}}}", tool_call.function.arguments); + let wrapped = format!("{{{}}}", raw_arguments); if let Ok(v) = serde_json::from_str::(&wrapped) { return format_json_value(&v); } // If all else fails, return the raw arguments if they're not empty - let trimmed = tool_call.function.arguments.trim(); + let trimmed = raw_arguments.trim(); if !trimmed.is_empty() { return trimmed.to_string(); } // Last resort - format!("function_name={}", tool_call.function.name) + format!("function_name={}", tool_call.name) } fn format_json_value(value: &Value) -> String { @@ -2750,7 +2744,7 @@ pub fn extract_command_purpose(command: &str, outside_title: &str) -> String { // Helper function to get command name for the outside title pub fn get_command_type_name(tool_call: &ToolCall) -> String { - match strip_tool_name(&tool_call.function.name) { + match strip_tool_name(&tool_call.name) { "create_file" => "Create file".to_string(), "create" => "Create".to_string(), "edit_file" => "Edit file".to_string(), @@ -2761,8 +2755,7 @@ pub fn get_command_type_name(tool_call: &ToolCall) -> String { "list_directory" => "List directory".to_string(), "search_files" => "Search files".to_string(), "dynamic_subagent_task" => { - let args = - serde_json::from_str::(&tool_call.function.arguments).ok(); + let args = Some(&tool_call.arguments); let desc = args .as_ref() .and_then(|a| a.get("description").and_then(|v| v.as_str())) @@ -2780,7 +2773,7 @@ pub fn get_command_type_name(tool_call: &ToolCall) -> String { "ask_user" => "Ask User".to_string(), _ => { // Convert function name to title case - strip_tool_name(&tool_call.function.name) + strip_tool_name(&tool_call.name) .replace("_", " ") .split_whitespace() .map(|word| { @@ -2818,11 +2811,9 @@ mod tests { for (input, expected) in test_cases { let tool_call = ToolCall { id: "test".to_string(), - r#type: "function".to_string(), - function: FunctionCall { - name: "test".to_string(), - arguments: input.to_string(), - }, + name: "test".to_string(), + arguments: serde_json::from_str(input) + .unwrap_or_else(|_| serde_json::Value::String(input.to_string())), metadata: None, }; From f529ea81bfdd836cc178e555329d12f440727c82 Mon Sep 17 00:00:00 2001 From: Ahmed Hesham Abdelkader <23265119+ahmedhesham6@users.noreply.github.com> Date: Sat, 9 May 2026 00:35:42 +0200 Subject: [PATCH 2/5] chore(tui): fix clippy test warnings --- tui/src/services/detect_term.rs | 8 +-- tui/src/services/markdown_renderer.rs | 8 +-- tui/src/services/textarea.rs | 72 +++++++++++++-------------- 3 files changed, 44 insertions(+), 44 deletions(-) diff --git a/tui/src/services/detect_term.rs b/tui/src/services/detect_term.rs index eaafdfd9d..ddf2dad2e 100644 --- a/tui/src/services/detect_term.rs +++ b/tui/src/services/detect_term.rs @@ -896,12 +896,12 @@ mod tests { // White background (15) should be light // This would require setting env var, so we test the parsing logic let colorfgbg_white = "0;15"; - let bg_str = colorfgbg_white.split(';').last(); + let bg_str = colorfgbg_white.split(';').next_back(); assert_eq!(bg_str, Some("15")); // Light gray background (7) should be light let colorfgbg_light_gray = "0;7"; - let bg_str = colorfgbg_light_gray.split(';').last(); + let bg_str = colorfgbg_light_gray.split(';').next_back(); assert_eq!(bg_str, Some("7")); } @@ -909,12 +909,12 @@ mod tests { fn test_detect_theme_from_colorfgbg_dark() { // Black background (0) should be dark let colorfgbg_black = "15;0"; - let bg_str = colorfgbg_black.split(';').last(); + let bg_str = colorfgbg_black.split(';').next_back(); assert_eq!(bg_str, Some("0")); // Other colors should default to dark let colorfgbg_blue = "15;4"; - let bg_str = colorfgbg_blue.split(';').last(); + let bg_str = colorfgbg_blue.split(';').next_back(); assert_eq!(bg_str, Some("4")); } diff --git a/tui/src/services/markdown_renderer.rs b/tui/src/services/markdown_renderer.rs index ec56385c0..f58203f75 100644 --- a/tui/src/services/markdown_renderer.rs +++ b/tui/src/services/markdown_renderer.rs @@ -1908,9 +1908,9 @@ mod tests { let style = MarkdownStyle::adaptive(); // Verify that the style has proper colors set - assert!(matches!(style.text_style.fg, Some(_))); - assert!(matches!(style.h1_style.fg, Some(_))); - assert!(matches!(style.code_style.fg, Some(_))); + assert!(style.text_style.fg.is_some()); + assert!(style.h1_style.fg.is_some()); + assert!(style.code_style.fg.is_some()); } #[test] @@ -2004,6 +2004,6 @@ mod tests { assert!(!lines.is_empty()); // Should not panic and should produce some output - assert!(lines.len() > 0); + assert!(!lines.is_empty()); } } diff --git a/tui/src/services/textarea.rs b/tui/src/services/textarea.rs index 751813f6a..6142459fa 100644 --- a/tui/src/services/textarea.rs +++ b/tui/src/services/textarea.rs @@ -2012,53 +2012,53 @@ mod tests { } 14 => { // Try inserting inside an existing element (should clamp to boundary) - if let Some(payload) = elem_texts.choose(&mut rng).cloned() { - if let Some(start) = ta.text().find(&payload) { - let end = start + payload.len(); - if end - start > 2 { - let pos = rng.random_range(start + 1..end - 1); - let ins = rand_grapheme(&mut rng); - ta.insert_str_at(pos, &ins); - } + if let Some(payload) = elem_texts.choose(&mut rng).cloned() + && let Some(start) = ta.text().find(&payload) + { + let end = start + payload.len(); + if end - start > 2 { + let pos = rng.random_range(start + 1..end - 1); + let ins = rand_grapheme(&mut rng); + ta.insert_str_at(pos, &ins); } } } 15 => { // Replace a range that intersects an element -> whole element should be replaced - if let Some(payload) = elem_texts.choose(&mut rng).cloned() { - if let Some(start) = ta.text().find(&payload) { - let end = start + payload.len(); - // Create an intersecting range [start-δ, end-δ2) - let mut s = start.saturating_sub(rng.random_range(0..=2)); - let mut e = (end + rng.random_range(0..=2)).min(ta.text().len()); - // Align to char boundaries to satisfy String::replace_range contract - let txt = ta.text(); - while s > 0 && !txt.is_char_boundary(s) { - s -= 1; - } - while e < txt.len() && !txt.is_char_boundary(e) { - e += 1; - } - if s < e { - // Small replacement text - let mut srep = String::new(); - for _ in 0..rng.random_range(0..=2) { - srep.push_str(&rand_grapheme(&mut rng)); - } - ta.replace_range(s..e, &srep); + if let Some(payload) = elem_texts.choose(&mut rng).cloned() + && let Some(start) = ta.text().find(&payload) + { + let end = start + payload.len(); + // Create an intersecting range [start-δ, end-δ2) + let mut s = start.saturating_sub(rng.random_range(0..=2)); + let mut e = (end + rng.random_range(0..=2)).min(ta.text().len()); + // Align to char boundaries to satisfy String::replace_range contract + let txt = ta.text(); + while s > 0 && !txt.is_char_boundary(s) { + s -= 1; + } + while e < txt.len() && !txt.is_char_boundary(e) { + e += 1; + } + if s < e { + // Small replacement text + let mut srep = String::new(); + for _ in 0..rng.random_range(0..=2) { + srep.push_str(&rand_grapheme(&mut rng)); } + ta.replace_range(s..e, &srep); } } } 16 => { // Try setting the cursor to a position inside an element; it should clamp out - if let Some(payload) = elem_texts.choose(&mut rng).cloned() { - if let Some(start) = ta.text().find(&payload) { - let end = start + payload.len(); - if end - start > 2 { - let pos = rng.random_range(start + 1..end - 1); - ta.set_cursor(pos); - } + if let Some(payload) = elem_texts.choose(&mut rng).cloned() + && let Some(start) = ta.text().find(&payload) + { + let end = start + payload.len(); + if end - start > 2 { + let pos = rng.random_range(start + 1..end - 1); + ta.set_cursor(pos); } } } From efebb6241b73132a9f4960a401f93df163db8fe7 Mon Sep 17 00:00:00 2001 From: Ahmed Hesham Abdelkader <23265119+ahmedhesham6@users.noreply.github.com> Date: Thu, 14 May 2026 11:46:32 +0200 Subject: [PATCH 3/5] Refactor: Improve tool result extraction and handling This commit refactors how tool results are extracted and handled in several areas: - `cli/src/commands/agent/run/checkpoint.rs`: - Introduces constants for known tool result markers (e.g., `TOOL_CALL_CANCELLED`, `TOOL_CALL_REJECTED`). - Implements `tool_result_status_from_content` to correctly infer the `ToolCallResultStatus` from various content formats, including explicit markers and structured error objects. - Enhances `tool_result_content_string` to handle non-string JSON content more robustly. - Adds tests to verify the extraction of cancelled and various error statuses for tool results. - `cli/src/commands/agent/run/tooling.rs`: - Modifies `run_tool_call` to use a new `mcp_tool_arguments` function for parsing tool call arguments. - `mcp_tool_arguments` now returns a `CallToolResult` error if the arguments are not a JSON object or null, improving argument validation. - Adds a helper `json_type_name` for better error messages. - Includes a test to ensure --- cli/src/commands/agent/run/checkpoint.rs | 157 ++++++++++++++++++++++- cli/src/commands/agent/run/tooling.rs | 66 +++++++++- tui/src/services/message.rs | 78 ++++++----- 3 files changed, 264 insertions(+), 37 deletions(-) diff --git a/cli/src/commands/agent/run/checkpoint.rs b/cli/src/commands/agent/run/checkpoint.rs index 526749600..39870e8a5 100644 --- a/cli/src/commands/agent/run/checkpoint.rs +++ b/cli/src/commands/agent/run/checkpoint.rs @@ -6,6 +6,15 @@ use stakpak_shared::models::agent_runtime::{ToolCallResult, ToolCallResultStatus use stakpak_tui::{InputEvent, LoadingOperation}; use uuid::Uuid; +const CANCELLED_TOOL_RESULT_MARKER: &str = "TOOL_CALL_CANCELLED"; +const ERROR_TOOL_RESULT_MARKERS: [&str; 5] = [ + "TOOL_CALL_REJECTED", + "INVALID_ARGUMENTS", + "MCP_TOOL_CALL_ERROR", + "MCP_ERROR", + "UNEXPECTED_RESPONSE", +]; + pub async fn get_checkpoint_messages( client: &dyn AgentProvider, checkpoint_id: &str, @@ -95,8 +104,8 @@ pub async fn extract_checkpoint_messages_and_tool_calls( input_tx, InputEvent::ToolResult(ToolCallResult { call: tool_call, - result: content.to_string(), - status: ToolCallResultStatus::Success, + result: tool_result_content_string(&content), + status: tool_result_status_from_content(&content), }), ) .await; @@ -172,6 +181,66 @@ pub fn extract_checkpoint_id_from_messages(messages: &[Message]) -> Option String { + content + .as_str() + .map(ToString::to_string) + .unwrap_or_else(|| content.to_string()) +} + +fn tool_result_status_from_content(content: &serde_json::Value) -> ToolCallResultStatus { + if value_contains_marker(content, CANCELLED_TOOL_RESULT_MARKER) { + return ToolCallResultStatus::Cancelled; + } + + if ERROR_TOOL_RESULT_MARKERS + .iter() + .any(|marker| value_contains_marker(content, marker)) + { + return ToolCallResultStatus::Error; + } + + let serde_json::Value::Object(obj) = content else { + return ToolCallResultStatus::Success; + }; + + if obj + .get("is_error") + .or_else(|| obj.get("isError")) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) + || obj.contains_key("error") + || obj.contains_key("errors") + { + return ToolCallResultStatus::Error; + } + + match obj.get("status").and_then(serde_json::Value::as_str) { + Some(status) if status.eq_ignore_ascii_case("cancelled") => ToolCallResultStatus::Cancelled, + Some(status) + if status.eq_ignore_ascii_case("error") + || status.eq_ignore_ascii_case("failed") + || status.eq_ignore_ascii_case("rejected") => + { + ToolCallResultStatus::Error + } + _ => ToolCallResultStatus::Success, + } +} + +fn value_contains_marker(value: &serde_json::Value, marker: &str) -> bool { + match value { + serde_json::Value::String(text) => text.contains(marker), + serde_json::Value::Array(values) => values + .iter() + .any(|value| value_contains_marker(value, marker)), + serde_json::Value::Object(map) => map + .values() + .any(|value| value_contains_marker(value, marker)), + _ => false, + } +} + /// Resumes a session from a checkpoint, loading messages and tool calls pub async fn resume_session_from_checkpoint( client: &dyn AgentProvider, @@ -207,8 +276,27 @@ pub async fn resume_session_from_checkpoint( #[cfg(test)] mod tests { use super::*; + use serde_json::json; use tokio::sync::mpsc; + fn assistant_tool_call_message(tool_call_id: &str) -> Message { + Message::new( + Role::Assistant, + MessageContent::Parts(vec![ContentPart::tool_call( + tool_call_id, + "test_tool", + json!({"path": "README.md"}), + )]), + ) + } + + fn tool_result_message(tool_call_id: &str, content: serde_json::Value) -> Message { + Message::new( + Role::Tool, + MessageContent::Parts(vec![ContentPart::tool_result(tool_call_id, content)]), + ) + } + #[tokio::test] async fn extract_checkpoint_messages_replays_user_text_without_injected_context_blocks() { let (input_tx, mut input_rx) = mpsc::channel(8); @@ -241,4 +329,69 @@ mod tests { event => panic!("unexpected input event: {event:?}"), } } + + #[tokio::test] + async fn extract_checkpoint_messages_replays_cancelled_tool_result_status() { + let (input_tx, mut input_rx) = mpsc::channel(8); + let messages = vec![ + assistant_tool_call_message("tc_1"), + tool_result_message("tc_1", json!("TOOL_CALL_CANCELLED")), + Message::new(Role::Assistant, "done".to_string()), + ]; + + extract_checkpoint_messages_and_tool_calls("checkpoint-id", &input_tx, messages) + .await + .expect("checkpoint extraction"); + + match input_rx.recv().await.expect("input event") { + InputEvent::ToolResult(result) => { + assert_eq!(result.status, ToolCallResultStatus::Cancelled); + assert_eq!(result.result, "TOOL_CALL_CANCELLED"); + } + event => panic!("unexpected input event: {event:?}"), + } + } + + #[tokio::test] + async fn extract_checkpoint_messages_replays_rejected_tool_result_status() { + let (input_tx, mut input_rx) = mpsc::channel(8); + let messages = vec![ + assistant_tool_call_message("tc_1"), + tool_result_message("tc_1", json!("TOOL_CALL_REJECTED")), + Message::new(Role::Assistant, "done".to_string()), + ]; + + extract_checkpoint_messages_and_tool_calls("checkpoint-id", &input_tx, messages) + .await + .expect("checkpoint extraction"); + + match input_rx.recv().await.expect("input event") { + InputEvent::ToolResult(result) => { + assert_eq!(result.status, ToolCallResultStatus::Error); + assert_eq!(result.result, "TOOL_CALL_REJECTED"); + } + event => panic!("unexpected input event: {event:?}"), + } + } + + #[tokio::test] + async fn extract_checkpoint_messages_replays_structured_error_tool_result_status() { + let (input_tx, mut input_rx) = mpsc::channel(8); + let messages = vec![ + assistant_tool_call_message("tc_1"), + tool_result_message("tc_1", json!({"isError": true, "message": "bad args"})), + Message::new(Role::Assistant, "done".to_string()), + ]; + + extract_checkpoint_messages_and_tool_calls("checkpoint-id", &input_tx, messages) + .await + .expect("checkpoint extraction"); + + match input_rx.recv().await.expect("input event") { + InputEvent::ToolResult(result) => { + assert_eq!(result.status, ToolCallResultStatus::Error); + } + event => panic!("unexpected input event: {event:?}"), + } + } } diff --git a/cli/src/commands/agent/run/tooling.rs b/cli/src/commands/agent/run/tooling.rs index a5abbad2e..cb1eceecd 100644 --- a/cli/src/commands/agent/run/tooling.rs +++ b/cli/src/commands/agent/run/tooling.rs @@ -1,6 +1,6 @@ use rmcp::model::{ CallToolRequestParam, CallToolResult, CancelledNotification, CancelledNotificationParam, - ServerResult, + Content, ServerResult, }; use stakpak_api::AgentProvider; use stakpak_api::storage::ListSessionsQuery; @@ -45,8 +45,10 @@ pub async fn run_tool_call( let tool_exists = tools.iter().any(|tool| tool.name == *tool_name); if tool_exists { - // Parse arguments safely - let arguments = tool_call.arguments.as_object().cloned(); + let arguments = match mcp_tool_arguments(tool_call) { + Ok(arguments) => arguments, + Err(error) => return Ok(Some(error)), + }; // Call tool and handle errors gracefully let metadata = Some({ @@ -163,3 +165,61 @@ pub async fn run_tool_call( Ok(None) } + +fn mcp_tool_arguments( + tool_call: &stakai::ToolCall, +) -> Result>, CallToolResult> { + match &tool_call.arguments { + serde_json::Value::Object(arguments) => Ok(Some(arguments.clone())), + serde_json::Value::Null => Ok(None), + other => Err(CallToolResult::error(vec![ + Content::text("INVALID_ARGUMENTS"), + Content::text(format!( + "Tool '{}' arguments must be a JSON object, got {}", + tool_call.name, + json_type_name(other) + )), + ])), + } +} + +fn json_type_name(value: &serde_json::Value) -> &'static str { + match value { + serde_json::Value::Null => "null", + serde_json::Value::Bool(_) => "boolean", + serde_json::Value::Number(_) => "number", + serde_json::Value::String(_) => "string", + serde_json::Value::Array(_) => "array", + serde_json::Value::Object(_) => "object", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn mcp_tool_arguments_rejects_string_arguments() { + let tool_call = stakai::ToolCall { + id: "tc_1".to_string(), + name: "test_tool".to_string(), + arguments: serde_json::Value::String(r#"{"path":"README.md"}"#.to_string()), + metadata: None, + }; + + let error = mcp_tool_arguments(&tool_call).expect_err("string arguments should be invalid"); + let content: Vec<&str> = error + .content + .iter() + .filter_map(|content| content.raw.as_text().map(|text| text.text.as_str())) + .collect::>(); + + assert_eq!(error.is_error, Some(true)); + assert_eq!(content.first(), Some(&"INVALID_ARGUMENTS")); + assert!( + content + .get(1) + .is_some_and(|message| message.contains("must be a JSON object")) + ); + } +} diff --git a/tui/src/services/message.rs b/tui/src/services/message.rs index f9290154e..3a986ef59 100644 --- a/tui/src/services/message.rs +++ b/tui/src/services/message.rs @@ -2508,11 +2508,15 @@ pub fn extract_full_command_arguments(tool_call: &ToolCall) -> String { return "Ask user questions...".to_string(); } - if !tool_call.arguments.is_null() { - return format_json_value(&tool_call.arguments); - } + let raw_arguments = match &tool_call.arguments { + Value::String(raw_arguments) => raw_arguments.as_str(), + Value::Null => "", + _ => return format_json_value(&tool_call.arguments), + }; - let raw_arguments = tool_call.arguments.as_str().unwrap_or(""); + if let Ok(value) = serde_json::from_str::(raw_arguments) { + return format_json_value(&value); + } // If arguments are a malformed JSON string, try regex patterns let patterns = vec![ @@ -2794,34 +2798,44 @@ mod tests { use super::*; #[test] - fn test_various_formats() { - // Test cases based on your examples - let test_cases = vec![ - (r#"{"path":"."}"#, "path=."), - (r#"{"confidence":1.0}"#, "confidence=1.0"), - (r#"{"command":"ls -la"}"#, "command=ls -la"), - ( - r#"{"action":"view","target":"file.txt"}"#, - "action=view, target=file.txt", - ), - (r#"path: ".", mode: "list""#, "path=., mode=list"), - ("", "function_name=test"), - ]; - - for (input, expected) in test_cases { - let tool_call = ToolCall { - id: "test".to_string(), - name: "test".to_string(), - arguments: serde_json::from_str(input) - .unwrap_or_else(|_| serde_json::Value::String(input.to_string())), - metadata: None, - }; + fn extract_full_command_arguments_recovers_from_string_arguments() { + let tool_call = ToolCall { + id: "test".to_string(), + name: "test".to_string(), + arguments: serde_json::Value::String(r#"path: ".", mode: "list""#.to_string()), + metadata: None, + }; - let result = extract_full_command_arguments(&tool_call); - println!( - "Input: '{}' -> Output: '{}' (Expected: '{}')", - input, result, expected - ); - } + let result = extract_full_command_arguments(&tool_call); + + assert_eq!(result, "path = ., mode = list"); + } + + #[test] + fn extract_full_command_arguments_parses_json_string_arguments() { + let tool_call = ToolCall { + id: "test".to_string(), + name: "test".to_string(), + arguments: serde_json::Value::String(r#"{"files":["a.txt","b.txt"]}"#.to_string()), + metadata: None, + }; + + let result = extract_full_command_arguments(&tool_call); + + assert_eq!(result, "files = [a.txt, b.txt]"); + } + + #[test] + fn extract_full_command_arguments_formats_object_arguments_directly() { + let tool_call = ToolCall { + id: "test".to_string(), + name: "test".to_string(), + arguments: serde_json::json!({"command": "ls -la"}), + metadata: None, + }; + + let result = extract_full_command_arguments(&tool_call); + + assert_eq!(result, "command = ls -la"); } } From 00598287e6fbe4f448bcd1c8c3fb86bbad91e4fd Mon Sep 17 00:00:00 2001 From: Shehab Atia <89648315+shehab299@users.noreply.github.com> Date: Tue, 26 May 2026 09:35:50 +0300 Subject: [PATCH 4/5] fix(ai) don't treat empty messages as ContentParts Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- libs/api/src/client/provider.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/libs/api/src/client/provider.rs b/libs/api/src/client/provider.rs index 45eb4e4a3..fa9c76c4b 100644 --- a/libs/api/src/client/provider.rs +++ b/libs/api/src/client/provider.rs @@ -957,7 +957,10 @@ impl AgentClient { let message = if tool_calls.is_empty() { Message::new(stakai::Role::Assistant, text) } else { - let mut parts = vec![ContentPart::text(text)]; + let mut parts: Vec = Vec::new(); + if !text.is_empty() { + parts.push(ContentPart::text(text)); + } parts.extend(tool_calls); Message::new(stakai::Role::Assistant, parts) }; From 7fba64af65d3eaa0baf5fd5b2d4bdb4a59f2c5d9 Mon Sep 17 00:00:00 2001 From: Ahmed Hesham Abdelkader <23265119+ahmedhesham6@users.noreply.github.com> Date: Wed, 3 Jun 2026 13:13:01 +0300 Subject: [PATCH 5/5] fixes --- cli/src/commands/acp/server.rs | 6 +- cli/src/commands/agent/run/checkpoint.rs | 62 ++++++- cli/src/commands/agent/run/stream.rs | 51 +++++- libs/api/src/client/provider.rs | 168 +++++++++++++++++- libs/api/src/lib.rs | 1 + libs/api/src/message_compat.rs | 211 +++++++++++++++++++++++ libs/api/src/stakpak/models.rs | 5 +- libs/api/src/storage.rs | 84 ++++++++- 8 files changed, 567 insertions(+), 21 deletions(-) create mode 100644 libs/api/src/message_compat.rs diff --git a/cli/src/commands/acp/server.rs b/cli/src/commands/acp/server.rs index 890c942b4..1187194d1 100644 --- a/cli/src/commands/acp/server.rs +++ b/cli/src/commands/acp/server.rs @@ -1,5 +1,5 @@ use crate::commands::agent::run::helpers::{system_message, user_message}; -use crate::commands::agent::run::stream::ToolCallAccumulator; +use crate::commands::agent::run::stream::{ToolCallAccumulator, tool_call_content_part}; use crate::config::AppConfig; use agent_client_protocol::{ self as acp, Client as AcpClient, ModelInfo, SessionModelState, SessionNotification, @@ -1305,9 +1305,7 @@ impl StakpakAcpAgent { if !text.is_empty() { parts.push(ContentPart::text(text)); } - parts.extend(final_tool_calls.into_iter().map(|tool_call| { - ContentPart::tool_call(tool_call.id, tool_call.name, tool_call.arguments) - })); + parts.extend(final_tool_calls.into_iter().map(tool_call_content_part)); Message::new(Role::Assistant, MessageContent::Parts(parts)) }; diff --git a/cli/src/commands/agent/run/checkpoint.rs b/cli/src/commands/agent/run/checkpoint.rs index 39870e8a5..d0bc18a80 100644 --- a/cli/src/commands/agent/run/checkpoint.rs +++ b/cli/src/commands/agent/run/checkpoint.rs @@ -47,11 +47,7 @@ pub async fn extract_checkpoint_messages_and_tool_calls( .find(|message| message.role != Role::User && message.role != Role::Tool) && last_message.role == Role::Assistant { - last_message.content = MessageContent::Text(format!( - "{}\n{}", - last_message.content.text().unwrap_or_default(), - checkpoint_id - )); + append_checkpoint_id_to_message(last_message, checkpoint_id); } for message in &checkpoint_messages { @@ -169,6 +165,22 @@ pub async fn extract_checkpoint_messages_and_tool_calls( Ok((checkpoint_messages, pending_tool_calls)) } +fn append_checkpoint_id_to_message(message: &mut Message, checkpoint_id: &str) { + let checkpoint_tag = format!("{checkpoint_id}"); + + match &mut message.content { + MessageContent::Text(text) => { + if !text.is_empty() { + text.push('\n'); + } + text.push_str(&checkpoint_tag); + } + MessageContent::Parts(parts) => { + parts.push(ContentPart::text(format!("\n{checkpoint_tag}"))); + } + } +} + pub fn extract_checkpoint_id_from_messages(messages: &[Message]) -> Option { messages .last() @@ -330,6 +342,46 @@ mod tests { } } + #[tokio::test] + async fn extract_checkpoint_messages_preserves_pending_tool_calls_when_injecting_checkpoint_id() + { + let (input_tx, mut input_rx) = mpsc::channel(8); + let messages = vec![Message::new( + Role::Assistant, + MessageContent::Parts(vec![ + ContentPart::text("I need to inspect the file."), + ContentPart::tool_call("tc_1", "test_tool", json!({"path": "README.md"})), + ]), + )]; + + let (checkpoint_messages, pending_tool_calls) = + extract_checkpoint_messages_and_tool_calls("checkpoint-id", &input_tx, messages) + .await + .expect("checkpoint extraction"); + + assert_eq!(pending_tool_calls.len(), 1); + assert_eq!(pending_tool_calls[0].id, "tc_1"); + + let parts = checkpoint_messages[0].parts(); + assert!( + parts + .iter() + .any(|part| matches!(part, ContentPart::ToolCall { id, .. } if id == "tc_1")) + ); + assert!(checkpoint_messages[0].text().is_some_and(|text| { + text.contains("I need to inspect the file.") + && text.contains("checkpoint-id") + })); + + match input_rx.recv().await.expect("input event") { + InputEvent::StreamAssistantMessage(_, content) => { + assert!(content.contains("I need to inspect the file.")); + assert!(content.contains("checkpoint-id")); + } + event => panic!("unexpected input event: {event:?}"), + } + } + #[tokio::test] async fn extract_checkpoint_messages_replays_cancelled_tool_result_status() { let (input_tx, mut input_rx) = mpsc::channel(8); diff --git a/cli/src/commands/agent/run/stream.rs b/cli/src/commands/agent/run/stream.rs index dd6b5668b..ef1de00c6 100644 --- a/cli/src/commands/agent/run/stream.rs +++ b/cli/src/commands/agent/run/stream.rs @@ -119,6 +119,24 @@ impl ToolCallAccumulator { } } +pub(crate) fn tool_call_content_part(tool_call: ToolCall) -> ContentPart { + let ToolCall { + id, + name, + arguments, + metadata, + } = tool_call; + let mut part = ContentPart::tool_call(id, name, arguments); + if let ContentPart::ToolCall { + metadata: part_metadata, + .. + } = &mut part + { + *part_metadata = metadata; + } + part +} + pub async fn process_responses_stream( stream: impl Stream>, input_tx: &tokio::sync::mpsc::Sender, @@ -226,9 +244,7 @@ pub async fn process_responses_stream( if !text.is_empty() { parts.push(ContentPart::text(text)); } - parts.extend(final_tool_calls.into_iter().map(|tool_call| { - ContentPart::tool_call(tool_call.id, tool_call.name, tool_call.arguments) - })); + parts.extend(final_tool_calls.into_iter().map(tool_call_content_part)); Message::new(stakai::Role::Assistant, MessageContent::Parts(parts)) }; @@ -331,4 +347,33 @@ mod tests { serde_json::Value::String("{\"key\":\"value\"}".to_string()) ); } + + #[tokio::test] + async fn preserves_tool_call_metadata_in_final_assistant_message() { + let (input_tx, mut input_rx) = tokio::sync::mpsc::channel(100); + let metadata = serde_json::json!({"thought_signature": "sig-123"}); + let responses = vec![event(StreamEvent::ToolCallEnd { + id: "tool-1".to_string(), + name: "my_function".to_string(), + arguments: serde_json::json!({"key": "value"}), + metadata: Some(metadata.clone()), + })]; + + let test_stream = stream::iter(responses); + tokio::spawn(async move { while input_rx.recv().await.is_some() {} }); + + let response = process_responses_stream(test_stream, &input_tx) + .await + .unwrap(); + let tool_call_metadata = response + .message + .parts() + .into_iter() + .find_map(|part| match part { + ContentPart::ToolCall { metadata, .. } => metadata, + _ => None, + }); + + assert_eq!(tool_call_metadata, Some(metadata)); + } } diff --git a/libs/api/src/client/provider.rs b/libs/api/src/client/provider.rs index 45eb4e4a3..9ff104a0c 100644 --- a/libs/api/src/client/provider.rs +++ b/libs/api/src/client/provider.rs @@ -656,9 +656,7 @@ fn response_to_message(response: &stakai::GenerateResponse) -> Message { for content in &response.content { match content { ResponseContent::Text { text } => parts.push(ContentPart::text(text.clone())), - ResponseContent::Reasoning { reasoning } => { - parts.push(ContentPart::text(format!("[Reasoning: {reasoning}]"))); - } + ResponseContent::Reasoning { .. } => {} ResponseContent::ToolCall(tool_call) => { let mut part = ContentPart::tool_call( tool_call.id.clone(), @@ -673,6 +671,10 @@ fn response_to_message(response: &stakai::GenerateResponse) -> Message { } } + if parts.is_empty() { + return Message::new(stakai::Role::Assistant, ""); + } + if parts.len() == 1 && let ContentPart::Text { text, .. } = &parts[0] { @@ -945,9 +947,11 @@ impl AgentClient { } => { usage = final_usage.clone(); } - StreamEvent::Start { .. } - | StreamEvent::ReasoningDelta { .. } - | StreamEvent::Error { .. } => {} + StreamEvent::Error { message } => { + let _ = tx.send(Err(message.clone())).await; + return Err(message.clone()); + } + StreamEvent::Start { .. } | StreamEvent::ReasoningDelta { .. } => {} } if tx.send(Ok(StreamMessage::Event(event))).await.is_err() { break; @@ -957,7 +961,10 @@ impl AgentClient { let message = if tool_calls.is_empty() { Message::new(stakai::Role::Assistant, text) } else { - let mut parts = vec![ContentPart::text(text)]; + let mut parts = Vec::new(); + if !text.is_empty() { + parts.push(ContentPart::text(text)); + } parts.extend(tool_calls); Message::new(stakai::Role::Assistant, parts) }; @@ -1045,3 +1052,150 @@ impl AgentClient { Ok(response.text()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::client::stakai::StakAIClient; + use crate::local::storage::LocalStorage; + use futures_util::stream; + use stakai::provider::Provider; + use stakai::{ + FinishReason, GenerateRequest, GenerateResponse, GenerateStream, Headers, MessageContent, + ProviderRegistry, Role, Usage, + }; + use stakpak_shared::hooks::HookContext; + use std::sync::Arc; + + #[derive(Clone)] + struct FakeProvider { + events: Vec, + response: GenerateResponse, + } + + #[async_trait::async_trait] + impl Provider for FakeProvider { + fn provider_id(&self) -> &str { + "fake" + } + + fn build_headers(&self, _custom_headers: Option<&Headers>) -> Headers { + Headers::new() + } + + async fn generate(&self, _request: GenerateRequest) -> stakai::Result { + Ok(self.response.clone()) + } + + async fn stream(&self, _request: GenerateRequest) -> stakai::Result { + let events = self.events.clone().into_iter().map(Ok::<_, stakai::Error>); + Ok(GenerateStream::new(Box::pin(stream::iter(events)))) + } + } + + fn generate_response(content: Vec) -> GenerateResponse { + GenerateResponse { + content, + usage: Usage::default(), + finish_reason: FinishReason::stop(), + metadata: None, + warnings: None, + } + } + + async fn agent_client_with_fake_provider( + events: Vec, + response: GenerateResponse, + ) -> AgentClient { + let registry = ProviderRegistry::new().register("fake", FakeProvider { events, response }); + let stakai = StakAIClient::with_registry(registry).expect("fake stakai client"); + let session_storage = Arc::new( + LocalStorage::new(":memory:") + .await + .expect("in-memory local storage"), + ); + + AgentClient { + stakai, + stakpak_api: None, + session_storage, + hook_registry: Arc::new(stakpak_shared::hooks::HookRegistry::::default()), + stakpak: None, + } + } + + async fn run_completion_with_stream_events( + events: Vec, + ) -> Result { + let client = agent_client_with_fake_provider(events, generate_response(Vec::new())).await; + let model = Model::custom("fake-model", "fake"); + let messages = vec![Message::new(Role::User, "hello")]; + let mut state = AgentState::new(model.clone(), messages.clone(), None, None); + state.set_llm_input(Some(GenerateRequest::new(model, messages))); + let mut ctx = HookContext::new(None, state); + let (tx, _rx) = mpsc::channel(16); + + client.run_agent_completion(&mut ctx, Some(tx)).await + } + + #[test] + fn response_to_message_drops_reasoning_content() { + let response = generate_response(vec![ + ResponseContent::Text { + text: "visible".to_string(), + }, + ResponseContent::Reasoning { + reasoning: "hidden chain of thought".to_string(), + }, + ]); + + let message = response_to_message(&response); + + assert_eq!(message.text().as_deref(), Some("visible")); + } + + #[test] + fn response_to_message_drops_reasoning_only_content() { + let response = generate_response(vec![ResponseContent::Reasoning { + reasoning: "hidden chain of thought".to_string(), + }]); + + let message = response_to_message(&response); + + assert_eq!(message.text().as_deref(), Some("")); + } + + #[tokio::test] + async fn streaming_completion_treats_stream_error_event_as_failure() { + let result = run_completion_with_stream_events(vec![StreamEvent::Error { + message: "provider overloaded".to_string(), + }]) + .await; + + assert!(result.is_err()); + assert!( + result + .err() + .is_some_and(|err| err.contains("provider overloaded")) + ); + } + + #[tokio::test] + async fn streaming_completion_omits_empty_text_part_for_tool_only_messages() { + let message = run_completion_with_stream_events(vec![StreamEvent::ToolCallEnd { + id: "tc_1".to_string(), + name: "test_tool".to_string(), + arguments: serde_json::json!({"path": "README.md"}), + metadata: None, + }]) + .await + .expect("stream completion"); + + let MessageContent::Parts(parts) = message.content else { + panic!("expected parts content"); + }; + + assert_eq!(parts.len(), 1); + assert!(matches!(parts[0], ContentPart::ToolCall { .. })); + } +} diff --git a/libs/api/src/lib.rs b/libs/api/src/lib.rs index e20b2a4ad..b49a4a25c 100644 --- a/libs/api/src/lib.rs +++ b/libs/api/src/lib.rs @@ -8,6 +8,7 @@ use uuid::Uuid; pub mod client; pub mod commands; pub mod local; +mod message_compat; pub mod models; pub mod stakpak; pub mod storage; diff --git a/libs/api/src/message_compat.rs b/libs/api/src/message_compat.rs new file mode 100644 index 000000000..8d54786e1 --- /dev/null +++ b/libs/api/src/message_compat.rs @@ -0,0 +1,211 @@ +use serde::{Deserialize, Deserializer}; +use serde_json::Value; +use stakai::{ContentPart, ImageDetail, Message, MessageContent, Role}; + +pub(crate) fn deserialize_messages<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let values = Option::>::deserialize(deserializer)?.unwrap_or_default(); + values + .into_iter() + .map(deserialize_message) + .collect::, _>>() + .map_err(serde::de::Error::custom) +} + +fn deserialize_message(value: Value) -> Result { + if requires_legacy_migration(&value) { + return migrate_legacy_message(&value) + .ok_or_else(|| "Invalid legacy checkpoint message".to_string()); + } + + serde_json::from_value::(value.clone()) + .or_else(|error| migrate_legacy_message(&value).ok_or_else(|| error.to_string())) +} + +fn requires_legacy_migration(value: &Value) -> bool { + let Some(obj) = value.as_object() else { + return false; + }; + + obj.contains_key("tool_calls") + || obj.contains_key("tool_call_id") + || obj.get("role").and_then(Value::as_str) == Some("developer") + || obj + .get("content") + .is_some_and(content_contains_legacy_parts) +} + +fn content_contains_legacy_parts(content: &Value) -> bool { + let Value::Array(parts) = content else { + return false; + }; + + parts.iter().any(|part| { + let Some(part) = part.as_object() else { + return false; + }; + part.contains_key("image_url") + || part.get("type").and_then(Value::as_str) == Some("image_url") + }) +} + +fn migrate_legacy_message(value: &Value) -> Option { + let obj = value.as_object()?; + let role = legacy_role( + obj.get("role") + .and_then(Value::as_str) + .unwrap_or("assistant"), + ); + let name = obj + .get("name") + .and_then(Value::as_str) + .map(ToString::to_string); + + if role == Role::Tool + && let Some(tool_call_id) = obj.get("tool_call_id").and_then(Value::as_str) + { + let mut message = Message::new( + role, + MessageContent::Parts(vec![ContentPart::tool_result( + tool_call_id, + legacy_tool_result_content(obj.get("content")), + )]), + ); + message.name = name; + return Some(message); + } + + let mut parts = legacy_content_parts(obj.get("content")); + let legacy_tool_calls = legacy_tool_calls(obj.get("tool_calls")); + let has_legacy_tool_calls = !legacy_tool_calls.is_empty(); + parts.extend(legacy_tool_calls); + + let content = match parts.as_slice() { + [] => MessageContent::Text(String::new()), + [ContentPart::Text { text, .. }] if !has_legacy_tool_calls => { + MessageContent::Text(text.clone()) + } + _ => MessageContent::Parts(parts), + }; + + let mut message = Message::new(role, content); + message.name = name; + Some(message) +} + +fn legacy_role(role: &str) -> Role { + match role { + "system" => Role::System, + "developer" => Role::System, + "user" => Role::User, + "assistant" => Role::Assistant, + "tool" => Role::Tool, + _ => Role::User, + } +} + +fn legacy_content_parts(content: Option<&Value>) -> Vec { + match content { + Some(Value::String(text)) if !text.is_empty() => vec![ContentPart::text(text.clone())], + Some(Value::Array(parts)) => parts.iter().filter_map(legacy_content_part).collect(), + _ => Vec::new(), + } +} + +fn legacy_content_part(part: &Value) -> Option { + if let Ok(part) = serde_json::from_value::(part.clone()) { + return Some(part); + } + + let obj = part.as_object()?; + + if let Some(text) = obj.get("text").and_then(Value::as_str) { + return Some(ContentPart::text(text.to_string())); + } + + let image_url = obj.get("image_url")?.as_object()?; + let url = image_url.get("url")?.as_str()?.to_string(); + let detail = image_url + .get("detail") + .and_then(Value::as_str) + .and_then(legacy_image_detail); + + Some(ContentPart::Image { + url, + detail, + provider_options: None, + }) +} + +fn legacy_image_detail(detail: &str) -> Option { + match detail { + "low" => Some(ImageDetail::Low), + "high" => Some(ImageDetail::High), + "auto" => Some(ImageDetail::Auto), + _ => None, + } +} + +fn legacy_tool_calls(tool_calls: Option<&Value>) -> Vec { + let Some(Value::Array(tool_calls)) = tool_calls else { + return Vec::new(); + }; + + tool_calls + .iter() + .filter_map(|tool_call| { + let obj = tool_call.as_object()?; + let id = obj.get("id").and_then(Value::as_str)?.to_string(); + let function = obj.get("function")?.as_object()?; + let name = function + .get("name") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); + let arguments = legacy_tool_arguments(function.get("arguments")); + let metadata = obj.get("metadata").cloned(); + + let mut part = ContentPart::tool_call(id, name, arguments); + if let ContentPart::ToolCall { + metadata: part_metadata, + .. + } = &mut part + { + *part_metadata = metadata; + } + Some(part) + }) + .collect() +} + +fn legacy_tool_arguments(arguments: Option<&Value>) -> Value { + match arguments { + Some(Value::String(raw)) => { + serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.clone())) + } + Some(value) => value.clone(), + None => Value::Object(serde_json::Map::new()), + } +} + +fn legacy_tool_result_content(content: Option<&Value>) -> Value { + match content { + Some(Value::String(text)) => Value::String(text.clone()), + Some(Value::Array(parts)) => { + let text = parts + .iter() + .filter_map(|part| { + part.as_object() + .and_then(|obj| obj.get("text")) + .and_then(Value::as_str) + }) + .collect::>() + .join("\n"); + Value::String(text) + } + Some(Value::Null) | None => Value::String(String::new()), + Some(value) => value.clone(), + } +} diff --git a/libs/api/src/stakpak/models.rs b/libs/api/src/stakpak/models.rs index 54ee541f2..707bd5b92 100644 --- a/libs/api/src/stakpak/models.rs +++ b/libs/api/src/stakpak/models.rs @@ -88,7 +88,10 @@ pub struct CheckpointSummary { /// Checkpoint state containing messages #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct CheckpointState { - #[serde(default)] + #[serde( + default, + deserialize_with = "crate::message_compat::deserialize_messages" + )] pub messages: Vec, /// Optional metadata for context trimming state, etc. #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/libs/api/src/storage.rs b/libs/api/src/storage.rs index dd5a043ef..e2cff2376 100644 --- a/libs/api/src/storage.rs +++ b/libs/api/src/storage.rs @@ -283,7 +283,10 @@ pub struct CheckpointSummary { /// Checkpoint state containing messages #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct CheckpointState { - #[serde(default)] + #[serde( + default, + deserialize_with = "crate::message_compat::deserialize_messages" + )] pub messages: Vec, /// Optional metadata for context trimming state, etc. #[serde(default, skip_serializing_if = "Option::is_none")] @@ -494,3 +497,82 @@ pub struct ToolUsageCounts { pub failed: u32, pub aborted: u32, } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn checkpoint_state_deserializes_legacy_assistant_tool_calls() { + let state: CheckpointState = serde_json::from_value(json!({ + "messages": [ + { + "role": "assistant", + "content": "I will inspect the file.", + "tool_calls": [ + { + "id": "tc_1", + "type": "function", + "function": { + "name": "view", + "arguments": "{\"path\":\"README.md\"}" + }, + "metadata": { + "thought_signature": "sig-123" + } + } + ] + } + ] + })) + .expect("legacy checkpoint state should deserialize"); + + let parts = state.messages[0].parts(); + assert!( + parts + .iter() + .any(|part| matches!(part, stakai::ContentPart::Text { text, .. } if text == "I will inspect the file.")) + ); + assert!(parts.iter().any(|part| { + matches!( + part, + stakai::ContentPart::ToolCall { + id, + name, + arguments, + metadata, + .. + } if id == "tc_1" + && name == "view" + && arguments == &json!({"path": "README.md"}) + && metadata.as_ref() == Some(&json!({"thought_signature": "sig-123"})) + ) + })); + } + + #[test] + fn checkpoint_state_deserializes_legacy_tool_result_messages() { + let state: CheckpointState = serde_json::from_value(json!({ + "messages": [ + { + "role": "tool", + "content": "file contents", + "tool_call_id": "tc_1" + } + ] + })) + .expect("legacy checkpoint state should deserialize"); + + let parts = state.messages[0].parts(); + assert_eq!(parts.len(), 1); + assert!(matches!( + &parts[0], + stakai::ContentPart::ToolResult { + tool_call_id, + content, + .. + } if tool_call_id == "tc_1" && content == "file contents" + )); + } +}