Skip to main content

zeph_core/agent/
utils.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use zeph_llm::provider::{LlmProvider, Message, MessagePart, Role};
5
6use super::{Agent, CODE_CONTEXT_PREFIX};
7use crate::channel::Channel;
8use crate::metrics::{MetricsSnapshot, SECURITY_EVENT_CAP, SecurityEvent};
9use zeph_common::SecurityEventCategory;
10use zeph_tools::FilterStats;
11
12impl<C: Channel> Agent<C> {
13    /// Read the community-detection failure counter from `SemanticMemory` and update metrics.
14    pub fn sync_community_detection_failures(&self) {
15        if let Some(memory) = self.services.memory.persistence.memory.as_ref() {
16            let failures = memory.community_detection_failures();
17            self.update_metrics(|m| {
18                m.graph_community_detection_failures = failures;
19            });
20        }
21    }
22
23    /// Sync all graph counters (extraction count/failures) from `SemanticMemory` to metrics.
24    pub fn sync_graph_extraction_metrics(&self) {
25        if let Some(memory) = self.services.memory.persistence.memory.as_ref() {
26            let count = memory.graph_extraction_count();
27            let failures = memory.graph_extraction_failures();
28            self.update_metrics(|m| {
29                m.graph_extraction_count = count;
30                m.graph_extraction_failures = failures;
31            });
32        }
33    }
34
35    /// Fetch entity/edge/community counts from the graph store and write to metrics.
36    pub async fn sync_graph_counts(&self) {
37        let Some(memory) = self.services.memory.persistence.memory.as_ref() else {
38            return;
39        };
40        let Some(store) = memory.graph_store.as_ref() else {
41            return;
42        };
43        let (entities, edges, communities) = tokio::join!(
44            store.entity_count(),
45            store.active_edge_count(),
46            store.community_count()
47        );
48        self.update_metrics(|m| {
49            m.graph_entities_total = entities.unwrap_or(0).cast_unsigned();
50            m.graph_edges_total = edges.unwrap_or(0).cast_unsigned();
51            m.graph_communities_total = communities.unwrap_or(0).cast_unsigned();
52        });
53    }
54
55    /// Perform a real health check on the vector store and update metrics.
56    pub async fn check_vector_store_health(&self, backend_name: &str) {
57        let connected = match self.services.memory.persistence.memory.as_ref() {
58            Some(m) => m.is_vector_store_connected().await,
59            None => false,
60        };
61        let name = backend_name.to_owned();
62        self.update_metrics(|m| {
63            m.qdrant_available = connected;
64            m.vector_backend = name;
65        });
66    }
67
68    /// Fetch compression-guidelines metadata from `SQLite` and write to metrics.
69    ///
70    /// Only fetches version and `created_at`; does not load the full guidelines text.
71    /// Feature-gated: compiled only when `compression-guidelines` is enabled.
72    pub async fn sync_guidelines_status(&self) {
73        let Some(memory) = self.services.memory.persistence.memory.as_ref() else {
74            return;
75        };
76        let cid = self.services.memory.persistence.conversation_id;
77        match memory.sqlite().load_compression_guidelines_meta(cid).await {
78            Ok((version, created_at)) => {
79                #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
80                let version_u32 = u32::try_from(version).unwrap_or(0);
81                self.update_metrics(|m| {
82                    m.guidelines_version = version_u32;
83                    m.guidelines_updated_at = created_at;
84                });
85            }
86            Err(e) => {
87                tracing::warn!("failed to sync guidelines status: {e:#}");
88            }
89        }
90    }
91
92    pub(super) fn record_filter_metrics(&mut self, fs: &FilterStats) {
93        let saved = fs.estimated_tokens_saved() as u64;
94        let raw = (fs.raw_chars / 4) as u64;
95        let confidence = fs.confidence;
96        let was_filtered = fs.filtered_chars < fs.raw_chars;
97        self.update_metrics(|m| {
98            m.filter_raw_tokens += raw;
99            m.filter_saved_tokens += saved;
100            m.filter_applications += 1;
101            m.filter_total_commands += 1;
102            if was_filtered {
103                m.filter_filtered_commands += 1;
104            }
105            if let Some(c) = confidence {
106                match c {
107                    zeph_tools::FilterConfidence::Full => {
108                        m.filter_confidence_full += 1;
109                    }
110                    zeph_tools::FilterConfidence::Partial => {
111                        m.filter_confidence_partial += 1;
112                    }
113                    zeph_tools::FilterConfidence::Fallback => {
114                        m.filter_confidence_fallback += 1;
115                    }
116                    _ => {}
117                }
118            }
119        });
120    }
121
122    pub(super) fn update_metrics(&self, f: impl FnOnce(&mut MetricsSnapshot)) {
123        if let Some(ref tx) = self.runtime.metrics.metrics_tx {
124            let elapsed = self.runtime.lifecycle.start_time.elapsed().as_secs();
125            tx.send_modify(|m| {
126                m.uptime_seconds = elapsed;
127                f(m);
128            });
129        }
130    }
131
132    /// Publish the effective context window limit from the active provider's budget into
133    /// [`MetricsSnapshot::context_max_tokens`].
134    ///
135    /// Call after the provider pool is constructed (builder) and on every successful `/provider`
136    /// switch so the TUI context gauge always reflects the active provider's window.
137    /// When no budget is configured the field is set to `0`, which the gauge renders as `"—"`.
138    pub(crate) fn publish_context_budget(&self) {
139        let max_tokens = self
140            .context_manager
141            .budget
142            .as_ref()
143            .map_or(0, |b| b.max_tokens() as u64);
144        self.update_metrics(|m| m.context_max_tokens = max_tokens);
145    }
146
147    /// Flush `metrics.pending_timings` into the rolling window and publish to the metrics snapshot.
148    ///
149    /// Call once per turn after all four phases have written to `pending_timings`.
150    /// Resets `pending_timings` to default after flushing.
151    pub(super) fn flush_turn_timings(&mut self) {
152        let timings = std::mem::take(&mut self.runtime.metrics.pending_timings);
153        tracing::debug!(
154            prepare_context_ms = timings.prepare_context_ms,
155            llm_chat_ms = timings.llm_chat_ms,
156            tool_exec_ms = timings.tool_exec_ms,
157            persist_message_ms = timings.persist_message_ms,
158            "turn timings"
159        );
160
161        if self.runtime.metrics.timing_window.len() >= 10 {
162            self.runtime.metrics.timing_window.pop_front();
163        }
164        self.runtime
165            .metrics
166            .timing_window
167            .push_back(timings.clone());
168
169        let count = self.runtime.metrics.timing_window.len();
170        let mut avg = crate::metrics::TurnTimings::default();
171        let mut max = crate::metrics::TurnTimings::default();
172        for t in &self.runtime.metrics.timing_window {
173            avg.prepare_context_ms = avg.prepare_context_ms.saturating_add(t.prepare_context_ms);
174            avg.llm_chat_ms = avg.llm_chat_ms.saturating_add(t.llm_chat_ms);
175            avg.tool_exec_ms = avg.tool_exec_ms.saturating_add(t.tool_exec_ms);
176            avg.persist_message_ms = avg.persist_message_ms.saturating_add(t.persist_message_ms);
177
178            max.prepare_context_ms = max.prepare_context_ms.max(t.prepare_context_ms);
179            max.llm_chat_ms = max.llm_chat_ms.max(t.llm_chat_ms);
180            max.tool_exec_ms = max.tool_exec_ms.max(t.tool_exec_ms);
181            max.persist_message_ms = max.persist_message_ms.max(t.persist_message_ms);
182        }
183        let n = count as u64;
184        avg.prepare_context_ms /= n;
185        avg.llm_chat_ms /= n;
186        avg.tool_exec_ms /= n;
187        avg.persist_message_ms /= n;
188
189        let total_ms = timings
190            .prepare_context_ms
191            .saturating_add(timings.llm_chat_ms)
192            .saturating_add(timings.tool_exec_ms)
193            .saturating_add(timings.persist_message_ms);
194
195        self.update_metrics(|m| {
196            m.last_turn_timings = timings;
197            m.avg_turn_timings = avg;
198            m.max_turn_timings = max;
199            m.timing_sample_count = n;
200        });
201
202        if let Some(ref recorder) = self.runtime.metrics.histogram_recorder {
203            recorder.observe_turn_duration(std::time::Duration::from_millis(total_ms));
204        }
205    }
206
207    /// Push the current classifier metrics snapshot into `MetricsSnapshot`.
208    ///
209    /// Call this after any classifier invocation (injection, PII, feedback) so the TUI panel
210    /// reflects the latest p50/p95 values. No-op when classifier metrics are not configured.
211    pub(super) fn push_classifier_metrics(&self) {
212        if let Some(ref m) = self.runtime.metrics.classifier_metrics {
213            let snapshot = m.snapshot();
214            self.update_metrics(|ms| ms.classifier = snapshot);
215        }
216    }
217
218    pub(super) fn push_security_event(
219        &self,
220        category: SecurityEventCategory,
221        source: &str,
222        detail: impl Into<String>,
223    ) {
224        if let Some(ref tx) = self.runtime.metrics.metrics_tx {
225            let event = SecurityEvent::new(category, source, detail);
226            let elapsed = self.runtime.lifecycle.start_time.elapsed().as_secs();
227            tx.send_modify(|m| {
228                m.uptime_seconds = elapsed;
229                if m.security_events.len() >= SECURITY_EVENT_CAP {
230                    m.security_events.pop_front();
231                }
232                m.security_events.push_back(event);
233            });
234        }
235    }
236
237    pub(super) fn recompute_prompt_tokens(&mut self) {
238        self.runtime.providers.cached_prompt_tokens = self
239            .msg
240            .messages
241            .iter()
242            .map(|m| self.runtime.metrics.token_counter.count_message_tokens(m) as u64)
243            .sum();
244    }
245
246    pub(super) fn push_message(&mut self, msg: Message) {
247        self.runtime.providers.cached_prompt_tokens +=
248            self.runtime
249                .metrics
250                .token_counter
251                .count_message_tokens(&msg) as u64;
252        if msg.role == zeph_llm::provider::Role::Assistant {
253            self.services.session.last_assistant_at = Some(std::time::Instant::now());
254        }
255        self.msg.messages.push(msg);
256        // Detect MagicDoc headers in tool output after pushing the message.
257        self.detect_magic_docs_in_messages();
258    }
259
260    pub(crate) fn record_cost_and_cache(&self, input_tokens: u64, output_tokens: u64) {
261        let (cache_write, cache_read) = self.provider.last_cache_usage().unwrap_or((0, 0));
262
263        if let Some(ref tracker) = self.runtime.metrics.cost_tracker {
264            let provider_name = if self.runtime.config.active_provider_name.is_empty() {
265                self.provider.name()
266            } else {
267                self.runtime.config.active_provider_name.as_str()
268            };
269            tracker.record_usage(
270                provider_name,
271                self.provider.provider_kind_str(),
272                &self.runtime.config.model_name,
273                input_tokens,
274                cache_read,
275                cache_write,
276                output_tokens,
277            );
278            let breakdown = tracker.provider_breakdown();
279            self.update_metrics(|m| {
280                m.cost_spent_cents = tracker.current_spend();
281                m.cache_creation_tokens += cache_write;
282                m.cache_read_tokens += cache_read;
283                m.provider_cost_breakdown = breakdown;
284            });
285        } else if cache_write > 0 || cache_read > 0 {
286            self.update_metrics(|m| {
287                m.cache_creation_tokens += cache_write;
288                m.cache_read_tokens += cache_read;
289            });
290        }
291    }
292
293    pub(crate) fn record_successful_task(&self) {
294        if let Some(ref tracker) = self.runtime.metrics.cost_tracker {
295            tracker.record_successful_task();
296            self.update_metrics(|m| {
297                m.cost_cps_cents = tracker.cps();
298                m.cost_successful_tasks = tracker.successful_tasks();
299            });
300        }
301    }
302
303    /// Extract a redacted preview of the last assistant message.
304    ///
305    /// Walks `self.msg.messages` in reverse to find the most recent `Role::Assistant`
306    /// message, takes up to `max_chars` Unicode scalar values from `message.content`,
307    /// and applies [`crate::redact::scrub_content`] to redact any secrets.
308    ///
309    /// Returns an empty string when no assistant message exists in the current turn.
310    pub(super) fn last_assistant_preview(&self, max_chars: usize) -> String {
311        let raw = self
312            .msg
313            .messages
314            .iter()
315            .rev()
316            .find(|m| m.role == Role::Assistant)
317            .map_or("", |m| m.content.as_str());
318
319        if raw.is_empty() {
320            return String::new();
321        }
322
323        // Truncate to max_chars before redaction to bound redaction work.
324        let truncated: &str = if raw.chars().count() > max_chars {
325            let end = raw
326                .char_indices()
327                .nth(max_chars)
328                .map_or(raw.len(), |(i, _)| i);
329            &raw[..end]
330        } else {
331            raw
332        };
333
334        crate::redact::scrub_content(truncated).into_owned()
335    }
336
337    /// Inject pre-formatted code context into the message list.
338    /// The caller is responsible for retrieving and formatting the text.
339    pub fn inject_code_context(&mut self, text: &str) {
340        self.remove_code_context_messages();
341        if text.is_empty() || self.msg.messages.len() <= 1 {
342            return;
343        }
344        let content = format!("{CODE_CONTEXT_PREFIX}{text}");
345        self.msg.messages.insert(
346            1,
347            Message::from_parts(
348                Role::System,
349                vec![MessagePart::CodeContext { text: content }],
350            ),
351        );
352    }
353
354    #[must_use]
355    pub fn context_messages(&self) -> &[Message] {
356        &self.msg.messages
357    }
358
359    /// Truncate stale tool result content in old messages to bound in-memory growth.
360    ///
361    /// After the LLM has seen and responded to tool output, the full content is no longer
362    /// needed in the hot message list (it is already persisted to `SQLite`). Truncating keeps
363    /// the in-process message vec small across long sessions.
364    ///
365    /// Skips the last 2 messages so the LLM retains full context for the next turn.
366    ///
367    /// Truncated variants: `MessagePart::ToolResult` (content) and `MessagePart::ToolOutput` (body).
368    pub(super) fn truncate_old_tool_results(&mut self) {
369        const LIMIT: usize = 2048;
370        const SUFFIX: &str = "…[truncated]";
371
372        let len = self.msg.messages.len();
373        if len <= 2 {
374            return;
375        }
376        for msg in &mut self.msg.messages[..len - 2] {
377            for part in &mut msg.parts {
378                match part {
379                    MessagePart::ToolResult { content, .. } if content.len() > LIMIT => {
380                        content.truncate(content.floor_char_boundary(LIMIT));
381                        content.push_str(SUFFIX);
382                    }
383                    MessagePart::ToolOutput { body, .. } if body.len() > LIMIT => {
384                        body.truncate(body.floor_char_boundary(LIMIT));
385                        body.push_str(SUFFIX);
386                    }
387                    _ => {}
388                }
389            }
390        }
391    }
392}
393
394#[cfg(test)]
395mod tests {
396    use super::super::agent_tests::{
397        MockChannel, MockToolExecutor, create_test_registry, mock_provider,
398    };
399    use super::*;
400    use zeph_llm::provider::{MessageMetadata, MessagePart};
401
402    #[test]
403    fn push_message_increments_cached_tokens() {
404        let provider = mock_provider(vec![]);
405        let channel = MockChannel::new(vec![]);
406        let registry = create_test_registry();
407        let executor = MockToolExecutor::no_tools();
408        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
409
410        let before = agent.runtime.providers.cached_prompt_tokens;
411        let msg = Message {
412            role: Role::User,
413            content: "hello world!!".to_string(),
414            parts: vec![],
415            metadata: MessageMetadata::default(),
416        };
417        let expected_delta = agent
418            .runtime
419            .metrics
420            .token_counter
421            .count_message_tokens(&msg) as u64;
422        agent.push_message(msg);
423        assert_eq!(
424            agent.runtime.providers.cached_prompt_tokens,
425            before + expected_delta
426        );
427    }
428
429    #[test]
430    fn recompute_prompt_tokens_matches_sum() {
431        let provider = mock_provider(vec![]);
432        let channel = MockChannel::new(vec![]);
433        let registry = create_test_registry();
434        let executor = MockToolExecutor::no_tools();
435        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
436
437        agent.msg.messages.push(Message {
438            role: Role::User,
439            content: "1234".to_string(),
440            parts: vec![],
441            metadata: MessageMetadata::default(),
442        });
443        agent.msg.messages.push(Message {
444            role: Role::Assistant,
445            content: "5678".to_string(),
446            parts: vec![],
447            metadata: MessageMetadata::default(),
448        });
449
450        agent.recompute_prompt_tokens();
451
452        let expected: u64 = agent
453            .msg
454            .messages
455            .iter()
456            .map(|m| agent.runtime.metrics.token_counter.count_message_tokens(m) as u64)
457            .sum();
458        assert_eq!(agent.runtime.providers.cached_prompt_tokens, expected);
459    }
460
461    #[test]
462    fn inject_code_context_into_messages_with_existing_content() {
463        let provider = mock_provider(vec![]);
464        let channel = MockChannel::new(vec![]);
465        let registry = create_test_registry();
466        let executor = MockToolExecutor::no_tools();
467        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
468
469        // Add a user message so we have more than 1 message
470        agent.push_message(Message {
471            role: Role::User,
472            content: "question".to_string(),
473            parts: vec![],
474            metadata: MessageMetadata::default(),
475        });
476
477        agent.inject_code_context("some code here");
478
479        let found = agent.msg.messages.iter().any(|m| {
480            m.parts.iter().any(|p| {
481                matches!(p, MessagePart::CodeContext { text } if text.contains("some code here"))
482            })
483        });
484        assert!(found, "code context should be injected into messages");
485    }
486
487    #[test]
488    fn inject_code_context_empty_text_is_noop() {
489        let provider = mock_provider(vec![]);
490        let channel = MockChannel::new(vec![]);
491        let registry = create_test_registry();
492        let executor = MockToolExecutor::no_tools();
493        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
494
495        agent.push_message(Message {
496            role: Role::User,
497            content: "question".to_string(),
498            parts: vec![],
499            metadata: MessageMetadata::default(),
500        });
501        let count_before = agent.msg.messages.len();
502
503        agent.inject_code_context("");
504
505        // No code context message inserted for empty text
506        assert_eq!(agent.msg.messages.len(), count_before);
507    }
508
509    #[test]
510    fn inject_code_context_with_single_message_is_noop() {
511        let provider = mock_provider(vec![]);
512        let channel = MockChannel::new(vec![]);
513        let registry = create_test_registry();
514        let executor = MockToolExecutor::no_tools();
515        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
516        // Only system prompt → len == 1 → inject should be noop
517        let count_before = agent.msg.messages.len();
518
519        agent.inject_code_context("some code");
520
521        assert_eq!(agent.msg.messages.len(), count_before);
522    }
523
524    #[test]
525    fn context_messages_returns_all_messages() {
526        let provider = mock_provider(vec![]);
527        let channel = MockChannel::new(vec![]);
528        let registry = create_test_registry();
529        let executor = MockToolExecutor::no_tools();
530        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
531
532        agent.push_message(Message {
533            role: Role::User,
534            content: "test".to_string(),
535            parts: vec![],
536            metadata: MessageMetadata::default(),
537        });
538
539        assert_eq!(agent.context_messages().len(), agent.msg.messages.len());
540    }
541
542    #[test]
543    fn truncate_old_tool_results_truncates_stale_content() {
544        let provider = mock_provider(vec![]);
545        let channel = MockChannel::new(vec![]);
546        let registry = create_test_registry();
547        let executor = MockToolExecutor::no_tools();
548        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
549
550        let big_content = "x".repeat(4096);
551
552        // Message 0 (old) — should be truncated.
553        agent.msg.messages.push(Message {
554            role: Role::User,
555            content: String::new(),
556            parts: vec![MessagePart::ToolResult {
557                tool_use_id: "id1".to_string(),
558                content: big_content.clone(),
559                is_error: false,
560            }],
561            metadata: MessageMetadata::default(),
562        });
563        // Message 1 (old) — ToolOutput should also be truncated.
564        agent.msg.messages.push(Message {
565            role: Role::User,
566            content: String::new(),
567            parts: vec![MessagePart::ToolOutput {
568                tool_name: "shell".into(),
569                body: big_content.clone(),
570                compacted_at: None,
571            }],
572            metadata: MessageMetadata::default(),
573        });
574        // Message 2 (recent) — must NOT be truncated.
575        agent.msg.messages.push(Message {
576            role: Role::Assistant,
577            content: "reply".to_string(),
578            parts: vec![MessagePart::ToolResult {
579                tool_use_id: "id3".to_string(),
580                content: big_content.clone(),
581                is_error: false,
582            }],
583            metadata: MessageMetadata::default(),
584        });
585        // Message 3 (most recent) — must NOT be truncated.
586        agent.msg.messages.push(Message {
587            role: Role::User,
588            content: "last".to_string(),
589            parts: vec![MessagePart::ToolResult {
590                tool_use_id: "id4".to_string(),
591                content: big_content.clone(),
592                is_error: false,
593            }],
594            metadata: MessageMetadata::default(),
595        });
596
597        // Agent::new inserts a system prompt at index 0, so our messages are at 1..=4.
598        let base = agent.msg.messages.len() - 4;
599
600        agent.truncate_old_tool_results();
601
602        // Old ToolResult truncated.
603        if let MessagePart::ToolResult { content, .. } = &agent.msg.messages[base].parts[0] {
604            assert!(
605                content.ends_with("…[truncated]"),
606                "msg[base] should be truncated"
607            );
608            assert!(content.len() <= 2048 + 16);
609        } else {
610            panic!("expected ToolResult at msg[base]");
611        }
612
613        // Old ToolOutput truncated.
614        if let MessagePart::ToolOutput { body, .. } = &agent.msg.messages[base + 1].parts[0] {
615            assert!(
616                body.ends_with("…[truncated]"),
617                "msg[base+1] should be truncated"
618            );
619        } else {
620            panic!("expected ToolOutput at msg[base+1]");
621        }
622
623        // Recent messages untouched.
624        if let MessagePart::ToolResult { content, .. } = &agent.msg.messages[base + 2].parts[0] {
625            assert_eq!(content.len(), 4096, "msg[base+2] should NOT be truncated");
626        } else {
627            panic!("expected ToolResult at msg[base+2]");
628        }
629        if let MessagePart::ToolResult { content, .. } = &agent.msg.messages[base + 3].parts[0] {
630            assert_eq!(content.len(), 4096, "msg[base+3] should NOT be truncated");
631        } else {
632            panic!("expected ToolResult at msg[base+3]");
633        }
634    }
635
636    #[test]
637    fn truncate_old_tool_results_noop_when_few_messages() {
638        let provider = mock_provider(vec![]);
639        let channel = MockChannel::new(vec![]);
640        let registry = create_test_registry();
641        let executor = MockToolExecutor::no_tools();
642        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
643
644        let big = "y".repeat(4096);
645        agent.msg.messages.push(Message {
646            role: Role::User,
647            content: String::new(),
648            parts: vec![MessagePart::ToolResult {
649                tool_use_id: "id".to_string(),
650                content: big.clone(),
651                is_error: false,
652            }],
653            metadata: MessageMetadata::default(),
654        });
655        agent.msg.messages.push(Message {
656            role: Role::Assistant,
657            content: "ok".to_string(),
658            parts: vec![MessagePart::ToolResult {
659                tool_use_id: "id2".to_string(),
660                content: big.clone(),
661                is_error: false,
662            }],
663            metadata: MessageMetadata::default(),
664        });
665
666        // Agent::new inserts a system prompt at index 0; our messages are at 1 and 2.
667        let len_before = agent.msg.messages.len();
668        agent.truncate_old_tool_results();
669
670        // Neither message truncated — both fall in the last-2 window (len=3, skip last 2).
671        assert_eq!(agent.msg.messages.len(), len_before);
672        if let MessagePart::ToolResult { content, .. } =
673            &agent.msg.messages[len_before - 2].parts[0]
674        {
675            assert_eq!(
676                content.len(),
677                4096,
678                "second-to-last should not be truncated"
679            );
680        } else {
681            panic!("expected ToolResult");
682        }
683        if let MessagePart::ToolResult { content, .. } =
684            &agent.msg.messages[len_before - 1].parts[0]
685        {
686            assert_eq!(content.len(), 4096, "last should not be truncated");
687        } else {
688            panic!("expected ToolResult");
689        }
690    }
691
692    fn make_timings(ctx: u64, llm: u64, tool: u64, persist: u64) -> crate::metrics::TurnTimings {
693        crate::metrics::TurnTimings {
694            prepare_context_ms: ctx,
695            llm_chat_ms: llm,
696            tool_exec_ms: tool,
697            persist_message_ms: persist,
698        }
699    }
700
701    fn agent_with_metrics_watch() -> (
702        Agent<MockChannel>,
703        tokio::sync::watch::Receiver<crate::metrics::MetricsSnapshot>,
704    ) {
705        let provider = mock_provider(vec![]);
706        let channel = MockChannel::new(vec![]);
707        let registry = create_test_registry();
708        let executor = MockToolExecutor::no_tools();
709        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
710
711        let (tx, rx) = tokio::sync::watch::channel(crate::metrics::MetricsSnapshot::default());
712        agent.runtime.metrics.metrics_tx = Some(tx);
713        (agent, rx)
714    }
715
716    // T1-a: single flush — last_turn_timings equals the flushed value, count == 1.
717    #[test]
718    fn flush_turn_timings_single_flush() {
719        let (mut agent, rx) = agent_with_metrics_watch();
720
721        agent.runtime.metrics.pending_timings = make_timings(10, 200, 50, 5);
722        agent.flush_turn_timings();
723
724        let snap = rx.borrow();
725        assert_eq!(snap.last_turn_timings.prepare_context_ms, 10);
726        assert_eq!(snap.last_turn_timings.llm_chat_ms, 200);
727        assert_eq!(snap.last_turn_timings.tool_exec_ms, 50);
728        assert_eq!(snap.last_turn_timings.persist_message_ms, 5);
729        assert_eq!(snap.timing_sample_count, 1);
730        // avg == last when sample_count == 1
731        assert_eq!(snap.avg_turn_timings.llm_chat_ms, 200);
732    }
733
734    // T1-b: pending_timings reset to default after flush.
735    #[test]
736    fn flush_turn_timings_resets_pending() {
737        let provider = mock_provider(vec![]);
738        let channel = MockChannel::new(vec![]);
739        let registry = create_test_registry();
740        let executor = MockToolExecutor::no_tools();
741        let mut agent = Agent::new(provider, channel, registry, None, 5, executor);
742
743        agent.runtime.metrics.pending_timings = make_timings(10, 200, 50, 5);
744        agent.flush_turn_timings();
745
746        let p = &agent.runtime.metrics.pending_timings;
747        assert_eq!(p.prepare_context_ms, 0);
748        assert_eq!(p.llm_chat_ms, 0);
749        assert_eq!(p.tool_exec_ms, 0);
750        assert_eq!(p.persist_message_ms, 0);
751    }
752
753    // T1-c: window capped at 10; avg and max computed correctly.
754    #[test]
755    fn flush_turn_timings_window_capped_at_10() {
756        let (mut agent, rx) = agent_with_metrics_watch();
757
758        // Push 12 turns: llm_chat_ms = i * 10 for i in 1..=12.
759        for i in 1_u64..=12 {
760            agent.runtime.metrics.pending_timings = make_timings(0, i * 10, 0, 0);
761            agent.flush_turn_timings();
762        }
763
764        let snap = rx.borrow();
765        // Window holds last 10: turns 3..=12, llm values 30..=120.
766        assert_eq!(snap.timing_sample_count, 10);
767        // max = 120
768        assert_eq!(snap.max_turn_timings.llm_chat_ms, 120);
769        // avg of 30,40,...,120 = (30+120)*10/2/10 = 75
770        assert_eq!(snap.avg_turn_timings.llm_chat_ms, 75);
771    }
772}