agentd/context/
mod.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! **Contexts**: the durable, self-compacting working memory of the root agent
3//! (`context/root`) and of every A2A conversation (`context/<contextId>`). A
4//! context is a versioned record of messages, a structured summary block, the
5//! loaded skill set, the working **plan**, the last preflight verdict and a
6//! token estimate. The runtime is its single writer and checkpoints it after
7//! every turn, so nothing here needs locking and a crash loses at most the
8//! turn in flight.
9//!
10//! The transcript representation here ([`Msg`]) is **serializable** (unlike
11//! the provider wire type) and converts to [`crate::wire::intel::Message`]
12//! at request time. Sub-modules: [`plan`] (the plan object), [`memory`] (the
13//! durable KV), [`compact`] (compaction planning + application), [`skills`]
14//! (skill catalogue/loaded set), [`tokens`] (estimates).
15
16pub mod compact;
17pub mod memory;
18pub mod plan;
19pub mod prompt;
20pub mod skills;
21pub mod tokens;
22
23use crate::state::{Durable, Kind, now_ms};
24use crate::store::StoreError;
25use crate::wire::intel::{Message, ToolCall};
26use serde::{Deserialize, Serialize};
27use serde_json::{Value, json};
28use std::collections::{BTreeMap, BTreeSet};
29
30/// The root context id: the one context that always exists, holding the root
31/// agent's own working memory rather than any caller's conversation.
32pub const ROOT: &str = "root";
33
34/// One transcript entry. `ts` is wall-clock ms; tool results keep the parsed
35/// JSON when the tool returned structured content (text otherwise).
36#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
37#[serde(tag = "role", rename_all = "snake_case")]
38pub enum Msg {
39    System {
40        text: String,
41        #[serde(default)]
42        ts: u64,
43    },
44    User {
45        text: String,
46        #[serde(default, skip_serializing_if = "Option::is_none")]
47        principal: Option<String>,
48        #[serde(default)]
49        ts: u64,
50    },
51    Assistant {
52        #[serde(default, skip_serializing_if = "Option::is_none")]
53        text: Option<String>,
54        #[serde(default, skip_serializing_if = "Vec::is_empty")]
55        tool_calls: Vec<ToolCall>,
56        #[serde(default)]
57        ts: u64,
58    },
59    Tool {
60        id: String,
61        name: String,
62        content: Value,
63        #[serde(default)]
64        is_error: bool,
65        #[serde(default)]
66        ts: u64,
67    },
68    /// A runtime note (a run finished, a subagent reported, an instruction
69    /// changed…) — rendered to the model as a system message.
70    Note {
71        text: String,
72        #[serde(default)]
73        ts: u64,
74    },
75}
76
77impl Msg {
78    pub fn system(text: impl Into<String>) -> Msg {
79        Msg::System {
80            text: text.into(),
81            ts: now_ms(),
82        }
83    }
84    pub fn user(text: impl Into<String>, principal: Option<String>) -> Msg {
85        Msg::User {
86            text: text.into(),
87            principal,
88            ts: now_ms(),
89        }
90    }
91    pub fn assistant(text: Option<String>, tool_calls: Vec<ToolCall>) -> Msg {
92        Msg::Assistant {
93            text,
94            tool_calls,
95            ts: now_ms(),
96        }
97    }
98    pub fn tool(
99        id: impl Into<String>,
100        name: impl Into<String>,
101        content: Value,
102        is_error: bool,
103    ) -> Msg {
104        Msg::Tool {
105            id: id.into(),
106            name: name.into(),
107            content,
108            is_error,
109            ts: now_ms(),
110        }
111    }
112    pub fn note(text: impl Into<String>) -> Msg {
113        Msg::Note {
114            text: text.into(),
115            ts: now_ms(),
116        }
117    }
118    pub fn ts(&self) -> u64 {
119        match self {
120            Msg::System { ts, .. }
121            | Msg::User { ts, .. }
122            | Msg::Assistant { ts, .. }
123            | Msg::Tool { ts, .. }
124            | Msg::Note { ts, .. } => *ts,
125        }
126    }
127    /// The provider wire message.
128    pub fn to_wire(&self) -> Message {
129        match self {
130            Msg::System { text, .. } => Message::System(text.clone()),
131            Msg::Note { text, .. } => Message::System(format!("[note] {text}")),
132            Msg::User { text, .. } => Message::User(text.clone()),
133            Msg::Assistant {
134                text, tool_calls, ..
135            } => Message::Assistant {
136                text: text.clone(),
137                tool_calls: tool_calls.clone(),
138            },
139            Msg::Tool {
140                id,
141                content,
142                is_error,
143                ..
144            } => Message::ToolResult {
145                id: id.clone(),
146                content: match content {
147                    Value::String(s) => s.clone(),
148                    other => other.to_string(),
149                },
150                is_error: *is_error,
151            },
152        }
153    }
154    /// A rough token estimate for this message.
155    pub fn est_tokens(&self) -> u64 {
156        let body = match self {
157            Msg::System { text, .. } | Msg::User { text, .. } | Msg::Note { text, .. } => {
158                tokens::estimate(text)
159            }
160            Msg::Assistant {
161                text, tool_calls, ..
162            } => {
163                tokens::estimate(text.as_deref().unwrap_or(""))
164                    + tool_calls
165                        .iter()
166                        .map(|c| tokens::estimate(&c.name) + tokens::estimate_value(&c.arguments))
167                        .sum::<u64>()
168            }
169            Msg::Tool { content, .. } => tokens::estimate_value(content),
170        };
171        body + tokens::MESSAGE_OVERHEAD
172    }
173    pub fn is_user(&self) -> bool {
174        matches!(self, Msg::User { .. })
175    }
176}
177
178/// The structured summary block a compaction produces: what the dropped
179/// messages amounted to, kept in fields so a later compaction can merge two
180/// summaries instead of summarising a summary.
181#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
182pub struct Summary {
183    #[serde(default, skip_serializing_if = "Vec::is_empty")]
184    pub goals: Vec<String>,
185    #[serde(default, skip_serializing_if = "Vec::is_empty")]
186    pub decisions: Vec<String>,
187    #[serde(default, skip_serializing_if = "Vec::is_empty")]
188    pub open: Vec<String>,
189    #[serde(default, skip_serializing_if = "Vec::is_empty")]
190    pub facts: Vec<String>,
191    /// Free-form narrative when the summarizer could not fill the fields.
192    #[serde(default, skip_serializing_if = "Option::is_none")]
193    pub narrative: Option<String>,
194    /// How many messages the summary stands for (cumulative).
195    #[serde(default)]
196    pub covers_messages: u64,
197    #[serde(default)]
198    pub updated: u64,
199}
200
201impl Summary {
202    pub fn is_empty(&self) -> bool {
203        self.goals.is_empty()
204            && self.decisions.is_empty()
205            && self.open.is_empty()
206            && self.facts.is_empty()
207            && self.narrative.as_deref().is_none_or(str::is_empty)
208    }
209    /// The block as the model sees it.
210    pub fn render(&self) -> String {
211        let mut out = String::from("Summary of earlier conversation:\n");
212        let sect = |out: &mut String, title: &str, items: &[String]| {
213            if !items.is_empty() {
214                out.push_str(title);
215                out.push('\n');
216                for i in items {
217                    out.push_str("- ");
218                    out.push_str(i);
219                    out.push('\n');
220                }
221            }
222        };
223        sect(&mut out, "Goals:", &self.goals);
224        sect(&mut out, "Decisions:", &self.decisions);
225        sect(&mut out, "Open items:", &self.open);
226        sect(&mut out, "Facts:", &self.facts);
227        if let Some(n) = &self.narrative
228            && !n.is_empty()
229        {
230            out.push_str(n);
231            out.push('\n');
232        }
233        out
234    }
235    /// Merge a newer summary over this one (lists appended + deduped, capped).
236    pub fn absorb(&mut self, newer: Summary) {
237        fn merge(into: &mut Vec<String>, more: Vec<String>) {
238            for m in more {
239                if !into.contains(&m) {
240                    into.push(m);
241                }
242            }
243            if into.len() > 32 {
244                let drop = into.len() - 32;
245                into.drain(0..drop);
246            }
247        }
248        merge(&mut self.goals, newer.goals);
249        merge(&mut self.decisions, newer.decisions);
250        merge(&mut self.open, newer.open);
251        merge(&mut self.facts, newer.facts);
252        if newer.narrative.is_some() {
253            self.narrative = newer.narrative;
254        }
255        self.covers_messages += newer.covers_messages;
256        self.updated = now_ms();
257    }
258}
259
260/// A loaded skill reference: its name plus the hash of the version that was
261/// loaded, so a skill edited on disk is detectable rather than silently
262/// diverging from what the transcript was built against.
263#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
264pub struct SkillRef {
265    pub name: String,
266    pub hash: String,
267}
268
269/// The kind of a context.
270#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
271#[serde(rename_all = "snake_case")]
272pub enum ContextKind {
273    #[default]
274    Root,
275    Conversation,
276}
277
278/// The durable context record: everything about a conversation that must
279/// survive a restart.
280#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
281pub struct ContextState {
282    #[serde(default)]
283    pub kind: ContextKind,
284    #[serde(default)]
285    pub version: u64,
286    #[serde(default)]
287    pub summary: Summary,
288    #[serde(default)]
289    pub messages: Vec<Msg>,
290    #[serde(default, skip_serializing_if = "Vec::is_empty")]
291    pub skills: Vec<SkillRef>,
292    #[serde(default, skip_serializing_if = "Option::is_none")]
293    pub plan: Option<plan::Plan>,
294    #[serde(default, skip_serializing_if = "Option::is_none")]
295    pub preflight: Option<Value>,
296    #[serde(default)]
297    pub est_tokens: u64,
298    #[serde(default)]
299    pub model_window: u64,
300    #[serde(default, skip_serializing_if = "Option::is_none")]
301    pub principal: Option<String>,
302    /// The A2A task the conversation's current work is attached to.
303    #[serde(default, skip_serializing_if = "Option::is_none")]
304    pub task: Option<String>,
305    #[serde(default)]
306    pub turns: u64,
307    #[serde(default)]
308    pub created: u64,
309    #[serde(default)]
310    pub updated: u64,
311    /// Whether the record has changed since its last checkpoint (never stored).
312    #[serde(skip)]
313    pub dirty: bool,
314}
315
316impl ContextState {
317    pub fn new(kind: ContextKind, model_window: u64) -> ContextState {
318        ContextState {
319            kind,
320            version: 1,
321            summary: Summary::default(),
322            messages: Vec::new(),
323            skills: Vec::new(),
324            plan: None,
325            preflight: None,
326            est_tokens: 0,
327            model_window,
328            principal: None,
329            task: None,
330            turns: 0,
331            created: now_ms(),
332            updated: now_ms(),
333            dirty: true,
334        }
335    }
336
337    pub fn append(&mut self, msg: Msg) {
338        self.est_tokens += msg.est_tokens();
339        self.messages.push(msg);
340        self.touch();
341    }
342
343    pub fn append_all(&mut self, msgs: impl IntoIterator<Item = Msg>) {
344        for m in msgs {
345            self.append(m);
346        }
347    }
348
349    pub fn touch(&mut self) {
350        self.updated = now_ms();
351        self.dirty = true;
352    }
353
354    /// Recompute the token estimate from scratch (after compaction / restore).
355    pub fn recount(&mut self) {
356        self.est_tokens = self.messages.iter().map(Msg::est_tokens).sum::<u64>()
357            + tokens::estimate(&self.summary.render())
358            + self
359                .plan
360                .as_ref()
361                .map(|p| tokens::estimate(&p.render()))
362                .unwrap_or(0);
363    }
364
365    /// Whether the compaction threshold is crossed (`compact_at` × window).
366    pub fn needs_compaction(&self, compact_at: f64) -> bool {
367        self.model_window > 0 && (self.est_tokens as f64) > compact_at * (self.model_window as f64)
368    }
369
370    /// The transcript slice a turn worker receives: summary block (if any) +
371    /// plan block (if any) as system messages, then the messages verbatim.
372    pub fn slice(&self) -> Vec<Msg> {
373        let mut out = Vec::with_capacity(self.messages.len() + 2);
374        if !self.summary.is_empty() {
375            out.push(Msg::system(self.summary.render()));
376        }
377        if let Some(p) = &self.plan {
378            out.push(Msg::system(p.render()));
379        }
380        out.extend(self.messages.iter().cloned());
381        out
382    }
383
384    /// The prompt slice for a turn: summary block (if any) + plan block (if
385    /// any) as system messages, then the messages verbatim.
386    pub fn to_wire(&self) -> Vec<Message> {
387        let mut out = Vec::with_capacity(self.messages.len() + 2);
388        if !self.summary.is_empty() {
389            out.push(Message::System(self.summary.render()));
390        }
391        if let Some(p) = &self.plan {
392            out.push(Message::System(p.render()));
393        }
394        out.extend(self.messages.iter().map(Msg::to_wire));
395        out
396    }
397
398    /// The loaded skill names.
399    pub fn skill_names(&self) -> BTreeSet<String> {
400        self.skills.iter().map(|s| s.name.clone()).collect()
401    }
402
403    pub fn load_skill(&mut self, name: &str, hash: &str, max_loaded: usize) -> Result<(), String> {
404        if let Some(s) = self.skills.iter_mut().find(|s| s.name == name) {
405            s.hash = hash.to_string();
406            self.touch();
407            return Ok(());
408        }
409        if self.skills.len() >= max_loaded {
410            return Err(format!(
411                "skills.max_loaded ({max_loaded}) reached; unload one first"
412            ));
413        }
414        self.skills.push(SkillRef {
415            name: name.to_string(),
416            hash: hash.to_string(),
417        });
418        self.touch();
419        Ok(())
420    }
421
422    pub fn unload_skill(&mut self, name: &str) -> bool {
423        let before = self.skills.len();
424        self.skills.retain(|s| s.name != name);
425        if self.skills.len() != before {
426            self.touch();
427            true
428        } else {
429            false
430        }
431    }
432}
433
434/// The in-memory registry of contexts + their durable mirror.
435pub struct Contexts {
436    map: BTreeMap<String, ContextState>,
437    model_window: u64,
438}
439
440impl Contexts {
441    pub fn new(model_window: u64) -> Contexts {
442        Contexts {
443            map: BTreeMap::new(),
444            model_window,
445        }
446    }
447
448    /// Adopt restored context envelopes.
449    pub fn restore(&mut self, envelopes: &[crate::store::Envelope]) -> Vec<String> {
450        let mut lost = Vec::new();
451        for env in envelopes {
452            match serde_json::from_value::<ContextState>(env.state.clone()) {
453                Ok(mut c) => {
454                    c.dirty = false;
455                    if c.model_window == 0 {
456                        c.model_window = self.model_window;
457                    }
458                    c.recount();
459                    self.map.insert(env.id.clone(), c);
460                }
461                Err(_) => lost.push(env.id.clone()),
462            }
463        }
464        lost
465    }
466
467    pub fn get(&self, id: &str) -> Option<&ContextState> {
468        self.map.get(id)
469    }
470    pub fn get_mut(&mut self, id: &str) -> Option<&mut ContextState> {
471        self.map.get_mut(id)
472    }
473    /// Get or create the root context.
474    pub fn root(&mut self) -> &mut ContextState {
475        let w = self.model_window;
476        self.map
477            .entry(ROOT.to_string())
478            .or_insert_with(|| ContextState::new(ContextKind::Root, w))
479    }
480    /// Get or create a conversation context.
481    pub fn conversation(&mut self, id: &str, principal: Option<&str>) -> &mut ContextState {
482        let w = self.model_window;
483        let c = self.map.entry(id.to_string()).or_insert_with(|| {
484            let mut c = ContextState::new(ContextKind::Conversation, w);
485            c.principal = principal.map(str::to_string);
486            c
487        });
488        if c.principal.is_none() && principal.is_some() {
489            c.principal = principal.map(str::to_string);
490        }
491        c
492    }
493    pub fn ids(&self) -> Vec<String> {
494        self.map.keys().cloned().collect()
495    }
496    pub fn len(&self) -> usize {
497        self.map.len()
498    }
499    /// The token estimate of the largest live context (for `agent_context_tokens`).
500    pub fn max_est_tokens(&self) -> u64 {
501        self.map.values().map(|c| c.est_tokens).max().unwrap_or(0)
502    }
503    pub fn is_empty(&self) -> bool {
504        self.map.is_empty()
505    }
506    pub fn remove(&mut self, id: &str) -> Option<ContextState> {
507        self.map.remove(id)
508    }
509
510    /// Checkpoint every dirty context — called after each turn and after each
511    /// compaction, which bounds a crash's loss to the work in flight.
512    ///
513    /// Returns the ids written. A clean context is skipped, so repeated calls
514    /// are cheap and the write count tracks real change.
515    pub fn checkpoint(&mut self, durable: &Durable) -> Result<Vec<String>, StoreError> {
516        let mut written = Vec::new();
517        for (id, c) in self.map.iter_mut() {
518            if !c.dirty {
519                continue;
520            }
521            crate::state::kill_point("context.before_put");
522            durable.put(
523                Kind::Context,
524                id,
525                serde_json::to_value(&*c).unwrap_or(Value::Null),
526                None,
527            )?;
528            c.dirty = false;
529            written.push(id.clone());
530        }
531        Ok(written)
532    }
533
534    /// A status view (`agent://conversations`).
535    pub fn status(&self) -> Value {
536        json!(
537            self.map
538                .iter()
539                .map(|(id, c)| {
540                    json!({
541                        "id": id, "kind": c.kind, "version": c.version, "messages": c.messages.len(),
542                        "est_tokens": c.est_tokens, "turns": c.turns, "principal": c.principal,
543                        "skills": c.skills.iter().map(|s| s.name.clone()).collect::<Vec<_>>(),
544                        "plan": c.plan.as_ref().map(|p| p.progress()),
545                        "updated": c.updated,
546                    })
547                })
548                .collect::<Vec<_>>()
549        )
550    }
551}
552
553#[cfg(test)]
554mod tests {
555    use super::*;
556    use crate::store::memory::MemoryStore;
557    use std::sync::Arc;
558
559    #[test]
560    fn messages_round_trip_and_convert_to_wire() {
561        let m = Msg::tool("c1", "memory.get", json!({"value": 1}), false);
562        let v = serde_json::to_value(&m).unwrap();
563        assert_eq!(v["role"], json!("tool"));
564        let back: Msg = serde_json::from_value(v).unwrap();
565        assert_eq!(back, m);
566        match back.to_wire() {
567            Message::ToolResult {
568                id,
569                content,
570                is_error,
571            } => {
572                assert_eq!(id, "c1");
573                assert_eq!(content, r#"{"value":1}"#);
574                assert!(!is_error);
575            }
576            other => panic!("{other:?}"),
577        }
578        assert!(
579            matches!(Msg::note("run finished").to_wire(), Message::System(s) if s.starts_with("[note]"))
580        );
581        assert!(Msg::user("hello world", None).est_tokens() > tokens::MESSAGE_OVERHEAD);
582    }
583
584    #[test]
585    fn contexts_checkpoint_dirty_only_and_restore() {
586        let mem = Arc::new(MemoryStore::new());
587        let d = Durable::new(
588            mem.clone(),
589            "agentd",
590            "i",
591            crate::state::Policy::default(),
592            None,
593        );
594        let mut cs = Contexts::new(100_000);
595        cs.root().append(Msg::user("hi", None));
596        cs.conversation("ctx-1", Some("user:a"))
597            .append(Msg::user("q", Some("user:a".into())));
598        let written = cs.checkpoint(&d).unwrap();
599        assert_eq!(written, vec!["ctx-1".to_string(), "root".to_string()]);
600        assert!(
601            cs.checkpoint(&d).unwrap().is_empty(),
602            "clean after checkpoint"
603        );
604        cs.get_mut("ctx-1")
605            .unwrap()
606            .append(Msg::assistant(Some("a".into()), vec![]));
607        assert_eq!(cs.checkpoint(&d).unwrap(), vec!["ctx-1".to_string()]);
608        // Restore into a fresh registry.
609        let restored = d.restore().unwrap();
610        let mut cs2 = Contexts::new(100_000);
611        assert!(cs2.restore(restored.of(Kind::Context)).is_empty());
612        assert_eq!(cs2.len(), 2);
613        let c = cs2.get("ctx-1").unwrap();
614        assert_eq!(c.messages.len(), 2);
615        assert_eq!(c.principal.as_deref(), Some("user:a"));
616        assert!(!c.dirty);
617        assert!(c.est_tokens > 0);
618    }
619
620    #[test]
621    fn summary_renders_and_absorbs() {
622        let mut s = Summary {
623            goals: vec!["ship".into()],
624            ..Default::default()
625        };
626        assert!(s.render().contains("Goals:\n- ship"));
627        s.absorb(Summary {
628            goals: vec!["ship".into(), "test".into()],
629            facts: vec!["x=1".into()],
630            covers_messages: 5,
631            ..Default::default()
632        });
633        assert_eq!(s.goals, vec!["ship".to_string(), "test".to_string()]);
634        assert_eq!(s.facts, vec!["x=1".to_string()]);
635        assert_eq!(s.covers_messages, 5);
636        assert!(!s.is_empty());
637    }
638
639    #[test]
640    fn skills_load_unload_and_caps() {
641        let mut c = ContextState::new(ContextKind::Conversation, 1000);
642        c.load_skill("a", "h1", 2).unwrap();
643        c.load_skill("b", "h2", 2).unwrap();
644        assert!(c.load_skill("c", "h3", 2).is_err());
645        c.load_skill("a", "h9", 2).unwrap();
646        assert_eq!(c.skills[0].hash, "h9");
647        assert!(c.unload_skill("a"));
648        assert!(!c.unload_skill("a"));
649        assert_eq!(c.skill_names().len(), 1);
650    }
651}
agentd/context/mod.rs

agentd/context/
mod.rs
