refactor: streamline RL trajectory recording

Author: Lixtt

core/src/agent/llm_turn.rs

@@ -3,7 +3,7 @@ use super::{AgentEvent, AgentLoop};
 use crate::hooks::{
     ErrorType, GenerateEndEvent, GenerateStartEvent, HookEvent, TokenUsageInfo, ToolCallInfo,
 };
-use crate::llm::{LlmResponse, Message, ToolCall};
+use crate::llm::{LlmResponse, Message, ToolCall, ToolDefinition};
 use anyhow::Context;
 use std::time::Duration;
 use tokio::sync::mpsc;
@@ -14,6 +14,16 @@ pub(super) struct LlmTurnOutput {
     pub(super) tool_calls: Vec<ToolCall>,
 }
 
+struct LlmCallRequest<'a> {
+    turn: usize,
+    messages: &'a [Message],
+    system: Option<&'a str>,
+    tools: &'a [ToolDefinition],
+    session_id: Option<&'a str>,
+    event_tx: &'a Option<mpsc::Sender<AgentEvent>>,
+    cancel_token: &'a tokio_util::sync::CancellationToken,
+}
+
 impl AgentLoop {
     pub(super) async fn execute_llm_turn(
         &self,
@@ -49,14 +59,15 @@ impl AgentLoop {
 
         let llm_start = std::time::Instant::now();
         let response = self
-            .call_llm_with_circuit_breaker(
+            .call_llm_with_circuit_breaker(LlmCallRequest {
                 turn,
-                &state.messages,
-                augmented_system.as_deref(),
+                messages: &state.messages,
+                system: augmented_system.as_deref(),
+                tools: &selected_tools,
                 session_id,
                 event_tx,
                 cancel_token,
-            )
+            })
             .await?;
 
         state.record_usage(&response.usage);
@@ -122,19 +133,14 @@ impl AgentLoop {
 
     async fn call_llm_with_circuit_breaker(
         &self,
-        turn: usize,
-        messages: &[Message],
-        system: Option<&str>,
-        session_id: Option<&str>,
-        event_tx: &Option<mpsc::Sender<AgentEvent>>,
-        cancel_token: &tokio_util::sync::CancellationToken,
+        request: LlmCallRequest<'_>,
     ) -> anyhow::Result<LlmResponse> {
         // Consult the host's BudgetGuard once per turn (not per retry).
         // A `Deny` bails out before the LLM is touched; a `SoftLimit`
         // surfaces a BudgetThresholdHit event and proceeds.
         if let Some(guard) = &self.config.budget_guard {
-            let sid = session_id.unwrap_or("");
-            let estimate = estimate_prompt_tokens(messages, system);
+            let sid = request.session_id.unwrap_or("");
+            let estimate = estimate_prompt_tokens(request.messages, request.system);
             match guard.check_before_llm(sid, estimate).await {
                 crate::budget::BudgetDecision::Allow => {}
                 crate::budget::BudgetDecision::SoftLimit {
@@ -143,7 +149,7 @@ impl AgentLoop {
                     limit,
                     message,
                 } => {
-                    if let Some(tx) = event_tx {
+                    if let Some(tx) = request.event_tx {
                         let _ = tx
                             .send(AgentEvent::BudgetThresholdHit {
                                 resource,
@@ -156,7 +162,7 @@ impl AgentLoop {
                     }
                 }
                 crate::budget::BudgetDecision::Deny { resource, reason } => {
-                    if let Some(tx) = event_tx {
+                    if let Some(tx) = request.event_tx {
                         let _ = tx
                             .send(AgentEvent::BudgetThresholdHit {
                                 resource: resource.clone(),
@@ -178,23 +184,31 @@ impl AgentLoop {
         loop {
             attempt += 1;
             let result = self
-                .call_llm(messages, system, event_tx, cancel_token)
+                .call_llm(
+                    request.messages,
+                    request.system,
+                    request.tools,
+                    request.event_tx,
+                    request.cancel_token,
+                )
                 .await;
             match result {
                 Ok(response) => {
                     if let Some(guard) = &self.config.budget_guard {
                         guard
-                            .record_after_llm(session_id.unwrap_or(""), &response.usage)
+                            .record_after_llm(request.session_id.unwrap_or(""), &response.usage)
                             .await;
                     }
                     return Ok(response);
                 }
-                Err(error) if cancel_token.is_cancelled() => {
+                Err(error) if request.cancel_token.is_cancelled() => {
                     anyhow::bail!(error);
                 }
-                Err(error) if attempt < threshold && (event_tx.is_none() || attempt == 1) => {
+                Err(error)
+                    if attempt < threshold && (request.event_tx.is_none() || attempt == 1) =>
+                {
                     tracing::warn!(
-                        turn = turn,
+                        turn = request.turn,
                         attempt = attempt,
                         threshold = threshold,
                         error = %error,
@@ -211,15 +225,15 @@ impl AgentLoop {
                     } else {
                         format!("LLM call failed: {}", error)
                     };
-                    tracing::error!(turn = turn, attempt = attempt, "{}", msg);
+                    tracing::error!(turn = request.turn, attempt = attempt, "{}", msg);
                     self.fire_on_error(
-                        session_id.unwrap_or(""),
+                        request.session_id.unwrap_or(""),
                         ErrorType::LlmFailure,
                         &msg,
-                        serde_json::json!({"turn": turn, "attempt": attempt}),
+                        serde_json::json!({"turn": request.turn, "attempt": attempt}),
                     )
                     .await;
-                    self.emit_error(event_tx, msg.clone()).await;
+                    self.emit_error(request.event_tx, msg.clone()).await;
                     anyhow::bail!(msg);
                 }
             }
@@ -365,23 +379,22 @@ impl AgentLoop {
     /// Streaming events (`TextDelta`, `ToolStart`) are forwarded to `event_tx`
     /// as they arrive. Non-streaming mode simply awaits the complete response.
     ///
-    /// Tool definitions are selected per turn by the centralized tool selector.
+    /// Tool definitions are selected once per turn by the centralized tool selector.
     ///
     /// Returns `Err` on any LLM API failure. The circuit breaker in
     /// `execute_loop` wraps this call with retry logic for non-streaming mode.
     async fn call_llm(
         &self,
         messages: &[Message],
         system: Option<&str>,
+        tools: &[ToolDefinition],
         event_tx: &Option<mpsc::Sender<AgentEvent>>,
         cancel_token: &tokio_util::sync::CancellationToken,
     ) -> anyhow::Result<LlmResponse> {
-        let tools = crate::tools::select_tools_for_messages(&self.config.tools, messages);
-
         if event_tx.is_some() {
             let mut stream_rx = match self
                 .llm_client
-                .complete_streaming(messages, system, &tools, cancel_token.clone())
+                .complete_streaming(messages, system, tools, cancel_token.clone())
                 .await
             {
                 Ok(rx) => rx,
@@ -396,7 +409,7 @@ impl AgentLoop {
                     );
                     return self
                         .llm_client
-                        .complete(messages, system, &tools)
+                        .complete(messages, system, tools)
                         .await
                         .with_context(|| {
                             format!(
@@ -447,7 +460,7 @@ impl AgentLoop {
             final_response.context("Stream ended without final response")
         } else {
             self.llm_client
-                .complete(messages, system, &tools)
+                .complete(messages, system, tools)
                 .await
                 .context("LLM call failed")
         }

core/src/agent/loop_runtime.rs

@@ -100,13 +100,15 @@ impl AgentLoop {
         let augmented_system = turn_context.augmented_system;
 
         self.config.rl_trajectory_recorder.record_execution_start(
-            session_id.unwrap_or(""),
-            &self.tool_context.workspace,
-            effective_prompt,
-            history,
-            augmented_system.as_deref(),
-            self.config.max_tool_rounds,
-            &format!("{:?}", self.config.planning_mode),
+            crate::rl_trajectory::ExecutionStartRecord {
+                session_id: session_id.unwrap_or(""),
+                workspace: &self.tool_context.workspace,
+                prompt: effective_prompt,
+                history,
+                system_prompt: augmented_system.as_deref(),
+                max_tool_rounds: self.config.max_tool_rounds,
+                planning_mode: &format!("{:?}", self.config.planning_mode),
+            },
         );
 
         // Add user message

core/src/agent_api/tests.rs

@@ -2700,15 +2700,15 @@ async fn test_session_options_with_rl_trajectory_records_jsonl() {
     session
         .config
         .rl_trajectory_recorder
-        .record_execution_start(
-            "sess-rl",
-            std::path::Path::new("/tmp/test-ws-rl-trajectory"),
-            "abcdef",
-            &[],
-            None,
-            16,
-            "disabled",
-        );
+        .record_execution_start(crate::rl_trajectory::ExecutionStartRecord {
+            session_id: "sess-rl",
+            workspace: std::path::Path::new("/tmp/test-ws-rl-trajectory"),
+            prompt: "abcdef",
+            history: &[],
+            system_prompt: None,
+            max_tool_rounds: 16,
+            planning_mode: "disabled",
+        });
 
     let content = std::fs::read_to_string(&trajectory_path).unwrap();
     let record: serde_json::Value = serde_json::from_str(content.lines().next().unwrap()).unwrap();

core/src/rl_trajectory.rs

@@ -161,6 +161,16 @@ pub struct RlTrajectoryRecorder {
     inner: Option<Arc<RlTrajectoryRecorderInner>>,
 }
 
+pub struct ExecutionStartRecord<'a> {
+    pub session_id: &'a str,
+    pub workspace: &'a Path,
+    pub prompt: &'a str,
+    pub history: &'a [Message],
+    pub system_prompt: Option<&'a str>,
+    pub max_tool_rounds: usize,
+    pub planning_mode: &'a str,
+}
+
 struct RlTrajectoryRecorderInner {
     config: RlTrajectoryConfig,
     context: RlTrajectoryContext,
@@ -228,29 +238,20 @@ impl RlTrajectoryRecorder {
         self.inner.is_some()
     }
 
-    pub fn record_execution_start(
-        &self,
-        session_id: &str,
-        workspace: &Path,
-        prompt: &str,
-        history: &[Message],
-        system_prompt: Option<&str>,
-        max_tool_rounds: usize,
-        planning_mode: &str,
-    ) {
+    pub fn record_execution_start(&self, record: ExecutionStartRecord<'_>) {
         let Some(inner) = &self.inner else {
             return;
         };
         let payload = json!({
-            "workspace": workspace.display().to_string(),
-            "prompt": inner.capture_text(prompt),
-            "history_message_count": history.len(),
-            "history": inner.capture_messages(history),
-            "system_prompt": system_prompt.map(|s| inner.capture_text(s)),
-            "max_tool_rounds": max_tool_rounds,
-            "planning_mode": planning_mode,
+            "workspace": record.workspace.display().to_string(),
+            "prompt": inner.capture_text(record.prompt),
+            "history_message_count": record.history.len(),
+            "history": inner.capture_messages(record.history),
+            "system_prompt": record.system_prompt.map(|s| inner.capture_text(s)),
+            "max_tool_rounds": record.max_tool_rounds,
+            "planning_mode": record.planning_mode,
         });
-        inner.record("execution_start", session_id, payload);
+        inner.record("execution_start", record.session_id, payload);
     }
 
     pub fn record_llm_request(
@@ -587,15 +588,15 @@ mod tests {
         let recorder =
             RlTrajectoryRecorder::from_config(Some(RlTrajectoryConfig::new(&path))).unwrap();
 
-        recorder.record_execution_start(
-            "sess-1",
-            Path::new("/workspace"),
-            "solve task",
-            &[],
-            Some("system"),
-            64,
-            "disabled",
-        );
+        recorder.record_execution_start(ExecutionStartRecord {
+            session_id: "sess-1",
+            workspace: Path::new("/workspace"),
+            prompt: "solve task",
+            history: &[],
+            system_prompt: Some("system"),
+            max_tool_rounds: 64,
+            planning_mode: "disabled",
+        });
         recorder.record_tool_result("sess-1", 1, "tool-1", "bash", "ok", 0, 3, &None, None);
 
         let lines = std::fs::read_to_string(path).unwrap();
@@ -615,15 +616,15 @@ mod tests {
         ))
         .unwrap();
 
-        recorder.record_execution_start(
-            "sess-1",
-            Path::new("/workspace"),
-            "abcdef",
-            &[],
-            None,
-            64,
-            "auto",
-        );
+        recorder.record_execution_start(ExecutionStartRecord {
+            session_id: "sess-1",
+            workspace: Path::new("/workspace"),
+            prompt: "abcdef",
+            history: &[],
+            system_prompt: None,
+            max_tool_rounds: 64,
+            planning_mode: "auto",
+        });
 
         let text = std::fs::read_to_string(path).unwrap();
         let record: Value = serde_json::from_str(text.lines().next().unwrap()).unwrap();