agentd/context/
mod.rs

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

agentd/context/
mod.rs
