From 0fbf0f8e1b9f623cc30595c6379e46459e5112d8 Mon Sep 17 00:00:00 2001 From: Revantark Date: Wed, 26 Aug 2026 11:18:14 +0530 Subject: [PATCH 1/5] add content parts to the message --- crates/llm/src/apis/chat_completions/codec.rs | 86 ++++++++++++++++++- crates/llm/src/lib.rs | 2 +- crates/llm/src/message.rs | 36 ++++++++ 3 files changed, 121 insertions(+), 3 deletions(-) diff --git a/crates/llm/src/apis/chat_completions/codec.rs b/crates/llm/src/apis/chat_completions/codec.rs index 5f87cf4..21c626a 100644 --- a/crates/llm/src/apis/chat_completions/codec.rs +++ b/crates/llm/src/apis/chat_completions/codec.rs @@ -25,7 +25,7 @@ struct Request<'a> { #[derive(Serialize)] struct WireMessage { role: &'static str, - content: Option, + content: Option, #[serde(skip_serializing_if = "Option::is_none")] tool_calls: Option>, #[serde(skip_serializing_if = "Option::is_none")] @@ -303,9 +303,17 @@ pub(crate) fn stop_reason_for_finish_reason(reason: Option<&str>) -> Option WireMessage { + let content = if let Some(parts) = &message.content_parts { + Some(serde_json::to_value(parts).expect("content parts must serialize")) + } else { + message + .content + .as_ref() + .map(|s| serde_json::Value::String(s.clone())) + }; WireMessage { role: role(message.role), - content: message.content.clone(), + content, tool_calls: message.tool_calls.as_ref().map(|calls| { calls .iter() @@ -491,4 +499,78 @@ mod tests { assert_eq!(stop_reason_for_finish_reason(Some(reason)), Some(expected)); } } + + #[test] + fn serializes_image_url_with_direct_url() { + use crate::{ContentPart, ImageUrl}; + + let messages = [Message::user_with_parts(vec![ + ContentPart::Text { + text: "What's in this image?".into(), + }, + ContentPart::Image { + image_url: ImageUrl { + url: "https://example.com/photo.jpg".into(), + }, + }, + ])]; + let options = RequestOptions::default(); + let body = serialize_request(&request("model-a", &messages, &[], &options)).unwrap(); + let json: serde_json::Value = serde_json::from_str(&body).unwrap(); + + let content = &json["messages"][0]["content"]; + assert!(content.is_array(), "content must be an array of parts"); + let parts = content.as_array().unwrap(); + assert_eq!(parts.len(), 2); + assert_eq!(parts[0]["type"], "text"); + assert_eq!(parts[0]["text"], "What's in this image?"); + assert_eq!(parts[1]["type"], "image_url"); + assert_eq!( + parts[1]["image_url"]["url"], + "https://example.com/photo.jpg" + ); + } + + #[test] + fn serializes_image_url_with_base64_data_uri() { + use crate::{ContentPart, ImageUrl}; + + let data_uri = "data:image/jpeg;base64,/9j/4AAQSkZJRg=="; + let messages = [Message::user_with_parts(vec![ + ContentPart::Text { + text: "Describe this local image".into(), + }, + ContentPart::Image { + image_url: ImageUrl { + url: data_uri.into(), + }, + }, + ])]; + let options = RequestOptions::default(); + let body = serialize_request(&request("model-a", &messages, &[], &options)).unwrap(); + let json: serde_json::Value = serde_json::from_str(&body).unwrap(); + + let content = &json["messages"][0]["content"]; + let parts = content.as_array().unwrap(); + assert_eq!(parts[1]["image_url"]["url"], data_uri); + } + + #[test] + fn serializes_text_only_content_parts_as_array() { + use crate::ContentPart; + + let messages = [Message::user_with_parts(vec![ContentPart::Text { + text: "hello".into(), + }])]; + let options = RequestOptions::default(); + let body = serialize_request(&request("model-a", &messages, &[], &options)).unwrap(); + let json: serde_json::Value = serde_json::from_str(&body).unwrap(); + + let content = &json["messages"][0]["content"]; + assert!(content.is_array()); + let parts = content.as_array().unwrap(); + assert_eq!(parts.len(), 1); + assert_eq!(parts[0]["type"], "text"); + assert_eq!(parts[0]["text"], "hello"); + } } diff --git a/crates/llm/src/lib.rs b/crates/llm/src/lib.rs index 69324df..71cb142 100644 --- a/crates/llm/src/lib.rs +++ b/crates/llm/src/lib.rs @@ -15,7 +15,7 @@ pub use api::LlmApi; pub use apis::ChatCompletionsApi; pub use error::LlmError; pub use event::{LlmEvent, LlmStream}; -pub use message::{Message, Role}; +pub use message::{ContentPart, ImageUrl, Message, Role}; pub use request::{ CompletionInput, Credential, LlmRequest, PromptCacheControl, PromptCacheControlType, PromptCacheTtl, ReasoningEffort, RequestOptions, diff --git a/crates/llm/src/message.rs b/crates/llm/src/message.rs index 0048e9f..bfd9ba4 100644 --- a/crates/llm/src/message.rs +++ b/crates/llm/src/message.rs @@ -11,11 +11,31 @@ pub enum Role { Tool, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ImageUrl { + pub url: String, +} + +/// A single content part within a multi-part message. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ContentPart { + Text { + text: String, + }, + #[serde(rename = "image_url")] + Image { + image_url: ImageUrl, + }, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Message { pub role: Role, pub content: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub content_parts: Option>, + #[serde(skip_serializing_if = "Option::is_none")] pub tool_calls: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub tool_call_id: Option, @@ -34,6 +54,18 @@ impl Message { Self::text(Role::User, content) } + pub fn user_with_parts(parts: Vec) -> Self { + Self { + role: Role::User, + content: None, + content_parts: (!parts.is_empty()).then_some(parts), + tool_calls: None, + tool_call_id: None, + reasoning: None, + reasoning_details: None, + } + } + pub fn assistant_with_reasoning( content: Option, reasoning: Option, @@ -42,6 +74,7 @@ impl Message { Self { role: Role::Assistant, content, + content_parts: None, tool_calls: None, tool_call_id: None, reasoning, @@ -58,6 +91,7 @@ impl Message { Self { role: Role::Assistant, content, + content_parts: None, tool_calls: Some(tool_calls), tool_call_id: None, reasoning, @@ -69,6 +103,7 @@ impl Message { Self { role: Role::Tool, content: Some(content.into()), + content_parts: None, tool_calls: None, tool_call_id: Some(tool_call_id.into()), reasoning: None, @@ -80,6 +115,7 @@ impl Message { Self { role, content: Some(content.into()), + content_parts: None, tool_calls: None, tool_call_id: None, reasoning: None, From ffefcbd7cd3eea226acb1147d744db66d22ba032 Mon Sep 17 00:00:00 2001 From: Revantark Date: Wed, 26 Aug 2026 16:29:20 +0530 Subject: [PATCH 2/5] split agent.rs and add image selection --- crates/agent/src/agent.rs | 1311 ---------------------- crates/agent/src/agent/builder.rs | 98 ++ crates/agent/src/agent/event.rs | 96 ++ crates/agent/src/agent/mod.rs | 108 ++ crates/agent/src/agent/persistence.rs | 75 ++ crates/agent/src/agent/prompt.rs | 243 ++++ crates/agent/src/agent/prompt_builder.rs | 60 + crates/agent/src/agent/tests.rs | 986 ++++++++++++++++ crates/agent/src/agent/tool_loop.rs | 169 +++ crates/agent/src/context.rs | 43 +- crates/agent/src/lib.rs | 2 +- crates/alan/src/core/chat.rs | 11 +- 12 files changed, 1881 insertions(+), 1321 deletions(-) delete mode 100644 crates/agent/src/agent.rs create mode 100644 crates/agent/src/agent/builder.rs create mode 100644 crates/agent/src/agent/event.rs create mode 100644 crates/agent/src/agent/mod.rs create mode 100644 crates/agent/src/agent/persistence.rs create mode 100644 crates/agent/src/agent/prompt.rs create mode 100644 crates/agent/src/agent/prompt_builder.rs create mode 100644 crates/agent/src/agent/tests.rs create mode 100644 crates/agent/src/agent/tool_loop.rs diff --git a/crates/agent/src/agent.rs b/crates/agent/src/agent.rs deleted file mode 100644 index f6d304b..0000000 --- a/crates/agent/src/agent.rs +++ /dev/null @@ -1,1311 +0,0 @@ -use crate::session::{Session, SessionError, SessionManager, StoreError}; -use crate::{ - AgentError, AgentMessage, AgentTool, Skill, build_system_prompt, context::AgentContext, -}; -use futures_util::StreamExt; -use llm::{ - CompletionInput, LlmEvent, LlmResponse, LlmResponseBuilder, Message, RequestOptions, ToolSpec, - Usage, -}; -use providers::{Model, ModelError}; -use std::{ - path::PathBuf, - sync::{ - Arc, - atomic::{AtomicBool, Ordering}, - }, -}; -use tokio::sync::{ - Mutex, - mpsc::{self, Receiver, Sender}, - watch, -}; - -const AGENT_EVENT_CAPACITY: usize = 128; - -/// Receiver for display-level events emitted by one agent prompt. -pub struct AgentStream { - receiver: Receiver>, - cancellation: watch::Sender, -} - -impl AgentStream { - pub async fn recv(&mut self) -> Option> { - self.receiver.recv().await - } - - pub fn try_recv( - &mut self, - ) -> Result, mpsc::error::TryRecvError> { - self.receiver.try_recv() - } - - pub fn abort(&self) { - let _ = self.cancellation.send(true); - } -} - -impl Drop for AgentStream { - fn drop(&mut self) { - let _ = self.cancellation.send(true); - } -} - -/// Display-level events emitted while an agent prompt runs. -#[derive(Debug, Clone, PartialEq)] -pub enum AgentEvent { - TextDelta(String), - ReasoningDelta(String), - ToolCallStarted { - id: String, - name: String, - arguments: String, - }, - ToolCallFinished { - id: String, - output: String, - }, - ToolCallFailed { - id: String, - error: String, - }, - Finished { - usage: Usage, - }, -} - -pub struct Agent { - model: Mutex, - context: Mutex, - plan_mode: AtomicBool, - max_tool_rounds: usize, - session_id: Mutex, - session_manager: Option>, - active_session: Mutex>, -} - -impl Agent { - pub fn builder(model: Model) -> AgentBuilder { - AgentBuilder { - model, - system_prompt: None, - skills: Vec::new(), - tools: Vec::new(), - max_tool_rounds: 100, - session_manager: None, - resumed_session: None, - } - } - - /// Buffered prompt: runs to completion and returns the final response. - pub async fn prompt(&self, content: impl Into) -> Result { - let content = content.into(); - if content.trim().is_empty() { - return Err(AgentError::Model(ModelError::Llm( - llm::LlmError::Configuration("empty prompt".into()), - ))); - } - - let model = self.model.lock().await; - let mut context = self.context.lock().await; - let plan_mode = self.plan_mode(); - let user_msg = AgentMessage::user(prompt_content(content, plan_mode)); - - self.ensure_session(&model).await?; - - self.append_context_message(&mut context, user_msg).await?; - - let (_cancellation, mut cancellation_receiver) = watch::channel(false); - let mut partial = String::new(); - let result = self - .run_with( - &model, - &mut context, - None, - &mut cancellation_receiver, - &mut partial, - ) - .await; - if result.is_err() && !partial.is_empty() { - let partial_msg = AgentMessage::Assistant(partial_response(&partial)); - self.append_context_message(&mut context, partial_msg) - .await?; - } - result - } - - /// Streaming prompt: returns a channel of incremental events. - /// - /// The agent runs in a background task. Tool calls execute internally - /// without surfacing provider-level fragments to the caller. - pub fn prompt_stream(self: &Arc, content: impl Into) -> AgentStream { - let (tx, receiver) = mpsc::channel(AGENT_EVENT_CAPACITY); - let (cancellation, cancellation_receiver) = watch::channel(false); - let agent = self.clone(); - let content = content.into(); - tokio::spawn(async move { - let result = agent.run_prompt(content, &tx, cancellation_receiver).await; - if let Err(error) = result { - let _ = tx.send(Err(error)).await; - } - }); - AgentStream { - receiver, - cancellation, - } - } - - pub async fn session_id(&self) -> Option { - self.active_session - .lock() - .await - .as_ref() - .map(|session| session.id.clone()) - } - - pub async fn set_model(&self, model: Model) { - *self.model.lock().await = model; - } - - pub fn set_plan_mode(&self, enabled: bool) { - self.plan_mode.store(enabled, Ordering::Release); - } - - pub fn plan_mode(&self) -> bool { - self.plan_mode.load(Ordering::Acquire) - } - - pub async fn messages(&self) -> Vec { - self.context.lock().await.messages.clone() - } - - pub async fn usage(&self) -> Usage { - self.context.lock().await.usage.clone() - } - - async fn run_prompt( - self: Arc, - content: String, - events: &Sender>, - mut cancellation: watch::Receiver, - ) -> Result { - if content.trim().is_empty() { - return Err(AgentError::Model(ModelError::Llm( - llm::LlmError::Configuration("empty prompt".into()), - ))); - } - - let model = self.model.lock().await; - let mut context = self.context.lock().await; - let plan_mode = self.plan_mode(); - let user_msg = AgentMessage::user(prompt_content(content, plan_mode)); - - self.ensure_session(&model).await?; - - self.append_context_message(&mut context, user_msg).await?; - - let mut partial = String::new(); - let result = self - .run_with( - &model, - &mut context, - Some(events), - &mut cancellation, - &mut partial, - ) - .await; - if result.is_err() && !partial.is_empty() { - let partial_msg = AgentMessage::Assistant(partial_response(&partial)); - self.append_context_message(&mut context, partial_msg) - .await?; - } - result - } - - async fn run_with( - &self, - model: &Model, - context: &mut AgentContext, - events: Option<&Sender>>, - cancellation: &mut watch::Receiver, - partial: &mut String, - ) -> Result { - let plan = self.plan_mode(); - for _ in 0..self.max_tool_rounds { - Self::check_cancelled(cancellation)?; - let session_id = self.session_id.lock().await.clone(); - let response = Self::stream_round( - session_id, - model, - context, - events, - cancellation, - partial, - plan, - ) - .await?; - - let calls: Vec<_> = response.tool_calls().cloned().collect(); - if let Some(usage) = response.usage.as_ref() { - context.usage.accumulate(usage); - self.persist_usage(&context.usage).await?; - } - - if calls.is_empty() { - let msg = AgentMessage::Assistant(response.clone()); - self.append_context_message(context, msg).await?; - if let Some(events) = events { - Self::send_event( - events, - Ok(AgentEvent::Finished { - usage: context.usage.clone(), - }), - cancellation, - ) - .await?; - } - return Ok(response); - } - - let assistant_msg = AgentMessage::Assistant(response); - self.append_context_message(context, assistant_msg).await?; - - for call in calls { - Self::check_cancelled(cancellation)?; - let tool_index = context - .tool_indexes - .get(&call.name) - .copied() - .ok_or_else(|| AgentError::ToolNotFound(call.name.clone()))?; - if plan && !context.tools[tool_index].read_only && call.name != "bash" { - return Err(AgentError::ToolNotFound(call.name.clone())); - } - let call_id = call.id.clone(); - if let Some(events) = events { - Self::send_event( - events, - Ok(AgentEvent::ToolCallStarted { - id: call_id.clone(), - name: call.name.clone(), - arguments: call.arguments.clone(), - }), - cancellation, - ) - .await?; - } - - let result = context.tools[tool_index].executor.execute(&call).await; - match result { - Ok(result) => { - let msg = AgentMessage::ToolResult { - tool_call_id: call_id.clone(), - content: result.clone(), - }; - self.append_context_message(context, msg).await?; - - if let Some(events) = events { - Self::send_event( - events, - Ok(AgentEvent::ToolCallFinished { - id: call_id, - output: tail_lines(&result, 5), - }), - cancellation, - ) - .await?; - } - Self::check_cancelled(cancellation)?; - } - Err(error) => { - let error = error.to_string(); - let msg = AgentMessage::ToolResult { - tool_call_id: call_id.clone(), - content: error.clone(), - }; - self.append_context_message(context, msg).await?; - - if let Some(events) = events { - Self::send_event( - events, - Ok(AgentEvent::ToolCallFailed { id: call_id, error }), - cancellation, - ) - .await?; - } - Self::check_cancelled(cancellation)?; - } - } - } - } - Err(AgentError::MaxToolRounds) - } - - async fn stream_round( - session_id: String, - model: &Model, - context: &AgentContext, - events: Option<&Sender>>, - cancellation: &mut watch::Receiver, - partial: &mut String, - plan: bool, - ) -> Result { - let tools: Vec<_> = context - .tools - .iter() - .filter(|tool| !plan || tool.read_only || tool.definition.name == "bash") - .map(|tool| ToolSpec::Function(tool.definition.clone())) - .collect(); - let messages = Self::build_messages(context); - let options = RequestOptions { - prompt_cache_key: Some(session_id.clone()), - session_id: Some(session_id), - ..RequestOptions::default() - }; - let mut stream = model - .stream(CompletionInput { - messages: &messages, - tools: &tools, - options: &options, - }) - .await?; - - let mut builder = LlmResponseBuilder::new(); - partial.clear(); - loop { - let next = tokio::select! { - result = stream.next() => result, - changed = cancellation.changed() => { - changed.map_err(|_| AgentError::Aborted)?; - if *cancellation.borrow() { - return Err(AgentError::Aborted); - } - continue; - } - }; - let Some(event) = next else { break }; - let event = event.map_err(ModelError::from)?; - builder.apply(&event).map_err(ModelError::from)?; - match &event { - LlmEvent::ReasoningDelta { reasoning, .. } => { - if let Some(events) = events { - Self::send_event( - events, - Ok(AgentEvent::ReasoningDelta(reasoning.clone())), - cancellation, - ) - .await?; - } - } - LlmEvent::TextDelta { text } => { - partial.push_str(text); - if let Some(events) = events { - Self::send_event( - events, - Ok(AgentEvent::TextDelta(text.clone())), - cancellation, - ) - .await?; - } - } - _ => {} - } - } - Ok(builder.finish().map_err(ModelError::from)?) - } - - fn check_cancelled(cancellation: &watch::Receiver) -> Result<(), AgentError> { - if *cancellation.borrow() { - Err(AgentError::Aborted) - } else { - Ok(()) - } - } - - async fn send_event( - events: &Sender>, - event: Result, - cancellation: &mut watch::Receiver, - ) -> Result<(), AgentError> { - tokio::select! { - result = events.send(event) => { - result.map_err(|_| AgentError::EventStreamClosed) - } - changed = cancellation.changed() => { - changed.map_err(|_| AgentError::Aborted)?; - Err(AgentError::Aborted) - } - } - } - - fn build_messages(context: &AgentContext) -> Vec { - let system = build_system_prompt(context.system_prompt.as_deref(), &context.skills); - let mut messages = - Vec::with_capacity(context.messages.len() + usize::from(system.is_some())); - if let Some(system) = system { - messages.push(Message::system(system)); - } - messages.extend(context.messages.iter().map(AgentMessage::to_llm)); - messages - } - - async fn ensure_session(&self, model: &Model) -> Result<(), AgentError> { - let mut active_session = self.active_session.lock().await; - if active_session.is_some() { - return Ok(()); - } - - let manager = match &self.session_manager { - Some(m) => m, - None => return Ok(()), - }; - - let pwd = std::env::current_dir().map_err(|e| { - AgentError::Session(SessionError::Store(StoreError::CreateDir { - dir: PathBuf::from("."), - source: e, - })) - })?; - - let session = manager - .create( - pwd, - model.info().provider.0.clone(), - model.info().id.clone(), - model.reasoning_effort(), - ) - .await?; - - *self.session_id.lock().await = session.id.clone(); - *active_session = Some(session); - - Ok(()) - } - - async fn append_context_message( - &self, - context: &mut AgentContext, - message: AgentMessage, - ) -> Result<(), AgentError> { - self.persist_message(&message).await?; - context.messages.push(message); - Ok(()) - } - - async fn persist_message(&self, message: &AgentMessage) -> Result<(), AgentError> { - let active_session = self.active_session.lock().await; - if let (Some(manager), Some(session)) = (&self.session_manager, &*active_session) { - manager - .append_message(&session.id, &session.pwd, message) - .await?; - } - Ok(()) - } - - async fn persist_usage(&self, usage: &Usage) -> Result<(), AgentError> { - let active_session = self.active_session.lock().await; - if let (Some(manager), Some(session)) = (&self.session_manager, &*active_session) { - manager - .append_usage(&session.id, &session.pwd, usage) - .await?; - } - Ok(()) - } -} - -fn prompt_content(content: impl Into, plan_mode: bool) -> String { - let content = content.into(); - if plan_mode { - format!("{content}\n\nPlan mode is on, do not edit any files.") - } else { - content - } -} - -fn partial_response(text: &str) -> LlmResponse { - LlmResponse { - content: vec![llm::ContentBlock::Text(text.to_owned())], - stop_reason: llm::StopReason::Aborted, - usage: None, - model: None, - reasoning: None, - reasoning_details: Vec::new(), - } -} - -fn tail_lines(output: &str, count: usize) -> String { - let mut lines: Vec<_> = output.lines().rev().take(count).collect(); - lines.reverse(); - lines.join("\n") -} - -pub struct AgentBuilder { - model: Model, - system_prompt: Option, - skills: Vec, - tools: Vec, - max_tool_rounds: usize, - session_manager: Option>, - resumed_session: Option, -} - -impl AgentBuilder { - pub fn system_prompt(mut self, prompt: impl Into) -> Self { - self.system_prompt = Some(prompt.into()); - self - } - - pub fn skill(mut self, skill: Skill) -> Self { - self.skills.push(skill); - self - } - - pub fn with_tools(mut self, tools: impl IntoIterator) -> Self { - self.tools.extend(tools); - self - } - - pub fn tool(self, tool: AgentTool) -> Self { - self.with_tools([tool]) - } - - pub fn max_tool_rounds(mut self, rounds: usize) -> Self { - self.max_tool_rounds = rounds; - self - } - - pub fn session_manager(mut self, manager: Arc) -> Self { - self.session_manager = Some(manager); - self - } - - pub fn resume_session(mut self, session: Session) -> Self { - self.resumed_session = Some(session); - self - } - - pub fn build(self) -> Result { - let mut session_id = uuid::Uuid::new_v4().to_string(); - let mut messages = Vec::new(); - let mut usage = Usage::default(); - let mut active_session = None; - - if let Some(session) = self.resumed_session { - if session.provider != self.model.info().provider.0 - || session.model != self.model.info().id - { - return Err(AgentError::Session(SessionError::InvalidHeader { - path: PathBuf::from(&session.id), - reason: format!( - "cannot resume session for model {} (provider {}) with bound model {} (provider {})", - session.model, - session.provider, - self.model.info().id, - self.model.info().provider.0 - ), - })); - } - session_id = session.id.clone(); - messages = session.messages.clone(); - usage = session.usage.clone(); - active_session = Some(session); - } - - let mut context = AgentContext::new(self.system_prompt, self.skills, self.tools); - context.hydrate(messages, usage); - - Ok(Agent { - model: Mutex::new(self.model), - context: Mutex::new(context), - plan_mode: AtomicBool::new(false), - max_tool_rounds: self.max_tool_rounds, - session_id: Mutex::new(session_id), - session_manager: self.session_manager, - active_session: Mutex::new(active_session), - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::session::{SessionManager, SessionRecord}; - use async_trait::async_trait; - use llm::{ContentBlock, LlmApi, LlmError, LlmEvent, StopReason}; - use providers::{ - ApiId, ModelCapabilities, ModelInfo, OpenRouterProvider, Provider, ProviderId, - }; - use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }; - - struct FakeApi; - - #[async_trait] - impl LlmApi for FakeApi { - async fn stream(&self, request: llm::LlmRequest<'_>) -> Result { - let user = request - .messages - .iter() - .rev() - .find_map(|message| message.content.clone()) - .unwrap_or_default(); - let response = LlmResponse { - content: vec![ContentBlock::Text(format!("echo: {user}"))], - stop_reason: StopReason::Stop, - usage: None, - model: Some(request.model_id.to_owned()), - reasoning: None, - reasoning_details: Vec::new(), - }; - let text = response.text(); - let model = response.model.clone(); - Ok(Box::pin(futures_util::stream::iter([ - Ok(llm::LlmEvent::TextDelta { text }), - Ok(llm::LlmEvent::Done { - stop_reason: StopReason::Stop, - usage: None, - model, - }), - ]))) - } - } - - struct FailAfterFirstApi { - calls: AtomicUsize, - } - - struct PendingAfterFirstApi { - calls: AtomicUsize, - } - - #[async_trait] - impl LlmApi for FailAfterFirstApi { - async fn stream(&self, request: llm::LlmRequest<'_>) -> Result { - if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { - return Ok(Box::pin(futures_util::stream::iter([ - Ok(LlmEvent::TextDelta { - text: "first response".into(), - }), - Ok(LlmEvent::Done { - stop_reason: StopReason::Stop, - usage: None, - model: Some(request.model_id.to_owned()), - }), - ]))); - } - - Ok(Box::pin(futures_util::stream::iter([ - Ok(LlmEvent::TextDelta { - text: "partial response".into(), - }), - Err(LlmError::InvalidResponse( - "stream ended without [DONE]".into(), - )), - ]))) - } - } - - #[async_trait] - impl LlmApi for PendingAfterFirstApi { - async fn stream(&self, request: llm::LlmRequest<'_>) -> Result { - if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { - return Ok(Box::pin(futures_util::stream::iter([ - Ok(LlmEvent::TextDelta { - text: "first response".into(), - }), - Ok(LlmEvent::Done { - stop_reason: StopReason::Stop, - usage: None, - model: Some(request.model_id.to_owned()), - }), - ]))); - } - - Ok(Box::pin(futures_util::stream::unfold( - 0, - |state| async move { - match state { - 0 => Some(( - Ok(LlmEvent::TextDelta { - text: "partial response".into(), - }), - 1, - )), - _ => futures_util::future::pending().await, - } - }, - ))) - } - } - - fn model_with_api(api: Arc) -> Model { - let info = ModelInfo { - provider: ProviderId::new("openrouter"), - id: "test".into(), - name: "Test".into(), - api: ApiId::ChatCompletions, - capabilities: ModelCapabilities::default(), - pricing: None, - }; - OpenRouterProvider::builder("key") - .with_models([info]) - .with_api(api) - .build() - .unwrap() - .bind("test") - .unwrap() - } - - fn model() -> Model { - model_with_api(Arc::new(FakeApi)) - } - - #[tokio::test] - async fn prompt_owns_history_and_system_prompt() { - let agent = Agent::builder(model()) - .system_prompt("Be helpful") - .build() - .unwrap(); - let response = agent.prompt("hello").await.unwrap(); - assert_eq!(response.text(), "echo: hello"); - assert_eq!(agent.messages().await.len(), 2); - } - - #[tokio::test] - async fn prompt_stream_emits_text_deltas_and_finished() { - let agent = Arc::new( - Agent::builder(model()) - .system_prompt("Be helpful") - .build() - .unwrap(), - ); - let mut rx = agent.prompt_stream("hello"); - - let mut events = Vec::new(); - while let Some(event) = rx.recv().await { - events.push(event.unwrap()); - } - - assert_eq!( - events, - vec![ - AgentEvent::TextDelta("echo: hello".into()), - AgentEvent::Finished { - usage: Usage::default(), - }, - ] - ); - assert_eq!(agent.messages().await.len(), 2); - } - - struct ReasoningApi; - - #[async_trait] - impl LlmApi for ReasoningApi { - async fn stream(&self, _request: llm::LlmRequest<'_>) -> Result { - Ok(Box::pin(futures_util::stream::iter([ - Ok(llm::LlmEvent::ReasoningDelta { - reasoning: "thinking...".into(), - details: vec![ - serde_json::json!({"type": "reasoning.text", "text": "thinking..."}), - ], - }), - Ok(llm::LlmEvent::TextDelta { - text: "done thinking".into(), - }), - Ok(llm::LlmEvent::Done { - stop_reason: StopReason::Stop, - usage: None, - model: Some("reasoning-model".into()), - }), - ]))) - } - } - - #[tokio::test] - async fn prompt_stream_emits_reasoning_and_text_deltas() { - let agent = Arc::new( - Agent::builder(model_with_api(Arc::new(ReasoningApi))) - .build() - .unwrap(), - ); - let mut rx = agent.prompt_stream("hello"); - - let mut events = Vec::new(); - while let Some(event) = rx.recv().await { - events.push(event.unwrap()); - } - - assert_eq!( - events, - vec![ - AgentEvent::ReasoningDelta("thinking...".into()), - AgentEvent::TextDelta("done thinking".into()), - AgentEvent::Finished { - usage: Usage::default(), - }, - ] - ); - let messages = agent.messages().await; - assert_eq!(messages.len(), 2); - if let AgentMessage::Assistant(resp) = &messages[1] { - assert_eq!(resp.reasoning.as_deref(), Some("thinking...")); - assert_eq!(resp.text(), "done thinking"); - } else { - panic!("expected assistant message"); - } - } - - #[tokio::test] - async fn streaming_error_preserves_previous_history() { - let api = Arc::new(FailAfterFirstApi { - calls: AtomicUsize::new(0), - }); - let agent = Arc::new(Agent::builder(model_with_api(api)).build().unwrap()); - - agent.prompt("first").await.unwrap(); - - let mut stream = agent.prompt_stream("second"); - let mut error = None; - while let Some(event) = stream.recv().await { - if let Err(error_event) = event { - error = Some(error_event); - break; - } - } - - assert_eq!( - error.unwrap().to_string(), - "invalid response: stream ended without [DONE]" - ); - let messages = agent.messages().await; - assert_eq!(messages.len(), 4); - assert!(matches!(&messages[0], AgentMessage::User(text) if text == "first")); - assert!( - matches!(&messages[1], AgentMessage::Assistant(response) if response.text() == "first response") - ); - assert!(matches!(&messages[2], AgentMessage::User(text) if text == "second")); - assert!( - matches!(&messages[3], AgentMessage::Assistant(response) if response.text() == "partial response") - ); - } - - #[tokio::test] - async fn abort_preserves_partial_assistant_output() { - let api = Arc::new(PendingAfterFirstApi { - calls: AtomicUsize::new(0), - }); - let agent = Arc::new(Agent::builder(model_with_api(api)).build().unwrap()); - - agent.prompt("first").await.unwrap(); - - let mut stream = agent.prompt_stream("second"); - assert!(matches!( - stream.recv().await, - Some(Ok(AgentEvent::TextDelta(text))) if text == "partial response" - )); - stream.abort(); - assert!(matches!( - stream.recv().await, - Some(Err(AgentError::Aborted)) - )); - - let messages = agent.messages().await; - assert_eq!(messages.len(), 4); - assert!(matches!(&messages[2], AgentMessage::User(text) if text == "second")); - assert!( - matches!(&messages[3], AgentMessage::Assistant(response) if response.text() == "partial response") - ); - } - - struct ToolCallingApi { - calls: AtomicUsize, - } - - #[async_trait] - impl LlmApi for ToolCallingApi { - async fn stream(&self, request: llm::LlmRequest<'_>) -> Result { - let call_count = self.calls.fetch_add(1, Ordering::SeqCst); - if call_count == 0 { - // First round: one tool call, no text content. - let response = LlmResponse { - content: vec![], - stop_reason: StopReason::ToolUse, - usage: Some(llm::Usage { - input_tokens: 10, - output_tokens: 5, - ..llm::Usage::default() - }), - model: Some(request.model_id.to_owned()), - reasoning: None, - reasoning_details: Vec::new(), - }; - Ok(Box::pin(futures_util::stream::iter([ - Ok(LlmEvent::ToolCallDelta { - index: 0, - id: Some("call-1".into()), - name: Some("bash".into()), - arguments: serde_json::json!({"command": "echo hi"}).to_string(), - }), - Ok(LlmEvent::Done { - stop_reason: StopReason::ToolUse, - usage: response.usage.clone(), - model: response.model.clone(), - }), - ]))) - } else { - // Second round: final text response after tool result. - let response = LlmResponse { - content: vec![ContentBlock::Text("done".into())], - stop_reason: StopReason::Stop, - usage: Some(llm::Usage { - input_tokens: 20, - output_tokens: 3, - ..llm::Usage::default() - }), - model: Some(request.model_id.to_owned()), - reasoning: None, - reasoning_details: Vec::new(), - }; - let text = response.text(); - let model = response.model.clone(); - Ok(Box::pin(futures_util::stream::iter([ - Ok(LlmEvent::TextDelta { text }), - Ok(LlmEvent::Done { - stop_reason: StopReason::Stop, - usage: response.usage.clone(), - model, - }), - ]))) - } - } - } - - #[tokio::test] - async fn building_agent_with_manager_creates_no_file() { - let root = std::env::temp_dir().join(format!("alan-plan3-build-{}", uuid::Uuid::new_v4())); - std::fs::create_dir_all(&root).unwrap(); - let manager = Arc::new(SessionManager::new(&root)); - let _agent = Agent::builder(model()) - .session_manager(manager) - .build() - .unwrap(); - - // No session file should exist before any prompt. - let entries: Vec<_> = std::fs::read_dir(&root) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - assert!( - entries.is_empty(), - "building must not create a session file" - ); - } - - #[tokio::test] - async fn first_buffered_prompt_creates_session_and_persists_messages() { - let root = - std::env::temp_dir().join(format!("alan-plan3-buffered-{}", uuid::Uuid::new_v4())); - std::fs::create_dir_all(&root).unwrap(); - let manager = Arc::new(SessionManager::new(&root)); - let agent = Agent::builder(model()) - .session_manager(manager.clone()) - .build() - .unwrap(); - - let response = agent.prompt("hello").await.unwrap(); - assert_eq!(response.text(), "echo: hello"); - - // One session file must exist under a pwd subdirectory. - let entries: Vec<_> = std::fs::read_dir(&root) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - assert_eq!(entries.len(), 1, "one pwd directory created"); - let pwd_dir = entries[0].path(); - let files: Vec<_> = std::fs::read_dir(&pwd_dir) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - assert_eq!(files.len(), 1, "one session file created"); - - // Load the session and verify messages. - let session_file = files[0].path(); - let content = std::fs::read_to_string(&session_file).unwrap(); - let lines: Vec<_> = content.lines().collect(); - assert_eq!(lines.len(), 3, "header + user + assistant"); - - let header: serde_json::Value = serde_json::from_str(lines[0]).unwrap(); - assert_eq!(header["type"], "session"); - - let user_record: serde_json::Value = serde_json::from_str(lines[1]).unwrap(); - assert_eq!(user_record["type"], "message"); - assert_eq!(user_record["message"]["kind"], "user"); - assert_eq!(user_record["message"]["content"], "hello"); - - let assistant_record: serde_json::Value = serde_json::from_str(lines[2]).unwrap(); - assert_eq!(assistant_record["type"], "message"); - assert_eq!(assistant_record["message"]["kind"], "assistant"); - } - - #[tokio::test] - async fn first_streaming_prompt_has_same_persistence_behavior() { - let root = - std::env::temp_dir().join(format!("alan-plan3-streaming-{}", uuid::Uuid::new_v4())); - std::fs::create_dir_all(&root).unwrap(); - let manager = Arc::new(SessionManager::new(&root)); - let agent = Arc::new( - Agent::builder(model()) - .session_manager(manager.clone()) - .build() - .unwrap(), - ); - - let mut rx = agent.prompt_stream("streaming hello"); - while let Some(event) = rx.recv().await { - let _ = event.unwrap(); - } - - let entries: Vec<_> = std::fs::read_dir(&root) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - assert_eq!(entries.len(), 1); - let pwd_dir = entries[0].path(); - let files: Vec<_> = std::fs::read_dir(&pwd_dir) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - assert_eq!(files.len(), 1); - - let content = std::fs::read_to_string(files[0].path()).unwrap(); - let lines: Vec<_> = content.lines().collect(); - assert_eq!(lines.len(), 3, "header + user + assistant for streaming"); - } - - #[tokio::test] - async fn provider_is_not_called_if_session_creation_fails() { - let _root = - std::env::temp_dir().join(format!("alan-plan3-fail-create-{}", uuid::Uuid::new_v4())); - // Use a non-existent, unwritable path so session creation fails. - let manager = Arc::new(SessionManager::new("/proc/self/mem/unwritable-dir")); - let agent = Agent::builder(model()) - .session_manager(manager) - .build() - .unwrap(); - - let result = agent.prompt("hello").await; - assert!(result.is_err(), "must fail when session creation fails"); - } - - #[tokio::test] - async fn tool_call_responses_and_results_are_persisted_in_order() { - let root = - std::env::temp_dir().join(format!("alan-plan3-tool-order-{}", uuid::Uuid::new_v4())); - std::fs::create_dir_all(&root).unwrap(); - let manager = Arc::new(SessionManager::new(&root)); - let api = Arc::new(ToolCallingApi { - calls: AtomicUsize::new(0), - }); - let agent = Agent::builder(model_with_api(api)) - .session_manager(manager.clone()) - .with_tools([AgentTool::new( - llm::ToolDefinition { - name: "bash".into(), - description: "Run a shell command".into(), - parameters: serde_json::json!({}), - }, - tools::BashExecutor, - )]) - .build() - .unwrap(); - - let response = agent.prompt("run echo hi").await.unwrap(); - assert_eq!(response.text(), "done"); - - let content = { - let entries: Vec<_> = std::fs::read_dir(&root) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - let pwd_dir = entries[0].path(); - let files: Vec<_> = std::fs::read_dir(&pwd_dir) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - std::fs::read_to_string(files[0].path()).unwrap() - }; - let lines: Vec<_> = content.lines().collect(); - // At minimum: header + user + assistant(tool_calls) + tool_result + assistant(final) - assert!( - lines.len() >= 5, - "expected at least 5 records, got {}", - lines.len() - ); - - let types: Vec<_> = lines - .iter() - .map(|l| { - let v: serde_json::Value = serde_json::from_str(l).unwrap(); - v["type"].as_str().unwrap().to_owned() - }) - .collect(); - assert_eq!(types[0], "session"); - assert_eq!(types[1], "message"); - assert_eq!(types[types.len() - 1], "message"); - // The last assistant message should be "done". - let last_assistant = &types[types.len() - 1]; - assert_eq!(last_assistant, "message"); - } - - #[tokio::test] - async fn aggregate_usage_from_multiple_rounds_is_persisted() { - let root = std::env::temp_dir().join(format!("alan-plan3-usage-{}", uuid::Uuid::new_v4())); - std::fs::create_dir_all(&root).unwrap(); - let manager = Arc::new(SessionManager::new(&root)); - let api = Arc::new(ToolCallingApi { - calls: AtomicUsize::new(0), - }); - let agent = Agent::builder(model_with_api(api)) - .session_manager(manager.clone()) - .with_tools([AgentTool::new( - llm::ToolDefinition { - name: "bash".into(), - description: "Run a shell command".into(), - parameters: serde_json::json!({}), - }, - tools::BashExecutor, - )]) - .build() - .unwrap(); - - let response = agent.prompt("run echo hi").await.unwrap(); - assert_eq!(response.text(), "done"); - - let entries: Vec<_> = std::fs::read_dir(&root) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - let pwd_dir = entries[0].path(); - let files: Vec<_> = std::fs::read_dir(&pwd_dir) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - let session_file = files[0].path(); - - // Load the session through the manager to verify usage. - let session = { - let content = std::fs::read_to_string(&session_file).unwrap(); - let lines: Vec<_> = content.lines().collect(); - let header_line = lines[0]; - let record = SessionRecord::parse(header_line).unwrap(); - let SessionRecord::Session { id, .. } = record else { - panic!("expected session header"); - }; - manager - .get_session(&id, &std::env::current_dir().unwrap()) - .await - .expect("load session for usage check") - }; - - // First round: 10 input + 5 output. Second round: 20 input + 3 output. - assert_eq!(session.usage.input_tokens, 30); - assert_eq!(session.usage.output_tokens, 8); - } - - #[tokio::test] - async fn resumed_agent_includes_restored_messages_in_first_request() { - let root = std::env::temp_dir().join(format!("alan-plan3-resume-{}", uuid::Uuid::new_v4())); - std::fs::create_dir_all(&root).unwrap(); - let manager = Arc::new(SessionManager::new(&root)); - - // Create a session, persist a user message, then load it back - // so the in-memory Session has the restored messages. - let session = manager - .create(&root, "openrouter", "test", None) - .await - .expect("create session"); - manager - .append_message( - &session.id, - &session.pwd, - &AgentMessage::user("restored message"), - ) - .await - .expect("append message"); - let session = manager - .get_session(&session.id, &root) - .await - .expect("load session with restored message"); - - let model = model(); - let agent = Agent::builder(model) - .session_manager(manager.clone()) - .resume_session(session) - .build() - .expect("build with resume"); - - let response = agent.prompt("new message").await.unwrap(); - assert_eq!(response.text(), "echo: new message"); - - // The restored message must be in the agent's history. - let messages = agent.messages().await; - assert!( - messages.len() >= 2, - "must include restored user + new user + assistant response, got {} messages", - messages.len() - ); - assert!( - matches!(&messages[0], AgentMessage::User(text) if text == "restored message"), - "first message must be the restored message, got {:?}", - messages[0] - ); - } - - #[tokio::test] - async fn request_uses_persisted_session_id_and_cache_key() { - let root = - std::env::temp_dir().join(format!("alan-plan3-session-id-{}", uuid::Uuid::new_v4())); - std::fs::create_dir_all(&root).unwrap(); - let manager = Arc::new(SessionManager::new(&root)); - let agent = Agent::builder(model()) - .session_manager(manager.clone()) - .build() - .unwrap(); - - agent.prompt("check session id").await.unwrap(); - - let entries: Vec<_> = std::fs::read_dir(&root) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - let pwd_dir = entries[0].path(); - let files: Vec<_> = std::fs::read_dir(&pwd_dir) - .unwrap() - .filter_map(|e| e.ok()) - .collect(); - let content = std::fs::read_to_string(files[0].path()).unwrap(); - let lines: Vec<_> = content.lines().collect(); - - // The header must contain the session id. - let header: serde_json::Value = serde_json::from_str(lines[0]).unwrap(); - let session_id = header["id"].as_str().unwrap(); - - // The prompt_cache_key and session_id in the request options must - // match the persisted session id. We verify this indirectly by - // confirming the session file name matches the session id. - let session_file_name = files[0].path().file_stem().unwrap().to_owned(); - assert_eq!(session_file_name, session_id); - } - - #[tokio::test] - async fn existing_no_manager_tests_still_pass() { - // Re-run the original no-manager buffered prompt test to confirm - // backward compatibility is preserved. - let agent = Agent::builder(model()).build().unwrap(); - let response = agent.prompt("hello").await.unwrap(); - assert_eq!(response.text(), "echo: hello"); - assert_eq!(agent.messages().await.len(), 2); - } -} diff --git a/crates/agent/src/agent/builder.rs b/crates/agent/src/agent/builder.rs new file mode 100644 index 0000000..9c2d5c0 --- /dev/null +++ b/crates/agent/src/agent/builder.rs @@ -0,0 +1,98 @@ +use crate::session::{Session, SessionError, SessionManager}; +use crate::{AgentError, AgentTool, Skill, context::AgentContext}; +use llm::Usage; +use providers::Model; +use std::{ + path::PathBuf, + sync::{Arc, atomic::AtomicBool}, +}; +use tokio::sync::Mutex; + +use super::Agent; + +pub struct AgentBuilder { + pub(super) model: Model, + pub(super) system_prompt: Option, + pub(super) skills: Vec, + pub(super) tools: Vec, + pub(super) max_tool_rounds: usize, + pub(super) session_manager: Option>, + pub(super) resumed_session: Option, +} + +impl AgentBuilder { + pub fn system_prompt(mut self, prompt: impl Into) -> Self { + self.system_prompt = Some(prompt.into()); + self + } + + pub fn skill(mut self, skill: Skill) -> Self { + self.skills.push(skill); + self + } + + pub fn with_tools(mut self, tools: impl IntoIterator) -> Self { + self.tools.extend(tools); + self + } + + pub fn tool(self, tool: AgentTool) -> Self { + self.with_tools([tool]) + } + + pub fn max_tool_rounds(mut self, rounds: usize) -> Self { + self.max_tool_rounds = rounds; + self + } + + pub fn session_manager(mut self, manager: Arc) -> Self { + self.session_manager = Some(manager); + self + } + + pub fn resume_session(mut self, session: Session) -> Self { + self.resumed_session = Some(session); + self + } + + pub fn build(self) -> Result { + let mut session_id = uuid::Uuid::new_v4().to_string(); + let mut messages = Vec::new(); + let mut usage = Usage::default(); + let mut active_session = None; + + if let Some(session) = self.resumed_session { + if session.provider != self.model.info().provider.0 + || session.model != self.model.info().id + { + return Err(AgentError::Session(SessionError::InvalidHeader { + path: PathBuf::from(&session.id), + reason: format!( + "cannot resume session for model {} (provider {}) with bound model {} (provider {})", + session.model, + session.provider, + self.model.info().id, + self.model.info().provider.0 + ), + })); + } + session_id = session.id.clone(); + messages = session.messages.clone(); + usage = session.usage.clone(); + active_session = Some(session); + } + + let mut context = AgentContext::new(self.system_prompt, self.skills, self.tools); + context.hydrate(messages, usage); + + Ok(Agent { + model: Mutex::new(self.model), + context: Mutex::new(context), + plan_mode: AtomicBool::new(false), + max_tool_rounds: self.max_tool_rounds, + session_id: Mutex::new(session_id), + session_manager: self.session_manager, + active_session: Mutex::new(active_session), + }) + } +} diff --git a/crates/agent/src/agent/event.rs b/crates/agent/src/agent/event.rs new file mode 100644 index 0000000..7f38eed --- /dev/null +++ b/crates/agent/src/agent/event.rs @@ -0,0 +1,96 @@ +use crate::AgentError; +use llm::{LlmResponse, Usage}; +use tokio::sync::{ + mpsc::{self, Receiver, Sender}, + watch, +}; + +/// Receiver for display-level events emitted by one agent prompt. +pub struct AgentStream { + pub(super) receiver: Receiver>, + pub(super) cancellation: watch::Sender, +} + +impl AgentStream { + pub async fn recv(&mut self) -> Option> { + self.receiver.recv().await + } + + pub fn try_recv( + &mut self, + ) -> Result, mpsc::error::TryRecvError> { + self.receiver.try_recv() + } + + pub fn abort(&self) { + let _ = self.cancellation.send(true); + } + + /// Drain all remaining events and return the final [`LlmResponse`]. + /// + /// Returns an error if the stream fails or is aborted before producing + /// a [`Finished`](AgentEvent::Finished) event. + pub async fn into_response(mut self) -> Result { + let mut response = None; + while let Some(result) = self.receiver.recv().await { + match result { + Ok(AgentEvent::Finished { response: resp, .. }) => { + response = Some(*resp); + break; + } + Ok(_other) => {} + Err(e) => return Err(e), + } + } + response.ok_or(AgentError::EventStreamClosed) + } +} + +impl Drop for AgentStream { + fn drop(&mut self) { + let _ = self.cancellation.send(true); + } +} + +/// Display-level events emitted while an agent prompt runs. +#[derive(Debug, Clone, PartialEq)] +pub enum AgentEvent { + TextDelta(String), + ReasoningDelta(String), + ToolCallStarted { + id: String, + name: String, + arguments: String, + }, + ToolCallFinished { + id: String, + output: String, + }, + ToolCallFailed { + id: String, + error: String, + }, + Finished { + usage: Usage, + response: Box, + }, +} + +/// Emit a display-level event, aborting if the consumer channel is closed +/// or the prompt has been cancelled. +pub(super) async fn emit_event( + events: Option<&Sender>>, + event: AgentEvent, + cancellation: &mut watch::Receiver, +) -> Result<(), AgentError> { + let Some(events) = events else { return Ok(()) }; + tokio::select! { + result = events.send(Ok(event)) => { + result.map_err(|_| AgentError::EventStreamClosed) + } + changed = cancellation.changed() => { + changed.map_err(|_| AgentError::Aborted)?; + Err(AgentError::Aborted) + } + } +} diff --git a/crates/agent/src/agent/mod.rs b/crates/agent/src/agent/mod.rs new file mode 100644 index 0000000..309296f --- /dev/null +++ b/crates/agent/src/agent/mod.rs @@ -0,0 +1,108 @@ +mod builder; +mod event; +mod persistence; +mod prompt; +mod prompt_builder; +mod tool_loop; + +#[cfg(test)] +mod tests; + +use crate::session::{Session, SessionManager}; +use crate::{AgentError, AgentMessage}; +use llm::Usage; +use providers::Model; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; +use tokio::sync::Mutex; + +pub use builder::AgentBuilder; +pub use event::{AgentEvent, AgentStream}; +pub use prompt_builder::PromptBuilder; + +pub(crate) const AGENT_EVENT_CAPACITY: usize = 128; + +pub struct Agent { + pub(super) model: Mutex, + pub(super) context: Mutex, + pub(super) plan_mode: AtomicBool, + pub(super) max_tool_rounds: usize, + /// Stable identifier used for LLM prompt caching. + /// When no session manager is configured this is a random UUID; + /// once a session is created it matches `active_session.id`. + pub(super) session_id: Mutex, + pub(super) session_manager: Option>, + pub(super) active_session: Mutex>, +} + +impl Agent { + pub fn builder(model: Model) -> AgentBuilder { + AgentBuilder { + model, + system_prompt: None, + skills: Vec::new(), + tools: Vec::new(), + max_tool_rounds: 100, + session_manager: None, + resumed_session: None, + } + } + + /// Start building a prompt request. + /// + /// Returns a [`PromptBuilder`] that can be configured with chained + /// setter calls, then passed to [`ask`](Self::ask) to execute. + pub fn prompt(&self) -> PromptBuilder { + PromptBuilder::new() + } + + /// Execute a prompt request and return an [`AgentStream`] for + /// receiving events. + /// + /// The agent runs the full prompt lifecycle (including tool-call + /// rounds) in a background task. Events are streamed through the + /// returned channel. + /// + /// Use [`AgentStream::into_response`] to drain the stream and + /// extract the final [`LlmResponse`](llm::LlmResponse). + pub fn ask(self: &Arc, builder: PromptBuilder) -> Result { + let content = builder.content.ok_or_else(|| { + AgentError::Model(providers::ModelError::Llm(llm::LlmError::Configuration( + "empty prompt".into(), + ))) + })?; + prompt::validate_not_empty(&content)?; + Ok(prompt::spawn_prompt_task( + self, + content, + builder.images, + builder.stream, + )) + } + + pub async fn session_id(&self) -> Option { + self.active_session + .lock() + .await + .as_ref() + .map(|session| session.id.clone()) + } + + pub fn set_plan_mode(&self, enabled: bool) { + self.plan_mode.store(enabled, Ordering::Release); + } + + pub fn plan_mode(&self) -> bool { + self.plan_mode.load(Ordering::Acquire) + } + + pub async fn messages(&self) -> Vec { + self.context.lock().await.messages.clone() + } + + pub async fn usage(&self) -> Usage { + self.context.lock().await.usage.clone() + } +} diff --git a/crates/agent/src/agent/persistence.rs b/crates/agent/src/agent/persistence.rs new file mode 100644 index 0000000..7fdfccb --- /dev/null +++ b/crates/agent/src/agent/persistence.rs @@ -0,0 +1,75 @@ +use crate::AgentError; +use crate::AgentMessage; +use crate::context::AgentContext; +use crate::session::{SessionError, StoreError}; +use llm::Usage; +use providers::Model; +use std::path::PathBuf; + +use super::Agent; + +/// Ensure a session exists. On the first call this creates and persists a +/// session file; subsequent calls are no-ops. +pub(super) async fn ensure_session(agent: &Agent, model: &Model) -> Result<(), AgentError> { + let mut active_session = agent.active_session.lock().await; + if active_session.is_some() { + return Ok(()); + } + + let manager = match &agent.session_manager { + Some(m) => m, + None => return Ok(()), + }; + + let pwd = std::env::current_dir().map_err(|e| { + AgentError::Session(SessionError::Store(StoreError::CreateDir { + dir: PathBuf::from("."), + source: e, + })) + })?; + + let session = manager + .create( + pwd, + model.info().provider.0.clone(), + model.info().id.clone(), + model.reasoning_effort(), + ) + .await?; + + *agent.session_id.lock().await = session.id.clone(); + *active_session = Some(session); + + Ok(()) +} + +/// Append a message to the in-memory context and persist it to disk. +pub(super) async fn append_context_message( + agent: &Agent, + context: &mut AgentContext, + message: AgentMessage, +) -> Result<(), AgentError> { + persist_message(agent, &message).await?; + context.messages.push(message); + Ok(()) +} + +async fn persist_message(agent: &Agent, message: &AgentMessage) -> Result<(), AgentError> { + let active_session = agent.active_session.lock().await; + if let (Some(manager), Some(session)) = (&agent.session_manager, &*active_session) { + manager + .append_message(&session.id, &session.pwd, message) + .await?; + } + Ok(()) +} + +pub(super) async fn persist_usage(agent: &Agent, usage: &Usage) -> Result<(), AgentError> { + let active_session = agent.active_session.lock().await; + if let (Some(manager), Some(session)) = (&agent.session_manager, &*active_session) { + manager + .append_usage(&session.id, &session.pwd, usage) + .await?; + } + Ok(()) +} diff --git a/crates/agent/src/agent/prompt.rs b/crates/agent/src/agent/prompt.rs new file mode 100644 index 0000000..3d39027 --- /dev/null +++ b/crates/agent/src/agent/prompt.rs @@ -0,0 +1,243 @@ +use super::Agent; +use super::event::{AgentEvent, emit_event}; +use crate::AgentMessage; +use crate::context::AgentContext; +use crate::{AgentError, AgentStream}; +use futures_util::StreamExt; +use llm::{ + CompletionInput, LlmEvent, LlmResponse, LlmResponseBuilder, Message, RequestOptions, ToolSpec, +}; +use providers::{Model, ModelError}; +use tokio::sync::{mpsc::Sender, watch}; + +/// Mutable state threaded through the entire prompt lifecycle. +pub(super) struct PromptCx<'a> { + pub events: Option<&'a Sender>>, + pub cancellation: &'a mut watch::Receiver, + pub partial: &'a mut String, + pub stream: bool, +} + +impl<'a> PromptCx<'a> { + /// Return `Err(Aborted)` if cancellation has been signalled. + pub(super) fn check_cancelled(&self) -> Result<(), AgentError> { + if *self.cancellation.borrow() { + Err(AgentError::Aborted) + } else { + Ok(()) + } + } +} + +pub(super) fn spawn_prompt_task( + agent: &std::sync::Arc, + content: String, + images: Vec, + stream: bool, +) -> AgentStream { + let (tx, receiver) = tokio::sync::mpsc::channel(super::AGENT_EVENT_CAPACITY); + let (cancellation, mut cancellation_receiver) = watch::channel(false); + let agent = agent.clone(); + + tokio::spawn(async move { + let user_msg = build_user_message(content, images, agent.plan_mode()); + let mut partial = String::new(); + let mut cx = PromptCx { + events: Some(&tx), + cancellation: &mut cancellation_receiver, + partial: &mut partial, + stream, + }; + let result = run_prompt_with_message(&agent, user_msg, &mut cx).await; + if let Err(error) = result { + let _ = tx.send(Err(error)).await; + } + }); + + AgentStream { + receiver, + cancellation, + } +} + +/// Run the full prompt lifecycle inside the spawned background task. +async fn run_prompt_with_message( + agent: &Agent, + user_msg: AgentMessage, + cx: &mut PromptCx<'_>, +) -> Result { + validate_prompt_message(&user_msg)?; + + let model = agent.model.lock().await; + let mut context = agent.context.lock().await; + + super::persistence::ensure_session(agent, &model).await?; + super::persistence::append_context_message(agent, &mut context, user_msg).await?; + + let result = super::tool_loop::run_with(agent, &model, &mut context, cx).await; + + save_partial_on_error(agent, &mut context, &result, cx.partial).await?; + result +} + +/// If the prompt failed but produced partial output, persist it so the +/// conversation history remains coherent. +async fn save_partial_on_error( + agent: &Agent, + context: &mut AgentContext, + result: &Result, + partial: &str, +) -> Result<(), AgentError> { + if result.is_err() && !partial.is_empty() { + super::persistence::append_context_message( + agent, + context, + AgentMessage::Assistant(partial_response(partial)), + ) + .await?; + } + Ok(()) +} + +/// Stream a single LLM round: send messages, receive response events, +/// and build the final [`LlmResponse`]. +pub(super) async fn stream_round( + session_id: String, + model: &Model, + context: &AgentContext, + cx: &mut PromptCx<'_>, + plan: bool, +) -> Result { + let tools: Vec<_> = context + .tools + .iter() + .filter(|tool| !plan || tool.read_only || tool.definition.name == "bash") + .map(|tool| ToolSpec::Function(tool.definition.clone())) + .collect(); + + let messages = build_messages(context); + let options = RequestOptions { + prompt_cache_key: Some(session_id.clone()), + session_id: Some(session_id), + ..RequestOptions::default() + }; + + let mut stream_resp = model + .stream(CompletionInput { + messages: &messages, + tools: &tools, + options: &options, + }) + .await?; + + let mut builder = LlmResponseBuilder::new(); + cx.partial.clear(); + + loop { + let next = tokio::select! { + result = stream_resp.next() => result, + changed = cx.cancellation.changed() => { + changed.map_err(|_| AgentError::Aborted)?; + if *cx.cancellation.borrow() { + return Err(AgentError::Aborted); + } + continue; + } + }; + + let Some(event) = next else { break }; + let event = event.map_err(ModelError::from)?; + builder.apply(&event).map_err(ModelError::from)?; + + if cx.stream { + match &event { + LlmEvent::ReasoningDelta { reasoning, .. } => { + emit_event( + cx.events, + AgentEvent::ReasoningDelta(reasoning.clone()), + cx.cancellation, + ) + .await?; + } + LlmEvent::TextDelta { text } => { + cx.partial.push_str(text); + emit_event( + cx.events, + AgentEvent::TextDelta(text.clone()), + cx.cancellation, + ) + .await?; + } + _ => {} + } + } else if let LlmEvent::TextDelta { text } = &event { + // Even in non-streaming mode we must track partial output for + // error recovery, but we suppress the event itself. + cx.partial.push_str(text); + } + } + + Ok(builder.finish().map_err(ModelError::from)?) +} + +fn build_messages(context: &AgentContext) -> Vec { + let system = crate::build_system_prompt(context.system_prompt.as_deref(), &context.skills); + let mut messages = Vec::with_capacity(context.messages.len() + usize::from(system.is_some())); + if let Some(system) = system { + messages.push(Message::system(system)); + } + messages.extend(context.messages.iter().map(AgentMessage::to_llm)); + messages +} + +/// Build a user message, applying the plan-mode suffix and optional images. +pub(super) fn build_user_message( + content: String, + images: Vec, + plan_mode: bool, +) -> AgentMessage { + let text = prompt_content(content, plan_mode); + if images.is_empty() { + AgentMessage::user(text) + } else { + AgentMessage::user_with_images(text, images) + } +} + +fn prompt_content(content: impl Into, plan_mode: bool) -> String { + let content = content.into(); + if plan_mode { + format!("{content}\n\nPlan mode is on, do not edit any files.") + } else { + content + } +} + +/// Validate that a user-facing prompt is non-empty. +pub(super) fn validate_not_empty(text: &str) -> Result<(), AgentError> { + if text.trim().is_empty() { + return Err(AgentError::Model(ModelError::Llm( + llm::LlmError::Configuration("empty prompt".into()), + ))); + } + Ok(()) +} + +/// Validate a user message (used by the streaming path after plan-mode suffix). +fn validate_prompt_message(message: &AgentMessage) -> Result<(), AgentError> { + if let AgentMessage::User { text, .. } = message { + validate_not_empty(text)?; + } + Ok(()) +} + +pub(super) fn partial_response(text: &str) -> LlmResponse { + LlmResponse { + content: vec![llm::ContentBlock::Text(text.to_owned())], + stop_reason: llm::StopReason::Aborted, + usage: None, + model: None, + reasoning: None, + reasoning_details: Vec::new(), + } +} diff --git a/crates/agent/src/agent/prompt_builder.rs b/crates/agent/src/agent/prompt_builder.rs new file mode 100644 index 0000000..88216ea --- /dev/null +++ b/crates/agent/src/agent/prompt_builder.rs @@ -0,0 +1,60 @@ +use llm::ImageUrl; + +/// Builder for constructing an agent prompt request. +/// +/// Created by [`Agent::prompt()`](super::Agent::prompt). Configure the +/// prompt with chained setter calls, then pass it to +/// [`Agent::ask()`](super::Agent::ask) to execute. +/// +/// # Examples +/// +/// ```ignore +/// // Streaming prompt +/// let stream = agent.ask(agent.prompt().content("hello").stream(true)).await?; +/// +/// // Buffered prompt with images +/// let response = agent.ask( +/// agent.prompt().content("describe").images(vec![img]) +/// ).await?.into_response().await?; +/// ``` +pub struct PromptBuilder { + pub(super) content: Option, + pub(super) images: Vec, + pub(super) stream: bool, +} + +impl PromptBuilder { + pub(super) fn new() -> Self { + Self { + content: None, + images: Vec::new(), + stream: false, + } + } + + /// Set the text content of the prompt. + pub fn content(mut self, text: impl Into) -> Self { + self.content = Some(text.into()); + self + } + + /// Add one or more images to the prompt. + pub fn images(mut self, images: impl IntoIterator) -> Self { + self.images.extend(images); + self + } + + /// Set the streaming mode. + /// + /// When `true`, the returned [`AgentStream`](super::AgentStream) emits + /// incremental [`TextDelta`](super::AgentEvent::TextDelta) and + /// [`ReasoningDelta`](super::AgentEvent::ReasoningDelta) events as the + /// model generates them. + /// + /// When `false` (the default), only tool-call events and the final + /// [`Finished`](super::AgentEvent::Finished) event are emitted. + pub fn stream(mut self, stream: bool) -> Self { + self.stream = stream; + self + } +} diff --git a/crates/agent/src/agent/tests.rs b/crates/agent/src/agent/tests.rs new file mode 100644 index 0000000..20fd50e --- /dev/null +++ b/crates/agent/src/agent/tests.rs @@ -0,0 +1,986 @@ +use super::*; +use crate::AgentTool; +use crate::session::{SessionManager, SessionRecord}; +use async_trait::async_trait; +use llm::{ContentBlock, LlmApi, LlmError, LlmEvent, LlmResponse, StopReason}; +use providers::{ApiId, ModelCapabilities, ModelInfo, OpenRouterProvider, Provider, ProviderId}; +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +struct FakeApi; + +#[async_trait] +impl LlmApi for FakeApi { + async fn stream(&self, request: llm::LlmRequest<'_>) -> Result { + let user = request + .messages + .iter() + .rev() + .find_map(|message| message.content.clone()) + .unwrap_or_default(); + let response = LlmResponse { + content: vec![ContentBlock::Text(format!("echo: {user}"))], + stop_reason: StopReason::Stop, + usage: None, + model: Some(request.model_id.to_owned()), + reasoning: None, + reasoning_details: Vec::new(), + }; + let text = response.text(); + let model = response.model.clone(); + Ok(Box::pin(futures_util::stream::iter([ + Ok(llm::LlmEvent::TextDelta { text }), + Ok(llm::LlmEvent::Done { + stop_reason: StopReason::Stop, + usage: None, + model, + }), + ]))) + } +} + +struct FailAfterFirstApi { + calls: AtomicUsize, +} + +struct PendingAfterFirstApi { + calls: AtomicUsize, +} + +#[async_trait] +impl LlmApi for FailAfterFirstApi { + async fn stream(&self, request: llm::LlmRequest<'_>) -> Result { + if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + return Ok(Box::pin(futures_util::stream::iter([ + Ok(LlmEvent::TextDelta { + text: "first response".into(), + }), + Ok(LlmEvent::Done { + stop_reason: StopReason::Stop, + usage: None, + model: Some(request.model_id.to_owned()), + }), + ]))); + } + + Ok(Box::pin(futures_util::stream::iter([ + Ok(LlmEvent::TextDelta { + text: "partial response".into(), + }), + Err(LlmError::InvalidResponse( + "stream ended without [DONE]".into(), + )), + ]))) + } +} + +#[async_trait] +impl LlmApi for PendingAfterFirstApi { + async fn stream(&self, request: llm::LlmRequest<'_>) -> Result { + if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + return Ok(Box::pin(futures_util::stream::iter([ + Ok(LlmEvent::TextDelta { + text: "first response".into(), + }), + Ok(LlmEvent::Done { + stop_reason: StopReason::Stop, + usage: None, + model: Some(request.model_id.to_owned()), + }), + ]))); + } + + Ok(Box::pin(futures_util::stream::unfold( + 0, + |state| async move { + match state { + 0 => Some(( + Ok(LlmEvent::TextDelta { + text: "partial response".into(), + }), + 1, + )), + _ => futures_util::future::pending().await, + } + }, + ))) + } +} + +fn model_with_api(api: Arc) -> Model { + let info = ModelInfo { + provider: ProviderId::new("openrouter"), + id: "test".into(), + name: "Test".into(), + api: ApiId::ChatCompletions, + capabilities: ModelCapabilities::default(), + pricing: None, + }; + OpenRouterProvider::builder("key") + .with_models([info]) + .with_api(api) + .build() + .unwrap() + .bind("test") + .unwrap() +} + +fn model() -> Model { + model_with_api(Arc::new(FakeApi)) +} + +// --------------------------------------------------------------------------- +// Helper: build an Arc with the given model and optional setup +// --------------------------------------------------------------------------- + +fn agent(model: Model) -> Arc { + Arc::new(Agent::builder(model).build().unwrap()) +} + +fn agent_with_manager(model: Model, manager: Arc) -> Arc { + Arc::new( + Agent::builder(model) + .session_manager(manager) + .build() + .unwrap(), + ) +} + +// --------------------------------------------------------------------------- +// PromptBuilder API tests +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn prompt_owns_history_and_system_prompt() { + let a = Arc::new( + Agent::builder(model()) + .system_prompt("Be helpful") + .build() + .unwrap(), + ); + let response = a + .ask(a.prompt().content("hello")) + .unwrap() + .into_response() + .await + .unwrap(); + assert_eq!(response.text(), "echo: hello"); + assert_eq!(a.messages().await.len(), 2); +} + +#[tokio::test] +async fn stream_false_suppresses_text_deltas() { + let a = Arc::new( + Agent::builder(model()) + .system_prompt("Be helpful") + .build() + .unwrap(), + ); + let mut rx = a.ask(a.prompt().content("hello").stream(false)).unwrap(); + + let mut events = Vec::new(); + while let Some(event) = rx.recv().await { + events.push(event.unwrap()); + } + + // stream=false suppresses TextDelta; only Finished should arrive. + assert_eq!(events.len(), 1); + assert!(matches!(&events[0], AgentEvent::Finished { .. })); + assert_eq!(a.messages().await.len(), 2); +} + +#[tokio::test] +async fn stream_true_emits_text_deltas_and_finished() { + let a = Arc::new( + Agent::builder(model()) + .system_prompt("Be helpful") + .build() + .unwrap(), + ); + let mut rx = a.ask(a.prompt().content("hello").stream(true)).unwrap(); + + let mut events = Vec::new(); + while let Some(event) = rx.recv().await { + events.push(event.unwrap()); + } + + assert!(matches!( + &events[0], + AgentEvent::TextDelta(text) if text == "echo: hello" + )); + assert!(matches!(events.last(), Some(AgentEvent::Finished { .. }))); + assert_eq!(a.messages().await.len(), 2); +} + +#[tokio::test] +async fn empty_prompt_content_returns_error() { + let a = agent(model()); + let result = a.ask(a.prompt().content(" ")); + assert!(result.is_err(), "must reject blank prompt"); +} + +#[tokio::test] +async fn missing_content_returns_error() { + let a = agent(model()); + let result = a.ask(a.prompt()); + assert!(result.is_err(), "must reject builder with no content"); +} + +#[tokio::test] +async fn into_response_extracts_final_response() { + let a = Arc::new( + Agent::builder(model()) + .system_prompt("Be helpful") + .build() + .unwrap(), + ); + let stream = a.ask(a.prompt().content("hello").stream(true)).unwrap(); + let response = stream.into_response().await.unwrap(); + assert_eq!(response.text(), "echo: hello"); + assert_eq!(a.messages().await.len(), 2); +} + +#[tokio::test] +async fn images_are_passed_through_builder() { + let api = Arc::new(ImageInspectApi); + let a = Arc::new(Agent::builder(model_with_api(api)).build().unwrap()); + + let images = vec![llm::ImageUrl { + url: "data:image/png;base64,abc123".into(), + }]; + let response = a + .ask(a.prompt().content("describe this").images(images)) + .unwrap() + .into_response() + .await + .unwrap(); + assert_eq!(response.text(), "text=true,images=1"); +} + +// --------------------------------------------------------------------------- +// Reasoning streaming +// --------------------------------------------------------------------------- + +struct ReasoningApi; + +#[async_trait] +impl LlmApi for ReasoningApi { + async fn stream(&self, _request: llm::LlmRequest<'_>) -> Result { + Ok(Box::pin(futures_util::stream::iter([ + Ok(llm::LlmEvent::ReasoningDelta { + reasoning: "thinking...".into(), + details: vec![serde_json::json!({ + "type": "reasoning.text", + "text": "thinking..." + })], + }), + Ok(llm::LlmEvent::TextDelta { + text: "done thinking".into(), + }), + Ok(llm::LlmEvent::Done { + stop_reason: StopReason::Stop, + usage: None, + model: Some("reasoning-model".into()), + }), + ]))) + } +} + +#[tokio::test] +async fn stream_true_emits_reasoning_and_text_deltas() { + let a = Arc::new( + Agent::builder(model_with_api(Arc::new(ReasoningApi))) + .build() + .unwrap(), + ); + let mut rx = a.ask(a.prompt().content("hello").stream(true)).unwrap(); + + let mut events = Vec::new(); + while let Some(event) = rx.recv().await { + events.push(event.unwrap()); + } + + assert!(matches!( + &events[0], + AgentEvent::ReasoningDelta(text) if text == "thinking..." + )); + assert!(matches!( + &events[1], + AgentEvent::TextDelta(text) if text == "done thinking" + )); + assert!(matches!(events.last(), Some(AgentEvent::Finished { .. }))); + + let messages = a.messages().await; + assert_eq!(messages.len(), 2); + if let AgentMessage::Assistant(resp) = &messages[1] { + assert_eq!(resp.reasoning.as_deref(), Some("thinking...")); + assert_eq!(resp.text(), "done thinking"); + } else { + panic!("expected assistant message"); + } +} + +#[tokio::test] +async fn stream_false_suppresses_reasoning_deltas() { + let a = Arc::new( + Agent::builder(model_with_api(Arc::new(ReasoningApi))) + .build() + .unwrap(), + ); + let mut rx = a.ask(a.prompt().content("hello").stream(false)).unwrap(); + + let mut events = Vec::new(); + while let Some(event) = rx.recv().await { + events.push(event.unwrap()); + } + + // ReasoningDelta should be suppressed. + assert_eq!(events.len(), 1); + assert!(matches!(&events[0], AgentEvent::Finished { .. })); +} + +// --------------------------------------------------------------------------- +// Error recovery +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn streaming_error_preserves_previous_history() { + let api = Arc::new(FailAfterFirstApi { + calls: AtomicUsize::new(0), + }); + let a = Arc::new(Agent::builder(model_with_api(api)).build().unwrap()); + + a.ask(a.prompt().content("first")) + .unwrap() + .into_response() + .await + .unwrap(); + + let mut stream = a.ask(a.prompt().content("second").stream(true)).unwrap(); + let mut error = None; + while let Some(event) = stream.recv().await { + if let Err(error_event) = event { + error = Some(error_event); + break; + } + } + + assert_eq!( + error.unwrap().to_string(), + "invalid response: stream ended without [DONE]" + ); + let messages = a.messages().await; + assert_eq!(messages.len(), 4); + assert!(matches!(&messages[0], AgentMessage::User { text, .. } if text == "first")); + assert!( + matches!(&messages[1], AgentMessage::Assistant(response) if response.text() == "first response") + ); + assert!(matches!(&messages[2], AgentMessage::User { text, .. } if text == "second")); + assert!( + matches!(&messages[3], AgentMessage::Assistant(response) if response.text() == "partial response") + ); +} + +#[tokio::test] +async fn abort_preserves_partial_assistant_output() { + let api = Arc::new(PendingAfterFirstApi { + calls: AtomicUsize::new(0), + }); + let a = Arc::new(Agent::builder(model_with_api(api)).build().unwrap()); + + a.ask(a.prompt().content("first")) + .unwrap() + .into_response() + .await + .unwrap(); + + let mut stream = a.ask(a.prompt().content("second").stream(true)).unwrap(); + assert!(matches!( + stream.recv().await, + Some(Ok(AgentEvent::TextDelta(text))) if text == "partial response" + )); + stream.abort(); + assert!(matches!( + stream.recv().await, + Some(Err(AgentError::Aborted)) + )); + + let messages = a.messages().await; + assert_eq!(messages.len(), 4); + assert!(matches!(&messages[2], AgentMessage::User { text, .. } if text == "second")); + assert!( + matches!(&messages[3], AgentMessage::Assistant(response) if response.text() == "partial response") + ); +} + +// --------------------------------------------------------------------------- +// Tool calling +// --------------------------------------------------------------------------- + +struct ToolCallingApi { + calls: AtomicUsize, +} + +#[async_trait] +impl LlmApi for ToolCallingApi { + async fn stream(&self, request: llm::LlmRequest<'_>) -> Result { + let call_count = self.calls.fetch_add(1, Ordering::SeqCst); + if call_count == 0 { + let response = LlmResponse { + content: vec![], + stop_reason: StopReason::ToolUse, + usage: Some(llm::Usage { + input_tokens: 10, + output_tokens: 5, + ..llm::Usage::default() + }), + model: Some(request.model_id.to_owned()), + reasoning: None, + reasoning_details: Vec::new(), + }; + Ok(Box::pin(futures_util::stream::iter([ + Ok(LlmEvent::ToolCallDelta { + index: 0, + id: Some("call-1".into()), + name: Some("bash".into()), + arguments: serde_json::json!({"command": "echo hi"}).to_string(), + }), + Ok(LlmEvent::Done { + stop_reason: StopReason::ToolUse, + usage: response.usage.clone(), + model: response.model.clone(), + }), + ]))) + } else { + let response = LlmResponse { + content: vec![ContentBlock::Text("done".into())], + stop_reason: StopReason::Stop, + usage: Some(llm::Usage { + input_tokens: 20, + output_tokens: 3, + ..llm::Usage::default() + }), + model: Some(request.model_id.to_owned()), + reasoning: None, + reasoning_details: Vec::new(), + }; + let text = response.text(); + let model = response.model.clone(); + Ok(Box::pin(futures_util::stream::iter([ + Ok(LlmEvent::TextDelta { text }), + Ok(LlmEvent::Done { + stop_reason: StopReason::Stop, + usage: response.usage.clone(), + model, + }), + ]))) + } + } +} + +// --------------------------------------------------------------------------- +// Session persistence +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn building_agent_with_manager_creates_no_file() { + let root = std::env::temp_dir().join(format!("alan-plan3-build-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let manager = Arc::new(SessionManager::new(&root)); + let _a = Arc::new( + Agent::builder(model()) + .session_manager(manager) + .build() + .unwrap(), + ); + + // No session file should exist before any prompt. + let entries: Vec<_> = std::fs::read_dir(&root) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + assert!( + entries.is_empty(), + "building must not create a session file" + ); +} + +#[tokio::test] +async fn first_prompt_creates_session_and_persists_messages() { + let root = std::env::temp_dir().join(format!("alan-plan3-buffered-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let manager = Arc::new(SessionManager::new(&root)); + let a = agent_with_manager(model(), manager.clone()); + + let response = a + .ask(a.prompt().content("hello")) + .unwrap() + .into_response() + .await + .unwrap(); + assert_eq!(response.text(), "echo: hello"); + + let entries: Vec<_> = std::fs::read_dir(&root) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + assert_eq!(entries.len(), 1, "one pwd directory created"); + let pwd_dir = entries[0].path(); + let files: Vec<_> = std::fs::read_dir(&pwd_dir) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + assert_eq!(files.len(), 1, "one session file created"); + + let session_file = files[0].path(); + let content = std::fs::read_to_string(&session_file).unwrap(); + let lines: Vec<_> = content.lines().collect(); + assert_eq!(lines.len(), 3, "header + user + assistant"); + + let header: serde_json::Value = serde_json::from_str(lines[0]).unwrap(); + assert_eq!(header["type"], "session"); + + let user_record: serde_json::Value = serde_json::from_str(lines[1]).unwrap(); + assert_eq!(user_record["type"], "message"); + assert_eq!(user_record["message"]["kind"], "user"); + assert_eq!(user_record["message"]["content"], "hello"); + + let assistant_record: serde_json::Value = serde_json::from_str(lines[2]).unwrap(); + assert_eq!(assistant_record["type"], "message"); + assert_eq!(assistant_record["message"]["kind"], "assistant"); +} + +#[tokio::test] +async fn provider_is_not_called_if_session_creation_fails() { + let _root = + std::env::temp_dir().join(format!("alan-plan3-fail-create-{}", uuid::Uuid::new_v4())); + let manager = Arc::new(SessionManager::new("/proc/self/mem/unwritable-dir")); + let a = agent_with_manager(model(), manager); + + let result = a + .ask(a.prompt().content("hello")) + .unwrap() + .into_response() + .await; + assert!(result.is_err(), "must fail when session creation fails"); +} + +#[tokio::test] +async fn tool_call_responses_and_results_are_persisted_in_order() { + let root = std::env::temp_dir().join(format!("alan-plan3-tool-order-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let manager = Arc::new(SessionManager::new(&root)); + let api = Arc::new(ToolCallingApi { + calls: AtomicUsize::new(0), + }); + let a = Arc::new( + Agent::builder(model_with_api(api)) + .session_manager(manager.clone()) + .with_tools([AgentTool::new( + llm::ToolDefinition { + name: "bash".into(), + description: "Run a shell command".into(), + parameters: serde_json::json!({}), + }, + tools::BashExecutor, + )]) + .build() + .unwrap(), + ); + + let response = a + .ask(a.prompt().content("run echo hi")) + .unwrap() + .into_response() + .await + .unwrap(); + assert_eq!(response.text(), "done"); + + let content = { + let entries: Vec<_> = std::fs::read_dir(&root) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + let pwd_dir = entries[0].path(); + let files: Vec<_> = std::fs::read_dir(&pwd_dir) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + std::fs::read_to_string(files[0].path()).unwrap() + }; + let lines: Vec<_> = content.lines().collect(); + assert!( + lines.len() >= 5, + "expected at least 5 records, got {}", + lines.len() + ); + + let types: Vec<_> = lines + .iter() + .map(|l| { + let v: serde_json::Value = serde_json::from_str(l).unwrap(); + v["type"].as_str().unwrap().to_owned() + }) + .collect(); + assert_eq!(types[0], "session"); + assert_eq!(types[1], "message"); + assert_eq!(types[types.len() - 1], "message"); +} + +#[tokio::test] +async fn aggregate_usage_from_multiple_rounds_is_persisted() { + let root = std::env::temp_dir().join(format!("alan-plan3-usage-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let manager = Arc::new(SessionManager::new(&root)); + let api = Arc::new(ToolCallingApi { + calls: AtomicUsize::new(0), + }); + let a = Arc::new( + Agent::builder(model_with_api(api)) + .session_manager(manager.clone()) + .with_tools([AgentTool::new( + llm::ToolDefinition { + name: "bash".into(), + description: "Run a shell command".into(), + parameters: serde_json::json!({}), + }, + tools::BashExecutor, + )]) + .build() + .unwrap(), + ); + + let response = a + .ask(a.prompt().content("run echo hi")) + .unwrap() + .into_response() + .await + .unwrap(); + assert_eq!(response.text(), "done"); + + let entries: Vec<_> = std::fs::read_dir(&root) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + let pwd_dir = entries[0].path(); + let files: Vec<_> = std::fs::read_dir(&pwd_dir) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + let session_file = files[0].path(); + + let session = { + let content = std::fs::read_to_string(&session_file).unwrap(); + let lines: Vec<_> = content.lines().collect(); + let header_line = lines[0]; + let record = SessionRecord::parse(header_line).unwrap(); + let SessionRecord::Session { id, .. } = record else { + panic!("expected session header"); + }; + manager + .get_session(&id, &std::env::current_dir().unwrap()) + .await + .expect("load session for usage check") + }; + + // First round: 10 input + 5 output. Second round: 20 input + 3 output. + assert_eq!(session.usage.input_tokens, 30); + assert_eq!(session.usage.output_tokens, 8); +} + +#[tokio::test] +async fn resumed_agent_includes_restored_messages_in_first_request() { + let root = std::env::temp_dir().join(format!("alan-plan3-resume-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let manager = Arc::new(SessionManager::new(&root)); + + let session = manager + .create(&root, "openrouter", "test", None) + .await + .expect("create session"); + manager + .append_message( + &session.id, + &session.pwd, + &AgentMessage::user("restored message"), + ) + .await + .expect("append message"); + let session = manager + .get_session(&session.id, &root) + .await + .expect("load session with restored message"); + + let m = model(); + let a = Arc::new( + Agent::builder(m) + .session_manager(manager.clone()) + .resume_session(session) + .build() + .expect("build with resume"), + ); + + let response = a + .ask(a.prompt().content("new message")) + .unwrap() + .into_response() + .await + .unwrap(); + assert_eq!(response.text(), "echo: new message"); + + let messages = a.messages().await; + assert!( + messages.len() >= 2, + "must include restored user + new user + assistant response, got {} messages", + messages.len() + ); + assert!( + matches!(&messages[0], AgentMessage::User { text, .. } if text == "restored message"), + "first message must be the restored message, got {:?}", + messages[0] + ); +} + +#[tokio::test] +async fn request_uses_persisted_session_id_and_cache_key() { + let root = std::env::temp_dir().join(format!("alan-plan3-session-id-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let manager = Arc::new(SessionManager::new(&root)); + let a = agent_with_manager(model(), manager); + + a.ask(a.prompt().content("check session id")) + .unwrap() + .into_response() + .await + .unwrap(); + + let entries: Vec<_> = std::fs::read_dir(&root) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + let pwd_dir = entries[0].path(); + let files: Vec<_> = std::fs::read_dir(&pwd_dir) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + let content = std::fs::read_to_string(files[0].path()).unwrap(); + let lines: Vec<_> = content.lines().collect(); + + let header: serde_json::Value = serde_json::from_str(lines[0]).unwrap(); + let session_id = header["id"].as_str().unwrap(); + let session_file_name = files[0].path().file_stem().unwrap().to_owned(); + assert_eq!(session_file_name, session_id); +} + +// --------------------------------------------------------------------------- +// Image tests +// --------------------------------------------------------------------------- + +struct ImageInspectApi; + +#[async_trait] +impl LlmApi for ImageInspectApi { + async fn stream(&self, request: llm::LlmRequest<'_>) -> Result { + let parts_summary = request + .messages + .iter() + .rev() + .find_map(|m| m.content_parts.as_ref()) + .map(|parts| { + let mut has_text = false; + let mut image_count = 0; + for part in parts { + match part { + llm::ContentPart::Text { .. } => has_text = true, + llm::ContentPart::Image { .. } => image_count += 1, + } + } + format!("text={has_text},images={image_count}") + }) + .unwrap_or_default(); + + let response = LlmResponse { + content: vec![ContentBlock::Text(parts_summary.clone())], + stop_reason: StopReason::Stop, + usage: None, + model: Some(request.model_id.to_owned()), + reasoning: None, + reasoning_details: Vec::new(), + }; + let text = response.text(); + let model = response.model.clone(); + Ok(Box::pin(futures_util::stream::iter([ + Ok(llm::LlmEvent::TextDelta { text }), + Ok(llm::LlmEvent::Done { + stop_reason: StopReason::Stop, + usage: None, + model, + }), + ]))) + } +} + +#[tokio::test] +async fn prompt_with_images_sends_content_parts() { + let api = Arc::new(ImageInspectApi); + let a = Arc::new(Agent::builder(model_with_api(api)).build().unwrap()); + + let images = vec![llm::ImageUrl { + url: "data:image/png;base64,abc123".into(), + }]; + let response = a + .ask(a.prompt().content("describe this").images(images)) + .unwrap() + .into_response() + .await + .unwrap(); + assert_eq!(response.text(), "text=true,images=1"); +} + +#[tokio::test] +async fn prompt_with_images_persists_images_in_session() { + let root = + std::env::temp_dir().join(format!("alan-plan5-img-persist-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let manager = Arc::new(SessionManager::new(&root)); + let api = Arc::new(ImageInspectApi); + let a = Arc::new( + Agent::builder(model_with_api(api)) + .session_manager(manager.clone()) + .build() + .unwrap(), + ); + + let images = vec![ + llm::ImageUrl { + url: "data:image/png;base64,abc".into(), + }, + llm::ImageUrl { + url: "https://example.com/photo.jpg".into(), + }, + ]; + let _response = a + .ask(a.prompt().content("describe these").images(images.clone())) + .unwrap() + .into_response() + .await + .unwrap(); + + let entries: Vec<_> = std::fs::read_dir(&root) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + let pwd_dir = entries[0].path(); + let files: Vec<_> = std::fs::read_dir(pwd_dir) + .unwrap() + .filter_map(|e| e.ok()) + .collect(); + let content = std::fs::read_to_string(files[0].path()).unwrap(); + let lines: Vec<_> = content.lines().collect(); + + let user_record: serde_json::Value = serde_json::from_str(lines[1]).unwrap(); + let msg = &user_record["message"]; + assert_eq!(msg["kind"], "user"); + assert_eq!(msg["content"], "describe these"); + let persisted_images = msg["images"].as_array().unwrap(); + assert_eq!(persisted_images.len(), 2); + assert_eq!(persisted_images[0]["url"], "data:image/png;base64,abc"); + assert_eq!(persisted_images[1]["url"], "https://example.com/photo.jpg"); +} + +#[tokio::test] +async fn session_images_roundtrip_through_reload() { + let root = + std::env::temp_dir().join(format!("alan-plan5-img-roundtrip-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let manager = Arc::new(SessionManager::new(&root)); + + let session = manager + .create(&root, "openrouter", "test", None) + .await + .expect("create session"); + let images = vec![llm::ImageUrl { + url: "https://example.com/pic.png".into(), + }]; + manager + .append_message( + &session.id, + &session.pwd, + &AgentMessage::user_with_images("look at this", images), + ) + .await + .expect("append image message"); + + let loaded = manager + .get_session(&session.id, &root) + .await + .expect("reload session"); + assert_eq!(loaded.messages.len(), 1); + match &loaded.messages[0] { + AgentMessage::User { text, images } => { + assert_eq!(text, "look at this"); + assert_eq!(images.len(), 1); + assert_eq!(images[0].url, "https://example.com/pic.png"); + } + other => panic!("expected User message, got {other:?}"), + } +} + +// --------------------------------------------------------------------------- +// Message / serialization tests +// --------------------------------------------------------------------------- + +#[test] +fn legacy_session_without_images_field_loads() { + let legacy = r#"{"kind":"user","content":"hello"}"#; + let msg: AgentMessage = serde_json::from_str(legacy).expect("legacy format deserializes"); + match msg { + AgentMessage::User { text, images } => { + assert_eq!(text, "hello"); + assert!(images.is_empty()); + } + other => panic!("expected User, got {other:?}"), + } +} + +#[test] +fn to_llm_user_with_no_images_is_plain_user_message() { + let msg = AgentMessage::user("hello"); + let llm_msg = msg.to_llm(); + assert_eq!(llm_msg.role, llm::Role::User); + assert_eq!(llm_msg.content.as_deref(), Some("hello")); + assert!(llm_msg.content_parts.is_none()); +} + +#[test] +fn to_llm_user_with_images_uses_content_parts() { + let msg = AgentMessage::user_with_images( + "describe", + vec![llm::ImageUrl { + url: "https://example.com/img.png".into(), + }], + ); + let llm_msg = msg.to_llm(); + assert_eq!(llm_msg.role, llm::Role::User); + assert!(llm_msg.content.is_none()); + let parts = llm_msg.content_parts.unwrap(); + assert_eq!(parts.len(), 2); + assert!(matches!(&parts[0], llm::ContentPart::Text { text } if text == "describe")); + assert!( + matches!(&parts[1], llm::ContentPart::Image { image_url } if image_url.url == "https://example.com/img.png") + ); +} + +#[test] +fn user_constructor_has_empty_images() { + let msg = AgentMessage::user("text"); + match msg { + AgentMessage::User { text, images } => { + assert_eq!(text, "text"); + assert!(images.is_empty()); + } + _ => panic!("expected User"), + } +} diff --git a/crates/agent/src/agent/tool_loop.rs b/crates/agent/src/agent/tool_loop.rs new file mode 100644 index 0000000..25f859a --- /dev/null +++ b/crates/agent/src/agent/tool_loop.rs @@ -0,0 +1,169 @@ +use crate::AgentError; +use crate::AgentMessage; +use crate::context::AgentContext; +use llm::LlmResponse; +use providers::Model; + +use super::Agent; +use super::event::{AgentEvent, emit_event}; +use super::prompt::PromptCx; + +/// Core agent loop: stream LLM responses and execute tool calls until the +/// model produces a final answer or `max_tool_rounds` is reached. +pub(super) async fn run_with( + agent: &Agent, + model: &Model, + context: &mut AgentContext, + cx: &mut PromptCx<'_>, +) -> Result { + let plan = agent.plan_mode(); + + for _ in 0..agent.max_tool_rounds { + cx.check_cancelled()?; + + let session_id = agent.session_id.lock().await.clone(); + let response = super::prompt::stream_round(session_id, model, context, cx, plan).await?; + + // Accumulate token usage across rounds. + if let Some(usage) = response.usage.as_ref() { + context.usage.accumulate(usage); + super::persistence::persist_usage(agent, &context.usage).await?; + } + + let calls: Vec<_> = response.tool_calls().cloned().collect(); + + if calls.is_empty() { + return finish_with_response(agent, context, response, cx).await; + } + + // Persist the assistant message that declares the tool calls. + super::persistence::append_context_message( + agent, + context, + AgentMessage::Assistant(response), + ) + .await?; + + handle_tool_calls(agent, calls, context, cx, plan).await?; + } + + Err(AgentError::MaxToolRounds) +} + +/// Emit the `Finished` event and return the final response. +async fn finish_with_response( + agent: &Agent, + context: &mut AgentContext, + response: LlmResponse, + cx: &mut PromptCx<'_>, +) -> Result { + super::persistence::append_context_message( + agent, + context, + AgentMessage::Assistant(response.clone()), + ) + .await?; + + emit_event( + cx.events, + AgentEvent::Finished { + usage: context.usage.clone(), + response: Box::new(response.clone()), + }, + cx.cancellation, + ) + .await?; + + Ok(response) +} + +/// Execute a batch of tool calls, appending each result to the context +/// and emitting start/finish/fail events. +async fn handle_tool_calls( + agent: &Agent, + calls: Vec, + context: &mut AgentContext, + cx: &mut PromptCx<'_>, + plan: bool, +) -> Result<(), AgentError> { + for call in calls { + cx.check_cancelled()?; + + let tool_index = context + .tool_indexes + .get(&call.name) + .copied() + .ok_or_else(|| AgentError::ToolNotFound(call.name.clone()))?; + + // In plan mode only read-only tools (and bash) may be invoked. + if plan && !context.tools[tool_index].read_only && call.name != "bash" { + return Err(AgentError::ToolNotFound(call.name.clone())); + } + + let call_id = call.id.clone(); + + emit_event( + cx.events, + AgentEvent::ToolCallStarted { + id: call_id.clone(), + name: call.name.clone(), + arguments: call.arguments.clone(), + }, + cx.cancellation, + ) + .await?; + + match context.tools[tool_index].executor.execute(&call).await { + Ok(output) => { + super::persistence::append_context_message( + agent, + context, + AgentMessage::ToolResult { + tool_call_id: call_id.clone(), + content: output.clone(), + }, + ) + .await?; + + emit_event( + cx.events, + AgentEvent::ToolCallFinished { + id: call_id, + output: tail_lines(&output, 5), + }, + cx.cancellation, + ) + .await?; + } + Err(error) => { + let error = error.to_string(); + super::persistence::append_context_message( + agent, + context, + AgentMessage::ToolResult { + tool_call_id: call_id.clone(), + content: error.clone(), + }, + ) + .await?; + + emit_event( + cx.events, + AgentEvent::ToolCallFailed { id: call_id, error }, + cx.cancellation, + ) + .await?; + } + } + + cx.check_cancelled()?; + } + + Ok(()) +} + +fn tail_lines(output: &str, count: usize) -> String { + let mut lines: Vec<_> = output.lines().rev().take(count).collect(); + lines.reverse(); + lines.join("\n") +} diff --git a/crates/agent/src/context.rs b/crates/agent/src/context.rs index 3f44e2e..29e2e6a 100644 --- a/crates/agent/src/context.rs +++ b/crates/agent/src/context.rs @@ -5,7 +5,10 @@ use std::collections::HashMap; #[derive(Debug, Clone, PartialEq)] pub enum AgentMessage { - User(String), + User { + text: String, + images: Vec, + }, Assistant(LlmResponse), ToolResult { tool_call_id: String, @@ -28,6 +31,8 @@ pub enum AgentMessage { enum SessionMessage { User { content: String, + #[serde(default)] + images: Vec, }, Assistant { response: LlmResponse, @@ -41,8 +46,9 @@ enum SessionMessage { impl From<&AgentMessage> for SessionMessage { fn from(message: &AgentMessage) -> Self { match message { - AgentMessage::User(content) => Self::User { - content: content.clone(), + AgentMessage::User { text, images } => Self::User { + content: text.clone(), + images: images.clone(), }, AgentMessage::Assistant(response) => Self::Assistant { response: response.clone(), @@ -61,7 +67,10 @@ impl From<&AgentMessage> for SessionMessage { impl From for AgentMessage { fn from(message: SessionMessage) -> Self { match message { - SessionMessage::User { content } => Self::User(content), + SessionMessage::User { content, images } => Self::User { + text: content, + images, + }, SessionMessage::Assistant { response } => Self::Assistant(response), SessionMessage::ToolResult { tool_call_id, @@ -88,12 +97,34 @@ impl<'de> Deserialize<'de> for AgentMessage { impl AgentMessage { pub fn user(content: impl Into) -> Self { - Self::User(content.into()) + Self::User { + text: content.into(), + images: Vec::new(), + } + } + + pub fn user_with_images(text: impl Into, images: Vec) -> Self { + Self::User { + text: text.into(), + images, + } } pub fn to_llm(&self) -> Message { match self { - Self::User(content) => Message::user(content), + Self::User { text, images } => { + if images.is_empty() { + Message::user(text) + } else { + let mut parts = vec![llm::ContentPart::Text { text: text.clone() }]; + for image in images { + parts.push(llm::ContentPart::Image { + image_url: image.clone(), + }); + } + Message::user_with_parts(parts) + } + } Self::Assistant(message) => { let text = message.text(); let calls: Vec = message.tool_calls().cloned().collect(); diff --git a/crates/agent/src/lib.rs b/crates/agent/src/lib.rs index 4f76856..0c2635e 100644 --- a/crates/agent/src/lib.rs +++ b/crates/agent/src/lib.rs @@ -5,7 +5,7 @@ mod session; mod skill; mod tool; -pub use agent::{Agent, AgentBuilder, AgentEvent, AgentStream}; +pub use agent::{Agent, AgentBuilder, AgentEvent, AgentStream, PromptBuilder}; pub use context::AgentMessage; pub use error::AgentError; pub use session::{SESSION_SCHEMA_VERSION, Session, SessionError, SessionManager, SessionRecord}; diff --git a/crates/alan/src/core/chat.rs b/crates/alan/src/core/chat.rs index a89d925..52eabc0 100644 --- a/crates/alan/src/core/chat.rs +++ b/crates/alan/src/core/chat.rs @@ -66,7 +66,7 @@ impl ChatController { for message in messages { match message { - agent::AgentMessage::User(text) => self.entries.push(Entry::Prompt(text)), + agent::AgentMessage::User { text, .. } => self.entries.push(Entry::Prompt(text)), agent::AgentMessage::Assistant(response) => { if let Some(reasoning) = response.reasoning.as_deref() && !reasoning.is_empty() @@ -146,7 +146,12 @@ impl ChatController { self.entries.push(Entry::Prompt(text.to_owned())); self.revision = self.revision.wrapping_add(1); self.busy = true; - self.stream = Some(self.agent.prompt_stream(text.to_owned())); + let builder = self.agent.prompt().content(text.to_owned()).stream(true); + self.stream = Some( + self.agent + .ask(builder) + .expect("prompt was already validated"), + ); } pub fn abort(&mut self) { @@ -216,7 +221,7 @@ impl ChatController { ToolStatus::Failed(error), ); } - AgentEvent::Finished { usage } => { + AgentEvent::Finished { usage, .. } => { changed |= Self::append_delta(&mut self.entries, &pending_text); pending_text.clear(); changed |= Self::ensure_response_entry(&mut self.entries); From 87de967aa9d5a356879332d889b384555635fda1 Mon Sep 17 00:00:00 2001 From: Revantark Date: Wed, 26 Aug 2026 21:24:43 +0530 Subject: [PATCH 3/5] let read tool to read images --- Cargo.lock | 1 + Cargo.toml | 1 + crates/agent/src/agent/tests.rs | 55 +++++++++ crates/agent/src/agent/tool_loop.rs | 27 ++++- crates/agent/src/context.rs | 32 ++++- crates/alan/src/core/chat.rs | 1 + crates/llm/src/message.rs | 17 +++ crates/tools/Cargo.toml | 1 + crates/tools/src/fs.rs | 179 ++++++++++++++++++++++++++-- crates/tools/src/shell.rs | 6 +- crates/tools/src/tool.rs | 13 +- 11 files changed, 311 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2b16a33..23af34f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2707,6 +2707,7 @@ name = "tools" version = "0.1.0" dependencies = [ "async-trait", + "base64", "llm", "serde_json", "thiserror", diff --git a/Cargo.toml b/Cargo.toml index d98cc0d..0075e0e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,3 +19,4 @@ tracing = "0.1.44" tracing-appender = "0.2" fs2 = "0.4" uuid = { version = "1", features = ["v4"] } +base64 = "0.22" diff --git a/crates/agent/src/agent/tests.rs b/crates/agent/src/agent/tests.rs index 20fd50e..89e87ea 100644 --- a/crates/agent/src/agent/tests.rs +++ b/crates/agent/src/agent/tests.rs @@ -945,6 +945,24 @@ fn legacy_session_without_images_field_loads() { } } +#[test] +fn legacy_tool_result_without_content_parts_loads() { + let legacy = r#"{"kind":"tool_result","tool_call_id":"call-1","content":"output"}"#; + let msg: AgentMessage = serde_json::from_str(legacy).expect("legacy format deserializes"); + match msg { + AgentMessage::ToolResult { + tool_call_id, + content, + content_parts, + } => { + assert_eq!(tool_call_id, "call-1"); + assert_eq!(content, "output"); + assert!(content_parts.is_empty()); + } + other => panic!("expected ToolResult, got {other:?}"), + } +} + #[test] fn to_llm_user_with_no_images_is_plain_user_message() { let msg = AgentMessage::user("hello"); @@ -984,3 +1002,40 @@ fn user_constructor_has_empty_images() { _ => panic!("expected User"), } } + +#[test] +fn to_llm_tool_result_with_parts_uses_content_parts() { + use llm::ContentPart; + let msg = AgentMessage::tool_result_with_parts( + "call-1", + "[image/png image, 100 bytes base64]", + vec![ContentPart::Image { + image_url: llm::ImageUrl { + url: "data:image/png;base64,iVBOR".into(), + }, + }], + ); + let llm_msg = msg.to_llm(); + assert_eq!(llm_msg.role, llm::Role::Tool); + assert_eq!(llm_msg.tool_call_id.as_deref(), Some("call-1")); + assert!(llm_msg.content.is_none()); + let parts = llm_msg.content_parts.unwrap(); + assert_eq!(parts.len(), 1); + assert!( + matches!(&parts[0], ContentPart::Image { image_url } if image_url.url == "data:image/png;base64,iVBOR") + ); +} + +#[test] +fn to_llm_tool_result_without_parts_uses_plain_content() { + let msg = AgentMessage::ToolResult { + tool_call_id: "call-2".into(), + content: "all good".into(), + content_parts: vec![], + }; + let llm_msg = msg.to_llm(); + assert_eq!(llm_msg.role, llm::Role::Tool); + assert_eq!(llm_msg.tool_call_id.as_deref(), Some("call-2")); + assert_eq!(llm_msg.content.as_deref(), Some("all good")); + assert!(llm_msg.content_parts.is_none()); +} diff --git a/crates/agent/src/agent/tool_loop.rs b/crates/agent/src/agent/tool_loop.rs index 25f859a..610340d 100644 --- a/crates/agent/src/agent/tool_loop.rs +++ b/crates/agent/src/agent/tool_loop.rs @@ -3,6 +3,7 @@ use crate::AgentMessage; use crate::context::AgentContext; use llm::LlmResponse; use providers::Model; +use tools::ToolOutput; use super::Agent; use super::event::{AgentEvent, emit_event}; @@ -115,13 +116,28 @@ async fn handle_tool_calls( match context.tools[tool_index].executor.execute(&call).await { Ok(output) => { + let (output_text, content_parts) = match output { + ToolOutput::Text(text) => (text, vec![]), + ToolOutput::Image { mime_type, data } => { + let description = + format!("[{mime_type} image, {} bytes base64]", data.len()); + let data_uri = format!("data:{mime_type};base64,{data}"); + ( + description, + vec![llm::ContentPart::Image { + image_url: llm::ImageUrl { url: data_uri }, + }], + ) + } + }; super::persistence::append_context_message( agent, context, - AgentMessage::ToolResult { - tool_call_id: call_id.clone(), - content: output.clone(), - }, + AgentMessage::tool_result_with_parts( + &call_id, + &output_text, + content_parts, + ), ) .await?; @@ -129,7 +145,7 @@ async fn handle_tool_calls( cx.events, AgentEvent::ToolCallFinished { id: call_id, - output: tail_lines(&output, 5), + output: tail_lines(&output_text, 5), }, cx.cancellation, ) @@ -143,6 +159,7 @@ async fn handle_tool_calls( AgentMessage::ToolResult { tool_call_id: call_id.clone(), content: error.clone(), + content_parts: vec![], }, ) .await?; diff --git a/crates/agent/src/context.rs b/crates/agent/src/context.rs index 29e2e6a..82634ae 100644 --- a/crates/agent/src/context.rs +++ b/crates/agent/src/context.rs @@ -1,5 +1,5 @@ use crate::{AgentTool, Skill}; -use llm::{LlmResponse, Message, ToolCall, Usage}; +use llm::{ContentPart, LlmResponse, Message, ToolCall, Usage}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -13,9 +13,9 @@ pub enum AgentMessage { ToolResult { tool_call_id: String, content: String, + content_parts: Vec, }, } - /// Serde representation of [`AgentMessage`]. /// /// Tagged explicitly with `kind` so the on-disk format does not depend on @@ -40,6 +40,8 @@ enum SessionMessage { ToolResult { tool_call_id: String, content: String, + #[serde(default)] + content_parts: Vec, }, } @@ -56,9 +58,11 @@ impl From<&AgentMessage> for SessionMessage { AgentMessage::ToolResult { tool_call_id, content, + content_parts, } => Self::ToolResult { tool_call_id: tool_call_id.clone(), content: content.clone(), + content_parts: content_parts.clone(), }, } } @@ -75,9 +79,11 @@ impl From for AgentMessage { SessionMessage::ToolResult { tool_call_id, content, + content_parts, } => Self::ToolResult { tool_call_id, content, + content_parts, }, } } @@ -110,6 +116,19 @@ impl AgentMessage { } } + /// Create a tool result with structured content parts (e.g. images). + pub fn tool_result_with_parts( + tool_call_id: impl Into, + content: impl Into, + content_parts: Vec, + ) -> Self { + Self::ToolResult { + tool_call_id: tool_call_id.into(), + content: content.into(), + content_parts, + } + } + pub fn to_llm(&self) -> Message { match self { Self::User { text, images } => { @@ -146,7 +165,14 @@ impl AgentMessage { Self::ToolResult { tool_call_id, content, - } => Message::tool_result(content, tool_call_id), + content_parts, + } => { + if content_parts.is_empty() { + Message::tool_result(content, tool_call_id) + } else { + Message::tool_result_with_parts(content_parts.clone(), tool_call_id) + } + } } } } diff --git a/crates/alan/src/core/chat.rs b/crates/alan/src/core/chat.rs index 52eabc0..8c085f1 100644 --- a/crates/alan/src/core/chat.rs +++ b/crates/alan/src/core/chat.rs @@ -90,6 +90,7 @@ impl ChatController { agent::AgentMessage::ToolResult { tool_call_id, content, + .. } => { if let Some(Entry::ToolCall { output, .. }) = self .entries diff --git a/crates/llm/src/message.rs b/crates/llm/src/message.rs index bfd9ba4..7a33d9b 100644 --- a/crates/llm/src/message.rs +++ b/crates/llm/src/message.rs @@ -111,6 +111,23 @@ impl Message { } } + /// Create a tool result message with structured content parts + /// (e.g. images). + pub fn tool_result_with_parts( + parts: Vec, + tool_call_id: impl Into, + ) -> Self { + Self { + role: Role::Tool, + content: None, + content_parts: (!parts.is_empty()).then_some(parts), + tool_calls: None, + tool_call_id: Some(tool_call_id.into()), + reasoning: None, + reasoning_details: None, + } + } + fn text(role: Role, content: impl Into) -> Self { Self { role, diff --git a/crates/tools/Cargo.toml b/crates/tools/Cargo.toml index e7f7c76..9403419 100644 --- a/crates/tools/Cargo.toml +++ b/crates/tools/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" [dependencies] async-trait = { workspace = true } +base64 = { workspace = true } llm = { workspace = true } serde_json = { workspace = true } thiserror = { workspace = true } diff --git a/crates/tools/src/fs.rs b/crates/tools/src/fs.rs index 00573aa..5ea0b73 100644 --- a/crates/tools/src/fs.rs +++ b/crates/tools/src/fs.rs @@ -1,5 +1,5 @@ use crate::args::{optional_u64, parse, required_string}; -use crate::tool::{ToolError, ToolExecutor}; +use crate::tool::{ToolError, ToolExecutor, ToolOutput}; use async_trait::async_trait; use llm::{ToolCall, ToolDefinition}; use serde_json::json; @@ -86,9 +86,14 @@ pub struct FileReadExecutor; #[async_trait] impl ToolExecutor for FileReadExecutor { - async fn execute(&self, call: &ToolCall) -> Result { + async fn execute(&self, call: &ToolCall) -> Result { let args = parse(call)?; let path = required_string(&args, "path")?; + + if let Some(mime_type) = image_mime_type(&path) { + return read_image(&path, mime_type).await; + } + let line_start = optional_u64(&args, "line_start", 1)?; let line_end = args .get("line_end") @@ -146,15 +151,48 @@ impl ToolExecutor for FileReadExecutor { line_number = line_number.saturating_add(1); } - Ok(content) + Ok(ToolOutput::Text(content)) + } +} + +/// Return the MIME type for supported image extensions, or `None` for +/// non-image files. +fn image_mime_type(path: &str) -> Option<&'static str> { + let ext = Path::new(path).extension()?.to_str()?.to_ascii_lowercase(); + match ext.as_str() { + "png" => Some("image/png"), + "jpg" | "jpeg" => Some("image/jpeg"), + "webp" => Some("image/webp"), + "gif" => Some("image/gif"), + _ => None, + } +} + +/// Read an image file and return it as base64-encoded `ToolOutput::Image`. +async fn read_image(path: &str, mime_type: &str) -> Result { + let bytes = tokio::fs::read(path) + .await + .map_err(|e| ToolError(format!("failed to read file {path}: {e}")))?; + + if bytes.len() > MAX_FILE_SIZE { + return Err(ToolError(format!( + "file content exceeds {MAX_FILE_SIZE} byte limit" + ))); } + + use base64::Engine; + let data = base64::engine::general_purpose::STANDARD.encode(&bytes); + Ok(ToolOutput::Image { + mime_type: mime_type.to_owned(), + data, + }) } pub struct FileEditExecutor; #[async_trait] impl ToolExecutor for FileEditExecutor { - async fn execute(&self, call: &ToolCall) -> Result { + async fn execute(&self, call: &ToolCall) -> Result { let args = parse(call)?; let path = required_string(&args, "path")?; let old_text = required_string(&args, "old_text")?; @@ -199,7 +237,7 @@ impl ToolExecutor for FileEditExecutor { .await .map_err(|e| ToolError(format!("failed to edit file {}: {e}", path)))?; - Ok(format!("File edited: {}", path)) + Ok(ToolOutput::Text(format!("File edited: {}", path))) } } @@ -234,7 +272,7 @@ pub struct FileWriteExecutor; #[async_trait] impl ToolExecutor for FileWriteExecutor { - async fn execute(&self, call: &ToolCall) -> Result { + async fn execute(&self, call: &ToolCall) -> Result { let args = parse(call)?; let path = required_string(&args, "path")?; let content = required_string(&args, "content")?; @@ -257,7 +295,7 @@ impl ToolExecutor for FileWriteExecutor { .await .map_err(|e| ToolError(format!("failed to write file {}: {e}", path)))?; - Ok(format!("File written: {}", path)) + Ok(ToolOutput::Text(format!("File written: {}", path))) } } @@ -294,6 +332,30 @@ mod tests { )) } + fn temp_path_with_ext(ext: &str) -> std::path::PathBuf { + static NEXT_ID: AtomicU64 = AtomicU64::new(0); + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before Unix epoch") + .as_nanos(); + let id = NEXT_ID.fetch_add(1, Ordering::Relaxed); + std::env::temp_dir().join(format!( + "alan-file-test-{}-{timestamp}-{id}.{ext}", + std::process::id() + )) + } + + /// Minimal 1x1 white PNG (67 bytes). + const MINIMAL_PNG: &[u8] = &[ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature + 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, // IHDR chunk + 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, + 0x53, 0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, // IDAT chunk + 0x54, 0x08, 0xD7, 0x63, 0xF8, 0xCF, 0xC0, 0x00, 0x00, 0x00, 0x02, 0x00, 0x01, 0xE2, 0x21, + 0xBC, 0x33, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, // IEND chunk + 0x44, 0xAE, 0x42, 0x60, 0x82, + ]; + #[tokio::test] async fn reads_selected_inclusive_line_range() { let path = temp_path(); @@ -304,7 +366,7 @@ mod tests { let call = tool_call(&path, r#""line_start":2,"line_end":3"#); let result = FileReadExecutor.execute(&call).await.unwrap(); - assert_eq!(result, "two\nthree\n"); + assert_eq!(result, ToolOutput::Text("two\nthree\n".into())); tokio::fs::remove_file(path).await.unwrap(); } @@ -318,7 +380,7 @@ mod tests { .await .unwrap(); - assert_eq!(result, "one\ntwo\n"); + assert_eq!(result, ToolOutput::Text("one\ntwo\n".into())); tokio::fs::remove_file(path).await.unwrap(); } @@ -356,7 +418,10 @@ mod tests { let call = edit_call(&path, "two", "TWO"); let result = FileEditExecutor.execute(&call).await.unwrap(); - assert_eq!(result, format!("File edited: {}", path.display())); + assert_eq!( + result, + ToolOutput::Text(format!("File edited: {}", path.display())) + ); assert_eq!( tokio::fs::read_to_string(&path).await.unwrap(), "one\nTWO\nthree\n" @@ -382,4 +447,98 @@ mod tests { ); tokio::fs::remove_file(path).await.unwrap(); } + + #[tokio::test] + async fn reads_png_as_image() { + let path = temp_path_with_ext("png"); + tokio::fs::write(&path, MINIMAL_PNG).await.unwrap(); + + let result = FileReadExecutor + .execute(&tool_call(&path, "")) + .await + .unwrap(); + + match &result { + ToolOutput::Image { mime_type, data } => { + assert_eq!(mime_type, "image/png"); + assert!(!data.is_empty()); + // Verify it's valid base64 that decodes back to the original bytes. + use base64::Engine; + let decoded = base64::engine::general_purpose::STANDARD + .decode(data) + .unwrap(); + assert_eq!(decoded, MINIMAL_PNG); + } + other => panic!("expected Image, got {other:?}"), + } + tokio::fs::remove_file(path).await.unwrap(); + } + + #[tokio::test] + async fn reads_jpg_extension_as_jpeg_image() { + let path = temp_path_with_ext("jpg"); + tokio::fs::write(&path, b"not-really-jpg").await.unwrap(); + + let result = FileReadExecutor + .execute(&tool_call(&path, "")) + .await + .unwrap(); + + match result { + ToolOutput::Image { mime_type, .. } => assert_eq!(mime_type, "image/jpeg"), + other => panic!("expected Image, got {other:?}"), + } + tokio::fs::remove_file(path).await.unwrap(); + } + + #[tokio::test] + async fn reads_jpeg_extension_as_jpeg_image() { + let path = temp_path_with_ext("jpeg"); + tokio::fs::write(&path, b"not-really-jpeg").await.unwrap(); + + let result = FileReadExecutor + .execute(&tool_call(&path, "")) + .await + .unwrap(); + + match result { + ToolOutput::Image { mime_type, .. } => assert_eq!(mime_type, "image/jpeg"), + other => panic!("expected Image, got {other:?}"), + } + tokio::fs::remove_file(path).await.unwrap(); + } + + #[tokio::test] + async fn reads_webp_as_image() { + let path = temp_path_with_ext("webp"); + tokio::fs::write(&path, b"not-really-webp").await.unwrap(); + + let result = FileReadExecutor + .execute(&tool_call(&path, "")) + .await + .unwrap(); + + match result { + ToolOutput::Image { mime_type, .. } => assert_eq!(mime_type, "image/webp"), + other => panic!("expected Image, got {other:?}"), + } + tokio::fs::remove_file(path).await.unwrap(); + } + + #[tokio::test] + async fn reads_gif_as_image() { + let path = temp_path_with_ext("gif"); + tokio::fs::write(&path, b"not-really-gif").await.unwrap(); + + let result = FileReadExecutor + .execute(&tool_call(&path, "")) + .await + .unwrap(); + + match result { + ToolOutput::Image { mime_type, .. } => assert_eq!(mime_type, "image/gif"), + other => panic!("expected Image, got {other:?}"), + } + tokio::fs::remove_file(path).await.unwrap(); + } } diff --git a/crates/tools/src/shell.rs b/crates/tools/src/shell.rs index 1108b47..0401d6f 100644 --- a/crates/tools/src/shell.rs +++ b/crates/tools/src/shell.rs @@ -1,5 +1,5 @@ use crate::args::{optional_u64, parse, required_string}; -use crate::tool::{ToolError, ToolExecutor}; +use crate::tool::{ToolError, ToolExecutor, ToolOutput}; use async_trait::async_trait; use llm::{ToolCall, ToolDefinition}; use serde_json::json; @@ -43,7 +43,7 @@ struct CommandOutput { #[async_trait] impl ToolExecutor for BashExecutor { - async fn execute(&self, call: &ToolCall) -> Result { + async fn execute(&self, call: &ToolCall) -> Result { let args = parse(call)?; let command = required_string(&args, "command")?; let timeout_seconds = optional_u64(&args, "timeout_seconds", DEFAULT_TIMEOUT_SECONDS)?; @@ -98,7 +98,7 @@ impl ToolExecutor for BashExecutor { let output = combine_output(&result.stdout, &result.stderr); if result.status.success() { - Ok(output) + Ok(ToolOutput::Text(output)) } else { Err(ToolError(format!( "command failed (exit {}): {output}", diff --git a/crates/tools/src/tool.rs b/crates/tools/src/tool.rs index a57f643..f180b9c 100644 --- a/crates/tools/src/tool.rs +++ b/crates/tools/src/tool.rs @@ -6,7 +6,18 @@ use thiserror::Error; #[error("tool execution failed: {0}")] pub struct ToolError(pub String); +/// The output produced by a tool executor. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ToolOutput { + Text(String), + Image { + mime_type: String, + /// Raw base64-encoded image data (no `data:` prefix). + data: String, + }, +} + #[async_trait] pub trait ToolExecutor: Send + Sync { - async fn execute(&self, call: &ToolCall) -> Result; + async fn execute(&self, call: &ToolCall) -> Result; } From 8f4938ee93529e37013ceb2b107e392a1dbe0b5b Mon Sep 17 00:00:00 2001 From: Revantark Date: Thu, 27 Aug 2026 11:02:50 +0530 Subject: [PATCH 4/5] add pasting of images into the editor --- Cargo.lock | 2 + crates/agent/src/agent/mod.rs | 2 +- crates/agent/src/agent/prompt.rs | 11 +- crates/alan/Cargo.toml | 2 + crates/alan/src/core/action.rs | 17 ++- crates/alan/src/core/chat.rs | 38 +++++-- crates/alan/src/core/controller.rs | 30 ++--- crates/alan/src/core/mod.rs | 2 +- crates/alan/src/main.rs | 3 + crates/alan/src/views/components/chat.rs | 19 ++-- crates/alan/src/views/components/footer.rs | 22 ++++ crates/alan/src/views/mod.rs | 124 +++++++++++++++++++-- crates/alan/src/views/theme.rs | 2 + crates/tools/src/fs.rs | 2 +- 14 files changed, 229 insertions(+), 47 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 23af34f..fc08116 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -42,9 +42,11 @@ dependencies = [ "anyhow", "arboard", "async-trait", + "base64", "crossterm", "futures-util", "ignore", + "image", "llm", "providers", "ratatui", diff --git a/crates/agent/src/agent/mod.rs b/crates/agent/src/agent/mod.rs index 309296f..7b6f7ba 100644 --- a/crates/agent/src/agent/mod.rs +++ b/crates/agent/src/agent/mod.rs @@ -73,7 +73,7 @@ impl Agent { "empty prompt".into(), ))) })?; - prompt::validate_not_empty(&content)?; + prompt::validate_prompt(&content, &builder.images)?; Ok(prompt::spawn_prompt_task( self, content, diff --git a/crates/agent/src/agent/prompt.rs b/crates/agent/src/agent/prompt.rs index 3d39027..94ec99e 100644 --- a/crates/agent/src/agent/prompt.rs +++ b/crates/agent/src/agent/prompt.rs @@ -213,9 +213,10 @@ fn prompt_content(content: impl Into, plan_mode: bool) -> String { } } -/// Validate that a user-facing prompt is non-empty. -pub(super) fn validate_not_empty(text: &str) -> Result<(), AgentError> { - if text.trim().is_empty() { +/// Validate that a prompt carries content: either non-empty text or at +/// least one attached image. +pub(super) fn validate_prompt(text: &str, images: &[llm::ImageUrl]) -> Result<(), AgentError> { + if text.trim().is_empty() && images.is_empty() { return Err(AgentError::Model(ModelError::Llm( llm::LlmError::Configuration("empty prompt".into()), ))); @@ -225,8 +226,8 @@ pub(super) fn validate_not_empty(text: &str) -> Result<(), AgentError> { /// Validate a user message (used by the streaming path after plan-mode suffix). fn validate_prompt_message(message: &AgentMessage) -> Result<(), AgentError> { - if let AgentMessage::User { text, .. } = message { - validate_not_empty(text)?; + if let AgentMessage::User { text, images } = message { + validate_prompt(text, images)?; } Ok(()) } diff --git a/crates/alan/Cargo.toml b/crates/alan/Cargo.toml index a6bfd2e..103cec3 100644 --- a/crates/alan/Cargo.toml +++ b/crates/alan/Cargo.toml @@ -18,6 +18,8 @@ strum = { workspace = true } async-trait = { workspace = true } unicode-width = "0.2" arboard = "=3.4.1" +base64 = { workspace = true } +image = { version = "0.25", default-features = false, features = ["png"] } tui-markdown = "0.3.9" tracing-subscriber = { workspace = true } tracing = { workspace = true } diff --git a/crates/alan/src/core/action.rs b/crates/alan/src/core/action.rs index 91d7b78..44f9524 100644 --- a/crates/alan/src/core/action.rs +++ b/crates/alan/src/core/action.rs @@ -9,6 +9,9 @@ pub enum Action { Backspace, Insert(char), Paste(String), + /// Explicit paste/attach request (Ctrl+V): attach a clipboard image if + /// present, otherwise paste clipboard text. + PasteOrAttachImage, ScrollUp, ScrollDown, MouseScrollUp, @@ -16,12 +19,24 @@ pub enum Action { TogglePlanMode, } +/// An image attached to the next prompt via clipboard paste. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ImageAttachment { + pub name: String, + pub mime_type: String, + /// Raw base64-encoded image data (no `data:` prefix). + pub base64_data: String, +} + /// Semantic commands emitted by frontend state and handled by application core. #[derive(Debug, Clone, PartialEq, Eq)] pub enum Command { Interrupt, Cancel, - Submit(String), + Submit { + text: String, + images: Vec, + }, MoveLoginSelection(isize), TogglePlanMode, } diff --git a/crates/alan/src/core/chat.rs b/crates/alan/src/core/chat.rs index 8c085f1..b4ab37b 100644 --- a/crates/alan/src/core/chat.rs +++ b/crates/alan/src/core/chat.rs @@ -1,5 +1,6 @@ //! Chat feature state and agent stream coordination. +use super::action::ImageAttachment; use super::Poll; use agent::{Agent, AgentEvent, AgentStream}; use llm::Usage; @@ -137,22 +138,41 @@ impl ChatController { self.revision = self.revision.wrapping_add(1); } - pub fn submit(&mut self, text: impl Into) { + pub fn submit(&mut self, text: impl Into, images: Vec) { let text = text.into(); let text = text.trim(); - if text.is_empty() || self.busy { + if (text.is_empty() && images.is_empty()) || self.busy { return; } self.entries.push(Entry::Prompt(text.to_owned())); self.revision = self.revision.wrapping_add(1); - self.busy = true; - let builder = self.agent.prompt().content(text.to_owned()).stream(true); - self.stream = Some( - self.agent - .ask(builder) - .expect("prompt was already validated"), - ); + let image_urls: Vec = images + .into_iter() + .map(|img| llm::ImageUrl { + url: format!("data:{};base64,{}", img.mime_type, img.base64_data), + }) + .collect(); + let builder = self + .agent + .prompt() + .content(text.to_owned()) + .images(image_urls) + .stream(true); + + match self.agent.ask(builder) { + Ok(stream) => { + self.busy = true; + self.stream = Some(stream); + } + Err(error) => { + // Replace the placeholder prompt with the failure so a + // validation mismatch between crates can never panic. + self.entries.pop(); + self.entries.push(Entry::Error(error.to_string())); + self.revision = self.revision.wrapping_add(1); + } + } } pub fn abort(&mut self) { diff --git a/crates/alan/src/core/controller.rs b/crates/alan/src/core/controller.rs index b803de5..31e97a1 100644 --- a/crates/alan/src/core/controller.rs +++ b/crates/alan/src/core/controller.rs @@ -1,6 +1,6 @@ //! UI-independent application coordinator. -use super::action::Command; +use super::action::{Command, ImageAttachment}; use super::chat::{ChatController, Entry}; use super::command::SlashCommand; use super::completion::CompletionController; @@ -134,8 +134,8 @@ impl Controller { } false } - Command::Submit(text) => { - self.submit(text); + Command::Submit { text, images } => { + self.submit(text, images); false } Command::MoveLoginSelection(delta) => { @@ -161,7 +161,7 @@ impl Controller { true } - pub fn submit(&mut self, text: String) { + pub fn submit(&mut self, text: String, images: Vec) { if self.overlay == Overlay::Login { self.login.submit(text); return; @@ -178,11 +178,11 @@ impl Controller { } let text = text.trim(); - if text.is_empty() || self.chat.is_busy() { + if (text.is_empty() && images.is_empty()) || self.chat.is_busy() { return; } - self.chat.submit(text.to_owned()); + self.chat.submit(text.to_owned(), images); } pub fn open_login(&mut self) { @@ -250,7 +250,7 @@ mod tests { #[tokio::test] async fn submit_streams_incremental_text() { let mut controller = make_controller(); - controller.submit("hi".into()); + controller.submit("hi".into(), vec![]); tokio::time::sleep(Duration::from_millis(50)).await; assert_eq!(controller.poll(), Poll::Finished); assert_eq!(controller.chat().len(), 2); @@ -300,7 +300,7 @@ mod tests { .unwrap(); let mut controller = Controller::new(Agent::builder(model).build().unwrap()); - controller.submit("hi".into()); + controller.submit("hi".into(), vec![]); tokio::time::sleep(Duration::from_millis(50)).await; assert_eq!(controller.poll(), Poll::Finished); assert_eq!(controller.chat().len(), 3); @@ -313,7 +313,7 @@ mod tests { #[test] fn submit_ignores_empty_and_busy() { let mut controller = make_controller(); - controller.submit(" ".into()); + controller.submit(" ".into(), vec![]); assert!(controller.chat().is_empty()); assert!(!controller.is_busy()); } @@ -321,7 +321,7 @@ mod tests { #[test] fn help_writes_to_the_transcript_without_reaching_the_agent() { let mut controller = make_controller(); - controller.submit("/help".into()); + controller.submit("/help".into(), vec![]); assert_eq!(controller.chat().len(), 1); assert!(matches!( @@ -337,8 +337,8 @@ mod tests { #[tokio::test] async fn info_is_not_absorbed_by_a_streaming_response() { let mut controller = make_controller(); - controller.submit("hi".into()); - controller.submit("/help".into()); + controller.submit("hi".into(), vec![]); + controller.submit("/help".into(), vec![]); tokio::time::sleep(Duration::from_millis(50)).await; controller.poll(); @@ -353,10 +353,10 @@ mod tests { let mut controller = make_controller(); assert!(!controller.plan_mode()); - controller.submit("/plan".into()); + controller.submit("/plan".into(), vec![]); assert!(controller.plan_mode()); - controller.submit("/plan".into()); + controller.submit("/plan".into(), vec![]); assert!(!controller.plan_mode()); } @@ -364,7 +364,7 @@ mod tests { #[tokio::test] async fn unknown_slash_text_is_sent_as_a_prompt() { let mut controller = make_controller(); - controller.submit("/logn".into()); + controller.submit("/logn".into(), vec![]); assert!(matches!(&controller.chat()[0], Entry::Prompt(text) if text == "/logn")); } diff --git a/crates/alan/src/core/mod.rs b/crates/alan/src/core/mod.rs index 9939d52..30a8c6d 100644 --- a/crates/alan/src/core/mod.rs +++ b/crates/alan/src/core/mod.rs @@ -7,7 +7,7 @@ pub mod completion; pub mod controller; pub mod login; -pub use action::{Action, Command}; +pub use action::{Action, Command, ImageAttachment}; pub use chat::Entry; pub use command::SlashCommand; #[cfg(test)] diff --git a/crates/alan/src/main.rs b/crates/alan/src/main.rs index abf1a27..398a662 100644 --- a/crates/alan/src/main.rs +++ b/crates/alan/src/main.rs @@ -255,6 +255,9 @@ fn action_from_event(event: &Event) -> Option { KeyCode::Char('c') if key.modifiers.contains(KeyModifiers::CONTROL) => { Action::Interrupt } + KeyCode::Char('v') if key.modifiers.contains(KeyModifiers::CONTROL) => { + Action::PasteOrAttachImage + } KeyCode::Tab | KeyCode::BackTab if key.modifiers.contains(KeyModifiers::SHIFT) || key.code == KeyCode::BackTab => diff --git a/crates/alan/src/views/components/chat.rs b/crates/alan/src/views/components/chat.rs index 234553b..7ed1684 100644 --- a/crates/alan/src/views/components/chat.rs +++ b/crates/alan/src/views/components/chat.rs @@ -156,13 +156,18 @@ fn wrap_entry(entry: &Entry, width: usize, content_width: usize) -> Vec { let mut lines = Vec::new(); lines.push(background_line("", width, theme::USER_FG, theme::USER_BG)); - for line in wrap_text(text, content_width) { - lines.push(background_line( - &line, - width, - theme::USER_FG, - theme::USER_BG, - )); + if text.is_empty() { + // Image-only submission: no text accompanied the attachment. + lines.push(indented_line("(attachment)", theme::MUTED_FG)); + } else { + for line in wrap_text(text, content_width) { + lines.push(background_line( + &line, + width, + theme::USER_FG, + theme::USER_BG, + )); + } } lines.push(background_line("", width, theme::USER_FG, theme::USER_BG)); lines diff --git a/crates/alan/src/views/components/footer.rs b/crates/alan/src/views/components/footer.rs index 05797e5..150457c 100644 --- a/crates/alan/src/views/components/footer.rs +++ b/crates/alan/src/views/components/footer.rs @@ -28,12 +28,14 @@ impl Component for Footer { frame.render_widget(background, area); let [ + attachment_area, _top_padding, status_area, _status_editor_gap, editor_area, _bottom_padding, ] = Layout::vertical([ + Constraint::Length(state.attachment_height()), Constraint::Length(1), Constraint::Length(1), Constraint::Length(1), @@ -42,6 +44,26 @@ impl Component for Footer { ]) .areas(area); + // Render attachment section when there are pending images. + if !state.attachments().is_empty() { + let mut lines: Vec> = vec![ + Line::from("\n"), + Line::from(Span::styled( + " Attachments", + Style::default().fg(theme::ATTACHMENT_FG).bold(), + )), + ]; + for attachment in state.attachments() { + lines.push(Line::from(Span::styled( + format!(" - {}", attachment.name), + Style::default().fg(theme::ATTACHMENT_FG), + ))); + } + let attachments = + Paragraph::new(Text::from(lines)).style(Style::default().bg(theme::ATTACHMENT_BG)); + frame.render_widget(attachments, attachment_area); + } + let (indicator, indicator_style, shortcuts) = if controller.is_busy() { ( " ● thinking", diff --git a/crates/alan/src/views/mod.rs b/crates/alan/src/views/mod.rs index 8a7133e..7013894 100644 --- a/crates/alan/src/views/mod.rs +++ b/crates/alan/src/views/mod.rs @@ -8,7 +8,10 @@ mod components; pub mod selection; mod theme; -use crate::core::{Action, Command, CompletionController, Controller, Overlay, Poll, SlashCommand}; +use crate::core::{ + Action, Command, CompletionController, Controller, ImageAttachment, Overlay, Poll, SlashCommand, +}; +use base64::Engine; use components::{Chat, Footer, Header, LoginOverlay}; use crossterm::event::{ Event, KeyCode, KeyEvent, KeyEventKind, KeyModifiers, MouseButton, MouseEvent, MouseEventKind, @@ -46,6 +49,8 @@ pub struct UiState { last_click: Option<(Instant, u16, u16)>, /// True when something changed since last draw and a redraw is needed. dirty: bool, + /// Images attached to the next prompt via clipboard paste. + attachments: Vec, } impl UiState { @@ -64,6 +69,7 @@ impl UiState { selection: None, last_click: None, dirty: true, + attachments: Vec::new(), } } @@ -78,10 +84,15 @@ impl UiState { Action::Submit => { self.follow_output = true; self.scroll_target = self.max_scroll; - Some(Command::Submit(std::mem::take(&mut self.input))) + let images = std::mem::take(&mut self.attachments); + Some(Command::Submit { + text: std::mem::take(&mut self.input), + images, + }) } Action::ClearInput => { self.input.clear(); + self.attachments.clear(); Some(Command::Cancel) } Action::Backspace => { @@ -96,6 +107,17 @@ impl UiState { self.input.push_str(&text); None } + Action::PasteOrAttachImage => { + if self.try_clipboard_image() { + None + } else { + // No image on the clipboard: fall back to pasting text. + match arboard::Clipboard::new().and_then(|mut c| c.get_text()) { + Ok(text) if !text.is_empty() => self.apply(Action::Paste(text), false), + _ => None, + } + } + } Action::ScrollUp => { if login_selection_active { Some(Command::MoveLoginSelection(-1)) @@ -185,6 +207,25 @@ impl UiState { { Some(Command::Interrupt) } + Event::Key(key) + if key.code == KeyCode::Char('v') + && key.modifiers.contains(KeyModifiers::CONTROL) => + { + // Bracketed paste only delivers text; images never arrive as + // an `Event::Paste`, so attaching needs an explicit trigger. + if !self.try_clipboard_image() { + // No image on the clipboard: fall back to pasting text. + let text = arboard::Clipboard::new().and_then(|mut c| c.get_text()); + if let Ok(text) = text + && !text.is_empty() + { + self.editor.insert_str(text); + self.dirty = true; + self.sync_completion(completion); + } + } + None + } Event::Key(key) if key.code == KeyCode::Char('u') && key.modifiers.contains(KeyModifiers::CONTROL) => @@ -205,9 +246,9 @@ impl UiState { } Event::Mouse(mouse) => self.handle_mouse_event(mouse, rendered_lines), Event::Paste(text) => { - self.editor.insert_str(text); self.dirty = true; self.sync_completion(completion); + self.editor.insert_str(text); None } event => { @@ -218,6 +259,7 @@ impl UiState { } }; self.sync_command_highlight(); + command } @@ -472,9 +514,10 @@ impl UiState { self.follow_output = true; self.scroll_target = self.max_scroll; let text = self.editor_text(); + let images = std::mem::take(&mut self.attachments); self.editor = Self::new_editor(); self.dirty = true; - Some(Command::Submit(text)) + Some(Command::Submit { text, images }) } /// The prompt soft-wraps at word boundaries and grows up to @@ -571,6 +614,61 @@ impl UiState { pub fn take_dirty(&mut self) -> bool { std::mem::take(&mut self.dirty) } + + /// Check the system clipboard for an image and add it as an attachment. + /// Returns true when an image was attached. + fn try_clipboard_image(&mut self) -> bool { + let img = match arboard::Clipboard::new().and_then(|mut clipboard| clipboard.get_image()) { + Ok(img) => img, + Err(error) => { + tracing::debug!(%error, "clipboard: no readable image"); + return false; + } + }; + if img.width == 0 || img.height == 0 { + tracing::debug!("clipboard: ignoring zero-size image"); + return false; + } + let (w, h) = (img.width, img.height); + let raw = img.into_owned_bytes(); + + let Some(rgba) = image::RgbaImage::from_raw(w as u32, h as u32, raw.into_owned()) else { + tracing::debug!(width = w, height = h, "clipboard: invalid image data"); + return false; + }; + let mut png_buf = std::io::Cursor::new(Vec::new()); + if let Err(error) = + image::DynamicImage::ImageRgba8(rgba).write_to(&mut png_buf, image::ImageFormat::Png) + { + tracing::debug!(%error, "clipboard: PNG encoding failed"); + return false; + } + let data = base64::engine::general_purpose::STANDARD.encode(png_buf.get_ref()); + + self.attachments.push(ImageAttachment { + name: format!("image-{}", self.attachments.len() + 1), + mime_type: "image/png".into(), + base64_data: data, + }); + tracing::debug!( + name = self.attachments.last().unwrap().name, + "clipboard: image attached" + ); + self.dirty = true; + true + } + + pub fn attachments(&self) -> &[ImageAttachment] { + &self.attachments + } + + pub fn attachment_height(&self) -> u16 { + if self.attachments.is_empty() { + 0 + } else { + 3 + self.attachments.len() as u16 + } + } } impl Default for UiState { @@ -628,7 +726,7 @@ impl AppView { let [header_area, chat_area, footer_area] = Layout::vertical([ Constraint::Length(1), Constraint::Min(1), - Constraint::Length(4 + state.editor_rows(editor_width)), + Constraint::Length(4 + state.editor_rows(editor_width) + state.attachment_height()), ]) .areas(frame.area()); @@ -1003,7 +1101,13 @@ mod tests { ), )); - assert_eq!(command, Some(Command::Submit("h".into()))); + assert_eq!( + command, + Some(Command::Submit { + text: "h".into(), + images: vec![] + }) + ); assert!(state.take_dirty()); } @@ -1097,7 +1201,13 @@ mod tests { // Popup closed: Enter submits again. let command = state.handle_event(key(KeyCode::Enter), &[], &mut completion); - assert_eq!(command, Some(Command::Submit("@".into()))); + assert_eq!( + command, + Some(Command::Submit { + text: "@".into(), + images: vec![] + }) + ); } #[test] diff --git a/crates/alan/src/views/theme.rs b/crates/alan/src/views/theme.rs index 7f71538..8647e1f 100644 --- a/crates/alan/src/views/theme.rs +++ b/crates/alan/src/views/theme.rs @@ -21,3 +21,5 @@ pub const PROMPT_GUTTER: u16 = 4; pub const SELECTION_BG: Color = Color::Rgb(58, 76, 107); pub const SELECTION_FG: Color = Color::White; pub const COMMAND_FG: Color = Color::Cyan; +pub const ATTACHMENT_BG: Color = Color::Rgb(50, 46, 28); +pub const ATTACHMENT_FG: Color = Color::Rgb(200, 190, 140); diff --git a/crates/tools/src/fs.rs b/crates/tools/src/fs.rs index 5ea0b73..1409bf1 100644 --- a/crates/tools/src/fs.rs +++ b/crates/tools/src/fs.rs @@ -11,7 +11,7 @@ const MAX_FILE_SIZE: usize = 1024 * 1024; pub fn file_read_definition() -> ToolDefinition { ToolDefinition { name: "read".into(), - description: "Read a file, optionally limited to an inclusive line range".into(), + description: "Read a file/image, optionally limited to an inclusive line range".into(), parameters: json!({ "type": "object", "properties": { From 544817e34a25751d32a15e2845b7d397f19722eb Mon Sep 17 00:00:00 2001 From: Revantark Date: Thu, 27 Aug 2026 16:51:24 +0530 Subject: [PATCH 5/5] fix fmt --- crates/agent/src/agent/tool_loop.rs | 6 +----- crates/alan/src/core/chat.rs | 2 +- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/crates/agent/src/agent/tool_loop.rs b/crates/agent/src/agent/tool_loop.rs index 610340d..243cbf4 100644 --- a/crates/agent/src/agent/tool_loop.rs +++ b/crates/agent/src/agent/tool_loop.rs @@ -133,11 +133,7 @@ async fn handle_tool_calls( super::persistence::append_context_message( agent, context, - AgentMessage::tool_result_with_parts( - &call_id, - &output_text, - content_parts, - ), + AgentMessage::tool_result_with_parts(&call_id, &output_text, content_parts), ) .await?; diff --git a/crates/alan/src/core/chat.rs b/crates/alan/src/core/chat.rs index b4ab37b..350ef56 100644 --- a/crates/alan/src/core/chat.rs +++ b/crates/alan/src/core/chat.rs @@ -1,7 +1,7 @@ //! Chat feature state and agent stream coordination. -use super::action::ImageAttachment; use super::Poll; +use super::action::ImageAttachment; use agent::{Agent, AgentEvent, AgentStream}; use llm::Usage; use std::sync::Arc;