diff --git a/libs/agent/chat/mod.rs b/libs/agent/chat/mod.rs index a22c499..362aa8d 100644 --- a/libs/agent/chat/mod.rs +++ b/libs/agent/chat/mod.rs @@ -8,6 +8,7 @@ use models::projects::project; use models::repos::repo; use models::rooms::{room, room_message}; use models::users::user; +use config::AppConfig; use std::collections::HashMap; use uuid::Uuid; @@ -29,6 +30,7 @@ pub type StreamCallback = Box< pub struct AiRequest { pub db: AppDatabase, pub cache: AppCache, + pub config: AppConfig, pub model: model::Model, pub project: project::Model, pub sender: user::Model, diff --git a/libs/agent/chat/service.rs b/libs/agent/chat/service.rs index 0230782..5b47ad3 100644 --- a/libs/agent/chat/service.rs +++ b/libs/agent/chat/service.rs @@ -321,13 +321,6 @@ impl ChatService { ); if has_tool_calls && tools_enabled { - // Send final text chunk - on_chunk(AiStreamChunk { - content: text_accumulated.clone(), - done: true, - }) - .await; - // Build ToolCall list from accumulated chunks let tool_calls: Vec<_> = tool_call_chunks .into_iter() @@ -340,12 +333,23 @@ impl ChatService { .collect(); if !tool_calls.is_empty() { + // Capture thinking text, send it as a completed chunk, then clear for the next turn + let thinking_text = text_accumulated.clone(); + if !thinking_text.is_empty() { + on_chunk(AiStreamChunk { + content: thinking_text.clone(), + done: true, + }) + .await; + } + text_accumulated.clear(); + // Append assistant message with tool calls to history messages.push(ChatCompletionRequestMessage::Assistant( ChatCompletionRequestAssistantMessage { content: Some( ChatCompletionRequestAssistantMessageContent::Text( - text_accumulated, + thinking_text, ), ), name: None, @@ -400,6 +404,7 @@ impl ChatService { let mut ctx = ToolContext::new( request.db.clone(), request.cache.clone(), + request.config.clone(), request.room.id, Some(request.sender.uid), ) diff --git a/libs/agent/tool/context.rs b/libs/agent/tool/context.rs index 220299f..7b993e3 100644 --- a/libs/agent/tool/context.rs +++ b/libs/agent/tool/context.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use db::cache::AppCache; use db::database::AppDatabase; +use config::AppConfig; use uuid::Uuid; use super::registry::ToolRegistry; @@ -21,6 +22,7 @@ pub struct ToolContext { struct Inner { pub db: AppDatabase, pub cache: AppCache, + pub config: AppConfig, pub room_id: Uuid, pub sender_id: Option, pub project_id: Uuid, @@ -32,11 +34,18 @@ struct Inner { } impl ToolContext { - pub fn new(db: AppDatabase, cache: AppCache, room_id: Uuid, sender_id: Option) -> Self { + pub fn new( + db: AppDatabase, + cache: AppCache, + config: AppConfig, + room_id: Uuid, + sender_id: Option, + ) -> Self { Self { inner: Arc::new(Inner { db, cache, + config, room_id, sender_id, project_id: Uuid::nil(), @@ -111,6 +120,11 @@ impl ToolContext { &self.inner.cache } + /// Application config. + pub fn config(&self) -> &AppConfig { + &self.inner.config + } + /// Room where the original message was sent. pub fn room_id(&self) -> Uuid { self.inner.room_id