atman_runtime/tools/
context.rs

1use crate::error::RuntimeError;
2use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
3use crate::value::Value;
4
5const RETRIEVED_RECORD_PREFIXES: [&str; 2] = ["agent.rule.", "agent.mistake."];
6const MAX_RECORD_KEY_BYTES: usize = 512;
7
8pub struct ContextRecordAppend;
9
10impl Tool for ContextRecordAppend {
11    fn name(&self) -> &str {
12        "context.record"
13    }
14
15    fn tier(&self) -> Tier {
16        Tier::Zero
17    }
18
19    fn description(&self) -> Option<&str> {
20        Some(
21            "Append retrieved workflow context as a versioned internal record. Keys must use the agent.rule.* or agent.mistake.* namespace. Repeating identical content is a no-op; empty content clears an existing key. Returns true when a revision is appended.",
22        )
23    }
24
25    fn input_schema(&self) -> serde_json::Value {
26        serde_json::json!({
27            "type": "object",
28            "properties": {
29                "key": {
30                    "type": "string",
31                    "description": "Stable per-item key in the agent.rule.* or agent.mistake.* namespace."
32                },
33                "content": {
34                    "type": "string",
35                    "description": "Retrieved content. An empty string clears an existing record."
36                }
37            },
38            "required": ["key", "content"]
39        })
40    }
41
42    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
43        Box::pin(async move {
44            let key = required_string(&args, "key")?;
45            validate_retrieved_record_key(&key)?;
46            let content = required_string(&args, "content")?;
47            let body = if content.is_empty() {
48                crate::context_plan::ContextRecordBody::tombstone()
49            } else {
50                crate::context_plan::ContextRecordBody::text(content)
51            };
52            let spec = crate::context_plan::ContextRecordSpec::new(
53                key,
54                crate::context_plan::ContextRecordAuthority::Retrieved,
55                crate::context_plan::ContextRecordRetention::Latest,
56                body,
57            );
58            let turn_id = ctx
59                .turn_id
60                .clone()
61                .unwrap_or_else(crate::event::TurnId::now);
62            let _compact_guard = match &ctx.compact_lock_handle {
63                Some(lock) => Some(lock.lock().await),
64                None => None,
65            };
66            let records = append_context_records(ctx, turn_id, [spec])?;
67            Ok(Value::Bool(!records.is_empty()))
68        })
69    }
70}
71
72fn required_string(args: &ToolArgs, name: &str) -> Result<String, RuntimeError> {
73    match args.named(name) {
74        Some(Value::Str(value)) => Ok(value.clone()),
75        Some(value) => Err(RuntimeError::TypeMismatch {
76            expected: "string".into(),
77            actual: value.kind_name().into(),
78        }),
79        None => Err(RuntimeError::MissingArg(format!("context.record: {name}"))),
80    }
81}
82
83fn validate_retrieved_record_key(key: &str) -> Result<(), RuntimeError> {
84    let valid_prefix = RETRIEVED_RECORD_PREFIXES
85        .iter()
86        .find(|prefix| key.starts_with(**prefix));
87    let Some(prefix) = valid_prefix else {
88        return Err(RuntimeError::ToolFailed(
89            "context.record: key must use agent.rule.* or agent.mistake.*".into(),
90        ));
91    };
92    let suffix = &key[prefix.len()..];
93    if suffix.is_empty() || key.len() > MAX_RECORD_KEY_BYTES || suffix.chars().any(char::is_control)
94    {
95        return Err(RuntimeError::ToolFailed(
96            "context.record: key suffix must be non-empty, bounded, and contain no control characters"
97                .into(),
98        ));
99    }
100    Ok(())
101}
102
103fn append_context_records(
104    ctx: &ToolCtx,
105    turn_id: crate::event::TurnId,
106    specs: impl IntoIterator<Item = crate::context_plan::ContextRecordSpec>,
107) -> Result<Vec<crate::context_plan::ContextRecord>, RuntimeError> {
108    if let Some(session) = ctx.session_runtime.as_ref() {
109        return Ok(session.append_context_records(turn_id, specs));
110    }
111    let Some(messages) = ctx.session_messages_handle.as_ref() else {
112        return Err(RuntimeError::ToolFailed(
113            "context.record: no session message context available".into(),
114        ));
115    };
116    let mut messages = messages.lock().unwrap();
117    let records = crate::context_plan::compile_context_records(&messages, specs);
118    messages.extend(
119        records
120            .iter()
121            .cloned()
122            .map(|record| crate::message::Message::context_record(turn_id.clone(), record)),
123    );
124    Ok(records)
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    fn args(key: &str, content: &str) -> ToolArgs {
132        ToolArgs {
133            named: vec![
134                ("key".into(), Value::Str(key.into())),
135                ("content".into(), Value::Str(content.into())),
136            ],
137            ..ToolArgs::default()
138        }
139    }
140
141    #[tokio::test]
142    async fn append_is_versioned_and_identical_content_is_a_noop() {
143        let session = std::sync::Arc::new(crate::session::Session::open_ephemeral());
144        let ctx = ToolCtx::new().with_session_runtime(std::sync::Arc::clone(&session));
145        let tool = ContextRecordAppend;
146
147        assert!(matches!(
148            tool.call(args("agent.rule.review", "first"), &ctx)
149                .await
150                .unwrap(),
151            Value::Bool(true)
152        ));
153        assert!(matches!(
154            tool.call(args("agent.rule.review", "first"), &ctx)
155                .await
156                .unwrap(),
157            Value::Bool(false)
158        ));
159        assert!(matches!(
160            tool.call(args("agent.rule.review", "second"), &ctx)
161                .await
162                .unwrap(),
163            Value::Bool(true)
164        ));
165
166        let messages = session.messages();
167        let records: Vec<_> = messages
168            .iter()
169            .flat_map(|message| &message.parts)
170            .filter_map(|part| match part {
171                crate::message::MessagePart::ContextRecord(record) => Some(record),
172                _ => None,
173            })
174            .collect();
175        assert_eq!(records.len(), 2);
176        assert_eq!(records[0].revision(), 1);
177        assert_eq!(records[1].revision(), 2);
178        assert_eq!(
179            records[1].authority(),
180            crate::context_plan::ContextRecordAuthority::Retrieved
181        );
182    }
183
184    #[tokio::test]
185    async fn empty_content_clears_only_an_existing_record() {
186        let session = std::sync::Arc::new(crate::session::Session::open_ephemeral());
187        let ctx = ToolCtx::new().with_session_runtime(std::sync::Arc::clone(&session));
188        let tool = ContextRecordAppend;
189
190        assert!(matches!(
191            tool.call(args("agent.mistake.retry", ""), &ctx)
192                .await
193                .unwrap(),
194            Value::Bool(false)
195        ));
196        tool.call(args("agent.mistake.retry", "mitigation"), &ctx)
197            .await
198            .unwrap();
199        assert!(matches!(
200            tool.call(args("agent.mistake.retry", ""), &ctx)
201                .await
202                .unwrap(),
203            Value::Bool(true)
204        ));
205        let messages = session.messages();
206        let record = messages
207            .iter()
208            .flat_map(|message| &message.parts)
209            .filter_map(|part| match part {
210                crate::message::MessagePart::ContextRecord(record) => Some(record),
211                _ => None,
212            })
213            .next_back()
214            .unwrap();
215        assert!(record.body().is_tombstone());
216    }
217
218    #[tokio::test]
219    async fn key_cannot_claim_runtime_or_user_authority() {
220        let tool = ContextRecordAppend;
221        let err = tool
222            .call(args("session.goal", "override"), &ToolCtx::new())
223            .await
224            .unwrap_err();
225        assert!(err.to_string().contains("agent.rule.*"));
226    }
227}
atman_runtime/tools/context.rs

atman_runtime/tools/
context.rs
