diff --git a/.github/workflows/dlcapt.yml b/.github/workflows/dlcapt.yml new file mode 100644 index 0000000..6a3b254 --- /dev/null +++ b/.github/workflows/dlcapt.yml @@ -0,0 +1,47 @@ +name: dlcapt + +on: + push: + paths: + - "crates/persisting-dlcapt/**" + - "crates/persisting-cli/**" + - "Cargo.toml" + - "Cargo.lock" + - ".github/workflows/dlcapt.yml" + pull_request: + paths: + - "crates/persisting-dlcapt/**" + - "crates/persisting-cli/**" + - "Cargo.toml" + - "Cargo.lock" + - ".github/workflows/dlcapt.yml" + +jobs: + check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + components: rustfmt, clippy + - name: Install protoc + run: sudo apt-get update && sudo apt-get install -y protobuf-compiler + - uses: Swatinem/rust-cache@v2 + - name: fmt + run: cargo fmt --check + - name: check + run: cargo check -p persisting-dlcapt --all-targets --all-features --locked + - name: test + run: cargo test -p persisting-dlcapt --all-targets --all-features --locked + - name: clippy + run: cargo clippy -p persisting-dlcapt --all-targets --all-features --locked -- -D warnings + - name: doc + run: RUSTDOCFLAGS="-D warnings" cargo doc -p persisting-dlcapt --all-features --no-deps --locked + - name: build standalone dlcapt + run: cargo build -p persisting-dlcapt --locked + - name: test persisting-cli dlcapt feature + run: | + metadata="$(cargo metadata --no-deps --format-version 1 --locked)" + target_directory="$(jq -r '.target_directory' <<<"${metadata}")" + DLCAPT_BIN="${target_directory}/debug/dlcapt" \ + cargo test -p persisting-cli --features dlcapt --locked diff --git a/.gitignore b/.gitignore index a7a8960..0d995dc 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,10 @@ __pycache__/ .venv/ *.abi3.so crates/persisting-engine/ds/ + +# persisting-dlcapt runtime / release artifacts +crates/persisting-dlcapt/var/ +crates/persisting-dlcapt/.capture/ +crates/persisting-dlcapt/tmp/ +.build-tmp/ +target/dlcapt/ diff --git a/Cargo.lock b/Cargo.lock index 5233bbe..65ddec7 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", @@ -5787,12 +5792,14 @@ dependencies = [ "dirs", "libloading", "persisting-capture", + "persisting-dlcapt", "persisting-engine", "persisting-proto", "reqwest 0.12.28", "ron", "serde", "serde_json", + "tempfile", "termimad", "tokio", "toml", @@ -5814,6 +5821,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 +6147,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/NOTICE b/NOTICE new file mode 100644 index 0000000..1f12864 --- /dev/null +++ b/NOTICE @@ -0,0 +1,13 @@ +persisting-dlcapt +================= + +This workspace includes `persisting-dlcapt`, migrated from +`capture/external/dlcapt` (source tree copy of first-party Rust sources). + +Target package: persisting-dlcapt +License: Apache-2.0 (workspace policy) +Third-party source files copied verbatim outside Cargo crates.io +dependencies: none identified at migration time. + +If additional vendored sources or incompatible licenses are discovered, +stop publication and obtain explicit clearance before release. diff --git a/crates/persisting-cli/Cargo.toml b/crates/persisting-cli/Cargo.toml index 7f9e67f..2a0ad6b 100644 --- a/crates/persisting-cli/Cargo.toml +++ b/crates/persisting-cli/Cargo.toml @@ -10,6 +10,10 @@ description = "Lightweight Persisting CLI (loads persisting-engine via dlopen)" name = "persisting" path = "src/main.rs" +[features] +default = [] +dlcapt = ["dep:persisting-dlcapt"] + [dependencies] anyhow = "1" chrono = "0.4" @@ -18,6 +22,7 @@ clap = { version = "4", features = ["derive", "env"] } dirs = "6" libloading = "0.8" persisting-capture = { path = "../persisting-capture" } +persisting-dlcapt = { path = "../persisting-dlcapt", optional = true } persisting-engine = { path = "../persisting-engine" } persisting-proto = { path = "../persisting-proto" } ron = "0.8" @@ -29,3 +34,11 @@ termimad = "0.34" toml = "0.8" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } + +[dev-dependencies] +tempfile = "3" + +[[test]] +name = "dlcapt_backend_e2e" +path = "tests/dlcapt_backend_e2e.rs" +required-features = ["dlcapt"] diff --git a/crates/persisting-cli/src/judge_manual.rs b/crates/persisting-cli/src/judge_manual.rs index a1086d7..1862225 100644 --- a/crates/persisting-cli/src/judge_manual.rs +++ b/crates/persisting-cli/src/judge_manual.rs @@ -332,7 +332,9 @@ mod tests { fn fixed_manual_scores_story_scope() { use persisting_capture::config::CaptureLevel; use persisting_capture::engine::Call; - use persisting_capture::sink::{llm_request_summary_record, llm_response_record_with_content}; + use persisting_capture::sink::{ + llm_request_summary_record, llm_response_record_with_content, + }; let call = Call { call_id: "c1".into(), @@ -368,9 +370,15 @@ mod tests { serde_json::to_string(&resp).unwrap(), ]; let story = story_from_replay_json(&records, "s", "s").unwrap(); - let scores = - fixed_manual_scores(&story, JudgeScope::Story, &["default".into()], 100, None, None) - .unwrap(); + let scores = fixed_manual_scores( + &story, + JudgeScope::Story, + &["default".into()], + 100, + None, + None, + ) + .unwrap(); assert_eq!(scores.len(), 1); assert_eq!(scores[0].score, 100); assert_eq!(scores[0].verdict, "pass"); @@ -381,7 +389,9 @@ mod tests { fn fixed_manual_scores_rejects_out_of_range() { use persisting_capture::config::CaptureLevel; use persisting_capture::engine::Call; - use persisting_capture::sink::{llm_request_summary_record, llm_response_record_with_content}; + use persisting_capture::sink::{ + llm_request_summary_record, llm_response_record_with_content, + }; let call = Call { call_id: "c1".into(), @@ -417,6 +427,14 @@ mod tests { serde_json::to_string(&resp).unwrap(), ]; let story = story_from_replay_json(&records, "s", "s").unwrap(); - assert!(fixed_manual_scores(&story, JudgeScope::Story, &["default".into()], 101, None, None).is_err()); + assert!(fixed_manual_scores( + &story, + JudgeScope::Story, + &["default".into()], + 101, + None, + None + ) + .is_err()); } } diff --git a/crates/persisting-cli/src/main.rs b/crates/persisting-cli/src/main.rs index 6da6daf..ee63094 100644 --- a/crates/persisting-cli/src/main.rs +++ b/crates/persisting-cli/src/main.rs @@ -635,19 +635,23 @@ enum TrajectoryCommand { JudgeStats(TrajectoryJudgeStatsArgs), } +#[derive(Debug, Clone, Copy, ValueEnum, Default)] +enum ProxyBackend { + #[default] + Capture, + Dlcapt, +} + /// `traj proxy` (foreground) or `traj proxy start|stop|list|status`. #[derive(Debug, Args)] #[command(args_conflicts_with_subcommands = true, after_long_help = PROXY_AFTER_HELP)] struct ProxyArgs { #[command(subcommand)] action: Option, + #[arg(long, value_enum, default_value_t = ProxyBackend::Capture)] + backend: ProxyBackend, /// Foreground proxy only: trajectory store root (`-o` with `traj proxy start` too). - #[arg( - long, - short = 'o', - value_name = "DIR", - env = "PERSISTING_CAPTURE_STORAGE" - )] + #[arg(long, short = 'o', value_name = "DIR")] output_dir: Option, /// Foreground proxy only: proxy TOML (`listen`, `models`, …). #[arg(long, short = 'c', value_name = "FILE")] @@ -656,8 +660,8 @@ struct ProxyArgs { #[arg(long)] debug: bool, /// Foreground proxy only: `md` (Markdown only) or `vortex` / `bin` (Vortex + live Markdown). - #[arg(long, short = 'f', value_enum, default_value_t = capture::CaptureFormat::Markdown)] - format: capture::CaptureFormat, + #[arg(long, short = 'f', value_enum)] + format: Option, } #[derive(Debug, Subcommand)] @@ -1109,10 +1113,24 @@ fn run_traj_import(lazy: &mut LazyEngine<'_>, args: &CaptureImportArgs) -> Resul } fn run_traj_proxy(lazy: &mut LazyEngine<'_>, args: &ProxyArgs) -> Result<()> { + match args.backend { + ProxyBackend::Capture => run_capture_proxy(lazy, args), + ProxyBackend::Dlcapt => run_dlcapt_proxy(args), + } +} + +fn capture_output_dir(args: &ProxyArgs) -> Option { + args.output_dir + .clone() + .or_else(|| std::env::var("PERSISTING_CAPTURE_STORAGE").ok()) +} + +fn run_capture_proxy(lazy: &mut LazyEngine<'_>, args: &ProxyArgs) -> Result<()> { + let format = args.format.unwrap_or(capture::CaptureFormat::Markdown); match &args.action { None => { - let output_dir = args - .output_dir + let capture_output_dir = capture_output_dir(args); + let output_dir = capture_output_dir .as_deref() .context("traj proxy requires -o ")?; let config = args @@ -1125,7 +1143,7 @@ fn run_traj_proxy(lazy: &mut LazyEngine<'_>, args: &ProxyArgs) -> Result<()> { output_dir: output_dir.to_string(), config: config.to_path_buf(), debug: args.debug, - format: args.format, + format, }, ) } @@ -1151,6 +1169,57 @@ fn run_traj_proxy(lazy: &mut LazyEngine<'_>, args: &ProxyArgs) -> Result<()> { } } +#[cfg(feature = "dlcapt")] +fn reject_dlcapt_capture_options(args: &ProxyArgs) -> Result<()> { + if args.output_dir.is_some() { + anyhow::bail!( + "-o is only supported by the capture backend; configure store_dir in the dlcapt TOML" + ); + } + if args.format.is_some() { + anyhow::bail!("-f is only supported by the capture backend"); + } + if args.debug { + anyhow::bail!("--debug is only supported by the capture backend"); + } + if args.action.is_some() { + anyhow::bail!("start, stop, list, and status are only supported by the capture backend"); + } + Ok(()) +} + +#[cfg(not(feature = "dlcapt"))] +fn run_dlcapt_proxy(_args: &ProxyArgs) -> Result<()> { + anyhow::bail!( + "dlcapt backend is not included in this build; rebuild persisting-cli with --features dlcapt" + ) +} + +#[cfg(feature = "dlcapt")] +fn init_dlcapt_cli_tracing_once() { + let _ = tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "dlcapt=info,tower_http=info".into()), + ) + .try_init(); +} + +#[cfg(feature = "dlcapt")] +fn run_dlcapt_proxy(args: &ProxyArgs) -> Result<()> { + reject_dlcapt_capture_options(args)?; + let config_path = args + .config + .as_deref() + .context("traj proxy --backend dlcapt requires -c ")?; + let config = persisting_dlcapt::config::ProxyConfig::load(config_path) + .with_context(|| format!("load dlcapt config {}", config_path.display()))?; + init_dlcapt_cli_tracing_once(); + tokio::runtime::Runtime::new() + .context("tokio runtime")? + .block_on(persisting_dlcapt::serve(config)) +} + struct TrajectoryAppendJob { storage: String, agent_id: String, @@ -2260,7 +2329,10 @@ fn run_traj_judge(lazy: &mut LazyEngine<'_>, args: &TrajectoryJudgeArgs) -> Resu } ); if let Some(score) = args.score { - eprintln!(" fixed score={score} scope={scope:?} rubrics={}", rubric_ids.join(",")); + eprintln!( + " fixed score={score} scope={scope:?} rubrics={}", + rubric_ids.join(",") + ); } let mut ok = 0usize; diff --git a/crates/persisting-cli/tests/dlcapt_backend.rs b/crates/persisting-cli/tests/dlcapt_backend.rs new file mode 100644 index 0000000..7f867bc --- /dev/null +++ b/crates/persisting-cli/tests/dlcapt_backend.rs @@ -0,0 +1,34 @@ +use std::process::Command; + +#[cfg(not(feature = "dlcapt"))] +#[test] +fn dlcapt_backend_without_feature_explains_how_to_enable_it() { + let output = Command::new(env!("CARGO_BIN_EXE_persisting")) + .args(["traj", "proxy", "--backend", "dlcapt", "-c", "proxy.toml"]) + .output() + .unwrap(); + assert!(!output.status.success()); + assert!(String::from_utf8_lossy(&output.stderr) + .contains("rebuild persisting-cli with --features dlcapt")); +} + +#[cfg(feature = "dlcapt")] +#[test] +fn dlcapt_backend_rejects_capture_only_output_directory() { + let output = Command::new(env!("CARGO_BIN_EXE_persisting")) + .args([ + "traj", + "proxy", + "--backend", + "dlcapt", + "-c", + "proxy.toml", + "-o", + "store", + ]) + .output() + .unwrap(); + assert!(!output.status.success()); + assert!(String::from_utf8_lossy(&output.stderr) + .contains("-o is only supported by the capture backend")); +} diff --git a/crates/persisting-cli/tests/dlcapt_backend_e2e.rs b/crates/persisting-cli/tests/dlcapt_backend_e2e.rs new file mode 100644 index 0000000..3df32d9 --- /dev/null +++ b/crates/persisting-cli/tests/dlcapt_backend_e2e.rs @@ -0,0 +1,225 @@ +use std::fs; +use std::net::TcpListener; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command, Stdio}; +use std::thread; +use std::time::{Duration, Instant}; + +use tempfile::TempDir; + +const ATTEMPTS: usize = 3; + +fn reserve_port() -> u16 { + TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() + .port() +} + +fn reserve_distinct_ports() -> (u16, u16) { + let public = reserve_port(); + loop { + let admin = reserve_port(); + if admin != public { + return (public, admin); + } + } +} + +fn write_config(dir: &TempDir, public_port: u16, admin_port: u16) -> PathBuf { + let path = dir.path().join("proxy.toml"); + fs::write( + &path, + format!( + r#" +listen = "127.0.0.1:{public_port}" +admin_listen = "127.0.0.1:{admin_port}" +store_dir = "{}" +agent_id = "cli-e2e" +session_header = "x-persisting-session-id" +default_session_id = "default" +preserve_raw = false +base_session_path = "/v1/sessions" + +[storage] +authoritative = "json_file" +also = ["md"] + +[[models]] +name = "*" +provider = "openai" +upstream_base_url = "https://example.invalid/v1" +api_key = "" +"#, + dir.path().join("store").display(), + ), + ) + .unwrap(); + path +} + +fn wait_for_ok(url: &str) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(10); + let client = reqwest::blocking::Client::new(); + + loop { + let last_error = match client.get(url).send() { + Ok(response) if response.status().is_success() => return Ok(()), + Ok(response) => format!("received HTTP {}", response.status()), + Err(error) => error.to_string(), + }; + + if Instant::now() >= deadline { + return Err(format!( + "service did not become ready: {url}; last error: {last_error}" + )); + } + thread::sleep(Duration::from_millis(50)); + } +} + +fn stop(child: Child) { + let mut child = child; + let _ = child.kill(); + let _ = child.wait_with_output(); +} + +fn failed_attempt( + attempt: usize, + public_url: &str, + admin_url: &str, + reason: impl std::fmt::Display, + child: Option, +) -> String { + let (status, stdout, stderr) = match child { + Some(mut child) => { + let _ = child.kill(); + match child.wait_with_output() { + Ok(output) => ( + output.status.to_string(), + String::from_utf8_lossy(&output.stdout).into_owned(), + String::from_utf8_lossy(&output.stderr).into_owned(), + ), + Err(error) => ( + "".into(), + String::new(), + format!("failed collecting child output: {error}"), + ), + } + } + None => ("".into(), String::new(), String::new()), + }; + + format!( + "attempt {attempt}/{ATTEMPTS}\n\ + public URL: {public_url}\n\ + admin URL: {admin_url}\n\ + readiness failure: {reason}\n\ + exit status: {status}\n\ + stdout:\n{stdout}\n\ + stderr:\n{stderr}" + ) +} + +fn run_case_with_retry(case_name: &str, mut start: F, mut verify: V) +where + F: FnMut(&TempDir, &Path) -> Result, + V: FnMut(&TempDir) -> Result<(), String>, +{ + let mut failures = Vec::new(); + + for attempt in 1..=ATTEMPTS { + let dir = tempfile::tempdir().unwrap(); + let (public, admin) = reserve_distinct_ports(); + let config = write_config(&dir, public, admin); + let public_url = format!("http://127.0.0.1:{public}/healthz"); + let admin_url = format!("http://127.0.0.1:{admin}/admin/sessions"); + + match start(&dir, &config) { + Ok(child) => match wait_for_ok(&public_url) + .and_then(|_| wait_for_ok(&admin_url)) + .and_then(|_| verify(&dir)) + { + Ok(()) => { + stop(child); + return; + } + Err(reason) => failures.push(failed_attempt( + attempt, + &public_url, + &admin_url, + reason, + Some(child), + )), + }, + Err(reason) => failures.push(failed_attempt( + attempt, + &public_url, + &admin_url, + reason, + None, + )), + } + } + + panic!( + "{case_name} failed after {ATTEMPTS} attempts:\n{}", + failures.join("\n\n") + ); +} + +fn dlcapt_bin() -> PathBuf { + std::env::var_os("DLCAPT_BIN") + .map(PathBuf::from) + .filter(|path| path.is_file()) + .expect("DLCAPT_BIN must point to a prebuilt dlcapt binary") +} + +#[test] +fn standalone_dlcapt_serves_health_and_admin() { + run_case_with_retry( + "standalone dlcapt", + |_, config| { + Command::new(dlcapt_bin()) + .arg(config) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .map_err(|error| error.to_string()) + }, + |_| Ok(()), + ); +} + +#[test] +fn cli_dlcapt_backend_ignores_capture_storage_environment() { + run_case_with_retry( + "persisting-cli dlcapt backend", + |dir, config| { + Command::new(env!("CARGO_BIN_EXE_persisting")) + .args([ + "traj", + "proxy", + "--backend", + "dlcapt", + "-c", + config.to_str().unwrap(), + ]) + .env( + "PERSISTING_CAPTURE_STORAGE", + dir.path().join("must-not-be-used"), + ) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .map_err(|error| error.to_string()) + }, + |dir| { + let capture_store = dir.path().join("must-not-be-used"); + (!capture_store.exists()) + .then_some(()) + .ok_or_else(|| format!("dlcapt unexpectedly created {}", capture_store.display())) + }, + ); +} 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..d7caa54 --- /dev/null +++ b/crates/persisting-dlcapt/README.md @@ -0,0 +1,64 @@ +# persisting-dlcapt + +`persisting-dlcapt` 是 Persisting workspace 内独立运行的 OpenAI 兼容代理与轨迹采集服务。 +第一阶段不依赖 `persisting-capture` 或 `persisting-engine`，以保持现有 dlcapt HTTP、session 和存储契约。 + +## Source and license + +The maintained runtime source is this `crates/persisting-dlcapt` package. It was +migrated from Capture `external/dlcapt`; that path describes the source +provenance, not the current runtime location. The package follows the +workspace's Apache-2.0 license policy; see the repository `NOTICE`. + +## Prepare a local configuration + +```bash +cd /path/to/Persisting +export DLCAPT_CONFIG="$HOME/.config/persisting/dlcapt-openclaw.toml" +export DLCAPT_STORE_DIR="$HOME/.local/share/persisting/dlcapt" +export DLCAPT_UPSTREAM_BASE_URL="https://your-upstream.example/v1" +mkdir -p "$(dirname "$DLCAPT_CONFIG")" +cp crates/persisting-dlcapt/config/proxy.openclaw-test.example.toml "$DLCAPT_CONFIG" +sed -i \ + -e "s|__STORE_DIR__|$DLCAPT_STORE_DIR|g" \ + -e "s|__UPSTREAM_BASE_URL__|$DLCAPT_UPSTREAM_BASE_URL|g" \ + "$DLCAPT_CONFIG" +``` + +CLI is positional only: `dlcapt `. There is no `serve -c`. + +## Run (development) + +```bash +cargo run -p persisting-dlcapt -- "$DLCAPT_CONFIG" +``` + +## Run through persisting-cli + +Build the optional backend first: + +```bash +cargo run -p persisting-cli --features dlcapt -- \ + traj proxy --backend dlcapt -c "$DLCAPT_CONFIG" +``` + +This is foreground-only. `store_dir`, storage sinks, model routes, and public/admin +listen addresses come from the dlcapt TOML; `-o`, `-f`, `--debug`, and daemon +actions remain capture-backend options. + +## Public / admin + +- Public: `listen` (default example `127.0.0.1:19081`) — `/healthz`, `/readyz`, `/v1/models`, chat/completions, session routes +- Admin: `admin_listen` (default example `127.0.0.1:19082`) — `/admin/sessions`, `/admin/errors` + +Session priority: URL `{id}` > Header > body `metadata.session_id` > `default_session_id`. + +## Storage + +Use `config/proxy.openclaw-test.example.toml` as the supported safe template. + +## Safety + +- Examples bind `127.0.0.1` only. Binding a wildcard all-interface address is for private deploys with firewall / reverse-proxy ACL; admin must not be exposed to untrusted networks. +- Trajectories may contain prompts, responses, and header-derived metadata — never commit `var/` or `.capture/`. +- Do not publish online/beta configs with real upstreams, buckets, or credentials. diff --git a/crates/persisting-dlcapt/config/proxy.openclaw-test.example.toml b/crates/persisting-dlcapt/config/proxy.openclaw-test.example.toml new file mode 100644 index 0000000..77f330c --- /dev/null +++ b/crates/persisting-dlcapt/config/proxy.openclaw-test.example.toml @@ -0,0 +1,34 @@ +# Rendered by Capture openclaw helper: replace __STORE_DIR__ and __UPSTREAM_BASE_URL__. +listen = "127.0.0.1:19081" +admin_listen = "127.0.0.1:19082" +store_dir = "__STORE_DIR__" +agent_id = "openclaw" +session_header = "x-persisting-session-id" +default_session_id = "default" +preserve_raw = false +base_session_path = "/v1/sessions" + +[storage] +authoritative = "json_file" +also = ["md"] + +[export.defaults] +env_name = "openclaw" +job_id = "dlcapt" + +[export.session_metadata] +source = "dlcapt-proxy" + +[[models]] +name = "kimi-k2.5" +display_name = "Kimi K2.5" +provider = "openai" +upstream_base_url = "__UPSTREAM_BASE_URL__" +api_key = "" + +[[models]] +name = "*" +display_name = "Fallback" +provider = "openai" +upstream_base_url = "__UPSTREAM_BASE_URL__" +api_key = "" 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..433a5a3 --- /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 + && 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..bf8b0a9 --- /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://s3.example.invalid: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..5464414 --- /dev/null +++ b/crates/persisting-dlcapt/src/config.rs @@ -0,0 +1,449 @@ +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 + && !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")) + && !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..f85849c --- /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) + && !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..8c030fc --- /dev/null +++ b/crates/persisting-dlcapt/src/dialogue/responses.rs @@ -0,0 +1,175 @@ +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") + && 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") + && let Some(text) = part.get("text").and_then(Value::as_str) + && !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) + && !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") + && 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() + && 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 new file mode 100644 index 0000000..0bcf617 --- /dev/null +++ b/crates/persisting-dlcapt/src/lib.rs @@ -0,0 +1,11 @@ +pub mod audit; +pub mod capture; +pub mod config; +pub mod dialogue; +pub mod proxy; +pub mod router; +pub mod service; +pub mod session; +pub mod tlv; + +pub use service::serve; diff --git a/crates/persisting-dlcapt/src/main.rs b/crates/persisting-dlcapt/src/main.rs new file mode 100644 index 0000000..b9f849b --- /dev/null +++ b/crates/persisting-dlcapt/src/main.rs @@ -0,0 +1,22 @@ +use anyhow::{Context, Result}; +use persisting_dlcapt::config::ProxyConfig; +use std::path::PathBuf; + +#[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()))?; + + persisting_dlcapt::serve(config).await +} diff --git a/crates/persisting-dlcapt/src/proxy.rs b/crates/persisting-dlcapt/src/proxy.rs new file mode 100644 index 0000000..5e9d05f --- /dev/null +++ b/crates/persisting-dlcapt/src/proxy.rs @@ -0,0 +1,728 @@ +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 +} + +// The capture event is assembled from both request and response paths; grouping +// its inputs would be a behavior-changing refactor during this mechanical migration. +#[allow(clippy::too_many_arguments)] +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() + && 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/service.rs b/crates/persisting-dlcapt/src/service.rs new file mode 100644 index 0000000..de0e82c --- /dev/null +++ b/crates/persisting-dlcapt/src/service.rs @@ -0,0 +1,34 @@ +use crate::config::ProxyConfig; +use crate::proxy::{AppState, build_admin_router, build_public_router}; +use anyhow::{Context, Result}; +use std::net::SocketAddr; +use tracing::info; + +pub async fn serve(config: ProxyConfig) -> Result<()> { + let state = AppState::new(config); + let public_addr: SocketAddr = state + .listen_addr() + .parse() + .context("invalid listen address")?; + let admin_addr: SocketAddr = state + .admin_listen_addr() + .parse() + .context("invalid admin listen address")?; + + let public_listener = tokio::net::TcpListener::bind(public_addr) + .await + .context("failed binding public listener")?; + let admin_listener = tokio::net::TcpListener::bind(admin_addr) + .await + .context("failed binding admin listener")?; + + info!("dlcapt public API listening on {public_addr}"); + info!("dlcapt admin API listening on {admin_addr}"); + + tokio::try_join!( + axum::serve(public_listener, build_public_router(state.clone())), + axum::serve(admin_listener, build_admin_router(state)), + ) + .context("proxy server exited with error")?; + Ok(()) +} diff --git a/crates/persisting-dlcapt/src/session/mod.rs b/crates/persisting-dlcapt/src/session/mod.rs new file mode 100644 index 0000000..ea06b1a --- /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()) + && 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..a824728 --- /dev/null +++ b/crates/persisting-dlcapt/src/session/resolver.rs @@ -0,0 +1,199 @@ +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, + ) && header != raw_path + { + conflicts.push(SessionConflict { + source: SessionSource::Header, + value: header, + }); + } + if let Some(body) = extract_body_metadata_session(ctx.body) + && 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 + && let Some(value) = header + { + conflicts.push(SessionConflict { + source: SessionSource::Header, + value, + }); + } + if source != SessionSource::BodyMetadata + && 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/config_examples.rs b/crates/persisting-dlcapt/tests/config_examples.rs new file mode 100644 index 0000000..5f14d7c --- /dev/null +++ b/crates/persisting-dlcapt/tests/config_examples.rs @@ -0,0 +1,82 @@ +use std::path::{Path, PathBuf}; + +const OPENCLAW_TEMPLATE: &str = "proxy.openclaw-test.example.toml"; +const ALLOWED_PLACEHOLDERS: [&str; 2] = ["__STORE_DIR__", "__UPSTREAM_BASE_URL__"]; + +fn config_root() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("config") +} + +fn config_files() -> Vec { + let root = config_root(); + let mut files: Vec = std::fs::read_dir(&root) + .unwrap_or_else(|_| panic!("missing config dir: {}", root.display())) + .filter_map(|e| e.ok().map(|e| e.path())) + .filter(|p| p.is_file()) + .collect(); + files.sort(); + files +} + +fn placeholders(value: &toml::Value, values: &mut Vec) { + match value { + toml::Value::String(value) if value.contains("__") => values.push(value.clone()), + toml::Value::Array(values_in_array) => { + for value in values_in_array { + placeholders(value, values); + } + } + toml::Value::Table(values_in_table) => { + for value in values_in_table.values() { + placeholders(value, values); + } + } + _ => {} + } +} + +#[test] +fn config_directory_contains_only_the_openclaw_template() { + assert_eq!( + config_files(), + vec![config_root().join(OPENCLAW_TEMPLATE)], + "the OpenClaw template is the only supported config example" + ); +} + +#[test] +fn openclaw_template_uses_only_safe_placeholders() { + let path = config_root().join(OPENCLAW_TEMPLATE); + let raw = std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("{e}")); + let parsed: toml::Value = toml::from_str(&raw).unwrap_or_else(|e| panic!("{e}")); + let mut found_placeholders = Vec::new(); + placeholders(&parsed, &mut found_placeholders); + found_placeholders.sort(); + assert!( + found_placeholders + .iter() + .all(|value| ALLOWED_PLACEHOLDERS.contains(&value.as_str())) + ); + for placeholder in ALLOWED_PLACEHOLDERS { + assert!(found_placeholders.iter().any(|value| value == placeholder)); + } + assert_eq!( + parsed.get("store_dir").and_then(toml::Value::as_str), + Some("__STORE_DIR__") + ); + let models = parsed + .get("models") + .and_then(toml::Value::as_array) + .expect("models array"); + for model in models { + assert_eq!( + model.get("upstream_base_url").and_then(toml::Value::as_str), + Some("__UPSTREAM_BASE_URL__") + ); + assert_eq!(model.get("api_key").and_then(toml::Value::as_str), Some("")); + } + assert!(raw.contains("kimi-k2.5")); + assert!(raw.contains("127.0.0.1:19081")); + assert!(!raw.contains("ailab-pj")); + assert!(!raw.contains("0.0.0.0")); +} diff --git a/crates/persisting-dlcapt/tests/http_proxy.rs b/crates/persisting-dlcapt/tests/http_proxy.rs new file mode 100644 index 0000000..e562fc9 --- /dev/null +++ b/crates/persisting-dlcapt/tests/http_proxy.rs @@ -0,0 +1,416 @@ +use axum::body::Body; +use axum::http::{Request, StatusCode, header}; +use axum::response::IntoResponse; +use axum::routing::post; +use axum::{Json, Router}; +use http_body_util::BodyExt; +use persisting_dlcapt::config::{ExportConfig, ModelRoute, ProxyConfig, StorageConfig}; +use persisting_dlcapt::proxy::{AppState, build_admin_router, build_public_router}; +use serde_json::{Value, json}; +use std::path::Path; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tempfile::tempdir; +use tokio::sync::oneshot; +use tower::ServiceExt; + +type Captured = Arc>>; + +async fn start_mock_upstream(captured: Captured) -> (String, oneshot::Sender<()>) { + let captured_chat = Arc::clone(&captured); + let captured_responses = Arc::clone(&captured); + let app = Router::new() + .route( + "/v1/chat/completions", + post(move |Json(body): Json| { + let captured_chat = Arc::clone(&captured_chat); + async move { + captured_chat + .lock() + .unwrap() + .push(("/v1/chat/completions".into(), body.clone())); + if body.get("stream").and_then(|v| v.as_bool()) == Some(true) { + ( + [(header::CONTENT_TYPE, "text/event-stream")], + "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":null}]}\n\n\ + data: [DONE]\n\n", + ) + .into_response() + } else { + Json(json!({ + "id": "chatcmpl-test", + "choices": [{ + "index": 0, + "message": {"role":"assistant","content":"ok"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens":1,"completion_tokens":1,"total_tokens":2} + })) + .into_response() + } + } + }), + ) + .route( + "/v1/responses", + post(move |Json(body): Json| { + let captured_responses = Arc::clone(&captured_responses); + async move { + captured_responses + .lock() + .unwrap() + .push(("/v1/responses".into(), body)); + Json(json!({ + "id": "resp-test", + "output": [{"content":[{"type":"output_text","text":"r-ok"}]}] + })) + .into_response() + } + }), + ); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(async { + let _ = shutdown_rx.await; + }) + .await + .unwrap(); + }); + (format!("http://{addr}/v1"), shutdown_tx) +} + +fn base_config(store_dir: &str, upstream: &str) -> ProxyConfig { + ProxyConfig { + listen: "127.0.0.1:0".to_string(), + admin_listen: "127.0.0.1:0".to_string(), + store_dir: store_dir.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: upstream.to_string(), + api_key: Some("".to_string()), + }, + ModelRoute { + name: "*".to_string(), + display_name: Some("Fallback".to_string()), + provider: "openai".to_string(), + upstream_base_url: upstream.to_string(), + api_key: Some("".to_string()), + }, + ], + } +} + +async fn body_bytes(response: axum::response::Response) -> bytes::Bytes { + response.into_body().collect().await.unwrap().to_bytes() +} + +async fn body_json(response: axum::response::Response) -> Value { + serde_json::from_slice(&body_bytes(response).await).unwrap() +} + +async fn wait_for_file(path: &Path) { + tokio::time::timeout(Duration::from_secs(5), async { + while !path.is_file() { + tokio::time::sleep(Duration::from_millis(25)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("capture artifact was not written: {}", path.display())); +} + +async fn closed_loopback_upstream() -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + drop(listener); + format!("http://{addr}/v1") +} + +#[tokio::test] +async fn health_ready_models_admin_sessions_contract() { + let dir = tempdir().unwrap(); + let upstream = closed_loopback_upstream().await; + let state = AppState::new(base_config(dir.path().to_str().unwrap(), &upstream)); + let public = build_public_router(state.clone()); + let admin = build_admin_router(state); + + let health = public + .clone() + .oneshot( + Request::builder() + .uri("/healthz") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(health.status(), StatusCode::OK); + assert_eq!(body_json(health).await["status"], "ok"); + + let ready = public + .clone() + .oneshot( + Request::builder() + .uri("/readyz") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(ready.status(), StatusCode::OK); + assert_eq!(body_json(ready).await["status"], "ready"); + + let models = public + .clone() + .oneshot( + Request::builder() + .uri("/v1/models") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(models.status(), StatusCode::OK); + let models_json = body_json(models).await; + assert!( + models_json["data"] + .as_array() + .unwrap() + .iter() + .any(|m| m["id"] == "kimi-k2.5") + ); + + let sessions = admin + .oneshot( + Request::builder() + .uri("/admin/sessions") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(sessions.status(), StatusCode::OK); + assert!(body_json(sessions).await.get("sessions").is_some()); +} + +#[tokio::test] +async fn flat_chat_completions_forwards_to_upstream() { + let dir = tempdir().unwrap(); + let captured: Captured = Arc::new(Mutex::new(Vec::new())); + let (upstream, shutdown) = start_mock_upstream(Arc::clone(&captured)).await; + let state = AppState::new(base_config(dir.path().to_str().unwrap(), &upstream)); + let app = build_public_router(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": "kimi-k2.5", + "messages": [{"role":"user","content":"ping"}], + "stream": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let captured_req = captured.lock().unwrap(); + assert_eq!(captured_req[0].0, "/v1/chat/completions"); + assert_eq!(captured_req[0].1["model"], "kimi-k2.5"); + let _ = shutdown.send(()); +} + +#[tokio::test] +async fn session_url_overrides_header_and_body_session() { + let dir = tempdir().unwrap(); + let captured: Captured = Arc::new(Mutex::new(Vec::new())); + let (upstream, shutdown) = start_mock_upstream(Arc::clone(&captured)).await; + let state = AppState::new(base_config(dir.path().to_str().unwrap(), &upstream)); + let app = build_public_router(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/sessions/session-a/chat/completions") + .header(header::CONTENT_TYPE, "application/json") + .header("x-persisting-session-id", "header-session") + .body(Body::from( + json!({ + "model": "kimi-k2.5", + "messages": [{"role":"user","content":"ping"}], + "metadata": {"session_id": "body-session"}, + "stream": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + { + let captured_req = captured.lock().unwrap(); + assert_eq!(captured_req[0].0, "/v1/chat/completions"); + assert_eq!(captured_req[0].1["model"], "kimi-k2.5"); + assert_eq!(captured_req[0].1["messages"][0]["role"], "user"); + assert_eq!(captured_req[0].1["messages"][0]["content"], "ping"); + assert_eq!(captured_req[0].1["stream"], false); + } + wait_for_file(&dir.path().join("session-a/trajectory.md")).await; + wait_for_file(&dir.path().join("session-a/session_steps.json")).await; + let _ = shutdown.send(()); +} + +#[tokio::test] +async fn session_responses_forwards_to_responses_path() { + let dir = tempdir().unwrap(); + let captured: Captured = Arc::new(Mutex::new(Vec::new())); + let (upstream, shutdown) = start_mock_upstream(Arc::clone(&captured)).await; + let state = AppState::new(base_config(dir.path().to_str().unwrap(), &upstream)); + let app = build_public_router(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/sessions/session-a/responses") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": "kimi-k2.5", + "input": "ping" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let captured_req = captured.lock().unwrap(); + assert_eq!(captured_req[0].0, "/v1/responses"); + let _ = shutdown.send(()); +} + +#[tokio::test] +async fn streaming_chat_returns_sse_and_writes_capture() { + let dir = tempdir().unwrap(); + let captured: Captured = Arc::new(Mutex::new(Vec::new())); + let (upstream, shutdown) = start_mock_upstream(Arc::clone(&captured)).await; + let state = AppState::new(base_config(dir.path().to_str().unwrap(), &upstream)); + let app = build_public_router(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/sessions/session-stream/chat/completions") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": "kimi-k2.5", + "messages": [{"role":"user","content":"ping"}], + "stream": true + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let ct = response + .headers() + .get(header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + assert!(ct.contains("text/event-stream"), "content-type={ct}"); + let body = String::from_utf8(body_bytes(response).await.to_vec()).unwrap(); + assert!( + body.contains( + "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":null}]}" + ), + "SSE delta was not relayed: {body}" + ); + assert!( + body.contains("data: [DONE]"), + "SSE completion was not relayed: {body}" + ); + { + let captured_req = captured.lock().unwrap(); + assert_eq!(captured_req[0].0, "/v1/chat/completions"); + assert_eq!(captured_req[0].1["stream"], true); + } + wait_for_file(&dir.path().join("session-stream/trajectory.md")).await; + wait_for_file(&dir.path().join("session-stream/session_steps.json")).await; + let _ = shutdown.send(()); +} + +#[tokio::test] +async fn upstream_unreachable_returns_502_and_admin_error() { + let dir = tempdir().unwrap(); + let upstream = closed_loopback_upstream().await; + let state = AppState::new(base_config(dir.path().to_str().unwrap(), &upstream)); + let public = build_public_router(state.clone()); + let admin = build_admin_router(state); + + let response = public + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from( + json!({ + "model": "kimi-k2.5", + "messages": [{"role":"user","content":"ping"}], + "stream": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + + let errors = admin + .oneshot( + Request::builder() + .uri("/admin/errors") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(errors.status(), StatusCode::OK); + let payload = body_json(errors).await; + let list = payload["recent_errors"].as_array().unwrap_or_else(|| { + panic!("expected recent_errors array, got {payload}"); + }); + assert!( + !list.is_empty(), + "admin errors should record upstream failure: {payload}" + ); +} 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("