diff --git a/.env b/.env new file mode 100644 index 0000000..4a5251d --- /dev/null +++ b/.env @@ -0,0 +1,6 @@ +GOOGLE_OAUTH_CLIENT_ID= +GOOGLE_OAUTH_CLIENT_SECRET= +ANTIGRAVITY_DEFAULT_PROJECT_ID= +UID= +GID= +USER= diff --git a/.gitignore b/.gitignore index 0892ad6..566248a 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,8 @@ /clawshell.toml /tarpaulin-report.json /tarpaulin-report.html +.env +.idea/ # npm platform binaries (added during release, not checked in) npm/clawshell-*/bin/clawshell diff --git a/Cargo.lock b/Cargo.lock index 3376b4e..eae84a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anstream" version = "0.6.21" @@ -92,6 +101,17 @@ dependencies = [ "wait-timeout", ] +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -252,6 +272,20 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chrono" +version = "0.4.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fac4744fb15ae8337dc853fee7fb3f4e48c0fbaa23d0afe49c447b4fab126118" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", +] + [[package]] name = "clap" version = "4.5.60" @@ -297,19 +331,26 @@ name = "clawshell" version = "0.1.1" dependencies = [ "assert_cmd", + "async-trait", "axum", + "base64", "bytes", + "chrono", "clap", "console 0.16.2", + "dotenvy", "futures-util", "http", "http-body-util", "inquire", "insta", "nix", + "oauth2", + "open", "predicates", + "rand 0.9.2", "regex", - "reqwest", + "reqwest 0.13.2", "rustls", "rustls-native-certs", "semver", @@ -319,11 +360,13 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "tokio", + "tokio-util", "toml", "tower", "tower-http", "tracing", "tracing-subscriber", + "urlencoding", "uuid", "vfs", "wiremock", @@ -526,6 +569,12 @@ dependencies = [ "litrs", ] +[[package]] +name = "dotenvy" +version = "0.15.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" + [[package]] name = "dunce" version = "1.0.5" @@ -894,6 +943,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", + "webpki-roots", ] [[package]] @@ -919,6 +969,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "2.1.1" @@ -1082,6 +1156,25 @@ dependencies = [ "serde", ] +[[package]] +name = "is-docker" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "928bae27f42bc99b60d9ac7334e3a21d10ad8f1835a4e12ec3ec0464765ed1b3" +dependencies = [ + "once_cell", +] + +[[package]] +name = "is-wsl" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "173609498df190136aa7dea1a91db051746d339e18476eed5ca40521f02d7aa5" +dependencies = [ + "is-docker", + "once_cell", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -1289,6 +1382,26 @@ dependencies = [ "libc", ] +[[package]] +name = "oauth2" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d" +dependencies = [ + "base64", + "chrono", + "getrandom 0.2.17", + "http", + "rand 0.8.5", + "reqwest 0.12.28", + "serde", + "serde_json", + "serde_path_to_error", + "sha2", + "thiserror 1.0.69", + "url", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -1301,6 +1414,17 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "open" +version = "5.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43bb73a7fa3799b198970490a51174027ba0d4ec504b03cd08caf513d40024bc" +dependencies = [ + "is-wsl", + "libc", + "pathdiff", +] + [[package]] name = "openssl-probe" version = "0.2.1" @@ -1330,6 +1454,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "pathdiff" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df94ce210e5bc13cb6651479fa48d14f601d9858cfe0467f43ae157023b938d3" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1445,7 +1575,7 @@ dependencies = [ "bytes", "getrandom 0.3.4", "lru-slab", - "rand", + "rand 0.9.2", "ring", "rustc-hash", "rustls", @@ -1486,14 +1616,35 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -1503,7 +1654,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", ] [[package]] @@ -1562,6 +1722,44 @@ version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a96887878f22d7bad8a3b6dc5b7440e0ada9a245242924394987b21cf2210a4c" +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64", + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "webpki-roots", +] + [[package]] name = "reqwest" version = "0.13.2" @@ -1588,6 +1786,7 @@ dependencies = [ "rustls-pki-types", "rustls-platform-verifier", "serde", + "serde_json", "serde_urlencoded", "sync_wrapper", "tokio", @@ -2352,8 +2551,15 @@ dependencies = [ "idna", "percent-encoding", "serde", + "serde_derive", ] +[[package]] +name = "urlencoding" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" + [[package]] name = "utf8_iter" version = "1.0.4" @@ -2585,6 +2791,15 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "webpki-roots" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22cfaf3c063993ff62e73cb4311efde4db1efb31ab78a3e5c457939ad5cc0bed" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "winapi" version = "0.3.9" @@ -2616,12 +2831,65 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.45.0" diff --git a/Cargo.toml b/Cargo.toml index 37a76bd..3d17cf3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,7 @@ description = "A security privileged process for the OpenClaw ecosystem." [dependencies] axum = "0.8.8" tokio = { version = "1.49", features = ["full"] } -reqwest = { version = "0.13.2", default-features = false, features = ["stream", "rustls", "form", "blocking"] } +reqwest = { version = "0.13.2", default-features = false, features = ["stream", "rustls", "form", "blocking", "json"] } rustls = { version = "0.23.36", default-features = false, features = ["ring", "std"] } serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.149" @@ -34,6 +34,15 @@ semver = "1" rustls-native-certs = "0.8.3" sha2 = "0.10.9" uuid = { version = "1.18.1", features = ["v4"] } +oauth2 = "5" +open = "5" +chrono = { version = "0.4", features = ["serde"] } +async-trait = "0.1" +base64 = "0.22" +rand = "0.9" +urlencoding = "2" +tokio-util = "0.7" +dotenvy = "0.15" [dev-dependencies] tokio = { version = "1.49", features = ["full", "test-util"] } diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..8a81d96 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,19 @@ +FROM openclaw +ARG UID=1000 +ARG GID=1000 +ARG USER=app +ENV UID=$UID +ENV GID=$GID +ENV USER=$USER + +ENV TERM=xterm-256color +ENV PATH="/home/node/nodeenv/bin:$PATH" +ENV CLAWSHELL_OAUTH_CALLBACK_HOST=0.0.0.0 +ENV CLAWSHELL_SERVER_HOST=0.0.0.0 +EXPOSE 18790 51121 +COPY target/release/clawshell /usr/local/bin/clawshell +RUN sudo useradd -r -m -o -u $UID -g $GID -s /bin/bash clawshell +USER $USER +WORKDIR /home/node +COPY .env /home/node/.env +ENTRYPOINT ["sudo", "-E", "env", "CLAWSHELL_OAUTH_CALLBACK_HOST=0.0.0.0", "PATH=/home/node/nodeenv/bin:/home/node/nodeenv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", "clawshell"] diff --git a/Dockerfile.openclaw b/Dockerfile.openclaw new file mode 100644 index 0000000..76ae4b2 --- /dev/null +++ b/Dockerfile.openclaw @@ -0,0 +1,23 @@ +FROM debian:13.3-slim +ARG UID=1000 +ARG GID=1000 +ARG USER=app +ARG VERSION=latest +ENV UID=$UID +ENV GID=$GID +ENV USER=$USER +ENV VERSION=$VERSION + +RUN groupadd -r -g $GID $USER +RUN useradd -r -m -u $UID -g $GID -s /bin/bash $USER + +RUN apt-get update -y && apt-get install -y nodejs npm nodeenv linux-headers-generic make g++ cmake git sudo ca-certificates && rm -rf /var/lib/apt/lists/* +RUN echo "$USER ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers +RUN mkdir /home/node && chown -R $UID:$GID /home/node + +USER $USER +RUN nodeenv --node=24.13.1 --npm=v11.8.0 /home/node/nodeenv +RUN cd /home/node && . ./nodeenv/bin/activate && npm config set prefix /home/node && npm install -g openclaw@$VERSION + +ENV PATH="/home/node/nodeenv/bin:$PATH" +ENTRYPOINT ["/home/node/nodeenv/bin/openclaw"] diff --git a/README.md b/README.md index d6d0f1c..310ff49 100644 --- a/README.md +++ b/README.md @@ -173,6 +173,12 @@ sudo clawshell migrate-config By default ClawShell listens on `127.0.0.1:18790`. +You can override the bind address at runtime with environment variables: + +```bash +CLAWSHELL_SERVER_HOST=0.0.0.0 CLAWSHELL_SERVER_PORT=17890 clawshell start --foreground +``` + ### Customized Configuration ClawShell reads its config from `/etc/clawshell/clawshell.toml`. You can view or edit it with: diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..b9d608e --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,44 @@ +services: + openclaw-gateway: + build: + context: . + dockerfile: Dockerfile.openclaw + image: openclaw + container_name: openclaw-gateway + command: ["gateway","run","--bind","lan","--port","18789"] + ports: + - "18789:18789" + environment: + UID: $UID + GID: $GID + USER: $USER + volumes: + - ./config/openclaw:/home/app/.openclaw + - ./config/openclaw:/home/$USER/.openclaw + restart: unless-stopped + + clawshell: + build: + context: . + dockerfile: Dockerfile + image: clawshell + container_name: clawshell + depends_on: + - openclaw-gateway + command: ["start", "--config", "/etc/clawshell/clawshell.toml", "--foreground"] + environment: + CLAWSHELL_SERVER_HOST: "0.0.0.0" + CLAWSHELL_OAUTH_CALLBACK_HOST: "0.0.0.0" + UID: $UID + GID: $GID + USER: $USER + ports: + - "18790:18790" + - "51121:51121" + volumes: + - ./config/clawshell:/etc/clawshell + restart: unless-stopped + +volumes: + config-openclaw: + config-clawshell: diff --git a/docs/architecture-comparison.md b/docs/architecture-comparison.md new file mode 100644 index 0000000..8b1ed6b --- /dev/null +++ b/docs/architecture-comparison.md @@ -0,0 +1,627 @@ +# Architecture Comparison: Current vs. Multi-Provider OAuth + +This document compares ClawShell's **current architecture** (static API keys only) +with the **proposed architecture** after adding multi-provider OAuth support. + +The first version (v1) implements two providers: **Codex (OpenAI)** and +**Antigravity (Google)**. OAuth providers are integrated into the existing +`clawshell onboard` wizard — no new CLI subcommands are added. + +--- + +## 1. High-Level Flow + +### Current (Static API Key) + +``` +┌──────────┐ Authorization: Bearer vk-001 ┌─────────────┐ Authorization: Bearer sk-real-... ┌──────────────┐ +│ │ ─────────────────────────────────► │ │ ──────────────────────────────────► │ │ +│ OpenClaw │ │ ClawShell │ │ OpenAI API │ +│ │ ◄───────────────────────────────── │ │ ◄────────────────────────────────── │ │ +└──────────┘ response └─────────────┘ response └──────────────┘ + │ + Lookup vk-001 + in BTreeMap + │ + ▼ + clawshell.toml + (static real_key) +``` + +**Characteristics:** +- One-time setup via `clawshell onboard`: paste API key +- Key never changes — no refresh needed +- Key lives on disk permanently in plaintext (protected by Unix file permissions) +- No external auth server interaction at runtime + +### Proposed (Multi-Provider OAuth) + +``` + clawshell onboard (same command, expanded menu) + ┌──────────────────────────────────────────────────┐ + │ │ + │ Select a model provider: │ + │ 1. OpenAI → prompt for API key │ + │ 2. OpenRouter → prompt for API key │ + │ 3. Anthropic → prompt for API key │ + │ 4. Codex / ChatGPT → OAuth browser flow │ ← NEW + │ 5. Antigravity / Google → OAuth browser flow │ ← NEW + │ │ + └──────────────────────────────────────────────────┘ + + RUNTIME (per request) +┌──────────┐ Bearer vk-001 ┌──────────────────────────────────────────────────┐ +│ │ ──────────────────► │ ClawShell │ +│ OpenClaw │ │ │ +│ │ ◄────────────────── │ 1. Lookup vk-001 → ResolvedKey { source, prov } │ +└──────────┘ response │ 2. KeySource? │ + │ ├── Static(key) → inject key (existing logic) │ + │ └── OAuth{provider_id} │ + │ ├── registry.inject_auth(id, headers) │ + │ ├── registry.prepare_request_body(id, b) │ + │ └── registry.upstream_url(id) │ + │ 3. Forward to upstream │ + │ 4. On 401 (OAuth only) → refresh + retry │ + └──────────────────────────────────────────────────┘ + + BACKGROUND (one task per active provider) + ┌─────────────────────────────────────────────────────────────┐ + │ codex: sleep(75% of ~8-day TTL) → auth.openai.com │ + │ antigravity: check 60s before expiry → googleapis.com │ + └─────────────────────────────────────────────────────────────┘ +``` + +--- + +## 2. Module Comparison + +### Current Module Map + +``` +src/ +├── main.rs CLI dispatch, daemon lifecycle +├── lib.rs AppState, build_router(), handle_request() +├── cli.rs Clap CLI definitions +├── config.rs TOML config model + validation +├── keys.rs KeyManager: virtual→real key BTreeMap lookup +├── dlp.rs DLP regex scanner (block/redact) +├── proxy.rs ProxyClient: upstream HTTP forwarding +├── onboard/ +│ ├── mod.rs Public API +│ ├── interactive.rs TUI wizard prompts +│ ├── types.rs OnboardConfig struct +│ ├── config_render.rs TOML generation +│ ├── credentials.rs API key detection +│ └── ... backup, openclaw_json, skills, etc. +├── process.rs PID file, privilege drop +├── tui.rs Terminal UI +└── platform/ + ├── mod.rs Platform dispatch + ├── linux.rs systemd, useradd + └── macos.rs launchctl, dscl +``` + +### Proposed Module Map + +``` +src/ +├── main.rs CLI dispatch + OAuthRegistry init + refresh tasks ← MODIFIED +├── lib.rs AppState + OAuthRegistry ← MODIFIED +├── cli.rs Clap CLI (UNCHANGED — no new subcommands) ← UNCHANGED +├── config.rs TOML config + [[oauth_providers]] + auth field ← MODIFIED +├── keys.rs KeyManager: virtual→KeySource (Static | OAuth{id}) ← MODIFIED +├── dlp.rs DLP regex scanner (block/redact) ← UNCHANGED +├── oauth/ ← NEW DIR +│ ├── mod.rs OAuthProvider trait, OAuthRegistry, OAuthTokens ← NEW +│ ├── codex.rs Codex (OpenAI): PKCE + device code ← NEW [v1] +│ ├── antigravity.rs Antigravity (Google): PKCE + headless URL ← NEW [v1] +│ └── storage.rs Per-provider token persistence ← NEW +├── proxy.rs ProxyClient + inject_auth + prepare_body + 401 ← MODIFIED +├── onboard/ +│ ├── mod.rs Public API ← UNCHANGED +│ ├── interactive.rs Provider menu + OAuth login branch ← MODIFIED +│ ├── types.rs OnboardConfig with AuthMethod enum ← MODIFIED +│ ├── config_render.rs TOML gen + [[oauth_providers]] rendering ← MODIFIED +│ ├── credentials.rs API key detection ← UNCHANGED +│ └── ... backup, openclaw_json, skills, etc. ← UNCHANGED +├── process.rs PID file, privilege drop ← UNCHANGED +├── tui.rs Terminal UI ← UNCHANGED +└── platform/ ← UNCHANGED +``` + +**Summary: 1 new directory with 4 files, 7 modified files, rest unchanged.** + +--- + +## 3. Core Abstraction: `OAuthProvider` Trait + +### How Providers Differ + +| Trait Method | Codex (OpenAI) | Antigravity (Google) | +|---------------------------|---------------------------------------------|------------------------------------------------| +| `id()` | `"codex"` | `"antigravity"` | +| `display_name()` | `"Codex (OpenAI)"` | `"Antigravity (Google)"` | +| `supports_device_code()` | `true` | `false` | +| `supports_headless_url()` | `false` | `true` | +| `login_browser()` | PKCE → `auth.openai.com` | PKCE → `accounts.google.com` + project discovery| +| `login_headless()` | Device code polling | Print URL, paste redirect back | +| `refresh()` | POST `auth.openai.com/oauth/token` | POST `oauth2.googleapis.com/token` | +| `inject_auth()` | `Authorization: Bearer ` | `Authorization: Bearer` + `X-Goog-Api-Client` + `Client-Metadata` | +| `prepare_request_body()` | `None` (pass-through) | `Some(wrapped)` (Gemini-style envelope) | +| `upstream_url()` | `None` (use `[upstream].base_url`) | `Some("cloudcode-pa.googleapis.com/...")` | + +--- + +## 4. Data Structures Comparison + +### `ResolvedKey` + +**Current:** + +```rust +pub struct ResolvedKey { + pub real_key: String, + pub provider: Provider, +} +``` + +**Proposed:** + +```rust +pub enum KeySource { + Static(String), + OAuth { provider_id: String }, +} + +pub struct ResolvedKey { + pub source: KeySource, + pub provider: Provider, +} +``` + +### `AppState` + +**Current:** + +```rust +pub struct AppState { + pub key_manager: Arc, + pub dlp_scanner: Arc, + pub proxy_client: Arc, +} +``` + +**Proposed:** + +```rust +pub struct AppState { + pub key_manager: Arc, + pub dlp_scanner: Arc, + pub proxy_client: Arc, + pub oauth_registry: Option>, +} +``` + +### `OnboardConfig` + +**Current:** + +```rust +pub struct OnboardConfig { + pub provider: String, + pub model: String, + pub real_api_key: String, // always required + pub virtual_api_key: String, + pub openclaw_config_path: PathBuf, + pub server_host: String, + pub server_port: u16, + pub email: Option, +} +``` + +**Proposed:** + +```rust +pub enum AuthMethod { + ApiKey { real_api_key: String }, + OAuth { provider_id: String }, // tokens stored during onboard flow +} + +pub struct OnboardConfig { + pub provider: String, + pub model: String, + pub auth: AuthMethod, // was: pub real_api_key: String + pub virtual_api_key: String, + pub openclaw_config_path: PathBuf, + pub server_host: String, + pub server_port: u16, + pub email: Option, +} +``` + +### `Config` / `KeyMapping` + +**Current:** + +```rust +pub struct KeyMapping { + pub virtual_key: String, + pub real_key: String, + pub provider: Provider, +} +``` + +**Proposed:** + +```rust +pub struct KeyMapping { + pub virtual_key: String, + pub real_key: Option, // optional when auth = "oauth" + pub provider: Provider, + pub auth: AuthMethod, // defaults to Static + pub oauth_provider: Option, // "codex" or "antigravity" +} +``` + +--- + +## 5. Configuration Comparison + +### Current `clawshell.toml` + +```toml +log_level = "info" + +[server] +host = "127.0.0.1" +port = 18790 + +[upstream] +base_url = "https://api.openai.com" + +[[keys]] +virtual_key = "vk-alice-001" +real_key = "sk-abc123..." +provider = "openai" + +[dlp] +scan_responses = true +patterns = [ + { name = "ssn", regex = '\b\d{3}-\d{2}-\d{4}\b', action = "redact" }, +] +``` + +### Proposed (Generated by `clawshell onboard` When OAuth Is Selected) + +```toml +log_level = "info" + +[server] +host = "127.0.0.1" +port = 18790 + +[upstream] +base_url = "https://api.openai.com" + +# OAuth-backed key — generated by onboard wizard +[[keys]] +virtual_key = "vk-chatgpt-001" +provider = "openai" +auth = "oauth" +oauth_provider = "codex" + +[[oauth_providers]] +provider = "codex" + +[dlp] +scan_responses = true +patterns = [ + { name = "ssn", regex = '\b\d{3}-\d{2}-\d{4}\b', action = "redact" }, +] +``` + +**When selecting API key providers, the output is identical to today.** + +--- + +## 6. Onboarding Flow Comparison + +This is the primary user-facing change — OAuth is integrated into the existing +`clawshell onboard` command, not a separate subcommand. + +### Current Onboarding Flow + +``` +$ sudo clawshell onboard + + Select a model provider: + ► OpenAI + OpenRouter + Anthropic + + Enter the model name: [gpt-5.2-chat-latest] + Enter the real API key: **************************** + Enter the virtual API key: [{clawshell-virtual-key-openai}] + ... (email, OpenClaw config, server settings) +``` + +### Proposed Onboarding Flow + +``` +$ sudo clawshell onboard + + Select a model provider: + ► OpenAI ← existing (API key) + OpenRouter ← existing (API key) + Anthropic ← existing (API key) + Codex / ChatGPT (OAuth) ← NEW + Antigravity / Google (OAuth) ← NEW + + ─── If user selects "Codex / ChatGPT (OAuth)" ────────── + + Enter the model name: [gpt-5.2-chat-latest] + + Opening browser for ChatGPT login... + (browser opens to auth.openai.com) + ✓ Login successful. Tokens saved. + + Enter the virtual API key: [{clawshell-virtual-key-codex}] + ... (email, OpenClaw config, server settings — unchanged) + + ─── If user selects "Antigravity / Google (OAuth)" ───── + + Enter the model name: [gemini-3-pro] + + Opening browser for Google login... + (browser opens to accounts.google.com) + ✓ Login successful. Project ID: proj-abc-123. Tokens saved. + + Enter the virtual API key: [{clawshell-virtual-key-antigravity}] + ... (email, OpenClaw config, server settings — unchanged) + + ─── If user selects "OpenAI" / "OpenRouter" / "Anthropic" ── + + (Identical to today — prompt for API key) +``` + +**In headless (SSH) environments:** + +``` + Codex: "Enter the device code shown in your browser: ___" + Antigravity: "Visit this URL, then paste the redirect URL here: ___" +``` + +--- + +## 7. CLI Commands Comparison + +### Current + +``` +clawshell start Start the proxy daemon +clawshell stop Stop the daemon +clawshell status Check daemon status +clawshell restart Restart the daemon +clawshell logs View/tail log file +clawshell config Display/edit config +clawshell onboard Interactive setup wizard +clawshell uninstall Remove ClawShell +clawshell version Print version +``` + +### Proposed + +``` +clawshell start Start daemon (+ spawn refresh tasks if OAuth configured) ← MODIFIED behavior +clawshell stop Stop the daemon ← UNCHANGED +clawshell status Check daemon status ← UNCHANGED +clawshell restart Restart the daemon ← UNCHANGED +clawshell logs View/tail log file ← UNCHANGED +clawshell config Display/edit config ← UNCHANGED +clawshell onboard Setup wizard (now with OAuth provider options) ← MODIFIED behavior +clawshell uninstall Remove ClawShell (+ remove OAuth token files) ← MODIFIED behavior +clawshell version Print version ← UNCHANGED +``` + +**No new subcommands.** The CLI interface is identical. Only the behavior of +`onboard`, `start`, and `uninstall` changes. + +--- + +## 8. Request Pipeline Comparison + +### Current: 5-Step Pipeline + +``` +Step 1 Extract Authorization header → extract_virtual_key() +Step 2 Resolve virtual key → resolve() → ResolvedKey { real_key, provider } +Step 3 Buffer request body +Step 4 DLP scan request body +Step 5 Forward to upstream → forward(real_key, provider) +Step 6 Optional DLP scan response +``` + +### Proposed: Pipeline With Provider-Aware Branching + +``` +Step 1 Extract Authorization header → extract_virtual_key() +Step 2 Resolve virtual key → resolve() → ResolvedKey { source, provider } + + ┌─── Static path (unchanged) ─────────────────────────────────────────────┐ + │ Step 3 Buffer body → DLP scan → Forward with static key │ + └─────────────────────────────────────────────────────────────────────────┘ + + ┌─── OAuth path (NEW) ───────────────────────────────────────────────────┐ + │ Step 3 Get access token via OAuthRegistry │ + │ Step 4 Buffer body → DLP scan │ + │ Step 5 Provider-specific prep: │ + │ Codex: inject_auth (Bearer) + pass-through body │ + │ Antigravity: inject_auth (Bearer + headers) + wrap body │ + │ Step 6 Forward to provider-resolved upstream │ + │ Step 7 On 401: refresh token → retry once │ + │ Step 8 Optional DLP scan response │ + └─────────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## 9. Upstream Request Comparison + +### Codex (OpenAI) — Thin Provider + +``` +After ClawShell: + POST /v1/chat/completions HTTP/1.1 ← same path as static key + Authorization: Bearer eyJ...access_token ← OAuth token + Content-Type: application/json + + {"model":"gpt-4o","messages":[...]} ← same body (pass-through) + +Upstream: api.openai.com +``` + +### Antigravity (Google) — Thick Provider + +``` +After ClawShell: + POST /v1internal:streamGenerateContent?alt=sse HTTP/1.1 ← different path + Authorization: Bearer ya29...access_token + X-Goog-Api-Client: google-cloud-sdk vscode_cloudshelleditor/0.1 + Client-Metadata: {"ideType":"ANTIGRAVITY",...} + Content-Type: application/json + + { "project": "proj-abc-123", "model": "gemini-3-pro", ← wrapped body + "request": { "contents": [...] } } + +Upstream: cloudcode-pa.googleapis.com +``` + +--- + +## 10. Credential Lifecycle Comparison + +### Current: Static Key + +``` + clawshell onboard Runtime clawshell uninstall +┌───────────────────┐ ┌──────────────────────┐ ┌───────────────────┐ +│ Paste API key │ │ Key loaded at startup │ │ Deletes config │ +│ → clawshell.toml │ │ Never changes │ │ Key is gone │ +│ │ │ No background tasks │ │ │ +└───────────────────┘ └──────────────────────┘ └───────────────────┘ +``` + +### Proposed: OAuth Token (Per Provider) + +``` + clawshell onboard Runtime clawshell uninstall +┌─────────────────────┐ ┌──────────────────────────────┐ ┌──────────────────┐ +│ Select Codex → │ │ Per-provider refresh tasks: │ │ Deletes config + │ +│ browser opens → │ │ codex: sleep(75% of ~8d TTL)│ │ oauth/ directory │ +│ tokens saved to │ │ antigravity: check 60s early│ │ │ +│ oauth/codex.json │ │ │ │ Tokens are gone │ +│ │ │ On 401: refresh + retry │ │ │ +│ OR │ │ │ │ Re-onboard to │ +│ │ │ Providers are independent │ │ login again │ +│ Select Antigravity →│ │ │ │ │ +│ browser opens → │ │ │ │ │ +│ project ID found → │ │ │ │ │ +│ oauth/antigravity. │ │ │ │ │ +│ json │ │ │ │ │ +└─────────────────────┘ └──────────────────────────────┘ └──────────────────┘ +``` + +--- + +## 11. Security Model Comparison + +### Current + +``` +/etc/clawshell/clawshell.toml (0600) + static API keys — permanent, manual revocation only +``` + +### Proposed (Additions) + +``` +/etc/clawshell/oauth/ (0700) +├── codex.json (0600) — ~8-day access token, single-use refresh +└── antigravity.json (0600) — ~1-hour access token, standard refresh + +Improvements: short-lived tokens, auto-rotation, revocable from provider dashboard +New surface: token files on disk (same 0600 mitigation), ephemeral callback servers, + network dependency on auth servers for refresh +``` + +--- + +## 12. Daemon Lifecycle Comparison + +### Current Startup + +``` +main() + ├── load config + ├── build AppState { KeyManager, DlpScanner, ProxyClient } + ├── bind socket → drop privileges → write PID + └── serve (no background tasks) +``` + +### Proposed Startup + +``` +main() + ├── load config + ├── if [[oauth_providers]]: ← NEW + │ ├── instantiate providers + │ ├── load tokens from oauth/.json + │ └── create OAuthRegistry + ├── build AppState { ..., OAuthRegistry? } + ├── bind socket → drop privileges → write PID + ├── if OAuthRegistry: ← NEW + │ └── spawn per-provider refresh tasks + └── serve +``` + +--- + +## 13. File Layout Comparison + +### Current + +``` +/etc/clawshell/ +├── clawshell.toml config (0600) +└── config.json onboarding metadata (0600) +``` + +### Proposed + +``` +/etc/clawshell/ +├── clawshell.toml config (0600) +├── config.json onboarding metadata (0600) +└── oauth/ token directory (0700) ← NEW + ├── codex.json OpenAI OAuth tokens (0600) ← NEW [v1] + └── antigravity.json Google OAuth tokens (0600) ← NEW [v1] +``` + +--- + +## 14. Summary Table + +| Aspect | Current | Proposed | +|-----------------------|---------------------------------|---------------------------------------------------| +| Auth methods | Static API keys only | Static + Codex OAuth + Antigravity OAuth | +| CLI commands | 9 commands | Same 9 commands (no new subcommands) | +| Onboard menu | OpenAI, OpenRouter, Anthropic | + Codex/ChatGPT, Antigravity/Google | +| Credential lifetime | Permanent | Static: permanent; Codex: ~8d; Antigravity: ~1hr | +| Background tasks | None | One refresh task per active OAuth provider | +| External auth calls | None | HTTPS to provider auth servers (refresh only) | +| New files on disk | — | `oauth/codex.json`, `oauth/antigravity.json` | +| New Rust modules | — | `oauth/` (mod, codex, antigravity, storage) | +| Modified modules | — | 7 files (main, lib, config, keys, proxy, onboard×2)| +| New dependencies | — | `oauth2`, `open`, `chrono`, `async-trait` | +| Config format | No OAuth section | `[[oauth_providers]]` + `[[keys]].auth` | +| Backward compatible | — | Yes — no OAuth config = identical behavior | diff --git a/docs/codex-oauth-plan.md b/docs/codex-oauth-plan.md new file mode 100644 index 0000000..d53e15b --- /dev/null +++ b/docs/codex-oauth-plan.md @@ -0,0 +1,729 @@ +# Multi-Provider OAuth Integration Plan for ClawShell + +## Executive Summary + +This document proposes adding a **multi-provider OAuth framework** to ClawShell, +enabling users to authenticate with subscription-based accounts instead of (or +alongside) static API keys. The first version implements two OAuth providers: +**Codex (OpenAI)** and **Antigravity (Google)**. + +OAuth providers are integrated into the existing `clawshell onboard` wizard — +no new CLI subcommands are needed. Users select a provider from the same menu +that already shows OpenAI, OpenRouter, and Anthropic. + +### Provider Roadmap + +| Provider | Status | Upstream API | Auth Via | +|---------------------------|------------------|--------------------------------------|-------------------------------| +| **Codex (OpenAI)** | v1 — Implement | `api.openai.com` | ChatGPT Plus/Pro subscription | +| **Antigravity (Google)** | v1 — Implement | `cloudcode-pa.googleapis.com` | Google account | +| **Claude (Anthropic)** | Blocked by ToS | `api.anthropic.com` | Claude Pro/Max subscription | + +> **Note on Claude OAuth:** Anthropic explicitly banned OAuth token usage in +> third-party tools as of January 2026. Their updated "Authentication and credential +> use" policy states that OAuth tokens from Free, Pro, and Max plans are authorized +> **exclusively for Claude Code and Claude.ai**. This provider cannot be implemented +> until Anthropic changes their policy. See Section 3.5. + +> **Note on Antigravity ToS:** There are reports of Google blocking accounts using +> third-party Antigravity auth plugins. The Antigravity ToS (as of 2026-02-18) states +> their service cannot be used with third-party products. This risk is documented in +> Section 10 but does not block implementation. + +--- + +## 1. OAuth Providers — Technical Details + +### 1.1 Codex OAuth (OpenAI) + +OpenAI's OAuth 2.0 + PKCE flow used by the Codex CLI to authenticate ChatGPT +subscribers. + +| Parameter | Value | +|------------------------|-------------------------------------------------------| +| Authorization endpoint | `https://auth.openai.com/authorize` | +| Token endpoint | `https://auth.openai.com/oauth/token` | +| Client ID | `app_EMoamEEZ73f0CkXaXp7hrann` | +| Redirect URI | `http://localhost:/auth/callback` | +| Scopes | `openid profile email offline_access` | +| PKCE method | S256 | +| Token refresh interval | ~8 days | +| Device code flow | Supported | +| Token injection | `Authorization: Bearer ` | +| API format | OpenAI-native (pass-through, no body transformation) | + +**Tokens produced:** access token (short-lived JWT), refresh token (long-lived, +single-use), ID token (user identity claims). + +### 1.2 Antigravity OAuth (Google) + +Google's OAuth 2.0 + PKCE flow used by the Antigravity IDE to authenticate Google +account holders. + +| Parameter | Value | +|------------------------|--------------------------------------------------------------| +| Authorization endpoint | `https://accounts.google.com/o/oauth2/auth` | +| Token endpoint | `https://oauth2.googleapis.com/token` | +| Client ID | Antigravity OAuth client (configurable) | +| Redirect URI | `http://localhost:/oauth-callback` | +| Scopes | `openid profile email https://www.googleapis.com/auth/cloud-platform` | +| Additional scopes | `auth/cclog`, `auth/experimentsandconfigs` | +| Access type | `offline` (enables refresh token) | +| PKCE method | S256 | +| Prompt | `consent` (forces consent screen) | +| Device code flow | No — headless fallback via copy/paste URL | +| Token injection | `Authorization: Bearer ` + extra headers | +| API format | Gemini-style (requires request body wrapping) | + +**Tokens produced:** access token, refresh token, with associated project_id and +account metadata. + +**API Endpoints (with fallback):** + +| Tier | Base URL | +|-------------|----------------------------------------------------------------| +| Production | `https://cloudcode-pa.googleapis.com` | +| Daily | `https://daily-cloudcode-pa.sandbox.googleapis.com` | +| Alt Prod | `https://codeassist.googleapis.com/v1` | + +**Required Headers (beyond Bearer token):** +- `X-Goog-Api-Client: google-cloud-sdk vscode_cloudshelleditor/0.1` +- `Client-Metadata: {"ideType":"ANTIGRAVITY","platform":"","pluginType":"GEMINI"}` +- `User-Agent: antigravity/1.15.8 ` + +**Key paths:** +- Generate content: `/v1internal:generateContent` +- Streaming: `/v1internal:streamGenerateContent?alt=sse` + +**Token storage fields:** `email`, `accessToken`, `refreshToken`, `expiresAt`, +`projectId`, `tier`, `rateLimitedUntil`, `lastUsed`. + +**Available models:** Gemini 3 Pro/Flash, Claude Sonnet 4.6, Claude Opus 4.6 +(Thinking), GPT-OSS 120B. + +### 1.3 Key Differences Between Providers + +| Aspect | Codex (OpenAI) | Antigravity (Google) | +|-----------------------|-----------------------------------|------------------------------------------| +| Auth server | `auth.openai.com` | `accounts.google.com` | +| Token endpoint | `auth.openai.com/oauth/token` | `oauth2.googleapis.com/token` | +| API request format | OpenAI-native | Gemini-style (wrapped body) | +| Extra headers needed | None | `X-Goog-Api-Client`, `Client-Metadata` | +| Upstream routing | Single endpoint | 3-tier fallback | +| Project ID | Not required | Required (per-account `projectId`) | +| Headless support | Device code (interactive) | Copy/paste URL (manual) | +| Multi-account | Single account | Up to 10 accounts with rotation | +| Token refresh check | Background (75% TTL) | Pre-request (60s before expiry) | +| Request body changes | Pass-through | Wrap with project metadata | + +### 1.4 Claude OAuth (Anthropic) — Blocked + +**Not implementable.** Anthropic deployed a technical block on January 9, 2026 +rejecting all OAuth tokens from non-Claude-Code clients. Policy formalized +~February 17-18, 2026. Error: "This credential is only authorized for use with +Claude Code." + +--- + +## 2. Why Add Multi-Provider OAuth? + +| Benefit | Detail | +|---------------------------------|-----------------------------------------------------------------| +| No API key required | Users with subscriptions can use their existing accounts | +| Broader user base | Many users have subscriptions but not API keys | +| Better security | OAuth tokens are short-lived and revocable vs. static keys | +| Multi-model access | Antigravity gives access to Gemini, Claude, and GPT-OSS models | +| Cost savings | Subscription usage avoids separate per-token API charges | + +--- + +## 3. Feasibility Assessment + +### 3.1 Compatible — Low Risk + +| Aspect | Why it works | +|----------------------|--------------------------------------------------------------------| +| HTTP proxy model | ClawShell already intercepts and rewrites auth headers | +| Provider abstraction | `Provider` enum and `ProxyClient` already branch on provider type | +| Config system | TOML config is extensible — add `[[oauth_providers]]` table | +| Onboard wizard | Already has a provider selection menu — just add new entries | +| Rust ecosystem | `oauth2` crate handles PKCE; `open` for browser; `reqwest` for API | +| Daemon architecture | Token refresh can run as background `tokio` tasks | + +### 3.2 Challenges + +| Challenge | Mitigation | +|------------------------------------|----------------------------------------------------------------| +| Browser needed for initial login | Headless fallback for both providers | +| Token storage security | Store in `/etc/clawshell/oauth/` with 0600 perms | +| Token refresh in a daemon | Background tokio task per provider; proactive refresh | +| Provider-specific API formats | Trait-based abstraction with `prepare_request()` per provider | +| Antigravity request body wrapping | `AntigravityProvider` handles Gemini-style body transformation | +| Client ID stability | All OAuth parameters configurable per provider | +| ToS restrictions | Monitor each provider's policy; disable if 3P banned | + +### 3.3 Open Questions — Codex (v1) + +1. Does OpenAI's API accept ChatGPT OAuth access tokens on the standard + `/v1/chat/completions` and `/v1/responses` endpoints? +2. Are there rate limits or model restrictions specific to OAuth-authenticated requests? +3. Is the Codex client ID (`app_EMoamEEZ73f0CkXaXp7hrann`) stable for third-party use? + +### 3.4 Open Questions — Antigravity (v1) + +1. Does the Antigravity API return standard error codes or Google-specific ones? +2. What is the exact token TTL (for refresh scheduling)? +3. What `projectId` assignment flow is needed on first login? +4. Are there per-account rate limits beyond what the plugin documents? + +### 3.5 Claude OAuth — Status + +**Not implementable.** Anthropic deployed a technical block on January 9, 2026. +Policy updated ~February 17-18, 2026. There is no known workaround. + +--- + +## 4. Architecture + +### 4.1 Current Flow (API Key Only) + +``` +OpenClaw ──► ClawShell (virtual key → real API key) ──► OpenAI / Anthropic API +``` + +### 4.2 Proposed Flow (Multi-Provider OAuth) + +``` + clawshell onboard + ┌──────────────────────────────────────────┐ + │ │ + │ Select a model provider: │ + │ 1. OpenAI (API key) │ ← existing + │ 2. OpenRouter (API key) │ ← existing + │ 3. Anthropic (API key) │ ← existing + │ 4. Codex / ChatGPT (OAuth login) │ ← NEW + │ 5. Antigravity / Google (OAuth login) │ ← NEW + │ │ + │ If 1-3: prompt for API key (unchanged) │ + │ If 4: open browser → auth.openai.com │ + │ If 5: open browser → accounts.google │ + │ │ + └──────────────────────────────────────────┘ + + RUNTIME +┌──────────┐ Bearer vk-001 ┌───────────────────────────┐ +│ │ ──────────────────► │ ClawShell │ +│ OpenClaw │ │ │ +│ │ ◄────────────────── │ Lookup vk-001 → KeySource │ +└──────────┘ response │ │ │ + │ ┌────┴────┐ │ + │ Static OAuth{id} │ + │ │ │ │ + │ ▼ ▼ │ + │ real_key OAuthRegistry │ + │ │ ┌──────────────┐ │ + │ │ │ provider_id? │ │ + │ │ │ codex ──────►│───│──► api.openai.com + │ │ │ antigravity─►│───│──► cloudcode-pa.googleapis.com + │ │ └──────────────┘ │ + │ └────┬────┘ │ + │ ▼ │ + │ Forward to upstream │ + └─────────────────────────────┘ + + BACKGROUND + ┌───────────────────────────────┐ + │ Refresh Task: codex │ sleep(75% of TTL) + ├───────────────────────────────┤ + │ Refresh Task: antigravity │ check 60s before expiry + └───────────────────────────────┘ +``` + +### 4.3 Module Map + +``` +src/ +├── oauth/ +│ ├── mod.rs ← NEW: OAuthProvider trait, OAuthRegistry, shared types +│ ├── codex.rs ← NEW: Codex (OpenAI) provider [v1] +│ ├── antigravity.rs ← NEW: Antigravity (Google) provider [v1] +│ └── storage.rs ← NEW: Per-provider token persistence +├── lib.rs ← MODIFY: AppState gains OAuthRegistry +├── cli.rs ← UNCHANGED (no new subcommands) +├── config.rs ← MODIFY: add [[oauth_providers]] config section +├── keys.rs ← MODIFY: ResolvedKey gains OAuth{provider_id} +├── proxy.rs ← MODIFY: provider.inject_auth() + prepare_request() + 401-retry +├── main.rs ← MODIFY: initialize OAuthRegistry, start refresh tasks +├── onboard/ +│ ├── interactive.rs ← MODIFY: add OAuth providers to menu, run OAuth flow +│ ├── types.rs ← MODIFY: OnboardConfig supports OAuth auth method +│ ├── config_render.rs ← MODIFY: render [[oauth_providers]] + auth="oauth" in TOML +│ └── (rest unchanged) +└── (rest unchanged) +``` + +--- + +## 5. Detailed Design + +### 5.1 The `OAuthProvider` Trait + +The core abstraction enabling multiple providers: + +```rust +#[async_trait] +pub trait OAuthProvider: Send + Sync + std::fmt::Debug { + /// Unique identifier (e.g., "codex", "antigravity"). + fn id(&self) -> &str; + + /// Display name (e.g., "Codex (OpenAI)", "Antigravity (Google)"). + fn display_name(&self) -> &str; + + /// Execute browser-based OAuth login flow. + async fn login_browser(&self, callback_port: u16) -> Result; + + /// Execute headless login flow (device code or copy/paste URL). + async fn login_headless(&self) -> Result; + + /// Refresh the access token using the refresh token. + async fn refresh(&self, refresh_token: &str) -> Result; + + /// Inject provider-specific auth headers into the request. + fn inject_auth(&self, headers: &mut HeaderMap, access_token: &str) -> Result<(), OAuthError>; + + /// Optionally transform the request body for provider-specific formats. + /// Returns None for pass-through (Codex); Some(wrapped) for Antigravity. + fn prepare_request_body( + &self, body: &[u8], tokens: &OAuthTokens, + ) -> Result>, OAuthError> { + let _ = (body, tokens); + Ok(None) + } + + /// Resolve the upstream URL for this provider. + /// Returns None to use the configured [upstream] URL (Codex). + fn upstream_url(&self, tokens: &OAuthTokens) -> Option { + let _ = tokens; + None + } + + /// Whether this provider supports device code flow. + fn supports_device_code(&self) -> bool { false } + + /// Whether this provider supports headless copy/paste URL fallback. + fn supports_headless_url(&self) -> bool { false } +} +``` + +### 5.2 Codex Provider (`codex.rs`) + +```rust +#[derive(Debug)] +pub struct CodexProvider { + client_id: String, + auth_url: String, + token_url: String, + scopes: Vec, + http_client: reqwest::Client, +} + +impl OAuthProvider for CodexProvider { + fn id(&self) -> &str { "codex" } + fn display_name(&self) -> &str { "Codex (OpenAI)" } + fn supports_device_code(&self) -> bool { true } + + fn inject_auth(&self, headers: &mut HeaderMap, token: &str) -> Result<(), OAuthError> { + headers.insert(AUTHORIZATION, format!("Bearer {}", token).parse()?); + Ok(()) + } + // prepare_request_body: default (None — pass-through) + // upstream_url: default (None — use [upstream].base_url) +} +``` + +### 5.3 Antigravity Provider (`antigravity.rs`) + +```rust +#[derive(Debug)] +pub struct AntigravityProvider { + client_id: String, + auth_url: String, + token_url: String, + scopes: Vec, + http_client: reqwest::Client, + endpoints: Vec, +} + +impl OAuthProvider for AntigravityProvider { + fn id(&self) -> &str { "antigravity" } + fn display_name(&self) -> &str { "Antigravity (Google)" } + fn supports_headless_url(&self) -> bool { true } + + fn inject_auth(&self, headers: &mut HeaderMap, token: &str) -> Result<(), OAuthError> { + headers.insert(AUTHORIZATION, format!("Bearer {}", token).parse()?); + headers.insert("x-goog-api-client", + "google-cloud-sdk vscode_cloudshelleditor/0.1".parse()?); + headers.insert("client-metadata", + r#"{"ideType":"ANTIGRAVITY","platform":"LINUX","pluginType":"GEMINI"}"#.parse()?); + Ok(()) + } + + fn prepare_request_body(&self, body: &[u8], tokens: &OAuthTokens) + -> Result>, OAuthError> { + let project_id = tokens.extra.get("project_id") + .and_then(|v| v.as_str()) + .ok_or(OAuthError::LoginFailed("Missing project_id".into()))?; + let wrapped = wrap_antigravity_request(body, project_id)?; + Ok(Some(wrapped)) + } + + fn upstream_url(&self, _tokens: &OAuthTokens) -> Option { + Some(format!("{}/v1internal:streamGenerateContent?alt=sse", + self.endpoints.last().unwrap_or(&self.endpoints[0]))) + } +} +``` + +### 5.4 `OAuthRegistry` + +```rust +#[derive(Debug)] +pub struct OAuthRegistry { + providers: BTreeMap>, + tokens: Arc>>, + storage: TokenStorage, +} + +impl OAuthRegistry { + pub async fn current_access_token(&self, provider_id: &str) -> Result; + pub async fn inject_auth(&self, provider_id: &str, headers: &mut HeaderMap) + -> Result<(), OAuthError>; + pub async fn prepare_request_body(&self, provider_id: &str, body: &[u8]) + -> Result>, OAuthError>; + pub async fn upstream_url(&self, provider_id: &str) -> Result, OAuthError>; + pub async fn refresh(&self, provider_id: &str) -> Result<(), OAuthError>; + pub fn spawn_refresh_tasks(&self, cancel: CancellationToken); +} +``` + +### 5.5 Token Storage (`storage.rs`) + +Per-provider token files under `/etc/clawshell/oauth/`: + +``` +/etc/clawshell/oauth/ +├── codex.json +│ { "access_token": "...", "refresh_token": "...", "expires_at": "...", +│ "account_id": "...", "extra": {} } +└── antigravity.json + { "access_token": "...", "refresh_token": "...", "expires_at": "...", + "account_id": "user@gmail.com", + "extra": { "project_id": "...", "tier": "...", "email": "..." } } +``` + +All files mode 0600, directory mode 0700, owned by `clawshell` user. + +### 5.6 Onboarding Changes (`onboard/interactive.rs`) + +The existing provider menu at line 326-331 currently shows: + +```rust +let provider_options = match existing.provider.as_deref() { + Some("anthropic") => vec!["Anthropic", "OpenAI", "OpenRouter"], + Some("openrouter") => vec!["OpenRouter", "OpenAI", "Anthropic"], + _ => vec!["OpenAI", "OpenRouter", "Anthropic"], +}; +``` + +This becomes: + +```rust +let provider_options = match existing.provider.as_deref() { + Some("anthropic") => vec!["Anthropic", "OpenAI", "OpenRouter", + "Codex / ChatGPT (OAuth)", "Antigravity / Google (OAuth)"], + Some("codex") => vec!["Codex / ChatGPT (OAuth)", "OpenAI", "OpenRouter", + "Anthropic", "Antigravity / Google (OAuth)"], + Some("antigravity") => vec!["Antigravity / Google (OAuth)", "OpenAI", "OpenRouter", + "Anthropic", "Codex / ChatGPT (OAuth)"], + // ... + _ => vec!["OpenAI", "OpenRouter", "Anthropic", + "Codex / ChatGPT (OAuth)", "Antigravity / Google (OAuth)"], +}; +``` + +**After provider selection**, the flow branches: + +``` +If provider is "openai" | "openrouter" | "anthropic": + → existing flow: prompt for API key, model, virtual key, etc. + +If provider is "codex": + → prompt for model name (default: models available via ChatGPT) + → detect headless environment (SSH_CONNECTION, etc.) + → if headless: run device code flow + → else: run browser PKCE flow → auth.openai.com + → store tokens to /etc/clawshell/oauth/codex.json + → prompt for virtual key + → continue with OpenClaw config, server settings, etc. + +If provider is "antigravity": + → prompt for model name (default: gemini-3-pro) + → detect headless environment + → if headless: print auth URL, prompt for redirect URL paste + → else: run browser PKCE flow → accounts.google.com + → discover project_id via loadCodeAssist + → store tokens to /etc/clawshell/oauth/antigravity.json + → prompt for virtual key + → continue with OpenClaw config, server settings, etc. +``` + +**The real API key prompt is skipped entirely for OAuth providers.** The +`OnboardConfig` struct changes: + +```rust +pub enum AuthMethod { + ApiKey { real_api_key: String }, + OAuth { provider_id: String }, // tokens already stored by the onboard flow +} + +pub struct OnboardConfig { + pub provider: String, + pub model: String, + pub auth: AuthMethod, // was: pub real_api_key: String + pub virtual_api_key: String, + pub openclaw_config_path: PathBuf, + pub server_host: String, + pub server_port: u16, + pub email: Option, +} +``` + +**Re-onboard behavior:** When a user runs `clawshell onboard` again and a previous +OAuth config exists, the wizard detects it (from `config.json` and the presence of +token files) and offers to re-authenticate or keep the existing tokens. + +### 5.7 Config Rendering Changes (`onboard/config_render.rs`) + +When the user selects an OAuth provider, the generated `clawshell.toml` includes: + +```toml +[[keys]] +virtual_key = "vk-chatgpt-001" +provider = "openai" +auth = "oauth" +oauth_provider = "codex" + +[[oauth_providers]] +provider = "codex" +``` + +And `config.json` stores `"provider": "codex"` (or `"antigravity"`) for re-onboard +detection. + +### 5.8 Config Changes (`config.rs`) + +```rust +pub struct Config { + pub server: ServerConfig, + pub upstream: UpstreamConfig, + pub keys: Vec, + pub dlp: DlpConfig, + pub log_level: String, + #[serde(default)] + pub oauth_providers: Vec, +} + +pub struct KeyMapping { + pub virtual_key: String, + pub real_key: Option, // optional when auth = "oauth" + pub provider: Provider, + #[serde(default)] + pub auth: AuthMethod, // defaults to Static + pub oauth_provider: Option, // "codex" or "antigravity" +} + +#[derive(Default)] +pub enum AuthMethod { #[default] Static, OAuth } + +pub struct OAuthProviderConfig { + pub provider: String, + #[serde(default = "default_true")] + pub enabled: bool, + pub client_id: Option, + pub auth_url: Option, + pub token_url: Option, + pub scopes: Option>, + pub callback_port: Option, +} +``` + +### 5.9 Key Resolution Changes (`keys.rs`) + +```rust +pub enum KeySource { + Static(String), + OAuth { provider_id: String }, +} + +pub struct ResolvedKey { + pub source: KeySource, + pub provider: Provider, +} +``` + +### 5.10 Proxy Changes (`proxy.rs`) + +```rust +match resolved.source { + KeySource::Static(ref key) => { + // existing logic + } + KeySource::OAuth { ref provider_id } => { + oauth_registry.inject_auth(provider_id, &mut req_headers).await?; + let body = match oauth_registry.prepare_request_body(provider_id, &body).await? { + Some(transformed) => Bytes::from(transformed), + None => body, + }; + let upstream = match oauth_registry.upstream_url(provider_id).await? { + Some(url) => url, + None => default_upstream_url(provider), + }; + // send, handle 401 → refresh + retry once + } +} +``` + +### 5.11 AppState Changes (`lib.rs`) + +```rust +pub struct AppState { + pub key_manager: Arc, + pub dlp_scanner: Arc, + pub proxy_client: Arc, + pub oauth_registry: Option>, +} +``` + +--- + +## 6. New Dependencies + +| Crate | Purpose | Size Impact | +|---------------|----------------------------------------------|-------------| +| `oauth2` | OAuth 2.0 client with PKCE support | Moderate | +| `open` | Open browser for auth URL (cross-platform) | Tiny | +| `chrono` | Token expiry math | Small | +| `base64` | PKCE verifier encoding (may be transitive) | Tiny | +| `async-trait` | Trait async methods (if not Rust 1.85+) | Small | + +--- + +## 7. Security Considerations + +| Concern | Mitigation | +|----------------------------|------------------------------------------------------------------| +| Token files on disk | Per-provider files in `/etc/clawshell/oauth/` with 0600 perms | +| Token in memory | `Arc>` — same threat model as current keys | +| Refresh token theft | Single-use rotation (Codex); standard rotation (Antigravity) | +| PKCE | Both providers use S256 — prevents code interception | +| Callback server exposure | `127.0.0.1` only; ephemeral; shuts down after one use | +| Provider client IDs | Configurable per provider in `[[oauth_providers]]` | +| ToS compliance | Monitor each provider's policy; documented risks | +| Provider isolation | Separate token files — compromise of one doesn't affect others | +| Antigravity extra headers | Injected server-side; client never sees them | + +--- + +## 8. Testing Strategy + +| Layer | Approach | +|--------------|-------------------------------------------------------------------| +| Unit | Mock `OAuthProvider` trait impls; test PKCE generation | +| Unit | Test `CodexProvider` and `AntigravityProvider` independently | +| Unit | Test `OAuthRegistry` with mock providers | +| Unit | Test `TokenStorage` with temp directories | +| Unit | Test Antigravity request body wrapping | +| Integration | `wiremock`: mock `auth.openai.com` for Codex | +| Integration | `wiremock`: mock `oauth2.googleapis.com` for Antigravity | +| Config | Snapshot tests for TOML with Codex / Antigravity / both / none | +| Onboard | Test `OnboardConfig` generation for OAuth vs API key paths | +| E2E | Manual: `clawshell onboard` → select Codex → proxy request | +| E2E | Manual: `clawshell onboard` → select Antigravity → proxy request | +| Existing | All existing tests must pass (OAuth is opt-in) | + +--- + +## 9. Implementation Phases + +### Phase 1: OAuth Framework (Medium Effort) +1. Add `oauth2`, `open`, `chrono` dependencies to `Cargo.toml`. +2. Create `src/oauth/mod.rs` — `OAuthProvider` trait, `OAuthTokens`, `OAuthError`. +3. Create `src/oauth/storage.rs` — per-provider token persistence. +4. Create `OAuthRegistry` with provider registration, token management, refresh tasks. +5. Unit tests with mock providers. + +### Phase 2: Codex Provider (Medium Effort) +1. Create `src/oauth/codex.rs` — browser PKCE flow + device code flow. +2. Implement `inject_auth()` (Bearer token). +3. Unit + integration tests with `wiremock`. + +### Phase 3: Antigravity Provider (Medium Effort) +1. Create `src/oauth/antigravity.rs` — browser PKCE flow + headless fallback. +2. Implement `inject_auth()` (Bearer + Google-specific headers). +3. Implement `prepare_request_body()` (Gemini-style wrapping). +4. Implement `upstream_url()` (endpoint resolution). +5. Implement project ID discovery via `loadCodeAssist`. +6. Unit + integration tests. + +### Phase 4: Config & Key Integration (Small Effort) +1. Add `[[oauth_providers]]` to `config.rs`. +2. Add `AuthMethod` enum and `oauth_provider` field to `KeyMapping`. +3. Extend `ResolvedKey` / `KeySource` in `keys.rs`. +4. Update `proxy.rs` — dispatch to `inject_auth()` + `prepare_request_body()` + 401-retry. +5. Wire `OAuthRegistry` into `AppState` in `lib.rs`. + +### Phase 5: Onboarding Integration (Medium Effort) +1. Add "Codex / ChatGPT (OAuth)" and "Antigravity / Google (OAuth)" to provider menu + in `onboard/interactive.rs`. +2. Add OAuth login flow branch (skip API key prompt, run browser/headless flow). +3. Update `OnboardConfig` with `AuthMethod` enum in `onboard/types.rs`. +4. Update `config_render.rs` to generate `[[oauth_providers]]` and `auth = "oauth"`. +5. Handle re-onboard detection (existing OAuth tokens). + +### Phase 6: Documentation & Polish +1. Update README. +2. Update example config. +3. Document ToS considerations per provider. + +--- + +## 10. Risks and Mitigations + +| Risk | Likelihood | Impact | Mitigation | +|------------------------------------------------|------------|--------|-----------------------------------------------| +| OpenAI changes Codex client ID | Medium | High | Make configurable; monitor changes | +| Google blocks Antigravity 3P usage | High | High | Make configurable; document risk; feature flag | +| OAuth tokens rejected on standard endpoints | Low | High | Verify during Phase 2/3; abort if incompatible | +| Antigravity API format changes | Medium | Medium | Version-pin User-Agent; test against live API | +| Rate limits differ for OAuth vs. API key | Medium | Medium | Document limitation; let users choose method | +| Token refresh fails silently | Low | Medium | Aggressive logging; prompt re-onboard | +| Anthropic maintains Claude OAuth ban | Very High | Low | Already accounted for — not implementing | + +--- + +## 11. Backward Compatibility + +Fully opt-in. No `[[oauth_providers]]` = identical behavior to today. The existing +provider options (OpenAI, OpenRouter, Anthropic) in `clawshell onboard` work exactly +as before. No new CLI subcommands — no change to the command interface. + +--- + +## 12. Decisions Required + +1. **Codex endpoint compatibility:** Verify ChatGPT OAuth tokens work on standard OpenAI API. +2. **Codex client ID policy:** Use Codex CLI's client ID or register our own? +3. **Antigravity client ID:** Use the known Antigravity client ID or register? +4. **Antigravity body transformation:** Should ClawShell translate OpenAI-format to + Gemini-style, or require clients to send Gemini-format directly? +5. **ToS risk acceptance:** Proceed with documented risk, or defer Antigravity? diff --git a/docs/docker.md b/docs/docker.md new file mode 100644 index 0000000..0a501c2 --- /dev/null +++ b/docs/docker.md @@ -0,0 +1,193 @@ +# Running ClawShell in Docker + +## Build + +Create a `.env` file in the project root with your credentials (required for +Antigravity/Google OAuth): + +``` +GOOGLE_OAUTH_CLIENT_ID=your-client-id.apps.googleusercontent.com +GOOGLE_OAUTH_CLIENT_SECRET=your-client-secret +``` + +Then build the release binary and Docker image: + +```bash +cargo build --release +docker build -t clawshell . +``` + +The `.env` file is baked into the image at `/etc/clawshell/.env` and loaded +automatically at runtime — no need to pass `--env-file` or `-e` flags. + +> **Security note:** The `.env` file is embedded in the image. Do not push +> the image to a public registry if it contains sensitive credentials. + +## Onboarding + +Run the interactive onboard wizard to generate configuration: + +```bash +docker run --rm -it clawshell onboard +``` + +If you're onboarding with Antigravity/Google OAuth, publish the callback port: + +```bash +docker run --rm -it \ + -p 51121:51121 \ + clawshell onboard +``` + +This creates the configuration files inside the container. To persist them, +mount a volume for `/etc/clawshell`: + +```bash +docker run --rm -it -v clawshell-config:/etc/clawshell clawshell onboard +``` + +The wizard will prompt you to select a provider (OpenAI, OpenRouter, Anthropic, +Codex/ChatGPT OAuth, or Antigravity/Google OAuth), a model, and an API key or +OAuth login. + +## Running the proxy + +Start ClawShell in the foreground with the persisted configuration: + +```bash +docker run -d \ + --name clawshell \ + -p 18790:18790 \ + -v clawshell-config:/etc/clawshell \ + clawshell start --foreground +``` + +The proxy listens on port `18790` by default. The `--foreground` flag is +required in Docker (no daemonization). + +### Binding to all interfaces + +By default ClawShell listens on `127.0.0.1`, which is unreachable from outside +the container. Set the host to `0.0.0.0` in your `clawshell.toml`: + +```toml +[server] +host = "0.0.0.0" +port = 18790 +``` + +Or pass it during onboard when prompted for the server host. + +You can also override server bind host/port at runtime: + +```bash +docker run -d \ + --name clawshell \ + -p 17890:17890 \ + -e CLAWSHELL_SERVER_HOST=0.0.0.0 \ + -e CLAWSHELL_SERVER_PORT=17890 \ + -v clawshell-config:/etc/clawshell \ + clawshell start --foreground +``` + +## Configuration volume + +All ClawShell state lives under `/etc/clawshell`: + +| Path | Purpose | +|---------------------------------|--------------------------------------| +| `/etc/clawshell/clawshell.toml` | Main configuration file | +| `/etc/clawshell/config.json` | Onboard metadata | +| `/etc/clawshell/oauth/` | OAuth token files (0600 perms) | +| `/etc/clawshell/.env` | Google OAuth credentials (from build)| + +Use a named volume (`clawshell-config`) or a bind mount to persist these across +container restarts. + +## Environment variables + +The `.env` file is copied into the image at build time and loaded automatically. +You can also override values at runtime if needed: + +```bash +docker run --rm -it \ + -e GOOGLE_OAUTH_CLIENT_ID=different-id.apps.googleusercontent.com \ + clawshell onboard +``` + +Runtime `-e` flags take precedence over the baked-in `.env` file. + +### Runtime server bind overrides + +| Variable | Description | +|--------------------------|-------------------------------------------| +| `CLAWSHELL_SERVER_HOST` | Overrides `[server].host` (e.g. `0.0.0.0`) | +| `CLAWSHELL_SERVER_PORT` | Overrides `[server].port` (e.g. `17890`) | + +### Required variables for Antigravity / Google OAuth + +| Variable | Description | +|-------------------------------|----------------------------| +| `GOOGLE_OAUTH_CLIENT_ID` | Google OAuth client ID | +| `GOOGLE_OAUTH_CLIENT_SECRET` | Google OAuth client secret | + +These are not needed for other providers (OpenAI, OpenRouter, Anthropic, Codex). + +## OAuth providers + +### Codex / ChatGPT (OAuth) + +Uses device code flow — no browser required inside the container. The wizard +prints a URL and a one-time code. Open the URL on any device, enter the code, +and the container receives the tokens automatically. + +No extra environment variables are needed for Codex. + +### Antigravity / Google (OAuth) + +Requires `GOOGLE_OAUTH_CLIENT_ID` and `GOOGLE_OAUTH_CLIENT_SECRET` (provided +via the `.env` file baked into the image at build time). + +Uses a localhost callback flow on port `51121`. During onboarding: + +1. Run the container with `-p 51121:51121`. +2. Open the printed Google authorization URL. +3. Complete consent; Google redirects to `http://localhost:51121/oauth-callback...`. + +The `Dockerfile` sets `CLAWSHELL_OAUTH_CALLBACK_HOST=0.0.0.0` so the callback +listener inside the container accepts the published port. + +## Stopping + +```bash +docker stop clawshell +``` + +## Example: full setup + +```bash +# 1. Create .env with Google OAuth credentials (skip if not using Antigravity) +cat > .env << 'EOF' +GOOGLE_OAUTH_CLIENT_ID=your-client-id.apps.googleusercontent.com +GOOGLE_OAUTH_CLIENT_SECRET=your-client-secret +EOF + +# 2. Build +cargo build --release +docker build -t clawshell . + +# 3. Onboard (interactive — creates config in the volume) +# Add -p 51121:51121 if using Antigravity/Google OAuth. +docker run --rm -it -v clawshell-config:/etc/clawshell -p 51121:51121 clawshell onboard + +# 4. Run +docker run -d \ + --name clawshell \ + --restart unless-stopped \ + -p 18790:18790 \ + -v clawshell-config:/etc/clawshell \ + clawshell start --foreground + +# 5. Verify +curl http://localhost:18790/health +``` diff --git a/src/app.rs b/src/app.rs index 45a4f1d..d7fc640 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1,17 +1,18 @@ -use crate::config::{Config, Provider}; +use crate::config::{Config, KeyAuthMethod, Provider}; use crate::dlp::DlpScanner; use crate::email::{ EmailAccountCredentials, EmailGetMessageRequest, EmailListMessagesRequest, EmailMessageContent, EmailMessageMetadata, EmailPolicy, EmailService, EmailServiceError, ImapEmailService, normalize_sender_rule, }; -use crate::keys::{KeyManager, ResolvedKey}; +use crate::keys::{KeyManager, KeySource, ResolvedKey}; +use crate::oauth::OAuthRegistry; use crate::proxy::ProxyClient; use axum::Router; use axum::body::Body; use axum::extract::{DefaultBodyLimit, Path, Query, Request, State}; -use axum::http::StatusCode; +use axum::http::{HeaderMap, Method, StatusCode, Uri}; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; use axum::routing::{any, get}; @@ -28,6 +29,7 @@ pub struct AppState { pub key_manager: Arc, pub dlp_scanner: Arc, pub proxy_client: Arc, + pub oauth_registry: Arc, pub email_enabled: bool, pub email_policy: Option, pub email_accounts: Arc>, @@ -35,7 +37,15 @@ pub struct AppState { } impl AppState { + #[allow(dead_code)] pub fn from_config(config: &Config) -> Result { + Self::from_config_with_registry(config, None) + } + + pub fn from_config_with_registry( + config: &Config, + oauth_registry: Option, + ) -> Result { let mut upstream_urls = BTreeMap::new(); upstream_urls.insert(Provider::Openai, config.upstream_url(Provider::Openai)); upstream_urls.insert( @@ -47,19 +57,29 @@ impl AppState { config.upstream_url(Provider::Anthropic), ); - let key_mappings = config - .key_map() - .iter() - .map(|(virtual_key, (real_key, provider))| { - ( - virtual_key.clone(), - ResolvedKey { - real_key: real_key.clone(), - provider: *provider, - }, - ) - }) - .collect(); + // Build key mappings for both static and OAuth keys + let mut key_mappings: BTreeMap = BTreeMap::new(); + + for key in &config.keys { + let source = match key.auth { + KeyAuthMethod::Static => KeySource::Static { + real_key: key.real_key.clone().unwrap_or_default(), + }, + KeyAuthMethod::OAuth => KeySource::OAuth { + provider_id: key.oauth_provider.clone().unwrap_or_default(), + }, + }; + key_mappings.insert( + key.virtual_key.clone(), + ResolvedKey { + source, + provider: key.provider, + }, + ); + } + + let oauth_registry = + oauth_registry.unwrap_or_else(|| OAuthRegistry::new(Default::default())); let email_policy = if config.email.enabled { config.email.mode.map(|mode| { @@ -112,6 +132,7 @@ impl AppState { upstream_urls, config.upstream.anthropic_version.clone(), )), + oauth_registry: Arc::new(oauth_registry), email_enabled: config.email.enabled, email_policy, email_accounts: Arc::new(email_accounts), @@ -493,7 +514,7 @@ async fn handle_request( ); error_response(StatusCode::UNAUTHORIZED, "Unknown API key") })?; - let real_key = resolved.real_key.clone(); + let source = resolved.source.clone(); let provider = resolved.provider; debug!( @@ -563,27 +584,54 @@ async fn handle_request( "Forwarding request to upstream" ); - let response = state - .proxy_client - .forward( - method.clone(), - &uri, - headers, - &real_key, - body_bytes, - provider, - ) - .await - .map_err(|e| { - error!( - method = %method, - path = %path, - virtual_key = %virtual_key, - error = %e, - "Proxy error" - ); - e.into_response() - })?; + let response = match source { + KeySource::Static { real_key } => { + state + .proxy_client + .forward( + method.clone(), + &uri, + headers, + &real_key, + body_bytes, + provider, + ) + .await + .map_err(|e| { + error!( + method = %method, + path = %path, + virtual_key = %virtual_key, + error = %e, + "Proxy error" + ); + e.into_response() + })? + } + KeySource::OAuth { provider_id } => { + forward_oauth_request( + &state, + method.clone(), + &uri, + headers, + body_bytes, + provider, + &provider_id, + ) + .await + .map_err(|e| { + error!( + method = %method, + path = %path, + virtual_key = %virtual_key, + oauth_provider = %provider_id, + error = %e, + "OAuth proxy error" + ); + error_response(StatusCode::BAD_GATEWAY, &format!("OAuth proxy error: {e}")) + })? + } + }; // 5. DLP scan on response body (redact all PII before returning to client) let response = if state.dlp_scanner.scan_responses() { @@ -627,14 +675,18 @@ async fn handle_request( Response::from_parts(parts, Body::from(body)) } } else { - warn!( + debug!( method = %method, path = %path, virtual_key = %virtual_key, - "Streaming response (SSE) — DLP scanning is not supported for streaming responses; \ - PII in streamed content will not be redacted" + "Streaming response (SSE) — wrapping with DLP SSE scanner" ); - response + let (parts, body) = response.into_parts(); + let dlp_body = crate::translate::wrap_body_with_dlp_sse_stream( + body, + state.dlp_scanner.clone(), + ); + Response::from_parts(parts, dlp_body) } } else { trace!("Response DLP scanning disabled"); @@ -644,6 +696,260 @@ async fn handle_request( Ok(response) } +async fn forward_oauth_request( + state: &AppState, + method: Method, + uri: &Uri, + headers: HeaderMap, + body_bytes: Bytes, + provider: Provider, + oauth_provider_id: &str, +) -> Result { + // 1. Inject auth headers + let mut auth_headers = HeaderMap::new(); + state + .oauth_registry + .inject_auth(oauth_provider_id, &mut auth_headers) + .await + .map_err(|e| format!("OAuth auth injection failed: {e}"))?; + + // 2. Optionally transform the body + let body = match state + .oauth_registry + .prepare_request_body(oauth_provider_id, &body_bytes) + .await + .map_err(|e| format!("OAuth body preparation failed: {e}"))? + { + Some(transformed) => Bytes::from(transformed), + None => body_bytes.clone(), + }; + + // 2b. Check if path needs rewriting (e.g., /v1/chat/completions → /v1/responses) + let original_path = uri.path().to_string(); + let rewritten_path = state + .oauth_registry + .rewrite_request_path(oauth_provider_id, &original_path) + .map_err(|e| format!("OAuth path rewrite failed: {e}"))?; + let needs_translation = state + .oauth_registry + .needs_response_translation(oauth_provider_id, &original_path) + .map_err(|e| format!("OAuth translation check failed: {e}"))?; + let response_format = state + .oauth_registry + .response_format(oauth_provider_id, &original_path) + .map_err(|e| format!("OAuth response format check failed: {e}"))?; + // Check the transformed body for stream flag (fixups may force stream: true) + let stream_requested = serde_json::from_slice::(&body) + .ok() + .and_then(|v| v.get("stream")?.as_bool()) + .unwrap_or(false); + + let effective_uri = if let Some(ref new_path) = rewritten_path { + build_rewritten_uri(uri, new_path)? + } else { + uri.clone() + }; + + if rewritten_path.is_some() { + debug!( + oauth_provider = %oauth_provider_id, + original_path = %original_path, + effective_path = %effective_uri.path(), + "Rewrote request path for OAuth provider" + ); + } + + // 3. Optionally get upstream URL override + let upstream_url = state + .oauth_registry + .upstream_url(oauth_provider_id) + .await + .map_err(|e| format!("OAuth upstream URL resolution failed: {e}"))?; + + // 4. Forward the request + let response = state + .proxy_client + .forward_oauth( + method.clone(), + &effective_uri, + headers.clone(), + body.clone(), + provider, + auth_headers.clone(), + upstream_url.as_deref(), + ) + .await + .map_err(|e| format!("OAuth forward failed: {e}"))?; + + // 5. If we got a 401, refresh the token and retry once + if response.status() == StatusCode::UNAUTHORIZED { + info!( + oauth_provider = %oauth_provider_id, + effective_path = %effective_uri.path(), + "Got 401 from upstream, attempting token refresh and retry" + ); + if let Err(e) = state.oauth_registry.refresh(oauth_provider_id).await { + warn!( + oauth_provider = %oauth_provider_id, + error = %e, + "Token refresh failed after 401" + ); + return maybe_translate_response(response, needs_translation, stream_requested, response_format).await; + } + + // Re-inject auth with refreshed token + let mut retry_auth_headers = HeaderMap::new(); + state + .oauth_registry + .inject_auth(oauth_provider_id, &mut retry_auth_headers) + .await + .map_err(|e| format!("OAuth retry auth injection failed: {e}"))?; + + // Optionally re-transform the body (tokens may have changed affecting body) + let retry_body = match state + .oauth_registry + .prepare_request_body(oauth_provider_id, &body_bytes) + .await + .map_err(|e| format!("OAuth retry body preparation failed: {e}"))? + { + Some(transformed) => Bytes::from(transformed), + None => body_bytes, + }; + + let retry_response = state + .proxy_client + .forward_oauth( + method, + &effective_uri, + headers, + retry_body, + provider, + retry_auth_headers, + upstream_url.as_deref(), + ) + .await + .map_err(|e| format!("OAuth retry forward failed: {e}"))?; + + if retry_response.status() == StatusCode::UNAUTHORIZED { + warn!( + oauth_provider = %oauth_provider_id, + effective_path = %effective_uri.path(), + "Retry after token refresh still returned 401" + ); + } + + return maybe_translate_response(retry_response, needs_translation, stream_requested, response_format).await; + } + + // Log error response bodies for debugging upstream issues + if response.status().is_client_error() || response.status().is_server_error() { + let status = response.status(); + let (parts, body) = response.into_parts(); + let body_bytes_resp = body + .collect() + .await + .map(|b| b.to_bytes()) + .unwrap_or_default(); + if let Ok(body_str) = std::str::from_utf8(&body_bytes_resp) { + warn!( + oauth_provider = %oauth_provider_id, + effective_path = %effective_uri.path(), + status = %status, + response_body = %body_str, + "Upstream returned error" + ); + } + let response = Response::from_parts(parts, Body::from(body_bytes_resp)); + return maybe_translate_response(response, needs_translation, stream_requested, response_format).await; + } + + maybe_translate_response(response, needs_translation, stream_requested, response_format).await +} + +/// Optionally translate an upstream response back to chat/completions format. +async fn maybe_translate_response( + response: Response, + needs_translation: bool, + stream_requested: bool, + response_format: Option, +) -> Result { + // Use response_format if available; fall back to needs_translation for backwards compat + let format = match response_format { + Some(f) => f, + None if needs_translation => crate::oauth::ResponseFormat::ResponsesApi, + None => return Ok(response), + }; + + let is_streaming = stream_requested + || response + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .is_some_and(|ct| ct.contains("text/event-stream")); + + debug!(format = ?format, is_streaming, "maybe_translate_response: translating response"); + + if is_streaming { + let (parts, body) = response.into_parts(); + let translated_body = match format { + crate::oauth::ResponseFormat::ResponsesApi => { + debug!("Wrapping streaming response with ResponsesApi translator"); + crate::translate::wrap_body_with_translate_stream(body) + } + crate::oauth::ResponseFormat::GeminiSse => { + debug!("Wrapping streaming response with GeminiSse translator"); + crate::translate::wrap_body_with_gemini_translate_stream(body) + } + }; + return Ok(Response::from_parts(parts, translated_body)); + } + + // Non-streaming: only translate successful responses + let status = response.status(); + if !status.is_success() { + return Ok(response); + } + + let (mut parts, body) = response.into_parts(); + let body_bytes = body + .collect() + .await + .map_err(|e| format!("failed to read response body for translation: {e}"))? + .to_bytes(); + + match format { + crate::oauth::ResponseFormat::ResponsesApi => { + match crate::translate::responses_to_chat_completion(&body_bytes) { + Ok(translated) => { + parts.headers.remove("content-length"); + Ok(Response::from_parts(parts, Body::from(translated))) + } + Err(e) => { + warn!(error = %e, "Response translation failed, returning original"); + Ok(Response::from_parts(parts, Body::from(body_bytes))) + } + } + } + crate::oauth::ResponseFormat::GeminiSse => { + // Non-streaming Gemini responses are not expected; pass through + Ok(Response::from_parts(parts, Body::from(body_bytes))) + } + } +} + +/// Build a new URI with a rewritten path, preserving query string. +/// Incoming axum URIs are path-only (no scheme/authority), so we build path-only too. +fn build_rewritten_uri(original: &Uri, new_path: &str) -> Result { + let path_and_query = if let Some(query) = original.query() { + format!("{new_path}?{query}") + } else { + new_path.to_string() + }; + path_and_query + .parse::() + .map_err(|e| format!("failed to build rewritten URI: {e}")) +} + fn error_response(status: StatusCode, message: &str) -> Response { let body = serde_json::json!({ "error": message }); (status, axum::Json(body)).into_response() diff --git a/src/app/tests.rs b/src/app/tests.rs index f143ffc..925db69 100644 --- a/src/app/tests.rs +++ b/src/app/tests.rs @@ -16,7 +16,8 @@ use crate::email::{ EmailAccountCredentials, EmailListMessagesResponse, EmailMessageContent, EmailMessageMetadata, EmailPolicy, EmailService, }; -use crate::keys::{KeyManager, ResolvedKey}; +use crate::keys::{KeyManager, KeySource, ResolvedKey}; +use crate::oauth::OAuthRegistry; use crate::proxy::ProxyClient; fn make_app(upstream_url: &str) -> axum::Router { @@ -24,14 +25,14 @@ fn make_app(upstream_url: &str) -> axum::Router { key_map.insert( "vk-test-1".to_string(), ResolvedKey { - real_key: "sk-real-1".to_string(), + source: KeySource::Static { real_key: "sk-real-1".to_string() }, provider: Provider::Openai, }, ); key_map.insert( "vk-test-2".to_string(), ResolvedKey { - real_key: "sk-real-2".to_string(), + source: KeySource::Static { real_key: "sk-real-2".to_string() }, provider: Provider::Openai, }, ); @@ -65,6 +66,7 @@ fn make_app(upstream_url: &str) -> axum::Router { upstream_urls, "2023-06-01".to_string(), )), + oauth_registry: Arc::new(OAuthRegistry::new(Default::default())), email_enabled: false, email_policy: None, email_accounts: Arc::new(BTreeMap::new()), @@ -79,14 +81,14 @@ fn make_app_with_anthropic(upstream_url: &str) -> axum::Router { key_map.insert( "vk-test-1".to_string(), ResolvedKey { - real_key: "sk-real-1".to_string(), + source: KeySource::Static { real_key: "sk-real-1".to_string() }, provider: Provider::Openai, }, ); key_map.insert( "vk-ant-1".to_string(), ResolvedKey { - real_key: "sk-ant-real-1".to_string(), + source: KeySource::Static { real_key: "sk-ant-real-1".to_string() }, provider: Provider::Anthropic, }, ); @@ -102,6 +104,7 @@ fn make_app_with_anthropic(upstream_url: &str) -> axum::Router { upstream_urls, "2023-06-01".to_string(), )), + oauth_registry: Arc::new(OAuthRegistry::new(Default::default())), email_enabled: false, email_policy: None, email_accounts: Arc::new(BTreeMap::new()), @@ -561,7 +564,10 @@ real_key = "sk-real-1" let config = Config::parse(toml_str).unwrap(); let state = AppState::from_config(&config).unwrap(); let resolved = state.key_manager.resolve("vk-1").unwrap(); - assert_eq!(resolved.real_key, "sk-real-1"); + match &resolved.source { + KeySource::Static { real_key } => assert_eq!(real_key, "sk-real-1"), + _ => panic!("expected Static key source"), + } assert_eq!(resolved.provider, Provider::Openai); assert!(state.key_manager.resolve("vk-unknown").is_none()); } @@ -587,10 +593,16 @@ provider = "anthropic" let config = Config::parse(toml_str).unwrap(); let state = AppState::from_config(&config).unwrap(); let oai = state.key_manager.resolve("vk-oai").unwrap(); - assert_eq!(oai.real_key, "sk-oai-key"); + match &oai.source { + KeySource::Static { real_key } => assert_eq!(real_key, "sk-oai-key"), + _ => panic!("expected Static key source"), + } assert_eq!(oai.provider, Provider::Openai); let ant = state.key_manager.resolve("vk-ant").unwrap(); - assert_eq!(ant.real_key, "sk-ant-key"); + match &ant.source { + KeySource::Static { real_key } => assert_eq!(real_key, "sk-ant-key"), + _ => panic!("expected Static key source"), + } assert_eq!(ant.provider, Provider::Anthropic); } @@ -664,7 +676,7 @@ async fn test_proxy_error_on_unreachable_upstream() { [( "vk-1".to_string(), ResolvedKey { - real_key: "sk-1".to_string(), + source: KeySource::Static { real_key: "sk-1".to_string() }, provider: Provider::Openai, }, )] @@ -681,6 +693,7 @@ async fn test_proxy_error_on_unreachable_upstream() { }, "2023-06-01".to_string(), )), + oauth_registry: Arc::new(OAuthRegistry::new(Default::default())), email_enabled: false, email_policy: None, email_accounts: Arc::new(BTreeMap::new()), @@ -803,7 +816,7 @@ async fn test_anthropic_dlp_blocks_sensitive_data() { key_map.insert( "vk-ant-dlp".to_string(), ResolvedKey { - real_key: "sk-ant-key".to_string(), + source: KeySource::Static { real_key: "sk-ant-key".to_string() }, provider: Provider::Anthropic, }, ); @@ -825,6 +838,7 @@ async fn test_anthropic_dlp_blocks_sensitive_data() { upstream_urls, "2023-06-01".to_string(), )), + oauth_registry: Arc::new(OAuthRegistry::new(Default::default())), email_enabled: false, email_policy: None, email_accounts: Arc::new(BTreeMap::new()), @@ -898,14 +912,14 @@ async fn test_openai_and_openrouter_keys_map_to_distinct_real_keys() { key_map.insert( "vk-openai".to_string(), ResolvedKey { - real_key: "sk-openai-real".to_string(), + source: KeySource::Static { real_key: "sk-openai-real".to_string() }, provider: Provider::Openai, }, ); key_map.insert( "vk-openrouter".to_string(), ResolvedKey { - real_key: "sk-openrouter-real".to_string(), + source: KeySource::Static { real_key: "sk-openrouter-real".to_string() }, provider: Provider::Openrouter, }, ); @@ -922,6 +936,7 @@ async fn test_openai_and_openrouter_keys_map_to_distinct_real_keys() { upstream_urls, "2023-06-01".to_string(), )), + oauth_registry: Arc::new(OAuthRegistry::new(Default::default())), email_enabled: false, email_policy: None, email_accounts: Arc::new(BTreeMap::new()), @@ -956,7 +971,7 @@ fn make_app_with_redact(upstream_url: &str) -> axum::Router { key_map.insert( "vk-test-1".to_string(), ResolvedKey { - real_key: "sk-real-1".to_string(), + source: KeySource::Static { real_key: "sk-real-1".to_string() }, provider: Provider::Openai, }, ); @@ -990,6 +1005,7 @@ fn make_app_with_redact(upstream_url: &str) -> axum::Router { upstream_urls, "2023-06-01".to_string(), )), + oauth_registry: Arc::new(OAuthRegistry::new(Default::default())), email_enabled: false, email_policy: None, email_accounts: Arc::new(BTreeMap::new()), @@ -1175,7 +1191,7 @@ async fn test_response_dlp_disabled() { key_map.insert( "vk-test-1".to_string(), ResolvedKey { - real_key: "sk-real-1".to_string(), + source: KeySource::Static { real_key: "sk-real-1".to_string() }, provider: Provider::Openai, }, ); @@ -1194,6 +1210,7 @@ async fn test_response_dlp_disabled() { upstream_urls, "2023-06-01".to_string(), )), + oauth_registry: Arc::new(OAuthRegistry::new(Default::default())), email_enabled: false, email_policy: None, email_accounts: Arc::new(BTreeMap::new()), @@ -1442,8 +1459,7 @@ async fn test_non_utf8_body_passes_through() { async fn test_streaming_response_with_dlp_enabled_passes_through() { let mock_server = MockServer::start().await; - // SSE response — should pass through when DLP scanning is enabled - // because streaming responses cannot be scanned (exercises lib.rs lines 261-268) + // SSE response with clean content — should pass through DLP scanning unchanged let sse_body = "data: {\"content\":\"hello world\"}\n\ndata: [DONE]\n\n"; Mock::given(method("POST")) .and(path("/v1/chat/completions")) @@ -1484,6 +1500,51 @@ async fn test_streaming_response_with_dlp_enabled_passes_through() { assert!(body_str.contains("[DONE]")); } +#[tokio::test] +async fn test_streaming_response_dlp_redacts_pii_in_sse() { + let mock_server = MockServer::start().await; + + // SSE response with PII in delta.content — DLP should redact it + let chunk = serde_json::json!({ + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "choices": [{ + "index": 0, + "delta": { "content": "Contact user@example.com for help" }, + "finish_reason": null, + }] + }); + let sse_body = format!("data: {}\n\ndata: [DONE]\n\n", serde_json::to_string(&chunk).unwrap()); + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_raw(sse_body, "text/event-stream"), + ) + .mount(&mock_server) + .await; + + let app = make_app_with_redact(&mock_server.uri()); + let body = r#"{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"Hi"}]}"#; + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer vk-test-1") + .header("content-type", "application/json") + .body(Body::from(body)) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = resp.into_body().collect().await.unwrap().to_bytes(); + let body_str = std::str::from_utf8(&body).unwrap(); + assert!(body_str.contains("[REDACTED:email]"), "PII should be redacted in streaming SSE"); + assert!(!body_str.contains("user@example.com"), "Original email should be gone"); + assert!(body_str.contains("[DONE]"), "Stream should still end with [DONE]"); +} + fn make_email_app( policy: EmailPolicy, email_accounts: BTreeMap, @@ -1500,6 +1561,7 @@ fn make_email_app( upstream_urls, "2023-06-01".to_string(), )), + oauth_registry: Arc::new(OAuthRegistry::new(Default::default())), email_enabled: true, email_policy: Some(policy), email_accounts: Arc::new(email_accounts), diff --git a/src/config.rs b/src/config.rs index 6c619ca..73ece6b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,7 @@ use regex::Regex; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; +use std::env::VarError; use std::path::Path; #[derive(Debug, Default, Deserialize, Serialize, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] @@ -37,6 +38,8 @@ pub struct Config { pub email: EmailConfig, #[serde(default = "default_log_level")] pub log_level: String, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub oauth_providers: Vec, } fn default_log_level() -> String { @@ -60,6 +63,9 @@ fn default_port() -> u16 { 18790 } +const SERVER_HOST_ENV: &str = "CLAWSHELL_SERVER_HOST"; +const SERVER_PORT_ENV: &str = "CLAWSHELL_SERVER_PORT"; + #[derive(Debug, Deserialize, Serialize, Clone)] #[serde(deny_unknown_fields)] pub struct UpstreamConfig { @@ -81,13 +87,33 @@ fn default_openai_base_url() -> String { "https://api.openai.com".to_string() } +/// How a key mapping authenticates: static API key or OAuth provider. +#[derive(Debug, Default, Deserialize, Serialize, Clone, Copy, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum KeyAuthMethod { + /// Static API key (the default, existing behavior). + #[default] + Static, + /// OAuth provider supplies the access token at runtime. + OAuth, +} + #[derive(Debug, Deserialize, Serialize, Clone)] #[serde(deny_unknown_fields)] pub struct KeyMapping { pub virtual_key: String, - pub real_key: String, + /// Required when auth = "static" (or omitted). Optional when auth = "oauth". + #[serde(default, skip_serializing_if = "Option::is_none")] + pub real_key: Option, #[serde(default)] pub provider: Provider, + /// Authentication method for this key. Defaults to "static". + #[serde(default)] + pub auth: KeyAuthMethod, + /// Which OAuth provider supplies the token (e.g. "codex"). + /// Required when auth = "oauth". + #[serde(default, skip_serializing_if = "Option::is_none")] + pub oauth_provider: Option, } #[derive(Debug, Default, Deserialize, Serialize, Clone, Copy, PartialEq, Eq)] @@ -220,10 +246,50 @@ impl Config { Regex::new(&pattern.regex) .map_err(|e| format!("Invalid DLP regex for '{}': {}", pattern.name, e))?; } + self.validate_keys()?; self.validate_email()?; Ok(()) } + fn validate_keys(&self) -> Result<(), Box> { + for key in &self.keys { + match key.auth { + KeyAuthMethod::Static => { + if key.real_key.is_none() { + return Err(format!( + "key '{}': real_key is required when auth = \"static\"", + key.virtual_key + ) + .into()); + } + } + KeyAuthMethod::OAuth => { + if key.oauth_provider.as_ref().is_none_or(|p| p.trim().is_empty()) { + return Err(format!( + "key '{}': oauth_provider is required when auth = \"oauth\"", + key.virtual_key + ) + .into()); + } + // Verify the referenced OAuth provider exists in config + let provider_id = key.oauth_provider.as_ref().unwrap(); + if !self + .oauth_providers + .iter() + .any(|p| p.provider == *provider_id) + { + return Err(format!( + "key '{}': oauth_provider '{}' not found in [[oauth_providers]]", + key.virtual_key, provider_id + ) + .into()); + } + } + } + } + Ok(()) + } + fn validate_email(&self) -> Result<(), Box> { let email = &self.email; @@ -315,10 +381,32 @@ impl Config { Ok(()) } + /// Returns a map of static key mappings: virtual_key → (real_key, provider). + /// OAuth-backed keys are excluded. + #[allow(dead_code)] pub fn key_map(&self) -> BTreeMap { self.keys .iter() - .map(|k| (k.virtual_key.clone(), (k.real_key.clone(), k.provider))) + .filter(|k| k.auth == KeyAuthMethod::Static) + .filter_map(|k| { + k.real_key + .clone() + .map(|rk| (k.virtual_key.clone(), (rk, k.provider))) + }) + .collect() + } + + /// Returns a map of OAuth key mappings: virtual_key → (oauth_provider_id, provider). + #[allow(dead_code)] + pub fn oauth_key_map(&self) -> BTreeMap { + self.keys + .iter() + .filter(|k| k.auth == KeyAuthMethod::OAuth) + .filter_map(|k| { + k.oauth_provider + .clone() + .map(|op| (k.virtual_key.clone(), (op, k.provider))) + }) .collect() } @@ -341,6 +429,61 @@ impl Config { pub fn listen_addr(&self) -> String { format!("{}:{}", self.server.host, self.server.port) } + + pub fn resolved_listen_addr(&self) -> Result> { + let host = resolve_server_host_override(&self.server.host)?; + let port = resolve_server_port_override(self.server.port)?; + Ok(format!("{host}:{port}")) + } +} + +fn resolve_server_host_override(default_host: &str) -> Result> { + resolve_server_host_override_from_var(default_host, std::env::var(SERVER_HOST_ENV)) +} + +fn resolve_server_host_override_from_var( + default_host: &str, + env_value: Result, +) -> Result> { + match env_value { + Ok(value) => { + let trimmed = value.trim(); + if trimmed.is_empty() { + return Err(format!("{SERVER_HOST_ENV} cannot be empty").into()); + } + Ok(trimmed.to_string()) + } + Err(VarError::NotPresent) => Ok(default_host.to_string()), + Err(VarError::NotUnicode(_)) => { + Err(format!("{SERVER_HOST_ENV} must be valid UTF-8").into()) + } + } +} + +fn resolve_server_port_override(default_port: u16) -> Result> { + resolve_server_port_override_from_var(default_port, std::env::var(SERVER_PORT_ENV)) +} + +fn resolve_server_port_override_from_var( + default_port: u16, + env_value: Result, +) -> Result> { + match env_value { + Ok(value) => { + let trimmed = value.trim(); + if trimmed.is_empty() { + return Err(format!("{SERVER_PORT_ENV} cannot be empty").into()); + } + trimmed.parse::().map_err(|_| { + format!("{SERVER_PORT_ENV} must be a valid port (0-65535), got '{trimmed}'") + .into() + }) + } + Err(VarError::NotPresent) => Ok(default_port), + Err(VarError::NotUnicode(_)) => { + Err(format!("{SERVER_PORT_ENV} must be valid UTF-8").into()) + } + } } pub(crate) fn validate_sender_rule(rule: &str) -> Result<(), String> { @@ -498,6 +641,7 @@ mod tests { Ok(()) } + #[test] fn test_valid_config_fixtures() { let paths = @@ -850,4 +994,62 @@ imap_port = 0 .contains("email.accounts[].imap_port must be greater than 0") ); } + + #[test] + fn test_resolved_listen_addr_uses_config_without_env() { + let cfg = r#" +[server] +host = "127.0.0.1" +port = 3000 + +[upstream] +openai_base_url = "https://api.openai.com" +"#; + let parsed = Config::parse(cfg).expect("config should parse"); + assert_eq!(parsed.listen_addr(), "127.0.0.1:3000"); + } + + #[test] + fn test_resolve_server_host_override_uses_default_when_unset() { + let host = resolve_server_host_override_from_var("127.0.0.1", Err(VarError::NotPresent)) + .expect("host should use default"); + assert_eq!(host, "127.0.0.1"); + } + + #[test] + fn test_resolve_server_host_override_accepts_env() { + let host = resolve_server_host_override_from_var("127.0.0.1", Ok("0.0.0.0".to_string())) + .expect("host override should be accepted"); + assert_eq!(host, "0.0.0.0"); + } + + #[test] + fn test_resolve_server_host_override_rejects_empty_env() { + let err = resolve_server_host_override_from_var("127.0.0.1", Ok(" ".to_string())) + .unwrap_err(); + assert!( + err.to_string() + .contains("CLAWSHELL_SERVER_HOST cannot be empty") + ); + } + + #[test] + fn test_resolve_server_port_override_uses_default_when_unset() { + let port = + resolve_server_port_override_from_var(3000, Err(VarError::NotPresent)).unwrap(); + assert_eq!(port, 3000); + } + + #[test] + fn test_resolve_server_port_override_accepts_env() { + let port = resolve_server_port_override_from_var(3000, Ok("17890".to_string())).unwrap(); + assert_eq!(port, 17890); + } + + #[test] + fn test_resolve_server_port_override_rejects_invalid_env() { + let err = resolve_server_port_override_from_var(3000, Ok("not-a-port".to_string())) + .unwrap_err(); + assert!(err.to_string().contains("CLAWSHELL_SERVER_PORT must be a valid port")); + } } diff --git a/src/keys.rs b/src/keys.rs index 1610cb7..b64beee 100644 --- a/src/keys.rs +++ b/src/keys.rs @@ -3,9 +3,15 @@ use crate::config::Provider; use std::collections::BTreeMap; use tracing::{debug, trace}; +#[derive(Debug, Clone)] +pub enum KeySource { + Static { real_key: String }, + OAuth { provider_id: String }, +} + #[derive(Debug, Clone)] pub struct ResolvedKey { - pub real_key: String, + pub source: KeySource, pub provider: Provider, } @@ -58,14 +64,16 @@ impl KeyManager { mod tests { use super::*; - fn make_map(entries: Vec<(&str, &str, Provider)>) -> BTreeMap { + fn make_static_map(entries: Vec<(&str, &str, Provider)>) -> BTreeMap { entries .into_iter() .map(|(vk, rk, p)| { ( vk.to_string(), ResolvedKey { - real_key: rk.to_string(), + source: KeySource::Static { + real_key: rk.to_string(), + }, provider: p, }, ) @@ -93,10 +101,13 @@ mod tests { #[test] fn test_resolve_existing_key() { - let map = make_map(vec![("vk-1", "sk-real-1", Provider::Openai)]); + let map = make_static_map(vec![("vk-1", "sk-real-1", Provider::Openai)]); let km = KeyManager::new(map); let resolved = km.resolve("vk-1").unwrap(); - assert_eq!(resolved.real_key, "sk-real-1"); + match &resolved.source { + KeySource::Static { real_key } => assert_eq!(real_key, "sk-real-1"), + KeySource::OAuth { .. } => panic!("expected Static"), + } assert_eq!(resolved.provider, Provider::Openai); } @@ -108,21 +119,51 @@ mod tests { #[test] fn test_multiple_virtual_to_same_real() { - let map = make_map(vec![ + let map = make_static_map(vec![ ("vk-1", "sk-shared", Provider::Openai), ("vk-2", "sk-shared", Provider::Openai), ]); let km = KeyManager::new(map); - assert_eq!(km.resolve("vk-1").unwrap().real_key, "sk-shared"); - assert_eq!(km.resolve("vk-2").unwrap().real_key, "sk-shared"); + match &km.resolve("vk-1").unwrap().source { + KeySource::Static { real_key } => assert_eq!(real_key, "sk-shared"), + _ => panic!("expected Static"), + } + match &km.resolve("vk-2").unwrap().source { + KeySource::Static { real_key } => assert_eq!(real_key, "sk-shared"), + _ => panic!("expected Static"), + } } #[test] fn test_resolve_anthropic_provider() { - let map = make_map(vec![("vk-ant", "sk-ant-key", Provider::Anthropic)]); + let map = make_static_map(vec![("vk-ant", "sk-ant-key", Provider::Anthropic)]); let km = KeyManager::new(map); let resolved = km.resolve("vk-ant").unwrap(); - assert_eq!(resolved.real_key, "sk-ant-key"); + match &resolved.source { + KeySource::Static { real_key } => assert_eq!(real_key, "sk-ant-key"), + _ => panic!("expected Static"), + } assert_eq!(resolved.provider, Provider::Anthropic); } + + #[test] + fn test_resolve_oauth_key() { + let mut map = BTreeMap::new(); + map.insert( + "vk-oauth".to_string(), + ResolvedKey { + source: KeySource::OAuth { + provider_id: "codex".to_string(), + }, + provider: Provider::Openai, + }, + ); + let km = KeyManager::new(map); + let resolved = km.resolve("vk-oauth").unwrap(); + match &resolved.source { + KeySource::OAuth { provider_id } => assert_eq!(provider_id, "codex"), + _ => panic!("expected OAuth"), + } + assert_eq!(resolved.provider, Provider::Openai); + } } diff --git a/src/main.rs b/src/main.rs index 5e716d6..83ba5b7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,11 +9,14 @@ mod dlp; mod email; mod keys; mod migration; +#[allow(dead_code)] +mod oauth; mod onboard; mod openclaw_cli; mod platform; mod process; mod proxy; +mod translate; mod tui; use clap::Parser; @@ -555,7 +558,10 @@ async fn cmd_start_inner(config_path: &str) -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result> { + use crate::oauth::{ + OAuthRegistry, TokenStorage, + codex::CodexProvider, + }; + use std::sync::Arc; + + let storage = TokenStorage::new(PathBuf::from("/etc/clawshell/oauth")); + let mut registry = OAuthRegistry::new(storage); + + for provider_config in &config.oauth_providers { + if !provider_config.enabled { + continue; + } + match provider_config.provider.as_str() { + "codex" => { + let provider = CodexProvider::from_config(provider_config); + registry.register(Arc::new(provider)); + } + other => { + return Err(format!("Unknown OAuth provider type: '{other}'").into()); + } + } + } + + // Load persisted tokens from disk + registry.load_tokens().await?; + + Ok(registry) +} + fn cmd_stop() -> Result<(), Box> { tui::print_banner("Stop"); ensure_default_config_migrated_if_present()?; @@ -1051,13 +1100,27 @@ fn cmd_onboard() -> Result<(), Box> { let toml_content = onboard::generate_clawshell_config(&ob_config); std::fs::write(&toml_config_path, &toml_content)?; - let config_json = serde_json::json!({ - "real_api_key": ob_config.real_api_key, - "virtual_api_key": ob_config.virtual_api_key, - "provider": ob_config.provider, - "model": ob_config.model, - "openclaw_config_path": ob_config.openclaw_config_path.to_string_lossy(), - }); + let config_json = match &ob_config.auth_method { + crate::onboard::OnboardAuthMethod::OAuth { provider_id } => { + serde_json::json!({ + "auth_method": "oauth", + "oauth_provider": provider_id, + "virtual_api_key": ob_config.virtual_api_key, + "provider": ob_config.provider, + "model": ob_config.model, + "openclaw_config_path": ob_config.openclaw_config_path.to_string_lossy(), + }) + } + crate::onboard::OnboardAuthMethod::StaticKey => { + serde_json::json!({ + "real_api_key": ob_config.real_api_key, + "virtual_api_key": ob_config.virtual_api_key, + "provider": ob_config.provider, + "model": ob_config.model, + "openclaw_config_path": ob_config.openclaw_config_path.to_string_lossy(), + }) + } + }; std::fs::write(&config_file, serde_json::to_string_pretty(&config_json)?)?; // Set permissions on config files diff --git a/src/oauth/codex.rs b/src/oauth/codex.rs new file mode 100644 index 0000000..e34da5c --- /dev/null +++ b/src/oauth/codex.rs @@ -0,0 +1,714 @@ +use super::{OAuthError, OAuthProvider, OAuthTokens}; +use async_trait::async_trait; +use axum::http::header::AUTHORIZATION; +use axum::http::HeaderMap; +use chrono::Utc; +use std::collections::BTreeMap; +use tracing::{debug, info}; + +const DEFAULT_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; +const DEFAULT_AUTH_URL: &str = "https://auth.openai.com/authorize"; +const DEFAULT_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; +const DEFAULT_SCOPES: &[&str] = &["openid", "profile", "email", "offline_access"]; + +#[derive(Debug)] +pub struct CodexProvider { + client_id: String, + auth_url: String, + token_url: String, + scopes: Vec, + http_client: reqwest::Client, +} + +impl CodexProvider { + pub fn new( + client_id: Option<&str>, + auth_url: Option<&str>, + token_url: Option<&str>, + scopes: Option<&[String]>, + ) -> Self { + Self { + client_id: client_id.unwrap_or(DEFAULT_CLIENT_ID).to_string(), + auth_url: auth_url.unwrap_or(DEFAULT_AUTH_URL).to_string(), + token_url: token_url.unwrap_or(DEFAULT_TOKEN_URL).to_string(), + scopes: scopes + .map(|s| s.to_vec()) + .unwrap_or_else(|| DEFAULT_SCOPES.iter().map(|s| s.to_string()).collect()), + http_client: reqwest::Client::builder() + .user_agent(format!( + "ClawShell/{} (https://github.com/nicholasgasior/clawshell)", + env!("CARGO_PKG_VERSION") + )) + .build() + .expect("failed to build HTTP client"), + } + } + + pub fn from_config(config: &super::OAuthProviderConfig) -> Self { + Self::new( + config.client_id.as_deref(), + config.auth_url.as_deref(), + config.token_url.as_deref(), + config.scopes.as_deref(), + ) + } + + async fn exchange_code( + &self, + code: &str, + code_verifier: &str, + redirect_uri: &str, + ) -> Result { + let params = [ + ("grant_type", "authorization_code"), + ("client_id", &self.client_id), + ("code", code), + ("code_verifier", code_verifier), + ("redirect_uri", redirect_uri), + ]; + + let resp = self + .http_client + .post(&self.token_url) + .form(¶ms) + .send() + .await?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(OAuthError::LoginFailed(format!( + "token exchange failed ({status}): {body}" + ))); + } + + let json: serde_json::Value = resp.json().await?; + parse_token_response(&json) + } + + async fn exchange_refresh_token( + &self, + refresh_token: &str, + ) -> Result { + let params = [ + ("grant_type", "refresh_token"), + ("client_id", &self.client_id), + ("refresh_token", refresh_token), + ]; + + let resp = self + .http_client + .post(&self.token_url) + .form(¶ms) + .send() + .await?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(OAuthError::RefreshFailed(format!( + "refresh failed ({status}): {body}" + ))); + } + + let json: serde_json::Value = resp.json().await?; + parse_token_response(&json) + } + + /// Poll OpenAI's custom device-auth token endpoint until user authorises. + /// Returns (authorization_code, code_verifier) on success. + async fn poll_device_auth( + &self, + device_auth_id: &str, + user_code: &str, + interval: u64, + ) -> Result<(String, String), OAuthError> { + let url = self.device_auth_base_url() + "/token"; + let max_wait = std::time::Duration::from_secs(15 * 60); + let start = std::time::Instant::now(); + + loop { + tokio::time::sleep(std::time::Duration::from_secs(interval)).await; + + if start.elapsed() > max_wait { + return Err(OAuthError::LoginFailed( + "device code polling timed out (15 min)".to_string(), + )); + } + + let body = serde_json::json!({ + "device_auth_id": device_auth_id, + "user_code": user_code, + }); + + let resp = self + .http_client + .post(&url) + .json(&body) + .send() + .await?; + + if resp.status().is_success() { + let json: serde_json::Value = resp.json().await?; + let auth_code = json + .get("authorization_code") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + OAuthError::LoginFailed( + "missing authorization_code in device-auth response".to_string(), + ) + })? + .to_string(); + let code_verifier = json + .get("code_verifier") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + OAuthError::LoginFailed( + "missing code_verifier in device-auth response".to_string(), + ) + })? + .to_string(); + return Ok((auth_code, code_verifier)); + } + + // 403 / 404 = authorization still pending + let status = resp.status(); + if status == reqwest::StatusCode::FORBIDDEN + || status == reqwest::StatusCode::NOT_FOUND + { + debug!("Device code authorization pending ({status})"); + continue; + } + + let text = resp.text().await.unwrap_or_default(); + return Err(OAuthError::LoginFailed(format!( + "device-auth polling failed ({status}): {text}" + ))); + } + } + + /// Base URL for OpenAI's custom device-auth API, derived from `auth_url`. + fn device_auth_base_url(&self) -> String { + // auth_url is e.g. "https://auth.openai.com/authorize" + // We need "https://auth.openai.com/api/accounts/deviceauth" + let base = self + .auth_url + .trim_end_matches("/authorize") + .trim_end_matches('/'); + format!("{base}/api/accounts/deviceauth") + } +} + +fn parse_token_response(json: &serde_json::Value) -> Result { + let access_token = json + .get("access_token") + .and_then(|v| v.as_str()) + .ok_or_else(|| OAuthError::LoginFailed("missing access_token in response".to_string()))? + .to_string(); + + let refresh_token = json + .get("refresh_token") + .and_then(|v| v.as_str()) + .map(String::from); + + let id_token = json + .get("id_token") + .and_then(|v| v.as_str()) + .map(String::from); + + let expires_at = json + .get("expires_in") + .and_then(|v| v.as_i64()) + .map(|secs| Utc::now() + chrono::Duration::seconds(secs)); + + Ok(OAuthTokens { + access_token, + refresh_token, + id_token, + expires_at, + account_id: None, + extra: BTreeMap::new(), + }) +} + +fn generate_pkce() -> (String, String) { + use base64::Engine; + use sha2::{Digest, Sha256}; + + let verifier_bytes: [u8; 32] = rand::random(); + let verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(verifier_bytes); + + let mut hasher = Sha256::new(); + hasher.update(verifier.as_bytes()); + let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hasher.finalize()); + + (verifier, challenge) +} + +#[async_trait] +impl OAuthProvider for CodexProvider { + fn id(&self) -> &str { + "codex" + } + + fn display_name(&self) -> &str { + "Codex / ChatGPT (OAuth)" + } + + fn supports_device_code(&self) -> bool { + true + } + + async fn login_browser(&self, callback_port: u16) -> Result { + let (verifier, challenge) = generate_pkce(); + let redirect_uri = format!("http://localhost:{callback_port}/auth/callback"); + let state: String = uuid::Uuid::new_v4().to_string(); + + let auth_url = format!( + "{}?response_type=code&client_id={}&redirect_uri={}&scope={}&code_challenge={}&code_challenge_method=S256&state={}", + self.auth_url, + urlencoding::encode(&self.client_id), + urlencoding::encode(&redirect_uri), + urlencoding::encode(&self.scopes.join(" ")), + urlencoding::encode(&challenge), + urlencoding::encode(&state), + ); + + info!("Opening browser for Codex OAuth login"); + if let Err(e) = open::that(&auth_url) { + return Err(OAuthError::LoginFailed(format!( + "failed to open browser: {e}. Visit this URL manually: {auth_url}" + ))); + } + + // Start a temporary HTTP server to receive the callback + let (code, received_state) = + wait_for_oauth_callback(callback_port).await.map_err(|e| { + OAuthError::LoginFailed(format!("callback server failed: {e}")) + })?; + + if received_state != state { + return Err(OAuthError::LoginFailed( + "OAuth state mismatch — possible CSRF".to_string(), + )); + } + + self.exchange_code(&code, &verifier, &redirect_uri).await + } + + async fn login_headless(&self) -> Result { + // Step 1: Request a user code from OpenAI's device-auth endpoint + let usercode_url = self.device_auth_base_url() + "/usercode"; + let body = serde_json::json!({ "client_id": self.client_id }); + + let resp = self + .http_client + .post(&usercode_url) + .json(&body) + .send() + .await?; + + if !resp.status().is_success() { + let status = resp.status(); + let text = resp.text().await.unwrap_or_default(); + return Err(OAuthError::LoginFailed(format!( + "device code request failed ({status}): {text}" + ))); + } + + let json: serde_json::Value = resp.json().await?; + + let device_auth_id = json + .get("device_auth_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + OAuthError::LoginFailed("missing device_auth_id in response".to_string()) + })?; + + let user_code = json + .get("user_code") + .or_else(|| json.get("usercode")) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + OAuthError::LoginFailed("missing user_code in response".to_string()) + })?; + + let interval = json + .get("interval") + .and_then(|v| v.as_u64().or_else(|| v.as_str().and_then(|s| s.parse().ok()))) + .unwrap_or(5); + + // Verification URL for the user + let base = self + .auth_url + .trim_end_matches("/authorize") + .trim_end_matches('/'); + let verification_url = format!("{base}/codex/device"); + + println!(); + println!(" Visit: {verification_url}"); + println!(" Enter code: {user_code}"); + println!(); + + // Step 2: Poll until user authorises, get authorization_code + code_verifier + let (auth_code, code_verifier) = + self.poll_device_auth(device_auth_id, user_code, interval).await?; + + // Step 3: Exchange authorization_code for tokens via the standard token endpoint + let redirect_uri = format!("{base}/deviceauth/callback"); + self.exchange_code(&auth_code, &code_verifier, &redirect_uri) + .await + } + + async fn refresh(&self, refresh_token: &str) -> Result { + self.exchange_refresh_token(refresh_token).await + } + + fn inject_auth( + &self, + headers: &mut HeaderMap, + access_token: &str, + ) -> Result<(), OAuthError> { + headers.insert( + AUTHORIZATION, + format!("Bearer {access_token}").parse()?, + ); + // ChatGPT backend requires Accept header for SSE streaming + headers.insert( + axum::http::header::ACCEPT, + "text/event-stream".parse().unwrap(), + ); + Ok(()) + } + + fn prepare_request_body( + &self, + body: &[u8], + _tokens: &OAuthTokens, + ) -> Result>, OAuthError> { + // Only translate if the body is JSON with a "messages" field + let Ok(parsed) = serde_json::from_slice::(body) else { + return Ok(None); + }; + if parsed.get("messages").is_none() { + return Ok(None); + } + match crate::translate::chat_completions_to_responses(body) { + Ok(translated) => Ok(Some(fixup_for_chatgpt_backend(&translated))), + Err(e) => Err(OAuthError::LoginFailed(format!( + "request translation failed: {e}" + ))), + } + } + + fn upstream_url(&self, _tokens: &OAuthTokens) -> Option { + Some("https://chatgpt.com/backend-api/codex".to_string()) + } + + fn rewrite_request_path(&self, path: &str) -> Option { + if path == "/v1/chat/completions" { + Some("/responses".to_string()) + } else { + None + } + } + + fn needs_response_translation(&self, original_path: &str) -> bool { + original_path == "/v1/chat/completions" + } + + fn response_format(&self, original_path: &str) -> Option { + if original_path == "/v1/chat/completions" { + Some(super::ResponseFormat::ResponsesApi) + } else { + None + } + } +} + +/// Wait for an OAuth callback on a local HTTP server. +/// Returns (code, state). +async fn wait_for_oauth_callback( + port: u16, +) -> Result<(String, String), Box> { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind(super::callback_bind_addr(port)).await?; + let (mut stream, _) = listener.accept().await?; + + let mut buf = vec![0u8; 4096]; + let n = stream.read(&mut buf).await?; + let request = String::from_utf8_lossy(&buf[..n]); + + // Parse the GET request for code and state query params + let path = request + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .unwrap_or(""); + + let query = path.split('?').nth(1).unwrap_or(""); + let mut code = String::new(); + let mut state = String::new(); + + for param in query.split('&') { + if let Some((key, value)) = param.split_once('=') { + match key { + "code" => code = urlencoding::decode(value).unwrap_or_default().to_string(), + "state" => state = urlencoding::decode(value).unwrap_or_default().to_string(), + _ => {} + } + } + } + + let response = "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n\ +

