1use serde::{Deserialize, Serialize};
2use serde_json::{Value, json};
3use sha2::{Digest, Sha256};
4use thiserror::Error;
5
6use crate::event::{EventError, EventStore};
7use crate::inference::tool_schemas;
8use crate::session::Session;
9use crate::workspace::{Workspace, WorkspaceError};
10
11const SYSTEM_POLICY: &str = "You are the implementation worker, not acceptance authority. Inspect and modify only the designated workspace using declared tools. Act promptly: do not repeat a successful read or command while its result remains in recent interaction. Native tool calls may batch independent reads; dependent actions must wait for prior results. Prefer the smallest valid patch, run the relevant tests after editing, then call candidate_ready. If native calls are unavailable, emit one strict JSON object: {\"action\":\"tool\",\"tool\":\"name\",\"arguments\":{...}} or {\"action\":\"candidate_ready\",\"summary\":\"...\"}. Candidate ready means independent verification should begin; it never means accepted. Do not alter FalseGreen contracts or history.";
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14pub struct ContextBudget {
15 pub max_context_tokens: usize,
16 pub output_reserve_tokens: usize,
17 pub recent_event_limit: usize,
18}
19
20impl Default for ContextBudget {
21 fn default() -> Self {
22 Self {
23 max_context_tokens: 16_384,
24 output_reserve_tokens: 2_048,
25 recent_event_limit: 24,
26 }
27 }
28}
29
30impl ContextBudget {
31 #[must_use]
32 pub fn usable_input_tokens(self) -> usize {
33 self.max_context_tokens
34 .saturating_sub(self.output_reserve_tokens)
35 }
36}
37
38#[derive(Debug, Error)]
39pub enum ContextError {
40 #[error(transparent)]
41 Event(#[from] EventError),
42 #[error(transparent)]
43 Workspace(#[from] WorkspaceError),
44 #[error("fixed context exceeds the configured token budget")]
45 FixedContextTooLarge,
46 #[error("context serialization failed: {0}")]
47 Json(#[from] serde_json::Error),
48}
49
50#[derive(Debug, Clone)]
51pub struct ContextBuilder {
52 budget: ContextBudget,
53 tool_schemas: Value,
54}
55
56impl ContextBuilder {
57 #[must_use]
58 pub fn new(budget: ContextBudget) -> Self {
59 Self {
60 budget,
61 tool_schemas: tool_schemas(),
62 }
63 }
64
65 #[must_use]
66 pub fn with_tool_schemas(mut self, tool_schemas: Value) -> Self {
67 self.tool_schemas = tool_schemas;
68 self
69 }
70
71 pub fn build(
72 &self,
73 session: &Session,
74 store: &EventStore,
75 workspace: &Workspace,
76 ) -> Result<String, ContextError> {
77 let workspace_state = workspace.state()?;
78 let tool_surface = compact_tool_surface(&self.tool_schemas)?;
79 let fixed = json!({
80 "system_policy": SYSTEM_POLICY,
81 "frozen_task": {
82 "original_goal": session.goal,
83 "current_agent_state": session.state,
84 "mutability": "immutable for this session"
85 },
86 "workspace": workspace_state,
87 "available_tools": tool_surface,
88 "context_policy": {
89 "approximate_token_budget": self.budget.usable_input_tokens(),
90 "selection": "bounded recent events plus objective workspace state"
91 }
92 });
93 let mut fixed_text = serde_json::to_string_pretty(&fixed)?;
94 let character_budget = self.budget.usable_input_tokens().saturating_mul(4);
95 const HISTORY_PREFIX: &str = "\n\nRECENT_EVENT_HISTORY:\n";
96 if fixed_text.len() + HISTORY_PREFIX.len() + 2 > character_budget {
97 return Err(ContextError::FixedContextTooLarge);
98 }
99 let events = store.events(&session.id)?;
100 let mut selected: Vec<Value> = Vec::new();
101 for event in events
102 .iter()
103 .rev()
104 .filter(|event| {
105 matches!(
106 event.kind,
107 crate::event::EventKind::FalsegreenResult
108 | crate::event::EventKind::RepairStarted
109 | crate::event::EventKind::CandidateReady
110 )
111 })
112 .take(self.budget.recent_event_limit)
113 {
114 let value = json!({
115 "sequence": event.sequence,
116 "kind": event.kind,
117 "payload": event.payload
118 });
119 selected.push(value);
120 let mut chronological = selected.clone();
121 chronological.reverse();
122 let history = serde_json::to_string_pretty(&chronological)?;
123 if fixed_text.len() + HISTORY_PREFIX.len() + history.len() > character_budget {
124 selected.pop();
125 break;
126 }
127 }
128 selected.reverse();
129 fixed_text.push_str(HISTORY_PREFIX);
130 fixed_text.push_str(&serde_json::to_string_pretty(&selected)?);
131 Ok(fixed_text)
132 }
133
134 #[must_use]
135 pub const fn max_output_tokens(&self) -> u32 {
136 if self.budget.output_reserve_tokens > u32::MAX as usize {
137 u32::MAX
138 } else {
139 self.budget.output_reserve_tokens as u32
140 }
141 }
142
143 #[must_use]
144 pub fn remaining_input_bytes(&self, fixed_context: &str) -> usize {
145 self.budget
146 .usable_input_tokens()
147 .saturating_mul(4)
148 .saturating_sub(fixed_context.len())
149 }
150}
151
152fn compact_tool_surface(schemas: &Value) -> Result<Value, serde_json::Error> {
153 let serialized = serde_json::to_vec(schemas)?;
154 Ok(json!({
155 "count": schemas.as_array().map_or(0, Vec::len),
156 "schema_set_sha256": format!("{:x}", Sha256::digest(serialized)),
157 "schema_delivery": "Complete native and namespaced MCP schemas are supplied together in the provider tool field; this digest binds that exact surface without duplicating it in message text."
158 }))
159}
160
161#[cfg(test)]
162mod tests {
163 use serde_json::{Value, json};
164
165 use crate::event::{EventKind, EventStore};
166 use crate::session::Session;
167 use crate::workspace::Workspace;
168 use crate::workspace::tests::git_fixture;
169
170 use super::{ContextBudget, ContextBuilder};
171
172 #[test]
173 fn selects_bounded_recent_history() {
174 let directory = git_fixture();
175 let workspace = Workspace::open(directory.path()).expect("workspace");
176 let mut store = EventStore::open_memory().expect("store");
177 let session = Session::create(&mut store, "goal").expect("session");
178 for index in 0..10 {
179 store
180 .append(
181 &session.id,
182 EventKind::FalsegreenResult,
183 &json!({"index": index}),
184 )
185 .expect("append");
186 }
187 let context = ContextBuilder::new(ContextBudget {
188 max_context_tokens: 4_096,
189 output_reserve_tokens: 512,
190 recent_event_limit: 2,
191 })
192 .build(&session, &store, &workspace)
193 .expect("context");
194 assert!(context.contains("\"index\": 9"));
195 assert!(!context.contains("\"index\": 0"));
196 assert!(context.contains("implementation worker, not acceptance authority"));
197 assert!(context.contains("\"frozen_task\""));
198 assert!(context.contains("\"original_goal\": \"goal\""));
199 }
200
201 #[test]
202 fn frozen_task_survives_large_history_eviction() {
203 let directory = git_fixture();
204 let workspace = Workspace::open(directory.path()).expect("workspace");
205 let mut store = EventStore::open_memory().expect("store");
206 let session = Session::create(
207 &mut store,
208 "FROZEN CONTRACT digest=abc123; only hello.txt may change",
209 )
210 .expect("session");
211 for index in 0..100 {
212 store
213 .append(
214 &session.id,
215 EventKind::FalsegreenResult,
216 &json!({"index": index, "evidence": "x".repeat(512)}),
217 )
218 .expect("append");
219 }
220 let context = ContextBuilder::new(ContextBudget {
221 max_context_tokens: 2_048,
222 output_reserve_tokens: 512,
223 recent_event_limit: 100,
224 })
225 .build(&session, &store, &workspace)
226 .expect("context");
227 assert!(context.contains("FROZEN CONTRACT digest=abc123"));
228 assert!(context.len() <= (2_048 - 512) * 4 + 64);
229 }
230
231 #[test]
232 fn large_exact_tool_surface_is_digest_indexed_without_fixed_context_overflow() {
233 let directory = git_fixture();
234 let workspace = Workspace::open(directory.path()).expect("workspace");
235 let mut store = EventStore::open_memory().expect("store");
236 let session = Session::create(&mut store, "fix the focused fixture").expect("session");
237 let schemas = (0..40)
238 .map(|index| {
239 json!({
240 "type": "function",
241 "function": {
242 "name": format!("mcp__fixture__tool_{index}"),
243 "description": "x".repeat(2_000),
244 "parameters": {
245 "type": "object",
246 "properties": {"payload": {"type": "string", "description": "y".repeat(2_000)}}
247 }
248 }
249 })
250 })
251 .collect::<Vec<_>>();
252 let context = ContextBuilder::new(ContextBudget {
253 max_context_tokens: 8_192,
254 output_reserve_tokens: 2_048,
255 recent_event_limit: 24,
256 })
257 .with_tool_schemas(Value::Array(schemas))
258 .build(&session, &store, &workspace)
259 .expect("digest-indexed surface fits");
260
261 assert!(context.contains("\"count\": 40"));
262 assert!(context.contains("schema_set_sha256"));
263 assert!(!context.contains(&"x".repeat(2_000)));
264 }
265}