agent_runtime_kernel/
context.rs

1use runtime_extension_api::ContextItem;
2use runtime_kernel_api::KernelFailure;
3use runtime_ports::{ConversationMessage, ModelMessage, ModelRole};
4
5/// Deterministic initial context assembly. Full request budgeting (including
6/// Tool schemas and output reservation) is enforced immediately before every
7/// model invocation by the compaction pipeline.
8pub struct ContextManager {
9    max_messages: usize,
10}
11
12impl ContextManager {
13    pub fn new(max_messages: usize) -> Result<Self, KernelFailure> {
14        if max_messages < 2 {
15            return Err(KernelFailure::new(
16                "CONTEXT_LIMIT_INVALID",
17                "context requires at least 2 messages",
18            ));
19        }
20        Ok(Self { max_messages })
21    }
22
23    pub fn assemble(
24        &self,
25        system_prompt: &str,
26        mut context_items: Vec<ContextItem>,
27        history: Vec<ConversationMessage>,
28        user_prompt: &str,
29    ) -> Result<Vec<ModelMessage>, KernelFailure> {
30        let system = self.render_system(system_prompt, &mut context_items)?;
31
32        let mut selected = Vec::new();
33        let history_slots = self.max_messages.saturating_sub(2);
34        for message in history.iter().rev().take(history_slots) {
35            selected.push(message.clone());
36        }
37        selected.reverse();
38
39        let dropped = history.len().saturating_sub(selected.len());
40        if dropped > 0 && history_slots > 0 {
41            let checkpoint = structured_checkpoint(dropped, &history);
42            if selected.len() >= history_slots && !selected.is_empty() {
43                selected.remove(0);
44            }
45            selected.insert(
46                0,
47                ConversationMessage {
48                    role: ModelRole::System,
49                    text: checkpoint,
50                },
51            );
52        }
53
54        let mut messages = Vec::with_capacity(selected.len() + 2);
55        messages.push(ModelMessage::text(ModelRole::System, system));
56        messages.extend(
57            selected
58                .into_iter()
59                .map(|message| ModelMessage::text(message.role, message.text)),
60        );
61        messages.push(ModelMessage::text(ModelRole::User, user_prompt));
62        Ok(messages)
63    }
64
65    pub fn refresh_system(
66        &self,
67        messages: &mut [ModelMessage],
68        system_prompt: &str,
69        mut context_items: Vec<ContextItem>,
70    ) -> Result<(), KernelFailure> {
71        let system = self.render_system(system_prompt, &mut context_items)?;
72        let first = messages.first_mut().ok_or_else(|| {
73            KernelFailure::new("CONTEXT_TRANSCRIPT_INVALID", "model transcript is empty")
74        })?;
75        if first.role != ModelRole::System {
76            return Err(KernelFailure::new(
77                "CONTEXT_TRANSCRIPT_INVALID",
78                "model transcript must begin with system context",
79            ));
80        }
81        *first = ModelMessage::text(ModelRole::System, system);
82        Ok(())
83    }
84
85    fn render_system(
86        &self,
87        system_prompt: &str,
88        context_items: &mut [ContextItem],
89    ) -> Result<String, KernelFailure> {
90        context_items.sort_by(|left, right| {
91            right
92                .required
93                .cmp(&left.required)
94                .then_with(|| right.priority.cmp(&left.priority))
95                .then_with(|| left.source.cmp(&right.source))
96        });
97
98        let mut system = system_prompt.trim().to_string();
99        for item in context_items {
100            if !item.content.trim().is_empty() {
101                system.push_str("\n\n<context source=\"");
102                system.push_str(&item.source);
103                system.push_str("\">\n");
104                system.push_str(item.content.trim());
105                system.push_str("\n</context>");
106            }
107        }
108
109        Ok(system)
110    }
111}
112
113fn structured_checkpoint(dropped: usize, history: &[ConversationMessage]) -> String {
114    let goal = history
115        .iter()
116        .find(|message| message.role == ModelRole::User)
117        .map(|message| truncate(&message.text, 512))
118        .unwrap_or_else(|| "unknown".to_string());
119    format!(
120        "<checkpoint>\ngoal: {goal}\ndone: {dropped} earlier conversation messages compacted\nstate: later messages remain authoritative\nwatch: verify assumptions before destructive actions\nnext: continue from the most recent user request\n</checkpoint>"
121    )
122}
123
124fn truncate(value: &str, max: usize) -> String {
125    if value.len() <= max {
126        return value.to_string();
127    }
128    let mut boundary = max;
129    while !value.is_char_boundary(boundary) {
130        boundary -= 1;
131    }
132    format!("{}…", &value[..boundary])
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn keeps_required_edges_and_compacts_old_history() {
141        let manager = ContextManager::new(4).unwrap();
142        let history = (0..10)
143            .map(|index| ConversationMessage {
144                role: if index % 2 == 0 {
145                    ModelRole::User
146                } else {
147                    ModelRole::Assistant
148                },
149                text: format!("message-{index}"),
150            })
151            .collect();
152        let messages = manager
153            .assemble("system", Vec::new(), history, "current")
154            .unwrap();
155        assert_eq!(messages.first().unwrap().role, ModelRole::System);
156        assert_eq!(messages.last().unwrap().role, ModelRole::User);
157        assert!(messages.len() <= 5);
158        assert!(format!("{:?}", messages).contains("checkpoint"));
159    }
160}
agent_runtime_kernel/context.rs

agent_runtime_kernel/
context.rs