Login successful!

You can close this tab.

"; + stream.write_all(response.as_bytes()).await?; + stream.shutdown().await?; + + if code.is_empty() { + return Err("no authorization code in callback".into()); + } + + Ok((code, state)) +} + +/// Apply ChatGPT backend-specific fixups to the translated request body: +/// - Strip provider prefix from model (e.g. "openai/gpt-5.2-codex" → "gpt-5.2-codex") +/// - Set `store: false` (required by ChatGPT backend) +/// - Set `stream: true` (required by ChatGPT backend) +fn fixup_for_chatgpt_backend(body: &[u8]) -> Vec { + let Ok(mut parsed) = serde_json::from_slice::(body) else { + return body.to_vec(); + }; + if let Some(model) = parsed.get("model").and_then(|v| v.as_str()) { + if let Some(stripped) = model.strip_prefix("openai/") { + parsed["model"] = serde_json::Value::String(stripped.to_string()); + } + } + parsed["store"] = serde_json::Value::Bool(false); + parsed["stream"] = serde_json::Value::Bool(true); + // Codex backend does not support max_output_tokens + if let Some(obj) = parsed.as_object_mut() { + obj.remove("max_output_tokens"); + } + serde_json::to_vec(&parsed).unwrap_or_else(|_| body.to_vec()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_token_response() { + let json = serde_json::json!({ + "access_token": "eyJ...", + "refresh_token": "v1.MjQ...", + "id_token": "eyJhbG...", + "expires_in": 3600, + "token_type": "Bearer" + }); + + let tokens = parse_token_response(&json).unwrap(); + assert_eq!(tokens.access_token, "eyJ..."); + assert_eq!(tokens.refresh_token.as_deref(), Some("v1.MjQ...")); + assert_eq!(tokens.id_token.as_deref(), Some("eyJhbG...")); + assert!(tokens.expires_at.is_some()); + } + + #[test] + fn test_parse_token_response_missing_access_token() { + let json = serde_json::json!({ + "refresh_token": "v1.MjQ...", + }); + + let result = parse_token_response(&json); + assert!(result.is_err()); + } + + #[test] + fn test_parse_token_response_minimal() { + let json = serde_json::json!({ + "access_token": "minimal" + }); + + let tokens = parse_token_response(&json).unwrap(); + assert_eq!(tokens.access_token, "minimal"); + assert!(tokens.refresh_token.is_none()); + assert!(tokens.id_token.is_none()); + assert!(tokens.expires_at.is_none()); + } + + #[test] + fn test_generate_pkce() { + let (verifier, challenge) = generate_pkce(); + assert!(!verifier.is_empty()); + assert!(!challenge.is_empty()); + assert_ne!(verifier, challenge); + + // Verify challenge is S256 of verifier + use base64::Engine; + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(verifier.as_bytes()); + let expected = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hasher.finalize()); + assert_eq!(challenge, expected); + } + + #[test] + fn test_codex_provider_defaults() { + let provider = CodexProvider::new(None, None, None, None); + assert_eq!(provider.id(), "codex"); + assert_eq!(provider.display_name(), "Codex / ChatGPT (OAuth)"); + assert!(provider.supports_device_code()); + assert!(!provider.supports_headless_url()); + assert_eq!(provider.client_id, DEFAULT_CLIENT_ID); + } + + #[test] + fn test_codex_provider_custom() { + let provider = CodexProvider::new( + Some("custom-client"), + Some("https://custom.auth/authorize"), + Some("https://custom.auth/token"), + Some(&["openid".to_string()]), + ); + assert_eq!(provider.client_id, "custom-client"); + assert_eq!(provider.auth_url, "https://custom.auth/authorize"); + assert_eq!(provider.token_url, "https://custom.auth/token"); + assert_eq!(provider.scopes, vec!["openid"]); + } + + #[test] + fn test_inject_auth() { + let provider = CodexProvider::new(None, None, None, None); + let mut headers = HeaderMap::new(); + provider.inject_auth(&mut headers, "test-token").unwrap(); + assert_eq!( + headers.get("authorization").unwrap().to_str().unwrap(), + "Bearer test-token" + ); + } + + #[test] + fn test_prepare_request_body_translates_chat() { + let provider = CodexProvider::new(None, None, None, None); + let tokens = OAuthTokens { + access_token: "t".to_string(), + refresh_token: None, + id_token: None, + expires_at: None, + account_id: None, + extra: BTreeMap::new(), + }; + let body = serde_json::json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}] + }); + let result = provider + .prepare_request_body(body.to_string().as_bytes(), &tokens) + .unwrap(); + assert!(result.is_some()); + let parsed: serde_json::Value = serde_json::from_slice(&result.unwrap()).unwrap(); + assert!(parsed.get("input").is_some()); + assert!(parsed.get("messages").is_none()); + } + + #[test] + fn test_prepare_request_body_passthrough_non_chat() { + let provider = CodexProvider::new(None, None, None, None); + let tokens = OAuthTokens { + access_token: "t".to_string(), + refresh_token: None, + id_token: None, + expires_at: None, + account_id: None, + extra: BTreeMap::new(), + }; + // No "messages" field → passthrough + let body = serde_json::json!({"model": "gpt-4o", "input": "hello"}); + let result = provider + .prepare_request_body(body.to_string().as_bytes(), &tokens) + .unwrap(); + assert!(result.is_none()); + + // Non-JSON → passthrough + let result = provider + .prepare_request_body(b"not json", &tokens) + .unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_rewrite_path_chat_completions() { + let provider = CodexProvider::new(None, None, None, None); + assert_eq!( + provider.rewrite_request_path("/v1/chat/completions"), + Some("/responses".to_string()) + ); + } + + #[test] + fn test_rewrite_path_other() { + let provider = CodexProvider::new(None, None, None, None); + assert_eq!(provider.rewrite_request_path("/v1/models"), None); + assert_eq!(provider.rewrite_request_path("/v1/responses"), None); + assert_eq!(provider.rewrite_request_path("/responses"), None); + } + + #[test] + fn test_needs_translation_chat_completions() { + let provider = CodexProvider::new(None, None, None, None); + assert!(provider.needs_response_translation("/v1/chat/completions")); + } + + #[test] + fn test_needs_translation_other() { + let provider = CodexProvider::new(None, None, None, None); + assert!(!provider.needs_response_translation("/v1/models")); + assert!(!provider.needs_response_translation("/v1/responses")); + } + + #[test] + fn test_upstream_url_chatgpt() { + let provider = CodexProvider::new(None, None, None, None); + let tokens = OAuthTokens { + access_token: "t".to_string(), + refresh_token: None, + id_token: None, + expires_at: None, + account_id: None, + extra: BTreeMap::new(), + }; + assert_eq!( + provider.upstream_url(&tokens), + Some("https://chatgpt.com/backend-api/codex".to_string()) + ); + } + + #[test] + fn test_fixup_strips_model_prefix() { + let body = serde_json::to_vec(&serde_json::json!({ + "model": "openai/gpt-5.2-codex", + "input": [{"role": "user", "content": "hi"}] + })) + .unwrap(); + let result = fixup_for_chatgpt_backend(&body); + let parsed: serde_json::Value = serde_json::from_slice(&result).unwrap(); + assert_eq!(parsed["model"], "gpt-5.2-codex"); + assert_eq!(parsed["store"], false); + } + + #[test] + fn test_fixup_sets_store_false() { + let body = serde_json::to_vec(&serde_json::json!({ + "model": "gpt-4o-mini", + "input": [{"role": "user", "content": "hi"}] + })) + .unwrap(); + let result = fixup_for_chatgpt_backend(&body); + let parsed: serde_json::Value = serde_json::from_slice(&result).unwrap(); + assert_eq!(parsed["model"], "gpt-4o-mini"); + assert_eq!(parsed["store"], false); + } +} diff --git a/src/oauth/mod.rs b/src/oauth/mod.rs new file mode 100644 index 0000000..58a4cd7 --- /dev/null +++ b/src/oauth/mod.rs @@ -0,0 +1,712 @@ +mod storage; + +pub mod codex; + +pub use storage::TokenStorage; + +use async_trait::async_trait; +use axum::http::HeaderMap; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; +use std::fmt; +use std::sync::Arc; +use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error, info, warn}; + +/// Error type for OAuth operations. +#[derive(Debug, thiserror::Error)] +pub enum OAuthError { + #[error("login failed: {0}")] + LoginFailed(String), + + #[error("token refresh failed: {0}")] + RefreshFailed(String), + + #[error("no tokens available for provider '{0}'")] + NoTokens(String), + + #[error("token expired for provider '{0}'")] + TokenExpired(String), + + #[error("provider not found: {0}")] + ProviderNotFound(String), + + #[error("header error: {0}")] + HeaderError(String), + + #[error("http error: {0}")] + HttpError(#[from] reqwest::Error), + + #[error("io error: {0}")] + IoError(#[from] std::io::Error), + + #[error("json error: {0}")] + JsonError(#[from] serde_json::Error), + + #[error("storage error: {0}")] + StorageError(String), +} + +impl From for OAuthError { + fn from(e: axum::http::header::InvalidHeaderValue) -> Self { + OAuthError::HeaderError(e.to_string()) + } +} + +/// Tokens obtained from an OAuth provider. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OAuthTokens { + pub access_token: String, + pub refresh_token: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub id_token: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub expires_at: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub account_id: Option, + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub extra: BTreeMap, +} + +impl OAuthTokens { + pub fn is_expired(&self) -> bool { + self.expires_at + .is_some_and(|exp| exp <= Utc::now()) + } + + pub fn expires_in_secs(&self) -> Option { + self.expires_at + .map(|exp| (exp - Utc::now()).num_seconds()) + } +} + +/// The core trait that each OAuth provider implements. +#[async_trait] +pub trait OAuthProvider: Send + Sync + fmt::Debug { + /// Unique identifier (e.g., "codex"). + fn id(&self) -> &str; + + /// Display name (e.g., "Codex (OpenAI)"). + fn display_name(&self) -> &str; + + /// Execute browser-based OAuth login flow. + async fn login_browser(&self, callback_port: u16) -> Result; + + /// Execute headless login flow (device code or copy/paste URL). + async fn login_headless(&self) -> Result; + + /// Refresh the access token using the refresh token. + async fn refresh(&self, refresh_token: &str) -> Result; + + /// Inject provider-specific auth headers into the request. + fn inject_auth(&self, headers: &mut HeaderMap, access_token: &str) -> Result<(), OAuthError>; + + /// Optionally transform the request body for provider-specific formats. + /// Returns None for pass-through (Codex). + fn prepare_request_body( + &self, + _body: &[u8], + _tokens: &OAuthTokens, + ) -> Result>, OAuthError> { + Ok(None) + } + + /// Resolve the upstream URL for this provider. + /// Returns None to use the configured [upstream] URL (Codex). + fn upstream_url(&self, _tokens: &OAuthTokens) -> Option { + None + } + + /// Whether this provider supports device code flow. + fn supports_device_code(&self) -> bool { + false + } + + /// Whether this provider supports headless copy/paste URL fallback. + fn supports_headless_url(&self) -> bool { + false + } + + /// Enrich tokens with provider-specific state if missing (e.g., project ID discovery). + /// Returns `Some(enriched)` if tokens were updated, `None` if no changes needed. + /// Called before `prepare_request_body` to ensure tokens are ready for use. + async fn enrich_tokens(&self, _tokens: &OAuthTokens) -> Result, OAuthError> { + Ok(None) + } + + /// Optionally rewrite the request path (e.g., `/v1/chat/completions` → `/v1/responses`). + /// Returns `None` to use the original path unchanged. + fn rewrite_request_path(&self, _path: &str) -> Option { + None + } + + /// Whether responses from the upstream need to be translated back + /// to match the original request format. + fn needs_response_translation(&self, _original_path: &str) -> bool { + false + } + + /// What format the upstream response is in, for translation purposes. + /// Returns `None` if no translation is needed (passthrough). + fn response_format(&self, _original_path: &str) -> Option { + None + } +} + +/// The format of upstream API responses, used to select the correct translator. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ResponseFormat { + /// OpenAI Responses API → translate to chat.completion format + ResponsesApi, + /// Google Gemini SSE → translate to chat.completion format + GeminiSse, +} + +/// Configuration for an OAuth provider from TOML. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OAuthProviderConfig { + pub provider: String, + #[serde(default = "default_true")] + pub enabled: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub auth_url: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub token_url: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scopes: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub callback_port: Option, +} + +fn default_true() -> bool { + true +} + +const CALLBACK_BIND_HOST_ENV: &str = "CLAWSHELL_OAUTH_CALLBACK_HOST"; +const DEFAULT_CALLBACK_BIND_HOST: &str = "127.0.0.1"; + +pub(crate) fn callback_bind_addr(port: u16) -> String { + let host = std::env::var(CALLBACK_BIND_HOST_ENV) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| DEFAULT_CALLBACK_BIND_HOST.to_string()); + format!("{host}:{port}") +} + +/// Manages multiple OAuth providers, their tokens, and per-provider refresh tasks. +#[derive(Debug)] +pub struct OAuthRegistry { + providers: BTreeMap>, + tokens: Arc>>, + storage: TokenStorage, +} + +impl OAuthRegistry { + pub fn new(storage: TokenStorage) -> Self { + Self { + providers: BTreeMap::new(), + tokens: Arc::new(RwLock::new(BTreeMap::new())), + storage, + } + } + + pub fn register(&mut self, provider: Arc) { + let id = provider.id().to_string(); + debug!(provider = %id, "Registering OAuth provider"); + self.providers.insert(id, provider); + } + + /// Load persisted tokens from disk for all registered providers. + pub async fn load_tokens(&self) -> Result<(), OAuthError> { + let mut tokens = self.tokens.write().await; + for id in self.providers.keys() { + match self.storage.load(id) { + Ok(Some(t)) => { + info!(provider = %id, expired = t.is_expired(), "Loaded OAuth tokens from disk"); + tokens.insert(id.clone(), t); + } + Ok(None) => { + debug!(provider = %id, "No persisted tokens found"); + } + Err(e) => { + warn!(provider = %id, error = %e, "Failed to load persisted tokens"); + } + } + } + Ok(()) + } + + /// Get the current access token for a provider, refreshing if expired. + pub async fn current_access_token(&self, provider_id: &str) -> Result { + { + let tokens = self.tokens.read().await; + if let Some(t) = tokens.get(provider_id) { + if !t.is_expired() { + return Ok(t.access_token.clone()); + } + } + } + // Token is expired or missing — try refreshing + self.refresh(provider_id).await?; + let tokens = self.tokens.read().await; + tokens + .get(provider_id) + .map(|t| t.access_token.clone()) + .ok_or_else(|| OAuthError::NoTokens(provider_id.to_string())) + } + + /// Inject auth headers for the given provider. + pub async fn inject_auth( + &self, + provider_id: &str, + headers: &mut HeaderMap, + ) -> Result<(), OAuthError> { + let token = self.current_access_token(provider_id).await?; + let provider = self + .providers + .get(provider_id) + .ok_or_else(|| OAuthError::ProviderNotFound(provider_id.to_string()))?; + provider.inject_auth(headers, &token) + } + + /// Prepare the request body for the given provider. + /// Calls `enrich_tokens` first to ensure provider-specific state is populated. + pub async fn prepare_request_body( + &self, + provider_id: &str, + body: &[u8], + ) -> Result>, OAuthError> { + let provider = self + .providers + .get(provider_id) + .ok_or_else(|| OAuthError::ProviderNotFound(provider_id.to_string()))?; + + // Enrich tokens on-demand if the provider needs it (e.g., project_id discovery) + { + let tokens = self.tokens.read().await; + let t = tokens + .get(provider_id) + .ok_or_else(|| OAuthError::NoTokens(provider_id.to_string()))?; + if let Some(enriched) = provider.enrich_tokens(t).await? { + drop(tokens); + info!(provider = %provider_id, "Enriched OAuth tokens with provider-specific state"); + if let Err(e) = self.storage.save(provider_id, &enriched) { + warn!(provider = %provider_id, error = %e, "Failed to persist enriched tokens"); + } + self.tokens + .write() + .await + .insert(provider_id.to_string(), enriched); + } + } + + let tokens = self.tokens.read().await; + let t = tokens + .get(provider_id) + .ok_or_else(|| OAuthError::NoTokens(provider_id.to_string()))?; + provider.prepare_request_body(body, t) + } + + /// Resolve the upstream URL for the given provider. + pub async fn upstream_url( + &self, + provider_id: &str, + ) -> Result, OAuthError> { + let provider = self + .providers + .get(provider_id) + .ok_or_else(|| OAuthError::ProviderNotFound(provider_id.to_string()))?; + let tokens = self.tokens.read().await; + let t = tokens + .get(provider_id) + .ok_or_else(|| OAuthError::NoTokens(provider_id.to_string()))?; + Ok(provider.upstream_url(t)) + } + + /// Refresh the access token for a specific provider. + pub async fn refresh(&self, provider_id: &str) -> Result<(), OAuthError> { + let provider = self + .providers + .get(provider_id) + .ok_or_else(|| OAuthError::ProviderNotFound(provider_id.to_string()))?; + + let refresh_token = { + let tokens = self.tokens.read().await; + tokens + .get(provider_id) + .and_then(|t| t.refresh_token.clone()) + .ok_or_else(|| { + OAuthError::RefreshFailed(format!( + "no refresh token for provider '{provider_id}'" + )) + })? + }; + + info!(provider = %provider_id, "Refreshing OAuth access token"); + let new_tokens = provider.refresh(&refresh_token).await?; + self.storage + .save(provider_id, &new_tokens) + .map_err(|e| OAuthError::StorageError(e.to_string()))?; + self.tokens + .write() + .await + .insert(provider_id.to_string(), new_tokens); + info!(provider = %provider_id, "OAuth token refreshed successfully"); + Ok(()) + } + + /// Store tokens after a successful login (called from onboard flow). + pub async fn store_tokens( + &self, + provider_id: &str, + tokens: OAuthTokens, + ) -> Result<(), OAuthError> { + self.storage + .save(provider_id, &tokens) + .map_err(|e| OAuthError::StorageError(e.to_string()))?; + self.tokens + .write() + .await + .insert(provider_id.to_string(), tokens); + Ok(()) + } + + /// Spawn background refresh tasks for all providers with tokens. + pub fn spawn_refresh_tasks(&self, cancel: CancellationToken) { + let tokens = Arc::clone(&self.tokens); + for (id, provider) in &self.providers { + let id = id.clone(); + let provider = Arc::clone(provider); + let tokens = Arc::clone(&tokens); + let storage = self.storage.clone(); + let cancel = cancel.clone(); + + tokio::spawn(async move { + loop { + let sleep_secs = { + let guard = tokens.read().await; + match guard.get(&id) { + Some(t) => { + let remaining = t.expires_in_secs().unwrap_or(3600); + // Refresh at 75% of TTL, minimum 60 seconds + (remaining * 3 / 4).max(60) + } + None => 3600, // no tokens yet, check hourly + } + }; + + debug!(provider = %id, sleep_secs, "OAuth refresh task sleeping"); + + tokio::select! { + _ = cancel.cancelled() => { + info!(provider = %id, "OAuth refresh task cancelled"); + return; + } + _ = tokio::time::sleep(std::time::Duration::from_secs(sleep_secs as u64)) => {} + } + + let refresh_token = { + let guard = tokens.read().await; + guard + .get(&id) + .and_then(|t| t.refresh_token.clone()) + }; + + let Some(refresh_token) = refresh_token else { + debug!(provider = %id, "No refresh token available, skipping refresh"); + continue; + }; + + match provider.refresh(&refresh_token).await { + Ok(new_tokens) => { + if let Err(e) = storage.save(&id, &new_tokens) { + error!(provider = %id, error = %e, "Failed to persist refreshed tokens"); + } + tokens.write().await.insert(id.clone(), new_tokens); + info!(provider = %id, "Background token refresh successful"); + } + Err(e) => { + error!(provider = %id, error = %e, "Background token refresh failed"); + } + } + } + }); + } + } + + pub fn rewrite_request_path( + &self, + provider_id: &str, + path: &str, + ) -> Result, OAuthError> { + let provider = self + .providers + .get(provider_id) + .ok_or_else(|| OAuthError::ProviderNotFound(provider_id.to_string()))?; + Ok(provider.rewrite_request_path(path)) + } + + pub fn needs_response_translation( + &self, + provider_id: &str, + original_path: &str, + ) -> Result { + let provider = self + .providers + .get(provider_id) + .ok_or_else(|| OAuthError::ProviderNotFound(provider_id.to_string()))?; + Ok(provider.needs_response_translation(original_path)) + } + + pub fn response_format( + &self, + provider_id: &str, + original_path: &str, + ) -> Result, OAuthError> { + let provider = self + .providers + .get(provider_id) + .ok_or_else(|| OAuthError::ProviderNotFound(provider_id.to_string()))?; + Ok(provider.response_format(original_path)) + } + + pub fn has_provider(&self, id: &str) -> bool { + self.providers.contains_key(id) + } + + pub fn provider_ids(&self) -> Vec { + self.providers.keys().cloned().collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Debug)] + struct MockProvider { + id: String, + } + + #[async_trait] + impl OAuthProvider for MockProvider { + fn id(&self) -> &str { + &self.id + } + fn display_name(&self) -> &str { + "Mock Provider" + } + async fn login_browser(&self, _callback_port: u16) -> Result { + Ok(OAuthTokens { + access_token: "mock-access".to_string(), + refresh_token: Some("mock-refresh".to_string()), + id_token: None, + expires_at: Some(Utc::now() + chrono::Duration::hours(1)), + account_id: None, + extra: BTreeMap::new(), + }) + } + async fn login_headless(&self) -> Result { + self.login_browser(0).await + } + async fn refresh(&self, _refresh_token: &str) -> Result { + Ok(OAuthTokens { + access_token: "refreshed-access".to_string(), + refresh_token: Some("new-refresh".to_string()), + id_token: None, + expires_at: Some(Utc::now() + chrono::Duration::hours(1)), + account_id: None, + extra: BTreeMap::new(), + }) + } + fn inject_auth( + &self, + headers: &mut HeaderMap, + access_token: &str, + ) -> Result<(), OAuthError> { + headers.insert( + axum::http::header::AUTHORIZATION, + format!("Bearer {access_token}").parse()?, + ); + Ok(()) + } + } + + #[test] + fn test_tokens_not_expired() { + let tokens = OAuthTokens { + access_token: "test".to_string(), + refresh_token: None, + id_token: None, + expires_at: Some(Utc::now() + chrono::Duration::hours(1)), + account_id: None, + extra: BTreeMap::new(), + }; + assert!(!tokens.is_expired()); + } + + #[test] + fn test_tokens_expired() { + let tokens = OAuthTokens { + access_token: "test".to_string(), + refresh_token: None, + id_token: None, + expires_at: Some(Utc::now() - chrono::Duration::hours(1)), + account_id: None, + extra: BTreeMap::new(), + }; + assert!(tokens.is_expired()); + } + + #[test] + fn test_tokens_no_expiry() { + let tokens = OAuthTokens { + access_token: "test".to_string(), + refresh_token: None, + id_token: None, + expires_at: None, + account_id: None, + extra: BTreeMap::new(), + }; + assert!(!tokens.is_expired()); + assert!(tokens.expires_in_secs().is_none()); + } + + #[tokio::test] + async fn test_registry_register_and_access() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + let mut registry = OAuthRegistry::new(storage); + + let provider = Arc::new(MockProvider { + id: "mock".to_string(), + }); + registry.register(provider); + assert!(registry.has_provider("mock")); + assert!(!registry.has_provider("other")); + } + + #[tokio::test] + async fn test_registry_store_and_retrieve_tokens() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + let mut registry = OAuthRegistry::new(storage); + + let provider = Arc::new(MockProvider { + id: "mock".to_string(), + }); + registry.register(provider); + + let tokens = OAuthTokens { + access_token: "test-access".to_string(), + refresh_token: Some("test-refresh".to_string()), + id_token: None, + expires_at: Some(Utc::now() + chrono::Duration::hours(1)), + account_id: None, + extra: BTreeMap::new(), + }; + registry.store_tokens("mock", tokens).await.unwrap(); + + let token = registry.current_access_token("mock").await.unwrap(); + assert_eq!(token, "test-access"); + } + + #[tokio::test] + async fn test_registry_refresh_expired_token() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + let mut registry = OAuthRegistry::new(storage); + + let provider = Arc::new(MockProvider { + id: "mock".to_string(), + }); + registry.register(provider); + + // Store an expired token + let tokens = OAuthTokens { + access_token: "expired-access".to_string(), + refresh_token: Some("test-refresh".to_string()), + id_token: None, + expires_at: Some(Utc::now() - chrono::Duration::hours(1)), + account_id: None, + extra: BTreeMap::new(), + }; + registry.store_tokens("mock", tokens).await.unwrap(); + + // Should auto-refresh + let token = registry.current_access_token("mock").await.unwrap(); + assert_eq!(token, "refreshed-access"); + } + + #[tokio::test] + async fn test_registry_inject_auth() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + let mut registry = OAuthRegistry::new(storage); + + let provider = Arc::new(MockProvider { + id: "mock".to_string(), + }); + registry.register(provider); + + let tokens = OAuthTokens { + access_token: "inject-test".to_string(), + refresh_token: Some("r".to_string()), + id_token: None, + expires_at: Some(Utc::now() + chrono::Duration::hours(1)), + account_id: None, + extra: BTreeMap::new(), + }; + registry.store_tokens("mock", tokens).await.unwrap(); + + let mut headers = HeaderMap::new(); + registry.inject_auth("mock", &mut headers).await.unwrap(); + assert_eq!( + headers.get("authorization").unwrap().to_str().unwrap(), + "Bearer inject-test" + ); + } + + #[tokio::test] + async fn test_registry_provider_not_found() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + let registry = OAuthRegistry::new(storage); + + let result = registry.current_access_token("nonexistent").await; + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("provider not found")); + } + + #[tokio::test] + async fn test_registry_load_tokens_from_disk() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + + // Pre-persist tokens + let tokens = OAuthTokens { + access_token: "disk-token".to_string(), + refresh_token: Some("disk-refresh".to_string()), + id_token: None, + expires_at: Some(Utc::now() + chrono::Duration::hours(1)), + account_id: None, + extra: BTreeMap::new(), + }; + storage.save("mock", &tokens).unwrap(); + + let mut registry = OAuthRegistry::new(storage); + let provider = Arc::new(MockProvider { + id: "mock".to_string(), + }); + registry.register(provider); + registry.load_tokens().await.unwrap(); + + let token = registry.current_access_token("mock").await.unwrap(); + assert_eq!(token, "disk-token"); + } +} diff --git a/src/oauth/storage.rs b/src/oauth/storage.rs new file mode 100644 index 0000000..c5f5ba6 --- /dev/null +++ b/src/oauth/storage.rs @@ -0,0 +1,190 @@ +use super::OAuthTokens; +use std::path::PathBuf; +use tracing::debug; + +/// Per-provider token persistence under a directory (e.g. `/etc/clawshell/oauth/`). +#[derive(Debug, Clone)] +pub struct TokenStorage { + dir: PathBuf, +} + +impl Default for TokenStorage { + fn default() -> Self { + Self { + dir: PathBuf::from("/etc/clawshell/oauth"), + } + } +} + +impl TokenStorage { + pub fn new(dir: PathBuf) -> Self { + Self { dir } + } + + pub fn dir(&self) -> &PathBuf { + &self.dir + } + + fn token_path(&self, provider_id: &str) -> PathBuf { + self.dir.join(format!("{provider_id}.json")) + } + + /// Save tokens for a provider, creating the directory if needed. + pub fn save(&self, provider_id: &str, tokens: &OAuthTokens) -> Result<(), std::io::Error> { + std::fs::create_dir_all(&self.dir)?; + let path = self.token_path(provider_id); + let content = serde_json::to_string_pretty(tokens) + .map_err(|e| std::io::Error::other(format!("failed to serialize tokens: {e}")))?; + std::fs::write(&path, content)?; + + // Set file permissions to 0600 (owner read/write only) + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?; + } + + debug!(provider = %provider_id, path = %path.display(), "OAuth tokens saved"); + Ok(()) + } + + /// Load tokens for a provider, returning None if the file doesn't exist. + pub fn load(&self, provider_id: &str) -> Result, std::io::Error> { + let path = self.token_path(provider_id); + if !path.exists() { + return Ok(None); + } + let content = std::fs::read_to_string(&path)?; + let tokens: OAuthTokens = serde_json::from_str(&content) + .map_err(|e| std::io::Error::other(format!("failed to parse tokens: {e}")))?; + debug!(provider = %provider_id, path = %path.display(), "OAuth tokens loaded"); + Ok(Some(tokens)) + } + + /// Remove tokens for a provider. + pub fn remove(&self, provider_id: &str) -> Result<(), std::io::Error> { + let path = self.token_path(provider_id); + if path.exists() { + std::fs::remove_file(&path)?; + debug!(provider = %provider_id, path = %path.display(), "OAuth tokens removed"); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + use std::collections::BTreeMap; + + fn test_tokens() -> OAuthTokens { + OAuthTokens { + access_token: "access-123".to_string(), + refresh_token: Some("refresh-456".to_string()), + id_token: None, + expires_at: Some(Utc::now() + chrono::Duration::hours(1)), + account_id: Some("user@test.com".to_string()), + extra: BTreeMap::new(), + } + } + + #[test] + fn test_save_and_load() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + + let tokens = test_tokens(); + storage.save("test-provider", &tokens).unwrap(); + + let loaded = storage.load("test-provider").unwrap().unwrap(); + assert_eq!(loaded.access_token, "access-123"); + assert_eq!(loaded.refresh_token.as_deref(), Some("refresh-456")); + assert_eq!(loaded.account_id.as_deref(), Some("user@test.com")); + } + + #[test] + fn test_load_nonexistent() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + + let loaded = storage.load("nonexistent").unwrap(); + assert!(loaded.is_none()); + } + + #[test] + fn test_remove() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + + let tokens = test_tokens(); + storage.save("removable", &tokens).unwrap(); + assert!(storage.load("removable").unwrap().is_some()); + + storage.remove("removable").unwrap(); + assert!(storage.load("removable").unwrap().is_none()); + } + + #[test] + fn test_remove_nonexistent() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + // Should not error + storage.remove("nonexistent").unwrap(); + } + + #[test] + fn test_creates_directory() { + let dir = tempfile::tempdir().unwrap(); + let nested = dir.path().join("a").join("b").join("c"); + let storage = TokenStorage::new(nested.clone()); + + let tokens = test_tokens(); + storage.save("test", &tokens).unwrap(); + assert!(nested.join("test.json").exists()); + } + + #[test] + fn test_tokens_with_extra_fields() { + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + + let mut tokens = test_tokens(); + tokens.extra.insert( + "project_id".to_string(), + serde_json::json!("proj-abc-123"), + ); + tokens + .extra + .insert("tier".to_string(), serde_json::json!("production")); + + storage.save("antigravity", &tokens).unwrap(); + + let loaded = storage.load("antigravity").unwrap().unwrap(); + assert_eq!( + loaded.extra.get("project_id").unwrap().as_str().unwrap(), + "proj-abc-123" + ); + assert_eq!( + loaded.extra.get("tier").unwrap().as_str().unwrap(), + "production" + ); + } + + #[cfg(unix)] + #[test] + fn test_file_permissions() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempfile::tempdir().unwrap(); + let storage = TokenStorage::new(dir.path().to_path_buf()); + + let tokens = test_tokens(); + storage.save("perms-test", &tokens).unwrap(); + + let path = dir.path().join("perms-test.json"); + let metadata = std::fs::metadata(path).unwrap(); + let mode = metadata.permissions().mode() & 0o777; + assert_eq!(mode, 0o600); + } +} diff --git a/src/onboard/config_render.rs b/src/onboard/config_render.rs index 061b305..dfe30f8 100644 --- a/src/onboard/config_render.rs +++ b/src/onboard/config_render.rs @@ -1,4 +1,4 @@ -use super::types::{OnboardConfig, OnboardEmailMode}; +use super::types::{OnboardAuthMethod, OnboardConfig, OnboardEmailMode}; /// Return the default OpenClaw config path. pub fn default_openclaw_config_path() -> String { @@ -11,6 +11,47 @@ pub fn default_openclaw_config_path() -> String { /// Generate the ClawShell TOML configuration content with the given key mapping. pub fn generate_clawshell_config(config: &OnboardConfig) -> String { + let key_section = match &config.auth_method { + OnboardAuthMethod::OAuth { provider_id } => { + format!( + r#"[[keys]] +virtual_key = {virtual_key} +provider = {provider} +auth = "oauth" +oauth_provider = {oauth_provider} +"#, + virtual_key = toml_string(&config.virtual_api_key), + provider = toml_string(&config.provider), + oauth_provider = toml_string(provider_id), + ) + } + OnboardAuthMethod::StaticKey => { + format!( + r#"[[keys]] +virtual_key = {virtual_key} +real_key = {real_key} +provider = {provider} +"#, + virtual_key = toml_string(&config.virtual_api_key), + real_key = toml_string(&config.real_api_key), + provider = toml_string(&config.provider), + ) + } + }; + + let oauth_providers_section = match &config.auth_method { + OnboardAuthMethod::OAuth { provider_id } => { + format!( + r#" +[[oauth_providers]] +provider = {provider_id} +"#, + provider_id = toml_string(provider_id), + ) + } + OnboardAuthMethod::StaticKey => String::new(), + }; + let mut output = format!( r#"# ClawShell Configuration version = "{version}" @@ -25,11 +66,7 @@ openai_base_url = "https://api.openai.com" openrouter_base_url = "https://openrouter.ai/api" anthropic_base_url = "https://api.anthropic.com" -[[keys]] -virtual_key = {virtual_key} -real_key = {real_key} -provider = {provider} -[dlp] +{key_section}[dlp] scan_responses = true patterns = [ {{ name = "ssn", regex = '\\b\\d{{3}}-\\d{{2}}-\\d{{4}}\\b', action = "redact" }}, @@ -38,13 +75,12 @@ patterns = [ {{ name = "mastercard", regex = '\\b5[1-5][0-9]{{14}}\\b', action = "redact" }}, {{ name = "amex_card", regex = '\\b3[47][0-9]{{13}}\\b', action = "redact" }}, ] -"#, +{oauth_providers_section}"#, version = env!("CARGO_PKG_VERSION"), host = config.server_host, port = config.server_port, - virtual_key = toml_string(&config.virtual_api_key), - real_key = toml_string(&config.real_api_key), - provider = toml_string(&config.provider), + key_section = key_section, + oauth_providers_section = oauth_providers_section, ); if let Some(email) = &config.email { diff --git a/src/onboard/interactive.rs b/src/onboard/interactive.rs index e1afcfe..42823db 100644 --- a/src/onboard/interactive.rs +++ b/src/onboard/interactive.rs @@ -1,6 +1,6 @@ use super::config_render::default_openclaw_config_path; use super::credentials::detect_openclaw_api_key_for_provider; -use super::types::{OnboardConfig, OnboardEmailConfig, OnboardEmailMode}; +use super::types::{OnboardAuthMethod, OnboardConfig, OnboardEmailConfig, OnboardEmailMode}; use crate::email::{EmailAccountCredentials, ImapEmailService}; use crate::tui; @@ -69,6 +69,14 @@ fn load_existing_config_from_vfs(config_dir: &VfsPath) -> Option .get("openclaw_config_path") .and_then(|v| v.as_str()) .map(String::from); + existing.auth_method = json + .get("auth_method") + .and_then(|v| v.as_str()) + .map(String::from); + existing.oauth_provider = json + .get("oauth_provider") + .and_then(|v| v.as_str()) + .map(String::from); } // Read clawshell.toml for server host/port and optional Email settings @@ -259,6 +267,8 @@ struct ExistingConfig { openclaw_config_path: Option, server_host: Option, server_port: Option, + auth_method: Option, + oauth_provider: Option, email_enabled: Option, email_mode: Option, email_sender_rules: Vec, @@ -278,6 +288,8 @@ impl ExistingConfig { || self.openclaw_config_path.is_some() || self.server_host.is_some() || self.server_port.is_some() + || self.auth_method.is_some() + || self.oauth_provider.is_some() || self.email_enabled.is_some() || self.email_mode.is_some() || !self.email_sender_rules.is_empty() @@ -308,51 +320,56 @@ fn mask_secret(secret: &str) -> String { } } -/// Collect all onboarding information using the TUI (interactive terminal prompts). -/// If a previous configuration exists, its values are used as defaults. -pub fn collect_onboard_config_tui() -> Result> { - let existing = load_existing_config(); +/// Run the OAuth login flow for the given provider, persisting tokens. +fn run_oauth_login(provider_id: &str) -> Result<(), Box> { + use crate::oauth::codex::CodexProvider; + use crate::oauth::{OAuthProvider, TokenStorage}; - if existing.is_some() { - tui::print_success("Existing configuration detected — using as defaults."); - println!(); - } - - let existing = existing.unwrap_or_default(); + let provider: Box = match provider_id { + "codex" => Box::new(CodexProvider::new(None, None, None, None)), + other => return Err(format!("unknown OAuth provider: {other}").into()), + }; - tui::print_section("API Configuration"); + let storage = TokenStorage::default(); - // Provider selection — if existing, reorder so the existing choice is first - let provider_options = match existing.provider.as_deref() { - Some("anthropic") => vec!["Anthropic", "OpenAI", "OpenRouter"], - Some("openrouter") => vec!["OpenRouter", "OpenAI", "Anthropic"], - _ => vec!["OpenAI", "OpenRouter", "Anthropic"], + // Called from within #[tokio::main], so use block_in_place to avoid + // "Cannot start a runtime from within a runtime" panic. + let run_async = |fut: std::pin::Pin + Send>>| { + let handle = tokio::runtime::Handle::current(); + tokio::task::block_in_place(|| handle.block_on(fut)) }; - let provider_choice = tui::prompt_select("Select a model provider", provider_options)?; - let provider = match provider_choice { - "Anthropic" => "anthropic".to_string(), - "OpenRouter" => "openrouter".to_string(), - _ => "openai".to_string(), + + let tokens = if provider.supports_device_code() { + tui::print_info("Flow", "device code (no browser required)"); + run_async(Box::pin(provider.login_headless()))? + } else if provider.supports_headless_url() { + tui::print_info("Flow", "headless (copy URL, paste code)"); + run_async(Box::pin(provider.login_headless()))? + } else { + tui::print_info("Flow", "browser login"); + tui::print_warning("A browser window will open for you to authorize access."); + run_async(Box::pin(provider.login_browser(8400)))? }; - // Model name — use existing model or provider-specific default - let default_model = existing - .model - .as_deref() - .unwrap_or(match provider.as_str() { - "anthropic" => "claude-sonnet-4-5-20250929", - "openai" => "gpt-5.2-chat-latest", - "openrouter" => "openrouter/auto", - _ => unreachable!(), - }); - let model = tui::prompt_text("Enter the model name", Some(default_model))?; + storage.save(provider_id, &tokens)?; + tui::print_success("OAuth login successful — tokens saved."); + if let Some(acct) = tokens.account_id.as_deref() { + tui::print_info("Account", acct); + } - // Real API key — if ClawShell already has one, use it; otherwise try detecting from OpenClaw + Ok(()) +} + +/// Collect a static API key from the user (original flow). +fn collect_static_api_key( + provider: &str, + existing: &ExistingConfig, +) -> Result> { let is_first_onboard = existing.real_api_key.is_none(); let effective_existing_key = if !is_first_onboard { existing.real_api_key.clone() } else { - let key = detect_openclaw_api_key_for_provider(&provider); + let key = detect_openclaw_api_key_for_provider(provider); if key.is_some() { tui::print_warning( "An API key was detected from your OpenClaw config. \ @@ -364,15 +381,12 @@ pub fn collect_onboard_config_tui() -> Result Result Result> { + let existing = load_existing_config(); + + if existing.is_some() { + tui::print_success("Existing configuration detected — using as defaults."); + println!(); + } + + let existing = existing.unwrap_or_default(); + + tui::print_section("API Configuration"); + + // Provider selection + const MENU_OPENAI: &str = "OpenAI"; + const MENU_OPENROUTER: &str = "OpenRouter"; + const MENU_ANTHROPIC: &str = "Anthropic"; + const MENU_CODEX: &str = "Codex / ChatGPT (OAuth)"; + let all_options = [MENU_OPENAI, MENU_OPENROUTER, MENU_ANTHROPIC, MENU_CODEX]; + + // Reorder so the existing choice appears first + let preferred = match ( + existing.auth_method.as_deref(), + existing.oauth_provider.as_deref(), + existing.provider.as_deref(), + ) { + (Some("oauth"), Some("codex"), _) | (Some("oauth"), _, _) => Some(MENU_CODEX), + (_, _, Some("anthropic")) => Some(MENU_ANTHROPIC), + (_, _, Some("openrouter")) => Some(MENU_OPENROUTER), + (_, _, Some("openai")) => Some(MENU_OPENAI), + _ => None, + }; + let provider_options: Vec<&str> = if let Some(first) = preferred { + std::iter::once(first) + .chain(all_options.iter().copied().filter(|o| *o != first)) + .collect() + } else { + all_options.to_vec() + }; + + let provider_choice = tui::prompt_select("Select a model provider", provider_options)?; + + let (provider, auth_method) = match provider_choice { + MENU_ANTHROPIC => ("anthropic".to_string(), OnboardAuthMethod::StaticKey), + MENU_OPENROUTER => ("openrouter".to_string(), OnboardAuthMethod::StaticKey), + MENU_CODEX => ( + "openai".to_string(), + OnboardAuthMethod::OAuth { + provider_id: "codex".to_string(), + }, + ), + _ => ("openai".to_string(), OnboardAuthMethod::StaticKey), + }; + + // Model name — use existing model or provider/auth-specific default + let default_model = existing.model.as_deref().unwrap_or(match provider_choice { + MENU_ANTHROPIC => "claude-sonnet-4-5-20250929", + MENU_OPENROUTER => "openrouter/auto", + MENU_CODEX => "gpt-5.2-chat-latest", + _ => "gpt-5.2-chat-latest", // OpenAI default + }); + let model = tui::prompt_text("Enter the model name", Some(default_model))?; + + let real_api_key = match &auth_method { + OnboardAuthMethod::OAuth { provider_id } => { + // OAuth flow — run device code or browser login + tui::print_section("OAuth Login"); + tui::print_info("OAuth provider", provider_id); + + run_oauth_login(provider_id)?; + + // No static API key needed for OAuth + String::new() + } + OnboardAuthMethod::StaticKey => { + // Static key flow — same as before + collect_static_api_key(&provider, &existing)? + } + }; + // Virtual API key let fallback_virtual_key = format!("{{clawshell-virtual-key-{}}}", provider); let default_virtual = existing @@ -705,6 +803,7 @@ pub fn collect_onboard_config_tui() -> Result OnboardConfig { OnboardConfig { provider: "openai".to_string(), model: "gpt-5.2".to_string(), + auth_method: super::types::OnboardAuthMethod::StaticKey, real_api_key: "sk-real-key-123".to_string(), virtual_api_key: "{clawshell-virtual-key-openai}".to_string(), openclaw_config_path: PathBuf::from("/tmp/test-openclaw.json"), diff --git a/src/onboard/types.rs b/src/onboard/types.rs index bb18fa4..6c3933d 100644 --- a/src/onboard/types.rs +++ b/src/onboard/types.rs @@ -41,11 +41,26 @@ pub struct OpenclawFileRemovalPreview { pub removals: Vec, } +/// Authentication method chosen during onboarding. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub enum OnboardAuthMethod { + /// Static API key (the traditional approach). + #[default] + StaticKey, + /// OAuth provider supplies access tokens at runtime. + OAuth { + /// Provider identifier, e.g. "codex". + provider_id: String, + }, +} + /// Collected onboarding configuration from user prompts. #[derive(Debug, Clone)] pub struct OnboardConfig { pub provider: String, pub model: String, + pub auth_method: OnboardAuthMethod, + /// Set for `StaticKey`; empty for `OAuth`. pub real_api_key: String, pub virtual_api_key: String, pub openclaw_config_path: PathBuf, diff --git a/src/openclaw_cli.rs b/src/openclaw_cli.rs index 24bc57a..349157f 100644 --- a/src/openclaw_cli.rs +++ b/src/openclaw_cli.rs @@ -692,6 +692,7 @@ mod tests { onboard::OnboardConfig { provider: "openai".to_string(), model: "gpt-5".to_string(), + auth_method: onboard::OnboardAuthMethod::StaticKey, real_api_key: "real_key".to_string(), virtual_api_key: "virtual_key".to_string(), openclaw_config_path: PathBuf::from("/home/user/.openclaw/openclaw.json"), diff --git a/src/proxy.rs b/src/proxy.rs index f958a4f..4e9b77b 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -62,22 +62,7 @@ impl ProxyClient { "Preparing upstream request" ); - let mut req_headers = HeaderMap::new(); - for (name, value) in &headers { - let name_str = name.as_str().to_lowercase(); - // Skip hop-by-hop headers and the original auth header - if name_str == "host" - || name_str == "authorization" - || name_str == "connection" - || name_str == "content-length" - || name_str == "transfer-encoding" - || name_str == "x-api-key" - { - trace!(header = %name_str, "Skipping hop-by-hop/auth header"); - continue; - } - req_headers.insert(name.clone(), value.clone()); - } + let mut req_headers = filter_hop_by_hop_headers(&headers); trace!( forwarded_header_count = req_headers.len(), @@ -107,6 +92,74 @@ impl ProxyClient { } } + self.send_upstream(method, &upstream_url, req_headers, body) + .await + } + + /// Forward a request using OAuth-injected auth headers and optional overrides. + #[allow(clippy::too_many_arguments)] + pub async fn forward_oauth( + &self, + method: Method, + uri: &Uri, + original_headers: HeaderMap, + body: Bytes, + provider: Provider, + auth_headers: HeaderMap, + upstream_url_override: Option<&str>, + ) -> Result { + let upstream_url = if let Some(base) = upstream_url_override { + format!( + "{}{}", + base, + uri.path_and_query() + .map(|pq| pq.as_str()) + .unwrap_or(uri.path()) + ) + } else { + let base_url = self.upstream_urls.get(&provider).ok_or_else(|| { + ProxyError::Internal(format!("No upstream URL for provider {:?}", provider)) + })?; + format!( + "{}{}", + base_url, + uri.path_and_query() + .map(|pq| pq.as_str()) + .unwrap_or(uri.path()) + ) + }; + + debug!( + %upstream_url, + %method, + provider = ?provider, + body_size = body.len(), + "Preparing OAuth upstream request" + ); + + let mut req_headers = filter_hop_by_hop_headers(&original_headers); + + // Apply OAuth auth headers (these may include Authorization, x-goog-api-client, etc.) + for (name, value) in &auth_headers { + req_headers.insert(name.clone(), value.clone()); + } + + trace!( + forwarded_header_count = req_headers.len(), + "Filtered request headers (OAuth)" + ); + + self.send_upstream(method, &upstream_url, req_headers, body) + .await + } + + async fn send_upstream( + &self, + method: Method, + upstream_url: &str, + req_headers: HeaderMap, + body: Bytes, + ) -> Result { let reqwest_method = match method { Method::GET => reqwest::Method::GET, Method::POST => reqwest::Method::POST, @@ -124,7 +177,7 @@ impl ProxyClient { let upstream_resp = self .client - .request(reqwest_method, &upstream_url) + .request(reqwest_method, upstream_url) .headers(req_headers) .body(body) .send() @@ -137,7 +190,6 @@ impl ProxyClient { debug!( upstream_status = %status, - provider = ?provider, "Received upstream response" ); @@ -164,15 +216,9 @@ impl ProxyClient { let byte_stream = upstream_resp.bytes_stream().map_err(IoError::other); let body = Body::from_stream(byte_stream); - // Rebind the `status` var to clarify the type for human developer: - // it is guaranteed to be `StatusCode` due to the `.unwrap_or` in its assignment above. let status: StatusCode = status; let mut response = Response::builder().status(status); - // INVARIANT: the `status` variable is guaranteed to be `StatusCode`, - // so this `.unwrap` should never panic. *response.headers_mut().unwrap() = resp_headers; - // INVARIANT: the builder should always succeed since we just added a valid status code and headers, - // so this `.unwrap` should never panic. Ok(response.body(body).unwrap()) } else { // Buffer the full response @@ -186,18 +232,34 @@ impl ProxyClient { "Buffered upstream response body" ); - // Rebind the `status` var to clarify the type for human developer: - // it is guaranteed to be `StatusCode` due to the `.unwrap_or` in its assignment above. let status: StatusCode = status; let mut response = Response::builder().status(status); - // INVARIANT: the builder should always succeed since we just added a valid status code and headers, - // so this `.unwrap` should never panic. *response.headers_mut().unwrap() = resp_headers; Ok(response.body(Body::from(resp_body)).unwrap()) } } } +fn filter_hop_by_hop_headers(headers: &HeaderMap) -> HeaderMap { + let mut filtered = HeaderMap::new(); + for (name, value) in headers { + let name_str = name.as_str().to_lowercase(); + // Skip hop-by-hop headers and the original auth header + if name_str == "host" + || name_str == "authorization" + || name_str == "connection" + || name_str == "content-length" + || name_str == "transfer-encoding" + || name_str == "x-api-key" + { + trace!(header = %name_str, "Skipping hop-by-hop/auth header"); + continue; + } + filtered.insert(name.clone(), value.clone()); + } + filtered +} + #[derive(Debug)] pub enum ProxyError { Upstream(String), @@ -308,4 +370,19 @@ mod tests { let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); assert!(json["error"].as_str().unwrap().contains("TRACE")); } + + #[test] + fn test_filter_hop_by_hop_headers() { + let mut headers = HeaderMap::new(); + headers.insert("authorization", "Bearer vk-test".parse().unwrap()); + headers.insert("content-type", "application/json".parse().unwrap()); + headers.insert("host", "localhost".parse().unwrap()); + headers.insert("x-custom", "custom-value".parse().unwrap()); + + let filtered = filter_hop_by_hop_headers(&headers); + assert!(filtered.get("authorization").is_none()); + assert!(filtered.get("host").is_none()); + assert!(filtered.get("content-type").is_some()); + assert!(filtered.get("x-custom").is_some()); + } } diff --git a/src/translate.rs b/src/translate.rs new file mode 100644 index 0000000..1e547fc --- /dev/null +++ b/src/translate.rs @@ -0,0 +1,1180 @@ +use crate::dlp::DlpScanner; +use axum::body::Body; +use bytes::{Bytes, BytesMut}; +use futures_util::Stream; +use serde_json::Value; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; +use tracing::{debug, warn}; + +#[derive(Debug, thiserror::Error)] +pub enum TranslateError { + #[error("json error: {0}")] + Json(#[from] serde_json::Error), + + #[error("missing field: {0}")] + MissingField(&'static str), +} + +/// Fields that are compatible between chat/completions and responses API. +const PASSTHROUGH_FIELDS: &[&str] = &["model", "stream", "temperature", "top_p", "stop"]; + +/// Fields that must be stripped from chat/completions requests (not supported by responses API). +const STRIP_FIELDS: &[&str] = &[ + "frequency_penalty", + "presence_penalty", + "logprobs", + "top_logprobs", + "logit_bias", + "n", + "response_format", + "seed", + "service_tier", + "user", +]; + +/// Translate a `/v1/chat/completions` request body to `/v1/responses` format. +pub fn chat_completions_to_responses(body: &[u8]) -> Result, TranslateError> { + let req: Value = serde_json::from_slice(body)?; + let obj = req.as_object().ok_or(TranslateError::MissingField("root object"))?; + + let messages = obj + .get("messages") + .and_then(Value::as_array) + .ok_or(TranslateError::MissingField("messages"))?; + + let mut result = serde_json::Map::new(); + + // Separate system messages → instructions, rest → input + let mut system_parts: Vec<&str> = Vec::new(); + let mut input: Vec = Vec::new(); + + for msg in messages { + let role = msg.get("role").and_then(Value::as_str).unwrap_or(""); + if role == "system" { + if let Some(content) = msg.get("content").and_then(Value::as_str) { + system_parts.push(content); + } + } else { + input.push(convert_message_content(msg.clone())); + } + } + + // Codex responses API requires `instructions` even when empty + result.insert( + "instructions".to_string(), + Value::String(system_parts.join("\n")), + ); + result.insert("input".to_string(), Value::Array(input)); + + // Rename max_tokens → max_output_tokens + if let Some(max_tokens) = obj.get("max_tokens") { + result.insert("max_output_tokens".to_string(), max_tokens.clone()); + } + + // Pass through compatible fields + for &field in PASSTHROUGH_FIELDS { + if let Some(value) = obj.get(field) { + result.insert(field.to_string(), value.clone()); + } + } + + // Strip incompatible fields — they are simply not copied over. + // (No action needed since we build a new object.) + let _ = STRIP_FIELDS; // acknowledge the constant is used by design + + Ok(serde_json::to_vec(&Value::Object(result))?) +} + +/// Convert a chat/completions message to a Responses API input item. +/// - Adds `type: "message"` (required by Responses API) +/// - For user messages: converts content `type: "text"` → `type: "input_text"` +/// - For assistant messages: converts content `type: "text"` → `type: "output_text"` +/// - Converts content `type: "image_url"` → `type: "input_image"` +/// - String content is left as-is (the Responses API accepts string content directly). +fn convert_message_content(mut msg: Value) -> Value { + let role = msg.get("role").and_then(Value::as_str).unwrap_or(""); + let is_assistant = role == "assistant"; + + // Responses API requires "type": "message" on each input item + if let Some(obj) = msg.as_object_mut() { + if !obj.contains_key("type") { + obj.insert("type".to_string(), Value::String("message".to_string())); + } + } + + let Some(content) = msg.get_mut("content") else { + return msg; + }; + let Some(parts) = content.as_array_mut() else { + // String content — no conversion needed + return msg; + }; + for part in parts.iter_mut() { + let Some(obj) = part.as_object_mut() else { + continue; + }; + match obj.get("type").and_then(Value::as_str) { + Some("text") => { + let text_type = if is_assistant { "output_text" } else { "input_text" }; + obj.insert("type".to_string(), Value::String(text_type.to_string())); + } + Some("image_url") => { + obj.insert("type".to_string(), Value::String("input_image".to_string())); + } + _ => {} + } + } + msg +} + +/// Translate a `/v1/responses` response body to `/v1/chat/completions` format. +pub fn responses_to_chat_completion(body: &[u8]) -> Result, TranslateError> { + let resp: Value = serde_json::from_slice(body)?; + let obj = resp.as_object().ok_or(TranslateError::MissingField("root object"))?; + + let id = obj + .get("id") + .and_then(Value::as_str) + .unwrap_or("chatcmpl-translate"); + let model = obj + .get("model") + .and_then(Value::as_str) + .unwrap_or("unknown"); + + // Extract text content from output[].content[].text where type == "output_text" + let mut content_parts: Vec<&str> = Vec::new(); + if let Some(output) = obj.get("output").and_then(Value::as_array) { + for item in output { + if item.get("type").and_then(Value::as_str) == Some("message") { + if let Some(content) = item.get("content").and_then(Value::as_array) { + for part in content { + if part.get("type").and_then(Value::as_str) == Some("output_text") { + if let Some(text) = part.get("text").and_then(Value::as_str) { + content_parts.push(text); + } + } + } + } + } + } + } + let content = content_parts.join(""); + + // Map status → finish_reason + let finish_reason = match obj.get("status").and_then(Value::as_str) { + Some("completed") | None => "stop", + Some("incomplete") => "length", + Some("failed") => "stop", + Some(_) => "stop", + }; + + // Map usage + let usage = if let Some(u) = obj.get("usage") { + serde_json::json!({ + "prompt_tokens": u.get("input_tokens").and_then(Value::as_u64).unwrap_or(0), + "completion_tokens": u.get("output_tokens").and_then(Value::as_u64).unwrap_or(0), + "total_tokens": + u.get("input_tokens").and_then(Value::as_u64).unwrap_or(0) + + u.get("output_tokens").and_then(Value::as_u64).unwrap_or(0) + }) + } else { + serde_json::json!({ "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0 }) + }; + + let result = serde_json::json!({ + "id": id, + "object": "chat.completion", + "model": model, + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": content, + }, + "finish_reason": finish_reason, + }], + "usage": usage, + }); + + Ok(serde_json::to_vec(&result)?) +} + +/// Translate a single SSE line from Responses API format to chat.completion.chunk format. +/// +/// Returns `Some(line(s))` for events that map to chat completions output, +/// or `None` for events that should be suppressed. +/// +/// `response_id` and `model` are captured from early events and reused in later chunks. +pub fn translate_sse_line( + line: &str, + response_id: &mut Option, + model: &mut Option, +) -> Option { + // Pass through [DONE] + if line.starts_with("data: [DONE]") { + return Some(line.to_string()); + } + + // Only process data: lines with JSON + let json_str = line.strip_prefix("data: ")?; + + let event: Value = serde_json::from_str(json_str).ok()?; + let event_type = event.get("type").and_then(Value::as_str)?; + + match event_type { + "response.created" | "response.in_progress" => { + // Capture response ID and model from these early events + if let Some(resp) = event.get("response") { + if let Some(id) = resp.get("id").and_then(Value::as_str) { + *response_id = Some(id.to_string()); + } + if let Some(m) = resp.get("model").and_then(Value::as_str) { + *model = Some(m.to_string()); + } + } + None // suppress + } + + "response.output_text.delta" => { + let delta = event.get("delta").and_then(Value::as_str).unwrap_or(""); + let id = response_id.as_deref().unwrap_or("chatcmpl-translate"); + let m = model.as_deref().unwrap_or("unknown"); + let chunk = serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "model": m, + "choices": [{ + "index": 0, + "delta": { "content": delta }, + "finish_reason": null, + }] + }); + Some(format!("data: {}", serde_json::to_string(&chunk).unwrap_or_default())) + } + + "response.completed" => { + let id = response_id.as_deref().unwrap_or("chatcmpl-translate"); + let m = model.as_deref().unwrap_or("unknown"); + let final_chunk = serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "model": m, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": "stop", + }] + }); + Some(format!( + "data: {}\n\ndata: [DONE]", + serde_json::to_string(&final_chunk).unwrap_or_default() + )) + } + + "response.failed" => { + let id = response_id.as_deref().unwrap_or("chatcmpl-translate"); + let m = model.as_deref().unwrap_or("unknown"); + let final_chunk = serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "model": m, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": "stop", + }] + }); + Some(format!( + "data: {}\n\ndata: [DONE]", + serde_json::to_string(&final_chunk).unwrap_or_default() + )) + } + + "response.incomplete" => { + let id = response_id.as_deref().unwrap_or("chatcmpl-translate"); + let m = model.as_deref().unwrap_or("unknown"); + let final_chunk = serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "model": m, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": "length", + }] + }); + Some(format!( + "data: {}\n\ndata: [DONE]", + serde_json::to_string(&final_chunk).unwrap_or_default() + )) + } + + // Suppress all structural/metadata events + "response.output_text.done" + | "response.content_part.added" + | "response.content_part.done" + | "response.output_item.added" + | "response.output_item.done" => None, + + // Suppress any other unknown events + _ => None, + } +} + +/// A stream adapter that wraps an axum Body and translates Responses API SSE events +/// to chat.completion.chunk format. +pub struct TranslateStream { + inner: Pin> + Send>>, + buffer: BytesMut, + response_id: Option, + model: Option, + output_buffer: Vec, +} + +impl std::fmt::Debug for TranslateStream { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TranslateStream") + .field("buffer_len", &self.buffer.len()) + .field("response_id", &self.response_id) + .field("model", &self.model) + .finish() + } +} + +impl TranslateStream { + pub fn new(body: Body) -> Self { + use http_body_util::BodyStream; + use futures_util::StreamExt; + + let stream = BodyStream::new(body).filter_map(|result| async move { + match result { + Ok(frame) => frame.into_data().ok().map(Ok), + Err(e) => Some(Err(e)), + } + }); + + Self { + inner: Box::pin(stream), + buffer: BytesMut::new(), + response_id: None, + model: None, + output_buffer: Vec::new(), + } + } + + fn process_buffered_lines(&mut self) { + loop { + let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') else { + break; + }; + + let line_bytes = self.buffer.split_to(pos + 1); + let line = String::from_utf8_lossy(&line_bytes).trim().to_string(); + + if line.is_empty() { + self.output_buffer.extend_from_slice(b"\n"); + continue; + } + + let rid = &mut self.response_id; + let mdl = &mut self.model; + if let Some(translated) = translate_sse_line(&line, rid, mdl) { + self.output_buffer.extend_from_slice(translated.as_bytes()); + self.output_buffer.extend_from_slice(b"\n\n"); + } + } + } +} + +impl Stream for TranslateStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + loop { + // First, drain any pending output + if !this.output_buffer.is_empty() { + let data = std::mem::take(&mut this.output_buffer); + return Poll::Ready(Some(Ok(Bytes::from(data)))); + } + + // Poll the inner stream for more data + match this.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(chunk))) => { + this.buffer.extend_from_slice(&chunk); + this.process_buffered_lines(); + // Loop to check if we produced output + } + Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))), + Poll::Ready(None) => { + // Stream ended — process any remaining buffer + if !this.buffer.is_empty() { + let remaining = std::mem::take(&mut this.buffer); + let line = String::from_utf8_lossy(&remaining).trim().to_string(); + if !line.is_empty() { + if let Some(translated) = translate_sse_line( + &line, + &mut this.response_id, + &mut this.model, + ) { + return Poll::Ready(Some(Ok(Bytes::from( + format!("{translated}\n\n"), + )))); + } + } + } + return Poll::Ready(None); + } + Poll::Pending => return Poll::Pending, + } + } + } +} + +/// Wrap a Body in a TranslateStream and return a new Body. +pub fn wrap_body_with_translate_stream(body: Body) -> Body { + Body::from_stream(TranslateStream::new(body)) +} + +// --------------------------------------------------------------------------- +// Gemini SSE → OpenAI chat.completion.chunk translation +// --------------------------------------------------------------------------- + +/// Translate a single SSE line from Gemini streamGenerateContent format +/// to OpenAI chat.completion.chunk format. +/// +/// Gemini SSE events look like: +/// ```text +/// data: {"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"},...}],...} +/// ``` +/// +/// Returns `Some(line)` for data events, `None` for events to suppress. +pub fn translate_gemini_sse_line( + line: &str, + model: &mut Option, +) -> Option { + // Pass through [DONE] + if line.starts_with("data: [DONE]") { + return Some(line.to_string()); + } + + let json_str = line.strip_prefix("data: ")?; + let event: Value = serde_json::from_str(json_str).ok()?; + + // Cloudcode-pa wraps the Gemini payload in a "response" envelope + let inner = event.get("response").unwrap_or(&event); + + // Capture model from modelVersion if present + if let Some(m) = inner.get("modelVersion").and_then(Value::as_str) { + *model = Some(m.to_string()); + } + + let candidates = inner.get("candidates").and_then(Value::as_array)?; + let candidate = candidates.first()?; + + let finish_reason = candidate + .get("finishReason") + .and_then(Value::as_str); + + let parts = candidate + .get("content") + .and_then(|c| c.get("parts")) + .and_then(Value::as_array); + + let text = parts + .and_then(|p| p.first()) + .and_then(|p| p.get("text")) + .and_then(Value::as_str) + .unwrap_or(""); + + let m = model.as_deref().unwrap_or("unknown"); + let id = "chatcmpl-gemini"; + + // If there's a finish reason (STOP, MAX_TOKENS, etc.), emit final chunk + [DONE] + if let Some(reason) = finish_reason { + let mapped_reason = match reason { + "STOP" => "stop", + "MAX_TOKENS" => "length", + _ => "stop", + }; + + // Emit content delta if any, then finish chunk, then [DONE] + let mut result = String::new(); + if !text.is_empty() { + let content_chunk = serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "model": m, + "choices": [{ + "index": 0, + "delta": { "content": text }, + "finish_reason": null, + }] + }); + result.push_str(&format!("data: {}\n\n", serde_json::to_string(&content_chunk).unwrap_or_default())); + } + + let final_chunk = serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "model": m, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": mapped_reason, + }] + }); + result.push_str(&format!( + "data: {}\n\ndata: [DONE]", + serde_json::to_string(&final_chunk).unwrap_or_default() + )); + return Some(result); + } + + // Regular content delta + if text.is_empty() { + return None; + } + + let chunk = serde_json::json!({ + "id": id, + "object": "chat.completion.chunk", + "model": m, + "choices": [{ + "index": 0, + "delta": { "content": text }, + "finish_reason": null, + }] + }); + Some(format!("data: {}", serde_json::to_string(&chunk).unwrap_or_default())) +} + +/// Stream adapter for Gemini SSE → OpenAI chat.completion.chunk. +pub struct GeminiTranslateStream { + inner: Pin> + Send>>, + buffer: BytesMut, + model: Option, + output_buffer: Vec, +} + +impl GeminiTranslateStream { + pub fn new(body: Body) -> Self { + use http_body_util::BodyStream; + use futures_util::StreamExt; + + debug!("GeminiTranslateStream created — will translate Gemini SSE → OpenAI chat.completion.chunk"); + + let stream = BodyStream::new(body).filter_map(|result| async move { + match result { + Ok(frame) => frame.into_data().ok().map(Ok), + Err(e) => Some(Err(e)), + } + }); + + Self { + inner: Box::pin(stream), + buffer: BytesMut::new(), + model: None, + output_buffer: Vec::new(), + } + } + + fn process_buffered_lines(&mut self) { + loop { + let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') else { + break; + }; + + let line_bytes = self.buffer.split_to(pos + 1); + let line = String::from_utf8_lossy(&line_bytes).trim().to_string(); + + if line.is_empty() { + self.output_buffer.extend_from_slice(b"\n"); + continue; + } + + debug!(gemini_line = %line.chars().take(200).collect::(), "GeminiTranslateStream: incoming line"); + + if let Some(translated) = translate_gemini_sse_line(&line, &mut self.model) { + debug!(translated_preview = %translated.chars().take(200).collect::(), "GeminiTranslateStream: translated"); + self.output_buffer.extend_from_slice(translated.as_bytes()); + self.output_buffer.extend_from_slice(b"\n\n"); + } else { + debug!("GeminiTranslateStream: line produced no translation output"); + } + } + } +} + +impl Stream for GeminiTranslateStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + loop { + if !this.output_buffer.is_empty() { + let data = std::mem::take(&mut this.output_buffer); + return Poll::Ready(Some(Ok(Bytes::from(data)))); + } + + match this.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(chunk))) => { + this.buffer.extend_from_slice(&chunk); + this.process_buffered_lines(); + } + Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))), + Poll::Ready(None) => { + if !this.buffer.is_empty() { + let remaining = std::mem::take(&mut this.buffer); + let line = String::from_utf8_lossy(&remaining).trim().to_string(); + if !line.is_empty() { + if let Some(translated) = + translate_gemini_sse_line(&line, &mut this.model) + { + return Poll::Ready(Some(Ok(Bytes::from( + format!("{translated}\n\n"), + )))); + } + } + } + return Poll::Ready(None); + } + Poll::Pending => return Poll::Pending, + } + } + } +} + +/// Wrap a Body in a GeminiTranslateStream and return a new Body. +pub fn wrap_body_with_gemini_translate_stream(body: Body) -> Body { + Body::from_stream(GeminiTranslateStream::new(body)) +} + +// --------------------------------------------------------------------------- +// DLP scanning for SSE streams +// --------------------------------------------------------------------------- + +/// Apply DLP redaction to a single SSE `data:` line. +/// +/// Parses the JSON, extracts `choices[0].delta.content`, runs redaction on it, +/// and patches the JSON back if any PII was found. Returns the (possibly +/// modified) line. +/// +/// Lines that are not `data:` JSON or don't contain delta content are returned +/// unchanged. +pub fn redact_sse_data_line(line: &str, scanner: &DlpScanner) -> String { + // Only process data: lines with JSON + let Some(json_str) = line.strip_prefix("data: ") else { + return line.to_string(); + }; + + // Don't touch [DONE] + if json_str.starts_with("[DONE]") { + return line.to_string(); + } + + let Ok(mut event) = serde_json::from_str::(json_str) else { + return line.to_string(); + }; + + // Extract delta.content from choices[0] + let Some(content) = event + .get_mut("choices") + .and_then(Value::as_array_mut) + .and_then(|choices| choices.first_mut()) + .and_then(|choice| choice.get_mut("delta")) + .and_then(|delta| delta.get_mut("content")) + else { + return line.to_string(); + }; + + let Some(text) = content.as_str() else { + return line.to_string(); + }; + + let (redacted, redacted_names) = scanner.redact_all(text.as_bytes()); + if redacted_names.is_empty() { + return line.to_string(); + } + + warn!( + redacted_patterns = ?redacted_names, + "PII redacted from streaming SSE chunk" + ); + + let redacted_str = String::from_utf8_lossy(&redacted); + *content = Value::String(redacted_str.into_owned()); + format!("data: {}", serde_json::to_string(&event).unwrap_or_else(|_| json_str.to_string())) +} + +/// Stream adapter that applies DLP redaction to SSE data lines. +pub struct DlpSseStream { + inner: Pin> + Send>>, + buffer: BytesMut, + scanner: Arc, + output_buffer: Vec, +} + +impl DlpSseStream { + pub fn new(body: Body, scanner: Arc) -> Self { + use futures_util::StreamExt; + use http_body_util::BodyStream; + + let stream = BodyStream::new(body).filter_map(|result| async move { + match result { + Ok(frame) => frame.into_data().ok().map(Ok), + Err(e) => Some(Err(e)), + } + }); + + Self { + inner: Box::pin(stream), + buffer: BytesMut::new(), + scanner, + output_buffer: Vec::new(), + } + } + + fn process_buffered_lines(&mut self) { + loop { + let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') else { + break; + }; + + let line_bytes = self.buffer.split_to(pos + 1); + let line = String::from_utf8_lossy(&line_bytes).trim().to_string(); + + if line.is_empty() { + self.output_buffer.extend_from_slice(b"\n"); + continue; + } + + let redacted = redact_sse_data_line(&line, &self.scanner); + self.output_buffer.extend_from_slice(redacted.as_bytes()); + self.output_buffer.extend_from_slice(b"\n"); + } + } +} + +impl Stream for DlpSseStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + loop { + if !this.output_buffer.is_empty() { + let data = std::mem::take(&mut this.output_buffer); + return Poll::Ready(Some(Ok(Bytes::from(data)))); + } + + match this.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(chunk))) => { + this.buffer.extend_from_slice(&chunk); + this.process_buffered_lines(); + } + Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))), + Poll::Ready(None) => { + if !this.buffer.is_empty() { + let remaining = std::mem::take(&mut this.buffer); + let line = String::from_utf8_lossy(&remaining).trim().to_string(); + if !line.is_empty() { + let redacted = redact_sse_data_line(&line, &this.scanner); + return Poll::Ready(Some(Ok(Bytes::from( + format!("{redacted}\n"), + )))); + } + } + return Poll::Ready(None); + } + Poll::Pending => return Poll::Pending, + } + } + } +} + +/// Wrap a Body in a DlpSseStream for streaming DLP redaction. +pub fn wrap_body_with_dlp_sse_stream(body: Body, scanner: Arc) -> Body { + Body::from_stream(DlpSseStream::new(body, scanner)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_chat_to_responses_basic() { + let body = serde_json::json!({ + "model": "gpt-4o-mini", + "messages": [ + {"role": "user", "content": "say hi"} + ] + }); + let result = chat_completions_to_responses(body.to_string().as_bytes()).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["model"], "gpt-4o-mini"); + assert_eq!(parsed["instructions"], "", "instructions should be empty when no system messages"); + let input = parsed["input"].as_array().unwrap(); + assert_eq!(input.len(), 1); + assert_eq!(input[0]["type"], "message"); + assert_eq!(input[0]["role"], "user"); + assert_eq!(input[0]["content"], "say hi"); + assert!(parsed.get("messages").is_none()); + } + + #[test] + fn test_chat_to_responses_with_system() { + let body = serde_json::json!({ + "model": "gpt-4o", + "messages": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "hello"} + ] + }); + let result = chat_completions_to_responses(body.to_string().as_bytes()).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["instructions"], "You are helpful."); + let input = parsed["input"].as_array().unwrap(); + assert_eq!(input.len(), 1); + assert_eq!(input[0]["role"], "user"); + } + + #[test] + fn test_chat_to_responses_multiple_system() { + let body = serde_json::json!({ + "model": "gpt-4o", + "messages": [ + {"role": "system", "content": "Be concise."}, + {"role": "system", "content": "Use markdown."}, + {"role": "user", "content": "hello"} + ] + }); + let result = chat_completions_to_responses(body.to_string().as_bytes()).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["instructions"], "Be concise.\nUse markdown."); + } + + #[test] + fn test_chat_to_responses_max_tokens() { + let body = serde_json::json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 100 + }); + let result = chat_completions_to_responses(body.to_string().as_bytes()).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["max_output_tokens"], 100); + assert!(parsed.get("max_tokens").is_none()); + } + + #[test] + fn test_chat_to_responses_strips_unsupported() { + let body = serde_json::json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "frequency_penalty": 0.5, + "presence_penalty": 0.5, + "logprobs": true, + "top_logprobs": 5, + "logit_bias": {"123": 1}, + "n": 2, + "response_format": {"type": "json_object"}, + "seed": 42, + "service_tier": "default", + "user": "user-123" + }); + let result = chat_completions_to_responses(body.to_string().as_bytes()).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + for field in STRIP_FIELDS { + assert!(parsed.get(*field).is_none(), "field '{}' should be stripped", field); + } + } + + #[test] + fn test_chat_to_responses_passthrough() { + let body = serde_json::json!({ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "stream": true, + "temperature": 0.7, + "top_p": 0.9, + "stop": ["\n"] + }); + let result = chat_completions_to_responses(body.to_string().as_bytes()).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["model"], "gpt-4o-mini"); + assert_eq!(parsed["stream"], true); + assert_eq!(parsed["temperature"], 0.7); + assert_eq!(parsed["top_p"], 0.9); + assert_eq!(parsed["stop"], serde_json::json!(["\n"])); + } + + #[test] + fn test_responses_to_chat_completion_basic() { + let body = serde_json::json!({ + "id": "resp_abc123", + "model": "gpt-4o-mini", + "status": "completed", + "output": [{ + "type": "message", + "content": [{ + "type": "output_text", + "text": "Hello!" + }] + }], + "usage": { + "input_tokens": 10, + "output_tokens": 5 + } + }); + let result = responses_to_chat_completion(body.to_string().as_bytes()).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["id"], "resp_abc123"); + assert_eq!(parsed["object"], "chat.completion"); + assert_eq!(parsed["model"], "gpt-4o-mini"); + let choice = &parsed["choices"][0]; + assert_eq!(choice["message"]["role"], "assistant"); + assert_eq!(choice["message"]["content"], "Hello!"); + assert_eq!(choice["finish_reason"], "stop"); + } + + #[test] + fn test_responses_to_chat_completion_usage() { + let body = serde_json::json!({ + "id": "resp_abc", + "model": "gpt-4o", + "status": "completed", + "output": [{ + "type": "message", + "content": [{"type": "output_text", "text": "hi"}] + }], + "usage": { + "input_tokens": 50, + "output_tokens": 25 + } + }); + let result = responses_to_chat_completion(body.to_string().as_bytes()).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["usage"]["prompt_tokens"], 50); + assert_eq!(parsed["usage"]["completion_tokens"], 25); + assert_eq!(parsed["usage"]["total_tokens"], 75); + } + + #[test] + fn test_responses_to_chat_completion_incomplete() { + let body = serde_json::json!({ + "id": "resp_inc", + "model": "gpt-4o", + "status": "incomplete", + "output": [{ + "type": "message", + "content": [{"type": "output_text", "text": "partial"}] + }], + "usage": { "input_tokens": 10, "output_tokens": 5 } + }); + let result = responses_to_chat_completion(body.to_string().as_bytes()).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + + assert_eq!(parsed["choices"][0]["finish_reason"], "length"); + } + + #[test] + fn test_sse_delta() { + let event = serde_json::json!({ + "type": "response.output_text.delta", + "delta": "Hello" + }); + let line = format!("data: {}", event); + let mut response_id = Some("resp_123".to_string()); + let mut model = Some("gpt-4o-mini".to_string()); + let result = translate_sse_line(&line, &mut response_id, &mut model).unwrap(); + + assert!(result.starts_with("data: ")); + let json_str = result.strip_prefix("data: ").unwrap(); + let parsed: Value = serde_json::from_str(json_str).unwrap(); + + assert_eq!(parsed["object"], "chat.completion.chunk"); + assert_eq!(parsed["id"], "resp_123"); + assert_eq!(parsed["model"], "gpt-4o-mini"); + assert_eq!(parsed["choices"][0]["delta"]["content"], "Hello"); + assert!(parsed["choices"][0]["finish_reason"].is_null()); + } + + #[test] + fn test_sse_completed() { + let event = serde_json::json!({ + "type": "response.completed", + "response": {"id": "resp_456", "status": "completed"} + }); + let line = format!("data: {}", event); + let mut response_id = Some("resp_456".to_string()); + let mut model = Some("gpt-4o".to_string()); + let result = translate_sse_line(&line, &mut response_id, &mut model).unwrap(); + + // Should contain a final chunk with finish_reason: "stop" and then [DONE] + assert!(result.contains("\"finish_reason\":\"stop\"")); + assert!(result.contains("data: [DONE]")); + } + + #[test] + fn test_sse_meta_suppressed() { + let mut response_id = None; + let mut model = None; + + let created = serde_json::json!({ + "type": "response.created", + "response": {"id": "resp_789", "model": "gpt-4o"} + }); + let result = translate_sse_line( + &format!("data: {}", created), + &mut response_id, + &mut model, + ); + assert!(result.is_none()); + assert_eq!(response_id.as_deref(), Some("resp_789")); + assert_eq!(model.as_deref(), Some("gpt-4o")); + + let in_progress = serde_json::json!({ + "type": "response.in_progress", + "response": {"id": "resp_789"} + }); + let result = translate_sse_line( + &format!("data: {}", in_progress), + &mut response_id, + &mut model, + ); + assert!(result.is_none()); + + // Structural events should also be suppressed + let content_part = serde_json::json!({"type": "response.content_part.added"}); + let result = translate_sse_line( + &format!("data: {}", content_part), + &mut response_id, + &mut model, + ); + assert!(result.is_none()); + } + + #[test] + fn test_sse_done_passthrough() { + let mut response_id = None; + let mut model = None; + let result = translate_sse_line("data: [DONE]", &mut response_id, &mut model); + assert_eq!(result, Some("data: [DONE]".to_string())); + } + + #[test] + fn test_chat_to_responses_multipart_content_types() { + let body = serde_json::to_vec(&serde_json::json!({ + "model": "gpt-4o", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}} + ] + } + ] + })) + .unwrap(); + + let result = chat_completions_to_responses(&body).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + assert_eq!(parsed["input"][0]["type"], "message"); + let content = parsed["input"][0]["content"].as_array().unwrap(); + assert_eq!(content[0]["type"], "input_text"); + assert_eq!(content[0]["text"], "What is in this image?"); + assert_eq!(content[1]["type"], "input_image"); + } + + #[test] + fn test_chat_to_responses_string_content_unchanged() { + let body = serde_json::to_vec(&serde_json::json!({ + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "hello"} + ] + })) + .unwrap(); + + let result = chat_completions_to_responses(&body).unwrap(); + let parsed: Value = serde_json::from_slice(&result).unwrap(); + assert_eq!(parsed["input"][0]["type"], "message"); + assert_eq!(parsed["input"][0]["role"], "user"); + assert_eq!(parsed["input"][0]["content"], "hello"); + } + + // ----------------------------------------------------------------------- + // DLP SSE redaction tests + // ----------------------------------------------------------------------- + + fn test_dlp_scanner() -> DlpScanner { + use crate::config::{DlpAction, DlpPattern}; + DlpScanner::new( + &[ + DlpPattern { + name: "email".to_string(), + regex: r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b".to_string(), + action: DlpAction::Redact, + }, + DlpPattern { + name: "ssn".to_string(), + regex: r"\b\d{3}-\d{2}-\d{4}\b".to_string(), + action: DlpAction::Block, + }, + ], + true, + ) + .unwrap() + } + + #[test] + fn test_redact_sse_data_line_with_pii() { + let scanner = test_dlp_scanner(); + let line = r#"data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Contact user@example.com for info"},"finish_reason":null}]}"#; + let result = redact_sse_data_line(line, &scanner); + assert!(result.starts_with("data: "), "Should still be an SSE data line"); + assert!(result.contains("[REDACTED:email]"), "Email should be redacted"); + assert!(!result.contains("user@example.com"), "Original email should be gone"); + } + + #[test] + fn test_redact_sse_data_line_clean() { + let scanner = test_dlp_scanner(); + let line = r#"data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello world"},"finish_reason":null}]}"#; + let result = redact_sse_data_line(line, &scanner); + assert_eq!(result, line, "Clean content should pass through unchanged"); + } + + #[test] + fn test_redact_sse_data_line_done() { + let scanner = test_dlp_scanner(); + let result = redact_sse_data_line("data: [DONE]", &scanner); + assert_eq!(result, "data: [DONE]"); + } + + #[test] + fn test_redact_sse_data_line_no_delta_content() { + let scanner = test_dlp_scanner(); + let line = r#"data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}"#; + let result = redact_sse_data_line(line, &scanner); + assert_eq!(result, line, "Lines without delta.content pass through unchanged"); + } + + #[test] + fn test_redact_sse_data_line_non_data_line() { + let scanner = test_dlp_scanner(); + let result = redact_sse_data_line("event: message", &scanner); + assert_eq!(result, "event: message", "Non-data lines pass through unchanged"); + } +} diff --git a/tests/snapshots/config_fixtures__all_fields.snap b/tests/snapshots/config_fixtures__all_fields.snap index ea04598..23c1fca 100644 --- a/tests/snapshots/config_fixtures__all_fields.snap +++ b/tests/snapshots/config_fixtures__all_fields.snap @@ -13,9 +13,11 @@ keys: - virtual_key: vk-1 real_key: sk-real-1 provider: openai + auth: static - virtual_key: vk-2 real_key: sk-real-2 provider: anthropic + auth: static dlp: patterns: - name: ssn diff --git a/tests/snapshots/config_fixtures__empty_keys.snap b/tests/snapshots/config_fixtures__empty_keys.snap index df56020..2a114e0 100644 --- a/tests/snapshots/config_fixtures__empty_keys.snap +++ b/tests/snapshots/config_fixtures__empty_keys.snap @@ -13,6 +13,7 @@ keys: - virtual_key: "" real_key: "" provider: openai + auth: static dlp: patterns: [] scan_responses: true diff --git a/tests/snapshots/config_fixtures__key_missing_real_key.snap b/tests/snapshots/config_fixtures__key_missing_real_key.snap index 29f035d..5aaa8e8 100644 --- a/tests/snapshots/config_fixtures__key_missing_real_key.snap +++ b/tests/snapshots/config_fixtures__key_missing_real_key.snap @@ -2,8 +2,4 @@ source: tests/config_fixtures.rs expression: err.to_string() --- -TOML parse error at line 4, column 1 - | -4 | [[keys]] - | ^^^^^^^^ -missing field `real_key` +key 'vk-1': real_key is required when auth = "static"