Skip to main content

magi_code/context/
mod.rs

1#[cfg(test)]
2mod cache;
3mod replay;
4mod replay_cache;
5mod tokens;
6
7#[cfg(test)]
8pub(crate) use cache::conversation_cache_material;
9#[cfg(test)]
10pub(crate) use replay::{
11    ConversationReplay, REPLAY_TOOL_RESULT_OUTPUT_CHAR_LIMIT, build_conversation_replay,
12    build_conversation_replay_with_limits, compact_historical_tool_output,
13};
14pub(crate) use replay::{
15    ConversationReplayLimits, REPLAY_JSONL_MAX_BYTES, REPLAY_JSONL_MAX_LINES,
16    build_conversation_replay_from_events,
17};
18pub(crate) use replay_cache::ConversationReplayCache;
19#[cfg(test)]
20pub(crate) use replay_cache::ReplayCacheMetrics;
21#[cfg(test)]
22pub(crate) use tokens::estimate_provider_request_input_tokens;
23pub(crate) use tokens::estimate_text_tokens;
24pub use tokens::{
25    ContextBudget, ContextBudgetOverride, estimate_messages_tokens, usage_input_tokens,
26};
27pub(crate) use tokens::{
28    ContextTokenCount, count_supported_provider_conversation_item_tokens,
29    project_provider_conversation_items_tokens, project_provider_request_input_tokens,
30    project_text_tokens, token_encoding_family_for_model,
31};
32
33#[cfg(test)]
34pub(crate) use crate::providers::{ChatMessage, ProviderConversationItem, ProviderToolResult};
35#[cfg(test)]
36use crate::sessions::{
37    Session, SessionDiagnosticOperation, SessionReadDiagnostic, report_session_diagnostic,
38};
39use crate::sessions::{SessionEvent, SessionEventKind};
40use serde_json::Value;
41#[cfg(test)]
42pub(crate) use serde_json::json;
43
44#[cfg(test)]
45const RECENT_EVENTS_CONTEXT_MAX_EVENTS: usize = 10_000;
46#[cfg(test)]
47const RECENT_EVENTS_CONTEXT_MAX_BYTES: usize = 2 * 1024 * 1024;
48// Tail read cap matches retained-byte budget: one bounded read can fill recent-context window.
49#[cfg(test)]
50const RECENT_EVENTS_CONTEXT_MAX_READ_BYTES: usize = RECENT_EVENTS_CONTEXT_MAX_BYTES;
51
52#[cfg(test)]
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub(crate) struct RecentEventsContext {
55    pub text: String,
56    pub diagnostics: Vec<SessionReadDiagnostic>,
57}
58
59#[cfg(test)]
60pub(crate) fn recent_events_context(
61    session: &Session,
62    keep_recent_tokens: usize,
63) -> anyhow::Result<String> {
64    Ok(recent_events_context_report(session, keep_recent_tokens)?.text)
65}
66
67#[cfg(test)]
68pub(crate) fn recent_events_context_report(
69    session: &Session,
70    keep_recent_tokens: usize,
71) -> anyhow::Result<RecentEventsContext> {
72    let mut selected = Vec::new();
73    let mut tokens = 0;
74    let recent = session
75        .read_recent_events_tolerant_tail(
76            RECENT_EVENTS_CONTEXT_MAX_EVENTS,
77            RECENT_EVENTS_CONTEXT_MAX_BYTES,
78            RECENT_EVENTS_CONTEXT_MAX_READ_BYTES,
79        )
80        .inspect_err(|error| {
81            report_session_diagnostic(
82                SessionDiagnosticOperation::RecentContext,
83                session.path(),
84                error,
85            );
86        })?;
87    for event in recent.events.into_iter().rev() {
88        if !is_provider_visible_session_event(&event) {
89            continue;
90        }
91        let line = event_line(&event);
92        let line_tokens = estimate_text_tokens(&line);
93        if !selected.is_empty() && tokens + line_tokens > keep_recent_tokens {
94            break;
95        }
96        tokens += line_tokens;
97        selected.push(line);
98    }
99    selected.reverse();
100    Ok(RecentEventsContext {
101        text: selected.join("\n"),
102        diagnostics: recent.diagnostics,
103    })
104}
105
106#[cfg(test)]
107fn is_provider_visible_session_event(event: &SessionEvent) -> bool {
108    !is_local_only_session_event(event) && event.kind() != Some(SessionEventKind::ReasoningSummary)
109}
110
111pub(crate) fn is_local_only_session_event(event: &SessionEvent) -> bool {
112    event.kind().is_some_and(SessionEventKind::is_local_only)
113}
114
115pub(crate) fn event_line(event: &SessionEvent) -> String {
116    let compact_payload = if event.kind() == Some(SessionEventKind::Compaction) {
117        event
118            .payload
119            .get("summary")
120            .and_then(Value::as_str)
121            .unwrap_or("[invalid compaction summary]")
122            .to_string()
123    } else if let Some(text) = event.payload.get("text").and_then(Value::as_str) {
124        text.to_string()
125    } else if let Some(message) = event.payload.get("message").and_then(Value::as_str) {
126        message.to_string()
127    } else {
128        event.payload.to_string()
129    };
130    format!(
131        "{}: {}",
132        event.event_type,
133        compact_payload.replace('\n', "\\n")
134    )
135}
136
137pub(crate) fn canonical_json(value: &Value) -> String {
138    serde_json::to_string(value).unwrap_or_else(|_| "null".to_string())
139}
140
141#[cfg(test)]
142mod tests;