a3s_code_core/
rl_trajectory.rs

1//! RL trajectory recording primitives.
2//!
3//! The default runtime trace is intentionally lightweight and diagnostic. This
4//! module records an opt-in training-oriented JSONL stream that can reconstruct LLM
5//! turns, tool calls, tool observations, token usage, and termination reasons.
6//! It is opt-in: normal sessions pay only a cheap disabled-recorder branch.
7
8use crate::llm::{
9    ContentBlock, LlmResponse, Message, TokenLogProb, TokenUsage, ToolCall, ToolDefinition,
10    TopTokenLogProb,
11};
12use anyhow::{Context, Result};
13use serde::{Deserialize, Serialize};
14use serde_json::{json, Value};
15use std::fs::{File, OpenOptions};
16use std::io::Write;
17use std::path::{Path, PathBuf};
18use std::sync::{
19    atomic::{AtomicU64, Ordering},
20    Arc, Mutex,
21};
22
23pub const RL_TRAJECTORY_SCHEMA: &str = "a3s.rl_trajectory.v1";
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
26#[serde(rename_all = "snake_case")]
27pub enum RlTrajectoryMode {
28    #[default]
29    Off,
30    On,
31}
32
33impl RlTrajectoryMode {
34    pub fn parse(value: &str) -> Option<Self> {
35        match value.trim().to_ascii_lowercase().as_str() {
36            "" | "off" | "0" | "false" | "none" => Some(Self::Off),
37            "on" | "1" | "true" | "yes" | "enabled" | "rl" | "train" | "training"
38            | "trajectory" | "trace" | "debug" | "full" | "compact" => Some(Self::On),
39            _ => None,
40        }
41    }
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct RlTrajectoryConfig {
46    pub mode: RlTrajectoryMode,
47    pub path: PathBuf,
48    pub max_text_bytes: usize,
49    pub include_messages: bool,
50}
51
52impl RlTrajectoryConfig {
53    pub fn new(path: impl Into<PathBuf>) -> Self {
54        Self {
55            mode: RlTrajectoryMode::On,
56            path: path.into(),
57            max_text_bytes: default_max_text_bytes(RlTrajectoryMode::On),
58            include_messages: true,
59        }
60    }
61
62    pub fn with_mode(mut self, mode: RlTrajectoryMode) -> Self {
63        self.mode = mode;
64        self.max_text_bytes = default_max_text_bytes(mode);
65        self.include_messages = mode == RlTrajectoryMode::On;
66        self
67    }
68
69    pub fn with_max_text_bytes(mut self, max_text_bytes: usize) -> Self {
70        self.max_text_bytes = max_text_bytes;
71        self
72    }
73
74    pub fn with_include_messages(mut self, include_messages: bool) -> Self {
75        self.include_messages = include_messages;
76        self
77    }
78
79    pub fn from_env() -> Result<Option<Self>> {
80        let mode_env = env_first(&["A3S_CODE_TRAJECTORY_MODE", "A3S_CODE_RL_TRAJECTORY_MODE"]);
81        let path_env = env_first(&["A3S_CODE_TRAJECTORY_PATH", "A3S_CODE_RL_TRAJECTORY_PATH"]);
82        if mode_env.is_none() && path_env.is_none() {
83            return Ok(None);
84        }
85
86        let mode = match mode_env.as_deref() {
87            Some(raw) => RlTrajectoryMode::parse(raw)
88                .with_context(|| format!("invalid A3S_CODE_RL_TRAJECTORY_MODE: {raw}"))?,
89            None => RlTrajectoryMode::On,
90        };
91        if mode == RlTrajectoryMode::Off {
92            return Ok(None);
93        }
94
95        let path = path_env
96            .filter(|s| !s.trim().is_empty())
97            .with_context(|| "A3S_CODE_TRAJECTORY_PATH is required when trajectory mode is on")?;
98
99        let max_text_bytes = env_first(&[
100            "A3S_CODE_TRAJECTORY_MAX_TEXT_BYTES",
101            "A3S_CODE_RL_TRAJECTORY_MAX_TEXT_BYTES",
102        ])
103        .and_then(|value| value.parse::<usize>().ok())
104        .unwrap_or_else(|| default_max_text_bytes(mode));
105
106        let include_messages = env_first(&[
107            "A3S_CODE_TRAJECTORY_INCLUDE_MESSAGES",
108            "A3S_CODE_RL_TRAJECTORY_INCLUDE_MESSAGES",
109        ])
110        .and_then(|value| parse_bool(&value))
111        .unwrap_or(true);
112
113        Ok(Some(Self {
114            mode,
115            path: PathBuf::from(path),
116            max_text_bytes,
117            include_messages,
118        }))
119    }
120}
121
122#[derive(Debug, Clone, Default, Serialize, Deserialize)]
123pub struct RlTrajectoryContext {
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub run_id: Option<String>,
126    #[serde(skip_serializing_if = "Option::is_none")]
127    pub task_id: Option<String>,
128    #[serde(skip_serializing_if = "Option::is_none")]
129    pub group_id: Option<String>,
130    #[serde(skip_serializing_if = "Option::is_none")]
131    pub replica_id: Option<String>,
132    #[serde(skip_serializing_if = "Option::is_none")]
133    pub sample_id: Option<String>,
134}
135
136impl RlTrajectoryContext {
137    fn from_env() -> Self {
138        Self {
139            run_id: env_first(&["A3S_CODE_RL_RUN_ID", "A3S_CODE_RUN_ID", "A3S_RUN_ID"]),
140            task_id: env_first(&["A3S_CODE_RL_TASK_ID", "A3S_CODE_TASK_ID", "TASK_ID"]),
141            group_id: env_first(&["A3S_CODE_RL_GROUP_ID", "A3S_CODE_GROUP_ID"]),
142            replica_id: env_first(&["A3S_CODE_RL_REPLICA_ID", "A3S_CODE_REPLICA_ID"]),
143            sample_id: env_first(&["A3S_CODE_RL_SAMPLE_ID", "A3S_CODE_SAMPLE_ID"]),
144        }
145    }
146}
147
148#[derive(Debug, Clone, Serialize, Deserialize)]
149pub struct CapturedText {
150    pub byte_len: usize,
151    pub sha256: String,
152    pub truncated: bool,
153    #[serde(skip_serializing_if = "Option::is_none")]
154    pub text: Option<String>,
155    #[serde(skip_serializing_if = "Option::is_none")]
156    pub preview: Option<String>,
157}
158
159#[derive(Clone)]
160pub struct RlTrajectoryRecorder {
161    inner: Option<Arc<RlTrajectoryRecorderInner>>,
162}
163
164pub struct ExecutionStartRecord<'a> {
165    pub session_id: &'a str,
166    pub workspace: &'a Path,
167    pub prompt: &'a str,
168    pub history: &'a [Message],
169    pub system_prompt: Option<&'a str>,
170    pub max_tool_rounds: usize,
171    pub planning_mode: &'a str,
172}
173
174struct RlTrajectoryRecorderInner {
175    config: RlTrajectoryConfig,
176    context: RlTrajectoryContext,
177    sequence: AtomicU64,
178    file: Mutex<File>,
179}
180
181impl std::fmt::Debug for RlTrajectoryRecorder {
182    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
183        f.debug_struct("RlTrajectoryRecorder")
184            .field("enabled", &self.inner.is_some())
185            .finish()
186    }
187}
188
189impl Default for RlTrajectoryRecorder {
190    fn default() -> Self {
191        Self::disabled()
192    }
193}
194
195impl RlTrajectoryRecorder {
196    pub fn disabled() -> Self {
197        Self { inner: None }
198    }
199
200    pub fn from_config(config: Option<RlTrajectoryConfig>) -> Result<Self> {
201        let Some(config) = config else {
202            return Ok(Self::disabled());
203        };
204        if config.mode == RlTrajectoryMode::Off {
205            return Ok(Self::disabled());
206        }
207
208        if let Some(parent) = config.path.parent().filter(|p| !p.as_os_str().is_empty()) {
209            std::fs::create_dir_all(parent).with_context(|| {
210                format!(
211                    "failed to create RL trajectory directory {}",
212                    parent.display()
213                )
214            })?;
215        }
216        let file = OpenOptions::new()
217            .create(true)
218            .append(true)
219            .open(&config.path)
220            .with_context(|| {
221                format!(
222                    "failed to open RL trajectory JSONL {}",
223                    config.path.display()
224                )
225            })?;
226
227        Ok(Self {
228            inner: Some(Arc::new(RlTrajectoryRecorderInner {
229                config,
230                context: RlTrajectoryContext::from_env(),
231                sequence: AtomicU64::new(0),
232                file: Mutex::new(file),
233            })),
234        })
235    }
236
237    pub fn is_enabled(&self) -> bool {
238        self.inner.is_some()
239    }
240
241    pub fn record_execution_start(&self, record: ExecutionStartRecord<'_>) {
242        let Some(inner) = &self.inner else {
243            return;
244        };
245        let payload = json!({
246            "workspace": record.workspace.display().to_string(),
247            "prompt": inner.capture_text(record.prompt),
248            "history_message_count": record.history.len(),
249            "history": inner.capture_messages(record.history),
250            "system_prompt": record.system_prompt.map(|s| inner.capture_text(s)),
251            "max_tool_rounds": record.max_tool_rounds,
252            "planning_mode": record.planning_mode,
253        });
254        inner.record("execution_start", record.session_id, payload);
255    }
256
257    pub fn record_llm_request(
258        &self,
259        session_id: &str,
260        turn: usize,
261        messages: &[Message],
262        system: Option<&str>,
263        available_tools: &[ToolDefinition],
264        estimated_prompt_tokens: usize,
265    ) {
266        let Some(inner) = &self.inner else {
267            return;
268        };
269        let available_tool_names = available_tools
270            .iter()
271            .map(|tool| tool.name.as_str())
272            .collect::<Vec<_>>();
273        let payload = json!({
274            "turn": turn,
275            "messages_count": messages.len(),
276            "messages": inner.capture_messages(messages),
277            "system_prompt": system.map(|s| inner.capture_text(s)),
278            "available_tools": available_tool_names,
279            "tool_definitions": available_tools.iter().map(tool_definition_value).collect::<Vec<_>>(),
280            "estimated_prompt_tokens": estimated_prompt_tokens,
281        });
282        inner.record("llm_request", session_id, payload);
283    }
284
285    pub fn record_llm_response(
286        &self,
287        session_id: &str,
288        turn: usize,
289        response: &LlmResponse,
290        duration_ms: u64,
291    ) {
292        let Some(inner) = &self.inner else {
293            return;
294        };
295        let payload = json!({
296            "turn": turn,
297            "message": inner.capture_message(&response.message),
298            "response_text": inner.capture_text(&response.text()),
299            "reasoning_content": response.message.reasoning_content.as_ref().map(|s| inner.capture_text(s)),
300            "tool_calls": response.tool_calls().iter().map(tool_call_value).collect::<Vec<_>>(),
301            "token_logprobs": response.token_logprobs.iter().map(token_logprob_value).collect::<Vec<_>>(),
302            "usage": token_usage_value(&response.usage),
303            "stop_reason": response.stop_reason.clone(),
304            "meta": response.meta.clone(),
305            "duration_ms": duration_ms,
306        });
307        inner.record("llm_response", session_id, payload);
308    }
309
310    pub fn record_tool_call(&self, session_id: &str, turn: usize, tool_call: &ToolCall) {
311        let Some(inner) = &self.inner else {
312            return;
313        };
314        let payload = json!({
315            "turn": turn,
316            "tool_call_id": tool_call.id,
317            "tool": tool_call.name,
318            "args": tool_call.args,
319        });
320        inner.record("tool_call", session_id, payload);
321    }
322
323    #[allow(clippy::too_many_arguments)]
324    pub fn record_tool_result(
325        &self,
326        session_id: &str,
327        turn: usize,
328        tool_call_id: &str,
329        tool_name: &str,
330        output: &str,
331        exit_code: i32,
332        duration_ms: u64,
333        metadata: &Option<Value>,
334        error_kind: Option<String>,
335    ) {
336        let Some(inner) = &self.inner else {
337            return;
338        };
339        let payload = json!({
340            "turn": turn,
341            "tool_call_id": tool_call_id,
342            "tool": tool_name,
343            "success": exit_code == 0,
344            "exit_code": exit_code,
345            "duration_ms": duration_ms,
346            "output": inner.capture_text(output),
347            "metadata": metadata,
348            "error_kind": error_kind,
349        });
350        inner.record("tool_result", session_id, payload);
351    }
352
353    pub fn record_context_compacted(
354        &self,
355        session_id: &str,
356        before_messages: usize,
357        after_messages: &[Message],
358        percent_before: f32,
359    ) {
360        let Some(inner) = &self.inner else {
361            return;
362        };
363        let payload = json!({
364            "before_messages": before_messages,
365            "after_messages": after_messages.len(),
366            "percent_before": percent_before,
367            "messages": inner.capture_messages(after_messages),
368        });
369        inner.record("context_compacted", session_id, payload);
370    }
371
372    pub fn record_execution_end(
373        &self,
374        session_id: &str,
375        success: bool,
376        response_text: Option<&str>,
377        usage: Option<&TokenUsage>,
378        tool_calls_count: Option<usize>,
379        error_message: Option<&str>,
380    ) {
381        let Some(inner) = &self.inner else {
382            return;
383        };
384        let payload = json!({
385            "success": success,
386            "response_text": response_text.map(|s| inner.capture_text(s)),
387            "usage": usage.map(token_usage_value),
388            "tool_calls_count": tool_calls_count,
389            "error_message": error_message.map(|s| inner.capture_text(s)),
390        });
391        inner.record("execution_end", session_id, payload);
392    }
393}
394
395impl RlTrajectoryRecorderInner {
396    fn record(&self, event_type: &str, session_id: &str, payload: Value) {
397        let sequence = self.sequence.fetch_add(1, Ordering::Relaxed) + 1;
398        let record = json!({
399            "schema": RL_TRAJECTORY_SCHEMA,
400            "sequence": sequence,
401            "timestamp_ms": chrono::Utc::now().timestamp_millis(),
402            "event_type": event_type,
403            "session_id": session_id,
404            "mode": self.config.mode,
405            "context": self.context,
406            "payload": payload,
407        });
408
409        let line = match serde_json::to_string(&record) {
410            Ok(line) => line,
411            Err(err) => {
412                tracing::warn!(error = %err, "Failed to serialize RL trajectory record");
413                return;
414            }
415        };
416        let mut file = match self.file.lock() {
417            Ok(file) => file,
418            Err(poisoned) => poisoned.into_inner(),
419        };
420        if let Err(err) = writeln!(file, "{line}") {
421            tracing::warn!(error = %err, "Failed to write RL trajectory record");
422        }
423    }
424
425    fn capture_messages(&self, messages: &[Message]) -> Value {
426        if !self.config.include_messages {
427            return json!({
428                "included": false,
429                "count": messages.len(),
430                "roles": messages.iter().map(|m| m.role.as_str()).collect::<Vec<_>>(),
431            });
432        }
433        Value::Array(
434            messages
435                .iter()
436                .enumerate()
437                .map(|(index, message)| {
438                    let mut value = self.capture_message(message);
439                    if let Value::Object(ref mut object) = value {
440                        object.insert("index".to_string(), json!(index));
441                    }
442                    value
443                })
444                .collect(),
445        )
446    }
447
448    fn capture_message(&self, message: &Message) -> Value {
449        json!({
450            "role": message.role,
451            "content": message.content.iter().map(|block| self.capture_content_block(block)).collect::<Vec<_>>(),
452            "reasoning_content": message.reasoning_content.as_ref().map(|s| self.capture_text(s)),
453        })
454    }
455
456    fn capture_content_block(&self, block: &ContentBlock) -> Value {
457        match block {
458            ContentBlock::Text { text } => json!({
459                "type": "text",
460                "text": self.capture_text(text),
461            }),
462            ContentBlock::Image { source } => json!({
463                "type": "image",
464                "source": {
465                    "media_type": source.media_type,
466                    "data": self.capture_text(&source.data),
467                }
468            }),
469            ContentBlock::ToolUse { id, name, input } => json!({
470                "type": "tool_use",
471                "id": id,
472                "name": name,
473                "input": input,
474            }),
475            ContentBlock::ToolResult {
476                tool_use_id,
477                content,
478                is_error,
479            } => json!({
480                "type": "tool_result",
481                "tool_use_id": tool_use_id,
482                "content": self.capture_text(&content.as_text()),
483                "is_error": is_error,
484            }),
485        }
486    }
487
488    fn capture_text(&self, text: &str) -> CapturedText {
489        let byte_len = text.len();
490        let sha256 = sha256::digest(text);
491        let (captured, truncated) = truncate_utf8(text, self.config.max_text_bytes);
492        CapturedText {
493            byte_len,
494            sha256,
495            truncated,
496            text: Some(captured),
497            preview: None,
498        }
499    }
500}
501
502fn tool_call_value(tool_call: &ToolCall) -> Value {
503    json!({
504        "tool_call_id": tool_call.id,
505        "tool": tool_call.name,
506        "args": tool_call.args,
507    })
508}
509
510fn tool_definition_value(tool: &ToolDefinition) -> Value {
511    json!({
512        "name": tool.name,
513        "description": tool.description,
514        "parameters": tool.parameters,
515    })
516}
517
518fn token_logprob_value(token: &TokenLogProb) -> Value {
519    json!({
520        "token": token.token,
521        "logprob": token.logprob,
522        "bytes": token.bytes,
523        "top_logprobs": token.top_logprobs.iter().map(top_token_logprob_value).collect::<Vec<_>>(),
524    })
525}
526
527fn top_token_logprob_value(token: &TopTokenLogProb) -> Value {
528    json!({
529        "token": token.token,
530        "logprob": token.logprob,
531        "bytes": token.bytes,
532    })
533}
534
535fn token_usage_value(usage: &TokenUsage) -> Value {
536    json!({
537        "prompt_tokens": usage.prompt_tokens,
538        "completion_tokens": usage.completion_tokens,
539        "total_tokens": usage.total_tokens,
540        "cache_read_tokens": usage.cache_read_tokens,
541        "cache_write_tokens": usage.cache_write_tokens,
542    })
543}
544
545fn truncate_utf8(text: &str, max_bytes: usize) -> (String, bool) {
546    if text.len() <= max_bytes {
547        return (text.to_string(), false);
548    }
549    let mut end = max_bytes.min(text.len());
550    while end > 0 && !text.is_char_boundary(end) {
551        end -= 1;
552    }
553    (text[..end].to_string(), true)
554}
555
556fn default_max_text_bytes(mode: RlTrajectoryMode) -> usize {
557    match mode {
558        RlTrajectoryMode::Off => 0,
559        RlTrajectoryMode::On => 1024 * 1024,
560    }
561}
562
563fn parse_bool(value: &str) -> Option<bool> {
564    match value.trim().to_ascii_lowercase().as_str() {
565        "1" | "true" | "yes" | "on" => Some(true),
566        "0" | "false" | "no" | "off" => Some(false),
567        _ => None,
568    }
569}
570
571fn env_first(names: &[&str]) -> Option<String> {
572    names.iter().find_map(|name| {
573        std::env::var(name)
574            .ok()
575            .filter(|value| !value.trim().is_empty())
576    })
577}
578
579#[cfg(test)]
580mod tests {
581    use super::*;
582    use tempfile::tempdir;
583
584    #[test]
585    fn rl_recorder_writes_jsonl() {
586        let dir = tempdir().unwrap();
587        let path = dir.path().join("trajectory.jsonl");
588        let recorder =
589            RlTrajectoryRecorder::from_config(Some(RlTrajectoryConfig::new(&path))).unwrap();
590
591        recorder.record_execution_start(ExecutionStartRecord {
592            session_id: "sess-1",
593            workspace: Path::new("/workspace"),
594            prompt: "solve task",
595            history: &[],
596            system_prompt: Some("system"),
597            max_tool_rounds: 64,
598            planning_mode: "disabled",
599        });
600        recorder.record_tool_result("sess-1", 1, "tool-1", "bash", "ok", 0, 3, &None, None);
601
602        let lines = std::fs::read_to_string(path).unwrap();
603        assert_eq!(lines.lines().count(), 2);
604        let first: Value = serde_json::from_str(lines.lines().next().unwrap()).unwrap();
605        assert_eq!(first["schema"], RL_TRAJECTORY_SCHEMA);
606        assert_eq!(first["event_type"], "execution_start");
607        assert_eq!(first["session_id"], "sess-1");
608    }
609
610    #[test]
611    fn enabled_mode_records_text_with_truncation_flag() {
612        let dir = tempdir().unwrap();
613        let path = dir.path().join("trajectory.jsonl");
614        let recorder = RlTrajectoryRecorder::from_config(Some(
615            RlTrajectoryConfig::new(&path).with_max_text_bytes(3),
616        ))
617        .unwrap();
618
619        recorder.record_execution_start(ExecutionStartRecord {
620            session_id: "sess-1",
621            workspace: Path::new("/workspace"),
622            prompt: "abcdef",
623            history: &[],
624            system_prompt: None,
625            max_tool_rounds: 64,
626            planning_mode: "auto",
627        });
628
629        let text = std::fs::read_to_string(path).unwrap();
630        let record: Value = serde_json::from_str(text.lines().next().unwrap()).unwrap();
631        let prompt = &record["payload"]["prompt"];
632        assert!(prompt.get("sha256").is_some());
633        assert_eq!(prompt["text"], "abc");
634        assert_eq!(prompt["truncated"], true);
635    }
636
637    #[test]
638    fn llm_events_include_tool_definitions_and_token_logprobs() {
639        let dir = tempdir().unwrap();
640        let path = dir.path().join("trajectory.jsonl");
641        let recorder =
642            RlTrajectoryRecorder::from_config(Some(RlTrajectoryConfig::new(&path))).unwrap();
643
644        let tools = vec![ToolDefinition {
645            name: "bash".to_string(),
646            description: "Run a shell command".to_string(),
647            parameters: json!({
648                "type": "object",
649                "properties": {
650                    "cmd": { "type": "string" }
651                },
652                "required": ["cmd"]
653            }),
654        }];
655        recorder.record_llm_request("sess-1", 1, &[Message::user("hi")], None, &tools, 7);
656
657        recorder.record_llm_response(
658            "sess-1",
659            1,
660            &LlmResponse {
661                message: Message {
662                    role: "assistant".to_string(),
663                    content: vec![ContentBlock::Text {
664                        text: "hello".to_string(),
665                    }],
666                    reasoning_content: None,
667                },
668                usage: TokenUsage {
669                    prompt_tokens: 7,
670                    completion_tokens: 1,
671                    total_tokens: 8,
672                    cache_read_tokens: None,
673                    cache_write_tokens: None,
674                },
675                stop_reason: Some("stop".to_string()),
676                token_logprobs: vec![TokenLogProb {
677                    token: "hello".to_string(),
678                    logprob: -0.2,
679                    bytes: Some(vec![104, 101, 108, 108, 111]),
680                    top_logprobs: vec![TopTokenLogProb {
681                        token: "hi".to_string(),
682                        logprob: -1.2,
683                        bytes: Some(vec![104, 105]),
684                    }],
685                }],
686                meta: None,
687            },
688            42,
689        );
690
691        let lines = std::fs::read_to_string(path).unwrap();
692        let records = lines
693            .lines()
694            .map(|line| serde_json::from_str::<Value>(line).unwrap())
695            .collect::<Vec<_>>();
696        let request = records
697            .iter()
698            .find(|record| record["event_type"] == "llm_request")
699            .unwrap();
700        assert_eq!(request["payload"]["available_tools"][0], "bash");
701        assert_eq!(request["payload"]["tool_definitions"][0]["name"], "bash");
702        assert_eq!(
703            request["payload"]["tool_definitions"][0]["parameters"]["required"][0],
704            "cmd"
705        );
706
707        let response = records
708            .iter()
709            .find(|record| record["event_type"] == "llm_response")
710            .unwrap();
711        assert_eq!(response["payload"]["token_logprobs"][0]["token"], "hello");
712        assert_eq!(response["payload"]["token_logprobs"][0]["logprob"], -0.2);
713        assert_eq!(
714            response["payload"]["token_logprobs"][0]["top_logprobs"][0]["token"],
715            "hi"
716        );
717    }
718}
a3s_code_core/rl_trajectory.rs

a3s_code_core/
rl_trajectory.rs
