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        request_messages: &[ConversationMessage],
29    ) -> Result<Vec<ModelMessage>, KernelFailure> {
30        let system = self.render_system(system_prompt, &mut context_items)?;
31
32        if request_messages.is_empty() || request_messages.len() >= self.max_messages {
33            return Err(KernelFailure::new(
34                "REQUEST_TRANSCRIPT_INVALID",
35                "request messages must fit within the configured context message limit",
36            ));
37        }
38
39        let mut selected = Vec::new();
40        let history_slots = self.max_messages.saturating_sub(1 + request_messages.len());
41        for message in history.iter().rev().take(history_slots) {
42            selected.push(message.clone());
43        }
44        selected.reverse();
45
46        let dropped = history.len().saturating_sub(selected.len());
47        if dropped > 0 && history_slots > 0 {
48            let checkpoint = structured_checkpoint(dropped, &history);
49            if selected.len() >= history_slots && !selected.is_empty() {
50                selected.remove(0);
51            }
52            selected.insert(
53                0,
54                ConversationMessage {
55                    role: ModelRole::System,
56                    text: checkpoint,
57                },
58            );
59        }
60
61        let mut messages = Vec::with_capacity(selected.len() + request_messages.len() + 1);
62        messages.push(ModelMessage::text(ModelRole::System, system));
63        messages.extend(
64            selected
65                .into_iter()
66                .map(|message| ModelMessage::text(message.role, message.text)),
67        );
68        messages.extend(
69            request_messages
70                .iter()
71                .map(|message| ModelMessage::text(message.role, message.text.clone())),
72        );
73        Ok(messages)
74    }
75
76    pub fn refresh_system(
77        &self,
78        messages: &mut [ModelMessage],
79        system_prompt: &str,
80        mut context_items: Vec<ContextItem>,
81    ) -> Result<(), KernelFailure> {
82        let system = self.render_system(system_prompt, &mut context_items)?;
83        let first = messages.first_mut().ok_or_else(|| {
84            KernelFailure::new("CONTEXT_TRANSCRIPT_INVALID", "model transcript is empty")
85        })?;
86        if first.role != ModelRole::System {
87            return Err(KernelFailure::new(
88                "CONTEXT_TRANSCRIPT_INVALID",
89                "model transcript must begin with system context",
90            ));
91        }
92        *first = ModelMessage::text(ModelRole::System, system);
93        Ok(())
94    }
95
96    fn render_system(
97        &self,
98        system_prompt: &str,
99        context_items: &mut [ContextItem],
100    ) -> Result<String, KernelFailure> {
101        context_items.sort_by(|left, right| {
102            right
103                .required
104                .cmp(&left.required)
105                .then_with(|| right.priority.cmp(&left.priority))
106                .then_with(|| left.source.cmp(&right.source))
107        });
108
109        let mut system = system_prompt.trim().to_string();
110        for item in context_items {
111            if !item.content.trim().is_empty() {
112                system.push_str("\n\n<context source=\"");
113                system.push_str(&item.source);
114                system.push_str("\">\n");
115                system.push_str(item.content.trim());
116                system.push_str("\n</context>");
117            }
118        }
119
120        Ok(system)
121    }
122}
123
124fn structured_checkpoint(dropped: usize, history: &[ConversationMessage]) -> String {
125    let goal = history
126        .iter()
127        .find(|message| message.role == ModelRole::User)
128        .map(|message| truncate(&message.text, 512))
129        .unwrap_or_else(|| "unknown".to_string());
130    format!(
131        "<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>"
132    )
133}
134
135fn truncate(value: &str, max: usize) -> String {
136    if value.len() <= max {
137        return value.to_string();
138    }
139    let mut boundary = max;
140    while !value.is_char_boundary(boundary) {
141        boundary -= 1;
142    }
143    format!("{}…", &value[..boundary])
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    #[test]
151    fn keeps_required_edges_and_compacts_old_history() {
152        let manager = ContextManager::new(4).unwrap();
153        let history = (0..10)
154            .map(|index| ConversationMessage {
155                role: if index % 2 == 0 {
156                    ModelRole::User
157                } else {
158                    ModelRole::Assistant
159                },
160                text: format!("message-{index}"),
161            })
162            .collect();
163        let messages = manager
164            .assemble(
165                "system",
166                Vec::new(),
167                history,
168                &[ConversationMessage {
169                    role: ModelRole::User,
170                    text: "current".into(),
171                }],
172            )
173            .unwrap();
174        assert_eq!(messages.first().unwrap().role, ModelRole::System);
175        assert_eq!(messages.last().unwrap().role, ModelRole::User);
176        assert!(messages.len() <= 5);
177        assert!(format!("{:?}", messages).contains("checkpoint"));
178    }
179}
agent_runtime_kernel/context.rs

agent_runtime_kernel/
context.rs
