From 706a2560703241cbb2d23bed527835b7360cc214 Mon Sep 17 00:00:00 2001 From: fangtianshun Date: Thu, 16 Jul 2026 09:40:19 +0800 Subject: [PATCH 01/25] Add standalone dlcapt workspace crate stub Co-authored-by: Cursor --- Cargo.lock | 38 ++++++++++++++++++++++++++++ Cargo.toml | 1 + crates/persisting-dlcapt/Cargo.toml | 33 ++++++++++++++++++++++++ crates/persisting-dlcapt/README.md | 4 +++ crates/persisting-dlcapt/src/lib.rs | 1 + crates/persisting-dlcapt/src/main.rs | 4 +++ 6 files changed, 81 insertions(+) create mode 100644 crates/persisting-dlcapt/Cargo.toml create mode 100644 crates/persisting-dlcapt/README.md create mode 100644 crates/persisting-dlcapt/src/lib.rs create mode 100644 crates/persisting-dlcapt/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 5233bbe..b967a6d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4149,6 +4149,7 @@ dependencies = [ "prost", "prost-build", "prost-types", + "protobuf-src", "rand 0.9.4", "rayon", "roaring", @@ -4316,6 +4317,7 @@ dependencies = [ "num-traits", "prost", "prost-build", + "protobuf-src", "rand 0.9.4", "strum", "tokio", @@ -4353,6 +4355,7 @@ dependencies = [ "prost", "prost-build", "prost-types", + "protobuf-src", "tokio", "tracing", ] @@ -4427,6 +4430,7 @@ dependencies = [ "prost", "prost-build", "prost-types", + "protobuf-src", "rand 0.9.4", "rand_distr", "rangemap", @@ -4558,6 +4562,7 @@ dependencies = [ "prost", "prost-build", "prost-types", + "protobuf-src", "rand 0.9.4", "rangemap", "roaring", @@ -5814,6 +5819,30 @@ dependencies = [ "serde_json", ] +[[package]] +name = "persisting-dlcapt" +version = "0.1.0" +dependencies = [ + "anyhow", + "arrow", + "async-trait", + "axum", + "bytes", + "chrono", + "http-body-util", + "lance", + "reqwest 0.12.28", + "serde", + "serde_json", + "tempfile", + "tokio", + "tokio-stream", + "toml", + "tower", + "tracing", + "tracing-subscriber", +] + [[package]] name = "persisting-engine" version = "0.1.0" @@ -6116,6 +6145,15 @@ dependencies = [ "prost", ] +[[package]] +name = "protobuf-src" +version = "2.1.1+27.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6217c3504da19b85a3a4b2e9a5183d635822d83507ba0986624b5c05b83bfc40" +dependencies = [ + "cmake", +] + [[package]] name = "pulsing-actor" version = "0.1.2" diff --git a/Cargo.toml b/Cargo.toml index 98904ee..e29d338 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ members = [ "crates/persisting-core", "crates/persisting-capture", "crates/persisting-cli", + "crates/persisting-dlcapt", ] [workspace.package] diff --git a/crates/persisting-dlcapt/Cargo.toml b/crates/persisting-dlcapt/Cargo.toml new file mode 100644 index 0000000..eaea021 --- /dev/null +++ b/crates/persisting-dlcapt/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "persisting-dlcapt" +version.workspace = true +edition = "2024" +authors.workspace = true +license.workspace = true +description = "Standalone OpenAI-compatible LLM proxy with session trajectory capture" + +[[bin]] +name = "dlcapt" +path = "src/main.rs" + +[dependencies] +anyhow = "1" +arrow = "58.3.0" +async-trait = "0.1" +axum = "0.8" +chrono = { version = "0.4", features = ["clock"] } +lance = { version = "7.0.0", features = ["protoc", "aws"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "net", "io-util", "fs", "time"] } +tokio-stream = "0.1" +toml = "0.8" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } + +[dev-dependencies] +tempfile = "3" +tower = { version = "0.5", features = ["util"] } +http-body-util = "0.1" +bytes = "1" diff --git a/crates/persisting-dlcapt/README.md b/crates/persisting-dlcapt/README.md new file mode 100644 index 0000000..d126be8 --- /dev/null +++ b/crates/persisting-dlcapt/README.md @@ -0,0 +1,4 @@ +# persisting-dlcapt + +`persisting-dlcapt` 是 Persisting workspace 内独立运行的 OpenAI 兼容代理与轨迹采集服务。 +第一阶段不依赖 `persisting-capture` 或 `persisting-engine`，以保持现有 dlcapt HTTP、session 和存储契约。 diff --git a/crates/persisting-dlcapt/src/lib.rs b/crates/persisting-dlcapt/src/lib.rs new file mode 100644 index 0000000..6a814c2 --- /dev/null +++ b/crates/persisting-dlcapt/src/lib.rs @@ -0,0 +1 @@ +//! persisting-dlcapt library stub; replaced by migrated sources in Task 2. diff --git a/crates/persisting-dlcapt/src/main.rs b/crates/persisting-dlcapt/src/main.rs new file mode 100644 index 0000000..a0d32fe --- /dev/null +++ b/crates/persisting-dlcapt/src/main.rs @@ -0,0 +1,4 @@ +fn main() { + eprintln!("persisting-dlcapt stub; run Task 2 migration before use"); + std::process::exit(1); +} From 21bd47b14f33a7b1c87edb0e972b11f86af64d75 Mon Sep 17 00:00:00 2001 From: fangtianshun Date: Thu, 16 Jul 2026 09:48:14 +0800 Subject: [PATCH 02/25] Migrate dlcapt sources into persisting-dlcapt Co-authored-by: Cursor --- crates/persisting-dlcapt/src/audit.rs | 111 +++ crates/persisting-dlcapt/src/capture/event.rs | 48 ++ .../src/capture/field_registry.rs | 246 ++++++ crates/persisting-dlcapt/src/capture/mod.rs | 17 + .../src/capture/post_process.rs | 37 + .../src/capture/session_dir.rs | 66 ++ .../src/capture/sink/lance.rs | 207 +++++ .../persisting-dlcapt/src/capture/sink/mod.rs | 1 + .../src/capture/sink_router.rs | 245 ++++++ .../src/capture/step_record.rs | 31 + .../src/capture/step_table_writer.rs | 130 ++++ .../src/capture/writers/lance_crate.rs | 237 ++++++ .../src/capture/writers/lance_storage.rs | 135 ++++ .../src/capture/writers/mod.rs | 2 + crates/persisting-dlcapt/src/config.rs | 450 +++++++++++ crates/persisting-dlcapt/src/dialogue/mod.rs | 60 ++ .../src/dialogue/responses.rs | 176 +++++ crates/persisting-dlcapt/src/lib.rs | 9 +- crates/persisting-dlcapt/src/main.rs | 55 +- crates/persisting-dlcapt/src/proxy.rs | 726 ++++++++++++++++++ crates/persisting-dlcapt/src/router.rs | 56 ++ crates/persisting-dlcapt/src/session/mod.rs | 76 ++ .../src/session/normalize.rs | 53 ++ .../src/session/providers/body.rs | 14 + .../src/session/providers/default.rs | 8 + .../src/session/providers/header.rs | 43 ++ .../src/session/providers/mod.rs | 9 + .../src/session/providers/url_path.rs | 8 + .../persisting-dlcapt/src/session/resolver.rs | 200 +++++ crates/persisting-dlcapt/src/tlv.rs | 331 ++++++++ crates/persisting-dlcapt/tests/mvp_tests.rs | 494 ++++++++++++ scripts/cargo-with-protoc.sh | 33 + 32 files changed, 4310 insertions(+), 4 deletions(-) create mode 100644 crates/persisting-dlcapt/src/audit.rs create mode 100644 crates/persisting-dlcapt/src/capture/event.rs create mode 100644 crates/persisting-dlcapt/src/capture/field_registry.rs create mode 100644 crates/persisting-dlcapt/src/capture/mod.rs create mode 100644 crates/persisting-dlcapt/src/capture/post_process.rs create mode 100644 crates/persisting-dlcapt/src/capture/session_dir.rs create mode 100644 crates/persisting-dlcapt/src/capture/sink/lance.rs create mode 100644 crates/persisting-dlcapt/src/capture/sink/mod.rs create mode 100644 crates/persisting-dlcapt/src/capture/sink_router.rs create mode 100644 crates/persisting-dlcapt/src/capture/step_record.rs create mode 100644 crates/persisting-dlcapt/src/capture/step_table_writer.rs create mode 100644 crates/persisting-dlcapt/src/capture/writers/lance_crate.rs create mode 100644 crates/persisting-dlcapt/src/capture/writers/lance_storage.rs create mode 100644 crates/persisting-dlcapt/src/capture/writers/mod.rs create mode 100644 crates/persisting-dlcapt/src/config.rs create mode 100644 crates/persisting-dlcapt/src/dialogue/mod.rs create mode 100644 crates/persisting-dlcapt/src/dialogue/responses.rs create mode 100644 crates/persisting-dlcapt/src/proxy.rs create mode 100644 crates/persisting-dlcapt/src/router.rs create mode 100644 crates/persisting-dlcapt/src/session/mod.rs create mode 100644 crates/persisting-dlcapt/src/session/normalize.rs create mode 100644 crates/persisting-dlcapt/src/session/providers/body.rs create mode 100644 crates/persisting-dlcapt/src/session/providers/default.rs create mode 100644 crates/persisting-dlcapt/src/session/providers/header.rs create mode 100644 crates/persisting-dlcapt/src/session/providers/mod.rs create mode 100644 crates/persisting-dlcapt/src/session/providers/url_path.rs create mode 100644 crates/persisting-dlcapt/src/session/resolver.rs create mode 100644 crates/persisting-dlcapt/src/tlv.rs create mode 100644 crates/persisting-dlcapt/tests/mvp_tests.rs create mode 100755 scripts/cargo-with-protoc.sh diff --git a/crates/persisting-dlcapt/src/audit.rs b/crates/persisting-dlcapt/src/audit.rs new file mode 100644 index 0000000..afd66bf --- /dev/null +++ b/crates/persisting-dlcapt/src/audit.rs @@ -0,0 +1,111 @@ +use anyhow::{Context, Result}; +use chrono::Utc; +use serde_json::Value; +use std::path::PathBuf; +use tokio::fs; + +#[derive(Debug, Clone)] +pub struct AuditWriter { + store_root: PathBuf, +} + +#[derive(Debug, Clone)] +pub struct AuditRecord { + pub session_id: String, + pub model: String, + pub stream: bool, + pub status_code: u16, + pub request_body: Value, + pub response_summary: String, + pub response_text: Option, + pub response_raw: Option, + pub finish_reason: Option, + pub usage: Option, +} + +impl AuditWriter { + pub fn new(store_root: PathBuf) -> Self { + Self { store_root } + } + + pub async fn write(&self, record: AuditRecord) -> Result { + let run_id = Utc::now().format("%Y%m%d-%H%M%S-%3f").to_string(); + let run_dir = self + .store_root + .join(sanitize_path_segment(&record.session_id)) + .join(format!("run-{run_id}")); + fs::create_dir_all(&run_dir) + .await + .with_context(|| format!("failed creating audit dir: {}", run_dir.display()))?; + + let run_file = run_dir.join(format!("run-{run_id}.md")); + let body_pretty = serde_json::to_string_pretty(&record.request_body) + .with_context(|| "failed rendering request body".to_string())?; + + let usage_text = record + .usage + .as_ref() + .map(|value| serde_json::to_string(value).unwrap_or_else(|_| "{}".to_string())) + .unwrap_or_else(|| "null".to_string()); + + let finish_reason = record + .finish_reason + .unwrap_or_else(|| "unknown".to_string()); + let response_text = record + .response_text + .as_deref() + .unwrap_or("(empty response text)"); + let response_raw = record + .response_raw + .as_deref() + .unwrap_or("(empty response raw)"); + let content = format!( + "# run-{run_id}\n\n\ +session_id: {session}\n\ +model: {model}\n\ +stream: {stream}\n\ +status_code: {status_code}\n\ +finish_reason: {finish_reason}\n\ +usage: {usage_text}\n\ +response_summary: {response_summary}\n\n\ +## request\n\ +```json\n\ +{body_pretty}\n\ +```\n\n\ +## response_text\n\ +{response_text}\n\n\ +## response_raw\n\ +```text\n\ +{response_raw}\n\ +```\n", + session = record.session_id, + model = record.model, + stream = record.stream, + status_code = record.status_code, + response_summary = record.response_summary.replace('\n', " "), + ); + + fs::write(&run_file, content) + .await + .with_context(|| format!("failed writing audit file: {}", run_file.display()))?; + Ok(run_file) + } +} + +fn sanitize_path_segment(raw: &str) -> String { + let sanitized: String = raw + .chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { + ch + } else { + '-' + } + }) + .collect(); + if sanitized.is_empty() { + "unknown".to_string() + } else { + sanitized + } +} diff --git a/crates/persisting-dlcapt/src/capture/event.rs b/crates/persisting-dlcapt/src/capture/event.rs new file mode 100644 index 0000000..20805cb --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/event.rs @@ -0,0 +1,48 @@ +use crate::dialogue::InferenceEndpoint; +use chrono::{DateTime, Utc}; +use serde_json::Value; +use std::collections::BTreeMap; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FieldSink { + TopLevel, + Extensions, + Capture, +} + +#[derive(Debug, Clone)] +pub struct FieldPatch { + pub sink: FieldSink, + pub value: Value, +} + +#[derive(Debug, Clone, Default)] +pub struct CaptureMeta { + pub finish_reason: Option, + pub usage: Option, + pub segment_kind: Option, +} + +#[derive(Debug, Clone)] +pub struct CaptureEvent { + pub call_id: String, + pub session_id: String, + pub agent_id: String, + pub step_id: u64, + pub turn: u64, + pub endpoint: InferenceEndpoint, + pub request_path: String, + pub model: String, + pub request: Value, + pub request_headers: BTreeMap, + pub response_raw: Value, + pub response_text: Option, + pub stream: bool, + pub status_code: u16, + pub completed_at: DateTime, + pub metadata: BTreeMap, + pub field_patches: BTreeMap, + pub capture_meta: CaptureMeta, + pub user_seq: u64, + pub assistant_seq: u64, +} diff --git a/crates/persisting-dlcapt/src/capture/field_registry.rs b/crates/persisting-dlcapt/src/capture/field_registry.rs new file mode 100644 index 0000000..f764734 --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/field_registry.rs @@ -0,0 +1,246 @@ +use crate::capture::event::{CaptureEvent, FieldSink}; +use crate::capture::step_record::StepRecord; +use crate::config::{ExportConfig, ExportDefaults}; +use crate::dialogue::InferenceEndpoint; +use chrono::Utc; +use serde_json::{Value, json}; + +pub struct FieldRegistry { + defaults: ExportDefaults, + max_steps_per_session: u64, +} + +impl FieldRegistry { + pub fn from_export(export: &ExportConfig) -> Self { + Self { + defaults: export.defaults.clone(), + max_steps_per_session: export.max_steps_per_session, + } + } +} + +pub fn materialize_session_step( + event: &CaptureEvent, + registry: &FieldRegistry, + session_metadata: &serde_json::Map, + run_bucket: &str, +) -> StepRecord { + let defaults = ®istry.defaults; + let step_id = event.step_id as i64; + + let mut extensions = serde_json::Map::new(); + for (key, value) in session_metadata { + extensions.insert(key.clone(), value.clone()); + } + for (key, value) in &event.metadata { + extensions.insert(key.clone(), value.clone()); + } + for (key, patch) in &event.field_patches { + if patch.sink == FieldSink::Extensions { + extensions.insert(key.clone(), patch.value.clone()); + } + } + + let mut capture_obj = serde_json::Map::new(); + capture_obj.insert("call_id".to_string(), json!(event.call_id)); + if let Some(reason) = &event.capture_meta.finish_reason { + capture_obj.insert("finish_reason".to_string(), json!(reason)); + } + if let Some(usage) = &event.capture_meta.usage { + capture_obj.insert("usage".to_string(), usage.clone()); + } + if let Some(kind) = &event.capture_meta.segment_kind { + capture_obj.insert("segment_kind".to_string(), json!(kind)); + } + for (key, patch) in &event.field_patches { + if patch.sink == FieldSink::Capture { + capture_obj.insert(key.clone(), patch.value.clone()); + } + } + + let messages = normalize_messages(event.endpoint, &event.request); + let response = build_response(event); + + let mut group_id = defaults.group_id.clone(); + let mut step_reward = defaults.step_reward; + let mut reward = defaults.reward; + let mut is_terminal = defaults.is_terminal; + let mut is_trainable = defaults.is_trainable; + let mut env_name = defaults.env_name.clone(); + let mut job_id = defaults.job_id.clone(); + + for (key, patch) in &event.field_patches { + if patch.sink != FieldSink::TopLevel { + continue; + } + match key.as_str() { + "group_id" if patch.value.is_string() => { + group_id = patch.value.as_str().unwrap_or_default().to_string(); + } + "step_reward" if patch.value.is_number() => { + step_reward = patch.value.as_f64().unwrap_or(step_reward); + } + "reward" if patch.value.is_number() => { + reward = patch.value.as_f64().unwrap_or(reward); + } + "is_terminal" if patch.value.is_boolean() => { + is_terminal = patch.value.as_bool().unwrap_or(is_terminal); + } + "is_trainable" if patch.value.is_boolean() => { + is_trainable = patch.value.as_bool().unwrap_or(is_trainable); + } + "env_name" if patch.value.is_string() => { + env_name = patch.value.as_str().unwrap_or_default().to_string(); + } + "job_id" if patch.value.is_string() => { + job_id = patch.value.as_str().unwrap_or_default().to_string(); + } + _ => {} + } + } + + let is_truncated = step_id as u64 >= registry.max_steps_per_session; + let is_session_completed = is_terminal || is_truncated; + let created_at = format_gateway_timestamp(event.completed_at); + let id = format!("{job_id}:{}:{step_id}", event.session_id); + + let extensions_json = if extensions.is_empty() { + None + } else { + Some(serde_json::to_string(&Value::Object(extensions)).unwrap_or_else(|_| "{}".to_string())) + }; + let capture_json = if capture_obj.is_empty() { + None + } else { + Some( + serde_json::to_string(&Value::Object(capture_obj)).unwrap_or_else(|_| "{}".to_string()), + ) + }; + + StepRecord { + id, + session_id: event.session_id.clone(), + step_id, + job_id, + agent_id: event.agent_id.clone(), + group_id, + env_name, + llm_model: event.model.clone(), + step_reward, + reward, + is_terminal, + is_truncated, + is_session_completed, + is_trainable, + created_at, + messages_json: serde_json::to_string(&messages).unwrap_or_else(|_| "[]".to_string()), + response_json: serde_json::to_string(&response).unwrap_or_else(|_| "{}".to_string()), + env_state_json: serde_json::to_string(&defaults.env_state) + .unwrap_or_else(|_| "{}".to_string()), + extensions_json, + capture_json, + run_bucket: run_bucket.to_string(), + call_id: event.call_id.clone(), + source_export_id: None, + } +} + +fn normalize_messages(endpoint: InferenceEndpoint, request: &Value) -> Value { + match endpoint { + InferenceEndpoint::ChatCompletions => request + .get("messages") + .cloned() + .unwrap_or_else(|| json!([])), + InferenceEndpoint::Responses => { + if let Some(input) = request.get("input") { + json!([{"role": "user", "content": input}]) + } else { + json!([]) + } + } + } +} + +fn build_response(event: &CaptureEvent) -> Value { + match event.endpoint { + InferenceEndpoint::ChatCompletions => event + .response_raw + .get("choices") + .and_then(|c| c.as_array()) + .and_then(|choices| choices.first()) + .and_then(|choice| choice.get("message")) + .cloned() + .unwrap_or_else(|| { + if let Some(text) = &event.response_text { + json!({"role": "assistant", "content": text}) + } else { + json!({"role": "assistant", "content": null}) + } + }), + InferenceEndpoint::Responses => { + if let Some(text) = &event.response_text { + json!({"role": "assistant", "content": text}) + } else { + json!({"role": "assistant", "content": null}) + } + } + } +} + +fn format_gateway_timestamp(dt: chrono::DateTime) -> String { + dt.format("%Y-%m-%d %H:%M:%S%.6f%:z").to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::capture::event::CaptureMeta; + use crate::config::ExportConfig; + use std::collections::BTreeMap; + + fn sample_event() -> CaptureEvent { + CaptureEvent { + call_id: "call-1".to_string(), + session_id: "sess-1".to_string(), + agent_id: "openclaw".to_string(), + step_id: 1, + turn: 1, + endpoint: InferenceEndpoint::ChatCompletions, + request_path: "/v1/chat/completions".to_string(), + model: "kimi".to_string(), + request: json!({"messages": [{"role": "user", "content": "hi"}]}), + request_headers: BTreeMap::new(), + response_raw: json!({ + "choices": [{"message": {"role": "assistant", "content": "hello"}}] + }), + response_text: Some("hello".to_string()), + stream: false, + status_code: 200, + completed_at: Utc::now(), + metadata: BTreeMap::new(), + field_patches: BTreeMap::new(), + capture_meta: CaptureMeta { + finish_reason: Some("stop".to_string()), + usage: Some(json!({"total_tokens": 10})), + segment_kind: None, + }, + user_seq: 0, + assistant_seq: 1, + } + } + + #[test] + fn materialize_produces_composite_id() { + let registry = FieldRegistry::from_export(&ExportConfig::default()); + let record = materialize_session_step( + &sample_event(), + ®istry, + &serde_json::Map::new(), + "2026-06-16", + ); + assert_eq!(record.id, "dlcapt:sess-1:1"); + assert_eq!(record.step_id, 1); + assert!(record.messages_json.contains("hi")); + assert!(record.response_json.contains("hello")); + } +} diff --git a/crates/persisting-dlcapt/src/capture/mod.rs b/crates/persisting-dlcapt/src/capture/mod.rs new file mode 100644 index 0000000..07bc6ed --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/mod.rs @@ -0,0 +1,17 @@ +pub mod event; +pub mod field_registry; +pub mod post_process; +pub mod session_dir; +pub mod sink; +pub mod sink_router; +pub mod step_record; +pub mod step_table_writer; +pub mod writers; + +pub use event::{CaptureEvent, CaptureMeta, FieldPatch, FieldSink}; +pub use field_registry::{FieldRegistry, materialize_session_step}; +pub use post_process::{PostProcessor, PostProcessorChain}; +pub use session_dir::{SessionLayout, resolve_session_layout, resolve_session_layout_with_bucket}; +pub use sink_router::CaptureSinkRouter; +pub use step_record::StepRecord; +pub use step_table_writer::{LanceStepRow, StepTableWriter, step_record_to_lance_row}; diff --git a/crates/persisting-dlcapt/src/capture/post_process.rs b/crates/persisting-dlcapt/src/capture/post_process.rs new file mode 100644 index 0000000..2345ec7 --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/post_process.rs @@ -0,0 +1,37 @@ +use crate::capture::event::CaptureEvent; +use anyhow::Result; + +pub trait PostProcessor: Send + Sync { + fn name(&self) -> &'static str; + fn process(&self, _event: &mut CaptureEvent) -> Result<()> { + Ok(()) + } +} + +pub struct PostProcessorChain { + processors: Vec>, +} + +impl PostProcessorChain { + pub fn new(processors: Vec>) -> Self { + Self { processors } + } + + pub fn empty() -> Self { + Self { + processors: Vec::new(), + } + } + + pub fn apply(&self, event: &mut CaptureEvent) { + for processor in &self.processors { + if let Err(err) = processor.process(event) { + tracing::warn!( + processor = processor.name(), + error = %err, + "post-processor failed; skipping" + ); + } + } + } +} diff --git a/crates/persisting-dlcapt/src/capture/session_dir.rs b/crates/persisting-dlcapt/src/capture/session_dir.rs new file mode 100644 index 0000000..9eebb49 --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/session_dir.rs @@ -0,0 +1,66 @@ +use chrono::{DateTime, Utc}; + +/// Resolved storage layout for a session (relative to `store_dir`). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionLayout { + pub session_dir: String, + pub run_bucket: String, +} + +/// Compute `session_dir` and `run_bucket` per spec §3.7.1. +pub fn resolve_session_layout( + storage_session_id: &str, + default_session_id: &str, + now: DateTime, +) -> SessionLayout { + let today = now.format("%Y-%m-%d").to_string(); + if storage_session_id == default_session_id { + SessionLayout { + session_dir: format!("default/{today}"), + run_bucket: today, + } + } else { + SessionLayout { + session_dir: storage_session_id.to_string(), + run_bucket: today, + } + } +} + +/// For real sessions, `run_bucket` is the first-write UTC date (may differ from `now`). +pub fn resolve_session_layout_with_bucket( + storage_session_id: &str, + default_session_id: &str, + now: DateTime, + existing_run_bucket: Option<&str>, +) -> SessionLayout { + let mut layout = resolve_session_layout(storage_session_id, default_session_id, now); + if storage_session_id != default_session_id { + if let Some(bucket) = existing_run_bucket.filter(|b| !b.is_empty()) { + layout.run_bucket = bucket.to_string(); + } + } + layout +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + + #[test] + fn default_session_uses_date_bucket_in_path() { + let now = Utc.with_ymd_and_hms(2026, 6, 16, 12, 0, 0).unwrap(); + let layout = resolve_session_layout("default", "default", now); + assert_eq!(layout.session_dir, "default/2026-06-16"); + assert_eq!(layout.run_bucket, "2026-06-16"); + } + + #[test] + fn real_session_uses_flat_dir() { + let now = Utc.with_ymd_and_hms(2026, 6, 16, 12, 0, 0).unwrap(); + let layout = resolve_session_layout("abc-123", "default", now); + assert_eq!(layout.session_dir, "abc-123"); + assert_eq!(layout.run_bucket, "2026-06-16"); + } +} diff --git a/crates/persisting-dlcapt/src/capture/sink/lance.rs b/crates/persisting-dlcapt/src/capture/sink/lance.rs new file mode 100644 index 0000000..071cd12 --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/sink/lance.rs @@ -0,0 +1,207 @@ +use crate::capture::step_record::StepRecord; +use crate::capture::step_table_writer::{StepTableWriter, build_step_table_writer}; +use crate::config::{LanceStorageConfig, ProxyConfig}; +use anyhow::{Context, Result}; +use chrono::Utc; +use serde_json::json; +use std::fs::OpenOptions; +use std::io::Write; +use std::path::PathBuf; +use std::sync::Arc; + +pub struct LanceSink { + writer: Arc, + fail_open: bool, + dead_letter_path: PathBuf, +} + +impl LanceSink { + pub fn from_config(config: &ProxyConfig) -> Result { + let lance_cfg = &config.storage.lance; + let writer = build_step_table_writer(lance_cfg)?; + let dead_letter_path = if lance_cfg.dead_letter_path.is_empty() { + PathBuf::from(&config.store_dir).join(".capture/lance_dead_letter.jsonl") + } else if PathBuf::from(&lance_cfg.dead_letter_path).is_absolute() { + PathBuf::from(&lance_cfg.dead_letter_path) + } else { + PathBuf::from(&config.store_dir).join(&lance_cfg.dead_letter_path) + }; + + Ok(Self { + writer, + fail_open: lance_cfg.fail_open, + dead_letter_path, + }) + } + + pub fn from_parts( + lance_cfg: &LanceStorageConfig, + store_dir: &str, + writer: Arc, + ) -> Result { + let dead_letter_path = if lance_cfg.dead_letter_path.is_empty() { + PathBuf::from(store_dir).join(".capture/lance_dead_letter.jsonl") + } else if PathBuf::from(&lance_cfg.dead_letter_path).is_absolute() { + PathBuf::from(&lance_cfg.dead_letter_path) + } else { + PathBuf::from(store_dir).join(&lance_cfg.dead_letter_path) + }; + + Ok(Self { + writer, + fail_open: lance_cfg.fail_open, + dead_letter_path, + }) + } + + pub async fn append(&self, record: &StepRecord) -> Result<()> { + match self.writer.append(record).await { + Ok(()) => Ok(()), + Err(err) if self.fail_open => { + self.write_dead_letter(record, &err)?; + Ok(()) + } + Err(err) => Err(err), + } + } + + fn write_dead_letter(&self, record: &StepRecord, error: &anyhow::Error) -> Result<()> { + if let Some(parent) = self.dead_letter_path.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("create dead letter dir {}", parent.display()))?; + } + + let line = json!({ + "ts": Utc::now().to_rfc3339(), + "error": error.to_string(), + "backend": "lance", + "record": record, + }); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&self.dead_letter_path) + .with_context(|| format!("open dead letter {}", self.dead_letter_path.display()))?; + serde_json::to_writer(&mut file, &line).context("serialize dead letter entry")?; + file.write_all(b"\n").context("write dead letter newline")?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::capture::step_record::StepRecord; + use crate::capture::step_table_writer::StepTableWriter; + use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct FailingWriter; + + #[async_trait] + impl StepTableWriter for FailingWriter { + async fn append(&self, _record: &StepRecord) -> Result<()> { + anyhow::bail!("simulated lance failure") + } + } + + struct CountingWriter { + count: AtomicUsize, + } + + #[async_trait] + impl StepTableWriter for CountingWriter { + async fn append(&self, _record: &StepRecord) -> Result<()> { + self.count.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + + fn test_record() -> StepRecord { + StepRecord { + id: "dlcapt:sess:1".to_string(), + session_id: "sess".to_string(), + step_id: 1, + job_id: "dlcapt".to_string(), + agent_id: "openclaw".to_string(), + group_id: String::new(), + env_name: "openclaw".to_string(), + llm_model: "m".to_string(), + step_reward: 0.0, + reward: 0.0, + is_terminal: false, + is_truncated: false, + is_session_completed: false, + is_trainable: true, + created_at: "2026-06-16T00:00:00Z".to_string(), + messages_json: "[]".to_string(), + response_json: "{}".to_string(), + env_state_json: "{}".to_string(), + extensions_json: None, + capture_json: None, + run_bucket: "2026-06-16".to_string(), + call_id: "call-1".to_string(), + source_export_id: None, + } + } + + #[tokio::test] + async fn fail_open_writes_dead_letter_and_returns_ok() { + let temp = tempfile::tempdir().expect("tempdir"); + let cfg = LanceStorageConfig { + fail_open: true, + dead_letter_path: ".capture/lance_dead_letter.jsonl".to_string(), + ..LanceStorageConfig::default() + }; + let sink = LanceSink::from_parts( + &cfg, + temp.path().to_string_lossy().as_ref(), + Arc::new(FailingWriter), + ) + .expect("sink"); + + sink.append(&test_record()).await.expect("fail_open append"); + + let dead_letter = temp.path().join(".capture/lance_dead_letter.jsonl"); + assert!(dead_letter.exists()); + let text = std::fs::read_to_string(dead_letter).expect("read dead letter"); + assert!(text.contains("simulated lance failure")); + assert!(text.contains("dlcapt:sess:1")); + } + + #[tokio::test] + async fn fail_closed_propagates_error() { + let temp = tempfile::tempdir().expect("tempdir"); + let cfg = LanceStorageConfig { + fail_open: false, + ..LanceStorageConfig::default() + }; + let sink = LanceSink::from_parts( + &cfg, + temp.path().to_string_lossy().as_ref(), + Arc::new(FailingWriter), + ) + .expect("sink"); + + let err = sink.append(&test_record()).await.unwrap_err(); + assert!(err.to_string().contains("simulated lance failure")); + } + + #[tokio::test] + async fn successful_append_delegates_to_writer() { + let temp = tempfile::tempdir().expect("tempdir"); + let counter = Arc::new(CountingWriter { + count: AtomicUsize::new(0), + }); + let cfg = LanceStorageConfig::default(); + let sink = LanceSink::from_parts( + &cfg, + temp.path().to_string_lossy().as_ref(), + Arc::clone(&counter) as Arc, + ) + .expect("sink"); + + sink.append(&test_record()).await.expect("append"); + assert_eq!(counter.count.load(Ordering::SeqCst), 1); + } +} diff --git a/crates/persisting-dlcapt/src/capture/sink/mod.rs b/crates/persisting-dlcapt/src/capture/sink/mod.rs new file mode 100644 index 0000000..63f862b --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/sink/mod.rs @@ -0,0 +1 @@ +pub mod lance; diff --git a/crates/persisting-dlcapt/src/capture/sink_router.rs b/crates/persisting-dlcapt/src/capture/sink_router.rs new file mode 100644 index 0000000..b3fb877 --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/sink_router.rs @@ -0,0 +1,245 @@ +use crate::capture::event::CaptureEvent; +use crate::capture::field_registry::{FieldRegistry, materialize_session_step}; +use crate::capture::session_dir::resolve_session_layout_with_bucket; +use crate::capture::sink::lance::LanceSink; +use crate::capture::step_record::StepRecord; +use crate::config::ProxyConfig; +use crate::tlv::{MdSinkInput, TlvWriter}; +use anyhow::{Context, Result}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use tokio::fs; +use tokio::sync::Mutex; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct SessionStepsEnvelope { + format_version: u32, + session_id: String, + session_dir: String, + agent_id: String, + run_bucket: String, + source: String, + authoritative: String, + #[serde(default)] + session_metadata: serde_json::Map, + #[serde(default)] + session_steps: Vec, +} + +pub struct CaptureSinkRouter { + store_root: PathBuf, + config: Arc, + registry: FieldRegistry, + tlv: TlvWriter, + write_lock: Arc>, + lance_sink: Option, +} + +impl CaptureSinkRouter { + pub fn new( + config: Arc, + tlv: TlvWriter, + write_lock: Arc>, + ) -> Result { + let registry = FieldRegistry::from_export(&config.export); + let lance_sink = if config.storage.lance_enabled() { + Some(LanceSink::from_config(&config)?) + } else { + None + }; + Ok(Self { + store_root: PathBuf::from(&config.store_dir), + config, + registry, + tlv, + write_lock, + lance_sink, + }) + } + + pub async fn dispatch(&self, event: CaptureEvent) -> Result<()> { + let _guard = self.write_lock.lock().await; + let now = Utc::now(); + let json_path = self.session_steps_path(&event.session_id, now).await?; + let existing_bucket = self.read_envelope_run_bucket(&json_path).await?; + let layout = resolve_session_layout_with_bucket( + &event.session_id, + &self.config.default_session_id, + now, + existing_bucket.as_deref(), + ); + + let session_metadata = self.config.export.session_metadata.clone(); + let record = materialize_session_step( + &event, + &self.registry, + &session_metadata, + &layout.run_bucket, + ); + + if self.should_write_json_authoritative() { + self.append_json_file(&event.session_id, &layout, &record, "json_file") + .await?; + } else if self.should_write_json_cache() { + self.append_json_file(&event.session_id, &layout, &record, "lance") + .await?; + } + + if let Some(lance_sink) = &self.lance_sink { + lance_sink.append(&record).await?; + } + + if self.should_write_md() { + self.append_md(&event).await?; + } + + Ok(()) + } + + fn should_write_json_authoritative(&self) -> bool { + self.config.storage.authoritative == "json_file" + } + + fn should_write_json_cache(&self) -> bool { + self.config.storage.json_cache_enabled() + } + + fn should_write_md(&self) -> bool { + self.config.storage.also.contains(&"md".to_string()) + || self.config.storage.authoritative == "md" + } + + async fn session_steps_path( + &self, + session_id: &str, + now: chrono::DateTime, + ) -> Result { + let layout = resolve_session_layout_with_bucket( + session_id, + &self.config.default_session_id, + now, + None, + ); + Ok(self + .store_root + .join(&layout.session_dir) + .join("session_steps.json")) + } + + async fn read_envelope_run_bucket(&self, path: &Path) -> Result> { + if !path.exists() { + return Ok(None); + } + let text = fs::read_to_string(path) + .await + .with_context(|| format!("read {}", path.display()))?; + let envelope: SessionStepsEnvelope = + serde_json::from_str(&text).context("parse session_steps.json envelope")?; + Ok(Some(envelope.run_bucket)) + } + + async fn append_json_file( + &self, + session_id: &str, + layout: &crate::capture::session_dir::SessionLayout, + record: &StepRecord, + authoritative: &str, + ) -> Result<()> { + let dir = self.store_root.join(&layout.session_dir); + fs::create_dir_all(&dir) + .await + .with_context(|| format!("create {}", dir.display()))?; + let path = dir.join("session_steps.json"); + + let mut envelope = if path.exists() { + let text = fs::read_to_string(&path).await?; + serde_json::from_str::(&text) + .unwrap_or_else(|_| self.fresh_envelope(session_id, layout, authoritative)) + } else { + self.fresh_envelope(session_id, layout, authoritative) + }; + + envelope + .session_steps + .push(step_record_to_json_element(record)?); + let content = serde_json::to_string_pretty(&envelope).context("serialize envelope")?; + fs::write(&path, content) + .await + .with_context(|| format!("write {}", path.display()))?; + Ok(()) + } + + fn fresh_envelope( + &self, + session_id: &str, + layout: &crate::capture::session_dir::SessionLayout, + authoritative: &str, + ) -> SessionStepsEnvelope { + SessionStepsEnvelope { + format_version: 1, + session_id: session_id.to_string(), + session_dir: layout.session_dir.clone(), + agent_id: self.config.agent_id.clone(), + run_bucket: layout.run_bucket.clone(), + source: "dlcapt-proxy".to_string(), + authoritative: authoritative.to_string(), + session_metadata: self.config.export.session_metadata.clone(), + session_steps: Vec::new(), + } + } + + async fn append_md(&self, event: &CaptureEvent) -> Result<()> { + let input = MdSinkInput::from_capture_event(event); + if let Some(record) = input.to_tlv_record() { + self.tlv + .append_turn_internal(record) + .await + .context("md sink tlv append")?; + } + Ok(()) + } +} + +fn step_record_to_json_element(record: &StepRecord) -> Result { + let extensions = record + .extensions_json + .as_ref() + .map(|s| serde_json::from_str::(s)) + .transpose()? + .unwrap_or(json!({})); + let capture = record + .capture_json + .as_ref() + .map(|s| serde_json::from_str::(s)) + .transpose()? + .unwrap_or(json!({})); + let messages: Value = serde_json::from_str(&record.messages_json).unwrap_or(json!([])); + let response: Value = serde_json::from_str(&record.response_json).unwrap_or(json!({})); + let env_state: Value = serde_json::from_str(&record.env_state_json).unwrap_or(json!({})); + + Ok(json!({ + "id": record.id, + "source_export_id": record.source_export_id, + "session_id": record.session_id, + "step_id": record.step_id, + "job_id": record.job_id, + "group_id": record.group_id, + "env_name": record.env_name, + "llm_model": record.llm_model, + "messages": messages, + "response": response, + "step_reward": record.step_reward, + "reward": record.reward, + "env_state": env_state, + "is_terminal": record.is_terminal, + "is_truncated": record.is_truncated, + "is_session_completed": record.is_session_completed, + "is_trainable": record.is_trainable, + "created_at": record.created_at, + "extensions": extensions, + "_capture": capture, + })) +} diff --git a/crates/persisting-dlcapt/src/capture/step_record.rs b/crates/persisting-dlcapt/src/capture/step_record.rs new file mode 100644 index 0000000..811b15e --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/step_record.rs @@ -0,0 +1,31 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct StepRecord { + pub id: String, + pub session_id: String, + pub step_id: i64, + pub job_id: String, + pub agent_id: String, + pub group_id: String, + pub env_name: String, + pub llm_model: String, + pub step_reward: f64, + pub reward: f64, + pub is_terminal: bool, + pub is_truncated: bool, + pub is_session_completed: bool, + pub is_trainable: bool, + pub created_at: String, + pub messages_json: String, + pub response_json: String, + pub env_state_json: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub extensions_json: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub capture_json: Option, + pub run_bucket: String, + pub call_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub source_export_id: Option, +} diff --git a/crates/persisting-dlcapt/src/capture/step_table_writer.rs b/crates/persisting-dlcapt/src/capture/step_table_writer.rs new file mode 100644 index 0000000..3d1caf9 --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/step_table_writer.rs @@ -0,0 +1,130 @@ +use crate::capture::step_record::StepRecord; +use crate::config::LanceStorageConfig; +use anyhow::{Result, bail}; +use async_trait::async_trait; +use std::sync::Arc; + +pub use crate::capture::writers::lance_crate::LanceCrateWriter; + +/// 单表 session_steps 追加写；一步一行。 +#[async_trait] +pub trait StepTableWriter: Send + Sync { + async fn append(&self, record: &StepRecord) -> Result<()>; +} + +pub fn build_step_table_writer(cfg: &LanceStorageConfig) -> Result> { + match cfg.backend.as_str() { + "lance" => Ok(Arc::new(LanceCrateWriter::new(cfg)?)), + other => bail!("unsupported lance backend: {other}"), + } +} + +/// Lance 行映射（与 tests/lance/session_steps_import.py `step_to_lance_row` 对齐）。 +#[derive(Debug, Clone, PartialEq)] +pub struct LanceStepRow { + pub id: String, + pub session_id: String, + pub step_id: i64, + pub job_id: String, + pub group_id: String, + pub env_name: String, + pub llm_model: String, + pub messages_json: String, + pub response_json: String, + pub step_reward: f64, + pub reward: f64, + pub env_state_json: String, + pub is_terminal: bool, + pub is_truncated: bool, + pub is_session_completed: bool, + pub is_trainable: bool, + pub created_at: String, + pub agent_id: String, + pub root_session: String, + pub extensions_json: Option, + pub capture_json: Option, + pub call_id: String, + pub source_export_id: Option, +} + +pub fn step_record_to_lance_row(record: &StepRecord) -> LanceStepRow { + LanceStepRow { + id: record.id.clone(), + session_id: record.session_id.clone(), + step_id: record.step_id, + job_id: record.job_id.clone(), + group_id: record.group_id.clone(), + env_name: record.env_name.clone(), + llm_model: record.llm_model.clone(), + messages_json: record.messages_json.clone(), + response_json: record.response_json.clone(), + step_reward: record.step_reward, + reward: record.reward, + env_state_json: record.env_state_json.clone(), + is_terminal: record.is_terminal, + is_truncated: record.is_truncated, + is_session_completed: record.is_session_completed, + is_trainable: record.is_trainable, + created_at: record.created_at.clone(), + agent_id: record.agent_id.clone(), + root_session: record.run_bucket.clone(), + extensions_json: record.extensions_json.clone(), + capture_json: record.capture_json.clone(), + call_id: record.call_id.clone(), + source_export_id: record.source_export_id, + } +} + +pub fn lance_dataset_uri(db_uri: &str, table_name: &str) -> String { + crate::capture::writers::lance_storage::lance_dataset_uri(db_uri, table_name) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::capture::step_record::StepRecord; + + fn fixture_record() -> StepRecord { + StepRecord { + id: "dlcapt:sess-import-test:1".to_string(), + session_id: "sess-import-test".to_string(), + step_id: 1, + job_id: "dlcapt".to_string(), + agent_id: "openclaw".to_string(), + group_id: String::new(), + env_name: "openclaw".to_string(), + llm_model: "kimi-k2.5".to_string(), + step_reward: 0.0, + reward: 0.0, + is_terminal: false, + is_truncated: false, + is_session_completed: false, + is_trainable: true, + created_at: "2026-06-16 09:57:27.641681+00:00".to_string(), + messages_json: r#"[{"role":"user","content":"ping"}]"#.to_string(), + response_json: r#"{"role":"assistant","content":"pong"}"#.to_string(), + env_state_json: "{}".to_string(), + extensions_json: Some(r#"{"source":"dlcapt-proxy"}"#.to_string()), + capture_json: Some( + r#"{"call_id":"call-import-test-1","finish_reason":"stop"}"#.to_string(), + ), + run_bucket: "2026-06-16".to_string(), + call_id: "call-import-test-1".to_string(), + source_export_id: None, + } + } + + #[test] + fn step_record_to_lance_row_maps_root_session_and_agent_id() { + let row = step_record_to_lance_row(&fixture_record()); + assert_eq!(row.id, "dlcapt:sess-import-test:1"); + assert_eq!(row.root_session, "2026-06-16"); + assert_eq!(row.agent_id, "openclaw"); + assert_eq!(row.call_id, "call-import-test-1"); + assert_eq!( + row.extensions_json.as_deref(), + Some(r#"{"source":"dlcapt-proxy"}"#) + ); + assert!(row.source_export_id.is_none()); + } +} diff --git a/crates/persisting-dlcapt/src/capture/writers/lance_crate.rs b/crates/persisting-dlcapt/src/capture/writers/lance_crate.rs new file mode 100644 index 0000000..1112ea2 --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/writers/lance_crate.rs @@ -0,0 +1,237 @@ +use crate::capture::step_record::StepRecord; +use crate::capture::step_table_writer::{LanceStepRow, StepTableWriter, step_record_to_lance_row}; +use crate::capture::writers::lance_storage::{ + build_object_store_params, lance_dataset_uri, open_dataset, write_params_with_store, +}; +use crate::config::LanceStorageConfig; +use anyhow::{Context, Result}; +use arrow::array::{BooleanArray, Float64Array, Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatchIterator; +use async_trait::async_trait; +use lance::Dataset; +use lance::dataset::{WriteMode, WriteParams}; +use lance::io::ObjectStoreParams; +use std::sync::Arc; + +pub struct LanceCrateWriter { + dataset_uri: String, + schema: Arc, + store_params: Option, +} + +impl LanceCrateWriter { + pub fn new(cfg: &LanceStorageConfig) -> Result { + Ok(Self { + dataset_uri: lance_dataset_uri(&cfg.db_uri, &cfg.table_name), + schema: Arc::new(session_steps_schema()), + store_params: build_object_store_params(cfg), + }) + } + + fn write_params(&self, mode: WriteMode) -> WriteParams { + WriteParams { + mode, + ..write_params_with_store(self.store_params.clone()) + } + } + + async fn append_async(&self, row: &LanceStepRow) -> Result<()> { + if !self.dataset_uri.starts_with("s3://") + && let Some(parent) = std::path::Path::new(&self.dataset_uri).parent() + { + std::fs::create_dir_all(parent) + .with_context(|| format!("create lance db dir {}", parent.display()))?; + } + + let batch = lance_row_to_batch(row, self.schema.clone())?; + let batches = RecordBatchIterator::new(std::iter::once(Ok(batch)), self.schema.clone()); + + if let Some(mut dataset) = open_dataset(&self.dataset_uri, &self.store_params) + .await + .context("probe lance dataset")? + { + dataset + .append(batches, Some(self.write_params(WriteMode::Append))) + .await + .context("append to lance dataset")?; + } else { + Dataset::write( + batches, + self.dataset_uri.as_str(), + Some(self.write_params(WriteMode::Create)), + ) + .await + .context("create lance dataset")?; + } + + Ok(()) + } +} + +#[async_trait] +impl StepTableWriter for LanceCrateWriter { + async fn append(&self, record: &StepRecord) -> Result<()> { + let row = step_record_to_lance_row(record); + self.append_async(&row).await + } +} + +pub fn session_steps_schema() -> Schema { + Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("session_id", DataType::Utf8, false), + Field::new("step_id", DataType::Int64, false), + Field::new("job_id", DataType::Utf8, false), + Field::new("group_id", DataType::Utf8, false), + Field::new("env_name", DataType::Utf8, false), + Field::new("llm_model", DataType::Utf8, false), + Field::new("messages_json", DataType::Utf8, false), + Field::new("response_json", DataType::Utf8, false), + Field::new("step_reward", DataType::Float64, false), + Field::new("reward", DataType::Float64, false), + Field::new("env_state_json", DataType::Utf8, false), + Field::new("is_terminal", DataType::Boolean, false), + Field::new("is_truncated", DataType::Boolean, false), + Field::new("is_session_completed", DataType::Boolean, false), + Field::new("is_trainable", DataType::Boolean, false), + Field::new("created_at", DataType::Utf8, false), + Field::new("agent_id", DataType::Utf8, false), + Field::new("root_session", DataType::Utf8, false), + Field::new("extensions_json", DataType::Utf8, true), + Field::new("capture_json", DataType::Utf8, true), + Field::new("call_id", DataType::Utf8, false), + Field::new("source_export_id", DataType::Int64, true), + ]) +} + +fn lance_row_to_batch(row: &LanceStepRow, schema: Arc) -> Result { + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(vec![row.id.as_str()])), + Arc::new(StringArray::from(vec![row.session_id.as_str()])), + Arc::new(Int64Array::from(vec![row.step_id])), + Arc::new(StringArray::from(vec![row.job_id.as_str()])), + Arc::new(StringArray::from(vec![row.group_id.as_str()])), + Arc::new(StringArray::from(vec![row.env_name.as_str()])), + Arc::new(StringArray::from(vec![row.llm_model.as_str()])), + Arc::new(StringArray::from(vec![row.messages_json.as_str()])), + Arc::new(StringArray::from(vec![row.response_json.as_str()])), + Arc::new(Float64Array::from(vec![row.step_reward])), + Arc::new(Float64Array::from(vec![row.reward])), + Arc::new(StringArray::from(vec![row.env_state_json.as_str()])), + Arc::new(BooleanArray::from(vec![row.is_terminal])), + Arc::new(BooleanArray::from(vec![row.is_truncated])), + Arc::new(BooleanArray::from(vec![row.is_session_completed])), + Arc::new(BooleanArray::from(vec![row.is_trainable])), + Arc::new(StringArray::from(vec![row.created_at.as_str()])), + Arc::new(StringArray::from(vec![row.agent_id.as_str()])), + Arc::new(StringArray::from(vec![row.root_session.as_str()])), + Arc::new(nullable_string_array(&row.extensions_json)), + Arc::new(nullable_string_array(&row.capture_json)), + Arc::new(StringArray::from(vec![row.call_id.as_str()])), + Arc::new(nullable_i64_array(row.source_export_id)), + ], + ) + .context("build lance record batch")?; + Ok(batch) +} + +fn nullable_string_array(value: &Option) -> StringArray { + match value { + Some(text) => StringArray::from(vec![Some(text.as_str())]), + None => StringArray::from(vec![None as Option<&str>]), + } +} + +fn nullable_i64_array(value: Option) -> Int64Array { + Int64Array::from(vec![value]) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::capture::step_record::StepRecord; + use crate::capture::writers::lance_storage::local_dataset_path; + use tempfile::tempdir; + + fn fixture_record() -> StepRecord { + StepRecord { + id: "dlcapt:sess-import-test:1".to_string(), + session_id: "sess-import-test".to_string(), + step_id: 1, + job_id: "dlcapt".to_string(), + agent_id: "openclaw".to_string(), + group_id: String::new(), + env_name: "openclaw".to_string(), + llm_model: "kimi-k2.5".to_string(), + step_reward: 0.0, + reward: 0.0, + is_terminal: false, + is_truncated: false, + is_session_completed: false, + is_trainable: true, + created_at: "2026-06-16 09:57:27.641681+00:00".to_string(), + messages_json: r#"[{"role":"user","content":"ping"}]"#.to_string(), + response_json: r#"{"role":"assistant","content":"pong"}"#.to_string(), + env_state_json: "{}".to_string(), + extensions_json: Some(r#"{"source":"dlcapt-proxy"}"#.to_string()), + capture_json: Some( + r#"{"call_id":"call-import-test-1","finish_reason":"stop"}"#.to_string(), + ), + run_bucket: "2026-06-16".to_string(), + call_id: "call-import-test-1".to_string(), + source_export_id: None, + } + } + + #[tokio::test] + async fn append_creates_and_appends_rows() { + let dir = tempdir().expect("tempdir"); + let cfg = LanceStorageConfig { + db_uri: dir.path().to_string_lossy().to_string(), + table_name: "session_steps".to_string(), + ..LanceStorageConfig::default() + }; + let writer = LanceCrateWriter::new(&cfg).expect("writer"); + let record = fixture_record(); + writer.append(&record).await.expect("first append"); + writer + .append(&StepRecord { + step_id: 2, + id: "dlcapt:sess-import-test:2".to_string(), + ..record.clone() + }) + .await + .expect("second append"); + + let dataset_uri = local_dataset_path(&cfg.db_uri, &cfg.table_name); + assert!(dataset_uri.exists()); + + let dataset = Dataset::open(dataset_uri.to_string_lossy().as_ref()) + .await + .expect("open"); + let count = dataset.count_rows(None).await.expect("count"); + assert_eq!(count, 2); + } + + #[test] + fn s3_writer_builds_object_store_params() { + let cfg = LanceStorageConfig { + db_uri: "s3://my-bucket/capture-prod".to_string(), + s3: Some(crate::config::LanceS3Config { + region: "cn-north-1".to_string(), + endpoint: None, + allow_http: Some(true), + }), + ..LanceStorageConfig::default() + }; + let writer = LanceCrateWriter::new(&cfg).expect("writer"); + assert_eq!( + writer.dataset_uri, + "s3://my-bucket/capture-prod/session_steps.lance" + ); + assert!(writer.store_params.is_some()); + } +} diff --git a/crates/persisting-dlcapt/src/capture/writers/lance_storage.rs b/crates/persisting-dlcapt/src/capture/writers/lance_storage.rs new file mode 100644 index 0000000..668e751 --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/writers/lance_storage.rs @@ -0,0 +1,135 @@ +use crate::config::LanceStorageConfig; +use lance::Dataset; +use lance::Error as LanceError; +use lance::dataset::builder::DatasetBuilder; +use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +pub fn is_s3_db_uri(db_uri: &str) -> bool { + db_uri.starts_with("s3://") +} + +pub fn lance_dataset_uri(db_uri: &str, table_name: &str) -> String { + let base = db_uri.trim_end_matches('/'); + format!("{base}/{table_name}.lance") +} + +pub fn local_dataset_path(db_uri: &str, table_name: &str) -> PathBuf { + PathBuf::from(lance_dataset_uri(db_uri, table_name)) +} + +pub fn build_object_store_params(cfg: &LanceStorageConfig) -> Option { + if !is_s3_db_uri(&cfg.db_uri) { + return None; + } + + let opts = cfg.storage_options(); + if opts.is_empty() { + return Some(ObjectStoreParams::default()); + } + + Some(ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options(opts))), + ..Default::default() + }) +} + +pub fn write_params_with_store( + store_params: Option, +) -> lance::dataset::WriteParams { + lance::dataset::WriteParams { + store_params, + ..Default::default() + } +} + +pub async fn open_dataset( + uri: &str, + store_params: &Option, +) -> lance::Result> { + if !is_s3_db_uri(uri) { + let path = Path::new(uri); + if !path.exists() { + return Ok(None); + } + return Dataset::open(uri).await.map(Some); + } + + let mut builder = DatasetBuilder::from_uri(uri); + if let Some(params) = store_params + && let Some(opts) = params.storage_options() + { + builder = builder.with_storage_options(opts.clone()); + } + + match builder.load().await { + Ok(dataset) => Ok(Some(dataset)), + Err(err) if is_dataset_missing(&err) => Ok(None), + Err(err) => Err(err), + } +} + +pub fn is_dataset_missing(err: &LanceError) -> bool { + matches!( + err, + LanceError::DatasetNotFound { .. } | LanceError::NotFound { .. } + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::LanceS3Config; + + #[test] + fn lance_dataset_uri_joins_table_name() { + assert_eq!( + lance_dataset_uri("s3://bucket/prefix", "session_steps"), + "s3://bucket/prefix/session_steps.lance" + ); + assert_eq!( + lance_dataset_uri("../../var/lance/local", "session_steps"), + "../../var/lance/local/session_steps.lance" + ); + } + + #[test] + fn storage_options_include_region_and_endpoint() { + let cfg = LanceStorageConfig { + db_uri: "s3://bucket/prefix".to_string(), + s3: Some(LanceS3Config { + region: "cn-north-1".to_string(), + endpoint: Some("https://minio.local".to_string()), + allow_http: Some(true), + }), + ..LanceStorageConfig::default() + }; + let opts = cfg.storage_options(); + assert_eq!( + opts.get("aws_region").map(String::as_str), + Some("cn-north-1") + ); + assert_eq!( + opts.get("aws_endpoint").map(String::as_str), + Some("https://minio.local") + ); + assert_eq!(opts.get("allow_http").map(String::as_str), Some("true")); + assert!(build_object_store_params(&cfg).is_some()); + } + + #[test] + fn storage_options_set_allow_http_for_http_endpoint() { + let cfg = LanceStorageConfig { + db_uri: "s3://bucket/prefix".to_string(), + s3: Some(LanceS3Config { + region: "cn-north-1".to_string(), + endpoint: Some("http://ssd1.h.pjlab.org.cn:8060".to_string()), + allow_http: Some(true), + }), + ..LanceStorageConfig::default() + }; + let opts = cfg.storage_options(); + assert_eq!(opts.get("allow_http").map(String::as_str), Some("true")); + } +} diff --git a/crates/persisting-dlcapt/src/capture/writers/mod.rs b/crates/persisting-dlcapt/src/capture/writers/mod.rs new file mode 100644 index 0000000..0f51e73 --- /dev/null +++ b/crates/persisting-dlcapt/src/capture/writers/mod.rs @@ -0,0 +1,2 @@ +pub mod lance_crate; +pub mod lance_storage; diff --git a/crates/persisting-dlcapt/src/config.rs b/crates/persisting-dlcapt/src/config.rs new file mode 100644 index 0000000..f6d2705 --- /dev/null +++ b/crates/persisting-dlcapt/src/config.rs @@ -0,0 +1,450 @@ +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; +use std::collections::HashMap; +use std::fs; +use std::path::Path; +use tracing::warn; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProxyConfig { + #[serde(default = "default_listen")] + pub listen: String, + #[serde(default = "default_admin_listen")] + pub admin_listen: String, + #[serde(default = "default_store_dir")] + pub store_dir: String, + #[serde(default = "default_agent_id")] + pub agent_id: String, + #[serde(default = "default_session_header")] + pub session_header: String, + #[serde(default)] + pub session_header_aliases: Vec, + #[serde(default = "default_session_id")] + pub default_session_id: String, + #[serde(default = "default_preserve_raw")] + pub preserve_raw: bool, + #[serde(default = "default_base_session_path")] + pub base_session_path: String, + #[serde(default)] + pub models: Vec, + #[serde(default)] + pub storage: StorageConfig, + #[serde(default)] + pub export: ExportConfig, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ModelRoute { + pub name: String, + pub provider: String, + pub upstream_base_url: String, + pub api_key: Option, + pub display_name: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StorageConfig { + #[serde(default = "default_authoritative")] + pub authoritative: String, + #[serde(default = "default_also_md")] + pub also: Vec, + #[serde(default)] + pub json_cache: JsonCacheConfig, + #[serde(default)] + pub lance: LanceStorageConfig, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LanceStorageConfig { + #[serde(default, alias = "uri")] + pub db_uri: String, + #[serde(default = "default_lance_table_name")] + pub table_name: String, + #[serde(default = "default_lance_backend")] + pub backend: String, + #[serde(default = "default_lance_mode")] + pub mode: String, + #[serde(default = "default_fail_open")] + pub fail_open: bool, + #[serde(default = "default_dead_letter_path")] + pub dead_letter_path: String, + #[serde(default)] + pub s3: Option, + #[serde(default = "default_batch_size")] + pub batch_size: u32, + #[serde(default)] + pub flush_interval_ms: u64, + #[serde(default)] + pub write_timeout_ms: u64, + #[serde(default)] + pub async_writer: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LanceS3Config { + pub region: String, + #[serde(default)] + pub endpoint: Option, + /// When unset, defaults to true if `endpoint` uses `http://`. + #[serde(default)] + pub allow_http: Option, +} + +impl LanceS3Config { + fn allow_http_enabled(&self) -> bool { + if let Some(allow_http) = self.allow_http { + return allow_http; + } + self.endpoint + .as_ref() + .is_some_and(|endpoint| endpoint.to_ascii_lowercase().starts_with("http://")) + } + + pub fn to_storage_options(&self) -> HashMap { + let mut opts = HashMap::new(); + if !self.region.is_empty() { + opts.insert("aws_region".to_string(), self.region.clone()); + } + if let Some(endpoint) = &self.endpoint { + if !endpoint.is_empty() { + opts.insert("aws_endpoint".to_string(), endpoint.clone()); + } + } + if self.allow_http_enabled() { + opts.insert("allow_http".to_string(), "true".to_string()); + } + opts + } +} + +impl LanceStorageConfig { + pub fn is_s3(&self) -> bool { + self.db_uri.starts_with("s3://") + } + + pub fn storage_options(&self) -> HashMap { + let mut opts = self + .s3 + .as_ref() + .map(LanceS3Config::to_storage_options) + .unwrap_or_default(); + if !opts.contains_key("aws_region") + && let Ok(region) = std::env::var("AWS_REGION") + && !region.is_empty() + { + opts.insert("aws_region".to_string(), region); + } + opts + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct JsonCacheConfig { + #[serde(default)] + pub enabled: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExportConfig { + #[serde(default = "default_max_steps")] + pub max_steps_per_session: u64, + #[serde(default = "default_messages_drift")] + pub messages_drift: String, + #[serde(default)] + pub defaults: ExportDefaults, + #[serde(default)] + pub session_metadata: Map, +} + +impl Default for ExportConfig { + fn default() -> Self { + Self { + max_steps_per_session: default_max_steps(), + messages_drift: default_messages_drift(), + defaults: ExportDefaults::default(), + session_metadata: Map::new(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExportDefaults { + #[serde(default = "default_env_name")] + pub env_name: String, + #[serde(default = "default_job_id")] + pub job_id: String, + #[serde(default)] + pub group_id: String, + #[serde(default)] + pub step_reward: f64, + #[serde(default)] + pub reward: f64, + #[serde(default = "default_env_state")] + pub env_state: Value, + #[serde(default)] + pub is_terminal: bool, + #[serde(default = "default_is_trainable")] + pub is_trainable: bool, +} + +impl Default for ExportDefaults { + fn default() -> Self { + Self { + env_name: default_env_name(), + job_id: default_job_id(), + group_id: String::new(), + step_reward: 0.0, + reward: 0.0, + env_state: default_env_state(), + is_terminal: false, + is_trainable: true, + } + } +} + +impl StorageConfig { + pub fn lance_enabled(&self) -> bool { + self.authoritative == "lance" || self.also.iter().any(|s| s == "lance") + } + + pub fn json_cache_enabled(&self) -> bool { + self.authoritative == "lance" + && (self.also.iter().any(|s| s == "json_cache") || self.json_cache.enabled) + } + + pub fn validate(&self) -> Result<()> { + if !VALID_AUTHORITATIVE.contains(&self.authoritative.as_str()) { + bail!( + "storage.authoritative must be one of: {}", + VALID_AUTHORITATIVE.join(", ") + ); + } + + for token in &self.also { + if !VALID_ALSO_TOKENS.contains(&token.as_str()) { + bail!( + "storage.also contains unknown token '{token}'; allowed: {}", + VALID_ALSO_TOKENS.join(", ") + ); + } + } + + if self.authoritative != "lance" + && (self.also.iter().any(|s| s == "json_cache") || self.json_cache.enabled) + { + bail!("storage.json_cache is only valid when storage.authoritative = \"lance\""); + } + + if self.lance_enabled() { + if self.lance.db_uri.trim().is_empty() { + bail!("storage.lance.db_uri is required when Lance sink is enabled"); + } + if self.lance.backend != "lance" { + bail!( + "storage.lance.backend must be \"lance\" in P0 (got \"{}\")", + self.lance.backend + ); + } + if !VALID_LANCE_MODES.contains(&self.lance.mode.as_str()) { + bail!( + "storage.lance.mode must be one of: {}", + VALID_LANCE_MODES.join(", ") + ); + } + if self.lance.is_s3() { + let has_region = self + .lance + .s3 + .as_ref() + .is_some_and(|s| !s.region.trim().is_empty()) + || std::env::var("AWS_REGION") + .map(|r| !r.trim().is_empty()) + .unwrap_or(false); + if !has_region { + bail!( + "storage.lance.s3.region is required when db_uri uses s3:// (unless AWS_REGION is set)" + ); + } + } + } + + if self.authoritative == "lance" && self.also.iter().any(|s| s == "lance") { + warn!( + "storage.also contains redundant \"lance\" when authoritative is already \"lance\"" + ); + } + + Ok(()) + } +} + +impl ProxyConfig { + pub fn load(path: &Path) -> Result { + let raw = fs::read_to_string(path) + .with_context(|| format!("failed reading config file: {}", path.display()))?; + let mut parsed: ProxyConfig = + toml::from_str(&raw).with_context(|| "failed parsing proxy.toml".to_string())?; + if let Ok(uri) = + std::env::var("DLCAPT_LANCE_DB_URI").or_else(|_| std::env::var("CAPTURE_LANCE_URI")) + { + if !uri.trim().is_empty() { + parsed.storage.lance.db_uri = uri; + } + } + parsed.validate()?; + Ok(parsed) + } + + fn validate(&self) -> Result<()> { + if self.models.is_empty() { + bail!("proxy config must contain at least one [[models]] route"); + } + + for model in &self.models { + if model.name.trim().is_empty() { + bail!("model route name cannot be empty"); + } + if model.provider.trim().is_empty() { + bail!("model provider cannot be empty"); + } + if model.upstream_base_url.trim().is_empty() { + bail!("model upstream_base_url cannot be empty"); + } + } + + self.storage.validate()?; + + Ok(()) + } + + pub fn session_settings(&self) -> crate::session::SessionIdSettings { + crate::session::SessionIdSettings { + default_session_id: self.default_session_id.clone(), + preserve_raw: self.preserve_raw, + session_header: self.session_header.clone(), + session_header_aliases: self.session_header_aliases.clone(), + } + } +} + +fn default_listen() -> String { + "127.0.0.1:19081".to_string() +} + +fn default_admin_listen() -> String { + "127.0.0.1:19082".to_string() +} + +fn default_store_dir() -> String { + "store".to_string() +} + +fn default_agent_id() -> String { + "openclaw".to_string() +} + +fn default_session_header() -> String { + "x-persisting-session-id".to_string() +} + +fn default_session_id() -> String { + "default".to_string() +} + +fn default_preserve_raw() -> bool { + false +} + +fn default_base_session_path() -> String { + "/v1/sessions".to_string() +} + +fn default_authoritative() -> String { + "json_file".to_string() +} + +fn default_also_md() -> Vec { + vec!["md".to_string()] +} + +fn default_max_steps() -> u64 { + 1000 +} + +fn default_messages_drift() -> String { + "trust_request".to_string() +} + +fn default_env_name() -> String { + "openclaw".to_string() +} + +fn default_job_id() -> String { + "dlcapt".to_string() +} + +fn default_env_state() -> Value { + Value::Object(Map::new()) +} + +fn default_is_trainable() -> bool { + true +} + +const VALID_AUTHORITATIVE: &[&str] = &["lance", "json_file", "md"]; +const VALID_ALSO_TOKENS: &[&str] = &["lance", "md", "json_cache"]; +const VALID_LANCE_MODES: &[&str] = &["create", "append", "overwrite"]; + +impl Default for StorageConfig { + fn default() -> Self { + Self { + authoritative: default_authoritative(), + also: default_also_md(), + json_cache: JsonCacheConfig::default(), + lance: LanceStorageConfig::default(), + } + } +} + +impl Default for LanceStorageConfig { + fn default() -> Self { + Self { + db_uri: String::new(), + table_name: default_lance_table_name(), + backend: default_lance_backend(), + mode: default_lance_mode(), + fail_open: default_fail_open(), + dead_letter_path: default_dead_letter_path(), + s3: None, + batch_size: default_batch_size(), + flush_interval_ms: 0, + write_timeout_ms: 0, + async_writer: false, + } + } +} + +fn default_lance_table_name() -> String { + "session_steps".to_string() +} + +fn default_lance_backend() -> String { + "lance".to_string() +} + +fn default_lance_mode() -> String { + "append".to_string() +} + +fn default_fail_open() -> bool { + true +} + +fn default_dead_letter_path() -> String { + ".capture/lance_dead_letter.jsonl".to_string() +} + +fn default_batch_size() -> u32 { + 1 +} diff --git a/crates/persisting-dlcapt/src/dialogue/mod.rs b/crates/persisting-dlcapt/src/dialogue/mod.rs new file mode 100644 index 0000000..b39e98b --- /dev/null +++ b/crates/persisting-dlcapt/src/dialogue/mod.rs @@ -0,0 +1,60 @@ +mod responses; + +use serde_json::Value; + +pub use responses::{ + extract_user_from_responses_input, summarize_responses_json_response, + summarize_responses_sse_response, +}; + +pub fn extract_last_user_message(body: &Value) -> Option { + let messages = body.get("messages")?.as_array()?; + for msg in messages.iter().rev() { + if msg.get("role").and_then(Value::as_str) != Some("user") { + continue; + } + if let Some(text) = msg.get("content").and_then(Value::as_str) { + let trimmed = text.trim(); + if !trimmed.is_empty() { + return Some(text.to_string()); + } + continue; + } + if let Some(parts) = msg.get("content").and_then(Value::as_array) { + let mut texts = Vec::new(); + for part in parts { + if let Some(text) = part.get("text").and_then(Value::as_str) { + if !text.trim().is_empty() { + texts.push(text.to_string()); + } + } + } + if !texts.is_empty() { + return Some(texts.join("\n")); + } + } + } + None +} + +pub fn extract_user_text(endpoint: InferenceEndpoint, body: &Value) -> Option { + match endpoint { + InferenceEndpoint::ChatCompletions => extract_last_user_message(body), + InferenceEndpoint::Responses => extract_user_from_responses_input(body), + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InferenceEndpoint { + ChatCompletions, + Responses, +} + +impl InferenceEndpoint { + pub fn upstream_suffix(self) -> &'static str { + match self { + Self::ChatCompletions => "chat/completions", + Self::Responses => "responses", + } + } +} diff --git a/crates/persisting-dlcapt/src/dialogue/responses.rs b/crates/persisting-dlcapt/src/dialogue/responses.rs new file mode 100644 index 0000000..abac22b --- /dev/null +++ b/crates/persisting-dlcapt/src/dialogue/responses.rs @@ -0,0 +1,176 @@ +use serde_json::Value; + +pub fn extract_user_from_responses_input(body: &Value) -> Option { + let input = body.get("input")?; + match input { + Value::String(text) if !text.trim().is_empty() => Some(text.to_string()), + Value::Array(items) => extract_tail_user_text(items), + _ => None, + } +} + +fn extract_tail_user_text(items: &[Value]) -> Option { + for item in items.iter().rev() { + let item_type = item.get("type").and_then(Value::as_str).unwrap_or(""); + if item_type == "function_call_output" { + continue; + } + if item.get("role").and_then(Value::as_str) == Some("user") { + if let Some(text) = content_to_text(item.get("content")) { + return Some(text); + } + } + if let Some(text) = content_to_text(Some(item)) { + return Some(text); + } + } + None +} + +fn content_to_text(content: Option<&Value>) -> Option { + let content = content?; + match content { + Value::String(text) if !text.trim().is_empty() => Some(text.to_string()), + Value::Array(parts) => { + let mut texts = Vec::new(); + for part in parts { + if let Some(text) = part.get("text").and_then(Value::as_str) { + if !text.trim().is_empty() { + texts.push(text.to_string()); + } + } else if part.get("type").and_then(Value::as_str) == Some("input_text") { + if let Some(text) = part.get("text").and_then(Value::as_str) { + if !text.trim().is_empty() { + texts.push(text.to_string()); + } + } + } + } + if texts.is_empty() { + None + } else { + Some(texts.join("\n")) + } + } + Value::Object(map) if map.contains_key("content") => content_to_text(map.get("content")), + _ => None, + } +} + +pub fn summarize_responses_json_response(value: &Value) -> (Option, Option) { + let usage = value.get("usage").cloned(); + if let Some(text) = value.get("output_text").and_then(Value::as_str) { + if !text.is_empty() { + return (Some(text.to_string()), usage); + } + } + + let mut texts = Vec::new(); + if let Some(output) = value.get("output").and_then(Value::as_array) { + for item in output { + match item.get("type").and_then(Value::as_str) { + Some("message") => { + if let Some(parts) = item.get("content").and_then(Value::as_array) { + for part in parts { + if part.get("type").and_then(Value::as_str) == Some("output_text") { + if let Some(text) = part.get("text").and_then(Value::as_str) { + texts.push(text.to_string()); + } + } + } + } + } + Some("output_text") => { + if let Some(text) = item.get("text").and_then(Value::as_str) { + texts.push(text.to_string()); + } + } + _ => {} + } + } + } + + let response_text = if texts.is_empty() { + None + } else { + Some(texts.join("\n")) + }; + (response_text, usage) +} + +pub fn summarize_responses_sse_response(raw: &str) -> (Option, Option) { + let mut content = String::new(); + let mut usage = None; + + for line in raw.lines() { + let trimmed = line.trim(); + if !trimmed.starts_with("data:") { + continue; + } + let data = trimmed.trim_start_matches("data:").trim(); + if data.is_empty() || data == "[DONE]" { + continue; + } + let chunk = match serde_json::from_str::(data) { + Ok(value) => value, + Err(_) => continue, + }; + + match chunk.get("type").and_then(Value::as_str) { + Some("response.output_text.delta") => { + if let Some(delta) = chunk.get("delta").and_then(Value::as_str) { + content.push_str(delta); + } + } + Some("response.output_text.done") => { + if content.is_empty() { + if let Some(text) = chunk.get("text").and_then(Value::as_str) { + content.push_str(text); + } + } + } + _ => {} + } + + if usage.is_none() { + usage = chunk + .get("response") + .and_then(|response| response.get("usage")) + .cloned() + .or_else(|| chunk.get("usage").cloned()); + } + } + + let response_text = if content.is_empty() { + None + } else { + Some(content) + }; + (response_text, usage) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn extract_user_from_string_input() { + let body = json!({"input": "hello responses"}); + assert_eq!( + extract_user_from_responses_input(&body).as_deref(), + Some("hello responses") + ); + } + + #[test] + fn summarize_responses_json_should_read_output_text() { + let body = json!({ + "output_text": "assistant reply", + "usage": {"total_tokens": 10} + }); + let (text, usage) = summarize_responses_json_response(&body); + assert_eq!(text.as_deref(), Some("assistant reply")); + assert!(usage.is_some()); + } +} diff --git a/crates/persisting-dlcapt/src/lib.rs b/crates/persisting-dlcapt/src/lib.rs index 6a814c2..e803d9e 100644 --- a/crates/persisting-dlcapt/src/lib.rs +++ b/crates/persisting-dlcapt/src/lib.rs @@ -1 +1,8 @@ -//! persisting-dlcapt library stub; replaced by migrated sources in Task 2. +pub mod audit; +pub mod capture; +pub mod config; +pub mod dialogue; +pub mod proxy; +pub mod router; +pub mod session; +pub mod tlv; diff --git a/crates/persisting-dlcapt/src/main.rs b/crates/persisting-dlcapt/src/main.rs index a0d32fe..32c961e 100644 --- a/crates/persisting-dlcapt/src/main.rs +++ b/crates/persisting-dlcapt/src/main.rs @@ -1,4 +1,53 @@ -fn main() { - eprintln!("persisting-dlcapt stub; run Task 2 migration before use"); - std::process::exit(1); +use anyhow::{Context, Result}; +use persisting_dlcapt::config::ProxyConfig; +use persisting_dlcapt::proxy::{AppState, build_admin_router, build_public_router}; +use std::net::SocketAddr; +use std::path::PathBuf; +use tracing::info; + +#[tokio::main] +async fn main() -> Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "dlcapt=info,tower_http=info".into()), + ) + .init(); + + let config_path = std::env::args() + .nth(1) + .unwrap_or_else(|| "config/proxy.toml".to_string()); + let config_path = PathBuf::from(config_path); + let config = ProxyConfig::load(&config_path) + .with_context(|| format!("failed loading config from {}", config_path.display()))?; + + let state = AppState::new(config); + let public_router = build_public_router(state.clone()); + let admin_router = build_admin_router(state.clone()); + + let public_addr: SocketAddr = state + .listen_addr() + .parse() + .with_context(|| "invalid listen address".to_string())?; + let admin_addr: SocketAddr = state + .admin_listen_addr() + .parse() + .with_context(|| "invalid admin_listen address".to_string())?; + + info!("dlcapt public API listening on {public_addr}"); + info!("dlcapt admin API listening on {admin_addr}"); + + let public_listener = tokio::net::TcpListener::bind(public_addr) + .await + .with_context(|| "failed binding public listener".to_string())?; + let admin_listener = tokio::net::TcpListener::bind(admin_addr) + .await + .with_context(|| "failed binding admin listener".to_string())?; + + let public_server = axum::serve(public_listener, public_router); + let admin_server = axum::serve(admin_listener, admin_router); + + tokio::try_join!(public_server, admin_server) + .with_context(|| "proxy server exited with error".to_string())?; + Ok(()) } diff --git a/crates/persisting-dlcapt/src/proxy.rs b/crates/persisting-dlcapt/src/proxy.rs new file mode 100644 index 0000000..749d8f1 --- /dev/null +++ b/crates/persisting-dlcapt/src/proxy.rs @@ -0,0 +1,726 @@ +use crate::capture::{CaptureEvent, CaptureMeta, CaptureSinkRouter, PostProcessorChain}; +use crate::config::ProxyConfig; +use crate::dialogue::InferenceEndpoint; +use crate::router::RouteTable; +use crate::session::{ + RequestContext, RouteSessionMode, SessionResolveError, SessionSource, resolve_session, +}; +use crate::tlv::new_call_id; +use axum::body::{Body, Bytes}; +use axum::extract::{Path, State}; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use chrono::Utc; +use serde_json::{Value, json}; +use std::collections::{BTreeMap, HashMap, VecDeque}; +use std::path::PathBuf; +use std::sync::Arc; +use tokio::sync::{Mutex, RwLock, mpsc}; +use tokio_stream::wrappers::ReceiverStream; + +#[derive(Debug, Default)] +struct SessionTrack { + requests: u64, + next_seq: u64, + next_turn: u64, +} + +#[derive(Debug, Clone)] +struct InferenceCall { + endpoint: InferenceEndpoint, + route_mode: RouteSessionMode, + path_session_id: Option, + request_path: String, +} + +#[derive(Clone)] +pub struct AppState { + client: reqwest::Client, + config: Arc, + routes: Arc, + capture_sink: Arc, + post_processors: Arc, + sessions: Arc>>, + errors: Arc>>, +} + +impl AppState { + pub fn new(config: ProxyConfig) -> Self { + let client = reqwest::Client::new(); + let routes = RouteTable::from_config(&config); + let tlv_root = PathBuf::from(config.store_dir.clone()); + let agent_id = config.agent_id.clone(); + let default_session_id = config.default_session_id.clone(); + let write_lock = Arc::new(Mutex::new(())); + let tlv = crate::tlv::TlvWriter::new( + tlv_root, + agent_id, + default_session_id, + Arc::clone(&write_lock), + ); + let config = Arc::new(config); + let capture_sink = Arc::new( + CaptureSinkRouter::new(Arc::clone(&config), tlv, write_lock) + .expect("failed initializing capture sink router"), + ); + Self { + client, + config, + routes: Arc::new(routes), + capture_sink, + post_processors: Arc::new(PostProcessorChain::empty()), + sessions: Arc::new(RwLock::new(HashMap::new())), + errors: Arc::new(Mutex::new(VecDeque::new())), + } + } + + pub fn listen_addr(&self) -> &str { + &self.config.listen + } + + pub fn admin_listen_addr(&self) -> &str { + &self.config.admin_listen + } + + async fn alloc_turn(&self, session_id: &str) -> (u64, u64, u64) { + let mut sessions = self.sessions.write().await; + let track = sessions + .entry(session_id.to_string()) + .or_insert_with(|| SessionTrack { + requests: 0, + next_seq: 0, + next_turn: 1, + }); + track.requests += 1; + let turn = track.next_turn; + let user_seq = track.next_seq; + let assistant_seq = track.next_seq + 1; + track.next_seq += 2; + track.next_turn += 1; + (user_seq, assistant_seq, turn) + } + + async fn push_error(&self, err: String) { + let mut errors = self.errors.lock().await; + if errors.len() >= 100 { + errors.pop_front(); + } + errors.push_back(err); + } +} + +pub fn build_public_router(state: AppState) -> Router { + Router::new() + .route("/healthz", get(healthz)) + .route("/readyz", get(readyz)) + .route("/v1/models", get(v1_models)) + .route("/v1/chat/completions", post(flat_chat_completions)) + .route( + "/v1/sessions/{session_id}/chat/completions", + post(session_chat_completions), + ) + .route( + "/v1/sessions/{session_id}/responses", + post(session_responses), + ) + .with_state(state) +} + +pub fn build_admin_router(state: AppState) -> Router { + Router::new() + .route("/healthz", get(healthz)) + .route("/readyz", get(readyz)) + .route("/admin/sessions", get(admin_sessions)) + .route("/admin/errors", get(admin_errors)) + .with_state(state) +} + +async fn healthz() -> impl IntoResponse { + Json(json!({"status": "ok"})) +} + +async fn readyz() -> impl IntoResponse { + Json(json!({"status": "ready"})) +} + +async fn admin_sessions(State(state): State) -> impl IntoResponse { + let sessions = state.sessions.read().await; + let mut rows = sessions + .iter() + .map(|(session_id, track)| { + json!({ + "session_id": session_id, + "requests": track.requests, + "turns": track.next_turn.saturating_sub(1), + }) + }) + .collect::>(); + rows.sort_by(|a, b| { + a.get("session_id") + .and_then(Value::as_str) + .unwrap_or_default() + .cmp( + b.get("session_id") + .and_then(Value::as_str) + .unwrap_or_default(), + ) + }); + Json(json!({ "sessions": rows })) +} + +async fn admin_errors(State(state): State) -> impl IntoResponse { + let errors = state.errors.lock().await; + let rows = errors.iter().cloned().collect::>(); + Json(json!({ "recent_errors": rows })) +} + +async fn v1_models(State(state): State) -> impl IntoResponse { + let data = state + .routes + .list_models() + .into_iter() + .map(|model| { + json!({ + "id": model.id, + "object": "model", + "owned_by": model.provider, + "display_name": model.display_name, + }) + }) + .collect::>(); + Json(json!({ + "object": "list", + "data": data + })) +} + +async fn flat_chat_completions( + State(state): State, + headers: HeaderMap, + Json(payload): Json, +) -> Response { + handle_inference( + state, + headers, + payload, + InferenceCall { + endpoint: InferenceEndpoint::ChatCompletions, + route_mode: RouteSessionMode::Flat, + path_session_id: None, + request_path: "/v1/chat/completions".to_string(), + }, + ) + .await +} + +async fn session_chat_completions( + State(state): State, + Path(session_id): Path, + headers: HeaderMap, + Json(payload): Json, +) -> Response { + let request_path = build_session_request_path( + &state.config.base_session_path, + &session_id, + InferenceEndpoint::ChatCompletions, + ); + handle_inference( + state, + headers, + payload, + InferenceCall { + endpoint: InferenceEndpoint::ChatCompletions, + route_mode: RouteSessionMode::SessionScoped, + path_session_id: Some(session_id), + request_path, + }, + ) + .await +} + +async fn session_responses( + State(state): State, + Path(session_id): Path, + headers: HeaderMap, + Json(payload): Json, +) -> Response { + let request_path = build_session_request_path( + &state.config.base_session_path, + &session_id, + InferenceEndpoint::Responses, + ); + handle_inference( + state, + headers, + payload, + InferenceCall { + endpoint: InferenceEndpoint::Responses, + route_mode: RouteSessionMode::SessionScoped, + path_session_id: Some(session_id), + request_path, + }, + ) + .await +} + +fn build_session_request_path( + base_session_path: &str, + session_id: &str, + endpoint: InferenceEndpoint, +) -> String { + format!( + "{}/{}/{}", + base_session_path.trim_end_matches('/'), + session_id, + endpoint.upstream_suffix() + ) +} + +async fn handle_inference( + state: AppState, + headers: HeaderMap, + payload: Value, + call: InferenceCall, +) -> Response { + let model = match payload.get("model").and_then(Value::as_str) { + Some(model) => model, + None => { + return error_response(StatusCode::BAD_REQUEST, "missing field: model"); + } + }; + + let route = match state.routes.resolve_model(model).cloned() { + Some(route) => route, + None => { + return error_response(StatusCode::BAD_REQUEST, "unknown model route"); + } + }; + + let ctx = RequestContext { + mode: call.route_mode, + path_session_id: call.path_session_id.as_deref(), + headers: &headers, + body: &payload, + }; + let resolved = match resolve_session(&ctx, &state.config.session_settings()) { + Ok(resolved) => resolved, + Err(SessionResolveError::MissingPathSessionId) => { + return error_response( + StatusCode::BAD_REQUEST, + "session_id must be supplied in URL path", + ); + } + Err(SessionResolveError::InvalidPathSessionId) => { + return error_response(StatusCode::BAD_REQUEST, "invalid session_id in URL path"); + } + Err(SessionResolveError::MissingSessionId | SessionResolveError::InvalidSessionId) => { + return error_response(StatusCode::BAD_REQUEST, "unable to resolve session_id"); + } + }; + + if resolved.source == SessionSource::Default { + tracing::warn!( + session_id = %resolved.storage_session_id, + "using default_session_id; inject session-scoped URL, header, or metadata.session_id" + ); + } + + let session_id = resolved.storage_session_id; + let upstream_url = format!( + "{}/{}", + route.upstream_base_url.trim_end_matches('/'), + call.endpoint.upstream_suffix() + ); + let mut request_builder = state.client.post(upstream_url).json(&payload); + if let Some(api_key) = route.api_key.as_ref().filter(|key| !key.trim().is_empty()) { + request_builder = request_builder.bearer_auth(api_key); + } + + let upstream_response = match request_builder.send().await { + Ok(response) => response, + Err(err) => { + state + .push_error(format!("upstream request failed: {err}")) + .await; + return error_response(StatusCode::BAD_GATEWAY, "upstream request failed"); + } + }; + + let status = StatusCode::from_u16(upstream_response.status().as_u16()) + .unwrap_or(StatusCode::BAD_GATEWAY); + let content_type = upstream_response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string); + let is_stream = payload + .get("stream") + .and_then(Value::as_bool) + .unwrap_or(false); + + if is_stream { + let (tx, rx) = mpsc::channel::>(16); + let stream_state = state.clone(); + let stream_session_id = session_id; + let stream_model = model.to_string(); + let stream_payload = payload; + let stream_call = call.clone(); + let mut stream_upstream = upstream_response; + let stream_status = status; + tokio::spawn(async move { + let mut raw = Vec::new(); + loop { + match stream_upstream.chunk().await { + Ok(Some(chunk)) => { + raw.extend_from_slice(&chunk); + if tx.send(Ok(chunk)).await.is_err() { + break; + } + } + Ok(None) => break, + Err(err) => { + stream_state + .push_error(format!("stream chunk read failed: {err}")) + .await; + let io_err = + std::io::Error::other(format!("stream chunk read failed: {err}")); + let _ = tx.send(Err(io_err)).await; + break; + } + } + } + + let raw_text = String::from_utf8_lossy(&raw).into_owned(); + let (response_text, usage, finish_reason) = + summarize_stream_response(stream_call.endpoint, &raw_text); + let (user_seq, assistant_seq, turn) = stream_state.alloc_turn(&stream_session_id).await; + let call_id = new_call_id(); + persist_inference( + &stream_state, + &stream_call, + stream_session_id, + stream_model, + stream_payload, + HeaderMap::new(), + stream_status.as_u16(), + true, + response_json_from_stream(stream_call.endpoint, &raw_text), + response_text, + usage, + finish_reason, + user_seq, + assistant_seq, + turn, + call_id, + ) + .await; + }); + + let relay_stream = ReceiverStream::new(rx); + let mut builder = Response::builder().status(status); + if let Some(content_type) = content_type.as_deref() { + builder = builder.header(axum::http::header::CONTENT_TYPE, content_type); + } else { + builder = builder.header(axum::http::header::CONTENT_TYPE, "text/event-stream"); + } + + return match builder.body(Body::from_stream(relay_stream)) { + Ok(response) => response, + Err(err) => { + state + .push_error(format!("building stream response failed: {err}")) + .await; + error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to build stream response", + ) + } + }; + } + + let body_text = match upstream_response.text().await { + Ok(text) => text, + Err(err) => { + state + .push_error(format!("read upstream body failed: {err}")) + .await; + return error_response(StatusCode::BAD_GATEWAY, "failed reading upstream response"); + } + }; + + let response_json = serde_json::from_str::(&body_text).unwrap_or_else(|_| { + json!({ + "raw_response": truncate_text(&body_text, 1200) + }) + }); + let (response_text, usage, finish_reason) = + summarize_json_response(call.endpoint, &response_json); + let (user_seq, assistant_seq, turn) = state.alloc_turn(&session_id).await; + let call_id = new_call_id(); + persist_inference( + &state, + &call, + session_id, + model.to_string(), + payload, + headers, + status.as_u16(), + false, + response_json, + response_text, + usage, + finish_reason, + user_seq, + assistant_seq, + turn, + call_id, + ) + .await; + + let mut builder = Response::builder().status(status); + if let Some(content_type) = content_type.as_deref() { + builder = builder.header(axum::http::header::CONTENT_TYPE, content_type); + } else { + builder = builder.header(axum::http::header::CONTENT_TYPE, "application/json"); + } + + match builder.body(Body::from(body_text)) { + Ok(response) => response, + Err(err) => { + state + .push_error(format!("building json response failed: {err}")) + .await; + error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "failed to build upstream response", + ) + } + } +} + +fn summarize_json_response( + endpoint: InferenceEndpoint, + value: &Value, +) -> (Option, Option, Option) { + match endpoint { + InferenceEndpoint::ChatCompletions => { + let finish_reason = value + .get("choices") + .and_then(Value::as_array) + .and_then(|choices| choices.first()) + .and_then(|choice| choice.get("finish_reason")) + .and_then(Value::as_str) + .map(ToString::to_string); + + let response_text = value + .get("choices") + .and_then(Value::as_array) + .and_then(|choices| choices.first()) + .and_then(|choice| { + choice + .get("message") + .and_then(|message| message.get("content")) + .and_then(Value::as_str) + .map(ToString::to_string) + .or_else(|| { + choice + .get("text") + .and_then(Value::as_str) + .map(ToString::to_string) + }) + }) + .or_else(|| { + value + .get("output_text") + .and_then(Value::as_str) + .map(ToString::to_string) + }); + (response_text, value.get("usage").cloned(), finish_reason) + } + InferenceEndpoint::Responses => { + let (response_text, usage) = crate::dialogue::summarize_responses_json_response(value); + (response_text, usage, None) + } + } +} + +fn summarize_stream_response( + endpoint: InferenceEndpoint, + raw: &str, +) -> (Option, Option, Option) { + match endpoint { + InferenceEndpoint::ChatCompletions => { + let (response_text, finish_reason, usage) = parse_chat_sse_response(raw); + (response_text, usage, finish_reason) + } + InferenceEndpoint::Responses => { + let (response_text, usage) = crate::dialogue::summarize_responses_sse_response(raw); + (response_text, usage, None) + } + } +} + +fn response_json_from_stream(endpoint: InferenceEndpoint, raw: &str) -> Value { + match endpoint { + InferenceEndpoint::ChatCompletions => { + let (response_text, finish_reason, usage) = parse_chat_sse_response(raw); + json!({ + "choices": [{ + "message": { + "role": "assistant", + "content": response_text, + }, + "finish_reason": finish_reason, + }], + "usage": usage, + }) + } + InferenceEndpoint::Responses => { + json!({"output_text": summarize_stream_response(endpoint, raw).0}) + } + } +} + +fn headers_to_map(headers: &HeaderMap) -> BTreeMap { + let mut out = BTreeMap::new(); + for (key, value) in headers.iter() { + if let Ok(text) = value.to_str() { + out.insert(key.as_str().to_string(), text.to_string()); + } + } + out +} + +async fn persist_inference( + state: &AppState, + call: &InferenceCall, + session_id: String, + model: String, + request: Value, + headers: HeaderMap, + status_code: u16, + stream: bool, + response_raw: Value, + response_text: Option, + usage: Option, + finish_reason: Option, + user_seq: u64, + assistant_seq: u64, + turn: u64, + call_id: String, +) { + let mut event = CaptureEvent { + call_id, + session_id, + agent_id: state.config.agent_id.clone(), + step_id: turn, + turn, + endpoint: call.endpoint, + request_path: call.request_path.clone(), + model, + request, + request_headers: headers_to_map(&headers), + response_raw, + response_text, + stream, + status_code, + completed_at: Utc::now(), + metadata: BTreeMap::new(), + field_patches: BTreeMap::new(), + capture_meta: CaptureMeta { + finish_reason, + usage, + segment_kind: None, + }, + user_seq, + assistant_seq, + }; + state.post_processors.apply(&mut event); + if let Err(err) = state.capture_sink.dispatch(event).await { + state + .push_error(format!("capture sink dispatch failed: {err}")) + .await; + } +} + +fn parse_chat_sse_response(raw: &str) -> (Option, Option, Option) { + let mut content = String::new(); + let mut finish_reason = None; + let mut usage = None; + + for line in raw.lines() { + let trimmed = line.trim(); + if !trimmed.starts_with("data:") { + continue; + } + let data = trimmed.trim_start_matches("data:").trim(); + if data.is_empty() || data == "[DONE]" { + continue; + } + + let chunk = match serde_json::from_str::(data) { + Ok(value) => value, + Err(_) => continue, + }; + + if let Some(delta_text) = chunk + .get("choices") + .and_then(Value::as_array) + .and_then(|choices| choices.first()) + .and_then(|choice| choice.get("delta")) + .and_then(|delta| delta.get("content")) + .and_then(Value::as_str) + { + content.push_str(delta_text); + } + + if content.is_empty() { + if let Some(message_text) = chunk + .get("choices") + .and_then(Value::as_array) + .and_then(|choices| choices.first()) + .and_then(|choice| choice.get("message")) + .and_then(|message| message.get("content")) + .and_then(Value::as_str) + { + content.push_str(message_text); + } + } + + if finish_reason.is_none() { + finish_reason = chunk + .get("choices") + .and_then(Value::as_array) + .and_then(|choices| choices.first()) + .and_then(|choice| choice.get("finish_reason")) + .and_then(Value::as_str) + .map(ToString::to_string); + } + + if usage.is_none() { + usage = chunk.get("usage").cloned(); + } + } + + let response_text = if content.is_empty() { + None + } else { + Some(content) + }; + (response_text, finish_reason, usage) +} + +fn truncate_text(text: &str, max_len: usize) -> String { + if text.chars().count() <= max_len { + return text.to_string(); + } + let mut truncated = text.chars().take(max_len).collect::(); + truncated.push_str("..."); + truncated +} + +fn error_response(status: StatusCode, message: &str) -> Response { + (status, Json(json!({ "error": message }))).into_response() +} diff --git a/crates/persisting-dlcapt/src/router.rs b/crates/persisting-dlcapt/src/router.rs new file mode 100644 index 0000000..367ea18 --- /dev/null +++ b/crates/persisting-dlcapt/src/router.rs @@ -0,0 +1,56 @@ +use crate::config::{ModelRoute, ProxyConfig}; +use std::collections::HashMap; + +#[derive(Debug, Clone)] +pub struct RouteTable { + exact: HashMap, + wildcard: Option, + all: Vec, +} + +impl RouteTable { + pub fn from_config(config: &ProxyConfig) -> Self { + let mut exact = HashMap::new(); + let mut wildcard = None; + + for route in &config.models { + if route.name == "*" { + wildcard = Some(route.clone()); + } else { + exact.insert(route.name.clone(), route.clone()); + } + } + + Self { + exact, + wildcard, + all: config.models.clone(), + } + } + + pub fn resolve_model(&self, model: &str) -> Option<&ModelRoute> { + self.exact.get(model).or(self.wildcard.as_ref()) + } + + pub fn list_models(&self) -> Vec { + self.all + .iter() + .filter(|route| route.name != "*") + .map(|route| ModelInfo { + id: route.name.clone(), + provider: route.provider.clone(), + display_name: route + .display_name + .clone() + .unwrap_or_else(|| route.name.clone()), + }) + .collect() + } +} + +#[derive(Debug, Clone)] +pub struct ModelInfo { + pub id: String, + pub provider: String, + pub display_name: String, +} diff --git a/crates/persisting-dlcapt/src/session/mod.rs b/crates/persisting-dlcapt/src/session/mod.rs new file mode 100644 index 0000000..47b9d19 --- /dev/null +++ b/crates/persisting-dlcapt/src/session/mod.rs @@ -0,0 +1,76 @@ +mod normalize; +mod providers; +mod resolver; + +pub use normalize::{DEFAULT_MAX_STEM_LEN, normalize_session_id}; +pub use resolver::{ + RequestContext, ResolvedSession, RouteSessionMode, SessionIdSettings, resolve_session, +}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SessionSource { + UrlPath, + Header, + BodyMetadata, + Default, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionConflict { + pub source: SessionSource, + pub value: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SessionResolveError { + MissingPathSessionId, + InvalidPathSessionId, + MissingSessionId, + InvalidSessionId, +} + +// Backward-compatible helpers for existing tests. +pub use providers::extract_header_session as extract_session_from_headers; + +pub fn session_id_from_body(body: &serde_json::Value) -> Option { + providers::extract_body_metadata_session(body) +} + +pub fn resolve_session_id( + header_session_id: Option<&str>, + body: &serde_json::Value, + default: &str, +) -> String { + resolve_session_with_source(header_session_id, body, default).0 +} + +pub fn resolve_session_with_source( + header_session_id: Option<&str>, + body: &serde_json::Value, + default: &str, +) -> (String, SessionSource) { + let mut headers = axum::http::HeaderMap::new(); + if let Some(value) = header_session_id.filter(|value| !value.trim().is_empty()) { + if let Ok(header_value) = value.parse() { + headers.insert("x-session-id", header_value); + } + } + + let settings = SessionIdSettings { + default_session_id: default.to_string(), + preserve_raw: false, + session_header: "x-persisting-session-id".to_string(), + session_header_aliases: vec![], + }; + let ctx = RequestContext { + mode: RouteSessionMode::Flat, + path_session_id: None, + headers: &headers, + body, + }; + + match resolve_session(&ctx, &settings) { + Ok(resolved) => (resolved.storage_session_id, resolved.source), + Err(_) => (default.to_string(), SessionSource::Default), + } +} diff --git a/crates/persisting-dlcapt/src/session/normalize.rs b/crates/persisting-dlcapt/src/session/normalize.rs new file mode 100644 index 0000000..f1b89a2 --- /dev/null +++ b/crates/persisting-dlcapt/src/session/normalize.rs @@ -0,0 +1,53 @@ +pub const DEFAULT_MAX_STEM_LEN: usize = 128; + +pub fn normalize_session_id(raw: &str, preserve_raw: bool, max_stem_len: usize) -> Option { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return None; + } + + let normalized = if preserve_raw { + trimmed.to_string() + } else { + sanitize_path_segment(trimmed) + }; + + if normalized.is_empty() { + return None; + } + + Some(normalized.chars().take(max_stem_len).collect()) +} + +fn sanitize_path_segment(raw: &str) -> String { + raw.chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { + ch + } else { + '-' + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn normalize_should_sanitize_unsafe_chars_by_default() { + assert_eq!( + normalize_session_id("feishu/main", false, DEFAULT_MAX_STEM_LEN).as_deref(), + Some("feishu-main") + ); + } + + #[test] + fn normalize_should_preserve_raw_when_enabled() { + assert_eq!( + normalize_session_id("feishu/main", true, DEFAULT_MAX_STEM_LEN).as_deref(), + Some("feishu/main") + ); + } +} diff --git a/crates/persisting-dlcapt/src/session/providers/body.rs b/crates/persisting-dlcapt/src/session/providers/body.rs new file mode 100644 index 0000000..e899513 --- /dev/null +++ b/crates/persisting-dlcapt/src/session/providers/body.rs @@ -0,0 +1,14 @@ +use serde_json::Value; + +pub fn extract_body_metadata_session(body: &Value) -> Option { + let session_id = body + .get("metadata") + .and_then(|meta| meta.get("session_id")) + .and_then(Value::as_str)?; + let trimmed = session_id.trim(); + if trimmed.is_empty() { + None + } else { + Some(trimmed.to_string()) + } +} diff --git a/crates/persisting-dlcapt/src/session/providers/default.rs b/crates/persisting-dlcapt/src/session/providers/default.rs new file mode 100644 index 0000000..b43771b --- /dev/null +++ b/crates/persisting-dlcapt/src/session/providers/default.rs @@ -0,0 +1,8 @@ +pub fn default_session_candidate(default_session_id: &str) -> Option { + let trimmed = default_session_id.trim(); + if trimmed.is_empty() { + None + } else { + Some(trimmed.to_string()) + } +} diff --git a/crates/persisting-dlcapt/src/session/providers/header.rs b/crates/persisting-dlcapt/src/session/providers/header.rs new file mode 100644 index 0000000..2d70b30 --- /dev/null +++ b/crates/persisting-dlcapt/src/session/providers/header.rs @@ -0,0 +1,43 @@ +use axum::http::HeaderMap; + +const SESSION_HEADER_ALIASES: &[&str] = &[ + "x-persisting-session-id", + "x-session-id", + "x-openclaw-session-id", +]; + +pub fn extract_header_session( + headers: &HeaderMap, + primary_header: &str, + extra_aliases: &[String], +) -> Option { + let mut candidates = Vec::with_capacity(SESSION_HEADER_ALIASES.len() + extra_aliases.len() + 1); + let primary = primary_header.trim().to_ascii_lowercase(); + if !primary.is_empty() { + candidates.push(primary); + } + for alias in SESSION_HEADER_ALIASES { + if !candidates.iter().any(|name| name == alias) { + candidates.push((*alias).to_string()); + } + } + for alias in extra_aliases { + let normalized = alias.trim().to_ascii_lowercase(); + if !normalized.is_empty() && !candidates.iter().any(|name| name == &normalized) { + candidates.push(normalized); + } + } + + for name in candidates { + if let Some(value) = headers + .get(name.as_str()) + .and_then(|header| header.to_str().ok()) + { + let trimmed = value.trim(); + if !trimmed.is_empty() { + return Some(trimmed.to_string()); + } + } + } + None +} diff --git a/crates/persisting-dlcapt/src/session/providers/mod.rs b/crates/persisting-dlcapt/src/session/providers/mod.rs new file mode 100644 index 0000000..5c2e6b2 --- /dev/null +++ b/crates/persisting-dlcapt/src/session/providers/mod.rs @@ -0,0 +1,9 @@ +mod body; +mod default; +mod header; +mod url_path; + +pub use body::extract_body_metadata_session; +pub use default::default_session_candidate; +pub use header::extract_header_session; +pub use url_path::extract_url_path_session; diff --git a/crates/persisting-dlcapt/src/session/providers/url_path.rs b/crates/persisting-dlcapt/src/session/providers/url_path.rs new file mode 100644 index 0000000..dd613aa --- /dev/null +++ b/crates/persisting-dlcapt/src/session/providers/url_path.rs @@ -0,0 +1,8 @@ +pub fn extract_url_path_session(raw_path_session_id: &str) -> Option { + let trimmed = raw_path_session_id.trim(); + if trimmed.is_empty() { + None + } else { + Some(trimmed.to_string()) + } +} diff --git a/crates/persisting-dlcapt/src/session/resolver.rs b/crates/persisting-dlcapt/src/session/resolver.rs new file mode 100644 index 0000000..7f607ba --- /dev/null +++ b/crates/persisting-dlcapt/src/session/resolver.rs @@ -0,0 +1,200 @@ +use axum::http::HeaderMap; +use serde_json::Value; + +use super::normalize::{DEFAULT_MAX_STEM_LEN, normalize_session_id}; +use super::providers::{ + default_session_candidate, extract_body_metadata_session, extract_header_session, + extract_url_path_session, +}; +use super::{SessionConflict, SessionResolveError, SessionSource}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RouteSessionMode { + Flat, + SessionScoped, +} + +pub struct RequestContext<'a> { + pub mode: RouteSessionMode, + pub path_session_id: Option<&'a str>, + pub headers: &'a HeaderMap, + pub body: &'a Value, +} + +pub struct SessionIdSettings { + pub default_session_id: String, + pub preserve_raw: bool, + pub session_header: String, + pub session_header_aliases: Vec, +} + +pub struct ResolvedSession { + pub storage_session_id: String, + pub source: SessionSource, + pub raw_value: String, + pub conflicts: Vec, +} + +pub fn resolve_session( + ctx: &RequestContext<'_>, + settings: &SessionIdSettings, +) -> Result { + match ctx.mode { + RouteSessionMode::SessionScoped => resolve_session_scoped(ctx, settings), + RouteSessionMode::Flat => resolve_session_flat(ctx, settings), + } +} + +fn resolve_session_scoped( + ctx: &RequestContext<'_>, + settings: &SessionIdSettings, +) -> Result { + let raw_path = ctx + .path_session_id + .and_then(extract_url_path_session) + .ok_or(SessionResolveError::MissingPathSessionId)?; + + let storage_session_id = + normalize_session_id(&raw_path, settings.preserve_raw, DEFAULT_MAX_STEM_LEN) + .ok_or(SessionResolveError::InvalidPathSessionId)?; + + let mut conflicts = Vec::new(); + if let Some(header) = extract_header_session( + ctx.headers, + &settings.session_header, + &settings.session_header_aliases, + ) { + if header != raw_path { + conflicts.push(SessionConflict { + source: SessionSource::Header, + value: header, + }); + } + } + if let Some(body) = extract_body_metadata_session(ctx.body) { + if body != raw_path { + conflicts.push(SessionConflict { + source: SessionSource::BodyMetadata, + value: body, + }); + } + } + + if !conflicts.is_empty() { + tracing::debug!( + storage_session_id = %storage_session_id, + conflict_count = conflicts.len(), + "session-scoped route ignored non-url session candidates" + ); + } + + Ok(ResolvedSession { + storage_session_id, + source: SessionSource::UrlPath, + raw_value: raw_path, + conflicts, + }) +} + +fn resolve_session_flat( + ctx: &RequestContext<'_>, + settings: &SessionIdSettings, +) -> Result { + let header = extract_header_session( + ctx.headers, + &settings.session_header, + &settings.session_header_aliases, + ); + let body = extract_body_metadata_session(ctx.body); + let default = default_session_candidate(&settings.default_session_id); + + let (raw_value, source) = if let Some(value) = header.as_deref() { + (value.to_string(), SessionSource::Header) + } else if let Some(value) = body.as_deref() { + (value.to_string(), SessionSource::BodyMetadata) + } else if let Some(value) = default.as_deref() { + (value.to_string(), SessionSource::Default) + } else { + return Err(SessionResolveError::MissingSessionId); + }; + + let storage_session_id = + normalize_session_id(&raw_value, settings.preserve_raw, DEFAULT_MAX_STEM_LEN) + .ok_or(SessionResolveError::InvalidSessionId)?; + + let mut conflicts = Vec::new(); + if source != SessionSource::Header { + if let Some(value) = header { + conflicts.push(SessionConflict { + source: SessionSource::Header, + value, + }); + } + } + if source != SessionSource::BodyMetadata { + if let Some(value) = body { + conflicts.push(SessionConflict { + source: SessionSource::BodyMetadata, + value, + }); + } + } + + Ok(ResolvedSession { + storage_session_id, + source, + raw_value, + conflicts, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::http::HeaderMap; + use serde_json::json; + + fn settings() -> SessionIdSettings { + SessionIdSettings { + default_session_id: "default".to_string(), + preserve_raw: false, + session_header: "x-persisting-session-id".to_string(), + session_header_aliases: vec![], + } + } + + #[test] + fn session_scoped_should_prefer_url_over_header() { + let mut headers = HeaderMap::new(); + headers.insert("x-session-id", "bbb".parse().unwrap()); + let body = json!({}); + let ctx = RequestContext { + mode: RouteSessionMode::SessionScoped, + path_session_id: Some("aaa"), + headers: &headers, + body: &body, + }; + + let resolved = resolve_session(&ctx, &settings()).expect("resolve session"); + assert_eq!(resolved.storage_session_id, "aaa"); + assert_eq!(resolved.source, SessionSource::UrlPath); + assert_eq!(resolved.conflicts.len(), 1); + assert_eq!(resolved.conflicts[0].value, "bbb"); + } + + #[test] + fn flat_route_should_prefer_header_then_body_then_default() { + let headers = HeaderMap::new(); + let body = json!({"metadata": {"session_id": "body-session"}}); + let ctx = RequestContext { + mode: RouteSessionMode::Flat, + path_session_id: None, + headers: &headers, + body: &body, + }; + + let resolved = resolve_session(&ctx, &settings()).expect("resolve session"); + assert_eq!(resolved.storage_session_id, "body-session"); + assert_eq!(resolved.source, SessionSource::BodyMetadata); + } +} diff --git a/crates/persisting-dlcapt/src/tlv.rs b/crates/persisting-dlcapt/src/tlv.rs new file mode 100644 index 0000000..1c30960 --- /dev/null +++ b/crates/persisting-dlcapt/src/tlv.rs @@ -0,0 +1,331 @@ +use crate::capture::CaptureEvent; +use crate::capture::session_dir::resolve_session_layout; +use anyhow::{Context, Result}; +use chrono::Utc; +use serde_json::{Map, Value, json}; +use std::collections::BTreeMap; +use std::path::PathBuf; +use std::sync::Arc; +use tokio::fs; +use tokio::sync::Mutex; + +const BLOCK_MARKER: &str = "\n\nmessage body\n\n"; +const BLOCK_FORMAT_VERSION: u64 = 1; + +#[derive(Debug, Clone)] +pub struct TlvWriter { + store_root: PathBuf, + agent_id: String, + default_session_id: String, + write_lock: Arc>, +} + +#[derive(Debug, Clone)] +pub struct TlvTurnRecord { + pub session_id: String, + pub agent_id: String, + pub model: String, + pub stream: bool, + pub status_code: u16, + pub user_text: Option, + pub assistant_text: Option, + pub usage: Option, + pub user_seq: u64, + pub assistant_seq: u64, + pub turn: u64, + pub call_id: String, + pub request_path: String, +} + +#[derive(Debug, Clone)] +pub struct MdSinkInput { + pub session_id: String, + pub agent_id: String, + pub model: String, + pub stream: bool, + pub status_code: u16, + pub user_text: Option, + pub assistant_text: Option, + pub usage: Option, + pub user_seq: u64, + pub assistant_seq: u64, + pub turn: u64, + pub call_id: String, + pub request_path: String, +} + +impl MdSinkInput { + pub fn from_capture_event(event: &CaptureEvent) -> Self { + Self { + session_id: event.session_id.clone(), + agent_id: event.agent_id.clone(), + model: event.model.clone(), + stream: event.stream, + status_code: event.status_code, + user_text: crate::dialogue::extract_user_text(event.endpoint, &event.request), + assistant_text: event.response_text.clone(), + usage: event.capture_meta.usage.clone(), + user_seq: event.user_seq, + assistant_seq: event.assistant_seq, + turn: event.turn, + call_id: event.call_id.clone(), + request_path: event.request_path.clone(), + } + } + + pub fn to_tlv_record(&self) -> Option { + let has_user = self + .user_text + .as_deref() + .is_some_and(|text| !text.is_empty()); + let has_assistant = self + .assistant_text + .as_deref() + .is_some_and(|text| !text.is_empty()); + if !has_user && !has_assistant { + return None; + } + Some(TlvTurnRecord { + session_id: self.session_id.clone(), + agent_id: self.agent_id.clone(), + model: self.model.clone(), + stream: self.stream, + status_code: self.status_code, + user_text: self.user_text.clone(), + assistant_text: self.assistant_text.clone(), + usage: self.usage.clone(), + user_seq: self.user_seq, + assistant_seq: self.assistant_seq, + turn: self.turn, + call_id: self.call_id.clone(), + request_path: self.request_path.clone(), + }) + } +} + +impl TlvWriter { + pub fn new( + store_root: PathBuf, + agent_id: String, + default_session_id: String, + write_lock: Arc>, + ) -> Self { + Self { + store_root, + agent_id, + default_session_id, + write_lock, + } + } + + pub fn write_lock(&self) -> Arc> { + Arc::clone(&self.write_lock) + } + + pub async fn append_turn(&self, record: TlvTurnRecord) -> Result { + let _guard = self.write_lock.lock().await; + self.append_turn_internal(record).await + } + + pub async fn append_turn_internal(&self, record: TlvTurnRecord) -> Result { + let now = Utc::now(); + let layout = resolve_session_layout(&record.session_id, &self.default_session_id, now); + let md_path = self + .store_root + .join(&layout.session_dir) + .join("trajectory.md"); + + if let Some(parent) = md_path.parent() { + fs::create_dir_all(parent) + .await + .with_context(|| format!("failed creating tlv dir: {}", parent.display()))?; + } + + let mut blocks = Vec::new(); + if let Some(user_text) = record.user_text.as_deref().filter(|text| !text.is_empty()) { + blocks.push(encode_block( + "user", + user_text, + &record, + record.user_seq, + "llm.request", + )?); + } + if let Some(assistant_text) = record + .assistant_text + .as_deref() + .filter(|text| !text.is_empty()) + { + let kind = if record.stream { + "llm.response.stream" + } else { + "llm.response" + }; + blocks.push(encode_block( + "assistant", + assistant_text, + &record, + record.assistant_seq, + kind, + )?); + } + + if blocks.is_empty() { + anyhow::bail!("tlv turn has no visible user or assistant content"); + } + + let turns = record.turn; + if !md_path.exists() { + let preamble = format_document_preamble(&record.session_id, &self.agent_id, turns)?; + let content = format!("{preamble}{}", blocks.join("")); + fs::write(&md_path, content) + .await + .with_context(|| format!("failed writing tlv file: {}", md_path.display()))?; + } else { + let existing = fs::read_to_string(&md_path) + .await + .with_context(|| format!("failed reading tlv file: {}", md_path.display()))?; + let updated = refresh_frontmatter_turns(&existing, turns)?; + let content = format!("{updated}{}", blocks.join("")); + fs::write(&md_path, content) + .await + .with_context(|| format!("failed appending tlv file: {}", md_path.display()))?; + } + + Ok(md_path) + } +} + +pub fn new_call_id() -> String { + let now = Utc::now(); + format!( + "call-{}-{}", + now.format("%Y%m%d%H%M%S"), + now.timestamp_subsec_micros() + ) +} + +fn encode_block( + speaker: &str, + body: &str, + record: &TlvTurnRecord, + seq: u64, + kind: &str, +) -> Result { + let timestamp = Utc::now().to_rfc3339(); + let mut fields = BTreeMap::new(); + fields.insert("agent_id".to_string(), json!(record.agent_id)); + fields.insert("call_id".to_string(), json!(record.call_id)); + fields.insert("kind".to_string(), json!(kind)); + fields.insert("model".to_string(), json!(record.model)); + fields.insert("path".to_string(), json!(record.request_path)); + fields.insert("role".to_string(), json!(speaker)); + fields.insert("seq".to_string(), json!(seq)); + fields.insert("session_id".to_string(), json!(record.session_id)); + fields.insert("source".to_string(), json!("dlcapt-proxy")); + fields.insert("timestamp".to_string(), json!(timestamp)); + fields.insert("trace_id".to_string(), json!(record.call_id)); + fields.insert("turn".to_string(), json!(record.turn)); + fields.insert("v".to_string(), json!(BLOCK_FORMAT_VERSION)); + + if speaker == "assistant" { + fields.insert("status".to_string(), json!(record.status_code)); + if let Some(usage) = record.usage.as_ref() { + if let Some(prompt_tokens) = usage.get("prompt_tokens") { + fields.insert("prompt_tokens".to_string(), prompt_tokens.clone()); + fields.insert("input_tokens".to_string(), prompt_tokens.clone()); + } + if let Some(completion_tokens) = usage.get("completion_tokens") { + fields.insert("completion_tokens".to_string(), completion_tokens.clone()); + fields.insert("output_tokens".to_string(), completion_tokens.clone()); + } + if let Some(total_tokens) = usage.get("total_tokens") { + fields.insert("total_tokens".to_string(), total_tokens.clone()); + } + } + } + + let header = json!({ + "type": "markdown", + "length": body.len(), + "fields": fields, + }); + let flat = flatten_block_header(&header)?; + let json_text = serde_json::to_string(&flat).context("serialize tlv block header")?; + Ok(format!( + "{BLOCK_MARKER}:{speaker} {json_text} -->\n\n{body}\n\n" + )) +} + +fn flatten_block_header(header: &Value) -> Result> { + let mut out = Map::new(); + if let Some(type_name) = header.get("type") { + out.insert("type".to_string(), type_name.clone()); + } + if let Some(length) = header.get("length") { + out.insert("length".to_string(), length.clone()); + } + if let Some(fields) = header.get("fields").and_then(Value::as_object).cloned() { + for (key, value) in fields { + out.insert(key, value); + } + } + Ok(out) +} + +fn format_document_preamble(session_id: &str, agent_id: &str, turns: u64) -> Result { + Ok(format!( + "---\n\ +format: persisting:1.0\n\ +block: |+\n\ + {BLOCK_LAYOUT}\n\ +session: {session_id}\n\ +agent: {agent_id}\n\ +turns: {turns}\n\ +client:\n\ + peer: ''\n\ + peer_port: 0\n\ + pid: 0\n\ + command: openclaw\n\ + machine_fp: ''\n\ +---\n\n" + )) +} + +fn refresh_frontmatter_turns(content: &str, turns: u64) -> Result { + let Some(end) = content.find("\n---\n\n") else { + anyhow::bail!("tlv file missing frontmatter terminator"); + }; + let preamble = &content[..end]; + let body = &content[end + "\n---\n\n".len()..]; + let updated_preamble = if preamble.contains("turns:") { + preamble + .lines() + .map(|line| { + if line.starts_with("turns:") { + format!("turns: {turns}") + } else { + line.to_string() + } + }) + .collect::>() + .join("\n") + } else { + format!("{preamble}\nturns: {turns}") + }; + Ok(format!("{updated_preamble}\n---\n\n{body}")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn refresh_frontmatter_turns_should_update_turn_count() { + let input = "---\nformat: persisting:1.0\nsession: abc\nturns: 1\n---\n\nbody\n"; + let updated = refresh_frontmatter_turns(input, 3).expect("refresh turns"); + assert!(updated.contains("turns: 3")); + assert!(updated.ends_with("body\n")); + } +} diff --git a/crates/persisting-dlcapt/tests/mvp_tests.rs b/crates/persisting-dlcapt/tests/mvp_tests.rs new file mode 100644 index 0000000..85ea439 --- /dev/null +++ b/crates/persisting-dlcapt/tests/mvp_tests.rs @@ -0,0 +1,494 @@ +use axum::http::HeaderMap; +use persisting_dlcapt::config::{ExportConfig, ModelRoute, ProxyConfig, StorageConfig}; +use persisting_dlcapt::dialogue::{InferenceEndpoint, extract_user_text}; +use persisting_dlcapt::router::RouteTable; +use persisting_dlcapt::session::{ + RequestContext, RouteSessionMode, SessionSource, extract_session_from_headers, resolve_session, + resolve_session_id, resolve_session_with_source, +}; +use persisting_dlcapt::tlv::{TlvTurnRecord, TlvWriter}; +use serde_json::json; +use std::sync::Arc; +use tempfile::tempdir; +use tokio::sync::Mutex; + +fn test_config() -> ProxyConfig { + ProxyConfig { + listen: "127.0.0.1:19081".to_string(), + admin_listen: "127.0.0.1:19082".to_string(), + store_dir: "store".to_string(), + agent_id: "openclaw".to_string(), + session_header: "x-persisting-session-id".to_string(), + session_header_aliases: vec![], + default_session_id: "default".to_string(), + preserve_raw: false, + base_session_path: "/v1/sessions".to_string(), + storage: StorageConfig::default(), + export: ExportConfig::default(), + models: vec![ + ModelRoute { + name: "kimi-k2.5".to_string(), + display_name: Some("Kimi K2.5".to_string()), + provider: "openai".to_string(), + upstream_base_url: "http://exact-upstream/v1".to_string(), + api_key: Some("exact-key".to_string()), + }, + ModelRoute { + name: "*".to_string(), + display_name: Some("Fallback".to_string()), + provider: "openai".to_string(), + upstream_base_url: "http://fallback-upstream/v1".to_string(), + api_key: Some("fallback-key".to_string()), + }, + ], + } +} + +#[test] +fn model_route_should_match_exact_then_wildcard() { + let config = test_config(); + let routes = RouteTable::from_config(&config); + let exact = routes + .resolve_model("kimi-k2.5") + .expect("should match exact"); + assert_eq!(exact.upstream_base_url, "http://exact-upstream/v1"); + + let fallback = routes + .resolve_model("unknown-model") + .expect("should match wildcard"); + assert_eq!(fallback.upstream_base_url, "http://fallback-upstream/v1"); +} + +#[test] +fn session_id_resolution_priority_should_be_header_then_body_then_default() { + let body = json!({ + "metadata": { + "session_id": "body-session" + } + }); + + let from_header = resolve_session_id(Some("header-session"), &body, "default-session"); + assert_eq!(from_header, "header-session"); + + let from_body = resolve_session_id(None, &body, "default-session"); + assert_eq!(from_body, "body-session"); + + let from_default = resolve_session_id(None, &json!({}), "default-session"); + assert_eq!(from_default, "default-session"); +} + +#[test] +fn session_header_aliases_should_be_checked_in_order() { + let mut headers = HeaderMap::new(); + headers.insert("x-openclaw-session-id", "openclaw-session".parse().unwrap()); + + let session = extract_session_from_headers(&headers, "x-persisting-session-id", &[]) + .expect("should resolve alias header"); + assert_eq!(session, "openclaw-session"); +} + +#[test] +fn configured_session_header_aliases_should_be_honored() { + let mut headers = HeaderMap::new(); + headers.insert("x-custom-session", "custom-session".parse().unwrap()); + + let session = extract_session_from_headers( + &headers, + "x-persisting-session-id", + &["x-custom-session".to_string()], + ) + .expect("should resolve configured alias"); + assert_eq!(session, "custom-session"); +} + +#[test] +fn resolve_session_with_source_should_mark_default_fallback() { + let (session_id, source) = resolve_session_with_source(None, &json!({}), "default"); + assert_eq!(session_id, "default"); + assert_eq!(source, SessionSource::Default); +} + +#[test] +fn session_scoped_resolver_should_normalize_url_and_ignore_header() { + let mut headers = HeaderMap::new(); + headers.insert("x-session-id", "bbb".parse().unwrap()); + let body = json!({}); + let ctx = RequestContext { + mode: RouteSessionMode::SessionScoped, + path_session_id: Some("aaa"), + headers: &headers, + body: &body, + }; + + let resolved = resolve_session(&ctx, &test_config().session_settings()).expect("resolve"); + assert_eq!(resolved.storage_session_id, "aaa"); + assert_eq!(resolved.source, SessionSource::UrlPath); + assert_eq!(resolved.conflicts.len(), 1); +} + +#[test] +fn responses_user_extraction_should_read_string_input() { + let body = json!({"input": "responses user"}); + assert_eq!( + extract_user_text(InferenceEndpoint::Responses, &body).as_deref(), + Some("responses user") + ); +} + +#[tokio::test] +async fn tlv_writer_should_append_same_session_into_one_markdown_file() { + let temp = tempdir().expect("temp dir"); + let writer = TlvWriter::new( + temp.path().to_path_buf(), + "openclaw".to_string(), + "default".to_string(), + Arc::new(Mutex::new(())), + ); + + let first = TlvTurnRecord { + session_id: "feishu-main-abc".to_string(), + agent_id: "openclaw".to_string(), + model: "kimi-k2.5".to_string(), + stream: false, + status_code: 200, + user_text: Some("hi".to_string()), + assistant_text: Some("hello".to_string()), + usage: Some(json!({"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13})), + user_seq: 0, + assistant_seq: 1, + turn: 1, + call_id: "call-test-1".to_string(), + request_path: "/v1/sessions/feishu-main-abc/chat/completions".to_string(), + }; + let second = TlvTurnRecord { + session_id: "feishu-main-abc".to_string(), + agent_id: "openclaw".to_string(), + model: "kimi-k2.5".to_string(), + stream: false, + status_code: 200, + user_text: Some("again".to_string()), + assistant_text: Some("sure".to_string()), + usage: Some(json!({"prompt_tokens": 12, "completion_tokens": 2, "total_tokens": 14})), + user_seq: 2, + assistant_seq: 3, + turn: 2, + call_id: "call-test-2".to_string(), + request_path: "/v1/sessions/feishu-main-abc/chat/completions".to_string(), + }; + + let first_path = writer.append_turn(first).await.expect("write first turn"); + let second_path = writer.append_turn(second).await.expect("write second turn"); + assert_eq!(first_path, second_path); + assert!(first_path.ends_with("feishu-main-abc/trajectory.md")); + + let content = tokio::fs::read_to_string(second_path) + .await + .expect("read markdown"); + assert!(content.contains("format: persisting:1.0")); + assert!(content.contains("session: feishu-main-abc")); + assert!(content.contains("turns: 2")); + assert!(content.contains("