diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 0e2232fcbf..1b31aece30 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1268,6 +1268,8 @@ "ring_0.17.14": "{\"dependencies\":[{\"default_features\":false,\"kind\":\"build\",\"name\":\"cc\",\"req\":\"^1.2.8\"},{\"default_features\":false,\"name\":\"cfg-if\",\"req\":\"^1.0.0\"},{\"name\":\"getrandom\",\"req\":\"^0.2.10\"},{\"default_features\":false,\"name\":\"libc\",\"req\":\"^0.2.148\",\"target\":\"cfg(all(any(all(target_arch = \\\"aarch64\\\", target_endian = \\\"little\\\"), all(target_arch = \\\"arm\\\", target_endian = \\\"little\\\")), any(target_os = \\\"android\\\", target_os = \\\"linux\\\")))\"},{\"default_features\":false,\"name\":\"libc\",\"req\":\"^0.2.155\",\"target\":\"cfg(all(all(target_arch = \\\"aarch64\\\", target_endian = \\\"little\\\"), target_vendor = \\\"apple\\\", any(target_os = \\\"ios\\\", target_os = \\\"macos\\\", target_os = \\\"tvos\\\", target_os = \\\"visionos\\\", target_os = \\\"watchos\\\")))\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"libc\",\"req\":\"^0.2.148\",\"target\":\"cfg(any(unix, windows, target_os = \\\"wasi\\\"))\"},{\"name\":\"untrusted\",\"req\":\"^0.9\"},{\"default_features\":false,\"features\":[\"std\"],\"kind\":\"dev\",\"name\":\"wasm-bindgen-test\",\"req\":\"^0.3.37\",\"target\":\"cfg(all(target_arch = \\\"wasm32\\\", target_os = \\\"unknown\\\"))\"},{\"features\":[\"Win32_Foundation\",\"Win32_System_Threading\"],\"name\":\"windows-sys\",\"req\":\"^0.52\",\"target\":\"cfg(all(all(target_arch = \\\"aarch64\\\", target_endian = \\\"little\\\"), target_os = \\\"windows\\\"))\"}],\"features\":{\"alloc\":[],\"default\":[\"alloc\",\"dev_urandom_fallback\"],\"dev_urandom_fallback\":[],\"less-safe-getrandom-custom-or-rdrand\":[],\"less-safe-getrandom-espidf\":[],\"slow_tests\":[],\"std\":[\"alloc\"],\"test_logging\":[],\"unstable-testing-arm-no-hw\":[],\"unstable-testing-arm-no-neon\":[],\"wasm32_unknown_unknown_js\":[\"getrandom/js\"]}}", "rmcp-macros_0.15.0": "{\"dependencies\":[{\"name\":\"darling\",\"req\":\"^0.23\"},{\"name\":\"proc-macro2\",\"req\":\"^1\"},{\"name\":\"quote\",\"req\":\"^1\"},{\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"features\":[\"full\"],\"name\":\"syn\",\"req\":\"^2\"}],\"features\":{}}", "rmcp_0.15.0": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"anyhow\",\"req\":\"^1.0\"},{\"name\":\"async-trait\",\"req\":\"^0.1.89\"},{\"kind\":\"dev\",\"name\":\"async-trait\",\"req\":\"^0.1\"},{\"name\":\"axum\",\"optional\":true,\"req\":\"^0.8\"},{\"name\":\"base64\",\"optional\":true,\"req\":\"^0.22\"},{\"name\":\"bytes\",\"optional\":true,\"req\":\"^1\"},{\"default_features\":false,\"features\":[\"serde\",\"clock\",\"std\",\"oldtime\"],\"name\":\"chrono\",\"req\":\"^0.4.38\",\"target\":\"cfg(all(target_family = \\\"wasm\\\", target_os = \\\"unknown\\\"))\"},{\"features\":[\"serde\"],\"name\":\"chrono\",\"req\":\"^0.4.38\",\"target\":\"cfg(not(all(target_family = \\\"wasm\\\", target_os = \\\"unknown\\\")))\"},{\"name\":\"futures\",\"req\":\"^0.3\"},{\"name\":\"http\",\"optional\":true,\"req\":\"^1\"},{\"name\":\"http-body\",\"optional\":true,\"req\":\"^1\"},{\"name\":\"http-body-util\",\"optional\":true,\"req\":\"^0.1\"},{\"default_features\":false,\"features\":[\"reqwest\"],\"name\":\"oauth2\",\"optional\":true,\"req\":\"^5.0\"},{\"name\":\"pastey\",\"optional\":true,\"req\":\"^0.2.0\"},{\"name\":\"pin-project-lite\",\"req\":\"^0.2\"},{\"features\":[\"tokio1\"],\"name\":\"process-wrap\",\"optional\":true,\"req\":\"^9.0\"},{\"name\":\"rand\",\"optional\":true,\"req\":\"^0.9\"},{\"default_features\":false,\"features\":[\"json\",\"stream\"],\"name\":\"reqwest\",\"optional\":true,\"req\":\"^0.12\"},{\"name\":\"rmcp-macros\",\"optional\":true,\"req\":\"^0.15.0\"},{\"features\":[\"chrono04\"],\"name\":\"schemars\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"chrono04\"],\"kind\":\"dev\",\"name\":\"schemars\",\"req\":\"^1.1.0\"},{\"features\":[\"derive\",\"rc\"],\"name\":\"serde\",\"req\":\"^1.0\"},{\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"name\":\"sse-stream\",\"optional\":true,\"req\":\"^0.2\"},{\"name\":\"thiserror\",\"req\":\"^2\"},{\"features\":[\"sync\",\"macros\",\"rt\",\"time\"],\"name\":\"tokio\",\"req\":\"^1\"},{\"features\":[\"full\"],\"kind\":\"dev\",\"name\":\"tokio\",\"req\":\"^1\"},{\"name\":\"tokio-stream\",\"optional\":true,\"req\":\"^0.1\"},{\"name\":\"tokio-util\",\"req\":\"^0.7\"},{\"name\":\"tower-service\",\"optional\":true,\"req\":\"^0.3\"},{\"name\":\"tracing\",\"req\":\"^0.1\"},{\"features\":[\"env-filter\",\"std\",\"fmt\"],\"kind\":\"dev\",\"name\":\"tracing-subscriber\",\"req\":\"^0.3\"},{\"name\":\"url\",\"optional\":true,\"req\":\"^2.4\"},{\"features\":[\"v4\"],\"name\":\"uuid\",\"optional\":true,\"req\":\"^1\"}],\"features\":{\"__reqwest\":[\"dep:reqwest\"],\"auth\":[\"dep:oauth2\",\"__reqwest\",\"dep:url\"],\"client\":[\"dep:tokio-stream\"],\"client-side-sse\":[\"dep:sse-stream\",\"dep:http\"],\"default\":[\"base64\",\"macros\",\"server\"],\"elicitation\":[\"dep:url\"],\"macros\":[\"dep:rmcp-macros\",\"dep:pastey\"],\"reqwest\":[\"__reqwest\",\"reqwest?/rustls-tls\"],\"reqwest-native-tls\":[\"__reqwest\",\"reqwest?/native-tls\"],\"reqwest-tls-no-provider\":[\"__reqwest\",\"reqwest?/rustls-tls-no-provider\"],\"schemars\":[\"dep:schemars\"],\"server\":[\"transport-async-rw\",\"dep:schemars\",\"dep:pastey\"],\"server-side-http\":[\"uuid\",\"dep:rand\",\"dep:tokio-stream\",\"dep:http\",\"dep:http-body\",\"dep:http-body-util\",\"dep:bytes\",\"dep:sse-stream\",\"dep:axum\",\"tower\"],\"tower\":[\"dep:tower-service\"],\"transport-async-rw\":[\"tokio/io-util\",\"tokio-util/codec\"],\"transport-child-process\":[\"transport-async-rw\",\"tokio/process\",\"dep:process-wrap\"],\"transport-io\":[\"transport-async-rw\",\"tokio/io-std\"],\"transport-streamable-http-client\":[\"client-side-sse\",\"transport-worker\"],\"transport-streamable-http-client-reqwest\":[\"transport-streamable-http-client\",\"__reqwest\"],\"transport-streamable-http-server\":[\"transport-streamable-http-server-session\",\"server-side-http\",\"transport-worker\"],\"transport-streamable-http-server-session\":[\"transport-async-rw\",\"dep:tokio-stream\"],\"transport-worker\":[\"dep:tokio-stream\"]}}", + "rmp-serde_1.3.1": "{\"dependencies\":[{\"name\":\"rmp\",\"req\":\"^0.8.14\"},{\"name\":\"serde\",\"req\":\"^1.0.228\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.228\"},{\"kind\":\"dev\",\"name\":\"serde_bytes\",\"req\":\"^0.11.19\"}],\"features\":{}}", + "rmp_0.8.15": "{\"dependencies\":[{\"default_features\":false,\"name\":\"num-traits\",\"req\":\"^0.2.19\"},{\"kind\":\"dev\",\"name\":\"quickcheck\",\"req\":\"^1.0.3\"}],\"features\":{\"default\":[\"std\"],\"std\":[\"num-traits/std\"]}}", "rsa_0.9.10": "{\"dependencies\":[{\"features\":[\"alloc\"],\"kind\":\"dev\",\"name\":\"base64ct\",\"req\":\"^1\"},{\"default_features\":false,\"name\":\"const-oid\",\"req\":\"^0.9\"},{\"default_features\":false,\"features\":[\"alloc\",\"oid\"],\"name\":\"digest\",\"req\":\"^0.10.5\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4.1\"},{\"default_features\":false,\"features\":[\"i128\",\"prime\",\"zeroize\"],\"name\":\"num-bigint\",\"package\":\"num-bigint-dig\",\"req\":\"^0.8.6\"},{\"default_features\":false,\"name\":\"num-integer\",\"req\":\"^0.1.39\"},{\"default_features\":false,\"features\":[\"libm\"],\"name\":\"num-traits\",\"req\":\"^0.2.9\"},{\"default_features\":false,\"features\":[\"alloc\",\"pkcs8\"],\"name\":\"pkcs1\",\"req\":\"^0.7.5\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"pkcs8\",\"req\":\"^0.10.2\"},{\"kind\":\"dev\",\"name\":\"proptest\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8\"},{\"kind\":\"dev\",\"name\":\"rand_chacha\",\"req\":\"^0.3\"},{\"default_features\":false,\"name\":\"rand_core\",\"req\":\"^0.6.4\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.6\"},{\"kind\":\"dev\",\"name\":\"rand_xorshift\",\"req\":\"^0.3\"},{\"default_features\":false,\"features\":[\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.184\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1.0.89\"},{\"default_features\":false,\"features\":[\"oid\"],\"name\":\"sha1\",\"optional\":true,\"req\":\"^0.10.5\"},{\"default_features\":false,\"features\":[\"oid\"],\"kind\":\"dev\",\"name\":\"sha1\",\"req\":\"^0.10.5\"},{\"default_features\":false,\"features\":[\"oid\"],\"name\":\"sha2\",\"optional\":true,\"req\":\"^0.10.6\"},{\"default_features\":false,\"features\":[\"oid\"],\"kind\":\"dev\",\"name\":\"sha2\",\"req\":\"^0.10.6\"},{\"default_features\":false,\"features\":[\"oid\"],\"kind\":\"dev\",\"name\":\"sha3\",\"req\":\"^0.10.7\"},{\"default_features\":false,\"features\":[\"alloc\",\"digest\",\"rand_core\"],\"name\":\"signature\",\"req\":\">2.0, <2.3\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"spki\",\"req\":\"^0.7.3\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2.1.1\"},{\"features\":[\"alloc\"],\"name\":\"zeroize\",\"req\":\"^1.5\"}],\"features\":{\"default\":[\"std\",\"pem\",\"u64_digit\"],\"getrandom\":[\"rand_core/getrandom\"],\"hazmat\":[],\"nightly\":[\"num-bigint/nightly\"],\"pem\":[\"pkcs1/pem\",\"pkcs8/pem\"],\"pkcs5\":[\"pkcs8/encryption\"],\"serde\":[\"dep:serde\",\"num-bigint/serde\"],\"std\":[\"digest/std\",\"pkcs1/std\",\"pkcs8/std\",\"rand_core/std\",\"signature/std\"],\"u64_digit\":[\"num-bigint/u64_digit\"]}}", "rtrb_0.3.3": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.8\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"crossbeam-utils\",\"req\":\"^0.8\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.10\"}],\"features\":{\"default\":[\"std\"],\"std\":[]}}", "rust-embed-impl_8.11.0": "{\"dependencies\":[{\"name\":\"proc-macro2\",\"req\":\"^1\"},{\"name\":\"quote\",\"req\":\"^1\"},{\"name\":\"rust-embed-utils\",\"req\":\"^8.11.0\"},{\"name\":\"shellexpand\",\"optional\":true,\"req\":\"^3\"},{\"default_features\":false,\"features\":[\"derive\",\"parsing\",\"proc-macro\",\"printing\"],\"name\":\"syn\",\"req\":\"^2\"},{\"name\":\"walkdir\",\"req\":\"^2.3.1\"}],\"features\":{\"compression\":[],\"debug-embed\":[],\"deterministic-timestamps\":[],\"include-exclude\":[\"rust-embed-utils/include-exclude\"],\"interpolate-folder-path\":[\"shellexpand\"],\"mime-guess\":[\"rust-embed-utils/mime-guess\"]}}", diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 9e7ed84b15..08f14210b7 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -449,7 +449,7 @@ dependencies = [ "objc2-foundation", "parking_lot", "percent-encoding", - "windows-sys 0.59.0", + "windows-sys 0.60.2", "wl-clipboard-rs", "x11rb", ] @@ -1363,7 +1363,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.1", ] [[package]] @@ -1640,6 +1640,17 @@ dependencies = [ "serde_with", ] +[[package]] +name = "codex-bridge" +version = "0.0.0" +dependencies = [ + "async-trait", + "rmp-serde", + "serde", + "thiserror 2.0.18", + "tokio", +] + [[package]] name = "codex-chatgpt" version = "0.0.0" @@ -2820,6 +2831,7 @@ version = "0.0.0" dependencies = [ "async-trait", "chrono", + "codex-bridge", "codex-git-utils", "codex-protocol", "codex-rollout", @@ -4110,7 +4122,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4355,7 +4367,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4916,7 +4928,7 @@ dependencies = [ "gobject-sys", "libc", "system-deps", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5360,7 +5372,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.2", "system-configuration", "tokio", "tower-service", @@ -5446,7 +5458,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.58.0", + "windows-core 0.62.2", ] [[package]] @@ -5866,7 +5878,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6712,7 +6724,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -7356,7 +7368,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7d8fae84b431384b68627d0f9b3b1245fcf9f46f6c0e3dc902e9dce64edd1967" dependencies = [ "libc", - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -8015,7 +8027,7 @@ dependencies = [ "quinn-udp", "rustc-hash 2.1.1", "rustls", - "socket2 0.5.10", + "socket2 0.6.2", "thiserror 2.0.18", "tokio", "tracing", @@ -8052,9 +8064,9 @@ dependencies = [ "cfg_aliases 0.2.1", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.2", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -8762,6 +8774,25 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "rmp" +version = "0.8.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ba8be72d372b2c9b35542551678538b562e7cf86c3315773cae48dfbfe7790c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "rmp-serde" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f81bee8c8ef9b577d1681a70ebbc962c232461e397b22c208c43c04b67a155" +dependencies = [ + "rmp", + "serde", +] + [[package]] name = "rsa" version = "0.9.10" @@ -8896,7 +8927,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.11.0", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -10355,7 +10386,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.1.3", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -11836,7 +11867,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index f82c3cce13..e2d5995b1e 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -14,6 +14,7 @@ members = [ "feedback", "features", "codex-backend-openapi-models", + "bridge", "code-mode", "cloud-requirements", "cloud-tasks", @@ -116,6 +117,7 @@ codex-arg0 = { path = "arg0" } codex-async-utils = { path = "async-utils" } codex-backend-client = { path = "backend-client" } codex-chatgpt = { path = "chatgpt" } +codex-bridge = { path = "bridge" } codex-cli = { path = "cli" } codex-client = { path = "codex-client" } codex-collaboration-mode-templates = { path = "collaboration-mode-templates" } @@ -279,6 +281,7 @@ ratatui-macros = "0.6.0" regex = "1.12.3" regex-lite = "0.1.8" reqwest = "0.12" +rmp-serde = "1.3.0" rmcp = { version = "0.15.0", default-features = false } runfiles = { git = "https://github.com/dzbarsky/rules_rust", rev = "b56cbaa8465e74127f1ea216f813cd377295ad81" } rustls = { version = "0.23", default-features = false, features = [ diff --git a/codex-rs/bridge/BUILD.bazel b/codex-rs/bridge/BUILD.bazel new file mode 100644 index 0000000000..bd397af7a4 --- /dev/null +++ b/codex-rs/bridge/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "bridge", + crate_name = "codex_bridge", +) diff --git a/codex-rs/bridge/Cargo.toml b/codex-rs/bridge/Cargo.toml new file mode 100644 index 0000000000..edfac88bfa --- /dev/null +++ b/codex-rs/bridge/Cargo.toml @@ -0,0 +1,19 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-bridge" +version.workspace = true + +[lib] +name = "codex_bridge" +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +async-trait = { workspace = true } +rmp-serde = { workspace = true } +serde = { workspace = true, features = ["derive"] } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["io-util", "macros", "net", "rt"] } diff --git a/codex-rs/bridge/src/client.rs b/codex-rs/bridge/src/client.rs new file mode 100644 index 0000000000..313032910c --- /dev/null +++ b/codex-rs/bridge/src/client.rs @@ -0,0 +1,170 @@ +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use serde::Serialize; +use serde::de::DeserializeOwned; + +use crate::BridgeEnvelope; +use crate::BridgeError; +use crate::BridgeFrame; +use crate::BridgeRequest; +use crate::BridgeResult; +use crate::BridgeTransport; +use crate::OpaqueFrame; + +/// Typed response returned by [`BridgeClient`]. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BridgeResponse { + /// Decoded response body. + pub body: T, + /// Opaque response frames associated with body fields. + pub opaque_frames: Vec, +} + +/// Reusable MsgPack bridge client over a caller-provided transport. +pub struct BridgeClient { + transport: T, + next_request_id: AtomicU64, +} + +impl BridgeClient { + /// Create a bridge client over the supplied transport. + pub fn new(transport: T) -> Self { + Self { + transport, + next_request_id: AtomicU64::new(1), + } + } +} + +impl BridgeClient +where + T: BridgeTransport, +{ + /// Call a bridge method using MsgPack for the typed body and raw frames for opaque fields. + pub async fn call( + &self, + method: &'static str, + request: BridgeRequest, + ) -> BridgeResult> + where + Req: Serialize + Send + Sync, + Resp: DeserializeOwned, + { + let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed); + let body = rmp_serde::to_vec_named(&request.body).map_err(|err| BridgeError::Codec { + message: err.to_string(), + })?; + let frame = BridgeFrame { + request_id, + method: method.to_string(), + body, + opaque_frames: request.opaque_frames, + }; + let response = self.transport.call(frame).await?; + decode_response(request_id, method, response) + } +} + +fn decode_response( + request_id: u64, + method: &'static str, + response: BridgeEnvelope, +) -> BridgeResult> +where + Resp: DeserializeOwned, +{ + if response.request_id != request_id { + return Err(BridgeError::InvalidResponse { + message: format!( + "expected request id {request_id}, received {}", + response.request_id + ), + }); + } + if response.method != method { + return Err(BridgeError::InvalidResponse { + message: format!("expected method `{method}`, received `{}`", response.method), + }); + } + if let Some(code) = response.error_code { + return Err(BridgeError::Remote { + method: method.to_string(), + code, + message: response.error_message.unwrap_or_default(), + }); + } + let Some(body) = response.body else { + return Err(BridgeError::InvalidResponse { + message: format!("method `{method}` returned no response body"), + }); + }; + let body = rmp_serde::from_slice(body.as_slice()).map_err(|err| BridgeError::Codec { + message: err.to_string(), + })?; + Ok(BridgeResponse { + body, + opaque_frames: response.opaque_frames, + }) +} + +#[cfg(test)] +mod tests { + use async_trait::async_trait; + use serde::Deserialize; + use serde::Serialize; + + use super::*; + + #[derive(Clone)] + struct EchoTransport; + + #[async_trait] + impl BridgeTransport for EchoTransport { + async fn call(&self, request: BridgeFrame) -> BridgeResult { + Ok(BridgeEnvelope { + request_id: request.request_id, + method: request.method, + body: Some(request.body), + opaque_frames: request.opaque_frames, + error_code: None, + error_message: None, + }) + } + } + + #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)] + struct Payload { + value: String, + } + + #[tokio::test] + async fn call_round_trips_msgpack_body_and_opaque_frames() -> BridgeResult<()> { + let client = BridgeClient::new(EchoTransport); + let response: BridgeResponse = client + .call( + "echo", + BridgeRequest::with_opaque_frames( + Payload { + value: "hello".to_string(), + }, + vec![OpaqueFrame { + field: "blob".to_string(), + codec: "raw".to_string(), + content_type: "application/octet-stream".to_string(), + bytes: b"world".to_vec(), + }], + ), + ) + .await?; + + assert_eq!( + response.body, + Payload { + value: "hello".to_string() + } + ); + assert_eq!(response.opaque_frames[0].bytes, b"world"); + Ok(()) + } +} diff --git a/codex-rs/bridge/src/error.rs b/codex-rs/bridge/src/error.rs new file mode 100644 index 0000000000..220655f727 --- /dev/null +++ b/codex-rs/bridge/src/error.rs @@ -0,0 +1,38 @@ +/// Result type returned by bridge operations. +pub type BridgeResult = Result; + +/// Errors surfaced by the local Rust-to-Python bridge. +#[derive(Debug, thiserror::Error)] +pub enum BridgeError { + /// MsgPack serialization or deserialization failed. + #[error("bridge messagepack error: {message}")] + Codec { + /// Codec error detail. + message: String, + }, + + /// The transport failed while reading or writing frames. + #[error("bridge transport error: {message}")] + Transport { + /// Transport error detail. + message: String, + }, + + /// The remote service returned an application error. + #[error("remote bridge method `{method}` failed: {message}")] + Remote { + /// Method that failed. + method: String, + /// Remote error code. + code: String, + /// Remote error detail. + message: String, + }, + + /// The remote response did not match the requested operation. + #[error("invalid bridge response: {message}")] + InvalidResponse { + /// Validation error detail. + message: String, + }, +} diff --git a/codex-rs/bridge/src/frame.rs b/codex-rs/bridge/src/frame.rs new file mode 100644 index 0000000000..e4bdc3282f --- /dev/null +++ b/codex-rs/bridge/src/frame.rs @@ -0,0 +1,75 @@ +use serde::Deserialize; +use serde::Serialize; + +/// A raw opaque byte frame associated with one field in a typed bridge message. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct OpaqueFrame { + /// Field path inside the typed request or response. + pub field: String, + /// Stable codec label used by Rust to encode/decode the bytes. + pub codec: String, + /// MIME-ish content type for diagnostics and non-Rust storage services. + pub content_type: String, + /// Raw bytes that Python should not deserialize as part of normal persistence. + pub bytes: Vec, +} + +/// A typed bridge request plus any opaque byte frames referenced by that request. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BridgeRequest { + /// Request body encoded in MsgPack inside the envelope. + pub body: T, + /// Opaque frames associated with body fields. + pub opaque_frames: Vec, +} + +impl BridgeRequest { + /// Create a request without opaque frames. + pub fn new(body: T) -> Self { + Self { + body, + opaque_frames: Vec::new(), + } + } + + /// Create a request with opaque frames. + pub fn with_opaque_frames(body: T, opaque_frames: Vec) -> Self { + Self { + body, + opaque_frames, + } + } +} + +/// A complete bridge frame as it is written to the transport. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BridgeFrame { + /// Request id chosen by the client. + pub request_id: u64, + /// Stable method name. + pub method: String, + /// MsgPack-encoded typed body. + pub body: Vec, + /// Opaque byte frames for large fields. + pub opaque_frames: Vec, +} + +/// Response envelope returned by the bridge service. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BridgeEnvelope { + /// Request id this response corresponds to. + pub request_id: u64, + /// Stable method name. + pub method: String, + /// MsgPack-encoded typed success body, if the call succeeded. + pub body: Option>, + /// Opaque byte frames for large response fields. + pub opaque_frames: Vec, + /// Remote error code, if the call failed. + pub error_code: Option, + /// Remote error message, if the call failed. + pub error_message: Option, +} diff --git a/codex-rs/bridge/src/lib.rs b/codex-rs/bridge/src/lib.rs new file mode 100644 index 0000000000..00299ed9f7 --- /dev/null +++ b/codex-rs/bridge/src/lib.rs @@ -0,0 +1,32 @@ +//! Reusable local bridge primitives for typed Rust-to-Python calls. +//! +//! The bridge intentionally separates a small MsgPack-encoded control message from large opaque +//! byte frames. Rust callers can keep strongly typed DTOs while Python services can persist or +//! forward opaque fields without deserializing them. + +mod client; +mod error; +mod frame; +mod opaque; +mod schema; +mod transport; + +pub use client::BridgeClient; +pub use client::BridgeResponse; +pub use error::BridgeError; +pub use error::BridgeResult; +pub use frame::BridgeEnvelope; +pub use frame::BridgeFrame; +pub use frame::BridgeRequest; +pub use frame::OpaqueFrame; +pub use opaque::decode_opaque_msgpack; +pub use opaque::encode_opaque_msgpack; +pub use schema::BridgeField; +pub use schema::BridgeMethod; +pub use schema::BridgeSchema; +pub use schema::BridgeType; +pub use schema::OpaqueField; +pub use transport::BridgeTransport; + +#[cfg(unix)] +pub use transport::UnixSocketBridgeTransport; diff --git a/codex-rs/bridge/src/opaque.rs b/codex-rs/bridge/src/opaque.rs new file mode 100644 index 0000000000..540d7e61d2 --- /dev/null +++ b/codex-rs/bridge/src/opaque.rs @@ -0,0 +1,47 @@ +use serde::Serialize; +use serde::de::DeserializeOwned; + +use crate::BridgeError; +use crate::BridgeResult; +use crate::OpaqueFrame; + +const MSGPACK_CONTENT_TYPE: &str = "application/vnd.codex.msgpack"; + +/// Encode a typed Rust value into an opaque MsgPack frame for Python to treat as bytes. +pub fn encode_opaque_msgpack(field: &str, codec: &str, value: &T) -> BridgeResult +where + T: Serialize, +{ + let bytes = rmp_serde::to_vec_named(value).map_err(|err| BridgeError::Codec { + message: format!("failed to encode opaque `{field}` payload: {err}"), + })?; + Ok(OpaqueFrame { + field: field.to_string(), + codec: codec.to_string(), + content_type: MSGPACK_CONTENT_TYPE.to_string(), + bytes, + }) +} + +/// Decode a typed Rust value from an opaque MsgPack frame returned by Python. +pub fn decode_opaque_msgpack(frames: &[OpaqueFrame], field: &str, codec: &str) -> BridgeResult +where + T: DeserializeOwned, +{ + let Some(frame) = frames.iter().find(|frame| frame.field == field) else { + return Err(BridgeError::InvalidResponse { + message: format!("bridge response did not include `{field}` opaque frame"), + }); + }; + if frame.codec != codec { + return Err(BridgeError::InvalidResponse { + message: format!( + "bridge response used unsupported `{field}` codec `{}`; expected `{codec}`", + frame.codec + ), + }); + } + rmp_serde::from_slice(frame.bytes.as_slice()).map_err(|err| BridgeError::Codec { + message: format!("failed to decode opaque `{field}` payload: {err}"), + }) +} diff --git a/codex-rs/bridge/src/schema.rs b/codex-rs/bridge/src/schema.rs new file mode 100644 index 0000000000..c51575c87b --- /dev/null +++ b/codex-rs/bridge/src/schema.rs @@ -0,0 +1,66 @@ +use serde::Deserialize; +use serde::Serialize; + +/// Language-neutral schema manifest generated from Rust bridge DTO declarations. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BridgeSchema { + /// Schema namespace. + pub namespace: String, + /// Schema version incremented when bridge DTOs change incompatibly. + pub version: u32, + /// Types exported to Python. + pub types: Vec, + /// Methods exported to Python service implementations. + pub methods: Vec, +} + +/// Exported bridge struct or enum. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BridgeType { + /// Type name. + pub name: String, + /// Either `struct` or `enum`. + pub kind: String, + /// Struct fields. + pub fields: Vec, + /// Enum variant names. + pub variants: Vec, +} + +/// Exported bridge struct field. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BridgeField { + /// Field name. + pub name: String, + /// Python-facing type expression. + pub python_type: String, + /// Whether the field is optional. + pub optional: bool, + /// Opaque field descriptor when Python should see bytes instead of the Rust semantic type. + pub opaque: Option, +} + +/// Metadata for a field carried as opaque bytes in Python. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct OpaqueField { + /// Stable codec label. + pub codec: String, + /// Rust type hidden behind the opaque Python bytes. + pub rust_type: String, +} + +/// Exported bridge method. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BridgeMethod { + /// Stable method name. + pub name: String, + /// Request type name. + pub request: String, + /// Response type name. + pub response: String, +} diff --git a/codex-rs/bridge/src/transport.rs b/codex-rs/bridge/src/transport.rs new file mode 100644 index 0000000000..1d35c967ff --- /dev/null +++ b/codex-rs/bridge/src/transport.rs @@ -0,0 +1,104 @@ +use async_trait::async_trait; + +use crate::BridgeEnvelope; +use crate::BridgeError; +use crate::BridgeFrame; +use crate::BridgeResult; + +/// Transport used by [`crate::BridgeClient`] to send one complete bridge frame. +#[async_trait] +pub trait BridgeTransport: Send + Sync { + /// Sends a request frame and returns the matching response envelope. + async fn call(&self, request: BridgeFrame) -> BridgeResult; +} + +#[cfg(unix)] +mod unix { + use std::path::PathBuf; + + use tokio::io::AsyncReadExt; + use tokio::io::AsyncWriteExt; + use tokio::net::UnixStream; + + use super::*; + + const MAX_FRAME_BYTES: usize = 128 * 1024 * 1024; + + /// Unix domain socket bridge transport. + #[derive(Clone, Debug)] + pub struct UnixSocketBridgeTransport { + path: PathBuf, + } + + impl UnixSocketBridgeTransport { + /// Create a transport that opens a fresh Unix socket connection for each call. + pub fn new(path: PathBuf) -> Self { + Self { path } + } + } + + #[async_trait] + impl BridgeTransport for UnixSocketBridgeTransport { + async fn call(&self, request: BridgeFrame) -> BridgeResult { + let mut stream = UnixStream::connect(self.path.as_path()) + .await + .map_err(transport_error)?; + write_msgpack(&mut stream, &request).await?; + let response = read_msgpack(&mut stream).await?; + Ok(response) + } + } + + async fn write_msgpack(stream: &mut UnixStream, value: &T) -> BridgeResult<()> + where + T: serde::Serialize, + { + let bytes = rmp_serde::to_vec_named(value).map_err(|err| BridgeError::Codec { + message: err.to_string(), + })?; + let len = u32::try_from(bytes.len()).map_err(|_| BridgeError::InvalidResponse { + message: format!("bridge request frame is too large: {} bytes", bytes.len()), + })?; + stream + .write_all(&len.to_be_bytes()) + .await + .map_err(transport_error)?; + stream.write_all(&bytes).await.map_err(transport_error)?; + stream.flush().await.map_err(transport_error)?; + Ok(()) + } + + async fn read_msgpack(stream: &mut UnixStream) -> BridgeResult + where + T: serde::de::DeserializeOwned, + { + let mut len_bytes = [0u8; 4]; + stream + .read_exact(&mut len_bytes) + .await + .map_err(transport_error)?; + let len = u32::from_be_bytes(len_bytes) as usize; + if len > MAX_FRAME_BYTES { + return Err(BridgeError::InvalidResponse { + message: format!("bridge response frame is too large: {len} bytes"), + }); + } + let mut bytes = vec![0u8; len]; + stream + .read_exact(&mut bytes) + .await + .map_err(transport_error)?; + rmp_serde::from_slice(bytes.as_slice()).map_err(|err| BridgeError::Codec { + message: err.to_string(), + }) + } + + fn transport_error(err: std::io::Error) -> BridgeError { + BridgeError::Transport { + message: err.to_string(), + } + } +} + +#[cfg(unix)] +pub use unix::UnixSocketBridgeTransport; diff --git a/codex-rs/thread-store/Cargo.toml b/codex-rs/thread-store/Cargo.toml index 6c1acee3cb..2fdebc2863 100644 --- a/codex-rs/thread-store/Cargo.toml +++ b/codex-rs/thread-store/Cargo.toml @@ -14,6 +14,7 @@ workspace = true [dependencies] async-trait = { workspace = true } chrono = { workspace = true, features = ["serde"] } +codex-bridge = { workspace = true } codex-git-utils = { workspace = true } codex-protocol = { workspace = true } codex-rollout = { workspace = true } diff --git a/codex-rs/thread-store/python/codex_thread_store_bridge.py b/codex-rs/thread-store/python/codex_thread_store_bridge.py new file mode 100644 index 0000000000..2078d72561 --- /dev/null +++ b/codex-rs/thread-store/python/codex_thread_store_bridge.py @@ -0,0 +1,225 @@ +# @generated by scripts/generate_thread_store_bridge_python.py +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Protocol + + +@dataclass(frozen=True, kw_only=True) +class OpaqueBytes: + codec: str + contentType: str + data: bytes + +@dataclass(frozen=True, kw_only=True) +class ThreadOwner: + chatgptUserId: str | None = None + accountId: str | None = None + authMode: str | None = None + +@dataclass(frozen=True, kw_only=True) +class RemoteThreadMetadata: + threadId: str + forkedFromId: str | None = None + owner: ThreadOwner + preview: str + name: str | None = None + modelProvider: str + model: str | None = None + serviceTier: Any | None = None + reasoningEffort: Any | None = None + createdAt: int + updatedAt: int + archivedAt: int | None = None + cwd: str + cliVersion: str + source: Any + agentNickname: str | None = None + agentRole: str | None = None + agentPath: str | None = None + gitInfo: Any | None = None + approvalMode: Any + sandboxPolicy: Any + tokenUsage: Any | None = None + firstUserMessage: str | None = None + memoryMode: str | None = None + +@dataclass(frozen=True, kw_only=True) +class ThreadIndexPatch: + updatedAt: int + firstUserMessage: str | None = None + preview: str | None = None + name: str | None = None + tokenUsage: Any | None = None + memoryMode: str | None = None + +@dataclass(frozen=True, kw_only=True) +class CreateThreadRequest: + thread: RemoteThreadMetadata + initialItems: OpaqueBytes + initialPayloadCodec: str + eventPersistenceMode: str + +@dataclass(frozen=True, kw_only=True) +class ResumeThreadRecorderRequest: + threadId: str + owner: ThreadOwner + includeArchived: bool + eventPersistenceMode: str + +@dataclass(frozen=True, kw_only=True) +class AppendThreadRequest: + threadId: str + owner: ThreadOwner + idempotencyKey: str | None = None + updatedAt: int | None = None + newThreadMemoryMode: str | None = None + eventPersistenceMode: str | None = None + indexPatch: ThreadIndexPatch + items: OpaqueBytes + payloadCodec: str + +@dataclass(frozen=True, kw_only=True) +class LoadThreadHistoryRequest: + threadId: str + owner: ThreadOwner + includeArchived: bool + +@dataclass(frozen=True, kw_only=True) +class LoadThreadHistoryResponse: + threadId: str + payloadCodec: str + history: OpaqueBytes + +@dataclass(frozen=True, kw_only=True) +class ReadThreadRequest: + threadId: str + owner: ThreadOwner + includeArchived: bool + includeHistory: bool + +@dataclass(frozen=True, kw_only=True) +class ReadThreadResponse: + thread: RemoteThreadMetadata + historyPayloadCodec: str | None = None + history: OpaqueBytes | None = None + +@dataclass(frozen=True, kw_only=True) +class ListThreadsRequest: + owner: ThreadOwner + pageSize: int + cursor: str | None = None + sortKey: str + allowedSources: list[Any] + modelProviders: list[str] | None = None + archived: bool + cwd: str | None = None + searchTerm: str | None = None + +@dataclass(frozen=True, kw_only=True) +class ListThreadsResponse: + items: list[RemoteThreadMetadata] + nextCursor: str | None = None + scanned: int | None = None + +@dataclass(frozen=True, kw_only=True) +class FindThreadByNameRequest: + owner: ThreadOwner + name: str + includeArchived: bool + cwd: str | None = None + allowedSources: list[Any] + modelProviders: list[str] | None = None + +@dataclass(frozen=True, kw_only=True) +class FindThreadByNameResponse: + thread: RemoteThreadMetadata | None = None + +@dataclass(frozen=True, kw_only=True) +class SetThreadNameRequest: + threadId: str + owner: ThreadOwner + name: str + +@dataclass(frozen=True, kw_only=True) +class UpdateThreadMetadataRequest: + threadId: str + owner: ThreadOwner + patch: Any + +@dataclass(frozen=True, kw_only=True) +class ArchiveThreadRequest: + threadId: str + owner: ThreadOwner + +@dataclass(frozen=True, kw_only=True) +class DynamicToolsRequest: + threadId: str + owner: ThreadOwner + +@dataclass(frozen=True, kw_only=True) +class DynamicToolsResponse: + dynamicTools: list[Any] | None = None + +@dataclass(frozen=True, kw_only=True) +class MemoryModeRequest: + threadId: str + owner: ThreadOwner + +@dataclass(frozen=True, kw_only=True) +class MemoryModeResponse: + memoryMode: str | None = None + +@dataclass(frozen=True, kw_only=True) +class SetMemoryModeRequest: + threadId: str + owner: ThreadOwner + memoryMode: str + +@dataclass(frozen=True, kw_only=True) +class ThreadSpawnEdgeRecord: + parentThreadId: str + childThreadId: str + status: str + +@dataclass(frozen=True, kw_only=True) +class ListThreadSpawnEdgesRequest: + threadId: str + owner: ThreadOwner + recursive: bool + status: str | None = None + +@dataclass(frozen=True, kw_only=True) +class ListThreadSpawnEdgesResponse: + edges: list[ThreadSpawnEdgeRecord] + +@dataclass(frozen=True, kw_only=True) +class FindThreadSpawnByPathRequest: + threadId: str + owner: ThreadOwner + recursive: bool + agentPath: str + +@dataclass(frozen=True, kw_only=True) +class FindThreadSpawnByPathResponse: + threadId: str | None = None + +class ThreadStoreBridgeService(Protocol): + async def create_thread(self, request: CreateThreadRequest) -> None: ... + async def resume_thread_recorder(self, request: ResumeThreadRecorderRequest) -> None: ... + async def append_thread_items(self, request: AppendThreadRequest) -> None: ... + async def load_thread_history(self, request: LoadThreadHistoryRequest) -> LoadThreadHistoryResponse: ... + async def read_thread(self, request: ReadThreadRequest) -> ReadThreadResponse: ... + async def list_threads(self, request: ListThreadsRequest) -> ListThreadsResponse: ... + async def find_thread_by_name(self, request: FindThreadByNameRequest) -> FindThreadByNameResponse: ... + async def set_thread_name(self, request: SetThreadNameRequest) -> None: ... + async def update_thread_metadata(self, request: UpdateThreadMetadataRequest) -> ReadThreadResponse: ... + async def archive_thread(self, request: ArchiveThreadRequest) -> None: ... + async def unarchive_thread(self, request: ArchiveThreadRequest) -> ReadThreadResponse: ... + async def dynamic_tools(self, request: DynamicToolsRequest) -> DynamicToolsResponse: ... + async def memory_mode(self, request: MemoryModeRequest) -> MemoryModeResponse: ... + async def set_memory_mode(self, request: SetMemoryModeRequest) -> None: ... + async def mark_memory_mode_polluted(self, request: MemoryModeRequest) -> None: ... + async def upsert_thread_spawn_edge(self, request: ThreadSpawnEdgeRecord) -> None: ... + async def list_thread_spawn_edges(self, request: ListThreadSpawnEdgesRequest) -> ListThreadSpawnEdgesResponse: ... + async def find_thread_spawn_by_path(self, request: FindThreadSpawnByPathRequest) -> FindThreadSpawnByPathResponse: ... diff --git a/codex-rs/thread-store/src/bin/codex-thread-store-bridge-schema.rs b/codex-rs/thread-store/src/bin/codex-thread-store-bridge-schema.rs new file mode 100644 index 0000000000..78b15d143b --- /dev/null +++ b/codex-rs/thread-store/src/bin/codex-thread-store-bridge-schema.rs @@ -0,0 +1,7 @@ +use codex_thread_store::thread_store_bridge_schema; + +fn main() -> Result<(), Box> { + let schema = thread_store_bridge_schema(); + println!("{}", serde_json::to_string_pretty(&schema)?); + Ok(()) +} diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 8cf93a7f62..fb82c675aa 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -7,6 +7,7 @@ mod error; mod local; mod recorder; +mod remote; mod store; mod types; @@ -14,6 +15,8 @@ pub use error::ThreadStoreError; pub use error::ThreadStoreResult; pub use local::LocalThreadStore; pub use recorder::ThreadRecorder; +pub use remote::RemoteThreadStore; +pub use remote::thread_store_bridge_schema; pub use store::ThreadStore; pub use types::AppendThreadItemsParams; pub use types::ArchiveThreadParams; diff --git a/codex-rs/thread-store/src/remote/mod.rs b/codex-rs/thread-store/src/remote/mod.rs new file mode 100644 index 0000000000..8e1f4f207c --- /dev/null +++ b/codex-rs/thread-store/src/remote/mod.rs @@ -0,0 +1,515 @@ +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::Utc; +use codex_bridge::BridgeClient; +use codex_bridge::BridgeRequest; +use codex_bridge::BridgeResponse; +use codex_bridge::BridgeTransport; +use codex_bridge::OpaqueFrame; +#[cfg(unix)] +use codex_bridge::UnixSocketBridgeTransport; +use codex_bridge::decode_opaque_msgpack; +use codex_bridge::encode_opaque_msgpack; +use codex_protocol::ThreadId; +use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; + +mod recorder; +mod schema; +mod types; + +use recorder::RemoteThreadRecorder; +pub use schema::thread_store_bridge_schema; +use types::AppendThreadRequest; +use types::ArchiveThreadRequest; +use types::CreateThreadRequest; +use types::DynamicToolsRequest; +use types::DynamicToolsResponse; +use types::FindThreadByNameRequest; +use types::FindThreadByNameResponse; +use types::FindThreadSpawnByPathRequest; +use types::FindThreadSpawnByPathResponse; +use types::ListThreadSpawnEdgesRequest; +use types::ListThreadSpawnEdgesResponse; +use types::ListThreadsRequest; +use types::ListThreadsResponse; +use types::LoadThreadHistoryRequest; +use types::LoadThreadHistoryResponse; +use types::MemoryModeRequest; +use types::MemoryModeResponse; +use types::ReadThreadRequest; +use types::ReadThreadResponse; +use types::ResumeThreadRecorderRequest; +use types::SetMemoryModeRequest; +use types::SetThreadNameRequest; +use types::StoredThreadPayload; +use types::ThreadSpawnEdgeRecord; +use types::UpdateThreadMetadataRequest; + +use crate::AppendThreadItemsParams; +use crate::ArchiveThreadParams; +use crate::CreateThreadParams; +use crate::DynamicToolsParams; +use crate::FindThreadByNameParams; +use crate::FindThreadSpawnByPathParams; +use crate::ListThreadSpawnEdgesParams; +use crate::ListThreadsParams; +use crate::LoadThreadHistoryParams; +use crate::ReadThreadParams; +use crate::ResolveLegacyPathParams; +use crate::ResumeThreadRecorderParams; +use crate::SetThreadMemoryModeParams; +use crate::SetThreadNameParams; +use crate::StoredThread; +use crate::StoredThreadHistory; +use crate::ThreadMemoryModeParams; +use crate::ThreadPage; +use crate::ThreadRecorder; +use crate::ThreadSpawnEdge; +use crate::ThreadStore; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; +use crate::UpdateThreadMetadataParams; + +const CREATE_THREAD: &str = "thread_store/create_thread"; +const RESUME_THREAD_RECORDER: &str = "thread_store/resume_thread_recorder"; +const APPEND_THREAD_ITEMS: &str = "thread_store/append_thread_items"; +const LOAD_THREAD_HISTORY: &str = "thread_store/load_thread_history"; +const READ_THREAD: &str = "thread_store/read_thread"; +const LIST_THREADS: &str = "thread_store/list_threads"; +const FIND_THREAD_BY_NAME: &str = "thread_store/find_thread_by_name"; +const SET_THREAD_NAME: &str = "thread_store/set_thread_name"; +const UPDATE_THREAD_METADATA: &str = "thread_store/update_thread_metadata"; +const ARCHIVE_THREAD: &str = "thread_store/archive_thread"; +const UNARCHIVE_THREAD: &str = "thread_store/unarchive_thread"; +const DYNAMIC_TOOLS: &str = "thread_store/dynamic_tools"; +const MEMORY_MODE: &str = "thread_store/memory_mode"; +const SET_MEMORY_MODE: &str = "thread_store/set_memory_mode"; +const MARK_MEMORY_MODE_POLLUTED: &str = "thread_store/mark_memory_mode_polluted"; +const UPSERT_THREAD_SPAWN_EDGE: &str = "thread_store/upsert_thread_spawn_edge"; +const LIST_THREAD_SPAWN_EDGES: &str = "thread_store/list_thread_spawn_edges"; +const FIND_THREAD_SPAWN_BY_PATH: &str = "thread_store/find_thread_spawn_by_path"; + +const ROLLOUT_ITEMS_CODEC: &str = "codex.rollout_items.msgpack.v1"; +const STORED_THREAD_HISTORY_CODEC: &str = "codex.stored_thread_history.msgpack.v1"; + +/// Remote [`ThreadStore`] implementation backed by a Python bridge service. +/// +/// This adapter keeps cloud-specific storage details behind the bridge. Python receives typed, +/// bounded metadata DTOs plus opaque MsgPack payload frames for rollout history. +pub struct RemoteThreadStore { + client: Arc>, +} + +impl Clone for RemoteThreadStore { + fn clone(&self) -> Self { + Self { + client: Arc::clone(&self.client), + } + } +} + +impl RemoteThreadStore { + /// Create a remote store from a bridge transport. + pub fn new(transport: T) -> Self { + Self { + client: Arc::new(BridgeClient::new(transport)), + } + } +} + +#[cfg(unix)] +impl RemoteThreadStore { + /// Create a remote store that talks to a Python bridge over a Unix domain socket. + pub fn from_unix_socket(path: PathBuf) -> Self { + Self::new(UnixSocketBridgeTransport::new(path)) + } +} + +impl RemoteThreadStore +where + T: BridgeTransport, +{ + async fn call( + &self, + method: &'static str, + request: BridgeRequest, + ) -> ThreadStoreResult> + where + Req: serde::Serialize + Send + Sync, + Resp: serde::de::DeserializeOwned, + { + self.client + .call(method, request) + .await + .map_err(remote_error) + } + + async fn call_empty(&self, method: &'static str, request: Req) -> ThreadStoreResult<()> + where + Req: serde::Serialize + Send + Sync, + { + let _: BridgeResponse<()> = self.call(method, BridgeRequest::new(request)).await?; + Ok(()) + } + + async fn append_items_inner( + &self, + params: AppendThreadItemsParams, + event_persistence_mode: Option, + ) -> ThreadStoreResult<()> { + let items_frame = encode_opaque_msgpack("items", ROLLOUT_ITEMS_CODEC, ¶ms.items) + .map_err(remote_error)?; + let request = AppendThreadRequest::from_params(params, event_persistence_mode); + let _: BridgeResponse<()> = self + .call( + APPEND_THREAD_ITEMS, + BridgeRequest::with_opaque_frames(request, vec![items_frame]), + ) + .await?; + Ok(()) + } +} + +#[async_trait] +impl ThreadStore for RemoteThreadStore +where + T: BridgeTransport + 'static, +{ + async fn create_thread( + &self, + params: CreateThreadParams, + ) -> ThreadStoreResult> { + let initial_items = initial_rollout_items(¶ms); + let initial_items_frame = + encode_opaque_msgpack("initialItems", ROLLOUT_ITEMS_CODEC, &initial_items) + .map_err(remote_error)?; + let thread_id = params.thread_id; + let owner = params.owner.clone(); + let event_persistence_mode = types::event_persistence_mode(params.event_persistence_mode); + let request = CreateThreadRequest::from_params(params); + let _: BridgeResponse<()> = self + .call( + CREATE_THREAD, + BridgeRequest::with_opaque_frames(request, vec![initial_items_frame]), + ) + .await?; + Ok(Box::new(RemoteThreadRecorder::new( + thread_id, + owner, + event_persistence_mode, + self.clone(), + ))) + } + + async fn resume_thread_recorder( + &self, + params: ResumeThreadRecorderParams, + ) -> ThreadStoreResult> { + let thread_id = params.thread_id; + let owner = params.owner.clone(); + let event_persistence_mode = types::event_persistence_mode(params.event_persistence_mode); + let _: BridgeResponse<()> = self + .call( + RESUME_THREAD_RECORDER, + BridgeRequest::new(ResumeThreadRecorderRequest::from_params(params)), + ) + .await?; + Ok(Box::new(RemoteThreadRecorder::new( + thread_id, + owner, + event_persistence_mode, + self.clone(), + ))) + } + + async fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreResult<()> { + self.append_items_inner(params, None).await + } + + async fn load_history( + &self, + params: LoadThreadHistoryParams, + ) -> ThreadStoreResult { + let response: BridgeResponse = self + .call( + LOAD_THREAD_HISTORY, + BridgeRequest::new(LoadThreadHistoryRequest::from_params(params)), + ) + .await?; + decode_history_from_response(response) + } + + async fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreResult { + let include_history = params.include_history; + let response: BridgeResponse = self + .call( + READ_THREAD, + BridgeRequest::new(ReadThreadRequest::from_params(params)), + ) + .await?; + let history = if include_history { + Some(decode_history_frame(response.opaque_frames.as_slice())?) + } else { + None + }; + Ok(response.body.thread.into_stored_thread(history)?) + } + + async fn list_threads(&self, params: ListThreadsParams) -> ThreadStoreResult { + let response: BridgeResponse = self + .call( + LIST_THREADS, + BridgeRequest::new(ListThreadsRequest::from_params(params)), + ) + .await?; + let mut items = Vec::with_capacity(response.body.items.len()); + for item in response.body.items { + items.push(item.into_stored_thread(None)?); + } + Ok(ThreadPage { + items, + next_cursor: response.body.next_cursor, + scanned: response.body.scanned, + }) + } + + async fn find_thread_by_name( + &self, + params: FindThreadByNameParams, + ) -> ThreadStoreResult> { + let response: BridgeResponse = self + .call( + FIND_THREAD_BY_NAME, + BridgeRequest::new(FindThreadByNameRequest::from_params(params)), + ) + .await?; + response + .body + .thread + .map(|thread| thread.into_stored_thread(None)) + .transpose() + } + + async fn set_thread_name(&self, params: SetThreadNameParams) -> ThreadStoreResult<()> { + self.call_empty(SET_THREAD_NAME, SetThreadNameRequest::from_params(params)) + .await + } + + async fn update_thread_metadata( + &self, + params: UpdateThreadMetadataParams, + ) -> ThreadStoreResult { + let response: BridgeResponse = self + .call( + UPDATE_THREAD_METADATA, + BridgeRequest::new(UpdateThreadMetadataRequest::from_params(params)), + ) + .await?; + response.body.thread.into_stored_thread(None) + } + + async fn archive_thread(&self, params: ArchiveThreadParams) -> ThreadStoreResult<()> { + self.call_empty(ARCHIVE_THREAD, ArchiveThreadRequest::from_params(params)) + .await + } + + async fn unarchive_thread( + &self, + params: ArchiveThreadParams, + ) -> ThreadStoreResult { + let response: BridgeResponse = self + .call( + UNARCHIVE_THREAD, + BridgeRequest::new(ArchiveThreadRequest::from_params(params)), + ) + .await?; + response.body.thread.into_stored_thread(None) + } + + async fn resolve_legacy_path( + &self, + _params: ResolveLegacyPathParams, + ) -> ThreadStoreResult> { + Ok(None) + } + + async fn dynamic_tools( + &self, + params: DynamicToolsParams, + ) -> ThreadStoreResult>> { + let response: BridgeResponse = self + .call( + DYNAMIC_TOOLS, + BridgeRequest::new(DynamicToolsRequest::from_params(params)), + ) + .await?; + Ok(response.body.dynamic_tools) + } + + async fn memory_mode( + &self, + params: ThreadMemoryModeParams, + ) -> ThreadStoreResult> { + let response: BridgeResponse = self + .call( + MEMORY_MODE, + BridgeRequest::new(MemoryModeRequest::from_params(params)), + ) + .await?; + Ok(response.body.memory_mode) + } + + async fn set_memory_mode(&self, params: SetThreadMemoryModeParams) -> ThreadStoreResult<()> { + self.call_empty(SET_MEMORY_MODE, SetMemoryModeRequest::from_params(params)) + .await + } + + async fn mark_memory_mode_polluted( + &self, + params: ThreadMemoryModeParams, + ) -> ThreadStoreResult<()> { + self.call_empty( + MARK_MEMORY_MODE_POLLUTED, + MemoryModeRequest::from_params(params), + ) + .await + } + + async fn upsert_thread_spawn_edge(&self, edge: ThreadSpawnEdge) -> ThreadStoreResult<()> { + self.call_empty( + UPSERT_THREAD_SPAWN_EDGE, + ThreadSpawnEdgeRecord::from_thread_spawn_edge(edge), + ) + .await + } + + async fn list_thread_spawn_edges( + &self, + params: ListThreadSpawnEdgesParams, + ) -> ThreadStoreResult> { + let response: BridgeResponse = self + .call( + LIST_THREAD_SPAWN_EDGES, + BridgeRequest::new(ListThreadSpawnEdgesRequest::from_params(params)), + ) + .await?; + response + .body + .edges + .into_iter() + .map(ThreadSpawnEdgeRecord::into_thread_spawn_edge) + .collect() + } + + async fn find_thread_spawn_by_path( + &self, + params: FindThreadSpawnByPathParams, + ) -> ThreadStoreResult> { + let response: BridgeResponse = self + .call( + FIND_THREAD_SPAWN_BY_PATH, + BridgeRequest::new(FindThreadSpawnByPathRequest::from_params(params)), + ) + .await?; + response.body.thread_id.map(parse_thread_id).transpose() + } + + fn supports_legacy_path(&self, _path: &Path) -> bool { + false + } +} + +pub(crate) async fn append_items_via_bridge( + store: &RemoteThreadStore, + thread_id: ThreadId, + owner: crate::ThreadOwner, + event_persistence_mode: String, + items: &[RolloutItem], +) -> ThreadStoreResult<()> +where + T: BridgeTransport, +{ + store + .append_items_inner( + AppendThreadItemsParams { + thread_id, + owner, + items: items.to_vec(), + idempotency_key: None, + updated_at: None, + new_thread_memory_mode: None, + }, + Some(event_persistence_mode), + ) + .await +} + +fn initial_rollout_items(params: &CreateThreadParams) -> Vec { + vec![RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: params.thread_id, + forked_from_id: params.forked_from_id, + timestamp: Utc::now().to_rfc3339(), + cwd: params.cwd.clone(), + originator: params.originator.clone(), + cli_version: params.cli_version.clone(), + source: params.source.clone(), + agent_nickname: None, + agent_role: None, + agent_path: None, + model_provider: Some(params.model_provider.clone()), + base_instructions: Some(params.base_instructions.clone()), + dynamic_tools: Some(params.dynamic_tools.clone()), + memory_mode: params.memory_mode.clone(), + }, + git: params.git_info.clone(), + })] +} + +fn decode_history_from_response( + response: BridgeResponse, +) -> ThreadStoreResult { + let mut history = decode_history_frame(response.opaque_frames.as_slice())?; + if history.thread_id.to_string() != response.body.thread_id { + history.thread_id = parse_thread_id(response.body.thread_id)?; + } + Ok(history) +} + +fn decode_history_frame(frames: &[OpaqueFrame]) -> ThreadStoreResult { + let payload: StoredThreadPayload = + decode_opaque_msgpack(frames, "history", STORED_THREAD_HISTORY_CODEC) + .map_err(remote_error)?; + payload.into_stored_thread_history() +} + +fn parse_thread_id(thread_id: String) -> ThreadStoreResult { + ThreadId::from_string(thread_id.as_str()).map_err(|err| ThreadStoreError::InvalidRequest { + message: format!("invalid remote thread id `{thread_id}`: {err}"), + }) +} + +fn remote_error(err: codex_bridge::BridgeError) -> ThreadStoreError { + match err { + codex_bridge::BridgeError::Remote { code, message, .. } => match code.as_str() { + "thread_not_found" => ThreadStoreError::InvalidRequest { message }, + "invalid_request" => ThreadStoreError::InvalidRequest { message }, + "conflict" => ThreadStoreError::Conflict { message }, + "unavailable" => ThreadStoreError::Unavailable { message }, + _ => ThreadStoreError::Internal { message }, + }, + codex_bridge::BridgeError::Transport { message } => { + ThreadStoreError::Unavailable { message } + } + codex_bridge::BridgeError::Codec { message } + | codex_bridge::BridgeError::InvalidResponse { message } => { + ThreadStoreError::Internal { message } + } + } +} diff --git a/codex-rs/thread-store/src/remote/recorder.rs b/codex-rs/thread-store/src/remote/recorder.rs new file mode 100644 index 0000000000..b991f0b4b0 --- /dev/null +++ b/codex-rs/thread-store/src/remote/recorder.rs @@ -0,0 +1,69 @@ +use async_trait::async_trait; +use codex_bridge::BridgeTransport; +use codex_protocol::ThreadId; +use codex_protocol::protocol::RolloutItem; + +use super::RemoteThreadStore; +use super::append_items_via_bridge; +use crate::ThreadOwner; +use crate::ThreadRecorder; +use crate::ThreadStoreResult; + +pub(crate) struct RemoteThreadRecorder { + thread_id: ThreadId, + owner: ThreadOwner, + event_persistence_mode: String, + store: RemoteThreadStore, +} + +impl RemoteThreadRecorder { + pub(crate) fn new( + thread_id: ThreadId, + owner: ThreadOwner, + event_persistence_mode: String, + store: RemoteThreadStore, + ) -> Self { + Self { + thread_id, + owner, + event_persistence_mode, + store, + } + } +} + +#[async_trait] +impl ThreadRecorder for RemoteThreadRecorder +where + T: BridgeTransport + 'static, +{ + fn thread_id(&self) -> ThreadId { + self.thread_id + } + + async fn record_items(&self, items: &[RolloutItem]) -> ThreadStoreResult<()> { + if items.is_empty() { + return Ok(()); + } + append_items_via_bridge( + &self.store, + self.thread_id, + self.owner.clone(), + self.event_persistence_mode.clone(), + items, + ) + .await + } + + async fn persist(&self) -> ThreadStoreResult<()> { + Ok(()) + } + + async fn flush(&self) -> ThreadStoreResult<()> { + Ok(()) + } + + async fn shutdown(&self) -> ThreadStoreResult<()> { + Ok(()) + } +} diff --git a/codex-rs/thread-store/src/remote/schema.rs b/codex-rs/thread-store/src/remote/schema.rs new file mode 100644 index 0000000000..31a5886210 --- /dev/null +++ b/codex-rs/thread-store/src/remote/schema.rs @@ -0,0 +1,402 @@ +use codex_bridge::BridgeField; +use codex_bridge::BridgeMethod; +use codex_bridge::BridgeSchema; +use codex_bridge::BridgeType; +use codex_bridge::OpaqueField; + +/// Returns the Rust-authored schema manifest used to generate Python bridge types. +pub fn thread_store_bridge_schema() -> BridgeSchema { + BridgeSchema { + namespace: "codex.thread_store.v1".to_string(), + version: 1, + types: vec![ + struct_type( + "OpaqueBytes", + vec![ + field("codec", "str"), + field("contentType", "str"), + field("data", "bytes"), + ], + ), + struct_type( + "ThreadOwner", + vec![ + optional_field("chatgptUserId", "str"), + optional_field("accountId", "str"), + optional_field("authMode", "str"), + ], + ), + struct_type( + "RemoteThreadMetadata", + vec![ + field("threadId", "str"), + optional_field("forkedFromId", "str"), + field("owner", "ThreadOwner"), + field("preview", "str"), + optional_field("name", "str"), + field("modelProvider", "str"), + optional_field("model", "str"), + optional_field("serviceTier", "Any"), + optional_field("reasoningEffort", "Any"), + field("createdAt", "int"), + field("updatedAt", "int"), + optional_field("archivedAt", "int"), + field("cwd", "str"), + field("cliVersion", "str"), + field("source", "Any"), + optional_field("agentNickname", "str"), + optional_field("agentRole", "str"), + optional_field("agentPath", "str"), + optional_field("gitInfo", "Any"), + field("approvalMode", "Any"), + field("sandboxPolicy", "Any"), + optional_field("tokenUsage", "Any"), + optional_field("firstUserMessage", "str"), + optional_field("memoryMode", "str"), + ], + ), + struct_type( + "ThreadIndexPatch", + vec![ + field("updatedAt", "int"), + optional_field("firstUserMessage", "str"), + optional_field("preview", "str"), + optional_field("name", "str"), + optional_field("tokenUsage", "Any"), + optional_field("memoryMode", "str"), + ], + ), + struct_type( + "CreateThreadRequest", + vec![ + field("thread", "RemoteThreadMetadata"), + opaque_field("initialItems", "rollout_items", "Vec"), + field("initialPayloadCodec", "str"), + field("eventPersistenceMode", "str"), + ], + ), + struct_type( + "ResumeThreadRecorderRequest", + vec![ + field("threadId", "str"), + field("owner", "ThreadOwner"), + field("includeArchived", "bool"), + field("eventPersistenceMode", "str"), + ], + ), + struct_type( + "AppendThreadRequest", + vec![ + field("threadId", "str"), + field("owner", "ThreadOwner"), + optional_field("idempotencyKey", "str"), + optional_field("updatedAt", "int"), + optional_field("newThreadMemoryMode", "str"), + optional_field("eventPersistenceMode", "str"), + field("indexPatch", "ThreadIndexPatch"), + opaque_field("items", "rollout_items", "Vec"), + field("payloadCodec", "str"), + ], + ), + struct_type( + "LoadThreadHistoryRequest", + vec![ + field("threadId", "str"), + field("owner", "ThreadOwner"), + field("includeArchived", "bool"), + ], + ), + struct_type( + "LoadThreadHistoryResponse", + vec![ + field("threadId", "str"), + field("payloadCodec", "str"), + opaque_field("history", "stored_thread_history", "StoredThreadHistory"), + ], + ), + struct_type( + "ReadThreadRequest", + vec![ + field("threadId", "str"), + field("owner", "ThreadOwner"), + field("includeArchived", "bool"), + field("includeHistory", "bool"), + ], + ), + struct_type( + "ReadThreadResponse", + vec![ + field("thread", "RemoteThreadMetadata"), + optional_field("historyPayloadCodec", "str"), + optional_opaque_field( + "history", + "stored_thread_history", + "StoredThreadHistory", + ), + ], + ), + struct_type( + "ListThreadsRequest", + vec![ + field("owner", "ThreadOwner"), + field("pageSize", "int"), + optional_field("cursor", "str"), + field("sortKey", "str"), + field("allowedSources", "list[Any]"), + optional_field("modelProviders", "list[str]"), + field("archived", "bool"), + optional_field("cwd", "str"), + optional_field("searchTerm", "str"), + ], + ), + struct_type( + "ListThreadsResponse", + vec![ + field("items", "list[RemoteThreadMetadata]"), + optional_field("nextCursor", "str"), + optional_field("scanned", "int"), + ], + ), + struct_type( + "FindThreadByNameRequest", + vec![ + field("owner", "ThreadOwner"), + field("name", "str"), + field("includeArchived", "bool"), + optional_field("cwd", "str"), + field("allowedSources", "list[Any]"), + optional_field("modelProviders", "list[str]"), + ], + ), + struct_type( + "FindThreadByNameResponse", + vec![optional_field("thread", "RemoteThreadMetadata")], + ), + struct_type( + "SetThreadNameRequest", + vec![ + field("threadId", "str"), + field("owner", "ThreadOwner"), + field("name", "str"), + ], + ), + struct_type( + "UpdateThreadMetadataRequest", + vec![ + field("threadId", "str"), + field("owner", "ThreadOwner"), + field("patch", "Any"), + ], + ), + struct_type( + "ArchiveThreadRequest", + vec![field("threadId", "str"), field("owner", "ThreadOwner")], + ), + struct_type( + "DynamicToolsRequest", + vec![field("threadId", "str"), field("owner", "ThreadOwner")], + ), + struct_type( + "DynamicToolsResponse", + vec![optional_field("dynamicTools", "list[Any]")], + ), + struct_type( + "MemoryModeRequest", + vec![field("threadId", "str"), field("owner", "ThreadOwner")], + ), + struct_type( + "MemoryModeResponse", + vec![optional_field("memoryMode", "str")], + ), + struct_type( + "SetMemoryModeRequest", + vec![ + field("threadId", "str"), + field("owner", "ThreadOwner"), + field("memoryMode", "str"), + ], + ), + struct_type( + "ThreadSpawnEdgeRecord", + vec![ + field("parentThreadId", "str"), + field("childThreadId", "str"), + field("status", "str"), + ], + ), + struct_type( + "ListThreadSpawnEdgesRequest", + vec![ + field("threadId", "str"), + field("owner", "ThreadOwner"), + field("recursive", "bool"), + optional_field("status", "str"), + ], + ), + struct_type( + "ListThreadSpawnEdgesResponse", + vec![field("edges", "list[ThreadSpawnEdgeRecord]")], + ), + struct_type( + "FindThreadSpawnByPathRequest", + vec![ + field("threadId", "str"), + field("owner", "ThreadOwner"), + field("recursive", "bool"), + field("agentPath", "str"), + ], + ), + struct_type( + "FindThreadSpawnByPathResponse", + vec![optional_field("threadId", "str")], + ), + ], + methods: vec![ + method("thread_store/create_thread", "CreateThreadRequest", "None"), + method( + "thread_store/resume_thread_recorder", + "ResumeThreadRecorderRequest", + "None", + ), + method( + "thread_store/append_thread_items", + "AppendThreadRequest", + "None", + ), + method( + "thread_store/load_thread_history", + "LoadThreadHistoryRequest", + "LoadThreadHistoryResponse", + ), + method( + "thread_store/read_thread", + "ReadThreadRequest", + "ReadThreadResponse", + ), + method( + "thread_store/list_threads", + "ListThreadsRequest", + "ListThreadsResponse", + ), + method( + "thread_store/find_thread_by_name", + "FindThreadByNameRequest", + "FindThreadByNameResponse", + ), + method( + "thread_store/set_thread_name", + "SetThreadNameRequest", + "None", + ), + method( + "thread_store/update_thread_metadata", + "UpdateThreadMetadataRequest", + "ReadThreadResponse", + ), + method( + "thread_store/archive_thread", + "ArchiveThreadRequest", + "None", + ), + method( + "thread_store/unarchive_thread", + "ArchiveThreadRequest", + "ReadThreadResponse", + ), + method( + "thread_store/dynamic_tools", + "DynamicToolsRequest", + "DynamicToolsResponse", + ), + method( + "thread_store/memory_mode", + "MemoryModeRequest", + "MemoryModeResponse", + ), + method( + "thread_store/set_memory_mode", + "SetMemoryModeRequest", + "None", + ), + method( + "thread_store/mark_memory_mode_polluted", + "MemoryModeRequest", + "None", + ), + method( + "thread_store/upsert_thread_spawn_edge", + "ThreadSpawnEdgeRecord", + "None", + ), + method( + "thread_store/list_thread_spawn_edges", + "ListThreadSpawnEdgesRequest", + "ListThreadSpawnEdgesResponse", + ), + method( + "thread_store/find_thread_spawn_by_path", + "FindThreadSpawnByPathRequest", + "FindThreadSpawnByPathResponse", + ), + ], + } +} + +fn struct_type(name: &str, fields: Vec) -> BridgeType { + BridgeType { + name: name.to_string(), + kind: "struct".to_string(), + fields, + variants: Vec::new(), + } +} + +fn field(name: &str, python_type: &str) -> BridgeField { + BridgeField { + name: name.to_string(), + python_type: python_type.to_string(), + optional: false, + opaque: None, + } +} + +fn optional_field(name: &str, python_type: &str) -> BridgeField { + BridgeField { + name: name.to_string(), + python_type: python_type.to_string(), + optional: true, + opaque: None, + } +} + +fn opaque_field(name: &str, codec: &str, rust_type: &str) -> BridgeField { + BridgeField { + name: name.to_string(), + python_type: "OpaqueBytes".to_string(), + optional: false, + opaque: Some(OpaqueField { + codec: codec.to_string(), + rust_type: rust_type.to_string(), + }), + } +} + +fn optional_opaque_field(name: &str, codec: &str, rust_type: &str) -> BridgeField { + BridgeField { + name: name.to_string(), + python_type: "OpaqueBytes".to_string(), + optional: true, + opaque: Some(OpaqueField { + codec: codec.to_string(), + rust_type: rust_type.to_string(), + }), + } +} + +fn method(name: &str, request: &str, response: &str) -> BridgeMethod { + BridgeMethod { + name: name.to_string(), + request: request.to_string(), + response: response.to_string(), + } +} diff --git a/codex-rs/thread-store/src/remote/types.rs b/codex-rs/thread-store/src/remote/types.rs new file mode 100644 index 0000000000..ee5d897057 --- /dev/null +++ b/codex-rs/thread-store/src/remote/types.rs @@ -0,0 +1,644 @@ +use std::path::PathBuf; + +use chrono::DateTime; +use chrono::Utc; +use codex_protocol::ThreadId; +use codex_protocol::config_types::ServiceTier; +use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_protocol::openai_models::ReasoningEffort; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::GitInfo; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TokenUsage; +use serde::Deserialize; +use serde::Serialize; + +use crate::AppendThreadItemsParams; +use crate::ArchiveThreadParams; +use crate::CreateThreadParams; +use crate::DynamicToolsParams; +use crate::FindThreadByNameParams; +use crate::FindThreadSpawnByPathParams; +use crate::ListThreadSpawnEdgesParams; +use crate::ListThreadsParams; +use crate::LoadThreadHistoryParams; +use crate::ReadThreadParams; +use crate::ResumeThreadRecorderParams; +use crate::SetThreadMemoryModeParams; +use crate::SetThreadNameParams; +use crate::StoredThread; +use crate::StoredThreadHistory; +use crate::ThreadEventPersistenceMode; +use crate::ThreadMetadataPatch; +use crate::ThreadOwner; +use crate::ThreadSortKey; +use crate::ThreadSpawnEdge; +use crate::ThreadSpawnEdgeStatus; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; +use crate::UpdateThreadMetadataParams; + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct CreateThreadRequest { + pub thread: RemoteThreadMetadata, + pub initial_payload_codec: String, + pub event_persistence_mode: String, +} + +impl CreateThreadRequest { + pub(crate) fn from_params(params: CreateThreadParams) -> Self { + Self { + thread: RemoteThreadMetadata::from_create_params(¶ms), + initial_payload_codec: super::ROLLOUT_ITEMS_CODEC.to_string(), + event_persistence_mode: event_persistence_mode(params.event_persistence_mode), + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ResumeThreadRecorderRequest { + pub thread_id: String, + pub owner: ThreadOwner, + pub include_archived: bool, + pub event_persistence_mode: String, +} + +impl ResumeThreadRecorderRequest { + pub(crate) fn from_params(params: ResumeThreadRecorderParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + include_archived: params.include_archived, + event_persistence_mode: event_persistence_mode(params.event_persistence_mode), + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AppendThreadRequest { + pub thread_id: String, + pub owner: ThreadOwner, + pub idempotency_key: Option, + pub updated_at: Option, + pub new_thread_memory_mode: Option, + pub event_persistence_mode: Option, + pub index_patch: ThreadIndexPatch, + pub payload_codec: String, +} + +impl AppendThreadRequest { + pub(crate) fn from_params( + params: AppendThreadItemsParams, + event_persistence_mode: Option, + ) -> Self { + let index_patch = ThreadIndexPatch::from_items( + params.items.as_slice(), + params.updated_at.unwrap_or_else(Utc::now), + params.new_thread_memory_mode.clone(), + ); + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + idempotency_key: params.idempotency_key, + updated_at: params.updated_at.map(|updated_at| updated_at.timestamp()), + new_thread_memory_mode: params.new_thread_memory_mode, + event_persistence_mode, + index_patch, + payload_codec: super::ROLLOUT_ITEMS_CODEC.to_string(), + } + } +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ThreadIndexPatch { + pub updated_at: i64, + pub first_user_message: Option, + pub preview: Option, + pub name: Option, + pub token_usage: Option, + pub memory_mode: Option, +} + +impl ThreadIndexPatch { + fn from_items( + items: &[RolloutItem], + updated_at: DateTime, + memory_mode: Option, + ) -> Self { + let mut patch = Self { + updated_at: updated_at.timestamp(), + memory_mode, + ..Self::default() + }; + for item in items { + match item { + RolloutItem::EventMsg(codex_protocol::protocol::EventMsg::UserMessage(event)) => { + if patch.first_user_message.is_none() { + patch.first_user_message = Some(event.message.clone()); + } + if patch.preview.is_none() { + patch.preview = Some(event.message.clone()); + } + } + RolloutItem::EventMsg(codex_protocol::protocol::EventMsg::ThreadNameUpdated( + event, + )) => { + patch.name.clone_from(&event.thread_name); + } + RolloutItem::EventMsg(codex_protocol::protocol::EventMsg::TokenCount(event)) => { + if let Some(info) = event.info.as_ref() { + patch.token_usage = Some(info.total_token_usage.clone()); + } + } + RolloutItem::SessionMeta(meta) => { + if patch.memory_mode.is_none() { + patch.memory_mode.clone_from(&meta.meta.memory_mode); + } + } + RolloutItem::ResponseItem(_) + | RolloutItem::Compacted(_) + | RolloutItem::TurnContext(_) + | RolloutItem::EventMsg(_) => {} + } + } + patch + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct LoadThreadHistoryRequest { + pub thread_id: String, + pub owner: ThreadOwner, + pub include_archived: bool, +} + +impl LoadThreadHistoryRequest { + pub(crate) fn from_params(params: LoadThreadHistoryParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + include_archived: params.include_archived, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct LoadThreadHistoryResponse { + pub thread_id: String, + pub payload_codec: String, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct StoredThreadPayload { + pub thread_id: String, + pub items: Vec, +} + +impl StoredThreadPayload { + pub(crate) fn into_stored_thread_history(self) -> ThreadStoreResult { + Ok(StoredThreadHistory { + thread_id: parse_thread_id(self.thread_id)?, + items: self.items, + }) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ReadThreadRequest { + pub thread_id: String, + pub owner: ThreadOwner, + pub include_archived: bool, + pub include_history: bool, +} + +impl ReadThreadRequest { + pub(crate) fn from_params(params: ReadThreadParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + include_archived: params.include_archived, + include_history: params.include_history, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ReadThreadResponse { + pub thread: RemoteThreadMetadata, + pub history_payload_codec: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ListThreadsRequest { + pub owner: ThreadOwner, + pub page_size: usize, + pub cursor: Option, + pub sort_key: String, + pub allowed_sources: Vec, + pub model_providers: Option>, + pub archived: bool, + pub cwd: Option, + pub search_term: Option, +} + +impl ListThreadsRequest { + pub(crate) fn from_params(params: ListThreadsParams) -> Self { + Self { + owner: params.owner, + page_size: params.page_size, + cursor: params.cursor, + sort_key: sort_key(params.sort_key), + allowed_sources: params.allowed_sources, + model_providers: params.model_providers, + archived: params.archived, + cwd: params.cwd, + search_term: params.search_term, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ListThreadsResponse { + pub items: Vec, + pub next_cursor: Option, + pub scanned: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct FindThreadByNameRequest { + pub owner: ThreadOwner, + pub name: String, + pub include_archived: bool, + pub cwd: Option, + pub allowed_sources: Vec, + pub model_providers: Option>, +} + +impl FindThreadByNameRequest { + pub(crate) fn from_params(params: FindThreadByNameParams) -> Self { + Self { + owner: params.owner, + name: params.name, + include_archived: params.include_archived, + cwd: params.cwd, + allowed_sources: params.allowed_sources, + model_providers: params.model_providers, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct FindThreadByNameResponse { + pub thread: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct SetThreadNameRequest { + pub thread_id: String, + pub owner: ThreadOwner, + pub name: String, +} + +impl SetThreadNameRequest { + pub(crate) fn from_params(params: SetThreadNameParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + name: params.name, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct UpdateThreadMetadataRequest { + pub thread_id: String, + pub owner: ThreadOwner, + pub patch: ThreadMetadataPatch, +} + +impl UpdateThreadMetadataRequest { + pub(crate) fn from_params(params: UpdateThreadMetadataParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + patch: params.patch, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ArchiveThreadRequest { + pub thread_id: String, + pub owner: ThreadOwner, +} + +impl ArchiveThreadRequest { + pub(crate) fn from_params(params: ArchiveThreadParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct DynamicToolsRequest { + pub thread_id: String, + pub owner: ThreadOwner, +} + +impl DynamicToolsRequest { + pub(crate) fn from_params(params: DynamicToolsParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct DynamicToolsResponse { + pub dynamic_tools: Option>, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct MemoryModeRequest { + pub thread_id: String, + pub owner: ThreadOwner, +} + +impl MemoryModeRequest { + pub(crate) fn from_params(params: crate::ThreadMemoryModeParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct MemoryModeResponse { + pub memory_mode: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct SetMemoryModeRequest { + pub thread_id: String, + pub owner: ThreadOwner, + pub memory_mode: String, +} + +impl SetMemoryModeRequest { + pub(crate) fn from_params(params: SetThreadMemoryModeParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + memory_mode: params.memory_mode, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ThreadSpawnEdgeRecord { + pub parent_thread_id: String, + pub child_thread_id: String, + pub status: String, +} + +impl ThreadSpawnEdgeRecord { + pub(crate) fn from_thread_spawn_edge(edge: ThreadSpawnEdge) -> Self { + Self { + parent_thread_id: edge.parent_thread_id.to_string(), + child_thread_id: edge.child_thread_id.to_string(), + status: spawn_status(edge.status), + } + } + + pub(crate) fn into_thread_spawn_edge(self) -> ThreadStoreResult { + Ok(ThreadSpawnEdge { + parent_thread_id: parse_thread_id(self.parent_thread_id)?, + child_thread_id: parse_thread_id(self.child_thread_id)?, + status: parse_spawn_status(self.status.as_str())?, + }) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ListThreadSpawnEdgesRequest { + pub thread_id: String, + pub owner: ThreadOwner, + pub recursive: bool, + pub status: Option, +} + +impl ListThreadSpawnEdgesRequest { + pub(crate) fn from_params(params: ListThreadSpawnEdgesParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + recursive: params.recursive, + status: params.status.map(spawn_status), + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ListThreadSpawnEdgesResponse { + pub edges: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct FindThreadSpawnByPathRequest { + pub thread_id: String, + pub owner: ThreadOwner, + pub recursive: bool, + pub agent_path: String, +} + +impl FindThreadSpawnByPathRequest { + pub(crate) fn from_params(params: FindThreadSpawnByPathParams) -> Self { + Self { + thread_id: params.thread_id.to_string(), + owner: params.owner, + recursive: params.recursive, + agent_path: params.agent_path, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct FindThreadSpawnByPathResponse { + pub thread_id: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct RemoteThreadMetadata { + pub thread_id: String, + pub forked_from_id: Option, + pub owner: ThreadOwner, + pub preview: String, + pub name: Option, + pub model_provider: String, + pub model: Option, + pub service_tier: Option, + pub reasoning_effort: Option, + pub created_at: i64, + pub updated_at: i64, + pub archived_at: Option, + pub cwd: PathBuf, + pub cli_version: String, + pub source: SessionSource, + pub agent_nickname: Option, + pub agent_role: Option, + pub agent_path: Option, + pub git_info: Option, + pub approval_mode: AskForApproval, + pub sandbox_policy: SandboxPolicy, + pub token_usage: Option, + pub first_user_message: Option, + pub memory_mode: Option, +} + +impl RemoteThreadMetadata { + fn from_create_params(params: &CreateThreadParams) -> Self { + let now = Utc::now().timestamp(); + Self { + thread_id: params.thread_id.to_string(), + forked_from_id: params.forked_from_id.map(|id| id.to_string()), + owner: params.owner.clone(), + preview: String::new(), + name: None, + model_provider: params.model_provider.clone(), + model: params.model.clone(), + service_tier: params.service_tier, + reasoning_effort: params.reasoning_effort, + created_at: now, + updated_at: now, + archived_at: None, + cwd: params.cwd.clone(), + cli_version: params.cli_version.clone(), + source: params.source.clone(), + agent_nickname: None, + agent_role: None, + agent_path: None, + git_info: params.git_info.clone(), + approval_mode: params.approval_mode, + sandbox_policy: params.sandbox_policy.clone(), + token_usage: None, + first_user_message: None, + memory_mode: params.memory_mode.clone(), + } + } + + pub(crate) fn into_stored_thread( + self, + history: Option, + ) -> ThreadStoreResult { + Ok(StoredThread { + thread_id: parse_thread_id(self.thread_id)?, + forked_from_id: self.forked_from_id.map(parse_thread_id).transpose()?, + legacy_path: None, + owner: self.owner, + preview: self.preview, + name: self.name, + model_provider: self.model_provider, + model: self.model, + service_tier: self.service_tier, + reasoning_effort: self.reasoning_effort, + created_at: datetime_from_unix_seconds(self.created_at, "createdAt")?, + updated_at: datetime_from_unix_seconds(self.updated_at, "updatedAt")?, + archived_at: self + .archived_at + .map(|archived_at| datetime_from_unix_seconds(archived_at, "archivedAt")) + .transpose()?, + cwd: self.cwd, + cli_version: self.cli_version, + source: self.source, + agent_nickname: self.agent_nickname, + agent_role: self.agent_role, + agent_path: self.agent_path, + git_info: self.git_info, + approval_mode: self.approval_mode, + sandbox_policy: self.sandbox_policy, + token_usage: self.token_usage, + first_user_message: self.first_user_message, + memory_mode: self.memory_mode, + history, + }) + } +} + +pub(crate) fn event_persistence_mode(mode: ThreadEventPersistenceMode) -> String { + match mode { + ThreadEventPersistenceMode::Limited => "limited", + ThreadEventPersistenceMode::Extended => "extended", + } + .to_string() +} + +fn sort_key(sort_key: ThreadSortKey) -> String { + match sort_key { + ThreadSortKey::CreatedAt => "created_at", + ThreadSortKey::UpdatedAt => "updated_at", + } + .to_string() +} + +fn spawn_status(status: ThreadSpawnEdgeStatus) -> String { + match status { + ThreadSpawnEdgeStatus::Open => "open", + ThreadSpawnEdgeStatus::Closed => "closed", + } + .to_string() +} + +fn parse_spawn_status(status: &str) -> ThreadStoreResult { + match status { + "open" => Ok(ThreadSpawnEdgeStatus::Open), + "closed" => Ok(ThreadSpawnEdgeStatus::Closed), + _ => Err(ThreadStoreError::InvalidRequest { + message: format!("invalid remote thread spawn edge status `{status}`"), + }), + } +} + +fn parse_thread_id(thread_id: String) -> ThreadStoreResult { + ThreadId::from_string(thread_id.as_str()).map_err(|err| ThreadStoreError::InvalidRequest { + message: format!("invalid remote thread id `{thread_id}`: {err}"), + }) +} + +fn datetime_from_unix_seconds(value: i64, field: &str) -> ThreadStoreResult> { + DateTime::from_timestamp(value, 0).ok_or_else(|| ThreadStoreError::InvalidRequest { + message: format!("invalid remote `{field}` timestamp `{value}`"), + }) +} diff --git a/scripts/generate_thread_store_bridge_python.py b/scripts/generate_thread_store_bridge_python.py new file mode 100644 index 0000000000..74bd0d037a --- /dev/null +++ b/scripts/generate_thread_store_bridge_python.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Generate Python bridge types from the Rust thread-store bridge schema. + +The Rust schema is the source of truth. By default this script invokes: + + cargo run -p codex-thread-store --bin codex-thread-store-bridge-schema + +from the codex-rs workspace, then writes a Python module containing dataclasses +for request/response DTOs and a Protocol for the service implementation. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +from pathlib import Path +from typing import Any + + +REPO_ROOT = Path(__file__).resolve().parents[1] +CODEX_RS = REPO_ROOT / "codex-rs" + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--schema-json", + type=Path, + help="Read a previously generated schema JSON file instead of invoking cargo.", + ) + parser.add_argument( + "--output", + type=Path, + default=CODEX_RS / "thread-store" / "python" / "codex_thread_store_bridge.py", + help="Python file to write.", + ) + args = parser.parse_args() + + schema = load_schema(args.schema_json) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(render_python(schema), encoding="utf-8") + + +def load_schema(schema_json: Path | None) -> dict[str, Any]: + if schema_json is not None: + return json.loads(schema_json.read_text(encoding="utf-8")) + + result = subprocess.run( + [ + "cargo", + "run", + "-p", + "codex-thread-store", + "--bin", + "codex-thread-store-bridge-schema", + ], + cwd=CODEX_RS, + check=True, + capture_output=True, + text=True, + ) + return json.loads(result.stdout) + + +def render_python(schema: dict[str, Any]) -> str: + lines: list[str] = [ + "# @generated by scripts/generate_thread_store_bridge_python.py", + "from __future__ import annotations", + "", + "from dataclasses import dataclass", + "from typing import Any, Protocol", + "", + "", + ] + + for type_def in schema["types"]: + if type_def["kind"] != "struct": + continue + lines.extend(render_dataclass(type_def)) + lines.append("") + + lines.extend(render_protocol(schema["methods"])) + return "\n".join(lines).rstrip() + "\n" + + +def render_dataclass(type_def: dict[str, Any]) -> list[str]: + lines = ["@dataclass(frozen=True, kw_only=True)", f"class {type_def['name']}:"] + fields = type_def["fields"] + if not fields: + lines.append(" pass") + return lines + + for field in fields: + py_type = field["pythonType"] + if field["optional"]: + py_type = f"{py_type} | None" + default = " = None" if field["optional"] else "" + lines.append(f" {field['name']}: {py_type}{default}") + return lines + + +def render_protocol(methods: list[dict[str, str]]) -> list[str]: + lines = ["class ThreadStoreBridgeService(Protocol):"] + if not methods: + lines.append(" pass") + return lines + + for method in methods: + py_name = ( + method["name"] + .removeprefix("thread_store/") + .replace("/", "_") + .replace("-", "_") + ) + request_type = method["request"] + response_type = "None" if method["response"] == "None" else method["response"] + lines.append( + f" async def {py_name}(self, request: {request_type}) -> {response_type}: ..." + ) + return lines + + +if __name__ == "__main__": + main()