Skip to main content

robit_agent/
agent.rs

1//! Agent — the event-driven loop that orchestrates LLM calls and tool execution.
2
3use async_openai::types::chat::{
4    ChatCompletionMessageToolCall, ChatCompletionMessageToolCalls,
5    ChatCompletionRequestAssistantMessage, ChatCompletionRequestMessage,
6    ChatCompletionRequestSystemMessage, ChatCompletionRequestToolMessage,
7    ChatCompletionRequestUserMessage, ChatCompletionRequestUserMessageContent,
8    ChatCompletionRequestUserMessageContentPart,
9    ChatCompletionRequestMessageContentPartText,
10    ChatCompletionRequestMessageContentPartImage,
11    FunctionCall,
12};
13
14// Import ImageUrl from wherever it is in async-openai 0.41
15use async_openai::types::chat::ImageUrl;
16use futures_util::StreamExt;
17use robit_ai::config::ContextConfig;
18use robit_ai::LlmClient;
19use std::any::Any;
20use std::collections::HashMap;
21use std::path::PathBuf;
22use std::sync::Arc;
23use tokio::sync::mpsc;
24
25use crate::context::{ContextManager, TruncationAction, TruncationResult};
26use crate::error::{AgentError, Result};
27use crate::event::{new_session_id, AgentEvent, FrontendMessage, MediaAttachment, SessionId};
28use crate::frontend::Frontend;
29use crate::media;
30use crate::prompt::PromptBuilder;
31use crate::skill::SkillRegistry;
32use crate::tool::async_runner::{AsyncTaskDone, AsyncTaskRunner};
33use crate::tool::task_registry::{AsyncTaskRecord, AsyncTaskStatus, TaskRegistry};
34use crate::tool::{ToolCallInfo, ToolContext, ToolImage, ToolRegistry, ToolResult};
35use tokio_util::sync::CancellationToken;
36
37// ============================================================================
38// AgentSession
39// ============================================================================
40
41/// A single conversation session with its own message history.
42pub struct AgentSession {
43    pub session_id: SessionId,
44    pub history: Vec<ChatCompletionRequestMessage>,
45    pub working_dir: PathBuf,
46    /// The last known exact prompt token count from API `usage.prompt_tokens`.
47    /// Used as a calibration anchor: `last_known_prompt_tokens` is the precise
48    /// token count for the first `snapshot_message_count` messages in history.
49    /// New messages appended after the snapshot are estimated incrementally.
50    /// Invalidated (set to `None`) when history is truncated/compressed.
51    pub last_known_prompt_tokens: Option<u32>,
52    /// Number of messages in `history` that `last_known_prompt_tokens` covers.
53    /// If `history.len() < snapshot_message_count`, the snapshot is stale
54    /// (truncation/compression happened) and calibration is invalid.
55    pub snapshot_message_count: usize,
56}
57
58impl AgentSession {
59    fn new(session_id: SessionId, working_dir: PathBuf, system_prompt: String) -> Self {
60        let system_msg = ChatCompletionRequestMessage::System(
61            ChatCompletionRequestSystemMessage {
62                content: system_prompt.into(),
63                name: None,
64            }
65            .into(),
66        );
67
68        Self {
69            session_id,
70            history: vec![system_msg],
71            working_dir,
72            last_known_prompt_tokens: None,
73            snapshot_message_count: 0,
74        }
75    }
76
77    /// Create session with pre-loaded history
78    pub fn with_history(
79        session_id: SessionId,
80        working_dir: PathBuf,
81        system_prompt: String,
82        history: Vec<ChatCompletionRequestMessage>,
83    ) -> Self {
84        // Create system message (new one with latest config)
85        let system_msg = ChatCompletionRequestMessage::System(
86            ChatCompletionRequestSystemMessage {
87                content: system_prompt.into(),
88                name: None,
89            }
90            .into(),
91        );
92
93        // Prepend new system message to history
94        let mut full_history = vec![system_msg];
95        full_history.extend(history);
96
97        Self {
98            session_id,
99            history: full_history,
100            working_dir,
101            last_known_prompt_tokens: None,
102            snapshot_message_count: 0,
103        }
104    }
105}
106
107// ============================================================================
108// Agent
109// ============================================================================
110
111/// The Agent orchestrates LLM calls and tool execution.
112pub struct Agent {
113    llm_client: Arc<LlmClient>,
114    tools: Arc<ToolRegistry>,
115    skills: Arc<SkillRegistry>,
116    sessions: HashMap<SessionId, AgentSession>,
117    default_session_id: SessionId,
118    context_manager: ContextManager,
119    frontend: Arc<dyn Frontend>,
120    auto_approve: bool,
121    /// Platform-specific extensions passed to ToolContext during tool execution.
122    extensions: HashMap<String, Arc<dyn Any + Send + Sync>>,
123    /// Pending truncation result that needs compression (handled at start of run loop).
124    pending_truncation: Option<(SessionId, crate::context::TruncationResult)>,
125    /// Handle for tools to submit async background work. Cloned into every
126    /// ToolContext; the matching `done_rx` is drained in `run()`.
127    async_runner: AsyncTaskRunner,
128    /// Receiver for completed async tasks. `Option` so `run()` can `take()` it
129    /// into a local, avoiding borrowing `self` across the `select!` loop body.
130    done_rx: Option<mpsc::Receiver<AsyncTaskDone>>,
131    /// In-flight async tasks keyed by task_id, with their cancel tokens.
132    pending_tasks: HashMap<String, PendingTask>,
133    /// Shared registry of async task statuses, read by the `query_task` tool
134    /// via ToolContext. Cloned per ToolContext (shared `Arc`).
135    task_registry: TaskRegistry,
136}
137
138/// Bookkeeping for one in-flight async task.
139struct PendingTask {
140    cancel: CancellationToken,
141    tool_name: String,
142}
143
144impl Agent {
145    /// Create a new Agent with the given dependencies.
146    pub fn new(
147        llm_client: Arc<LlmClient>,
148        tools: Arc<ToolRegistry>,
149        skills: Arc<SkillRegistry>,
150        frontend: Arc<dyn Frontend>,
151        context_config: Option<&ContextConfig>,
152        context_window: Option<u64>,
153        working_dir: PathBuf,
154        auto_approve: bool,
155        extensions: HashMap<String, Arc<dyn Any + Send + Sync>>,
156    ) -> Self {
157        let prompt_builder = PromptBuilder::with_working_dir(Some(&working_dir));
158        let context_manager = ContextManager::new(context_window, context_config);
159
160        // Build system prompt with skills. Tools are NOT listed in the prompt;
161        // they are exposed via the function-calling `tools` request parameter.
162        let skill_descs = skills.skill_descriptions();
163        let system_prompt = prompt_builder.build_system_prompt(&skill_descs, &working_dir);
164
165        // Create default session
166        let session_id = new_session_id();
167        let session = AgentSession::new(session_id.clone(), working_dir, system_prompt);
168
169        let mut sessions = HashMap::new();
170        sessions.insert(session_id.clone(), session);
171
172        let (done_tx, done_rx) = mpsc::channel::<AsyncTaskDone>(32);
173        let async_runner = AsyncTaskRunner::new(done_tx);
174        let task_registry = TaskRegistry::new();
175
176        Self {
177            llm_client,
178            tools,
179            skills,
180            sessions,
181            default_session_id: session_id,
182            context_manager,
183            frontend,
184            auto_approve,
185            extensions,
186            pending_truncation: None,
187            async_runner,
188            done_rx: Some(done_rx),
189            pending_tasks: HashMap::new(),
190            task_registry,
191        }
192    }
193
194    /// Create Agent with pre-loaded history (for resuming sessions)
195    pub fn with_history(
196        llm_client: Arc<LlmClient>,
197        tools: Arc<ToolRegistry>,
198        skills: Arc<SkillRegistry>,
199        frontend: Arc<dyn Frontend>,
200        context_config: Option<&ContextConfig>,
201        context_window: Option<u64>,
202        working_dir: PathBuf,
203        auto_approve: bool,
204        extensions: HashMap<String, Arc<dyn Any + Send + Sync>>,
205        session_id: SessionId,
206        history: Vec<ChatCompletionRequestMessage>,
207    ) -> Self {
208        tracing::info!(
209            "Agent::with_history: session_id={}, received {} history messages",
210            session_id,
211            history.len()
212        );
213        let prompt_builder = PromptBuilder::with_working_dir(Some(&working_dir));
214        let context_manager = ContextManager::new(context_window, context_config);
215
216        // Build system prompt with skills. Tools are NOT listed in the prompt;
217        // they are exposed via the function-calling `tools` request parameter.
218        let skill_descs = skills.skill_descriptions();
219        let system_prompt = prompt_builder.build_system_prompt(&skill_descs, &working_dir);
220
221        // Create session with history
222        let mut session = AgentSession::with_history(
223            session_id.clone(),
224            working_dir,
225            system_prompt,
226            history,
227        );
228
229        tracing::debug!(
230            "Agent::with_history: after adding system prompt, session history length = {}",
231            session.history.len()
232        );
233        // Sanitize history: remove image_url parts if the model doesn't support images.
234        // This prevents 400 errors from APIs that only accept text content.
235        let supports_images = llm_client.supports_images();
236        sanitize_history_for_model(&mut session.history, supports_images);
237        // Apply context truncation before starting
238        let truncation_result = context_manager.maybe_truncate(
239            &mut session.history,
240            session.last_known_prompt_tokens,
241            session.snapshot_message_count,
242        );
243        if truncation_result.rounds_removed > 0 {
244            tracing::info!(
245                "Agent::with_history: truncated {} rounds ({} messages), needs_compression={}",
246                truncation_result.rounds_removed,
247                truncation_result.messages_removed,
248                truncation_result.needs_compression
249            );
250        }
251        tracing::debug!(
252            "Agent::with_history: after truncation, session history length = {}",
253            session.history.len()
254        );
255
256        let pending_truncation = if truncation_result.needs_compression {
257            Some((session_id.clone(), truncation_result))
258        } else {
259            None
260        };
261
262        let mut sessions = HashMap::new();
263        sessions.insert(session_id.clone(), session);
264
265        let (done_tx, done_rx) = mpsc::channel::<AsyncTaskDone>(32);
266        let async_runner = AsyncTaskRunner::new(done_tx);
267        let task_registry = TaskRegistry::new();
268
269        Self {
270            llm_client,
271            tools,
272            skills,
273            sessions,
274            default_session_id: session_id,
275            context_manager,
276            frontend,
277            auto_approve,
278            extensions,
279            pending_truncation,
280            async_runner,
281            done_rx: Some(done_rx),
282            pending_tasks: HashMap::new(),
283            task_registry,
284        }
285    }
286
287    /// Run the agent's main event loop. Takes ownership of the message receiver.
288    /// Returns when the channel is closed or user types /exit.
289    pub async fn run(mut self, mut message_rx: mpsc::Receiver<FrontendMessage>) {
290        tracing::info!("Agent started, session: {}", self.default_session_id);
291
292        // Handle pending compression from with_history initialization.
293        // May need multiple compression rounds for long histories.
294        if self.pending_truncation.is_some() {
295            tracing::info!("=== Starting pending compression processing ===");
296            let session_id = self.default_session_id.clone();
297            let mut iterations = 0;
298            const MAX_COMPRESSION_ITERATIONS: usize = 20;
299
300            loop {
301                // Take one pending result, if any
302                let pending = self.pending_truncation.take();
303                let result = match pending {
304                    Some((_, r)) => r,
305                    None => break,
306                };
307
308                iterations += 1;
309                if iterations > MAX_COMPRESSION_ITERATIONS {
310                    tracing::warn!("Reached max compression iterations ({}), stopping", MAX_COMPRESSION_ITERATIONS);
311                    break;
312                }
313
314                tracing::info!("Compression iteration {}: action={:?}, removed_rounds={}, removed_msgs={}",
315                    iterations, result.action, result.rounds_removed, result.messages_removed);
316
317                // Apply the compression result (generate summary / merge)
318                if let Some(session) = self.sessions.get_mut(&session_id) {
319                    apply_compression_result(&self.llm_client, &mut session.history, &result).await;
320                    // Compression changed history — invalidate calibration
321                    session.last_known_prompt_tokens = None;
322                    session.snapshot_message_count = 0;
323                }
324
325                // Check if more compression is needed
326                let needs_more = if let Some(session) = self.sessions.get(&session_id) {
327                    let estimated = self.context_manager.estimate_context_tokens(
328                        &session.history,
329                        session.last_known_prompt_tokens,
330                        session.snapshot_message_count,
331                    );
332                    estimated > self.context_manager.truncation_threshold()
333                } else {
334                    false
335                };
336
337                if !needs_more {
338                    tracing::info!("Context below threshold after {} compression iterations", iterations);
339                    break;
340                }
341
342                // Do another round of truncation
343                if let Some(session) = self.sessions.get_mut(&session_id) {
344                    let next_result = self.context_manager.maybe_truncate(
345                        &mut session.history,
346                        session.last_known_prompt_tokens,
347                        session.snapshot_message_count,
348                    );
349                    if next_result.needs_compression {
350                        self.pending_truncation = Some((session_id.clone(), next_result));
351                    } else if next_result.messages_removed > 0 {
352                        // Truncation happened but no compression needed (e.g. discard)
353                        tracing::info!("Truncation without compression: {} messages removed", next_result.messages_removed);
354                        // Continue the loop to check if still over threshold
355                        self.pending_truncation = Some((session_id.clone(), next_result));
356                    } else {
357                        break;
358                    }
359                }
360            }
361
362            tracing::info!("=== Compression processing finished ({} iterations) ===", iterations);
363        } else {
364            tracing::debug!("No pending compression needed");
365        }
366
367        // Take done_rx out of self so the select! loop can borrow it without
368        // borrowing all of `self` (which would conflict with the &mut self
369        // calls made inside the message branch).
370        let mut done_rx = self
371            .done_rx
372            .take()
373            .expect("done_rx is consumed exactly once in run()");
374
375        loop {
376            tokio::select! {
377                msg = message_rx.recv() => {
378                    let Some(msg) = msg else { break; };
379                    match msg {
380                        FrontendMessage::UserInput { text, attachments } => {
381                            if text == "/exit" || text == "/quit" {
382                                break;
383                            }
384                            if text == "/clear" {
385                                self.clear_session();
386                                let _ = self
387                                    .frontend
388                                    .on_event(AgentEvent::TextDelta(
389                                        "\n[Conversation history cleared]\n".to_string(),
390                                    ))
391                                    .await;
392                                let _ = self.frontend.on_event(AgentEvent::TurnComplete).await;
393                                continue;
394                            }
395
396                            // Check for skill trigger
397                            if let Some((skill, args)) = self.skills.match_trigger(&text) {
398                                let skill = skill.clone();
399                                self.run_skill_turn(&skill, &args).await;
400                                continue;
401                            }
402
403                            self.run_turn(&text, attachments).await;
404                        }
405                        FrontendMessage::Cancel => {
406                            // Cancel all in-flight async tasks for this Agent.
407                            self.handle_cancel_all().await;
408                        }
409                        FrontendMessage::CancelTask { task_id } => {
410                            self.handle_cancel_task(&task_id).await;
411                        }
412                        FrontendMessage::ConfirmationResponse { .. } => {
413                            // Confirmation is handled via frontend.request_tool_confirmation()
414                            // within run_one_step. This variant is reserved for future use.
415                            tracing::warn!("Unexpected ConfirmationResponse outside tool confirmation");
416                        }
417                    }
418                }
419                done = done_rx.recv() => {
420                    let Some(done) = done else { break; };
421                    self.handle_async_done(done).await;
422                }
423            }
424        }
425
426        // ── Drain phase: cancel pending async tasks and collect their results ──
427        // Without this, `done_rx` is dropped when `self` goes out of scope,
428        // causing any in-flight async tasks (e.g. image generation) to have
429        // their results silently lost. We cancel first (fast), then give
430        // tasks a brief window to deliver their final status through done_rx.
431        if !self.pending_tasks.is_empty() {
432            let remaining = self.pending_tasks.len();
433            tracing::warn!(
434                "[async] Agent exiting with {} pending task(s), cancelling and draining...",
435                remaining
436            );
437            // Cancel all pending tasks so they finish quickly.
438            for (_, pending) in self.pending_tasks.drain() {
439                pending.cancel.cancel();
440            }
441            // Collect results from the done channel. Each task should respond
442            // to cancellation within seconds; use a bounded window per task.
443            let drain_deadline = tokio::time::Instant::now()
444                + tokio::time::Duration::from_secs(5);
445            while tokio::time::Instant::now() < drain_deadline {
446                match tokio::time::timeout(
447                    tokio::time::Duration::from_millis(500),
448                    done_rx.recv(),
449                )
450                .await
451                {
452                    Ok(Some(done)) => {
453                        tracing::info!(
454                            "[async] drained result after shutdown: task_id={}, tool={}, cancelled={}",
455                            done.task_id, done.tool_name, done.cancelled
456                        );
457                        self.handle_async_done(done).await;
458                    }
459                    Ok(None) => {
460                        // All senders dropped — no more results coming.
461                        tracing::debug!("[async] done_tx closed during drain");
462                        break;
463                    }
464                    Err(_) => {
465                        // Per-iteration timeout — loop back to check deadline.
466                    }
467                }
468            }
469            tracing::info!("[async] drain phase complete");
470        }
471
472        tracing::info!("Agent stopped");
473    }
474
475    /// Execute a single turn: user input -> LLM call(s) -> tool execution(s) -> response.
476    async fn run_turn(&mut self, user_input: &str, attachments: Vec<MediaAttachment>) {
477        let session_id = self.default_session_id.clone();
478
479        // Build user message first (to avoid borrow conflict)
480        let user_message = self.build_user_message(user_input, &attachments).await;
481
482        // Add user message to history
483        if let Some(session) = self.sessions.get_mut(&session_id) {
484            session.history.push(user_message);
485        }
486
487        // Run the agentic loop (may iterate if LLM calls tools).
488        self.run_agent_loop(&session_id).await;
489    }
490
491    /// Run the agentic loop: call LLM, execute tools, repeat until the LLM
492    /// produces a final response (no tool calls) or a safety limit is hit.
493    /// Shared by user-input turns and async-task-completion reinjection.
494    async fn run_agent_loop(&mut self, session_id: &SessionId) {
495        let max_tool_calls = self.context_manager.max_tool_calls_per_turn;
496        let max_iterations = 20;
497        let mut total_tool_calls = 0usize;
498        for iteration in 0..max_iterations {
499            match self.run_one_step(session_id).await {
500                Ok(0) => {
501                    let _ = self.frontend.on_event(AgentEvent::TurnComplete).await;
502                    return;
503                }
504                Ok(tool_call_count) => {
505                    total_tool_calls += tool_call_count;
506
507                    // Check against per-turn tool call limit
508                    if total_tool_calls >= max_tool_calls {
509                        tracing::warn!(
510                            "Tool call limit reached: {} >= {} (max_tool_calls_per_turn), forcing turn completion",
511                            total_tool_calls,
512                            max_tool_calls
513                        );
514                        let _ = self
515                            .frontend
516                            .on_event(AgentEvent::TextDelta(
517                                format!(
518                                    "\n\n[Tool call limit reached ({} calls). Please summarize progress and continue in the next message.]\n",
519                                    total_tool_calls
520                                ),
521                            ))
522                            .await;
523                        let _ = self.frontend.on_event(AgentEvent::TurnComplete).await;
524                        return;
525                    }
526
527                    tracing::debug!(
528                        "Iteration {}: {} tool calls executed (total: {}/{}), continuing loop",
529                        iteration,
530                        tool_call_count,
531                        total_tool_calls,
532                        max_tool_calls
533                    );
534                }
535                Err(e) => {
536                    let _ = self.frontend.on_event(AgentEvent::Error(e)).await;
537                    let _ = self.frontend.on_event(AgentEvent::TurnComplete).await;
538                    return;
539                }
540            }
541        }
542
543        // Safety limit
544        let _ = self
545            .frontend
546            .on_event(AgentEvent::Error(AgentError::InternalError(
547                format!("Max iterations reached ({})", max_iterations),
548            )))
549            .await;
550        let _ = self.frontend.on_event(AgentEvent::TurnComplete).await;
551    }
552
553    /// Run one step: call LLM, process response, execute tools.
554    /// Returns the number of tool calls executed (0 = turn complete, no tools called).
555    async fn run_one_step(&mut self, session_id: &SessionId) -> Result<usize> {
556        let session = self
557            .sessions
558            .get_mut(session_id)
559            .ok_or_else(|| AgentError::InternalError("Session not found".to_string()))?;
560
561        // Truncate context if needed
562        let truncation_result = self.context_manager.maybe_truncate(
563            &mut session.history,
564            session.last_known_prompt_tokens,
565            session.snapshot_message_count,
566        );
567
568        // Handle compression: generate actual summary / merge via LLM
569        if truncation_result.needs_compression {
570            apply_compression_result(&self.llm_client, &mut session.history, &truncation_result).await;
571            // Truncation changed history — invalidate calibration
572            session.last_known_prompt_tokens = None;
573            session.snapshot_message_count = 0;
574
575            tracing::info!(
576                "Compression completed: action={:?}, removed_rounds={}",
577                truncation_result.action, truncation_result.rounds_removed
578            );
579        } else if truncation_result.messages_removed > 0 {
580            // Messages were removed (e.g. discard) — also invalidate calibration
581            session.last_known_prompt_tokens = None;
582            session.snapshot_message_count = 0;
583
584            tracing::info!(
585                "Context truncated without compression: {} messages removed",
586                truncation_result.messages_removed
587            );
588        }
589
590        // Build tool schemas. Only sent when the model supports function
591        // calling (`supports_tools = true` in config): OpenAI-compatible
592        // providers that don't understand the `tools` parameter reject the
593        // whole request, and models without tool support can't act on them.
594        let tools_param = if self.llm_client.supports_tools() {
595            let tool_schemas = self.tools.tool_schemas();
596            if tool_schemas.is_empty() {
597                None
598            } else {
599                Some(tool_schemas)
600            }
601        } else {
602            None
603        };
604
605        // Log estimated token usage before call (uses calibrated estimation when available)
606        let estimated_prompt = self.context_manager.estimate_context_tokens(
607            &session.history,
608            session.last_known_prompt_tokens,
609            session.snapshot_message_count,
610        );
611        let calibration_tag = if session.last_known_prompt_tokens.is_some() {
612            "calibrated"
613        } else {
614            "heuristic"
615        };
616        tracing::info!(
617            "LLM call: ~{} prompt tokens ({}), {} messages",
618            estimated_prompt,
619            calibration_tag,
620            session.history.len(),
621        );
622
623        // Sanitize history right before API call — some providers reject
624        // image_url content even when supports_images was checked earlier
625        // (e.g. model switch mid-session, or stale history from DB).
626        if !self.llm_client.supports_images() {
627            sanitize_history_for_model(&mut session.history, false);
628        }
629
630        // Call LLM (streaming)
631        let mut stream = match self
632            .llm_client
633            .chat_stream(session.history.clone(), tools_param)
634            .await
635        {
636            Ok(s) => s,
637            Err(e) => {
638                tracing::error!("LLM chat_stream failed: {:?}", e);
639                return Err(e.into());
640            }
641        };
642        tracing::trace!("LLM stream obtained, starting to collect response");
643
644        // Collect streaming response
645        let mut full_text = String::new();
646        let mut tool_call_chunks: HashMap<usize, ToolCallAccumulator> = HashMap::new();
647        let mut api_usage: Option<async_openai::types::chat::CompletionUsage> = None;
648
649        let mut chunk_count = 0;
650        while let Some(chunk_result) = stream.next().await {
651            let chunk = match chunk_result {
652                Ok(c) => c,
653                Err(e) => {
654                    // Providers sometimes push `{"error": ...}` events (e.g.
655                    // content moderation) into the stream; recover the real
656                    // cause from the raw payload instead of surfacing a bare
657                    // JSON deserialization failure.
658                    let llm_error = robit_ai::LlmError::from_openai_error(e);
659                    tracing::error!("Stream chunk error: {}", llm_error);
660                    return Err(AgentError::LlmError(llm_error));
661                }
662            };
663            chunk_count += 1;
664
665            // Capture usage info if present in this chunk (some providers include it in final chunk)
666            if let Some(ref usage) = chunk.usage {
667                api_usage = Some(usage.clone());
668            }
669
670            if let Some(choice) = chunk.choices.first() {
671                // Text content
672                if let Some(content) = &choice.delta.content {
673                    full_text.push_str(content);
674                    let _ = self
675                        .frontend
676                        .on_event(AgentEvent::TextDelta(content.clone()))
677                        .await;
678                }
679
680                // Tool call deltas
681                if let Some(tool_calls) = &choice.delta.tool_calls {
682                    for tc in tool_calls {
683                        let acc = tool_call_chunks
684                            .entry(tc.index as usize)
685                            .or_insert_with(ToolCallAccumulator::new);
686
687                        if let Some(id) = &tc.id {
688                            // 只有当id非空时才更新
689                            if !id.is_empty() {
690                                acc.id = Some(id.clone());
691                            }
692                        }
693                        if let Some(function) = &tc.function {
694                            if let Some(name) = &function.name {
695                                // 只有当name非空时才更新
696                                if !name.is_empty() {
697                                    acc.name = Some(name.clone());
698                                }
699                            }
700                            if let Some(args) = &function.arguments {
701                                acc.arguments.push_str(args);
702                            }
703                        }
704                    }
705                }
706            }
707        }
708
709        tracing::debug!("Stream collection complete: {} chunks, {} chars of text", chunk_count, full_text.len());
710
711        // Assemble complete tool calls from chunks
712        let assembled_tool_calls: Vec<ChatCompletionMessageToolCall> = {
713            let mut indices: Vec<usize> = tool_call_chunks.keys().cloned().collect();
714            indices.sort();
715            indices
716                .into_iter()
717                .filter_map(|idx| tool_call_chunks.remove(&idx)?.into_tool_call())
718                .collect()
719        };
720
721        // Log token usage summary
722        let estimated_response = crate::context::estimate_tokens(&full_text);
723        if let Some(ref usage) = api_usage {
724            tracing::info!(
725                "LLM response: API usage = {} prompt + {} completion = {} total tokens. Estimated: ~{} prompt + ~{} response = ~{} total",
726                usage.prompt_tokens,
727                usage.completion_tokens,
728                usage.total_tokens,
729                estimated_prompt,
730                estimated_response,
731                estimated_prompt + estimated_response,
732            );
733        } else {
734            tracing::info!(
735                "LLM response: {} chars, ~{} estimated tokens ({} tool calls). API usage not available from streaming.",
736                full_text.len(),
737                estimated_response,
738                assembled_tool_calls.len(),
739            );
740        }
741
742        // Calibrate token estimation: store the API-reported prompt_tokens as a
743        // precise baseline. At this point session.history still reflects exactly
744        // what was sent to the API (assistant_msg hasn't been pushed yet), so
745        // prompt_tokens is the exact token count for session.history.
746        if let Some(ref usage) = api_usage {
747            session.last_known_prompt_tokens = Some(usage.prompt_tokens);
748            session.snapshot_message_count = session.history.len();
749            tracing::trace!(
750                "Token calibration updated: prompt_tokens={} at {} messages",
751                usage.prompt_tokens, session.history.len()
752            );
753        }
754
755        // Add assistant message to history
756        let content = if full_text.is_empty() {
757            None
758        } else {
759            Some(full_text.clone().into())
760        };
761        let tool_calls = if assembled_tool_calls.is_empty() {
762            None
763        } else {
764            Some(
765                assembled_tool_calls
766                    .clone()
767                    .into_iter()
768                    .map(ChatCompletionMessageToolCalls::Function)
769                    .collect(),
770            )
771        };
772
773        // Ensure we don't add an invalid assistant message to history
774        if content.is_some() || tool_calls.is_some() {
775            let assistant_msg = ChatCompletionRequestMessage::Assistant(
776                ChatCompletionRequestAssistantMessage {
777                    content,
778                    name: None,
779                    tool_calls,
780                    refusal: None,
781                    audio: None,
782                    #[allow(deprecated)]
783                    function_call: None,
784                }
785                .into(),
786            );
787
788            session.history.push(assistant_msg);
789        } else {
790            tracing::warn!("Not adding empty assistant message to history (no content and no tool calls)");
791        }
792
793        // If no tool calls, turn is complete
794        if assembled_tool_calls.is_empty() {
795            return Ok(0);
796        }
797
798        self.execute_tool_calls(session_id, &assembled_tool_calls).await
799    }
800
801    /// Execute the tool calls assembled from one LLM response and append the
802    /// results to the session history.
803    async fn execute_tool_calls(
804        &mut self,
805        session_id: &SessionId,
806        assembled_tool_calls: &[ChatCompletionMessageToolCall],
807    ) -> Result<usize> {
808        // First get the working_dir before any mutable borrow of sessions
809        let working_dir = {
810            let session = self
811                .sessions
812                .get(session_id)
813                .ok_or_else(|| AgentError::InternalError("Session not found".to_string()))?;
814            session.working_dir.clone()
815        };
816
817        // Images collected across the whole batch, injected as one user
818        // message after the loop.
819        let mut batch_images: Vec<ToolImage> = Vec::new();
820
821        // Execute each tool call
822        for (tc_idx, tc) in assembled_tool_calls.iter().enumerate() {
823            tracing::info!(
824                "Executing tool [{}/{}]: name='{}', id='{}', args={}",
825                tc_idx + 1,
826                assembled_tool_calls.len(),
827                tc.function.name,
828                tc.id,
829                truncate_for_log(&tc.function.arguments, 80)
830            );
831
832            let tc_info = ToolCallInfo {
833                id: tc.id.clone(),
834                name: tc.function.name.clone(),
835                arguments: tc.function.arguments.clone(),
836            };
837
838            // Notify frontend. Capture the result instead of `let _ =` so a
839            // failed delivery (closed/full channel, platform send error) is
840            // surfaced — a silent failure here is exactly the "no feedback"
841            // symptom we want to catch.
842            if let Err(e) = self
843                .frontend
844                .on_event(AgentEvent::ToolCallRequested {
845                    tool_call_id: tc_info.id.clone(),
846                    name: tc_info.name.clone(),
847                    arguments: tc_info.arguments.clone(),
848                })
849                .await
850            {
851                tracing::warn!(
852                    "[tool] ToolCallRequested delivery FAILED (user feedback may be lost): tool_call_id='{}', name='{}', error={}",
853                    tc_info.id,
854                    tc_info.name,
855                    e
856                );
857            }
858
859            // Check confirmation
860            let requires_confirm = self.tools.requires_confirmation(&tc.function.name);
861            let approved = if requires_confirm && !self.auto_approve {
862                tracing::trace!(
863                    "[tool] requesting user confirmation: tool_call_id='{}', name='{}'",
864                    tc_info.id,
865                    tc_info.name
866                );
867                match self.frontend.request_tool_confirmation(&tc_info).await {
868                    Ok(approved) => {
869                        tracing::trace!(
870                            "[tool] confirmation response: tool_call_id='{}', name='{}', approved={}",
871                            tc_info.id,
872                            tc_info.name,
873                            approved
874                        );
875                        approved
876                    }
877                    Err(e) => {
878                        tracing::warn!(
879                            "[tool] confirmation request failed: tool_call_id='{}', name='{}', error={}",
880                            tc_info.id,
881                            tc_info.name,
882                            e
883                        );
884                        return Err(e);
885                    }
886                }
887            } else {
888                tracing::trace!(
889                    "[tool] skipping confirmation (requires_confirm={}, auto_approve={})",
890                    requires_confirm,
891                    self.auto_approve
892                );
893                true
894            };
895
896            // Execute or reject
897            let result = if approved {
898                let args: serde_json::Value = serde_json::from_str(&tc.function.arguments)
899                    .unwrap_or(serde_json::Value::Null);
900
901                // Per-call cancellation token. Async tools pass a clone into
902                // `async_runner.submit`; if the tool goes async the Agent keeps
903                // this clone in `pending_tasks` so it can cancel the work later.
904                let cancel_token = CancellationToken::new();
905
906                let ctx = ToolContext {
907                    working_dir: working_dir.clone(),
908                    session_id: session_id.clone(),
909                    tool_call_id: tc.id.clone(),
910                    frontend: self.frontend.clone(),
911                    extensions: self.extensions.clone(),
912                    supports_images: self.llm_client.supports_images(),
913                    async_runner: self.async_runner.clone(),
914                    cancel_token: cancel_token.clone(),
915                    task_registry: self.task_registry.clone(),
916                };
917
918                let result = self.tools.execute(&tc.function.name, args, &ctx).await;
919                tracing::trace!(
920                    "[tool] execution returned: tool_call_id='{}', name='{}', is_pending={}, is_error={}, content_len={}",
921                    tc_info.id,
922                    tc_info.name,
923                    result.is_pending,
924                    result.is_error,
925                    result.content.len()
926                );
927
928                // If the tool went async, register the task so it can be
929                // tracked and cancelled. The placeholder content is still added
930                // to history below (as the tool message) so the LLM can keep
931                // working while the task runs.
932                if result.is_pending {
933                    if let Some(tid) = &result.pending_task_id {
934                        tracing::info!(
935                            "[async] task submitted: task_id={}, tool={}, tool_call_id={}",
936                            tid,
937                            tc.function.name,
938                            tc.id
939                        );
940                        self.pending_tasks.insert(
941                            tid.clone(),
942                            PendingTask {
943                                cancel: cancel_token,
944                                tool_name: tc.function.name.clone(),
945                            },
946                        );
947                        self.task_registry.register(AsyncTaskRecord {
948                            task_id: tid.clone(),
949                            tool_name: tc.function.name.clone(),
950                            tool_call_id: tc.id.clone(),
951                            session_id: session_id.clone(),
952                            status: AsyncTaskStatus::Pending,
953                            started_at: std::time::Instant::now(),
954                            result_summary: None,
955                        });
956                    } else {
957                        tracing::warn!(
958                            "[async] tool {} returned is_pending without pending_task_id",
959                            tc.function.name
960                        );
961                    }
962                }
963
964                result
965            } else {
966                tracing::trace!(
967                    "[tool] tool call rejected by user: tool_call_id='{}', name='{}'",
968                    tc_info.id,
969                    tc_info.name
970                );
971                ToolResult::error("User rejected this tool call")
972            };
973
974            // Truncate output
975            let raw_len = result.content.len();
976            let truncated_result = ToolResult {
977                content: self.context_manager.truncate_tool_output(&result.content),
978                is_error: result.is_error,
979                images: result.images.clone(),
980                is_pending: result.is_pending,
981                pending_task_id: result.pending_task_id.clone(),
982            };
983            if truncated_result.content.len() != raw_len {
984                tracing::trace!(
985                    "[tool] output truncated: tool_call_id='{}', name='{}', raw_len={}, truncated_len={}",
986                    tc_info.id,
987                    tc_info.name,
988                    raw_len,
989                    truncated_result.content.len()
990                );
991            }
992
993            // Notify frontend of result. Same rationale as above: capture
994            // delivery errors so a lost ToolCallResult is never silent.
995            if let Err(e) = self
996                .frontend
997                .on_event(AgentEvent::ToolCallResult {
998                    tool_call_id: tc.id.clone(),
999                    result: truncated_result.clone(),
1000                })
1001                .await
1002            {
1003                tracing::warn!(
1004                    "[tool] ToolCallResult delivery FAILED (user feedback may be lost): tool_call_id='{}', name='{}', error={}",
1005                    tc_info.id,
1006                    tc_info.name,
1007                    e
1008                );
1009            }
1010
1011            // Add tool result to history
1012            let tool_msg = ChatCompletionRequestMessage::Tool(
1013                ChatCompletionRequestToolMessage {
1014                    content: truncated_result.content.into(),
1015                    tool_call_id: tc.id.clone(),
1016                }
1017                .into(),
1018            );
1019
1020            let session = self
1021                .sessions
1022                .get_mut(session_id)
1023                .ok_or_else(|| AgentError::InternalError("Session not found".to_string()))?;
1024            session.history.push(tool_msg);
1025
1026            // Collect images from this tool result; they are injected after
1027            // ALL tool messages of the batch (below).
1028            batch_images.extend(truncated_result.images);
1029        }
1030
1031        // Inject the batch's collected images as a single multimodal user
1032        // message AFTER all tool messages. OpenAI protocol restricts tool
1033        // message content to text, so images travel in a separate user
1034        // message — but it must not interleave with the tool responses:
1035        // providers reject anything between an assistant `tool_calls`
1036        // message and its tool responses with a 400 error.
1037        if self.llm_client.supports_images() {
1038            if let Some(image_msg) = build_image_user_message(&batch_images) {
1039                let session = self
1040                    .sessions
1041                    .get_mut(session_id)
1042                    .ok_or_else(|| AgentError::InternalError("Session not found".to_string()))?;
1043                session.history.push(image_msg);
1044            }
1045        }
1046
1047        Ok(assembled_tool_calls.len())
1048    }
1049
1050    /// Clear the current session's history (keep system prompt).
1051    fn clear_session(&mut self) {
1052        if let Some(session) = self.sessions.get_mut(&self.default_session_id) {
1053            session.history.truncate(1);
1054        }
1055    }
1056
1057    /// Build a user message, potentially with images if model supports them.
1058    async fn build_user_message(
1059        &self,
1060        text: &str,
1061        attachments: &[MediaAttachment],
1062    ) -> ChatCompletionRequestMessage {
1063        // If model supports images and we have image attachments, build multimodal message
1064        if self.llm_client.supports_images()
1065            && !attachments.is_empty()
1066            && attachments.iter().any(|a| a.is_image())
1067        {
1068            self.build_multimodal_message(text, attachments)
1069                .await
1070        } else {
1071            // Fallback: add attachment descriptions to text
1072            let mut full_text = text.to_string();
1073            for attachment in attachments {
1074                full_text = format!("{}\n{}", full_text, attachment.describe());
1075            }
1076            ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage {
1077                content: full_text.into(),
1078                name: None,
1079            })
1080        }
1081    }
1082
1083    /// Build a multimodal message with text + images.
1084    async fn build_multimodal_message(
1085        &self,
1086        text: &str,
1087        attachments: &[MediaAttachment],
1088    ) -> ChatCompletionRequestMessage {
1089        let mut parts = vec![ChatCompletionRequestUserMessageContentPart::Text(
1090            ChatCompletionRequestMessageContentPartText {
1091                text: text.to_string(),
1092                prompt_cache_breakpoint: None,
1093            },
1094        )];
1095
1096        // Add images
1097        for attachment in attachments {
1098            if attachment.is_image() {
1099                // Download and encode as base64
1100                match media::download_and_encode_base64(
1101                    &attachment.url,
1102                    &attachment.content_type,
1103                )
1104                .await
1105                {
1106                    Ok(base64_url) => {
1107                        parts.push(ChatCompletionRequestUserMessageContentPart::ImageUrl(
1108                            ChatCompletionRequestMessageContentPartImage {
1109                                image_url: ImageUrl {
1110                                    url: base64_url,
1111                                    detail: None,
1112                                },
1113                                prompt_cache_breakpoint: None,
1114                            },
1115                        ));
1116                    }
1117                    Err(e) => {
1118                        tracing::warn!("Failed to encode image: {}", e);
1119                        // Fallback to description
1120                        let desc = attachment.describe();
1121                        let current_text = match &mut parts[0] {
1122                            ChatCompletionRequestUserMessageContentPart::Text(t) => &mut t.text,
1123                            _ => unreachable!(),
1124                        };
1125                        *current_text = format!("{}\n{}", current_text, desc);
1126                    }
1127                }
1128            } else {
1129                // Non-image: add description
1130                let desc = attachment.describe();
1131                let current_text = match &mut parts[0] {
1132                    ChatCompletionRequestUserMessageContentPart::Text(t) => &mut t.text,
1133                    _ => unreachable!(),
1134                };
1135                *current_text = format!("{}\n{}", current_text, desc);
1136            }
1137        }
1138
1139        ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage {
1140            content: ChatCompletionRequestUserMessageContent::Array(parts),
1141            name: None,
1142        })
1143    }
1144
1145    /// Execute a skill-triggered turn: inject skill content, then run the agent loop.
1146    ///
1147    /// The skill's full content is injected as a temporary system message and removed
1148    /// after the turn completes, so it doesn't occupy context in future turns.
1149    async fn run_skill_turn(&mut self, skill: &crate::skill::Skill, args: &str) {
1150        // Notify frontend
1151        let _ = self
1152            .frontend
1153            .on_event(AgentEvent::SkillTriggered {
1154                name: skill.frontmatter.name.clone(),
1155                description: skill.frontmatter.description.clone(),
1156            })
1157            .await;
1158
1159        let session_id = self.default_session_id.clone();
1160
1161        // Inject skill content as a system message
1162        let skill_message = format!(
1163            "## Skill: {}\n\n{}\n\n{}",
1164            skill.frontmatter.name,
1165            skill.frontmatter.description,
1166            skill.content
1167        );
1168
1169        let skill_msg = ChatCompletionRequestMessage::System(
1170            ChatCompletionRequestSystemMessage {
1171                content: skill_message.into(),
1172                name: Some(skill.frontmatter.name.clone()),
1173            }
1174            .into(),
1175        );
1176
1177        if let Some(session) = self.sessions.get_mut(&session_id) {
1178            session.history.push(skill_msg);
1179        }
1180
1181        // Add user message (args or default)
1182        let user_content = if args.is_empty() {
1183            "(User triggered skill, no additional arguments)".to_string()
1184        } else {
1185            args.to_string()
1186        };
1187
1188        if let Some(session) = self.sessions.get_mut(&session_id) {
1189            session.history.push(ChatCompletionRequestMessage::User(
1190                ChatCompletionRequestUserMessage {
1191                    content: user_content.into(),
1192                    name: None,
1193                }
1194                .into(),
1195            ));
1196        }
1197
1198        // Run the agentic loop
1199        let max_iterations = 20;
1200        let mut completed = false;
1201        for iteration in 0..max_iterations {
1202            match self.run_one_step(&session_id).await {
1203                Ok(tool_call_count) => {
1204                    if tool_call_count == 0 {
1205                        completed = true;
1206                        break;
1207                    }
1208                    tracing::debug!(
1209                        "Skill iteration {}: tool calls executed",
1210                        iteration
1211                    );
1212                }
1213                Err(e) => {
1214                    let _ = self.frontend.on_event(AgentEvent::Error(e)).await;
1215                    break;
1216                }
1217            }
1218        }
1219
1220        if !completed {
1221            let _ = self
1222                .frontend
1223                .on_event(AgentEvent::Error(AgentError::InternalError(
1224                    format!("Max iterations reached ({})", max_iterations),
1225                )))
1226                .await;
1227        }
1228
1229        let _ = self.frontend.on_event(AgentEvent::TurnComplete).await;
1230
1231        // Remove the injected skill system message to avoid polluting future turns
1232        if let Some(session) = self.sessions.get_mut(&session_id) {
1233            let skill_name = skill.frontmatter.name.clone();
1234            session.history.retain(|msg| {
1235                !matches!(
1236                    msg,
1237                    ChatCompletionRequestMessage::System(s)
1238                        if s.name.as_deref() == Some(&skill_name)
1239                )
1240            });
1241        }
1242    }
1243
1244    /// Handle a completed async background task: update tracking, notify the
1245    /// frontend, reinject the result into history, and wake the LLM.
1246    async fn handle_async_done(&mut self, done: AsyncTaskDone) {
1247        tracing::info!(
1248            "[async] task done: task_id={}, tool={}, session={}, cancelled={}, is_error={}",
1249            done.task_id,
1250            done.tool_name,
1251            done.session_id,
1252            done.cancelled,
1253            done.result.is_error
1254        );
1255
1256        // No longer in flight.
1257        self.pending_tasks.remove(&done.task_id);
1258
1259        // Update the registry with final status + a result summary.
1260        let status = if done.cancelled {
1261            AsyncTaskStatus::Cancelled
1262        } else if done.result.is_error {
1263            AsyncTaskStatus::Failed
1264        } else {
1265            AsyncTaskStatus::Completed
1266        };
1267        let summary = summarize_result(&done.result.content);
1268        self.task_registry
1269            .update(&done.task_id, status, Some(summary));
1270
1271        // Notify the frontend (TUI/GUI update task panels; chatbot usually
1272        // relies on the subsequent LLM reply delivered via TextDelta).
1273        let _ = self
1274            .frontend
1275            .on_event(AgentEvent::AsyncToolCompleted {
1276                task_id: done.task_id.clone(),
1277                tool_call_id: done.tool_call_id.clone(),
1278                result: done.result.clone(),
1279            })
1280            .await;
1281
1282        // Reinject into the owning session. If the session is gone (e.g. the
1283        // chatbot expired this Agent), drop the result - side effects like
1284        // saved files already happened.
1285        let session_id = done.session_id.clone();
1286        if !self.sessions.contains_key(&session_id) {
1287            tracing::error!(
1288                "[async] task {} (tool={}) finished but session {} not found; dropping result. \
1289                 This means the Agent exited or the session was cleaned up before the task completed. \
1290                 Result: {} chars, is_error={}, cancelled={}",
1291                done.task_id, done.tool_name, session_id,
1292                done.result.content.len(), done.result.is_error, done.cancelled
1293            );
1294            return;
1295        }
1296
1297        // Append a user-role notification. We do NOT mutate the original
1298        // placeholder tool message: the LLM may have already acted on it, and
1299        // rewriting history would break consistency.
1300        let notice = format!(
1301            "[后台任务完成通知] task_id={} (工具: {})\n{}",
1302            done.task_id, done.tool_name, done.result.content
1303        );
1304        if let Some(session) = self.sessions.get_mut(&session_id) {
1305            session.history.push(ChatCompletionRequestMessage::User(
1306                ChatCompletionRequestUserMessage {
1307                    content: notice.into(),
1308                    name: None,
1309                },
1310            ));
1311
1312            // Inject result images as a multimodal user message (same pattern
1313            // as sync tool results).
1314            if self.llm_client.supports_images() {
1315                if let Some(image_msg) = build_image_user_message(&done.result.images) {
1316                    session.history.push(image_msg);
1317                }
1318            }
1319        }
1320
1321        // Wake the LLM to process the notification.
1322        self.run_agent_loop(&session_id).await;
1323    }
1324
1325    /// Cancel a specific async task by id. The spawned task emits a cancelled
1326    /// `done` which flows through `handle_async_done` to update status.
1327    async fn handle_cancel_task(&mut self, task_id: &str) {
1328        match self.pending_tasks.remove(task_id) {
1329            Some(pending) => {
1330                tracing::info!(
1331                    "[async] cancelling task {} (tool={})",
1332                    task_id,
1333                    pending.tool_name
1334                );
1335                pending.cancel.cancel();
1336            }
1337            None => {
1338                tracing::warn!("[async] cancel request for unknown task {}", task_id);
1339            }
1340        }
1341    }
1342
1343    /// Cancel all in-flight async tasks for this Agent.
1344    async fn handle_cancel_all(&mut self) {
1345        let count = self.pending_tasks.len();
1346        if count == 0 {
1347            tracing::info!("[async] Cancel requested, no pending tasks");
1348            return;
1349        }
1350        tracing::info!("[async] cancelling all {} pending task(s)", count);
1351        for (_, pending) in self.pending_tasks.drain() {
1352            pending.cancel.cancel();
1353        }
1354    }
1355}
1356
1357impl Drop for Agent {
1358    fn drop(&mut self) {
1359        // Cancel any still-running async tasks so they don't outlive the Agent
1360        // (e.g. when a chatbot session expires and the Agent task is dropped).
1361        let count = self.pending_tasks.len();
1362        if count > 0 {
1363            tracing::info!(
1364                "[async] Agent dropped, cancelling {} pending task(s)",
1365                count
1366            );
1367            for (_, pending) in self.pending_tasks.drain() {
1368                pending.cancel.cancel();
1369            }
1370        }
1371    }
1372}
1373
1374// ============================================================================
1375// Summary generation (free function to avoid borrow conflicts)
1376// ============================================================================
1377
1378/// Generate a summary of removed conversation messages using the LLM.
1379/// Uses a non-streaming call to produce a 1-2 sentence summary.
1380/// Falls back to a static message on failure.
1381/// Apply a truncation result to the session history.
1382/// For NewSegment: generates a summary from removed messages and replaces the placeholder.
1383/// For MergeSegments: merges existing summary segments and replaces the placeholder.
1384/// For TruncateOnly: no-op.
1385async fn apply_compression_result(
1386    llm_client: &LlmClient,
1387    history: &mut [ChatCompletionRequestMessage],
1388    result: &TruncationResult,
1389) {
1390    if !result.needs_compression {
1391        return;
1392    }
1393
1394    let pos = result.insert_position;
1395    if pos >= history.len() {
1396        tracing::warn!("Insert position {} out of bounds (history len: {})", pos, history.len());
1397        return;
1398    }
1399
1400    let (content, name) = match &result.action {
1401        TruncationAction::NewSegment => {
1402            let summary = generate_summary(llm_client, &result.removed_messages).await;
1403            (
1404                format!("[Summary: {}]", summary),
1405                "summary_segment".to_string(),
1406            )
1407        }
1408        TruncationAction::MergeSegments { summaries, .. } => {
1409            let merged = merge_summaries(llm_client, summaries).await;
1410            // Determine new merge level from the placeholder's name
1411            let current_level = crate::context::get_merge_level(&history[pos]);
1412            let name = if current_level == 0 {
1413                "summary_segment".to_string()
1414            } else {
1415                format!("summary_segment_m{}", current_level)
1416            };
1417            (format!("[Summary: {}]", merged), name)
1418        }
1419        TruncationAction::TruncateOnly => return,
1420    };
1421
1422    tracing::info!("Compression applied at position {}: {}", pos, name);
1423
1424    history[pos] = ChatCompletionRequestMessage::User(
1425        ChatCompletionRequestUserMessage {
1426            content: content.into(),
1427            name: Some(name),
1428        }
1429    );
1430}
1431
1432/// Generate a short summary from removed full conversation rounds.
1433async fn generate_summary(
1434    llm_client: &LlmClient,
1435    removed_messages: &[ChatCompletionRequestMessage],
1436) -> String {
1437    tracing::debug!("Generating summary: removed_messages count = {}", removed_messages.len());
1438    let transcript = crate::context::format_removed_messages_as_transcript(removed_messages);
1439    tracing::debug!("Formatted transcript length: {} characters", transcript.len());
1440
1441    let system_prompt = "Summarize the following conversation transcript in 1-2 concise sentences. Focus on: what the user asked for, what actions were taken, and the outcomes. Be brief and factual.";
1442
1443    let messages = vec![
1444        ChatCompletionRequestMessage::System(
1445            ChatCompletionRequestSystemMessage {
1446                content: system_prompt.into(),
1447                name: None,
1448            }
1449        ),
1450        ChatCompletionRequestMessage::User(
1451            ChatCompletionRequestUserMessage {
1452                content: format!("Conversation transcript:\n\n{}", transcript).into(),
1453                name: None,
1454            }
1455        ),
1456    ];
1457
1458    tracing::info!("Calling LLM to generate summary...");
1459    match llm_client.chat(messages, None).await {
1460        Ok(response) => {
1461            tracing::info!("LLM responded successfully for summary generation");
1462            tracing::debug!("Number of choices in response: {}", response.choices.len());
1463            if let Some(choice) = response.choices.first() {
1464                tracing::debug!("Choice index: 0, has content: {}", choice.message.content.is_some());
1465                if let Some(content) = &choice.message.content {
1466                    let summary = content.trim().to_string();
1467                    if !summary.is_empty() {
1468                        tracing::info!("Successfully generated summary (length: {})", summary.len());
1469                        return summary;
1470                    }
1471                }
1472            }
1473            tracing::warn!("Summary generation returned empty response, using fallback");
1474            "Conversation history compressed.".to_string()
1475        }
1476        Err(e) => {
1477            tracing::error!("Summary generation failed with error: {}, using fallback", e);
1478            "Conversation history compressed.".to_string()
1479        }
1480    }
1481}
1482
1483/// Merge multiple existing summary segments into one coherent summary.
1484async fn merge_summaries(
1485    llm_client: &LlmClient,
1486    summaries: &[String],
1487) -> String {
1488    tracing::info!("Merging {} summary segments...", summaries.len());
1489
1490    let numbered: Vec<String> = summaries
1491        .iter()
1492        .enumerate()
1493        .map(|(i, s)| format!("[{}] {}", i + 1, s))
1494        .collect();
1495    let joined = numbered.join("\n\n");
1496
1497    let system_prompt = "You are given multiple conversation summaries from different time periods, ordered from oldest to newest. Merge them into a single concise summary (2-3 sentences) that preserves all key information.
1498
1499Key points to preserve:
1500- User goals and requests
1501- Important decisions made
1502- Technical context (file paths, APIs, architectures)
1503- Major outcomes and conclusions
1504
1505Do not simply concatenate — synthesize into a coherent narrative.";
1506
1507    let messages = vec![
1508        ChatCompletionRequestMessage::System(
1509            ChatCompletionRequestSystemMessage {
1510                content: system_prompt.into(),
1511                name: None,
1512            }
1513        ),
1514        ChatCompletionRequestMessage::User(
1515            ChatCompletionRequestUserMessage {
1516                content: format!("Summaries to merge:\n\n{}", joined).into(),
1517                name: None,
1518            }
1519        ),
1520    ];
1521
1522    match llm_client.chat(messages, None).await {
1523        Ok(response) => {
1524            if let Some(choice) = response.choices.first() {
1525                if let Some(content) = &choice.message.content {
1526                    let summary = content.trim().to_string();
1527                    if !summary.is_empty() {
1528                        tracing::info!("Successfully merged {} summaries (length: {})", summaries.len(), summary.len());
1529                        return summary;
1530                    }
1531                }
1532            }
1533            tracing::warn!("Summary merge returned empty response, using fallback");
1534            "Multiple earlier conversation segments merged.".to_string()
1535        }
1536        Err(e) => {
1537            tracing::error!("Summary merge failed with error: {}, using fallback", e);
1538            "Multiple earlier conversation segments merged.".to_string()
1539        }
1540    }
1541}
1542
1543// ============================================================================
1544// Helper types
1545// ============================================================================
1546
1547/// Accumulates streaming tool call chunks.
1548#[derive(Debug)]
1549struct ToolCallAccumulator {
1550    id: Option<String>,
1551    name: Option<String>,
1552    arguments: String,
1553}
1554
1555impl ToolCallAccumulator {
1556    fn new() -> Self {
1557        Self {
1558            id: None,
1559            name: None,
1560            arguments: String::new(),
1561        }
1562    }
1563
1564    /// Convert accumulated chunks into a complete tool call.
1565    fn into_tool_call(self) -> Option<ChatCompletionMessageToolCall> {
1566        let id = self.id?;
1567        let name = self.name?;
1568
1569        tracing::trace!(
1570            "Tool call assembled: id='{}', name='{}', args={}",
1571            id,
1572            name,
1573            truncate_for_log(&self.arguments, 80)
1574        );
1575
1576        Some(ChatCompletionMessageToolCall {
1577            id,
1578            function: FunctionCall {
1579                name,
1580                arguments: self.arguments,
1581            },
1582        })
1583    }
1584}
1585
1586/// Truncate a string to at most `max_chars` characters for log output,
1587/// appending a length note when truncated. Counts by `char` to avoid
1588/// splitting multi-byte UTF-8 sequences (safe for CJK text).
1589fn truncate_for_log(s: &str, max_chars: usize) -> String {
1590    let char_count = s.chars().count();
1591    if char_count <= max_chars {
1592        s.to_string()
1593    } else {
1594        let preview: String = s.chars().take(max_chars).collect();
1595        format!("{}... ({} chars total)", preview, char_count)
1596    }
1597}
1598
1599/// Build a multimodal user message carrying tool-result images, or `None` if
1600/// there are no images. Shared by sync tool results (`run_one_step`) and async
1601/// task-completion reinjection (`handle_async_done`).
1602fn build_image_user_message(images: &[ToolImage]) -> Option<ChatCompletionRequestMessage> {
1603    if images.is_empty() {
1604        return None;
1605    }
1606    let mut parts = vec![ChatCompletionRequestUserMessageContentPart::Text(
1607        ChatCompletionRequestMessageContentPartText {
1608            text: format!(
1609                "[工具返回的图片] {}",
1610                images
1611                    .iter()
1612                    .map(|i| i.label.as_str())
1613                    .collect::<Vec<_>>()
1614                    .join(", ")
1615            ),
1616            prompt_cache_breakpoint: None,
1617        },
1618    )];
1619    for img in images {
1620        parts.push(ChatCompletionRequestUserMessageContentPart::ImageUrl(
1621            ChatCompletionRequestMessageContentPartImage {
1622                image_url: ImageUrl {
1623                    url: img.data_url.clone(),
1624                    detail: None,
1625                },
1626                prompt_cache_breakpoint: None,
1627            },
1628        ));
1629    }
1630    Some(ChatCompletionRequestMessage::User(ChatCompletionRequestUserMessage {
1631        content: ChatCompletionRequestUserMessageContent::Array(parts),
1632        name: None,
1633    }))
1634}
1635
1636/// Sanitize session history for models that don't support image inputs.
1637///
1638/// When the current model has `supports_images = false`, any `image_url`
1639/// content parts in user messages are incompatible with the API and will
1640/// cause a 400 error. This function downgrades multimodal `Array` content
1641/// to plain `Text` by concatenating the text parts and discarding image
1642/// parts. Non-array (plain text) messages are left unchanged.
1643fn sanitize_history_for_model(
1644    history: &mut Vec<ChatCompletionRequestMessage>,
1645    supports_images: bool,
1646) {
1647    if supports_images {
1648        return;
1649    }
1650
1651    let mut sanitized_count = 0usize;
1652    for msg in history.iter_mut() {
1653        if let ChatCompletionRequestMessage::User(user_msg) = msg {
1654            if let ChatCompletionRequestUserMessageContent::Array(parts) = &user_msg.content {
1655                // Check if this message actually contains image parts
1656                let has_image = parts
1657                    .iter()
1658                    .any(|p| matches!(p, ChatCompletionRequestUserMessageContentPart::ImageUrl(_)));
1659                if has_image {
1660                    // Concatenate all text parts, skip image parts
1661                    let text: String = parts
1662                        .iter()
1663                        .filter_map(|p| {
1664                            if let ChatCompletionRequestUserMessageContentPart::Text(t) = p {
1665                                Some(t.text.as_str())
1666                            } else {
1667                                None
1668                            }
1669                        })
1670                        .collect::<Vec<_>>()
1671                        .join("\n");
1672
1673                    user_msg.content = ChatCompletionRequestUserMessageContent::Text(text);
1674                    sanitized_count += 1;
1675                }
1676            }
1677        }
1678    }
1679
1680    if sanitized_count > 0 {
1681        tracing::info!(
1682            "sanitize_history_for_model: downgraded {} message(s) with image_url to text \
1683             (model does not support images)",
1684            sanitized_count
1685        );
1686    }
1687}
1688
1689/// Truncate a task result to a bounded summary for the task registry.
1690fn summarize_result(content: &str) -> String {
1691    const MAX: usize = 500;
1692    let char_count = content.chars().count();
1693    if char_count <= MAX {
1694        content.to_string()
1695    } else {
1696        let truncated: String = content.chars().take(MAX).collect();
1697        format!("{}... (truncated, {} chars total)", truncated, char_count)
1698    }
1699}
1700
1701#[cfg(test)]
1702mod tests {
1703    use super::*;
1704    use crate::event::AgentEvent;
1705    use crate::frontend::Frontend;
1706    use crate::skill::SkillRegistry;
1707    use crate::tool::{Tool, ToolContext};
1708    use async_trait::async_trait;
1709    use robit_ai::config::{ModelConfig, ProviderConfig, RobitConfig};
1710    use serde_json::Value;
1711
1712    /// Frontend that swallows all events (no UI in tests).
1713    struct NoopFrontend;
1714
1715    #[async_trait]
1716    impl Frontend for NoopFrontend {
1717        async fn on_event(&self, _event: AgentEvent) -> Result<()> {
1718            Ok(())
1719        }
1720
1721        async fn request_tool_confirmation(&self, _info: &ToolCallInfo) -> Result<bool> {
1722            Ok(true)
1723        }
1724    }
1725
1726    /// A tool whose result always carries one image, mimicking `read` on an
1727    /// image file with a vision-capable model.
1728    struct ImageTool;
1729
1730    #[async_trait]
1731    impl Tool for ImageTool {
1732        fn name(&self) -> &str {
1733            "fake_image_tool"
1734        }
1735
1736        fn description(&self) -> &str {
1737            "Returns an image"
1738        }
1739
1740        fn parameters_schema(&self) -> Value {
1741            serde_json::json!({"type": "object", "properties": {}})
1742        }
1743
1744        fn requires_confirmation(&self) -> bool {
1745            false
1746        }
1747
1748        async fn execute(&self, _args: Value, _ctx: &ToolContext) -> Result<ToolResult> {
1749            Ok(ToolResult {
1750                content: "Image file: x.png".to_string(),
1751                is_error: false,
1752                images: vec![ToolImage {
1753                    data_url: "data:image/png;base64,Zm9v".to_string(),
1754                    label: "x.png".to_string(),
1755                }],
1756                is_pending: false,
1757                pending_task_id: None,
1758            })
1759        }
1760    }
1761
1762    /// Build an `LlmClient` for a vision model without contacting it (the
1763    /// base URL points at a closed port; only `supports_images()` is used).
1764    fn vision_llm_client() -> Arc<LlmClient> {
1765        let config = RobitConfig {
1766            default_model: Some("test/vision".to_string()),
1767            providers: HashMap::from([(
1768                "test".to_string(),
1769                ProviderConfig {
1770                    name: Some("Test".to_string()),
1771                    base_url: "http://127.0.0.1:1".to_string(),
1772                    api_key: "sk-test".to_string(),
1773                    models: vec![ModelConfig {
1774                        id: "vision".to_string(),
1775                        name: Some("Vision".to_string()),
1776                        context_window: None,
1777                        max_output_tokens: None,
1778                        temperature: None,
1779                        max_tokens: None,
1780                        supports_images: Some(true),
1781                        supports_tools: Some(true),
1782                    }],
1783                },
1784            )]),
1785            app: None,
1786            channels: None,
1787            default_image_model: None,
1788            image_providers: HashMap::new(),
1789        };
1790        Arc::new(LlmClient::from_config(&config, None).unwrap())
1791    }
1792
1793    fn tool_call(id: &str) -> ChatCompletionMessageToolCall {
1794        ChatCompletionMessageToolCall {
1795            id: id.to_string(),
1796            function: FunctionCall {
1797                name: "fake_image_tool".to_string(),
1798                arguments: "{}".to_string(),
1799            },
1800        }
1801    }
1802
1803    /// Regression test for the DeepSeek 400 "insufficient tool messages
1804    /// following tool_calls message": when a parallel tool-call batch returns
1805    /// images, the injected multimodal user message(s) must come AFTER all
1806    /// tool messages of the batch, never between them.
1807    #[tokio::test]
1808    async fn parallel_image_tool_results_keep_tool_messages_contiguous() {
1809        let mut tools = ToolRegistry::new();
1810        tools.register(ImageTool);
1811        let mut agent = Agent::new(
1812            vision_llm_client(),
1813            Arc::new(tools),
1814            Arc::new(SkillRegistry::new(vec![], &[])),
1815            Arc::new(NoopFrontend),
1816            None,
1817            None,
1818            PathBuf::from("."),
1819            true,
1820            HashMap::new(),
1821        );
1822
1823        let session_id = agent.default_session_id.clone();
1824        let calls = vec![tool_call("call_0"), tool_call("call_1"), tool_call("call_2")];
1825        let executed = agent.execute_tool_calls(&session_id, &calls).await.unwrap();
1826        assert_eq!(executed, 3);
1827
1828        let session = agent.sessions.get(&session_id).unwrap();
1829        // History layout: [system, tool, tool, tool, user(images)]. The
1830        // system prompt is message 0; skip it.
1831        assert_eq!(session.history.len(), 5, "3 tool messages + 1 image user message");
1832        let kinds: Vec<&str> = session
1833            .history
1834            .iter()
1835            .skip(1)
1836            .map(|m| match m {
1837                ChatCompletionRequestMessage::Tool(_) => "tool",
1838                ChatCompletionRequestMessage::User(_) => "user",
1839                ChatCompletionRequestMessage::Assistant(_) => "assistant",
1840                _ => "other",
1841            })
1842            .collect();
1843        assert_eq!(
1844            kinds,
1845            vec!["tool", "tool", "tool", "user"],
1846            "tool responses must be contiguous after the assistant tool_calls \
1847             message; image user message(s) go after the batch"
1848        );
1849    }
1850}