context_harness/agent_runtime/
context.rs

1//! Deterministic request projection. Compatibility remains the default; bounded
2//! projection is enabled only by an explicit resource budget.
3use super::*;
4use crate::agent_store::WorkingStateInspection;
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
8pub(super) struct SourceRange {
9    pub start: u64,
10    pub end: u64,
11}
12
13#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
14pub(super) struct ProjectionMetadata {
15    pub builder: String,
16    pub strategy: String,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub categories: Option<Vec<String>>,
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub message_count: Option<usize>,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub tool_count: Option<usize>,
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub max_context_bytes: Option<u64>,
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub request_bytes: Option<u64>,
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub source_message_count: Option<u64>,
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub included_message_count: Option<u64>,
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub omitted_message_count: Option<u64>,
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub included_exchange_count: Option<u64>,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub omitted_exchange_count: Option<u64>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub retained_source_ranges: Option<Vec<SourceRange>>,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub working_state_projection_version: Option<i64>,
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub working_state_revision: Option<i64>,
43    #[serde(skip_serializing_if = "Option::is_none")]
44    pub working_state_event_cursor: Option<i64>,
45}
46pub(super) struct ProjectedContext {
47    pub request: ModelRequest,
48    pub metadata: ProjectionMetadata,
49    pub retained_source: Option<ModelRequest>,
50    pub retained_indexes: Option<Vec<u64>>,
51}
52pub(super) trait RunContextBuilder: Send + Sync {
53    fn project(
54        &self,
55        system: &str,
56        input: &str,
57        tools: Vec<ModelTool>,
58        restored: Option<ModelRequest>,
59    ) -> Result<ProjectedContext>;
60    fn describe(&self, request: &ModelRequest) -> ProjectionMetadata;
61}
62pub(super) struct CompatibilityContextBuilder;
63impl RunContextBuilder for CompatibilityContextBuilder {
64    fn project(
65        &self,
66        system: &str,
67        input: &str,
68        tools: Vec<ModelTool>,
69        restored: Option<ModelRequest>,
70    ) -> Result<ProjectedContext> {
71        let request = restored.unwrap_or_else(|| ModelRequest {
72            messages: vec![
73                ModelMessage::System {
74                    content: system.into(),
75                },
76                ModelMessage::User {
77                    content: input.into(),
78                },
79            ],
80            tools,
81            ..Default::default()
82        });
83        request.validate()?;
84        Ok(ProjectedContext {
85            metadata: self.describe(&request),
86            request,
87            retained_source: None,
88            retained_indexes: None,
89        })
90    }
91    fn describe(&self, request: &ModelRequest) -> ProjectionMetadata {
92        ProjectionMetadata {
93            builder: "compatibility_v1".into(),
94            strategy: "complete_transcript".into(),
95            categories: Some(vec![
96                "system_prompt".into(),
97                "objective".into(),
98                "complete_transcript".into(),
99                "tool_declarations".into(),
100            ]),
101            message_count: Some(request.messages.len()),
102            tool_count: Some(request.tools.len()),
103            max_context_bytes: None,
104            request_bytes: None,
105            source_message_count: None,
106            included_message_count: None,
107            omitted_message_count: None,
108            included_exchange_count: None,
109            omitted_exchange_count: None,
110            retained_source_ranges: None,
111            working_state_projection_version: None,
112            working_state_revision: None,
113            working_state_event_cursor: None,
114        }
115    }
116}
117
118#[derive(Serialize)]
119struct WorkingStateCapsule<'a> {
120    r#type: &'static str,
121    projection_version: i64,
122    revision: i64,
123    event_cursor: i64,
124    snapshot: &'a crate::agent_store::WorkingStateSnapshot,
125}
126
127pub(super) fn bounded_project(
128    source: &ModelRequest,
129    max_bytes: u64,
130    working: &WorkingStateInspection,
131    logical_indexes: &[u64],
132    source_message_count: u64,
133    source_exchange_count: u64,
134) -> Result<ProjectedContext> {
135    ensure!(
136        source.messages.len() >= 2 && logical_indexes.len() == source.messages.len(),
137        "invalid bounded projection source"
138    );
139    let snapshot = working
140        .snapshot
141        .as_ref()
142        .context("working state snapshot is unavailable")?;
143    let projection_version = working
144        .projection_version
145        .context("working state projection version is unavailable")?;
146    let revision = working
147        .revision
148        .context("working state revision is unavailable")?;
149    let event_cursor = working
150        .event_cursor
151        .context("working state event cursor is unavailable")?;
152    let capsule = serde_json::to_string(&WorkingStateCapsule {
153        r#type: "working_state",
154        projection_version,
155        revision,
156        event_cursor,
157        snapshot,
158    })?;
159    let mut exchanges = Vec::new();
160    let mut cursor = 2;
161    while cursor < source.messages.len() {
162        let start = cursor;
163        ensure!(
164            matches!(source.messages[cursor], ModelMessage::Assistant { .. }),
165            "bounded source exchange must start with an assistant message"
166        );
167        cursor += 1;
168        while cursor < source.messages.len()
169            && matches!(source.messages[cursor], ModelMessage::Tool { .. })
170        {
171            cursor += 1;
172        }
173        exchanges.push((start, cursor));
174    }
175    let mut first = exchanges.len();
176    if !exchanges.is_empty() {
177        first -= 1;
178        loop {
179            let candidate = build_request(source, &capsule, &exchanges[first..]);
180            if serde_json::to_vec(&candidate)?.len() as u64 > max_bytes {
181                if first + 1 == exchanges.len() {
182                    anyhow::bail!("context_bytes");
183                }
184                first += 1;
185                break;
186            }
187            if first == 0 {
188                break;
189            }
190            first -= 1;
191        }
192    }
193    let request = build_request(source, &capsule, &exchanges[first..]);
194    let request_bytes = serde_json::to_vec(&request)?.len() as u64;
195    ensure!(request_bytes <= max_bytes, "context_bytes");
196    let retained: Vec<usize> = [0, 1]
197        .into_iter()
198        .chain(exchanges[first..].iter().flat_map(|(s, e)| *s..*e))
199        .collect();
200    let included = retained.len() as u64;
201    let mut retained_source = source.clone();
202    retained_source.messages = retained
203        .iter()
204        .map(|position| source.messages[*position].clone())
205        .collect();
206    let retained_indexes = retained
207        .iter()
208        .map(|position| logical_indexes[*position])
209        .collect();
210    Ok(ProjectedContext {
211        request,
212        retained_source: Some(retained_source),
213        retained_indexes: Some(retained_indexes),
214        metadata: ProjectionMetadata {
215            builder: "bounded_v1".into(),
216            strategy: "bounded_suffix_v1".into(),
217            categories: None,
218            message_count: None,
219            tool_count: None,
220            max_context_bytes: Some(max_bytes),
221            request_bytes: Some(request_bytes),
222            source_message_count: Some(source_message_count),
223            included_message_count: Some(included),
224            omitted_message_count: Some(source_message_count.saturating_sub(included)),
225            included_exchange_count: Some((exchanges.len() - first) as u64),
226            omitted_exchange_count: Some(
227                source_exchange_count.saturating_sub((exchanges.len() - first) as u64),
228            ),
229            retained_source_ranges: Some(coalesce(retained.iter().map(|p| logical_indexes[*p]))),
230            working_state_projection_version: Some(projection_version),
231            working_state_revision: Some(revision),
232            working_state_event_cursor: Some(event_cursor),
233        },
234    })
235}
236pub(super) fn exchange_count(request: &ModelRequest) -> u64 {
237    request
238        .messages
239        .iter()
240        .filter(|message| matches!(message, ModelMessage::Assistant { .. }))
241        .count() as u64
242}
243fn build_request(
244    source: &ModelRequest,
245    capsule: &str,
246    exchanges: &[(usize, usize)],
247) -> ModelRequest {
248    let mut messages = vec![
249        source.messages[0].clone(),
250        ModelMessage::System {
251            content: capsule.into(),
252        },
253        source.messages[1].clone(),
254    ];
255    for (start, end) in exchanges {
256        messages.extend_from_slice(&source.messages[*start..*end]);
257    }
258    ModelRequest {
259        messages,
260        tools: source.tools.clone(),
261        ..source.clone()
262    }
263}
264fn coalesce(indexes: impl IntoIterator<Item = u64>) -> Vec<SourceRange> {
265    let mut ranges: Vec<SourceRange> = Vec::new();
266    for index in indexes {
267        if let Some(last) = ranges.last_mut().filter(|r| r.end + 1 == index) {
268            last.end = index;
269        } else {
270            ranges.push(SourceRange {
271                start: index,
272                end: index,
273            });
274        }
275    }
276    ranges
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::agent_store::{
283        RunLifecycle, RunUsage, WorkingStateArtifacts, WorkingStateEventReference,
284        WorkingStateObjective, WorkingStateRun, WorkingStateSnapshot, WorkingStateStatus,
285    };
286    fn working() -> WorkingStateInspection {
287        WorkingStateInspection {
288            status: WorkingStateStatus::Current,
289            projection_version: Some(1),
290            revision: Some(7),
291            event_cursor: Some(9),
292            updated_at: Some(10),
293            snapshot: Some(WorkingStateSnapshot {
294                objective: WorkingStateObjective {
295                    kind: "agent_run".into(),
296                    run_id: "run".into(),
297                },
298                run: WorkingStateRun {
299                    lifecycle: RunLifecycle::Active,
300                    outcome: None,
301                    reason_code: None,
302                },
303                usage: RunUsage::default(),
304                latest_event: WorkingStateEventReference {
305                    sequence: 9,
306                    event_type: "tool.completed".into(),
307                },
308                artifacts: WorkingStateArtifacts {
309                    items: Vec::new(),
310                    total: 0,
311                    omitted: 0,
312                },
313            }),
314        }
315    }
316    #[test]
317    fn compatibility_builder_preserves_fresh_request_shape() {
318        let p = CompatibilityContextBuilder
319            .project("system", "objective", Vec::new(), None)
320            .unwrap();
321        assert_eq!(p.request.messages.len(), 2);
322        assert_eq!(p.metadata.builder, "compatibility_v1");
323    }
324    #[test]
325    fn compatibility_builder_does_not_rewrite_restored_requests() {
326        let r = ModelRequest {
327            messages: vec![ModelMessage::User {
328                content: "restored".into(),
329            }],
330            ..Default::default()
331        };
332        let p = CompatibilityContextBuilder
333            .project("ignored", "ignored", Vec::new(), Some(r.clone()))
334            .unwrap();
335        assert_eq!(
336            serde_json::to_value(p.request).unwrap(),
337            serde_json::to_value(r).unwrap()
338        );
339        assert_eq!(p.metadata.message_count, Some(1));
340    }
341    #[test]
342    fn bounded_builder_keeps_prefix_and_newest_whole_exchange() {
343        let source = ModelRequest {
344            messages: vec![
345                ModelMessage::System {
346                    content: "system".into(),
347                },
348                ModelMessage::User {
349                    content: "objective".into(),
350                },
351                ModelMessage::Assistant {
352                    text: "x".repeat(2_000),
353                    tool_calls: Vec::new(),
354                    continuation: None,
355                },
356                ModelMessage::Assistant {
357                    text: "newest".into(),
358                    tool_calls: Vec::new(),
359                    continuation: None,
360                },
361            ],
362            ..Default::default()
363        };
364        let projected = bounded_project(&source, 1_000, &working(), &[0, 1, 2, 3], 4, 2).unwrap();
365        assert_eq!(projected.metadata.omitted_exchange_count, Some(1));
366        assert_eq!(
367            projected.metadata.retained_source_ranges,
368            Some(vec![
369                SourceRange { start: 0, end: 1 },
370                SourceRange { start: 3, end: 3 }
371            ])
372        );
373        assert_eq!(
374            serde_json::to_vec(&projected.request).unwrap().len() as u64,
375            projected.metadata.request_bytes.unwrap()
376        );
377        assert!(
378            matches!(&projected.request.messages[1],ModelMessage::System{content} if content.contains("\"projection_version\":1") && !content.contains("updated_at"))
379        );
380    }
381}
context_harness/agent_runtime/context.rs

context_harness/agent_runtime/
context.rs
