zeph_agent_context/
helpers.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Pure helper functions for context assembly.
5//!
6//! These functions are called by `assembly.rs` in `zeph-core` (via a module alias)
7//! and by the [`crate::service::ContextService`] stubs that will be filled in during
8//! subsequent migration steps.
9//!
10//! All functions operate on [`crate::state::ContextAssemblyView`] instead of the
11//! `zeph-core`-internal `MemoryState`, keeping this crate free of `zeph-core` types.
12
13use std::fmt::Write as _;
14use std::time::Instant;
15
16use zeph_config::ContextFormat;
17use zeph_llm::provider::{Message, MessagePart, Role};
18use zeph_memory::{RetrievalFailureRecord, RetrievalFailureType, TokenCounter};
19
20use crate::error::ContextError;
21use crate::state::ContextAssemblyView;
22
23/// System message prefix for persona context injected into the system prompt.
24pub const PERSONA_PREFIX: &str = "[Persona context]\n";
25/// System message prefix for trajectory (past experience) context.
26pub const TRAJECTORY_PREFIX: &str = "[Past experience]\n";
27/// System message prefix for tree-based memory summaries.
28pub const TREE_MEMORY_PREFIX: &str = "[Memory summary]\n";
29/// System message prefix for reasoning strategy context.
30pub const REASONING_PREFIX: &str = "[Reasoning Strategy]\n";
31
32/// System message prefix for graph memory facts injected into context.
33pub const GRAPH_FACTS_PREFIX: &str = "[known facts]\n";
34/// System message prefix for semantic recall entries.
35pub const RECALL_PREFIX: &str = "[semantic recall]\n";
36/// System message prefix for session summary entries.
37pub const SUMMARY_PREFIX: &str = "[conversation summaries]\n";
38/// System message prefix for cross-session context entries.
39pub const CROSS_SESSION_PREFIX: &str = "[cross-session context]\n";
40
41/// System message prefix for past user corrections injected into context.
42pub const CORRECTIONS_PREFIX: &str = "[past corrections]\n";
43/// System message prefix for code-context (repo-map / file context) injections.
44pub const CODE_CONTEXT_PREFIX: &str = "[code context]\n";
45/// User message prefix for session digest summaries from the previous interaction.
46pub const SESSION_DIGEST_PREFIX: &str = "[Session digest from previous interaction]\n";
47/// System message prefix for LSP context notes (diagnostics, hover data, etc.).
48pub const LSP_NOTE_PREFIX: &str = "[lsp ";
49/// System message prefix for document RAG results.
50pub const DOCUMENT_RAG_PREFIX: &str = "## Relevant documents\n";
51
52/// Truncate `s` to at most `max_chars` Unicode scalar values.
53///
54/// Delegates to `zeph_common::text::truncate_to_chars` which respects UTF-8 boundaries.
55#[must_use]
56pub fn truncate_chars(s: &str, max_chars: usize) -> String {
57    zeph_common::text::truncate_to_chars(s, max_chars)
58}
59
60/// Format a user correction as a single bullet point for injection into the system prompt.
61///
62/// The `correction_text` must already be scrubbed by the caller before being passed here.
63/// Truncated to 200 characters to avoid inflating the context with verbose correction notes.
64#[must_use]
65pub fn format_correction_note(correction_text: &str) -> String {
66    format!(
67        "- Past user correction: \"{}\"",
68        truncate_chars(correction_text, 200)
69    )
70}
71
72/// Return the effective spreading-activation recall timeout in milliseconds.
73///
74/// A configured value of `0` would silently disable recall; this function clamps it to
75/// `100ms` and emits a warning so operators notice the misconfiguration without a crash.
76pub fn effective_recall_timeout_ms(configured: u64) -> u64 {
77    if configured == 0 {
78        tracing::warn!(
79            "recall_timeout_ms is 0, which would disable spreading activation recall; \
80             clamping to 100ms"
81        );
82        100
83    } else {
84        configured
85    }
86}
87
88/// Read-only inputs for [`fetch_semantic_recall_raw`]: the query, its retrieval
89/// limits/format, and the confidence threshold used to flag low-confidence recall for
90/// telemetry.
91///
92/// Distinct from [`crate::service::SemanticRecallParams`] (the service-level façade
93/// struct, which additionally carries tiered-retrieval provider/config fields) — this is
94/// the smaller subset of fields actually read by the flat (non-tiered) recall path.
95/// `memory` and `router` are kept as separate arguments on the function since they are
96/// resource handles rather than per-call query configuration.
97pub struct SemanticRecallRawParams<'a> {
98    /// Maximum number of memories to retrieve.
99    pub recall_limit: usize,
100    /// Format applied when serialising recalled memories.
101    pub context_format: ContextFormat,
102    /// Query string used for retrieval.
103    pub query: &'a str,
104    /// Maximum number of tokens the injected recall may consume.
105    pub token_budget: usize,
106    /// Token counter used to enforce `token_budget`.
107    pub tc: &'a TokenCounter,
108    /// When `Some(t)`, results with a top score below `t` are classified as
109    /// low-confidence and logged via the memory's retrieval failure logger.
110    pub low_confidence_threshold: Option<f32>,
111}
112
113/// Fetch semantically recalled messages using individual field arguments.
114///
115/// Raw-args variant used by [`fetch_semantic_recall`] and by
116/// [`crate::service::ContextService`]'s flat (non-tiered) recall path.
117///
118/// # Errors
119///
120/// Returns [`zeph_memory::MemoryError`] when the memory backend returns an error.
121#[tracing::instrument(
122    name = "agent_context.helpers.fetch_semantic_recall_raw",
123    skip_all,
124    err
125)]
126pub async fn fetch_semantic_recall_raw(
127    memory: Option<&zeph_memory::semantic::SemanticMemory>,
128    params: SemanticRecallRawParams<'_>,
129    router: Option<&dyn zeph_memory::AsyncMemoryRouter>,
130) -> Result<(Option<Message>, Option<f32>), zeph_memory::MemoryError> {
131    let Some(memory) = memory else {
132        return Ok((None, None));
133    };
134    if params.recall_limit == 0 || params.token_budget == 0 {
135        return Ok((None, None));
136    }
137
138    let t0 = Instant::now();
139    let recalled = if let Some(r) = router {
140        memory
141            .recall_routed_async(params.query, params.recall_limit, None, r, None)
142            .await?
143    } else {
144        memory
145            .recall(params.query, params.recall_limit, None)
146            .await?
147    };
148    let latency_ms = t0.elapsed().as_millis().try_into().unwrap_or(u64::MAX);
149
150    if recalled.is_empty() {
151        memory.log_retrieval_failure(RetrievalFailureRecord {
152            conversation_id: None,
153            turn_index: 0,
154            failure_type: RetrievalFailureType::NoHit,
155            retrieval_strategy: "semantic".to_owned(),
156            query_text: params.query.to_owned(),
157            query_len: params.query.len(),
158            top_score: None,
159            confidence_threshold: params.low_confidence_threshold,
160            result_count: 0,
161            latency_ms,
162            edge_types: None,
163            error_context: None,
164        });
165        return Ok((None, None));
166    }
167
168    let top_score = recalled.first().map(|r| r.score);
169
170    if let (Some(score), Some(threshold)) = (top_score, params.low_confidence_threshold)
171        && score < threshold
172    {
173        memory.log_retrieval_failure(RetrievalFailureRecord {
174            conversation_id: None,
175            turn_index: 0,
176            failure_type: RetrievalFailureType::LowConfidence,
177            retrieval_strategy: "semantic".to_owned(),
178            query_text: params.query.to_owned(),
179            query_len: params.query.len(),
180            top_score: Some(score),
181            confidence_threshold: Some(threshold),
182            result_count: recalled.len(),
183            latency_ms,
184            edge_types: None,
185            error_context: None,
186        });
187    }
188    let initial_cap = (params.recall_limit * 512).min(params.token_budget * 3);
189    let mut recall_text = String::with_capacity(initial_cap);
190    recall_text.push_str(RECALL_PREFIX);
191    let mut tokens_used = params.tc.count_tokens(&recall_text);
192
193    for item in &recalled {
194        if item.message.content.starts_with("[skipped]")
195            || item.message.content.starts_with("[stopped]")
196        {
197            continue;
198        }
199        let entry = match params.context_format {
200            ContextFormat::Structured => format_structured_recall_entry(item),
201            _ => format_plain_recall_entry(item),
202        };
203        let entry_tokens = params.tc.count_tokens(&entry);
204        if tokens_used + entry_tokens > params.token_budget {
205            break;
206        }
207        recall_text.push_str(&entry);
208        tokens_used += entry_tokens;
209    }
210
211    if tokens_used > params.tc.count_tokens(RECALL_PREFIX) {
212        Ok((
213            Some(Message::from_parts(
214                Role::System,
215                vec![MessagePart::Recall { text: recall_text }],
216            )),
217            top_score,
218        ))
219    } else {
220        Ok((None, None))
221    }
222}
223
224/// Fetch session summaries using individual field arguments.
225///
226/// Raw-args variant used by `zeph-core` test bridge methods and by [`fetch_summaries`].
227///
228/// # Errors
229///
230/// Returns [`zeph_memory::MemoryError`] when the memory backend returns an error.
231#[tracing::instrument(name = "agent_context.helpers.fetch_summaries_raw", skip_all, err)]
232pub async fn fetch_summaries_raw(
233    memory: Option<&zeph_memory::semantic::SemanticMemory>,
234    conversation_id: Option<zeph_memory::ConversationId>,
235    token_budget: usize,
236    tc: &TokenCounter,
237) -> Result<Option<Message>, zeph_memory::MemoryError> {
238    let (Some(memory), Some(cid)) = (memory, conversation_id) else {
239        return Ok(None);
240    };
241    if token_budget == 0 {
242        return Ok(None);
243    }
244
245    let summaries = memory.load_summaries(cid).await?;
246    if summaries.is_empty() {
247        return Ok(None);
248    }
249
250    let mut summary_text = String::from(SUMMARY_PREFIX);
251    let mut tokens_used = tc.count_tokens(&summary_text);
252
253    for summary in summaries.iter().rev() {
254        let first = summary.first_message_id.map_or(0, |m| m.0);
255        let last = summary.last_message_id.map_or(0, |m| m.0);
256        let entry = format!("- Messages {first}-{last}: {}\n", summary.content);
257        let cost = tc.count_tokens(&entry);
258        if tokens_used + cost > token_budget {
259            break;
260        }
261        summary_text.push_str(&entry);
262        tokens_used += cost;
263    }
264
265    if tokens_used > tc.count_tokens(SUMMARY_PREFIX) {
266        Ok(Some(Message::from_parts(
267            Role::System,
268            vec![MessagePart::Summary { text: summary_text }],
269        )))
270    } else {
271        Ok(None)
272    }
273}
274
275/// Fetch cross-session context summaries using individual field arguments.
276///
277/// Raw-args variant used by `zeph-core` test bridge methods and by [`fetch_cross_session`].
278///
279/// # Errors
280///
281/// Returns [`zeph_memory::MemoryError`] when the memory backend returns an error.
282#[tracing::instrument(name = "agent_context.helpers.fetch_cross_session_raw", skip_all, err)]
283pub async fn fetch_cross_session_raw(
284    memory: Option<&zeph_memory::semantic::SemanticMemory>,
285    conversation_id: Option<zeph_memory::ConversationId>,
286    cross_session_score_threshold: f32,
287    query: &str,
288    token_budget: usize,
289    tc: &TokenCounter,
290) -> Result<Option<Message>, zeph_memory::MemoryError> {
291    let (Some(memory), Some(cid)) = (memory, conversation_id) else {
292        return Ok(None);
293    };
294    if token_budget == 0 {
295        return Ok(None);
296    }
297
298    let results: Vec<_> = memory
299        .search_session_summaries(query, 5, Some(cid))
300        .await?
301        .into_iter()
302        .filter(|r| r.score >= cross_session_score_threshold)
303        .collect();
304    if results.is_empty() {
305        return Ok(None);
306    }
307
308    let mut text = String::from(CROSS_SESSION_PREFIX);
309    let mut tokens_used = tc.count_tokens(&text);
310
311    for item in &results {
312        let entry = format!("- {}\n", item.summary_text);
313        let cost = tc.count_tokens(&entry);
314        if tokens_used + cost > token_budget {
315            break;
316        }
317        text.push_str(&entry);
318        tokens_used += cost;
319    }
320
321    if tokens_used > tc.count_tokens(CROSS_SESSION_PREFIX) {
322        Ok(Some(Message::from_parts(
323            Role::System,
324            vec![MessagePart::CrossSession { text }],
325        )))
326    } else {
327        Ok(None)
328    }
329}
330
331/// Fetch semantically recalled messages for the given query and enforce the token budget.
332///
333/// Delegates to [`fetch_semantic_recall_raw`] using fields from `view`.
334///
335/// Returns `(None, None)` when memory is absent, recall is disabled, the budget is zero,
336/// or the recalled set is empty.
337///
338/// The second element of the tuple is the similarity score of the top recalled entry, used
339/// by the caller to track recall confidence for telemetry.
340///
341/// # Errors
342///
343/// Returns [`ContextError::Memory`] when the memory recall backend returns an error.
344#[tracing::instrument(name = "agent_context.helpers.fetch_semantic_recall", skip_all, err)]
345pub async fn fetch_semantic_recall(
346    view: &ContextAssemblyView<'_>,
347    query: &str,
348    token_budget: usize,
349    tc: &TokenCounter,
350    router: Option<&dyn zeph_memory::AsyncMemoryRouter>,
351) -> Result<(Option<Message>, Option<f32>), ContextError> {
352    fetch_semantic_recall_raw(
353        view.memory.as_deref(),
354        SemanticRecallRawParams {
355            recall_limit: view.recall_limit,
356            context_format: view.context_format,
357            query,
358            token_budget,
359            tc,
360            low_confidence_threshold: None,
361        },
362        router,
363    )
364    .await
365    .map_err(ContextError::Memory)
366}
367
368fn format_plain_recall_entry(item: &zeph_memory::RecalledMessage) -> String {
369    let role_label = match item.message.role {
370        Role::Assistant => "assistant",
371        Role::System => "system",
372        Role::User | _ => "user",
373    };
374    format!("- [{}] {}\n", role_label, item.message.content)
375}
376
377#[allow(clippy::map_unwrap_or)]
378fn format_structured_recall_entry(item: &zeph_memory::RecalledMessage) -> String {
379    let source = match item.message.role {
380        Role::Assistant => "assistant",
381        Role::System => "system",
382        Role::User | _ => "user",
383    };
384    // Use compacted_at as a proxy for message age when available; otherwise "unknown".
385    // A full timestamp lookup from SQLite would require an async DB call in the assembler
386    // and is deferred to a future enhancement (TODO: enhance when message timestamps are
387    // propagated into RecalledMessage).
388    let date = item
389        .message
390        .metadata
391        .compacted_at
392        .and_then(|ts| chrono::DateTime::from_timestamp(ts, 0))
393        .map(|dt| dt.format("%Y-%m-%d").to_string())
394        .unwrap_or_else(|| "unknown".to_owned());
395    format!(
396        "[Memory | {} | {} | relevance: {:.2}]\n{}\n",
397        source, date, item.score, item.message.content
398    )
399}
400
401/// Fetch session summaries for the current conversation and enforce the token budget.
402///
403/// Delegates to [`fetch_summaries_raw`] using fields from `view`.
404///
405/// Returns `None` when memory or the conversation ID is absent, the budget is zero,
406/// or no summaries exist yet.
407///
408/// # Errors
409///
410/// Returns [`ContextError::Memory`] when the memory backend returns an error.
411#[tracing::instrument(name = "agent_context.helpers.fetch_summaries", skip_all, err)]
412pub async fn fetch_summaries(
413    view: &ContextAssemblyView<'_>,
414    token_budget: usize,
415    tc: &TokenCounter,
416) -> Result<Option<Message>, ContextError> {
417    fetch_summaries_raw(
418        view.memory.as_deref(),
419        view.conversation_id,
420        token_budget,
421        tc,
422    )
423    .await
424    .map_err(ContextError::Memory)
425}
426
427/// Fetch cross-session context summaries for the given query and enforce the token budget.
428///
429/// Delegates to [`fetch_cross_session_raw`] using fields from `view`.
430///
431/// Results are filtered by `view.cross_session_score_threshold` before token counting,
432/// and the current conversation is excluded from the search results.
433///
434/// Returns `None` when memory or the conversation ID is absent, the budget is zero,
435/// no results exceed the threshold, or the result set is empty.
436///
437/// # Errors
438///
439/// Returns [`ContextError::Memory`] when the memory backend returns an error.
440#[tracing::instrument(name = "agent_context.helpers.fetch_cross_session", skip_all, err)]
441pub async fn fetch_cross_session(
442    view: &ContextAssemblyView<'_>,
443    query: &str,
444    token_budget: usize,
445    tc: &TokenCounter,
446) -> Result<Option<Message>, ContextError> {
447    fetch_cross_session_raw(
448        view.memory.as_deref(),
449        view.conversation_id,
450        view.cross_session_score_threshold,
451        query,
452        token_budget,
453        tc,
454    )
455    .await
456    .map_err(ContextError::Memory)
457}
458
459/// Budget state injected into the volatile system prompt section.
460///
461/// All fields are optional — omitted when the corresponding data source is unavailable.
462/// [`BudgetHint::format_xml`] returns `None` when all fields would be absent.
463///
464/// Callers should construct this from cost-tracker and tool-orchestrator state, then call
465/// `format_xml` and append the result to the system prompt when `Some`.
466pub struct BudgetHint {
467    /// Remaining daily budget in US cents, if a daily limit is configured.
468    pub remaining_cost_cents: Option<f64>,
469    /// Total daily budget in US cents, if a daily limit is configured.
470    pub total_budget_cents: Option<f64>,
471    /// Remaining tool-call iterations this turn.
472    pub remaining_tool_calls: usize,
473    /// Maximum allowed tool-call iterations per turn (0 = no limit configured).
474    pub max_tool_calls: usize,
475}
476
477impl BudgetHint {
478    /// Render the budget hint as an XML fragment for injection into the system prompt.
479    ///
480    /// Returns `None` when no meaningful budget data is available — callers must skip
481    /// injection rather than injecting an empty `<budget></budget>` block.
482    ///
483    /// # Examples
484    ///
485    /// ```
486    /// use zeph_agent_context::helpers::BudgetHint;
487    ///
488    /// let hint = BudgetHint {
489    ///     remaining_cost_cents: Some(50.0),
490    ///     total_budget_cents: Some(100.0),
491    ///     remaining_tool_calls: 8,
492    ///     max_tool_calls: 10,
493    /// };
494    /// let xml = hint.format_xml().unwrap();
495    /// assert!(xml.contains("<remaining_cost_cents>50.00</remaining_cost_cents>"));
496    /// assert!(xml.contains("<remaining_tool_calls>8</remaining_tool_calls>"));
497    /// ```
498    #[must_use]
499    pub fn format_xml(&self) -> Option<String> {
500        let has_cost = self.remaining_cost_cents.is_some();
501        // Always include tool call budget — max_tool_calls > 0 in any real config.
502        if !has_cost && self.max_tool_calls == 0 {
503            return None;
504        }
505        let mut s = String::from("<budget>");
506        if let Some(remaining) = self.remaining_cost_cents {
507            let _ = write!(
508                s,
509                "\n<remaining_cost_cents>{remaining:.2}</remaining_cost_cents>"
510            );
511        }
512        if let Some(total) = self.total_budget_cents {
513            let _ = write!(s, "\n<total_budget_cents>{total:.2}</total_budget_cents>");
514        }
515        if self.max_tool_calls > 0 {
516            let _ = write!(
517                s,
518                "\n<remaining_tool_calls>{}</remaining_tool_calls>",
519                self.remaining_tool_calls
520            );
521            let _ = write!(
522                s,
523                "\n<max_tool_calls>{}</max_tool_calls>",
524                self.max_tool_calls
525            );
526        }
527        s.push_str("\n</budget>");
528        Some(s)
529    }
530}
531
532#[cfg(test)]
533mod budget_hint_tests {
534    use super::*;
535
536    #[test]
537    fn format_xml_none_when_no_data() {
538        let hint = BudgetHint {
539            remaining_cost_cents: None,
540            total_budget_cents: None,
541            remaining_tool_calls: 0,
542            max_tool_calls: 0,
543        };
544        assert!(hint.format_xml().is_none());
545    }
546
547    #[test]
548    fn format_xml_with_cost_only() {
549        let hint = BudgetHint {
550            remaining_cost_cents: Some(25.5),
551            total_budget_cents: Some(100.0),
552            remaining_tool_calls: 0,
553            max_tool_calls: 0,
554        };
555        let xml = hint.format_xml().unwrap();
556        assert!(xml.contains("<remaining_cost_cents>25.50</remaining_cost_cents>"));
557        assert!(xml.contains("<total_budget_cents>100.00</total_budget_cents>"));
558    }
559
560    #[test]
561    fn format_xml_with_tool_calls_only() {
562        let hint = BudgetHint {
563            remaining_cost_cents: None,
564            total_budget_cents: None,
565            remaining_tool_calls: 3,
566            max_tool_calls: 10,
567        };
568        let xml = hint.format_xml().unwrap();
569        assert!(xml.contains("<remaining_tool_calls>3</remaining_tool_calls>"));
570        assert!(xml.contains("<max_tool_calls>10</max_tool_calls>"));
571    }
572
573    #[test]
574    fn format_xml_with_all_fields() {
575        let hint = BudgetHint {
576            remaining_cost_cents: Some(50.0),
577            total_budget_cents: Some(100.0),
578            remaining_tool_calls: 8,
579            max_tool_calls: 10,
580        };
581        let xml = hint.format_xml().unwrap();
582        assert!(xml.starts_with("<budget>"));
583        assert!(xml.ends_with("</budget>"));
584    }
585}
zeph_agent_context/helpers.rs

zeph_agent_context/
helpers.rs
