From 0631eb6e917ef4be771ecbff0d1e23dfae63a899 Mon Sep 17 00:00:00 2001 From: Li Xiangtian <1193027052@qq.com> Date: Wed, 24 Jun 2026 23:35:22 +0800 Subject: [PATCH 1/5] Add RL trajectory capture --- core/src/agent.rs | 4 + core/src/agent/execution_mode.rs | 16 + core/src/agent/execution_state.rs | 4 + core/src/agent/llm_turn.rs | 27 + core/src/agent/loop_runtime.rs | 10 + core/src/agent/tool_completion_runtime.rs | 15 + core/src/agent/tool_execution_runtime.rs | 26 + core/src/agent/tool_guard_runtime.rs | 23 + core/src/agent/tool_turn.rs | 8 +- core/src/agent_api.rs | 7 + core/src/agent_api/session_builder.rs | 7 + core/src/agent_api/session_options.rs | 11 + core/src/agent_api/tests.rs | 36 ++ core/src/lib.rs | 2 + core/src/rl_trajectory.rs | 601 ++++++++++++++++++++++ sdk/python/src/lib.rs | 105 +++- 16 files changed, 900 insertions(+), 2 deletions(-) create mode 100644 core/src/rl_trajectory.rs diff --git a/core/src/agent.rs b/core/src/agent.rs index ce3167c8..8624a567 100644 --- a/core/src/agent.rs +++ b/core/src/agent.rs @@ -91,6 +91,8 @@ pub(crate) struct AgentConfig { pub goal_tracking: bool, /// Optional hook engine for firing lifecycle events (PreToolUse, PostToolUse, etc.) pub hook_engine: Option>, + /// Optional structured JSONL trajectory recorder for RL training and service data capture. + pub rl_trajectory_recorder: crate::rl_trajectory::RlTrajectoryRecorder, /// Optional skill registry for tool permission enforcement pub skill_registry: Option>, /// When true, active skill `allowed-tools` restrict ordinary session tool calls. @@ -193,6 +195,7 @@ impl std::fmt::Debug for AgentConfig { .field("planning_mode", &self.planning_mode) .field("goal_tracking", &self.goal_tracking) .field("hook_engine", &self.hook_engine.is_some()) + .field("rl_trajectory", &self.rl_trajectory_recorder.is_enabled()) .field( "skill_registry", &self.skill_registry.as_ref().map(|r| r.len()), @@ -241,6 +244,7 @@ impl Default for AgentConfig { planning_mode: PlanningMode::default(), goal_tracking: false, hook_engine: None, + rl_trajectory_recorder: crate::rl_trajectory::RlTrajectoryRecorder::disabled(), skill_registry: Some(Arc::new(crate::skills::SkillRegistry::with_builtins())), enforce_active_skill_tool_restrictions: false, max_parse_retries: 2, diff --git a/core/src/agent/execution_mode.rs b/core/src/agent/execution_mode.rs index 99ada7bb..9cb090a7 100644 --- a/core/src/agent/execution_mode.rs +++ b/core/src/agent/execution_mode.rs @@ -190,6 +190,14 @@ impl AgentLoop { a3s.llm.total_tokens = r.usage.total_tokens, "a3s.agent.execute completed" ); + self.config.rl_trajectory_recorder.record_execution_end( + session_id.unwrap_or(""), + true, + Some(&r.text), + Some(&r.usage), + Some(r.tool_calls_count), + None, + ); self.fire_post_response( session_id.unwrap_or(""), &r.text, @@ -204,6 +212,14 @@ impl AgentLoop { error = %e, "a3s.agent.execute failed" ); + self.config.rl_trajectory_recorder.record_execution_end( + session_id.unwrap_or(""), + false, + None, + None, + None, + Some(&e.to_string()), + ); self.fire_on_error( session_id.unwrap_or(""), ErrorType::Other, diff --git a/core/src/agent/execution_state.rs b/core/src/agent/execution_state.rs index 5f8e8f1e..851252e0 100644 --- a/core/src/agent/execution_state.rs +++ b/core/src/agent/execution_state.rs @@ -70,6 +70,10 @@ impl ExecutionLoopState { self.turn } + pub(super) fn current_turn(&self) -> usize { + self.turn + } + pub(super) fn continuation_count(&self) -> u32 { self.continuation_count } diff --git a/core/src/agent/llm_turn.rs b/core/src/agent/llm_turn.rs index d727f608..309c3123 100644 --- a/core/src/agent/llm_turn.rs +++ b/core/src/agent/llm_turn.rs @@ -33,6 +33,20 @@ impl AgentLoop { "LLM completion started" ); + let selected_tool_names = + crate::tools::select_tools_for_messages(&self.config.tools, &state.messages) + .into_iter() + .map(|tool| tool.name) + .collect::>(); + self.config.rl_trajectory_recorder.record_llm_request( + session_id.unwrap_or(""), + turn, + &state.messages, + augmented_system.as_deref(), + &selected_tool_names, + estimate_prompt_tokens(&state.messages, augmented_system.as_deref()), + ); + self.fire_generate_start(session_id.unwrap_or(""), effective_prompt, augmented_system) .await; @@ -255,6 +269,12 @@ impl AgentLoop { a3s.llm.total_tokens = response.usage.total_tokens, "Turn token usage" ); + self.config.rl_trajectory_recorder.record_llm_response( + session_id.unwrap_or(""), + turn, + response, + llm_duration.as_millis() as u64, + ); } async fn emit_turn_end( @@ -318,6 +338,13 @@ impl AgentLoop { state.messages = compacted; } + self.config.rl_trajectory_recorder.record_context_compacted( + session_id.unwrap_or(""), + before_len, + &state.messages, + percent_before, + ); + if let Some(tx) = event_tx { tx.send(AgentEvent::ContextCompacted { session_id: session_id.unwrap_or("").to_string(), diff --git a/core/src/agent/loop_runtime.rs b/core/src/agent/loop_runtime.rs index 849714bc..913a7d28 100644 --- a/core/src/agent/loop_runtime.rs +++ b/core/src/agent/loop_runtime.rs @@ -99,6 +99,16 @@ impl AgentLoop { let effective_prompt = turn_context.effective_prompt.as_str(); 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), + ); + // Add user message if !msg_prompt.is_empty() { state.messages.push(Message::user(msg_prompt)); diff --git a/core/src/agent/tool_completion_runtime.rs b/core/src/agent/tool_completion_runtime.rs index 43921ed6..ea643949 100644 --- a/core/src/agent/tool_completion_runtime.rs +++ b/core/src/agent/tool_completion_runtime.rs @@ -50,6 +50,21 @@ impl AgentLoop { ) .await; + self.config.rl_trajectory_recorder.record_tool_result( + session_id.unwrap_or(""), + state.current_turn(), + &tool_call.id, + &tool_call.name, + &output, + normalized.exit_code, + tool_duration.as_millis() as u64, + &normalized.metadata, + normalized + .error_kind + .as_ref() + .map(|kind| format!("{kind:?}")), + ); + self.remember_tool_result( effective_prompt, &tool_call.name, diff --git a/core/src/agent/tool_execution_runtime.rs b/core/src/agent/tool_execution_runtime.rs index f0789c4a..ea106ad3 100644 --- a/core/src/agent/tool_execution_runtime.rs +++ b/core/src/agent/tool_execution_runtime.rs @@ -1,5 +1,6 @@ use super::tool_result_runtime::NormalizedToolResult; use super::{AgentEvent, AgentLoop, ToolCommand}; +use crate::llm::ToolCall; use crate::tools::{ToolContext, ToolStreamEvent}; use serde_json::Value; use std::sync::Arc; @@ -41,6 +42,16 @@ impl AgentLoop { event_tx: &Option>, ) -> (String, i32, bool, Option) { let call_id = format!("plan-{}-{}", tool_name, uuid::Uuid::new_v4()); + let synthetic_call = ToolCall { + id: call_id.clone(), + name: tool_name.to_string(), + args: args.clone(), + }; + self.config.rl_trajectory_recorder.record_tool_call( + session_id.unwrap_or(""), + 0, + &synthetic_call, + ); if let Some(tx) = event_tx { tx.send(AgentEvent::ToolStart { id: call_id.clone(), @@ -51,9 +62,24 @@ impl AgentLoop { } let ctx = self.tool_context_for_plan(session_id); + let started = std::time::Instant::now(); let normalized = NormalizedToolResult::from_execution( self.execute_tool_timed(tool_name, args, &ctx).await, ); + self.config.rl_trajectory_recorder.record_tool_result( + session_id.unwrap_or(""), + 0, + &call_id, + tool_name, + &normalized.output, + normalized.exit_code, + started.elapsed().as_millis() as u64, + &normalized.metadata, + normalized + .error_kind + .as_ref() + .map(|kind| format!("{kind:?}")), + ); if let Some(tx) = event_tx { tx.send(AgentEvent::ToolEnd { diff --git a/core/src/agent/tool_guard_runtime.rs b/core/src/agent/tool_guard_runtime.rs index f4981a87..92d1966f 100644 --- a/core/src/agent/tool_guard_runtime.rs +++ b/core/src/agent/tool_guard_runtime.rs @@ -13,6 +13,7 @@ impl AgentLoop { tool_call: &ToolCall, state: &mut ExecutionLoopState, event_tx: &Option>, + session_id: Option<&str>, ) -> anyhow::Result { if let Some((duplicate_count, error_msg)) = state.duplicate_tool_call( &tool_call.name, @@ -37,6 +38,17 @@ impl AgentLoop { state .messages .push(Message::tool_result(&tool_call.id, &error_msg, true)); + self.config.rl_trajectory_recorder.record_tool_result( + session_id.unwrap_or(""), + state.current_turn(), + &tool_call.id, + &tool_call.name, + &error_msg, + 1, + 0, + &None, + Some("duplicate_tool_call".to_string()), + ); return Ok(true); } @@ -68,6 +80,17 @@ impl AgentLoop { &parse_outcome.output, true, )); + self.config.rl_trajectory_recorder.record_tool_result( + session_id.unwrap_or(""), + state.current_turn(), + &tool_call.id, + &tool_call.name, + &parse_outcome.output, + 1, + 0, + &None, + Some("parse_error".to_string()), + ); if let Some(msg) = parse_outcome.fatal_message { tracing::error!("{}", msg); diff --git a/core/src/agent/tool_turn.rs b/core/src/agent/tool_turn.rs index 37146088..8e5b0bdf 100644 --- a/core/src/agent/tool_turn.rs +++ b/core/src/agent/tool_turn.rs @@ -43,6 +43,12 @@ impl AgentLoop { ) -> anyhow::Result<()> { state.record_tool_call(); let tool_start = std::time::Instant::now(); + let turn = state.current_turn(); + self.config.rl_trajectory_recorder.record_tool_call( + session_id.unwrap_or(""), + turn, + &tool_call, + ); tracing::info!( tool_name = tool_call.name.as_str(), @@ -51,7 +57,7 @@ impl AgentLoop { ); if self - .handle_tool_preflight_guard(&tool_call, state, event_tx) + .handle_tool_preflight_guard(&tool_call, state, event_tx, session_id) .await? { return Ok(()); diff --git a/core/src/agent_api.rs b/core/src/agent_api.rs index 80fa2c23..79f5319a 100644 --- a/core/src/agent_api.rs +++ b/core/src/agent_api.rs @@ -210,6 +210,13 @@ pub struct SessionOptions { /// tasks). `None` (default) keeps everything — fine for short /// sessions, a memory leak for hours-long cluster workloads. pub retention_limits: Option, + /// Optional structured JSONL trajectory config. + /// + /// When set, a3s-code records user prompts, LLM turns, tool calls, + /// tool observations, token usage, and execution end status for RL + /// training or service data collection. If unset, the same config can + /// be enabled by `A3S_CODE_TRAJECTORY_PATH`. + pub rl_trajectory: Option, /// Auto-save after each completed `send()` or default-history `stream()` call. pub auto_save: bool, /// Optional artifact retention limits for large tool/program outputs. diff --git a/core/src/agent_api/session_builder.rs b/core/src/agent_api/session_builder.rs index 268ab3dc..857ad2b3 100644 --- a/core/src/agent_api/session_builder.rs +++ b/core/src/agent_api/session_builder.rs @@ -138,6 +138,12 @@ pub(super) fn build_agent_session( let base = agent.config.clone(); let auto_delegation = resolve_auto_delegation_config(&agent.code_config, opts); + let rl_trajectory_config = match opts.rl_trajectory.clone() { + Some(config) => Some(config), + None => crate::rl_trajectory::RlTrajectoryConfig::from_env()?, + }; + let rl_trajectory_recorder = + crate::rl_trajectory::RlTrajectoryRecorder::from_config(rl_trajectory_config)?; let config = AgentConfig { prompt_slots, tools: tool_defs, @@ -181,6 +187,7 @@ pub(super) fn build_agent_session( agent_registry: Some(Arc::clone(&agent_registry)), max_execution_time_ms: opts.max_execution_time_ms.or(base.max_execution_time_ms), budget_guard: opts.budget_guard.clone().or(base.budget_guard.clone()), + rl_trajectory_recorder, host_env: opts .host_env .clone() diff --git a/core/src/agent_api/session_options.rs b/core/src/agent_api/session_options.rs index c11c7178..ce9a1f3b 100644 --- a/core/src/agent_api/session_options.rs +++ b/core/src/agent_api/session_options.rs @@ -42,6 +42,7 @@ impl std::fmt::Debug for SessionOptions { .field("memory_store", &self.memory_store.is_some()) .field("session_store", &self.session_store.is_some()) .field("session_id", &self.session_id) + .field("rl_trajectory", &self.rl_trajectory) .field("auto_save", &self.auto_save) .field("artifact_store_limits", &self.artifact_store_limits) .field("max_parse_retries", &self.max_parse_retries) @@ -368,6 +369,16 @@ impl SessionOptions { self } + /// Enable structured JSONL trajectory capture for this session. + /// + /// This is the preferred programmatic path for RL training and deployed + /// service data collection. Environment-only deployments can instead set + /// `A3S_CODE_TRAJECTORY_PATH`. + pub fn with_rl_trajectory(mut self, config: crate::rl_trajectory::RlTrajectoryConfig) -> Self { + self.rl_trajectory = Some(config); + self + } + /// Enable auto-save after each `send()` call pub fn with_auto_save(mut self, enabled: bool) -> Self { self.auto_save = enabled; diff --git a/core/src/agent_api/tests.rs b/core/src/agent_api/tests.rs index 57056f7d..480113ed 100644 --- a/core/src/agent_api/tests.rs +++ b/core/src/agent_api/tests.rs @@ -2909,6 +2909,42 @@ async fn test_active_skill_tool_restriction_option_defaults_and_overrides() { assert!(legacy_session.config.enforce_active_skill_tool_restrictions); } +#[tokio::test(flavor = "multi_thread")] +async fn test_session_options_with_rl_trajectory_records_jsonl() { + let dir = tempfile::TempDir::new().unwrap(); + let trajectory_path = dir.path().join("trajectory.jsonl"); + let agent = Agent::from_config(test_config()).await.unwrap(); + + let opts = SessionOptions::new().with_rl_trajectory( + crate::rl_trajectory::RlTrajectoryConfig::new(&trajectory_path).with_max_text_bytes(4), + ); + let session = agent + .session("/tmp/test-ws-rl-trajectory", Some(opts)) + .unwrap(); + + assert!(session.config.rl_trajectory_recorder.is_enabled()); + session + .config + .rl_trajectory_recorder + .record_execution_start( + "sess-rl", + std::path::Path::new("/tmp/test-ws-rl-trajectory"), + "abcdef", + &[], + None, + 16, + "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(); + assert_eq!(record["schema"], crate::rl_trajectory::RL_TRAJECTORY_SCHEMA); + assert_eq!(record["event_type"], "execution_start"); + assert_eq!(record["session_id"], "sess-rl"); + assert_eq!(record["payload"]["prompt"]["text"], "abcd"); + assert_eq!(record["payload"]["prompt"]["truncated"], true); +} + #[tokio::test(flavor = "multi_thread")] async fn test_session_max_parallel_tasks_config_and_override() { let mut config = test_config(); diff --git a/core/src/lib.rs b/core/src/lib.rs index ca80b303..f4614967 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -104,6 +104,7 @@ pub(crate) mod prompts; pub mod queue; pub mod retention; pub(crate) mod retry; +pub mod rl_trajectory; pub mod run; pub(crate) mod safety_gate; pub mod sandbox; @@ -148,6 +149,7 @@ pub use orchestration::{ WorkflowStepRecord, WORKFLOW_CHECKPOINT_SCHEMA_VERSION, }; pub use prompts::{AgentStyle, DetectionConfidence, PlanningMode, SystemPromptSlots}; +pub use rl_trajectory::{RlTrajectoryConfig, RlTrajectoryMode, RlTrajectoryRecorder}; pub use run::{ ActiveToolSnapshot, InMemoryRunStore, RunEventRecord, RunHandle, RunRecord, RunSnapshot, RunStatus, diff --git a/core/src/rl_trajectory.rs b/core/src/rl_trajectory.rs new file mode 100644 index 00000000..77c062fb --- /dev/null +++ b/core/src/rl_trajectory.rs @@ -0,0 +1,601 @@ +//! RL trajectory recording primitives. +//! +//! The default runtime trace is intentionally lightweight and diagnostic. This +//! module records an opt-in training-oriented JSONL stream that can reconstruct LLM +//! turns, tool calls, tool observations, token usage, and termination reasons. +//! It is opt-in: normal sessions pay only a cheap disabled-recorder branch. + +use crate::llm::{ContentBlock, LlmResponse, Message, TokenUsage, ToolCall}; +use anyhow::{Context, Result}; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::fs::{File, OpenOptions}; +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Mutex, +}; + +pub const RL_TRAJECTORY_SCHEMA: &str = "a3s.rl_trajectory.v1"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "snake_case")] +pub enum RlTrajectoryMode { + #[default] + Off, + On, +} + +impl RlTrajectoryMode { + pub fn parse(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "" | "off" | "0" | "false" | "none" => Some(Self::Off), + "on" | "1" | "true" | "yes" | "enabled" | "rl" | "train" | "training" + | "trajectory" | "trace" | "debug" | "full" | "compact" => Some(Self::On), + _ => None, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RlTrajectoryConfig { + pub mode: RlTrajectoryMode, + pub path: PathBuf, + pub max_text_bytes: usize, + pub include_messages: bool, +} + +impl RlTrajectoryConfig { + pub fn new(path: impl Into) -> Self { + Self { + mode: RlTrajectoryMode::On, + path: path.into(), + max_text_bytes: default_max_text_bytes(RlTrajectoryMode::On), + include_messages: true, + } + } + + pub fn with_mode(mut self, mode: RlTrajectoryMode) -> Self { + self.mode = mode; + self.max_text_bytes = default_max_text_bytes(mode); + self.include_messages = mode == RlTrajectoryMode::On; + self + } + + pub fn with_max_text_bytes(mut self, max_text_bytes: usize) -> Self { + self.max_text_bytes = max_text_bytes; + self + } + + pub fn with_include_messages(mut self, include_messages: bool) -> Self { + self.include_messages = include_messages; + self + } + + pub fn from_env() -> Result> { + let mode_env = env_first(&["A3S_CODE_TRAJECTORY_MODE", "A3S_CODE_RL_TRAJECTORY_MODE"]); + let path_env = env_first(&["A3S_CODE_TRAJECTORY_PATH", "A3S_CODE_RL_TRAJECTORY_PATH"]); + if mode_env.is_none() && path_env.is_none() { + return Ok(None); + } + + let mode = match mode_env.as_deref() { + Some(raw) => RlTrajectoryMode::parse(raw) + .with_context(|| format!("invalid A3S_CODE_RL_TRAJECTORY_MODE: {raw}"))?, + None => RlTrajectoryMode::On, + }; + if mode == RlTrajectoryMode::Off { + return Ok(None); + } + + let path = path_env + .filter(|s| !s.trim().is_empty()) + .with_context(|| "A3S_CODE_TRAJECTORY_PATH is required when trajectory mode is on")?; + + let max_text_bytes = env_first(&[ + "A3S_CODE_TRAJECTORY_MAX_TEXT_BYTES", + "A3S_CODE_RL_TRAJECTORY_MAX_TEXT_BYTES", + ]) + .and_then(|value| value.parse::().ok()) + .unwrap_or_else(|| default_max_text_bytes(mode)); + + let include_messages = env_first(&[ + "A3S_CODE_TRAJECTORY_INCLUDE_MESSAGES", + "A3S_CODE_RL_TRAJECTORY_INCLUDE_MESSAGES", + ]) + .and_then(|value| parse_bool(&value)) + .unwrap_or(true); + + Ok(Some(Self { + mode, + path: PathBuf::from(path), + max_text_bytes, + include_messages, + })) + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct RlTrajectoryContext { + #[serde(skip_serializing_if = "Option::is_none")] + pub run_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub task_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub group_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub replica_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub sample_id: Option, +} + +impl RlTrajectoryContext { + fn from_env() -> Self { + Self { + run_id: env_first(&["A3S_CODE_RL_RUN_ID", "A3S_CODE_RUN_ID", "A3S_RUN_ID"]), + task_id: env_first(&["A3S_CODE_RL_TASK_ID", "A3S_CODE_TASK_ID", "TASK_ID"]), + group_id: env_first(&["A3S_CODE_RL_GROUP_ID", "A3S_CODE_GROUP_ID"]), + replica_id: env_first(&["A3S_CODE_RL_REPLICA_ID", "A3S_CODE_REPLICA_ID"]), + sample_id: env_first(&["A3S_CODE_RL_SAMPLE_ID", "A3S_CODE_SAMPLE_ID"]), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CapturedText { + pub byte_len: usize, + pub sha256: String, + pub truncated: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub text: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub preview: Option, +} + +#[derive(Clone)] +pub struct RlTrajectoryRecorder { + inner: Option>, +} + +struct RlTrajectoryRecorderInner { + config: RlTrajectoryConfig, + context: RlTrajectoryContext, + sequence: AtomicU64, + file: Mutex, +} + +impl std::fmt::Debug for RlTrajectoryRecorder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RlTrajectoryRecorder") + .field("enabled", &self.inner.is_some()) + .finish() + } +} + +impl Default for RlTrajectoryRecorder { + fn default() -> Self { + Self::disabled() + } +} + +impl RlTrajectoryRecorder { + pub fn disabled() -> Self { + Self { inner: None } + } + + pub fn from_config(config: Option) -> Result { + let Some(config) = config else { + return Ok(Self::disabled()); + }; + if config.mode == RlTrajectoryMode::Off { + return Ok(Self::disabled()); + } + + if let Some(parent) = config.path.parent().filter(|p| !p.as_os_str().is_empty()) { + std::fs::create_dir_all(parent).with_context(|| { + format!( + "failed to create RL trajectory directory {}", + parent.display() + ) + })?; + } + let file = OpenOptions::new() + .create(true) + .append(true) + .open(&config.path) + .with_context(|| { + format!( + "failed to open RL trajectory JSONL {}", + config.path.display() + ) + })?; + + Ok(Self { + inner: Some(Arc::new(RlTrajectoryRecorderInner { + config, + context: RlTrajectoryContext::from_env(), + sequence: AtomicU64::new(0), + file: Mutex::new(file), + })), + }) + } + + pub fn is_enabled(&self) -> bool { + 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, + ) { + 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, + }); + inner.record("execution_start", session_id, payload); + } + + pub fn record_llm_request( + &self, + session_id: &str, + turn: usize, + messages: &[Message], + system: Option<&str>, + available_tools: &[String], + estimated_prompt_tokens: usize, + ) { + let Some(inner) = &self.inner else { + return; + }; + let payload = json!({ + "turn": turn, + "messages_count": messages.len(), + "messages": inner.capture_messages(messages), + "system_prompt": system.map(|s| inner.capture_text(s)), + "available_tools": available_tools, + "estimated_prompt_tokens": estimated_prompt_tokens, + }); + inner.record("llm_request", session_id, payload); + } + + pub fn record_llm_response( + &self, + session_id: &str, + turn: usize, + response: &LlmResponse, + duration_ms: u64, + ) { + let Some(inner) = &self.inner else { + return; + }; + let payload = json!({ + "turn": turn, + "message": inner.capture_message(&response.message), + "response_text": inner.capture_text(&response.text()), + "reasoning_content": response.message.reasoning_content.as_ref().map(|s| inner.capture_text(s)), + "tool_calls": response.tool_calls().iter().map(tool_call_value).collect::>(), + "usage": token_usage_value(&response.usage), + "stop_reason": response.stop_reason.clone(), + "meta": response.meta.clone(), + "duration_ms": duration_ms, + }); + inner.record("llm_response", session_id, payload); + } + + pub fn record_tool_call(&self, session_id: &str, turn: usize, tool_call: &ToolCall) { + let Some(inner) = &self.inner else { + return; + }; + let payload = json!({ + "turn": turn, + "tool_call_id": tool_call.id, + "tool": tool_call.name, + "args": tool_call.args, + }); + inner.record("tool_call", session_id, payload); + } + + #[allow(clippy::too_many_arguments)] + pub fn record_tool_result( + &self, + session_id: &str, + turn: usize, + tool_call_id: &str, + tool_name: &str, + output: &str, + exit_code: i32, + duration_ms: u64, + metadata: &Option, + error_kind: Option, + ) { + let Some(inner) = &self.inner else { + return; + }; + let payload = json!({ + "turn": turn, + "tool_call_id": tool_call_id, + "tool": tool_name, + "success": exit_code == 0, + "exit_code": exit_code, + "duration_ms": duration_ms, + "output": inner.capture_text(output), + "metadata": metadata, + "error_kind": error_kind, + }); + inner.record("tool_result", session_id, payload); + } + + pub fn record_context_compacted( + &self, + session_id: &str, + before_messages: usize, + after_messages: &[Message], + percent_before: f32, + ) { + let Some(inner) = &self.inner else { + return; + }; + let payload = json!({ + "before_messages": before_messages, + "after_messages": after_messages.len(), + "percent_before": percent_before, + "messages": inner.capture_messages(after_messages), + }); + inner.record("context_compacted", session_id, payload); + } + + pub fn record_execution_end( + &self, + session_id: &str, + success: bool, + response_text: Option<&str>, + usage: Option<&TokenUsage>, + tool_calls_count: Option, + error_message: Option<&str>, + ) { + let Some(inner) = &self.inner else { + return; + }; + let payload = json!({ + "success": success, + "response_text": response_text.map(|s| inner.capture_text(s)), + "usage": usage.map(token_usage_value), + "tool_calls_count": tool_calls_count, + "error_message": error_message.map(|s| inner.capture_text(s)), + }); + inner.record("execution_end", session_id, payload); + } +} + +impl RlTrajectoryRecorderInner { + fn record(&self, event_type: &str, session_id: &str, payload: Value) { + let sequence = self.sequence.fetch_add(1, Ordering::Relaxed) + 1; + let record = json!({ + "schema": RL_TRAJECTORY_SCHEMA, + "sequence": sequence, + "timestamp_ms": chrono::Utc::now().timestamp_millis(), + "event_type": event_type, + "session_id": session_id, + "mode": self.config.mode, + "context": self.context, + "payload": payload, + }); + + let line = match serde_json::to_string(&record) { + Ok(line) => line, + Err(err) => { + tracing::warn!(error = %err, "Failed to serialize RL trajectory record"); + return; + } + }; + let mut file = match self.file.lock() { + Ok(file) => file, + Err(poisoned) => poisoned.into_inner(), + }; + if let Err(err) = writeln!(file, "{line}") { + tracing::warn!(error = %err, "Failed to write RL trajectory record"); + } + } + + fn capture_messages(&self, messages: &[Message]) -> Value { + if !self.config.include_messages { + return json!({ + "included": false, + "count": messages.len(), + "roles": messages.iter().map(|m| m.role.as_str()).collect::>(), + }); + } + Value::Array( + messages + .iter() + .enumerate() + .map(|(index, message)| { + let mut value = self.capture_message(message); + if let Value::Object(ref mut object) = value { + object.insert("index".to_string(), json!(index)); + } + value + }) + .collect(), + ) + } + + fn capture_message(&self, message: &Message) -> Value { + json!({ + "role": message.role, + "content": message.content.iter().map(|block| self.capture_content_block(block)).collect::>(), + "reasoning_content": message.reasoning_content.as_ref().map(|s| self.capture_text(s)), + }) + } + + fn capture_content_block(&self, block: &ContentBlock) -> Value { + match block { + ContentBlock::Text { text } => json!({ + "type": "text", + "text": self.capture_text(text), + }), + ContentBlock::Image { source } => json!({ + "type": "image", + "source": { + "media_type": source.media_type, + "data": self.capture_text(&source.data), + } + }), + ContentBlock::ToolUse { id, name, input } => json!({ + "type": "tool_use", + "id": id, + "name": name, + "input": input, + }), + ContentBlock::ToolResult { + tool_use_id, + content, + is_error, + } => json!({ + "type": "tool_result", + "tool_use_id": tool_use_id, + "content": self.capture_text(&content.as_text()), + "is_error": is_error, + }), + } + } + + fn capture_text(&self, text: &str) -> CapturedText { + let byte_len = text.len(); + let sha256 = sha256::digest(text); + let (captured, truncated) = truncate_utf8(text, self.config.max_text_bytes); + CapturedText { + byte_len, + sha256, + truncated, + text: Some(captured), + preview: None, + } + } +} + +fn tool_call_value(tool_call: &ToolCall) -> Value { + json!({ + "tool_call_id": tool_call.id, + "tool": tool_call.name, + "args": tool_call.args, + }) +} + +fn token_usage_value(usage: &TokenUsage) -> Value { + json!({ + "prompt_tokens": usage.prompt_tokens, + "completion_tokens": usage.completion_tokens, + "total_tokens": usage.total_tokens, + "cache_read_tokens": usage.cache_read_tokens, + "cache_write_tokens": usage.cache_write_tokens, + }) +} + +fn truncate_utf8(text: &str, max_bytes: usize) -> (String, bool) { + if text.len() <= max_bytes { + return (text.to_string(), false); + } + let mut end = max_bytes.min(text.len()); + while end > 0 && !text.is_char_boundary(end) { + end -= 1; + } + (text[..end].to_string(), true) +} + +fn default_max_text_bytes(mode: RlTrajectoryMode) -> usize { + match mode { + RlTrajectoryMode::Off => 0, + RlTrajectoryMode::On => 1024 * 1024, + } +} + +fn parse_bool(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "1" | "true" | "yes" | "on" => Some(true), + "0" | "false" | "no" | "off" => Some(false), + _ => None, + } +} + +fn env_first(names: &[&str]) -> Option { + names.iter().find_map(|name| { + std::env::var(name) + .ok() + .filter(|value| !value.trim().is_empty()) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[test] + fn rl_recorder_writes_jsonl() { + let dir = tempdir().unwrap(); + let path = dir.path().join("trajectory.jsonl"); + 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_tool_result("sess-1", 1, "tool-1", "bash", "ok", 0, 3, &None, None); + + let lines = std::fs::read_to_string(path).unwrap(); + assert_eq!(lines.lines().count(), 2); + let first: Value = serde_json::from_str(lines.lines().next().unwrap()).unwrap(); + assert_eq!(first["schema"], RL_TRAJECTORY_SCHEMA); + assert_eq!(first["event_type"], "execution_start"); + assert_eq!(first["session_id"], "sess-1"); + } + + #[test] + fn enabled_mode_records_text_with_truncation_flag() { + let dir = tempdir().unwrap(); + let path = dir.path().join("trajectory.jsonl"); + let recorder = RlTrajectoryRecorder::from_config(Some( + RlTrajectoryConfig::new(&path).with_max_text_bytes(3), + )) + .unwrap(); + + recorder.record_execution_start( + "sess-1", + Path::new("/workspace"), + "abcdef", + &[], + None, + 64, + "auto", + ); + + let text = std::fs::read_to_string(path).unwrap(); + let record: Value = serde_json::from_str(text.lines().next().unwrap()).unwrap(); + let prompt = &record["payload"]["prompt"]; + assert!(prompt.get("sha256").is_some()); + assert_eq!(prompt["text"], "abc"); + assert_eq!(prompt["truncated"], true); + } +} diff --git a/sdk/python/src/lib.rs b/sdk/python/src/lib.rs index 1d114c35..1c7e910b 100644 --- a/sdk/python/src/lib.rs +++ b/sdk/python/src/lib.rs @@ -4818,7 +4818,10 @@ impl PyAutoDelegationConfig { fn __repr__(&self) -> String { format!( "AutoDelegationConfig(enabled={}, auto_parallel={}, min_confidence={}, max_tasks={})", - self.enabled, self.auto_parallel, self.min_confidence, self.max_tasks + self.enabled, + self.auto_parallel, + self.min_confidence, + self.max_tasks ) } } @@ -4968,6 +4971,15 @@ struct PySessionOptions { /// long-running cluster sessions to stop in-memory state from /// growing unboundedly. retention_limits: Option, + /// Structured JSONL trajectory path. When set, records user input, + /// LLM turns, tool calls, tool observations, token usage, and episode end. + trajectory_path: Option, + /// Trajectory mode: "on" or "off". Defaults to "on" when trajectory_path is set. + trajectory_mode: Option, + /// Max bytes retained for any single text field before truncation. + trajectory_max_text_bytes: Option, + /// Whether to include full message arrays in LLM request records. + trajectory_include_messages: Option, } impl Clone for PySessionOptions { @@ -5033,6 +5045,10 @@ impl Clone for PySessionOptions { retention_limits: pyo3::Python::with_gil(|py| { self.retention_limits.as_ref().map(|o| o.clone_ref(py)) }), + trajectory_path: self.trajectory_path.clone(), + trajectory_mode: self.trajectory_mode.clone(), + trajectory_max_text_bytes: self.trajectory_max_text_bytes, + trajectory_include_messages: self.trajectory_include_messages, } } } @@ -5088,6 +5104,10 @@ impl PySessionOptions { ahp_transport: None, budget_guard: None, retention_limits: None, + trajectory_path: None, + trajectory_mode: None, + trajectory_max_text_bytes: None, + trajectory_include_messages: None, } } @@ -5649,6 +5669,54 @@ impl PySessionOptions { self.retention_limits = value; } + /// Structured JSONL trajectory output path. + /// + /// When set, a3s-code records user input, LLM turns, tool calls, tool + /// observations, token usage, and execution end status. This is the + /// programmatic equivalent of ``A3S_CODE_TRAJECTORY_PATH``. + #[getter] + fn get_trajectory_path(&self) -> Option { + self.trajectory_path.clone() + } + + #[setter] + fn set_trajectory_path(&mut self, value: Option) { + self.trajectory_path = value; + } + + /// Trajectory mode: ``"on"`` or ``"off"``. + #[getter] + fn get_trajectory_mode(&self) -> Option { + self.trajectory_mode.clone() + } + + #[setter] + fn set_trajectory_mode(&mut self, value: Option) { + self.trajectory_mode = value; + } + + /// Max bytes retained for any single text field before truncation. + #[getter] + fn get_trajectory_max_text_bytes(&self) -> Option { + self.trajectory_max_text_bytes + } + + #[setter] + fn set_trajectory_max_text_bytes(&mut self, value: Option) { + self.trajectory_max_text_bytes = value; + } + + /// Whether LLM request records include full message arrays. + #[getter] + fn get_trajectory_include_messages(&self) -> Option { + self.trajectory_include_messages + } + + #[setter] + fn set_trajectory_include_messages(&mut self, value: Option) { + self.trajectory_include_messages = value; + } + /// Register an instruction skill programmatically. /// /// Instructions are injected into the system prompt at session start. @@ -6111,6 +6179,22 @@ fn build_rust_session_options(so: PySessionOptions) -> PyResult Date: Mon, 6 Jul 2026 22:52:56 +0800 Subject: [PATCH 2/5] feat: capture RL training logprobs --- core/src/agent/llm_turn.rs | 9 +- core/src/agent/tests.rs | 5 + core/src/agent_api.rs | 8 ++ core/src/agent_api/session_config.rs | 31 +++++ core/src/agent_api/session_options.rs | 15 +++ core/src/agent_api/tests.rs | 1 + core/src/llm/anthropic.rs | 2 + core/src/llm/factory.rs | 37 ++++++ core/src/llm/openai.rs | 156 ++++++++++++++++++++++++++ core/src/llm/structured_tests.rs | 3 + core/src/llm/tests.rs | 7 ++ core/src/llm/types.rs | 21 ++++ core/src/llm/zhipu.rs | 10 ++ core/src/planning/llm_planner.rs | 1 + core/src/rl_trajectory.rs | 122 +++++++++++++++++++- core/src/tools/skill.rs | 1 + core/src/tools/task.rs | 1 + sdk/node/generated.d.ts | 21 ++++ sdk/node/src/lib.rs | 73 ++++++++++++ sdk/python/src/lib.rs | 49 ++++++++ 20 files changed, 564 insertions(+), 9 deletions(-) diff --git a/core/src/agent/llm_turn.rs b/core/src/agent/llm_turn.rs index 309c3123..f8262122 100644 --- a/core/src/agent/llm_turn.rs +++ b/core/src/agent/llm_turn.rs @@ -33,17 +33,14 @@ impl AgentLoop { "LLM completion started" ); - let selected_tool_names = - crate::tools::select_tools_for_messages(&self.config.tools, &state.messages) - .into_iter() - .map(|tool| tool.name) - .collect::>(); + let selected_tools = + crate::tools::select_tools_for_messages(&self.config.tools, &state.messages); self.config.rl_trajectory_recorder.record_llm_request( session_id.unwrap_or(""), turn, &state.messages, augmented_system.as_deref(), - &selected_tool_names, + &selected_tools, estimate_prompt_tokens(&state.messages, augmented_system.as_deref()), ); diff --git a/core/src/agent/tests.rs b/core/src/agent/tests.rs index 465c699e..e1abbec8 100644 --- a/core/src/agent/tests.rs +++ b/core/src/agent/tests.rs @@ -352,6 +352,7 @@ impl MockLlmClient { cache_write_tokens: None, }, stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, } } @@ -399,6 +400,7 @@ impl MockLlmClient { cache_write_tokens: None, }, stop_reason: Some("tool_use".to_string()), + token_logprobs: Vec::new(), meta: None, } } @@ -1337,6 +1339,7 @@ async fn test_agent_hitl_multiple_tool_calls() { cache_write_tokens: None, }, stop_reason: Some("tool_use".to_string()), + token_logprobs: Vec::new(), meta: None, }, MockLlmClient::text_response("Both executed!"), @@ -1421,6 +1424,7 @@ async fn test_agent_hitl_partial_approval() { cache_write_tokens: None, }, stop_reason: Some("tool_use".to_string()), + token_logprobs: Vec::new(), meta: None, }, MockLlmClient::text_response("First worked, second rejected."), @@ -2809,6 +2813,7 @@ async fn test_agent_multiple_tools_single_turn() { cache_write_tokens: None, }, stop_reason: Some("tool_use".to_string()), + token_logprobs: Vec::new(), meta: None, }, MockLlmClient::text_response("Both commands ran"), diff --git a/core/src/agent_api.rs b/core/src/agent_api.rs index 79f5319a..ae24e3ca 100644 --- a/core/src/agent_api.rs +++ b/core/src/agent_api.rs @@ -217,6 +217,14 @@ pub struct SessionOptions { /// training or service data collection. If unset, the same config can /// be enabled by `A3S_CODE_TRAJECTORY_PATH`. pub rl_trajectory: Option, + /// Request token-level log probabilities from compatible LLM providers. + /// + /// This is off by default because many public providers reject logprob + /// requests with tool calls. Training/evaluation harnesses using compatible + /// OpenAI-style backends can enable it explicitly. + pub llm_logprobs: Option, + /// Number of alternative token logprobs to request per generated token. + pub llm_top_logprobs: Option, /// Auto-save after each completed `send()` or default-history `stream()` call. pub auto_save: bool, /// Optional artifact retention limits for large tool/program outputs. diff --git a/core/src/agent_api/session_config.rs b/core/src/agent_api/session_config.rs index a6e89379..76a65117 100644 --- a/core/src/agent_api/session_config.rs +++ b/core/src/agent_api/session_config.rs @@ -79,6 +79,22 @@ pub(super) fn resolve_session_llm_client( llm_config = llm_config.with_api_timeout(timeout_ms); } + let logprobs = opts + .llm_logprobs + .or_else(|| env_bool("A3S_CODE_LLM_LOGPROBS")) + .or_else(|| env_bool("A3S_CODE_OPENAI_LOGPROBS")); + if let Some(enabled) = logprobs { + llm_config = llm_config.with_logprobs(enabled); + } + + let top_logprobs = opts + .llm_top_logprobs + .or_else(|| env_usize("A3S_CODE_LLM_TOP_LOGPROBS")) + .or_else(|| env_usize("A3S_CODE_OPENAI_TOP_LOGPROBS")); + if let Some(top_logprobs) = top_logprobs { + llm_config = llm_config.with_top_logprobs(top_logprobs); + } + if let Some(session_id) = session_id { llm_config = llm_config.with_session_id(session_id); } @@ -86,6 +102,21 @@ pub(super) fn resolve_session_llm_client( Ok(crate::llm::create_client_with_config(llm_config)) } +fn env_bool(name: &str) -> Option { + let value = std::env::var(name).ok()?; + match value.trim().to_ascii_lowercase().as_str() { + "1" | "true" | "yes" | "on" => Some(true), + "0" | "false" | "no" | "off" => Some(false), + _ => None, + } +} + +fn env_usize(name: &str) -> Option { + std::env::var(name) + .ok() + .and_then(|value| value.trim().parse::().ok()) +} + pub(super) struct ResolvedSessionMemory { pub(super) memory: Arc, pub(super) init_warning: Option, diff --git a/core/src/agent_api/session_options.rs b/core/src/agent_api/session_options.rs index ce9a1f3b..0f391bb2 100644 --- a/core/src/agent_api/session_options.rs +++ b/core/src/agent_api/session_options.rs @@ -43,6 +43,8 @@ impl std::fmt::Debug for SessionOptions { .field("session_store", &self.session_store.is_some()) .field("session_id", &self.session_id) .field("rl_trajectory", &self.rl_trajectory) + .field("llm_logprobs", &self.llm_logprobs) + .field("llm_top_logprobs", &self.llm_top_logprobs) .field("auto_save", &self.auto_save) .field("artifact_store_limits", &self.artifact_store_limits) .field("max_parse_retries", &self.max_parse_retries) @@ -379,6 +381,19 @@ impl SessionOptions { self } + /// Request token-level log probabilities from compatible LLM providers. + pub fn with_llm_logprobs(mut self, enabled: bool) -> Self { + self.llm_logprobs = Some(enabled); + self + } + + /// Request up to `top_logprobs` alternative logprobs per generated token. + pub fn with_llm_top_logprobs(mut self, top_logprobs: usize) -> Self { + self.llm_logprobs = Some(true); + self.llm_top_logprobs = Some(top_logprobs); + self + } + /// Enable auto-save after each `send()` call pub fn with_auto_save(mut self, enabled: bool) -> Self { self.auto_save = enabled; diff --git a/core/src/agent_api/tests.rs b/core/src/agent_api/tests.rs index 480113ed..4058169c 100644 --- a/core/src/agent_api/tests.rs +++ b/core/src/agent_api/tests.rs @@ -35,6 +35,7 @@ impl StaticStreamingClient { cache_write_tokens: None, }, stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, } } diff --git a/core/src/llm/anthropic.rs b/core/src/llm/anthropic.rs index 5be11b60..2de7deca 100644 --- a/core/src/llm/anthropic.rs +++ b/core/src/llm/anthropic.rs @@ -245,6 +245,7 @@ impl AnthropicClient { cache_write_tokens: parsed.usage.cache_creation_input_tokens, }, stop_reason: Some(parsed.stop_reason), + token_logprobs: Vec::new(), meta: Some(LlmResponseMeta { provider: Some(self.provider_name.clone()), request_model: Some(self.model.clone()), @@ -581,6 +582,7 @@ impl AnthropicClient { }, usage: usage.clone(), stop_reason: stop_reason.clone(), + token_logprobs: Vec::new(), meta: Some(LlmResponseMeta { provider: Some(provider_name.clone()), request_model: Some(request_model.clone()), diff --git a/core/src/llm/factory.rs b/core/src/llm/factory.rs index a44db602..40d137df 100644 --- a/core/src/llm/factory.rs +++ b/core/src/llm/factory.rs @@ -30,6 +30,10 @@ pub struct LlmConfig { pub max_tokens: Option, /// Extended thinking budget in tokens (Anthropic only). pub thinking_budget: Option, + /// Request token-level log probabilities from OpenAI-compatible providers. + pub logprobs: Option, + /// Number of alternative logprobs per token when logprobs are requested. + pub top_logprobs: Option, /// When true, temperature is never sent to the API (e.g., o1 models). pub disable_temperature: bool, } @@ -52,6 +56,8 @@ impl std::fmt::Debug for LlmConfig { .field("temperature", &self.temperature) .field("max_tokens", &self.max_tokens) .field("thinking_budget", &self.thinking_budget) + .field("logprobs", &self.logprobs) + .field("top_logprobs", &self.top_logprobs) .field("disable_temperature", &self.disable_temperature) .finish() } @@ -76,6 +82,8 @@ impl LlmConfig { temperature: None, max_tokens: None, thinking_budget: None, + logprobs: None, + top_logprobs: None, disable_temperature: false, } } @@ -125,6 +133,17 @@ impl LlmConfig { self } + pub fn with_logprobs(mut self, enabled: bool) -> Self { + self.logprobs = Some(enabled); + self + } + + pub fn with_top_logprobs(mut self, top_logprobs: usize) -> Self { + self.logprobs = Some(true); + self.top_logprobs = Some(top_logprobs); + self + } + pub(crate) fn resolved_headers(&self) -> HashMap { let mut headers = self.headers.clone(); if let (Some(header_name), Some(session_id)) = (&self.session_id_header, &self.session_id) { @@ -192,6 +211,12 @@ pub fn create_client_with_config(config: LlmConfig) -> Arc { if let Some(max) = config.max_tokens { client = client.with_max_tokens(max); } + if let Some(enabled) = config.logprobs { + client = client.with_logprobs(enabled); + } + if let Some(top_logprobs) = config.top_logprobs { + client = client.with_top_logprobs(top_logprobs); + } Arc::new(client) } "glm" | "zhipu" | "bigmodel" => { @@ -210,6 +235,12 @@ pub fn create_client_with_config(config: LlmConfig) -> Arc { if let Some(max) = config.max_tokens { client = client.with_max_tokens(max); } + if let Some(enabled) = config.logprobs { + client = client.with_logprobs(enabled); + } + if let Some(top_logprobs) = config.top_logprobs { + client = client.with_top_logprobs(top_logprobs); + } Arc::new(client) } // OpenAI-compatible providers (deepseek, groq, together, ollama, etc.) @@ -238,6 +269,12 @@ pub fn create_client_with_config(config: LlmConfig) -> Arc { if let Some(max) = config.max_tokens { client = client.with_max_tokens(max); } + if let Some(enabled) = config.logprobs { + client = client.with_logprobs(enabled); + } + if let Some(top_logprobs) = config.top_logprobs { + client = client.with_top_logprobs(top_logprobs); + } Arc::new(client) } } diff --git a/core/src/llm/openai.rs b/core/src/llm/openai.rs index 822f949a..a6ffc5ae 100644 --- a/core/src/llm/openai.rs +++ b/core/src/llm/openai.rs @@ -25,6 +25,8 @@ pub struct OpenAiClient { pub(crate) headers: HashMap, pub(crate) temperature: Option, pub(crate) max_tokens: Option, + pub(crate) logprobs: bool, + pub(crate) top_logprobs: Option, pub(crate) http: Arc, pub(crate) retry_config: RetryConfig, } @@ -92,6 +94,8 @@ impl OpenAiClient { headers: HashMap::new(), temperature: None, max_tokens: None, + logprobs: false, + top_logprobs: None, http: default_http_client(), retry_config: RetryConfig::default(), } @@ -132,6 +136,17 @@ impl OpenAiClient { self } + pub fn with_logprobs(mut self, enabled: bool) -> Self { + self.logprobs = enabled; + self + } + + pub fn with_top_logprobs(mut self, top_logprobs: usize) -> Self { + self.logprobs = true; + self.top_logprobs = Some(top_logprobs); + self + } + pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self { self.retry_config = retry_config; self @@ -341,6 +356,12 @@ impl OpenAiClient { if let Some(max) = self.max_tokens { request["max_tokens"] = serde_json::json!(max); } + if self.logprobs { + request["logprobs"] = serde_json::json!(true); + if let Some(top_logprobs) = self.top_logprobs { + request["top_logprobs"] = serde_json::json!(top_logprobs); + } + } if !tools.is_empty() { request["tools"] = serde_json::json!(self.convert_tools(tools)); @@ -406,6 +427,11 @@ impl OpenAiClient { serde_json::from_str(&response).context("Failed to parse OpenAI response")?; let choice = parsed.choices.into_iter().next().context("No choices")?; + let token_logprobs = choice + .logprobs + .as_ref() + .map(openai_logprobs_to_token_logprobs) + .unwrap_or_default(); let mut content = vec![]; @@ -458,6 +484,7 @@ impl OpenAiClient { cache_write_tokens: None, }, stop_reason: choice.finish_reason, + token_logprobs, meta: Some(LlmResponseMeta { provider: Some(self.provider_name.clone()), request_model: Some(self.model.clone()), @@ -637,6 +664,7 @@ impl OpenAiClient { std::collections::BTreeMap::new(); let mut usage = TokenUsage::default(); let mut finish_reason = None; + let mut token_logprobs: Vec = Vec::new(); let mut response_id = None; let mut response_model = None; let mut response_object = None; @@ -695,6 +723,7 @@ impl OpenAiClient { }, usage: usage.clone(), stop_reason: std::mem::take(&mut finish_reason), + token_logprobs: std::mem::take(&mut token_logprobs), meta: Some(LlmResponseMeta { provider: Some(provider_name.clone()), request_model: Some(request_model.clone()), @@ -738,6 +767,11 @@ impl OpenAiClient { } if let Some(choice) = event.choices.into_iter().next() { + if let Some(logprobs) = choice.logprobs.as_ref() { + token_logprobs.extend( + openai_logprobs_to_token_logprobs(logprobs), + ); + } if let Some(reason) = choice.finish_reason { finish_reason = Some(reason); } @@ -922,6 +956,9 @@ impl OpenAiClient { .and_then(|d| d.cached_tokens); } if let Some(choice) = event.choices.into_iter().next() { + if let Some(logprobs) = choice.logprobs.as_ref() { + token_logprobs.extend(openai_logprobs_to_token_logprobs(logprobs)); + } if let Some(reason) = choice.finish_reason { finish_reason = Some(reason); } @@ -1013,6 +1050,9 @@ impl OpenAiClient { .and_then(|d| d.cached_tokens); if let Some(choice) = response.choices.into_iter().next() { + if let Some(logprobs) = choice.logprobs.as_ref() { + token_logprobs.extend(openai_logprobs_to_token_logprobs(logprobs)); + } finish_reason = choice.finish_reason; if let Some(text) = choice.message.content.filter(|text| !text.is_empty()) @@ -1079,6 +1119,7 @@ impl OpenAiClient { }, usage: usage.clone(), stop_reason: std::mem::take(&mut finish_reason), + token_logprobs: std::mem::take(&mut token_logprobs), meta: Some(LlmResponseMeta { provider: Some(provider_name.clone()), request_model: Some(request_model.clone()), @@ -1106,6 +1147,32 @@ impl OpenAiClient { } } +fn openai_logprobs_to_token_logprobs(logprobs: &OpenAiChoiceLogprobs) -> Vec { + logprobs + .content + .as_ref() + .map(|items| { + items + .iter() + .map(|item| TokenLogProb { + token: item.token.clone(), + logprob: item.logprob, + bytes: item.bytes.clone(), + top_logprobs: item + .top_logprobs + .iter() + .map(|top| TopTokenLogProb { + token: top.token.clone(), + logprob: top.logprob, + bytes: top.bytes.clone(), + }) + .collect(), + }) + .collect() + }) + .unwrap_or_default() +} + // OpenAI API response types (private) #[derive(Debug, Deserialize)] pub(crate) struct OpenAiResponse { @@ -1123,6 +1190,32 @@ pub(crate) struct OpenAiResponse { pub(crate) struct OpenAiChoice { pub(crate) message: OpenAiMessage, pub(crate) finish_reason: Option, + #[serde(default)] + pub(crate) logprobs: Option, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct OpenAiChoiceLogprobs { + #[serde(default)] + pub(crate) content: Option>, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct OpenAiTokenLogprob { + pub(crate) token: String, + pub(crate) logprob: f64, + #[serde(default)] + pub(crate) bytes: Option>, + #[serde(default)] + pub(crate) top_logprobs: Vec, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct OpenAiTopLogprob { + pub(crate) token: String, + pub(crate) logprob: f64, + #[serde(default)] + pub(crate) bytes: Option>, } #[derive(Debug, Deserialize)] @@ -1189,6 +1282,8 @@ pub(crate) struct OpenAiStreamChoice { pub(crate) message: Option, pub(crate) delta: Option, pub(crate) finish_reason: Option, + #[serde(default)] + pub(crate) logprobs: Option, } #[derive(Debug, Deserialize)] @@ -1338,6 +1433,28 @@ mod tests { assert!(resp.message.tool_calls().is_empty()); } + #[tokio::test] + async fn streaming_collects_token_logprobs() { + let chunks = vec![ + "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"logprobs\":{\"content\":[{\"token\":\"hello\",\"logprob\":-0.2,\"bytes\":[104,101,108,108,111],\"top_logprobs\":[{\"token\":\"hi\",\"logprob\":-1.2,\"bytes\":[104,105]}]}]}}]}\n\n" + .to_string(), + "data: {\"choices\":[{\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1,\"total_tokens\":2}}\n\n" + .to_string(), + "data: [DONE]\n\n".to_string(), + ]; + let resp = drain_to_done(&glm_client(chunks).with_logprobs(true)).await; + assert_eq!(resp.text(), "hello"); + assert_eq!(resp.token_logprobs.len(), 1); + assert_eq!(resp.token_logprobs[0].token, "hello"); + assert_eq!(resp.token_logprobs[0].logprob, -0.2); + assert_eq!( + resp.token_logprobs[0].bytes.as_deref(), + Some(&[104, 101, 108, 108, 111][..]) + ); + assert_eq!(resp.token_logprobs[0].top_logprobs[0].token, "hi"); + assert_eq!(resp.token_logprobs[0].top_logprobs[0].logprob, -1.2); + } + #[test] fn test_apply_directive_forced_function_tool_choice() { let mut req = serde_json::json!({ "model": "m" }); @@ -1410,6 +1527,45 @@ mod tests { let req = make_client().build_chat_request(&[Message::user("hi")], None, &[], None); assert!(req.get("tool_choice").is_none()); assert!(req.get("response_format").is_none()); + assert!(req.get("logprobs").is_none()); + assert!(req.get("top_logprobs").is_none()); + } + + #[test] + fn test_build_chat_request_includes_logprob_options_when_enabled() { + let req = make_client().with_top_logprobs(1).build_chat_request( + &[Message::user("hi")], + None, + &[], + None, + ); + assert_eq!(req["logprobs"], true); + assert_eq!(req["top_logprobs"], 1); + } + + #[test] + fn test_parse_openai_token_logprobs() { + let parsed = openai_logprobs_to_token_logprobs(&OpenAiChoiceLogprobs { + content: Some(vec![OpenAiTokenLogprob { + token: "hello".to_string(), + logprob: -0.25, + bytes: Some(vec![104, 101, 108, 108, 111]), + top_logprobs: vec![OpenAiTopLogprob { + token: "hi".to_string(), + logprob: -1.5, + bytes: Some(vec![104, 105]), + }], + }]), + }); + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].token, "hello"); + assert_eq!(parsed[0].logprob, -0.25); + assert_eq!( + parsed[0].bytes.as_deref(), + Some(&[104, 101, 108, 108, 111][..]) + ); + assert_eq!(parsed[0].top_logprobs[0].token, "hi"); + assert_eq!(parsed[0].top_logprobs[0].logprob, -1.5); } #[test] diff --git a/core/src/llm/structured_tests.rs b/core/src/llm/structured_tests.rs index 692e36e5..a2010b00 100644 --- a/core/src/llm/structured_tests.rs +++ b/core/src/llm/structured_tests.rs @@ -35,6 +35,7 @@ impl MockStructuredClient { cache_write_tokens: None, }, stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, } } @@ -58,6 +59,7 @@ impl MockStructuredClient { cache_write_tokens: None, }, stop_reason: Some("tool_use".to_string()), + token_logprobs: Vec::new(), meta: None, } } @@ -86,6 +88,7 @@ impl MockStructuredClient { cache_write_tokens: None, }, stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, } } diff --git a/core/src/llm/tests.rs b/core/src/llm/tests.rs index 0498b843..346d4f6b 100644 --- a/core/src/llm/tests.rs +++ b/core/src/llm/tests.rs @@ -330,6 +330,7 @@ mod tests { cache_write_tokens: None, }, stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, }; assert_eq!(response.text(), "Hello!"); @@ -352,6 +353,7 @@ mod tests { }, usage: TokenUsage::default(), stop_reason: Some("tool_use".to_string()), + token_logprobs: Vec::new(), meta: None, }; let calls = response.tool_calls(); @@ -567,6 +569,7 @@ mod extra_llm_tests { }, usage: TokenUsage::default(), stop_reason: None, + token_logprobs: Vec::new(), meta: None, }; assert_eq!(r.text(), "resp"); @@ -586,6 +589,7 @@ mod extra_llm_tests { }, usage: TokenUsage::default(), stop_reason: None, + token_logprobs: Vec::new(), meta: None, }; assert_eq!(r.tool_calls().len(), 1); @@ -1708,6 +1712,7 @@ mod extra_llm_tests2 { }, usage: TokenUsage::default(), stop_reason: None, + token_logprobs: Vec::new(), meta: None, }; assert_eq!(response.text(), "response text"); @@ -1727,6 +1732,7 @@ mod extra_llm_tests2 { }, usage: TokenUsage::default(), stop_reason: Some("tool_use".to_string()), + token_logprobs: Vec::new(), meta: None, }; let calls = response.tool_calls(); @@ -2140,6 +2146,7 @@ mod extra_llm_tests2 { message: Message::user("test"), usage: TokenUsage::default(), stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, }; let json = serde_json::to_string(&response).unwrap(); diff --git a/core/src/llm/types.rs b/core/src/llm/types.rs index b5dfa3dc..e406f8c5 100644 --- a/core/src/llm/types.rs +++ b/core/src/llm/types.rs @@ -394,6 +394,8 @@ pub struct LlmResponse { pub message: Message, pub usage: TokenUsage, pub stop_reason: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub token_logprobs: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub meta: Option, } @@ -420,6 +422,25 @@ pub struct TokenUsage { pub cache_write_tokens: Option, } +/// Token-level log probability emitted by an OpenAI-compatible backend. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TokenLogProb { + pub token: String, + pub logprob: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub bytes: Option>, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub top_logprobs: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TopTokenLogProb { + pub token: String, + pub logprob: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub bytes: Option>, +} + /// Tool call from LLM #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolCall { diff --git a/core/src/llm/zhipu.rs b/core/src/llm/zhipu.rs index 6808ad81..3ece1cf9 100644 --- a/core/src/llm/zhipu.rs +++ b/core/src/llm/zhipu.rs @@ -40,6 +40,16 @@ impl ZhipuClient { self } + pub fn with_logprobs(mut self, enabled: bool) -> Self { + self.0 = self.0.with_logprobs(enabled); + self + } + + pub fn with_top_logprobs(mut self, top_logprobs: usize) -> Self { + self.0 = self.0.with_top_logprobs(top_logprobs); + self + } + pub fn with_base_url(mut self, base_url: String) -> Self { self.0 = self.0.with_base_url(base_url); self diff --git a/core/src/planning/llm_planner.rs b/core/src/planning/llm_planner.rs index e1a4669c..a8433e6c 100644 --- a/core/src/planning/llm_planner.rs +++ b/core/src/planning/llm_planner.rs @@ -616,6 +616,7 @@ mod tests { }, usage: crate::llm::TokenUsage::default(), stop_reason: None, + token_logprobs: Vec::new(), meta: None, }) } diff --git a/core/src/rl_trajectory.rs b/core/src/rl_trajectory.rs index 77c062fb..c6be837e 100644 --- a/core/src/rl_trajectory.rs +++ b/core/src/rl_trajectory.rs @@ -5,7 +5,10 @@ //! turns, tool calls, tool observations, token usage, and termination reasons. //! It is opt-in: normal sessions pay only a cheap disabled-recorder branch. -use crate::llm::{ContentBlock, LlmResponse, Message, TokenUsage, ToolCall}; +use crate::llm::{ + ContentBlock, LlmResponse, Message, TokenLogProb, TokenUsage, ToolCall, ToolDefinition, + TopTokenLogProb, +}; use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; @@ -256,18 +259,23 @@ impl RlTrajectoryRecorder { turn: usize, messages: &[Message], system: Option<&str>, - available_tools: &[String], + available_tools: &[ToolDefinition], estimated_prompt_tokens: usize, ) { let Some(inner) = &self.inner else { return; }; + let available_tool_names = available_tools + .iter() + .map(|tool| tool.name.as_str()) + .collect::>(); let payload = json!({ "turn": turn, "messages_count": messages.len(), "messages": inner.capture_messages(messages), "system_prompt": system.map(|s| inner.capture_text(s)), - "available_tools": available_tools, + "available_tools": available_tool_names, + "tool_definitions": available_tools.iter().map(tool_definition_value).collect::>(), "estimated_prompt_tokens": estimated_prompt_tokens, }); inner.record("llm_request", session_id, payload); @@ -289,6 +297,7 @@ impl RlTrajectoryRecorder { "response_text": inner.capture_text(&response.text()), "reasoning_content": response.message.reasoning_content.as_ref().map(|s| inner.capture_text(s)), "tool_calls": response.tool_calls().iter().map(tool_call_value).collect::>(), + "token_logprobs": response.token_logprobs.iter().map(token_logprob_value).collect::>(), "usage": token_usage_value(&response.usage), "stop_reason": response.stop_reason.clone(), "meta": response.meta.clone(), @@ -497,6 +506,31 @@ fn tool_call_value(tool_call: &ToolCall) -> Value { }) } +fn tool_definition_value(tool: &ToolDefinition) -> Value { + json!({ + "name": tool.name, + "description": tool.description, + "parameters": tool.parameters, + }) +} + +fn token_logprob_value(token: &TokenLogProb) -> Value { + json!({ + "token": token.token, + "logprob": token.logprob, + "bytes": token.bytes, + "top_logprobs": token.top_logprobs.iter().map(top_token_logprob_value).collect::>(), + }) +} + +fn top_token_logprob_value(token: &TopTokenLogProb) -> Value { + json!({ + "token": token.token, + "logprob": token.logprob, + "bytes": token.bytes, + }) +} + fn token_usage_value(usage: &TokenUsage) -> Value { json!({ "prompt_tokens": usage.prompt_tokens, @@ -598,4 +632,86 @@ mod tests { assert_eq!(prompt["text"], "abc"); assert_eq!(prompt["truncated"], true); } + + #[test] + fn llm_events_include_tool_definitions_and_token_logprobs() { + let dir = tempdir().unwrap(); + let path = dir.path().join("trajectory.jsonl"); + let recorder = + RlTrajectoryRecorder::from_config(Some(RlTrajectoryConfig::new(&path))).unwrap(); + + let tools = vec![ToolDefinition { + name: "bash".to_string(), + description: "Run a shell command".to_string(), + parameters: json!({ + "type": "object", + "properties": { + "cmd": { "type": "string" } + }, + "required": ["cmd"] + }), + }]; + recorder.record_llm_request("sess-1", 1, &[Message::user("hi")], None, &tools, 7); + + recorder.record_llm_response( + "sess-1", + 1, + &LlmResponse { + message: Message { + role: "assistant".to_string(), + content: vec![ContentBlock::Text { + text: "hello".to_string(), + }], + reasoning_content: None, + }, + usage: TokenUsage { + prompt_tokens: 7, + completion_tokens: 1, + total_tokens: 8, + cache_read_tokens: None, + cache_write_tokens: None, + }, + stop_reason: Some("stop".to_string()), + token_logprobs: vec![TokenLogProb { + token: "hello".to_string(), + logprob: -0.2, + bytes: Some(vec![104, 101, 108, 108, 111]), + top_logprobs: vec![TopTokenLogProb { + token: "hi".to_string(), + logprob: -1.2, + bytes: Some(vec![104, 105]), + }], + }], + meta: None, + }, + 42, + ); + + let lines = std::fs::read_to_string(path).unwrap(); + let records = lines + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .collect::>(); + let request = records + .iter() + .find(|record| record["event_type"] == "llm_request") + .unwrap(); + assert_eq!(request["payload"]["available_tools"][0], "bash"); + assert_eq!(request["payload"]["tool_definitions"][0]["name"], "bash"); + assert_eq!( + request["payload"]["tool_definitions"][0]["parameters"]["required"][0], + "cmd" + ); + + let response = records + .iter() + .find(|record| record["event_type"] == "llm_response") + .unwrap(); + assert_eq!(response["payload"]["token_logprobs"][0]["token"], "hello"); + assert_eq!(response["payload"]["token_logprobs"][0]["logprob"], -0.2); + assert_eq!( + response["payload"]["token_logprobs"][0]["top_logprobs"][0]["token"], + "hi" + ); + } } diff --git a/core/src/tools/skill.rs b/core/src/tools/skill.rs index 43501f66..6fb92230 100644 --- a/core/src/tools/skill.rs +++ b/core/src/tools/skill.rs @@ -405,6 +405,7 @@ mod tests { cache_write_tokens: None, }, stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, } } diff --git a/core/src/tools/task.rs b/core/src/tools/task.rs index f59c6041..cceff579 100644 --- a/core/src/tools/task.rs +++ b/core/src/tools/task.rs @@ -1580,6 +1580,7 @@ mod tests { cache_write_tokens: None, }, stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, } } diff --git a/sdk/node/generated.d.ts b/sdk/node/generated.d.ts index 89bac4a7..6638646b 100644 --- a/sdk/node/generated.d.ts +++ b/sdk/node/generated.d.ts @@ -565,6 +565,27 @@ export interface SessionOptions { * Only applied when `model` is also set. Provider must support extended thinking. */ thinkingBudget?: number + /** + * Request token-level log probabilities from OpenAI-compatible backends. + * + * Providers that do not support logprobs may reject the request. + */ + llmLogprobs?: boolean + /** Number of top token log probabilities to request when logprobs are enabled. */ + llmTopLogprobs?: number + /** + * Structured JSONL trajectory path. + * + * When set, records user input, LLM turns, tool calls, tool observations, + * token usage, and execution end status. + */ + trajectoryPath?: string + /** Trajectory mode: "on" or "off". Defaults to "on" when `trajectoryPath` is set. */ + trajectoryMode?: string + /** Max bytes retained for any single text field before truncation. */ + trajectoryMaxTextBytes?: number + /** Whether LLM request records include full message arrays. */ + trajectoryIncludeMessages?: boolean /** * Enable continuation injection (default: true). * When enabled, the loop injects a follow-up prompt when the LLM stops without completing. diff --git a/sdk/node/src/lib.rs b/sdk/node/src/lib.rs index 582cb13d..8d8c38ab 100644 --- a/sdk/node/src/lib.rs +++ b/sdk/node/src/lib.rs @@ -1908,6 +1908,23 @@ pub struct SessionOptions { /// Extended thinking token budget (e.g. 10_000). Enables chain-of-thought reasoning. /// Only applied when `model` is also set. Provider must support extended thinking. pub thinking_budget: Option, + /// Request token-level log probabilities from OpenAI-compatible backends. + /// + /// Providers that do not support logprobs may reject the request. + pub llm_logprobs: Option, + /// Number of top token log probabilities to request when logprobs are enabled. + pub llm_top_logprobs: Option, + /// Structured JSONL trajectory path. + /// + /// When set, records user input, LLM turns, tool calls, tool observations, + /// token usage, and execution end status. + pub trajectory_path: Option, + /// Trajectory mode: "on" or "off". Defaults to "on" when `trajectoryPath` is set. + pub trajectory_mode: Option, + /// Max bytes retained for any single text field before truncation. + pub trajectory_max_text_bytes: Option, + /// Whether LLM request records include full message arrays. + pub trajectory_include_messages: Option, /// Enable continuation injection (default: true). /// When enabled, the loop injects a follow-up prompt when the LLM stops without completing. pub continuation_enabled: Option, @@ -2505,6 +2522,28 @@ fn js_session_options_to_rust(options: Option) -> napi::Result, + /// Request token-level log probabilities from OpenAI-compatible backends. + llm_logprobs: Option, + /// Number of top token log probabilities to request when logprobs are enabled. + llm_top_logprobs: Option, /// Enable continuation injection (default: True). /// When enabled, the loop injects a follow-up prompt when the LLM stops without completing. continuation_enabled: Option, @@ -5027,6 +5031,8 @@ impl Clone for PySessionOptions { circuit_breaker_threshold: self.circuit_breaker_threshold, temperature: self.temperature, thinking_budget: self.thinking_budget, + llm_logprobs: self.llm_logprobs, + llm_top_logprobs: self.llm_top_logprobs, continuation_enabled: self.continuation_enabled, max_continuation_turns: self.max_continuation_turns, max_execution_time_ms: self.max_execution_time_ms, @@ -5092,6 +5098,8 @@ impl PySessionOptions { circuit_breaker_threshold: None, temperature: None, thinking_budget: None, + llm_logprobs: None, + llm_top_logprobs: None, continuation_enabled: None, max_continuation_turns: None, max_execution_time_ms: None, @@ -5529,6 +5537,30 @@ impl PySessionOptions { self.thinking_budget = value; } + /// Request token-level log probabilities from OpenAI-compatible backends. + /// + /// Providers that do not support logprobs may reject the request. + #[getter] + fn get_llm_logprobs(&self) -> Option { + self.llm_logprobs + } + + #[setter] + fn set_llm_logprobs(&mut self, value: Option) { + self.llm_logprobs = value; + } + + /// Number of top token log probabilities to request. + #[getter] + fn get_llm_top_logprobs(&self) -> Option { + self.llm_top_logprobs + } + + #[setter] + fn set_llm_top_logprobs(&mut self, value: Option) { + self.llm_top_logprobs = value; + } + /// Enable or disable continuation injection (default: True). #[getter] fn get_continuation_enabled(&self) -> Option { @@ -6145,6 +6177,12 @@ fn build_rust_session_options(so: PySessionOptions) -> PyResult Date: Mon, 6 Jul 2026 23:24:44 +0800 Subject: [PATCH 3/5] refactor: streamline RL trajectory recording --- core/src/agent/llm_turn.rs | 73 ++++++++++++++++++++-------------- core/src/agent/loop_runtime.rs | 16 ++++---- core/src/agent_api/tests.rs | 18 ++++----- core/src/rl_trajectory.rs | 73 +++++++++++++++++----------------- 4 files changed, 98 insertions(+), 82 deletions(-) diff --git a/core/src/agent/llm_turn.rs b/core/src/agent/llm_turn.rs index f8262122..725197e9 100644 --- a/core/src/agent/llm_turn.rs +++ b/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, } +struct LlmCallRequest<'a> { + turn: usize, + messages: &'a [Message], + system: Option<&'a str>, + tools: &'a [ToolDefinition], + session_id: Option<&'a str>, + event_tx: &'a Option>, + 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>, - cancel_token: &tokio_util::sync::CancellationToken, + request: LlmCallRequest<'_>, ) -> anyhow::Result { // 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,7 +379,7 @@ 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. @@ -373,15 +387,14 @@ impl AgentLoop { &self, messages: &[Message], system: Option<&str>, + tools: &[ToolDefinition], event_tx: &Option>, cancel_token: &tokio_util::sync::CancellationToken, ) -> anyhow::Result { - 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") } diff --git a/core/src/agent/loop_runtime.rs b/core/src/agent/loop_runtime.rs index 913a7d28..e0074a71 100644 --- a/core/src/agent/loop_runtime.rs +++ b/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 diff --git a/core/src/agent_api/tests.rs b/core/src/agent_api/tests.rs index 4058169c..af15237e 100644 --- a/core/src/agent_api/tests.rs +++ b/core/src/agent_api/tests.rs @@ -2927,15 +2927,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(); diff --git a/core/src/rl_trajectory.rs b/core/src/rl_trajectory.rs index c6be837e..c9883e37 100644 --- a/core/src/rl_trajectory.rs +++ b/core/src/rl_trajectory.rs @@ -161,6 +161,16 @@ pub struct RlTrajectoryRecorder { inner: Option>, } +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(); From 3540577cc13544db3ebe662edf814c7fe0de819c Mon Sep 17 00:00:00 2001 From: Li Xiangtian <1193027052@qq.com> Date: Tue, 7 Jul 2026 00:35:46 +0800 Subject: [PATCH 4/5] fix: align RL capture branch with A3S main --- Cargo.lock | 42 ++++++++++----- core/Cargo.toml | 2 +- core/src/agent/memory_extraction_runtime.rs | 23 +++++++- core/src/agent/tests.rs | 1 + core/src/agent_api/tests.rs | 2 + core/src/memory.rs | 2 +- core/tests/test_workspace_manifest_agent.rs | 2 + sdk/node/Cargo.lock | 42 ++++++++++----- sdk/python/Cargo.lock | 58 ++++++++++++++++----- 9 files changed, 129 insertions(+), 45 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d312500e..e067d366 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -59,7 +59,7 @@ dependencies = [ "bytes", "cfb", "chrono", - "cron", + "cron 0.12.1", "dirs", "flate2", "futures", @@ -137,6 +137,7 @@ version = "0.4.0" dependencies = [ "async-trait", "chrono", + "cron 0.17.0", "dashmap", "futures-core", "num_cpus", @@ -153,7 +154,7 @@ dependencies = [ [[package]] name = "a3s-memory" -version = "0.1.2" +version = "0.1.1" dependencies = [ "anyhow", "async-trait", @@ -296,7 +297,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -307,7 +308,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1360,6 +1361,19 @@ dependencies = [ "once_cell", ] +[[package]] +name = "cron" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5dcd6f69605c2956916ce24e8af637b754964c9a83f4662d3a2361654cdba09" +dependencies = [ + "chrono", + "once_cell", + "phf 0.11.3", + "serde", + "winnow", +] + [[package]] name = "crossbeam-channel" version = "0.5.15" @@ -1707,7 +1721,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2407,7 +2421,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -2968,7 +2982,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3467,7 +3481,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.40", - "socket2 0.5.10", + "socket2 0.6.4", "thiserror 2.0.18", "tokio", "tracing", @@ -3504,7 +3518,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.4", "tracing", "windows-sys 0.60.2", ] @@ -3864,7 +3878,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4305,7 +4319,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -4469,10 +4483,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5336,7 +5350,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/core/Cargo.toml b/core/Cargo.toml index 800b3abd..0324a5b4 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -14,7 +14,7 @@ path = "src/lib.rs" [dependencies] # Internal crates a3s-common = { version = "0.1.1", path = "../../common" } -a3s-memory = { version = "0.1.2", path = "../../memory" } +a3s-memory = { version = "0.1.1", path = "../../memory" } a3s-lane = { version = "0.4", path = "../../lane" } a3s-search = { version = "1.2.3", path = "../../search", default-features = false, features = ["lightpanda"] } a3s-flow = { version = "0.4.1", path = "../../flow" } diff --git a/core/src/agent/memory_extraction_runtime.rs b/core/src/agent/memory_extraction_runtime.rs index d87bc8b5..dcb3c3a3 100644 --- a/core/src/agent/memory_extraction_runtime.rs +++ b/core/src/agent/memory_extraction_runtime.rs @@ -182,7 +182,7 @@ impl AgentLoop { let item = if let Some(existing) = similar_existing_memory(&memory, &item, &supersedes, &conflicts_with).await { - existing.merge_duplicate(item) + merge_duplicate_memory(existing, item) } else { item }; @@ -629,6 +629,27 @@ async fn delete_superseded_memories(memory: &Arc, supersedes: &[Str } } +fn merge_duplicate_memory(mut existing: MemoryItem, mut extracted: MemoryItem) -> MemoryItem { + extracted.id = existing.id; + extracted.timestamp = existing.timestamp; + extracted.importance = existing.importance.max(extracted.importance); + extracted.access_count = existing.access_count; + extracted.last_accessed = existing.last_accessed; + + for tag in existing.tags.drain(..) { + if !extracted.tags.contains(&tag) { + extracted.tags.push(tag); + } + } + + for (key, value) in existing.metadata.drain() { + extracted.metadata.entry(key).or_insert(value); + } + + extracted.content_lower = extracted.content.to_lowercase(); + extracted +} + async fn similar_existing_memory( memory: &Arc, item: &MemoryItem, diff --git a/core/src/agent/tests.rs b/core/src/agent/tests.rs index e1abbec8..1d161b95 100644 --- a/core/src/agent/tests.rs +++ b/core/src/agent/tests.rs @@ -372,6 +372,7 @@ impl MockLlmClient { cache_write_tokens: None, }, stop_reason: Some("stop".to_string()), + token_logprobs: Vec::new(), meta: None, } } diff --git a/core/src/agent_api/tests.rs b/core/src/agent_api/tests.rs index af15237e..9b060b1f 100644 --- a/core/src/agent_api/tests.rs +++ b/core/src/agent_api/tests.rs @@ -75,6 +75,7 @@ fn scripted_text_response(text: &str) -> LlmResponse { cache_write_tokens: None, }, stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, } } @@ -102,6 +103,7 @@ fn scripted_tool_call_response( cache_write_tokens: None, }, stop_reason: Some("tool_use".to_string()), + token_logprobs: Vec::new(), meta: None, } } diff --git a/core/src/memory.rs b/core/src/memory.rs index e06cf51b..7b6d76ab 100644 --- a/core/src/memory.rs +++ b/core/src/memory.rs @@ -208,7 +208,7 @@ impl AgentMemory { /// Store a memory and return the normalized item that was sent to storage. pub async fn remember_item(&self, item: MemoryItem) -> anyhow::Result { - let item = self.store.store_and_return(item).await?; + self.store.store(item.clone()).await?; let mut short_term = self.short_term.write().await; if let Some(existing) = short_term .iter_mut() diff --git a/core/tests/test_workspace_manifest_agent.rs b/core/tests/test_workspace_manifest_agent.rs index bcb6030e..46f4623d 100644 --- a/core/tests/test_workspace_manifest_agent.rs +++ b/core/tests/test_workspace_manifest_agent.rs @@ -107,6 +107,7 @@ fn text_response(text: &str) -> LlmResponse { cache_write_tokens: None, }, stop_reason: Some("end_turn".to_string()), + token_logprobs: Vec::new(), meta: None, } } @@ -130,6 +131,7 @@ fn tool_response(tool_id: &str, tool_name: &str, input: serde_json::Value) -> Ll cache_write_tokens: None, }, stop_reason: Some("tool_use".to_string()), + token_logprobs: Vec::new(), meta: None, } } diff --git a/sdk/node/Cargo.lock b/sdk/node/Cargo.lock index 71311f1b..7e400905 100644 --- a/sdk/node/Cargo.lock +++ b/sdk/node/Cargo.lock @@ -59,7 +59,7 @@ dependencies = [ "bytes", "cfb", "chrono", - "cron", + "cron 0.12.1", "dirs", "flate2", "futures", @@ -146,6 +146,7 @@ version = "0.4.0" dependencies = [ "async-trait", "chrono", + "cron 0.17.0", "dashmap", "futures-core", "num_cpus", @@ -162,7 +163,7 @@ dependencies = [ [[package]] name = "a3s-memory" -version = "0.1.2" +version = "0.1.1" dependencies = [ "anyhow", "async-trait", @@ -305,7 +306,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -316,7 +317,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1368,6 +1369,19 @@ dependencies = [ "once_cell", ] +[[package]] +name = "cron" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5dcd6f69605c2956916ce24e8af637b754964c9a83f4662d3a2361654cdba09" +dependencies = [ + "chrono", + "once_cell", + "phf 0.11.3", + "serde", + "winnow", +] + [[package]] name = "crossbeam-channel" version = "0.5.15" @@ -1707,7 +1721,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2387,7 +2401,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -2999,7 +3013,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3379,7 +3393,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.40", - "socket2 0.5.10", + "socket2 0.6.4", "thiserror 2.0.18", "tokio", "tracing", @@ -3416,7 +3430,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.4", "tracing", "windows-sys 0.60.2", ] @@ -3763,7 +3777,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4204,7 +4218,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -4368,10 +4382,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5167,7 +5181,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/sdk/python/Cargo.lock b/sdk/python/Cargo.lock index fce2c93c..f5a546bb 100644 --- a/sdk/python/Cargo.lock +++ b/sdk/python/Cargo.lock @@ -44,6 +44,7 @@ dependencies = [ "a3s-acl 0.2.0", "a3s-ahp", "a3s-common", + "a3s-flow", "a3s-lane", "a3s-memory", "a3s-search", @@ -58,7 +59,7 @@ dependencies = [ "bytes", "cfb", "chrono", - "cron", + "cron 0.12.1", "dirs", "flate2", "futures", @@ -123,12 +124,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "a3s-flow" +version = "0.4.1" +dependencies = [ + "async-trait", + "chrono", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 1.0.69", + "tokio", + "tracing", + "uuid", +] + [[package]] name = "a3s-lane" version = "0.4.0" dependencies = [ "async-trait", "chrono", + "cron 0.17.0", "dashmap", "futures-core", "num_cpus", @@ -145,7 +162,7 @@ dependencies = [ [[package]] name = "a3s-memory" -version = "0.1.2" +version = "0.1.1" dependencies = [ "anyhow", "async-trait", @@ -288,7 +305,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -299,7 +316,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1342,6 +1359,19 @@ dependencies = [ "once_cell", ] +[[package]] +name = "cron" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5dcd6f69605c2956916ce24e8af637b754964c9a83f4662d3a2361654cdba09" +dependencies = [ + "chrono", + "once_cell", + "phf 0.11.3", + "serde", + "winnow", +] + [[package]] name = "crossbeam-channel" version = "0.5.15" @@ -1671,7 +1701,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2351,7 +2381,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -2911,7 +2941,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3360,7 +3390,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.40", - "socket2 0.5.10", + "socket2 0.6.4", "thiserror 2.0.18", "tokio", "tracing", @@ -3397,7 +3427,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.4", "tracing", "windows-sys 0.60.2", ] @@ -3744,7 +3774,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4185,7 +4215,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -4355,10 +4385,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5160,7 +5190,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] From 1ff2b47927db8cfd6e9d97ed0a22651de2e57b6a Mon Sep 17 00:00:00 2001 From: Li Xiangtian <1193027052@qq.com> Date: Tue, 7 Jul 2026 00:39:46 +0800 Subject: [PATCH 5/5] ci: replace flow path dependency in standalone builds --- .github/setup-workspace.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/setup-workspace.sh b/.github/setup-workspace.sh index 98e48002..518167f6 100644 --- a/.github/setup-workspace.sh +++ b/.github/setup-workspace.sh @@ -14,6 +14,7 @@ sed -i.bak \ -e 's|a3s-memory = { version = "0.1.1", path = "../../memory" }|a3s-memory = "0.1.1"|' \ -e 's|a3s-lane = { version = "0.4", path = "../../lane" }|a3s-lane = "0.4"|' \ -e 's|a3s-search = { version = "1.2.3", path = "../../search", default-features = false, features = \["lightpanda"\] }|a3s-search = { version = "1.2.3", default-features = false, features = ["lightpanda"] }|' \ + -e 's|a3s-flow = { version = "0.4.1", path = "../../flow" }|a3s-flow = "0.4.1"|' \ -e 's|a3s-box-sdk = { version = "0.7", path = "../../box/src/sdk", optional = true }|a3s-box-sdk = { version = "0.7", optional = true }|' \ -e 's|a3s-ahp = { version = "2.4", path = "../../ahp", optional = true, features = \["http", "websocket", "unix-socket"\] }|a3s-ahp = { version = "2.4", optional = true, features = ["http", "websocket", "unix-socket"] }|' \ -e 's|a3s-ahp = { version = "2.4", path = "../../ahp", optional = true, features = \["http", "websocket"\] }|a3s-ahp = { version = "2.4", optional = true, features = ["http", "websocket"] }|' \