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                ..
480            } => json!({
481                "type": "tool_result",
482                "tool_use_id": tool_use_id,
483                "content": self.capture_text(&content.as_text()),
484                "is_error": is_error,
485            }),
486        }
487    }
488
489    fn capture_text(&self, text: &str) -> CapturedText {
490        let byte_len = text.len();
491        let sha256 = sha256::digest(text);
492        let (captured, truncated) = truncate_utf8(text, self.config.max_text_bytes);
493        CapturedText {
494            byte_len,
495            sha256,
496            truncated,
497            text: Some(captured),
498            preview: None,
499        }
500    }
501}
502
503fn tool_call_value(tool_call: &ToolCall) -> Value {
504    json!({
505        "tool_call_id": tool_call.id,
506        "tool": tool_call.name,
507        "args": tool_call.args,
508    })
509}
510
511fn tool_definition_value(tool: &ToolDefinition) -> Value {
512    json!({
513        "name": tool.name,
514        "description": tool.description,
515        "parameters": tool.parameters,
516    })
517}
518
519fn token_logprob_value(token: &TokenLogProb) -> Value {
520    json!({
521        "token": token.token,
522        "logprob": token.logprob,
523        "bytes": token.bytes,
524        "top_logprobs": token.top_logprobs.iter().map(top_token_logprob_value).collect::<Vec<_>>(),
525    })
526}
527
528fn top_token_logprob_value(token: &TopTokenLogProb) -> Value {
529    json!({
530        "token": token.token,
531        "logprob": token.logprob,
532        "bytes": token.bytes,
533    })
534}
535
536fn token_usage_value(usage: &TokenUsage) -> Value {
537    json!({
538        "prompt_tokens": usage.prompt_tokens,
539        "completion_tokens": usage.completion_tokens,
540        "total_tokens": usage.total_tokens,
541        "cache_read_tokens": usage.cache_read_tokens,
542        "cache_write_tokens": usage.cache_write_tokens,
543    })
544}
545
546fn truncate_utf8(text: &str, max_bytes: usize) -> (String, bool) {
547    if text.len() <= max_bytes {
548        return (text.to_string(), false);
549    }
550    let mut end = max_bytes.min(text.len());
551    while end > 0 && !text.is_char_boundary(end) {
552        end -= 1;
553    }
554    (text[..end].to_string(), true)
555}
556
557fn default_max_text_bytes(mode: RlTrajectoryMode) -> usize {
558    match mode {
559        RlTrajectoryMode::Off => 0,
560        RlTrajectoryMode::On => 1024 * 1024,
561    }
562}
563
564fn parse_bool(value: &str) -> Option<bool> {
565    match value.trim().to_ascii_lowercase().as_str() {
566        "1" | "true" | "yes" | "on" => Some(true),
567        "0" | "false" | "no" | "off" => Some(false),
568        _ => None,
569    }
570}
571
572fn env_first(names: &[&str]) -> Option<String> {
573    names.iter().find_map(|name| {
574        std::env::var(name)
575            .ok()
576            .filter(|value| !value.trim().is_empty())
577    })
578}
579
580#[cfg(test)]
581mod tests {
582    use super::*;
583    use tempfile::tempdir;
584
585    #[test]
586    fn rl_recorder_writes_jsonl() {
587        let dir = tempdir().unwrap();
588        let path = dir.path().join("trajectory.jsonl");
589        let recorder =
590            RlTrajectoryRecorder::from_config(Some(RlTrajectoryConfig::new(&path))).unwrap();
591
592        recorder.record_execution_start(ExecutionStartRecord {
593            session_id: "sess-1",
594            workspace: Path::new("/workspace"),
595            prompt: "solve task",
596            history: &[],
597            system_prompt: Some("system"),
598            max_tool_rounds: 64,
599            planning_mode: "disabled",
600        });
601        recorder.record_tool_result("sess-1", 1, "tool-1", "bash", "ok", 0, 3, &None, None);
602
603        let lines = std::fs::read_to_string(path).unwrap();
604        assert_eq!(lines.lines().count(), 2);
605        let first: Value = serde_json::from_str(lines.lines().next().unwrap()).unwrap();
606        assert_eq!(first["schema"], RL_TRAJECTORY_SCHEMA);
607        assert_eq!(first["event_type"], "execution_start");
608        assert_eq!(first["session_id"], "sess-1");
609    }
610
611    #[test]
612    fn enabled_mode_records_text_with_truncation_flag() {
613        let dir = tempdir().unwrap();
614        let path = dir.path().join("trajectory.jsonl");
615        let recorder = RlTrajectoryRecorder::from_config(Some(
616            RlTrajectoryConfig::new(&path).with_max_text_bytes(3),
617        ))
618        .unwrap();
619
620        recorder.record_execution_start(ExecutionStartRecord {
621            session_id: "sess-1",
622            workspace: Path::new("/workspace"),
623            prompt: "abcdef",
624            history: &[],
625            system_prompt: None,
626            max_tool_rounds: 64,
627            planning_mode: "auto",
628        });
629
630        let text = std::fs::read_to_string(path).unwrap();
631        let record: Value = serde_json::from_str(text.lines().next().unwrap()).unwrap();
632        let prompt = &record["payload"]["prompt"];
633        assert!(prompt.get("sha256").is_some());
634        assert_eq!(prompt["text"], "abc");
635        assert_eq!(prompt["truncated"], true);
636    }
637
638    #[test]
639    fn llm_events_include_tool_definitions_and_token_logprobs() {
640        let dir = tempdir().unwrap();
641        let path = dir.path().join("trajectory.jsonl");
642        let recorder =
643            RlTrajectoryRecorder::from_config(Some(RlTrajectoryConfig::new(&path))).unwrap();
644
645        let tools = vec![ToolDefinition {
646            name: "bash".to_string(),
647            description: "Run a shell command".to_string(),
648            parameters: json!({
649                "type": "object",
650                "properties": {
651                    "cmd": { "type": "string" }
652                },
653                "required": ["cmd"]
654            }),
655        }];
656        recorder.record_llm_request("sess-1", 1, &[Message::user("hi")], None, &tools, 7);
657
658        recorder.record_llm_response(
659            "sess-1",
660            1,
661            &LlmResponse {
662                message: Message {
663                    role: "assistant".to_string(),
664                    content: vec![ContentBlock::Text {
665                        text: "hello".to_string(),
666                    }],
667                    reasoning_content: None,
668                },
669                usage: TokenUsage {
670                    prompt_tokens: 7,
671                    completion_tokens: 1,
672                    total_tokens: 8,
673                    cache_read_tokens: None,
674                    cache_write_tokens: None,
675                },
676                stop_reason: Some("stop".to_string()),
677                token_logprobs: vec![TokenLogProb {
678                    token: "hello".to_string(),
679                    logprob: -0.2,
680                    bytes: Some(vec![104, 101, 108, 108, 111]),
681                    top_logprobs: vec![TopTokenLogProb {
682                        token: "hi".to_string(),
683                        logprob: -1.2,
684                        bytes: Some(vec![104, 105]),
685                    }],
686                }],
687                meta: None,
688            },
689            42,
690        );
691
692        let lines = std::fs::read_to_string(path).unwrap();
693        let records = lines
694            .lines()
695            .map(|line| serde_json::from_str::<Value>(line).unwrap())
696            .collect::<Vec<_>>();
697        let request = records
698            .iter()
699            .find(|record| record["event_type"] == "llm_request")
700            .unwrap();
701        assert_eq!(request["payload"]["available_tools"][0], "bash");
702        assert_eq!(request["payload"]["tool_definitions"][0]["name"], "bash");
703        assert_eq!(
704            request["payload"]["tool_definitions"][0]["parameters"]["required"][0],
705            "cmd"
706        );
707
708        let response = records
709            .iter()
710            .find(|record| record["event_type"] == "llm_response")
711            .unwrap();
712        assert_eq!(response["payload"]["token_logprobs"][0]["token"], "hello");
713        assert_eq!(response["payload"]["token_logprobs"][0]["logprob"], -0.2);
714        assert_eq!(
715            response["payload"]["token_logprobs"][0]["top_logprobs"][0]["token"],
716            "hi"
717        );
718    }
719}
a3s_code_core/rl_trajectory.rs

a3s_code_core/
rl_trajectory.rs
