From 90321cf61937d943501d219f797cd20a7b0672a3 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Sat, 21 Feb 2026 04:24:52 -0500 Subject: [PATCH 01/16] add the plan to support codex oauth --- docs/architecture-comparison.md | 627 +++++++++++++++++++++++++++ docs/codex-oauth-plan.md | 729 ++++++++++++++++++++++++++++++++ 2 files changed, 1356 insertions(+) create mode 100644 docs/architecture-comparison.md create mode 100644 docs/codex-oauth-plan.md 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? From af095c04770f44975d201bd3b654d526dcdf5710 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Sat, 21 Feb 2026 23:08:10 -0500 Subject: [PATCH 02/16] support codex and antigravity oauth --- .gitignore | 2 + Cargo.lock | 278 +++++++- Cargo.toml | 11 +- Dockerfile | 8 + docs/docker.md | 160 +++++ src/app.rs | 228 +++++-- src/app/tests.rs | 46 +- src/config.rs | 88 ++- src/keys.rs | 61 +- src/main.rs | 79 ++- src/oauth/antigravity.rs | 606 +++++++++++++++++ src/oauth/codex.rs | 550 ++++++++++++++++ src/oauth/mod.rs | 610 ++++++++++++++++++ src/oauth/storage.rs | 190 ++++++ src/onboard/config_render.rs | 56 +- src/onboard/interactive.rs | 193 ++++-- src/onboard/mod.rs | 5 +- src/onboard/test_support.rs | 1 + src/onboard/types.rs | 15 + src/openclaw_cli.rs | 1 + src/proxy.rs | 133 +++- .../config_fixtures__all_fields.snap | 2 + .../config_fixtures__empty_keys.snap | 1 + ...config_fixtures__key_missing_real_key.snap | 6 +- 24 files changed, 3170 insertions(+), 160 deletions(-) create mode 100644 Dockerfile create mode 100644 docs/docker.md create mode 100644 src/oauth/antigravity.rs create mode 100644 src/oauth/codex.rs create mode 100644 src/oauth/mod.rs create mode 100644 src/oauth/storage.rs 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..b1b62ac --- /dev/null +++ b/Dockerfile @@ -0,0 +1,8 @@ +FROM debian:sid-slim +RUN apt-get update && apt-get install -y ca-certificates && rm -rf /var/lib/apt/lists/* +ENV TERM=xterm-256color +EXPOSE 18790 +WORKDIR /etc/clawshell +COPY target/release/clawshell /usr/local/bin/clawshell +COPY .env .env +ENTRYPOINT ["clawshell"] diff --git a/docs/docker.md b/docs/docker.md new file mode 100644 index 0000000..5de6cd0 --- /dev/null +++ b/docs/docker.md @@ -0,0 +1,160 @@ +# 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 +``` + +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. + +## 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. + +### 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 copy/paste flow. The wizard prints a Google authorization URL. Open it +in your browser, authorize, then copy the authorization code from the result +page and paste it back into the terminal. + +## 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) +docker run --rm -it -v clawshell-config:/etc/clawshell 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 fe4e44c..16a1ce1 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::response::{IntoResponse, Response}; use axum::routing::{any, get}; use bytes::Bytes; @@ -27,6 +28,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>, @@ -34,7 +36,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( @@ -46,19 +56,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| { @@ -111,6 +131,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), @@ -459,7 +480,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!( @@ -529,27 +550,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() { @@ -620,6 +668,110 @@ 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(), + }; + + // 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(), + 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, + "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 Ok(response); + } + + // 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, + uri, + headers, + retry_body, + provider, + retry_auth_headers, + upstream_url.as_deref(), + ) + .await + .map_err(|e| format!("OAuth retry forward failed: {e}"))?; + + return Ok(retry_response); + } + + Ok(response) +} + 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..381c1be 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()), @@ -1500,6 +1517,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..25c9074 100644 --- a/src/config.rs +++ b/src/config.rs @@ -37,6 +37,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 { @@ -81,13 +83,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", "antigravity"). + /// 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 +242,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 +377,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() } 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..7dba9a2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,6 +9,8 @@ mod dlp; mod email; mod keys; mod migration; +#[allow(dead_code)] +mod oauth; mod onboard; mod openclaw_cli; mod platform; @@ -555,7 +557,10 @@ async fn cmd_start_inner(config_path: &str) -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result> { + use crate::oauth::{ + OAuthRegistry, TokenStorage, + antigravity::AntigravityProvider, + 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)); + } + "antigravity" => { + let provider = AntigravityProvider::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/antigravity.rs b/src/oauth/antigravity.rs new file mode 100644 index 0000000..ae500d6 --- /dev/null +++ b/src/oauth/antigravity.rs @@ -0,0 +1,606 @@ +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, warn}; + +const DEFAULT_AUTH_URL: &str = "https://accounts.google.com/o/oauth2/auth"; +const DEFAULT_TOKEN_URL: &str = "https://oauth2.googleapis.com/token"; +const DEFAULT_SCOPES: &[&str] = &[ + "openid", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", + "https://www.googleapis.com/auth/cloud-platform", +]; + +const ENDPOINT_PRODUCTION: &str = "https://cloudcode-pa.googleapis.com"; +const ENDPOINT_DAILY: &str = "https://daily-cloudcode-pa.sandbox.googleapis.com"; +const ENDPOINT_ALT: &str = "https://codeassist.googleapis.com/v1"; + +#[derive(Debug)] +pub struct AntigravityProvider { + client_id: String, + client_secret: String, + auth_url: String, + token_url: String, + scopes: Vec, + http_client: reqwest::Client, + endpoints: Vec, +} + +impl AntigravityProvider { + pub fn new( + client_id: Option<&str>, + auth_url: Option<&str>, + token_url: Option<&str>, + scopes: Option<&[String]>, + ) -> Self { + Self::new_with_secret(client_id, None, auth_url, token_url, scopes) + } + + pub fn new_with_secret( + client_id: Option<&str>, + client_secret: Option<&str>, + auth_url: Option<&str>, + token_url: Option<&str>, + scopes: Option<&[String]>, + ) -> Self { + // Load .env file if present (ignored if missing) + let _ = dotenvy::dotenv(); + + // Env vars take priority, then explicit constructor arguments. + // No hardcoded defaults — credentials must come from env, .env file, or config. + let resolved_client_id = std::env::var("GOOGLE_OAUTH_CLIENT_ID") + .ok() + .or_else(|| client_id.map(String::from)) + .expect("GOOGLE_OAUTH_CLIENT_ID env var or client_id argument is required"); + let resolved_client_secret = std::env::var("GOOGLE_OAUTH_CLIENT_SECRET") + .ok() + .or_else(|| client_secret.map(String::from)) + .expect("GOOGLE_OAUTH_CLIENT_SECRET env var or client_secret argument is required"); + + Self { + client_id: resolved_client_id, + client_secret: resolved_client_secret, + 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"), + endpoints: vec![ + ENDPOINT_PRODUCTION.to_string(), + ENDPOINT_DAILY.to_string(), + ENDPOINT_ALT.to_string(), + ], + } + } + + 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.as_str()), + ("client_secret", self.client_secret.as_str()), + ("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?; + let mut tokens = parse_google_token_response(&json)?; + + // Discover project ID after successful login + if let Err(e) = self.discover_project_id(&mut tokens).await { + warn!(error = %e, "Failed to discover Antigravity project ID"); + } + + Ok(tokens) + } + + async fn exchange_refresh_token( + &self, + refresh_token: &str, + ) -> Result { + let params = [ + ("grant_type", "refresh_token"), + ("client_id", self.client_id.as_str()), + ("client_secret", self.client_secret.as_str()), + ("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?; + // Google refresh responses may not include a new refresh token; + // the caller should preserve the original refresh token. + let mut tokens = parse_google_token_response(&json)?; + if tokens.refresh_token.is_none() { + tokens.refresh_token = Some(refresh_token.to_string()); + } + Ok(tokens) + } + + async fn discover_project_id(&self, tokens: &mut OAuthTokens) -> Result<(), OAuthError> { + let url = format!("{}/v1internal:loadCodeAssist", self.endpoints[0]); + + let resp: reqwest::Response = self + .http_client + .post(&url) + .header("Authorization", format!("Bearer {}", tokens.access_token)) + .header( + "x-goog-api-client", + "google-cloud-sdk vscode_cloudshelleditor/0.1", + ) + .json(&serde_json::json!({})) + .send() + .await?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(OAuthError::LoginFailed(format!( + "loadCodeAssist failed ({status}): {body}" + ))); + } + + let json: serde_json::Value = resp.json().await?; + if let Some(project_id) = json.get("projectId").and_then(|v| v.as_str()) { + tokens + .extra + .insert("project_id".to_string(), serde_json::json!(project_id)); + debug!(project_id, "Discovered Antigravity project ID"); + } + if let Some(tier) = json.get("tier").and_then(|v| v.as_str()) { + tokens + .extra + .insert("tier".to_string(), serde_json::json!(tier)); + } + + Ok(()) + } +} + +fn parse_google_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) +} + +/// Wrap an OpenAI-format request body into Antigravity/Gemini-style format. +pub fn wrap_antigravity_request( + body: &[u8], + project_id: &str, +) -> Result, OAuthError> { + let original: serde_json::Value = serde_json::from_slice(body).map_err(|e| { + OAuthError::LoginFailed(format!("failed to parse request body as JSON: {e}")) + })?; + + let model = original + .get("model") + .and_then(|v| v.as_str()) + .unwrap_or("gemini-2.0-flash"); + + let wrapped = serde_json::json!({ + "project": project_id, + "model": model, + "request": original, + }); + + serde_json::to_vec(&wrapped) + .map_err(|e| OAuthError::LoginFailed(format!("failed to serialize wrapped body: {e}"))) +} + +#[async_trait] +impl OAuthProvider for AntigravityProvider { + fn id(&self) -> &str { + "antigravity" + } + + fn display_name(&self) -> &str { + "Antigravity / Google (OAuth)" + } + + fn supports_headless_url(&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}/oauth-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={}&access_type=offline&prompt=consent", + 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 Antigravity/Google OAuth login"); + if let Err(e) = open::that(&auth_url) { + return Err(OAuthError::LoginFailed(format!( + "failed to open browser: {e}. Visit this URL manually:\n{auth_url}" + ))); + } + + 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 { + let (verifier, challenge) = generate_pkce(); + let redirect_uri = "https://codeassist.google.com/authcode".to_string(); + + let auth_url = format!( + "{}?response_type=code&client_id={}&redirect_uri={}&scope={}&code_challenge={}&code_challenge_method=S256&access_type=offline&prompt=consent", + self.auth_url, + urlencoding::encode(&self.client_id), + urlencoding::encode(&redirect_uri), + urlencoding::encode(&self.scopes.join(" ")), + urlencoding::encode(&challenge), + ); + + println!(); + println!(" Visit this URL to authenticate:"); + println!(" {auth_url}"); + println!(); + println!(" After authorizing, Google will show an authorization code."); + println!(" Copy and paste it below:"); + + let code = crate::tui::prompt_text("Authorization code", None) + .map_err(|e| OAuthError::LoginFailed(format!("failed to read code: {e}")))?; + + self.exchange_code(code.trim(), &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()?, + ); + 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 for Antigravity provider".to_string(), + ))?; + 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[0] + )) + } +} + +/// Wait for an OAuth callback on a local HTTP server (same as codex). +async fn wait_for_oauth_callback( + port: u16, +) -> Result<(String, String), Box> { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{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]); + + 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)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_provider() -> AntigravityProvider { + AntigravityProvider::new_with_secret( + Some("test-client-id"), + Some("test-client-secret"), + None, + None, + None, + ) + } + + #[test] + fn test_parse_google_token_response() { + let json = serde_json::json!({ + "access_token": "ya29.test", + "refresh_token": "1//test", + "expires_in": 3600, + "scope": "openid email profile", + "token_type": "Bearer" + }); + + let tokens = parse_google_token_response(&json).unwrap(); + assert_eq!(tokens.access_token, "ya29.test"); + assert_eq!(tokens.refresh_token.as_deref(), Some("1//test")); + assert!(tokens.expires_at.is_some()); + } + + #[test] + fn test_parse_google_token_response_no_refresh() { + let json = serde_json::json!({ + "access_token": "ya29.refreshed", + "expires_in": 3600 + }); + + let tokens = parse_google_token_response(&json).unwrap(); + assert_eq!(tokens.access_token, "ya29.refreshed"); + assert!(tokens.refresh_token.is_none()); + } + + #[test] + fn test_wrap_antigravity_request() { + let body = serde_json::json!({ + "model": "gemini-3-pro", + "messages": [{"role": "user", "content": "hello"}] + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + + let wrapped = wrap_antigravity_request(&body_bytes, "proj-abc-123").unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); + + assert_eq!(parsed["project"], "proj-abc-123"); + assert_eq!(parsed["model"], "gemini-3-pro"); + assert_eq!(parsed["request"]["messages"][0]["content"], "hello"); + } + + #[test] + fn test_wrap_antigravity_request_default_model() { + let body = serde_json::json!({ + "messages": [{"role": "user", "content": "test"}] + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + + let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); + + assert_eq!(parsed["model"], "gemini-2.0-flash"); + } + + #[test] + fn test_antigravity_provider_defaults() { + let provider = test_provider(); + assert_eq!(provider.id(), "antigravity"); + assert_eq!(provider.display_name(), "Antigravity / Google (OAuth)"); + assert!(provider.supports_headless_url()); + assert!(!provider.supports_device_code()); + } + + #[test] + fn test_inject_auth_headers() { + let provider = test_provider(); + let mut headers = HeaderMap::new(); + provider.inject_auth(&mut headers, "ya29.test").unwrap(); + + assert_eq!( + headers.get("authorization").unwrap().to_str().unwrap(), + "Bearer ya29.test" + ); + assert!(headers.get("x-goog-api-client").is_some()); + assert!(headers.get("client-metadata").is_some()); + } + + #[test] + fn test_prepare_request_body_with_project_id() { + let provider = test_provider(); + let mut tokens = OAuthTokens { + access_token: "t".to_string(), + refresh_token: None, + id_token: None, + expires_at: None, + account_id: None, + extra: BTreeMap::new(), + }; + tokens.extra.insert( + "project_id".to_string(), + serde_json::json!("proj-test"), + ); + + let body = serde_json::json!({"model": "gemini-3-pro", "messages": []}); + let result = provider + .prepare_request_body(&serde_json::to_vec(&body).unwrap(), &tokens) + .unwrap(); + assert!(result.is_some()); + + let parsed: serde_json::Value = serde_json::from_slice(&result.unwrap()).unwrap(); + assert_eq!(parsed["project"], "proj-test"); + } + + #[test] + fn test_prepare_request_body_missing_project_id() { + let provider = test_provider(); + 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": "gemini-3-pro"}); + let result = provider.prepare_request_body(&serde_json::to_vec(&body).unwrap(), &tokens); + assert!(result.is_err()); + } + + #[test] + fn test_upstream_url() { + let provider = test_provider(); + let tokens = OAuthTokens { + access_token: "t".to_string(), + refresh_token: None, + id_token: None, + expires_at: None, + account_id: None, + extra: BTreeMap::new(), + }; + + let url = provider.upstream_url(&tokens).unwrap(); + assert!(url.contains("cloudcode-pa.googleapis.com")); + assert!(url.contains("streamGenerateContent")); + } +} diff --git a/src/oauth/codex.rs b/src/oauth/codex.rs new file mode 100644 index 0000000..f521bd2 --- /dev/null +++ b/src/oauth/codex.rs @@ -0,0 +1,550 @@ +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()?, + ); + Ok(()) + } +} + +/// 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(format!("127.0.0.1:{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)) +} + +#[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_passthrough() { + 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 result = provider.prepare_request_body(b"test body", &tokens).unwrap(); + assert!(result.is_none()); // pass-through + } + + #[test] + fn test_upstream_url_none() { + 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!(provider.upstream_url(&tokens).is_none()); + } +} diff --git a/src/oauth/mod.rs b/src/oauth/mod.rs new file mode 100644 index 0000000..d77e493 --- /dev/null +++ b/src/oauth/mod.rs @@ -0,0 +1,610 @@ +mod storage; + +pub mod codex; +pub mod antigravity; + +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", "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> { + 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 + } +} + +/// 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 +} + +/// 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. + 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()))?; + 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 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..0e20d26 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,58 @@ 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::antigravity::AntigravityProvider; + 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)), + "antigravity" => Box::new(AntigravityProvider::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 +383,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 — 5 top-level menu items per plan + const MENU_OPENAI: &str = "OpenAI"; + const MENU_OPENROUTER: &str = "OpenRouter"; + const MENU_ANTHROPIC: &str = "Anthropic"; + const MENU_CODEX: &str = "Codex / ChatGPT (OAuth)"; + const MENU_ANTIGRAVITY: &str = "Antigravity / Google (OAuth)"; + + let all_options = [ + MENU_OPENAI, + MENU_OPENROUTER, + MENU_ANTHROPIC, + MENU_CODEX, + MENU_ANTIGRAVITY, + ]; + + // 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("antigravity"), _) => Some(MENU_ANTIGRAVITY), + (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(), + }, + ), + MENU_ANTIGRAVITY => ( + "openai".to_string(), + OnboardAuthMethod::OAuth { + provider_id: "antigravity".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", + MENU_ANTIGRAVITY => "gemini-2.0-flash", + _ => "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 +821,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..ed54460 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", "antigravity". + 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/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" From 47bfe320ab7147aad288b230a2039be2d22eafaf Mon Sep 17 00:00:00 2001 From: u20024804 Date: Sun, 22 Feb 2026 09:55:26 -0500 Subject: [PATCH 03/16] support codex and antigravity oauth --- src/app.rs | 132 ++++++- src/main.rs | 1 + src/oauth/antigravity.rs | 17 + src/oauth/codex.rs | 160 ++++++++- src/oauth/mod.rs | 64 ++++ src/translate.rs | 733 +++++++++++++++++++++++++++++++++++++++ 6 files changed, 1097 insertions(+), 10 deletions(-) create mode 100644 src/translate.rs diff --git a/src/app.rs b/src/app.rs index 16a1ce1..0ebadea 100644 --- a/src/app.rs +++ b/src/app.rs @@ -696,6 +696,36 @@ async fn forward_oauth_request( 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 stream_requested = serde_json::from_slice::(&body_bytes) + .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 @@ -708,7 +738,7 @@ async fn forward_oauth_request( .proxy_client .forward_oauth( method.clone(), - uri, + &effective_uri, headers.clone(), body.clone(), provider, @@ -722,6 +752,7 @@ async fn forward_oauth_request( 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 { @@ -730,7 +761,7 @@ async fn forward_oauth_request( error = %e, "Token refresh failed after 401" ); - return Ok(response); + return maybe_translate_response(response, needs_translation, stream_requested).await; } // Re-inject auth with refreshed token @@ -756,7 +787,7 @@ async fn forward_oauth_request( .proxy_client .forward_oauth( method, - uri, + &effective_uri, headers, retry_body, provider, @@ -766,10 +797,101 @@ async fn forward_oauth_request( .await .map_err(|e| format!("OAuth retry forward failed: {e}"))?; - return Ok(retry_response); + 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).await; } - Ok(response) + // 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).await; + } + + maybe_translate_response(response, needs_translation, stream_requested).await +} + +/// Optionally translate a Responses API response back to chat/completions format. +async fn maybe_translate_response( + response: Response, + needs_translation: bool, + stream_requested: bool, +) -> Result { + if !needs_translation { + 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")); + + if is_streaming { + let (parts, body) = response.into_parts(); + let translated_body = crate::translate::wrap_body_with_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 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))) + } + } +} + +/// 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 { diff --git a/src/main.rs b/src/main.rs index 7dba9a2..4c1670c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,7 @@ mod openclaw_cli; mod platform; mod process; mod proxy; +mod translate; mod tui; use clap::Parser; diff --git a/src/oauth/antigravity.rs b/src/oauth/antigravity.rs index ae500d6..282c2d9 100644 --- a/src/oauth/antigravity.rs +++ b/src/oauth/antigravity.rs @@ -167,6 +167,13 @@ impl AntigravityProvider { if tokens.refresh_token.is_none() { tokens.refresh_token = Some(refresh_token.to_string()); } + + // Re-discover project ID with the fresh access token. + // This also recovers from initial login discovery failures. + if let Err(e) = self.discover_project_id(&mut tokens).await { + warn!(error = %e, "Failed to discover Antigravity project ID during refresh"); + } + Ok(tokens) } @@ -381,6 +388,16 @@ impl OAuthProvider for AntigravityProvider { Ok(()) } + async fn enrich_tokens(&self, tokens: &OAuthTokens) -> Result, OAuthError> { + if tokens.extra.contains_key("project_id") { + return Ok(None); + } + info!("Antigravity tokens missing project_id, discovering on-demand"); + let mut enriched = tokens.clone(); + self.discover_project_id(&mut enriched).await?; + Ok(Some(enriched)) + } + fn prepare_request_body( &self, body: &[u8], diff --git a/src/oauth/codex.rs b/src/oauth/codex.rs index f521bd2..32443ef 100644 --- a/src/oauth/codex.rs +++ b/src/oauth/codex.rs @@ -373,8 +373,49 @@ impl OAuthProvider for CodexProvider { 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" + } } /// Wait for an OAuth callback on a local HTTP server. @@ -424,6 +465,22 @@ async fn wait_for_oauth_callback( 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) +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); + serde_json::to_vec(&parsed).unwrap_or_else(|_| body.to_vec()) +} + #[cfg(test)] mod tests { use super::*; @@ -520,7 +577,31 @@ mod tests { } #[test] - fn test_prepare_request_body_passthrough() { + 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(), @@ -530,12 +611,52 @@ mod tests { account_id: None, extra: BTreeMap::new(), }; - let result = provider.prepare_request_body(b"test body", &tokens).unwrap(); - assert!(result.is_none()); // pass-through + // 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_upstream_url_none() { + 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(), @@ -545,6 +666,35 @@ mod tests { account_id: None, extra: BTreeMap::new(), }; - assert!(provider.upstream_url(&tokens).is_none()); + 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 index d77e493..731c32b 100644 --- a/src/oauth/mod.rs +++ b/src/oauth/mod.rs @@ -129,6 +129,25 @@ pub trait OAuthProvider: Send + Sync + fmt::Debug { 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 + } } /// Configuration for an OAuth provider from TOML. @@ -230,6 +249,7 @@ impl OAuthRegistry { } /// 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, @@ -239,6 +259,26 @@ impl OAuthRegistry { .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) @@ -373,6 +413,30 @@ impl OAuthRegistry { } } + 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 has_provider(&self, id: &str) -> bool { self.providers.contains_key(id) } diff --git a/src/translate.rs b/src/translate.rs new file mode 100644 index 0000000..ee0e3bb --- /dev/null +++ b/src/translate.rs @@ -0,0 +1,733 @@ +use axum::body::Body; +use bytes::{Bytes, BytesMut}; +use futures_util::Stream; +use serde_json::Value; +use std::pin::Pin; +use std::task::{Context, Poll}; + +#[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())); + } + } + + if !system_parts.is_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 chat/completions content types to responses API content types. +/// - `type: "text"` → `type: "input_text"` +/// - `type: "image_url"` → `type: "input_image"` with `image_url` → `image_url` +/// String content is left as-is (the Responses API accepts string content directly). +fn convert_message_content(mut msg: Value) -> Value { + 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") => { + obj.insert("type".to_string(), Value::String("input_text".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)) +} + +#[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!(parsed.get("instructions").is_none()); + let input = parsed["input"].as_array().unwrap(); + assert_eq!(input.len(), 1); + 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(); + 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]["content"], "hello"); + } +} From 3631ae079aa3cf96455e1000a50b85a43bdd37d2 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Sun, 22 Feb 2026 12:20:49 -0500 Subject: [PATCH 04/16] support codex and antigravity oauth --- .env | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 .env diff --git a/.env b/.env new file mode 100644 index 0000000..201f16d --- /dev/null +++ b/.env @@ -0,0 +1,3 @@ +GOOGLE_OAUTH_CLIENT_ID= +GOOGLE_OAUTH_CLIENT_SECRET= +ANTIGRAVITY_DEFAULT_PROJECT_ID= From 9937d7108ef42d22a11332e4c893b978b8357793 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Mon, 23 Feb 2026 03:02:31 -0500 Subject: [PATCH 05/16] support codex and antigravity oauth --- src/app.rs | 57 ++-- src/oauth/antigravity.rs | 557 +++++++++++++++++++++++++++++++++++++-- src/oauth/codex.rs | 8 + src/oauth/mod.rs | 27 ++ src/translate.rs | 235 ++++++++++++++++- 5 files changed, 838 insertions(+), 46 deletions(-) diff --git a/src/app.rs b/src/app.rs index 0ebadea..89e9647 100644 --- a/src/app.rs +++ b/src/app.rs @@ -706,6 +706,10 @@ async fn forward_oauth_request( .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}"))?; let stream_requested = serde_json::from_slice::(&body_bytes) .ok() .and_then(|v| v.get("stream")?.as_bool()) @@ -761,7 +765,7 @@ async fn forward_oauth_request( error = %e, "Token refresh failed after 401" ); - return maybe_translate_response(response, needs_translation, stream_requested).await; + return maybe_translate_response(response, needs_translation, stream_requested, response_format).await; } // Re-inject auth with refreshed token @@ -805,7 +809,7 @@ async fn forward_oauth_request( ); } - return maybe_translate_response(retry_response, needs_translation, stream_requested).await; + return maybe_translate_response(retry_response, needs_translation, stream_requested, response_format).await; } // Log error response bodies for debugging upstream issues @@ -827,21 +831,25 @@ async fn forward_oauth_request( ); } let response = Response::from_parts(parts, Body::from(body_bytes_resp)); - return maybe_translate_response(response, needs_translation, stream_requested).await; + return maybe_translate_response(response, needs_translation, stream_requested, response_format).await; } - maybe_translate_response(response, needs_translation, stream_requested).await + maybe_translate_response(response, needs_translation, stream_requested, response_format).await } -/// Optionally translate a Responses API response back to chat/completions format. +/// 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 { - if !needs_translation { - return Ok(response); - } + // 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 @@ -850,9 +858,20 @@ async fn maybe_translate_response( .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 = crate::translate::wrap_body_with_translate_stream(body); + 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)); } @@ -869,13 +888,21 @@ async fn maybe_translate_response( .map_err(|e| format!("failed to read response body for translation: {e}"))? .to_bytes(); - match crate::translate::responses_to_chat_completion(&body_bytes) { - Ok(translated) => { - parts.headers.remove("content-length"); - Ok(Response::from_parts(parts, Body::from(translated))) + 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))) + } + } } - Err(e) => { - warn!(error = %e, "Response translation failed, returning original"); + crate::oauth::ResponseFormat::GeminiSse => { + // Non-streaming Gemini responses are not expected; pass through Ok(Response::from_parts(parts, Body::from(body_bytes))) } } diff --git a/src/oauth/antigravity.rs b/src/oauth/antigravity.rs index 282c2d9..4cc7c0c 100644 --- a/src/oauth/antigravity.rs +++ b/src/oauth/antigravity.rs @@ -6,13 +6,14 @@ use chrono::Utc; use std::collections::BTreeMap; use tracing::{debug, info, warn}; -const DEFAULT_AUTH_URL: &str = "https://accounts.google.com/o/oauth2/auth"; +const DEFAULT_AUTH_URL: &str = "https://accounts.google.com/o/oauth2/v2/auth"; const DEFAULT_TOKEN_URL: &str = "https://oauth2.googleapis.com/token"; const DEFAULT_SCOPES: &[&str] = &[ - "openid", + "https://www.googleapis.com/auth/cloud-platform", "https://www.googleapis.com/auth/userinfo.email", "https://www.googleapis.com/auth/userinfo.profile", - "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/cclog", + "https://www.googleapis.com/auth/experimentsandconfigs", ]; const ENDPOINT_PRODUCTION: &str = "https://cloudcode-pa.googleapis.com"; @@ -28,6 +29,7 @@ pub struct AntigravityProvider { scopes: Vec, http_client: reqwest::Client, endpoints: Vec, + default_project_id: Option, } impl AntigravityProvider { @@ -60,6 +62,7 @@ impl AntigravityProvider { .ok() .or_else(|| client_secret.map(String::from)) .expect("GOOGLE_OAUTH_CLIENT_SECRET env var or client_secret argument is required"); + let default_project_id = std::env::var("ANTIGRAVITY_DEFAULT_PROJECT_ID").ok(); Self { client_id: resolved_client_id, @@ -81,6 +84,7 @@ impl AntigravityProvider { ENDPOINT_DAILY.to_string(), ENDPOINT_ALT.to_string(), ], + default_project_id, } } @@ -188,24 +192,68 @@ impl AntigravityProvider { "x-goog-api-client", "google-cloud-sdk vscode_cloudshelleditor/0.1", ) - .json(&serde_json::json!({})) + .header( + "Client-Metadata", + r#"{"ideType":"IDE_UNSPECIFIED","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}"#, + ) + .json(&serde_json::json!({ + "metadata": { + "ideType": "IDE_UNSPECIFIED", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI" + } + })) .send() .await?; if !resp.status().is_success() { let status = resp.status(); let body = resp.text().await.unwrap_or_default(); + warn!( + %status, + "loadCodeAssist failed, attempting fallback to default project ID" + ); + debug!(response_body = %body, "loadCodeAssist error response"); + if let Some(ref default_id) = self.default_project_id { + tokens.extra.insert( + "project_id".to_string(), + serde_json::json!(default_id), + ); + info!(project_id = %default_id, "Using default Antigravity project ID from env"); + return Ok(()); + } return Err(OAuthError::LoginFailed(format!( - "loadCodeAssist failed ({status}): {body}" + "loadCodeAssist failed ({status}) and no ANTIGRAVITY_DEFAULT_PROJECT_ID set" ))); } let json: serde_json::Value = resp.json().await?; - if let Some(project_id) = json.get("projectId").and_then(|v| v.as_str()) { + + // Extract project ID — openclaw uses cloudaicompanionProject (string or {id: ...}) + let project_id = json + .get("cloudaicompanionProject") + .and_then(|v| { + v.as_str().map(String::from).or_else(|| { + v.get("id").and_then(|id| id.as_str()).map(String::from) + }) + }) + .or_else(|| { + json.get("projectId") + .and_then(|v| v.as_str()) + .map(String::from) + }); + + if let Some(pid) = project_id { + debug!(project_id = %pid, "Discovered Antigravity project ID"); tokens .extra - .insert("project_id".to_string(), serde_json::json!(project_id)); - debug!(project_id, "Discovered Antigravity project ID"); + .insert("project_id".to_string(), serde_json::json!(pid)); + } else if let Some(ref default_id) = self.default_project_id { + tokens.extra.insert( + "project_id".to_string(), + serde_json::json!(default_id), + ); + info!(project_id = %default_id, "API returned no project ID, using default from env"); } if let Some(tier) = json.get("tier").and_then(|v| v.as_str()) { tokens @@ -264,6 +312,13 @@ fn generate_pkce() -> (String, String) { } /// Wrap an OpenAI-format request body into Antigravity/Gemini-style format. +/// +/// Translates OpenAI chat/completions fields to Gemini generateContent fields: +/// - `messages` → `contents` + `systemInstruction` +/// - `max_tokens`/`max_completion_tokens` → `generationConfig.maxOutputTokens` +/// - `temperature`, `top_p`, `stop` → `generationConfig` +/// - `tools` (OpenAI function-calling) → Gemini `tools[].functionDeclarations` +/// - Strips OpenAI-only fields (`stream`, `stream_options`, `store`, etc.) pub fn wrap_antigravity_request( body: &[u8], project_id: &str, @@ -272,21 +327,239 @@ pub fn wrap_antigravity_request( OAuthError::LoginFailed(format!("failed to parse request body as JSON: {e}")) })?; - let model = original + let obj = original.as_object().ok_or_else(|| { + OAuthError::LoginFailed("request body is not a JSON object".to_string()) + })?; + + let raw_model = obj .get("model") .and_then(|v| v.as_str()) .unwrap_or("gemini-2.0-flash"); + // Strip provider prefix (e.g. "google/gemini-2.5-flash" → "gemini-2.5-flash") + let model = raw_model + .split_once('/') + .map(|(_, id)| id) + .unwrap_or(raw_model); + + let mut request = serde_json::Map::new(); + + // Translate messages → contents + systemInstruction + if let Some(messages) = obj.get("messages").and_then(|v| v.as_array()) { + let mut system_parts: Vec = Vec::new(); + let mut contents: Vec = Vec::new(); + + for msg in messages { + let role = msg.get("role").and_then(|v| v.as_str()).unwrap_or(""); + let parts = message_content_to_gemini_parts(msg); + + match role { + "system" => { + system_parts.extend(parts); + } + "assistant" => { + contents.push(serde_json::json!({ + "role": "model", + "parts": parts, + })); + } + _ => { + // "user", "tool", and anything else → keep role as-is + contents.push(serde_json::json!({ + "role": role, + "parts": parts, + })); + } + } + } + + if !system_parts.is_empty() { + request.insert( + "systemInstruction".to_string(), + serde_json::json!({ "parts": system_parts }), + ); + } + request.insert("contents".to_string(), serde_json::json!(contents)); + } + + // Translate generation parameters → generationConfig + let mut gen_config = serde_json::Map::new(); + if let Some(v) = obj.get("max_tokens").or(obj.get("max_completion_tokens")) { + gen_config.insert("maxOutputTokens".to_string(), v.clone()); + } + if let Some(v) = obj.get("temperature") { + gen_config.insert("temperature".to_string(), v.clone()); + } + if let Some(v) = obj.get("top_p") { + gen_config.insert("topP".to_string(), v.clone()); + } + if let Some(v) = obj.get("stop") { + // OpenAI: stop can be string or array; Gemini: stopSequences is always array + let sequences = if v.is_string() { + serde_json::json!([v]) + } else { + v.clone() + }; + gen_config.insert("stopSequences".to_string(), sequences); + } + if !gen_config.is_empty() { + request.insert( + "generationConfig".to_string(), + serde_json::Value::Object(gen_config), + ); + } + + // Translate tools (OpenAI function-calling → Gemini functionDeclarations) + if let Some(tools) = obj.get("tools").and_then(|v| v.as_array()) { + let mut func_decls: Vec = Vec::new(); + for tool in tools { + if tool.get("type").and_then(|v| v.as_str()) == Some("function") { + if let Some(func) = tool.get("function") { + let mut decl = serde_json::Map::new(); + if let Some(name) = func.get("name") { + decl.insert("name".to_string(), name.clone()); + } + if let Some(desc) = func.get("description") { + decl.insert("description".to_string(), desc.clone()); + } + if let Some(params) = func.get("parameters") { + decl.insert( + "parameters".to_string(), + sanitize_schema_for_gemini(params.clone()), + ); + } + func_decls.push(serde_json::Value::Object(decl)); + } + } + } + if !func_decls.is_empty() { + request.insert( + "tools".to_string(), + serde_json::json!([{ "functionDeclarations": func_decls }]), + ); + } + } + + debug!( + raw_model = raw_model, + model = model, + project = project_id, + "Antigravity request: wrapping for upstream" + ); let wrapped = serde_json::json!({ "project": project_id, "model": model, - "request": original, + "user_prompt_id": uuid::Uuid::new_v4().to_string(), + "request": request, }); serde_json::to_vec(&wrapped) .map_err(|e| OAuthError::LoginFailed(format!("failed to serialize wrapped body: {e}"))) } +/// Convert an OpenAI message's `content` field to Gemini `parts` array. +fn message_content_to_gemini_parts(msg: &serde_json::Value) -> Vec { + let Some(content) = msg.get("content") else { + return vec![]; + }; + + // String content → single text part + if let Some(text) = content.as_str() { + return vec![serde_json::json!({ "text": text })]; + } + + // Array content (multimodal) → convert each part + if let Some(parts) = content.as_array() { + return parts + .iter() + .filter_map(|part| { + match part.get("type").and_then(|v| v.as_str()) { + Some("text") => { + let text = part.get("text").and_then(|v| v.as_str()).unwrap_or(""); + Some(serde_json::json!({ "text": text })) + } + Some("image_url") => { + // Convert OpenAI image_url to Gemini inlineData or fileData + let url = part + .get("image_url") + .and_then(|v| v.get("url")) + .and_then(|v| v.as_str()) + .unwrap_or(""); + if let Some(rest) = url.strip_prefix("data:") { + // data URI → inlineData + if let Some((mime, data)) = rest.split_once(";base64,") { + return Some(serde_json::json!({ + "inlineData": { + "mimeType": mime, + "data": data, + } + })); + } + } + // URL → fileData + Some(serde_json::json!({ + "fileData": { + "fileUri": url, + } + })) + } + _ => None, + } + }) + .collect(); + } + + vec![] +} + +/// JSON Schema keywords that the Cloud Code Assist API rejects. +/// Matching openclaw's `GEMINI_UNSUPPORTED_SCHEMA_KEYWORDS`. +const GEMINI_SCHEMA_REJECTED: &[&str] = &[ + "patternProperties", + "additionalProperties", + "$schema", + "$id", + "$ref", + "$defs", + "definitions", + "examples", + "minLength", + "maxLength", + "minimum", + "maximum", + "multipleOf", + "pattern", + "format", + "minItems", + "maxItems", + "uniqueItems", + "minProperties", + "maxProperties", +]; + +/// Recursively strip JSON Schema fields that the Gemini API does not support. +fn sanitize_schema_for_gemini(mut value: serde_json::Value) -> serde_json::Value { + let Some(obj) = value.as_object_mut() else { + return value; + }; + + obj.retain(|key, _| !GEMINI_SCHEMA_REJECTED.contains(&key.as_str())); + + // Recurse into nested schemas + if let Some(props) = obj.get_mut("properties") { + if let Some(map) = props.as_object_mut() { + for v in map.values_mut() { + *v = sanitize_schema_for_gemini(v.clone()); + } + } + } + if let Some(items) = obj.get_mut("items") { + *items = sanitize_schema_for_gemini(items.clone()); + } + + value +} + #[async_trait] impl OAuthProvider for AntigravityProvider { fn id(&self) -> &str { @@ -301,9 +574,11 @@ impl OAuthProvider for AntigravityProvider { true } - async fn login_browser(&self, callback_port: u16) -> Result { + async fn login_browser(&self, _callback_port: u16) -> Result { + // This client ID requires a fixed redirect URI registered in Google's OAuth console. + const ANTIGRAVITY_CALLBACK_PORT: u16 = 51121; let (verifier, challenge) = generate_pkce(); - let redirect_uri = format!("http://localhost:{callback_port}/oauth-callback"); + let redirect_uri = format!("http://localhost:{ANTIGRAVITY_CALLBACK_PORT}/oauth-callback"); let state: String = uuid::Uuid::new_v4().to_string(); let auth_url = format!( @@ -324,7 +599,7 @@ impl OAuthProvider for AntigravityProvider { } let (code, received_state) = - wait_for_oauth_callback(callback_port).await.map_err(|e| { + wait_for_oauth_callback(ANTIGRAVITY_CALLBACK_PORT).await.map_err(|e| { OAuthError::LoginFailed(format!("callback server failed: {e}")) })?; @@ -338,30 +613,39 @@ impl OAuthProvider for AntigravityProvider { } async fn login_headless(&self) -> Result { + const ANTIGRAVITY_CALLBACK_PORT: u16 = 51121; let (verifier, challenge) = generate_pkce(); - let redirect_uri = "https://codeassist.google.com/authcode".to_string(); + let redirect_uri = format!("http://localhost:{ANTIGRAVITY_CALLBACK_PORT}/oauth-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&access_type=offline&prompt=consent", + "{}?response_type=code&client_id={}&redirect_uri={}&scope={}&code_challenge={}&code_challenge_method=S256&state={}&access_type=offline&prompt=consent", self.auth_url, urlencoding::encode(&self.client_id), urlencoding::encode(&redirect_uri), urlencoding::encode(&self.scopes.join(" ")), urlencoding::encode(&challenge), + urlencoding::encode(&state), ); println!(); println!(" Visit this URL to authenticate:"); println!(" {auth_url}"); println!(); - println!(" After authorizing, Google will show an authorization code."); - println!(" Copy and paste it below:"); - let code = crate::tui::prompt_text("Authorization code", None) - .map_err(|e| OAuthError::LoginFailed(format!("failed to read code: {e}")))?; + // Start a local HTTP server to receive the OAuth callback, then wait. + let (code, received_state) = + wait_for_oauth_callback(ANTIGRAVITY_CALLBACK_PORT).await.map_err(|e| { + OAuthError::LoginFailed(format!("callback server failed: {e}")) + })?; - self.exchange_code(code.trim(), &verifier, &redirect_uri) - .await + 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 refresh(&self, refresh_token: &str) -> Result { @@ -415,13 +699,47 @@ impl OAuthProvider for AntigravityProvider { } fn upstream_url(&self, _tokens: &OAuthTokens) -> Option { - Some(format!( - "{}/v1internal:streamGenerateContent?alt=sse", - self.endpoints[0] - )) + Some(self.endpoints[0].clone()) + } + + fn rewrite_request_path(&self, _path: &str) -> Option { + Some("/v1internal:streamGenerateContent?alt=sse".to_string()) + } + + fn response_format(&self, _original_path: &str) -> Option { + Some(super::ResponseFormat::GeminiSse) } } +/// Parse code and state from a pasted redirect URL (headless flow). +fn parse_redirect_url(input: &str) -> Result<(String, String), OAuthError> { + let query = input + .split('?') + .nth(1) + .ok_or_else(|| OAuthError::LoginFailed("no query string in redirect URL".to_string()))?; + + 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(), + _ => {} + } + } + } + + if code.is_empty() { + return Err(OAuthError::LoginFailed( + "no authorization code found in redirect URL".to_string(), + )); + } + + Ok((code, state)) +} + /// Wait for an OAuth callback on a local HTTP server (same as codex). async fn wait_for_oauth_callback( port: u16, @@ -522,7 +840,15 @@ mod tests { assert_eq!(parsed["project"], "proj-abc-123"); assert_eq!(parsed["model"], "gemini-3-pro"); - assert_eq!(parsed["request"]["messages"][0]["content"], "hello"); + + // messages should be translated to Gemini contents format + let contents = parsed["request"]["contents"].as_array().unwrap(); + assert_eq!(contents.len(), 1); + assert_eq!(contents[0]["role"], "user"); + assert_eq!(contents[0]["parts"][0]["text"], "hello"); + + // Raw OpenAI fields should not be present + assert!(parsed["request"].get("messages").is_none()); } #[test] @@ -538,6 +864,173 @@ mod tests { assert_eq!(parsed["model"], "gemini-2.0-flash"); } + #[test] + fn test_wrap_antigravity_system_message() { + let body = serde_json::json!({ + "model": "gemini-2.0-flash", + "messages": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "hi"} + ] + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + + let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); + + // System messages become systemInstruction + assert_eq!( + parsed["request"]["systemInstruction"]["parts"][0]["text"], + "You are helpful." + ); + // Only non-system messages in contents + let contents = parsed["request"]["contents"].as_array().unwrap(); + assert_eq!(contents.len(), 1); + assert_eq!(contents[0]["role"], "user"); + } + + #[test] + fn test_wrap_antigravity_assistant_role() { + let body = serde_json::json!({ + "model": "gemini-2.0-flash", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"} + ] + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + + let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); + + let contents = parsed["request"]["contents"].as_array().unwrap(); + assert_eq!(contents[1]["role"], "model"); + assert_eq!(contents[1]["parts"][0]["text"], "hello"); + } + + #[test] + fn test_wrap_antigravity_generation_config() { + let body = serde_json::json!({ + "model": "gemini-2.0-flash", + "messages": [{"role": "user", "content": "hi"}], + "max_completion_tokens": 1024, + "temperature": 0.7, + "top_p": 0.9, + "stop": ["\n"] + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + + let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); + + let gc = &parsed["request"]["generationConfig"]; + assert_eq!(gc["maxOutputTokens"], 1024); + assert_eq!(gc["temperature"], 0.7); + assert_eq!(gc["topP"], 0.9); + assert_eq!(gc["stopSequences"], serde_json::json!(["\n"])); + } + + #[test] + fn test_wrap_antigravity_tools() { + let body = serde_json::json!({ + "model": "gemini-2.0-flash", + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": {"type": "object", "properties": {}} + } + } + ] + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + + let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); + + let decls = &parsed["request"]["tools"][0]["functionDeclarations"]; + assert_eq!(decls[0]["name"], "get_weather"); + assert_eq!(decls[0]["description"], "Get weather"); + } + + #[test] + fn test_wrap_antigravity_strips_openai_fields() { + let body = serde_json::json!({ + "model": "gemini-2.0-flash", + "messages": [{"role": "user", "content": "hi"}], + "stream": true, + "stream_options": {"include_usage": true}, + "store": false + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + + let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); + + assert!(parsed["request"].get("stream").is_none()); + assert!(parsed["request"].get("stream_options").is_none()); + assert!(parsed["request"].get("store").is_none()); + assert!(parsed["request"].get("messages").is_none()); + } + + #[test] + fn test_sanitize_schema_strips_unsupported_fields() { + let schema = serde_json::json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "The name", + "patternProperties": {"^x-": {"type": "string"}}, + "additionalProperties": false, + "minLength": 1, + "maxLength": 100 + }, + "tags": { + "type": "array", + "items": { + "type": "string", + "$ref": "#/defs/Tag", + "default": "foo" + } + } + }, + "required": ["name"], + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "$defs": {} + }); + + let sanitized = sanitize_schema_for_gemini(schema); + + // Top-level rejected fields stripped + assert!(sanitized.get("$schema").is_none()); + assert!(sanitized.get("additionalProperties").is_none()); + assert!(sanitized.get("$defs").is_none()); + // Supported fields preserved + assert_eq!(sanitized["type"], "object"); + assert_eq!(sanitized["required"], serde_json::json!(["name"])); + + // Nested property rejected fields stripped + let name_prop = &sanitized["properties"]["name"]; + assert!(name_prop.get("patternProperties").is_none()); + assert!(name_prop.get("additionalProperties").is_none()); + assert!(name_prop.get("minLength").is_none()); + assert!(name_prop.get("maxLength").is_none()); + assert_eq!(name_prop["type"], "string"); + assert_eq!(name_prop["description"], "The name"); + + // Items rejected fields stripped, but non-rejected fields preserved + let items = &sanitized["properties"]["tags"]["items"]; + assert!(items.get("$ref").is_none()); + assert_eq!(items["type"], "string"); + // "default" is NOT in the rejected list, so it's preserved + assert_eq!(items["default"], "foo"); + } + #[test] fn test_antigravity_provider_defaults() { let provider = test_provider(); @@ -618,6 +1111,16 @@ mod tests { let url = provider.upstream_url(&tokens).unwrap(); assert!(url.contains("cloudcode-pa.googleapis.com")); - assert!(url.contains("streamGenerateContent")); + assert!(!url.contains("streamGenerateContent"), "upstream_url should be base only"); + } + + #[test] + fn test_rewrite_request_path() { + let provider = test_provider(); + let rewritten = provider.rewrite_request_path("/v1/chat/completions"); + assert_eq!( + rewritten.as_deref(), + Some("/v1internal:streamGenerateContent?alt=sse") + ); } } diff --git a/src/oauth/codex.rs b/src/oauth/codex.rs index 32443ef..2364ca1 100644 --- a/src/oauth/codex.rs +++ b/src/oauth/codex.rs @@ -416,6 +416,14 @@ impl OAuthProvider for CodexProvider { 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. diff --git a/src/oauth/mod.rs b/src/oauth/mod.rs index 731c32b..e4c68e7 100644 --- a/src/oauth/mod.rs +++ b/src/oauth/mod.rs @@ -148,6 +148,21 @@ pub trait OAuthProvider: Send + Sync + fmt::Debug { 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. @@ -437,6 +452,18 @@ impl OAuthRegistry { 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) } diff --git a/src/translate.rs b/src/translate.rs index ee0e3bb..c54a0de 100644 --- a/src/translate.rs +++ b/src/translate.rs @@ -4,6 +4,7 @@ use futures_util::Stream; use serde_json::Value; use std::pin::Pin; use std::task::{Context, Poll}; +use tracing::debug; #[derive(Debug, thiserror::Error)] pub enum TranslateError { @@ -82,11 +83,19 @@ pub fn chat_completions_to_responses(body: &[u8]) -> Result, TranslateEr Ok(serde_json::to_vec(&Value::Object(result))?) } -/// Convert chat/completions content types to responses API content types. -/// - `type: "text"` → `type: "input_text"` -/// - `type: "image_url"` → `type: "input_image"` with `image_url` → `image_url` -/// String content is left as-is (the Responses API accepts string content directly). +/// Convert a chat/completions message to a Responses API input item. +/// - Adds `type: "message"` (required by Responses API) +/// - Converts content `type: "text"` → `type: "input_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 { + // 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; }; @@ -420,6 +429,220 @@ 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)) +} + #[cfg(test)] mod tests { use super::*; @@ -439,6 +662,7 @@ mod tests { assert!(parsed.get("instructions").is_none()); 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()); @@ -710,6 +934,7 @@ mod tests { 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?"); @@ -728,6 +953,8 @@ mod tests { 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"); } } From a875a37e1d7499958c97d3d89843d76d6b4cfbed Mon Sep 17 00:00:00 2001 From: u20024804 Date: Mon, 23 Feb 2026 03:56:00 -0500 Subject: [PATCH 06/16] support codex and antigravity oauth --- src/app.rs | 15 ++- src/app/tests.rs | 48 +++++++++- src/oauth/codex.rs | 6 ++ src/translate.rs | 234 +++++++++++++++++++++++++++++++++++++++++++-- 4 files changed, 289 insertions(+), 14 deletions(-) diff --git a/src/app.rs b/src/app.rs index 89e9647..1c34b89 100644 --- a/src/app.rs +++ b/src/app.rs @@ -641,14 +641,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" + ); + let (parts, body) = response.into_parts(); + let dlp_body = crate::translate::wrap_body_with_dlp_sse_stream( + body, + state.dlp_scanner.clone(), ); - response + Response::from_parts(parts, dlp_body) } } else { trace!("Response DLP scanning disabled"); @@ -710,7 +714,8 @@ async fn forward_oauth_request( .oauth_registry .response_format(oauth_provider_id, &original_path) .map_err(|e| format!("OAuth response format check failed: {e}"))?; - let stream_requested = serde_json::from_slice::(&body_bytes) + // 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); diff --git a/src/app/tests.rs b/src/app/tests.rs index 381c1be..925db69 100644 --- a/src/app/tests.rs +++ b/src/app/tests.rs @@ -1459,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")) @@ -1501,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, diff --git a/src/oauth/codex.rs b/src/oauth/codex.rs index 2364ca1..4236919 100644 --- a/src/oauth/codex.rs +++ b/src/oauth/codex.rs @@ -476,6 +476,7 @@ async fn wait_for_oauth_callback( /// 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(); @@ -486,6 +487,11 @@ fn fixup_for_chatgpt_backend(body: &[u8]) -> Vec { } } 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()) } diff --git a/src/translate.rs b/src/translate.rs index c54a0de..1e547fc 100644 --- a/src/translate.rs +++ b/src/translate.rs @@ -1,10 +1,12 @@ +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; +use tracing::{debug, warn}; #[derive(Debug, thiserror::Error)] pub enum TranslateError { @@ -59,9 +61,11 @@ pub fn chat_completions_to_responses(body: &[u8]) -> Result, TranslateEr } } - if !system_parts.is_empty() { - result.insert("instructions".to_string(), Value::String(system_parts.join("\n"))); - } + // 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 @@ -85,10 +89,14 @@ pub fn chat_completions_to_responses(body: &[u8]) -> Result, TranslateEr /// Convert a chat/completions message to a Responses API input item. /// - Adds `type: "message"` (required by Responses API) -/// - Converts content `type: "text"` → `type: "input_text"` +/// - 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") { @@ -109,7 +117,8 @@ fn convert_message_content(mut msg: Value) -> Value { }; match obj.get("type").and_then(Value::as_str) { Some("text") => { - obj.insert("type".to_string(), Value::String("input_text".to_string())); + 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())); @@ -643,6 +652,153 @@ 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::*; @@ -659,7 +815,7 @@ mod tests { let parsed: Value = serde_json::from_slice(&result).unwrap(); assert_eq!(parsed["model"], "gpt-4o-mini"); - assert!(parsed.get("instructions").is_none()); + 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"); @@ -957,4 +1113,68 @@ mod tests { 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"); + } } From d6e90a8146f8282a5fe4321b23f1a728c8987216 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Tue, 24 Feb 2026 00:11:13 -0500 Subject: [PATCH 07/16] support codex and antigravity oauth --- Dockerfile | 16 ++++-- Dockerfile.openclaw | 10 ++++ README.md | 6 ++ docker-compose.yml | 35 ++++++++++++ docs/docker.md | 41 ++++++++++++-- src/config.rs | 118 +++++++++++++++++++++++++++++++++++++++ src/main.rs | 8 ++- src/oauth/antigravity.rs | 2 +- src/oauth/codex.rs | 2 +- src/oauth/mod.rs | 12 ++++ 10 files changed, 236 insertions(+), 14 deletions(-) create mode 100644 Dockerfile.openclaw create mode 100644 docker-compose.yml diff --git a/Dockerfile b/Dockerfile index b1b62ac..0378d7c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,8 +1,12 @@ -FROM debian:sid-slim -RUN apt-get update && apt-get install -y ca-certificates && rm -rf /var/lib/apt/lists/* +FROM openclaw ENV TERM=xterm-256color -EXPOSE 18790 -WORKDIR /etc/clawshell +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 -COPY .env .env -ENTRYPOINT ["clawshell"] +RUN sudo useradd clawshell +USER node +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..558abf5 --- /dev/null +++ b/Dockerfile.openclaw @@ -0,0 +1,10 @@ +FROM debian:13.3-slim +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 useradd -m node +RUN echo "node ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers +USER node +RUN nodeenv --node=24.13.1 --npm=v11.8.0 /home/node/nodeenv +RUN cd /home/node && . ./nodeenv/bin/activate && npm install -g openclaw@latest + +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..023c091 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,35 @@ +services: + openclaw-gateway: + build: + context: . + dockerfile: Dockerfile.openclaw + image: openclaw + container_name: openclaw-gateway + command: ["gateway","run","--bind","lan","--port","18789"] + ports: + - "18789:18789" + volumes: + - config-openclaw:/home/node/.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" + ports: + - "18790:18790" + - "51121:51121" + volumes: + - config-clawshell:/etc/clawshell + restart: unless-stopped + +volumes: + clawshell-config: diff --git a/docs/docker.md b/docs/docker.md index 5de6cd0..0a501c2 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -31,6 +31,14 @@ Run the interactive onboard wizard to generate configuration: 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`: @@ -70,6 +78,18 @@ 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`: @@ -97,6 +117,13 @@ docker run --rm -it \ 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 | @@ -121,9 +148,14 @@ No extra environment variables are needed for Codex. Requires `GOOGLE_OAUTH_CLIENT_ID` and `GOOGLE_OAUTH_CLIENT_SECRET` (provided via the `.env` file baked into the image at build time). -Uses a copy/paste flow. The wizard prints a Google authorization URL. Open it -in your browser, authorize, then copy the authorization code from the result -page and paste it back into the terminal. +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 @@ -145,7 +177,8 @@ cargo build --release docker build -t clawshell . # 3. Onboard (interactive — creates config in the volume) -docker run --rm -it -v clawshell-config:/etc/clawshell clawshell onboard +# 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 \ diff --git a/src/config.rs b/src/config.rs index 25c9074..7344b4c 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)] @@ -62,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 { @@ -425,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> { @@ -582,6 +641,7 @@ mod tests { Ok(()) } + #[test] fn test_valid_config_fixtures() { let paths = @@ -934,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/main.rs b/src/main.rs index 4c1670c..636816c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -578,8 +578,12 @@ async fn cmd_start_inner(config_path: &str) -> Result<(), Box Result<(), Box Result<(String, String), Box> { use tokio::io::{AsyncReadExt, AsyncWriteExt}; - let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}")).await?; + let listener = tokio::net::TcpListener::bind(super::callback_bind_addr(port)).await?; let (mut stream, _) = listener.accept().await?; let mut buf = vec![0u8; 4096]; diff --git a/src/oauth/codex.rs b/src/oauth/codex.rs index 4236919..e34da5c 100644 --- a/src/oauth/codex.rs +++ b/src/oauth/codex.rs @@ -433,7 +433,7 @@ async fn wait_for_oauth_callback( ) -> Result<(String, String), Box> { use tokio::io::{AsyncReadExt, AsyncWriteExt}; - let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{port}")).await?; + let listener = tokio::net::TcpListener::bind(super::callback_bind_addr(port)).await?; let (mut stream, _) = listener.accept().await?; let mut buf = vec![0u8; 4096]; diff --git a/src/oauth/mod.rs b/src/oauth/mod.rs index e4c68e7..4391d78 100644 --- a/src/oauth/mod.rs +++ b/src/oauth/mod.rs @@ -187,6 +187,18 @@ 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 { From b5771c34e7064c8ba747a37d05371bccab29114f Mon Sep 17 00:00:00 2001 From: u20024804 Date: Tue, 24 Feb 2026 09:08:05 -0500 Subject: [PATCH 08/16] support codex and antigravity oauth --- .env | 3 +++ Dockerfile | 10 ++++++++-- Dockerfile.openclaw | 14 ++++++++++++-- docker-compose.yml | 3 +++ 4 files changed, 26 insertions(+), 4 deletions(-) diff --git a/.env b/.env index 201f16d..4a5251d 100644 --- a/.env +++ b/.env @@ -1,3 +1,6 @@ GOOGLE_OAUTH_CLIENT_ID= GOOGLE_OAUTH_CLIENT_SECRET= ANTIGRAVITY_DEFAULT_PROJECT_ID= +UID= +GID= +USER= diff --git a/Dockerfile b/Dockerfile index 0378d7c..a59735e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,12 +1,18 @@ FROM openclaw +ARG UID=1000 +ARG GID=1000 +ARG USER=node +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 clawshell USER node 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"] +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 index 558abf5..2b03cb4 100644 --- a/Dockerfile.openclaw +++ b/Dockerfile.openclaw @@ -1,10 +1,20 @@ FROM debian:13.3-slim +ARG UID=1000 +ARG GID=1000 +ARG USER=node +ENV UID=$UID +ENV GID=$GID +ENV USER=$USER + +RUN groupadd -r -g $GID $USER && useradd -r -m -u $UID -g $GID -s /bin/bash $USER +RUN useradd -r -m -o -u $UID -g $GID -s /bin/bash node +RUN useradd -r -m -o -u $UID -g $GID -s /bin/bash clawshell 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 useradd -m node +RUN echo "$USER ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers RUN echo "node ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers USER node RUN nodeenv --node=24.13.1 --npm=v11.8.0 /home/node/nodeenv -RUN cd /home/node && . ./nodeenv/bin/activate && npm install -g openclaw@latest +RUN cd /home/node && . ./nodeenv/bin/activate && npm config set prefix /home/node && npm install -g openclaw@2026.2.23 ENV PATH="/home/node/nodeenv/bin:$PATH" ENTRYPOINT ["/home/node/nodeenv/bin/openclaw"] diff --git a/docker-compose.yml b/docker-compose.yml index 023c091..0484897 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -24,6 +24,9 @@ services: 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" From 4f4afde2f4eb32fb0879ead6a8d17bdabb538821 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Tue, 24 Feb 2026 09:09:24 -0500 Subject: [PATCH 09/16] .ginreo --- docker-compose.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/docker-compose.yml b/docker-compose.yml index 0484897..d28b3c5 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -8,8 +8,12 @@ services: command: ["gateway","run","--bind","lan","--port","18789"] ports: - "18789:18789" + environment: + UID: $UID + GID: $GID + USER: $USER volumes: - - config-openclaw:/home/node/.openclaw + - config-openclaw:/home/$USER/.openclaw restart: unless-stopped clawshell: From ae4e8631bcf583a13d3c1da7a282897cceb9bc78 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Tue, 24 Feb 2026 12:15:04 -0500 Subject: [PATCH 10/16] support codex and antigravity oauth --- src/oauth/antigravity.rs | 60 ++++++++++++++++++++++++++++++++++------ 1 file changed, 52 insertions(+), 8 deletions(-) diff --git a/src/oauth/antigravity.rs b/src/oauth/antigravity.rs index 025adcc..a3ae53b 100644 --- a/src/oauth/antigravity.rs +++ b/src/oauth/antigravity.rs @@ -335,11 +335,9 @@ pub fn wrap_antigravity_request( .get("model") .and_then(|v| v.as_str()) .unwrap_or("gemini-2.0-flash"); - // Strip provider prefix (e.g. "google/gemini-2.5-flash" → "gemini-2.5-flash") - let model = raw_model - .split_once('/') - .map(|(_, id)| id) - .unwrap_or(raw_model); + // Strip provider prefix (e.g. "google/gemini-2.5-flash" or + // "custom-provider/google/gemini-2.5-flash" → "gemini-2.5-flash") + let model = raw_model.rsplit('/').next().unwrap_or(raw_model); let mut request = serde_json::Map::new(); @@ -353,7 +351,8 @@ pub fn wrap_antigravity_request( let parts = message_content_to_gemini_parts(msg); match role { - "system" => { + // OpenAI/Codex "developer" has system-like semantics. + "system" | "developer" => { system_parts.extend(parts); } "assistant" => { @@ -362,10 +361,17 @@ pub fn wrap_antigravity_request( "parts": parts, })); } + "user" => { + contents.push(serde_json::json!({ + "role": "user", + "parts": parts, + })); + } _ => { - // "user", "tool", and anything else → keep role as-is + // Gemini contents only support "user"/"model" roles. + // Coerce tool/unknown roles to "user" to avoid INVALID_ARGUMENT. contents.push(serde_json::json!({ - "role": role, + "role": "user", "parts": parts, })); } @@ -864,6 +870,20 @@ mod tests { assert_eq!(parsed["model"], "gemini-2.0-flash"); } + #[test] + fn test_wrap_antigravity_request_strips_nested_provider_prefixes() { + let body = serde_json::json!({ + "model": "custom-clawshell-18790/google/gemini-2.0-flash", + "messages": [{"role": "user", "content": "hello"}] + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + + let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); + + assert_eq!(parsed["model"], "gemini-2.0-flash"); + } + #[test] fn test_wrap_antigravity_system_message() { let body = serde_json::json!({ @@ -889,6 +909,30 @@ mod tests { assert_eq!(contents[0]["role"], "user"); } + #[test] + fn test_wrap_antigravity_developer_message_maps_to_system_instruction() { + let body = serde_json::json!({ + "model": "gemini-2.0-flash", + "messages": [ + {"role": "developer", "content": "Follow this policy."}, + {"role": "user", "content": "hello"} + ] + }); + let body_bytes = serde_json::to_vec(&body).unwrap(); + + let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); + + assert_eq!( + parsed["request"]["systemInstruction"]["parts"][0]["text"], + "Follow this policy." + ); + let contents = parsed["request"]["contents"].as_array().unwrap(); + assert_eq!(contents.len(), 1); + assert_eq!(contents[0]["role"], "user"); + assert_eq!(contents[0]["parts"][0]["text"], "hello"); + } + #[test] fn test_wrap_antigravity_assistant_role() { let body = serde_json::json!({ From ec9959c92b27042f9b4ab1da607414a532105b34 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Wed, 25 Feb 2026 03:29:12 -0500 Subject: [PATCH 11/16] support codex and antigravity oauth --- Dockerfile | 4 ++-- Dockerfile.openclaw | 16 ++++++++++------ 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/Dockerfile b/Dockerfile index a59735e..b54be39 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,7 +1,7 @@ FROM openclaw ARG UID=1000 ARG GID=1000 -ARG USER=node +ARG USER=app ENV UID=$UID ENV GID=$GID ENV USER=$USER @@ -12,7 +12,7 @@ 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 -USER node +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 index 2b03cb4..e569848 100644 --- a/Dockerfile.openclaw +++ b/Dockerfile.openclaw @@ -1,20 +1,24 @@ FROM debian:13.3-slim ARG UID=1000 ARG GID=1000 -ARG USER=node +ARG USER=app +ARG VERSION=latest ENV UID=$UID ENV GID=$GID ENV USER=$USER +ENV VERSION=$VERSION -RUN groupadd -r -g $GID $USER && useradd -r -m -u $UID -g $GID -s /bin/bash $USER -RUN useradd -r -m -o -u $UID -g $GID -s /bin/bash node +RUN groupadd -r -g $GID $USER +RUN useradd -r -m -u $UID -g $GID -s /bin/bash $USER RUN useradd -r -m -o -u $UID -g $GID -s /bin/bash clawshell + 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 echo "node ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers -USER node +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@2026.2.23 +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"] From 49fc42fad086b06256988d8de86f0b00be496725 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Wed, 25 Feb 2026 04:09:24 -0500 Subject: [PATCH 12/16] support codex and antigravity oauth --- docker-compose.yml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index d28b3c5..b9d930e 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -11,9 +11,9 @@ services: environment: UID: $UID GID: $GID - USER: $USER + USER: app volumes: - - config-openclaw:/home/$USER/.openclaw + - config-openclaw:/home/app/.openclaw restart: unless-stopped clawshell: @@ -30,7 +30,7 @@ services: CLAWSHELL_OAUTH_CALLBACK_HOST: "0.0.0.0" UID: $UID GID: $GID - USER: $USER + USER: app ports: - "18790:18790" - "51121:51121" @@ -39,4 +39,5 @@ services: restart: unless-stopped volumes: - clawshell-config: + config-openclaw: + config-clawshell: From e164fa55ac4dea1587dc126a09df5d1606eb8ad8 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Wed, 25 Feb 2026 04:28:19 -0500 Subject: [PATCH 13/16] support codex and antigravity oauth --- docker-compose.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index b9d930e..5274d75 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -11,9 +11,10 @@ services: environment: UID: $UID GID: $GID - USER: app + USER: $USER volumes: - config-openclaw:/home/app/.openclaw + - config-openclaw:/home/$USER/.openclaw restart: unless-stopped clawshell: @@ -30,7 +31,7 @@ services: CLAWSHELL_OAUTH_CALLBACK_HOST: "0.0.0.0" UID: $UID GID: $GID - USER: app + USER: $USER ports: - "18790:18790" - "51121:51121" From 3f27209f0326bbd905aba83a3a7bb94c9351d582 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Wed, 25 Feb 2026 05:10:25 -0500 Subject: [PATCH 14/16] support codex and antigravity oauth --- Dockerfile | 1 + Dockerfile.openclaw | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index b54be39..8a81d96 100644 --- a/Dockerfile +++ b/Dockerfile @@ -12,6 +12,7 @@ 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 diff --git a/Dockerfile.openclaw b/Dockerfile.openclaw index e569848..76ae4b2 100644 --- a/Dockerfile.openclaw +++ b/Dockerfile.openclaw @@ -10,7 +10,6 @@ ENV VERSION=$VERSION RUN groupadd -r -g $GID $USER RUN useradd -r -m -u $UID -g $GID -s /bin/bash $USER -RUN useradd -r -m -o -u $UID -g $GID -s /bin/bash clawshell 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 From d8b3ba0150b7a2faccd44bad756c05ba5c38a1e9 Mon Sep 17 00:00:00 2001 From: u20024804 Date: Fri, 27 Feb 2026 07:21:24 -0500 Subject: [PATCH 15/16] hide Google Antigravity by default --- docker-compose.yml | 6 +++--- src/onboard/interactive.rs | 30 +++++++++++++++++++++--------- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index 5274d75..b9d608e 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -13,8 +13,8 @@ services: GID: $GID USER: $USER volumes: - - config-openclaw:/home/app/.openclaw - - config-openclaw:/home/$USER/.openclaw + - ./config/openclaw:/home/app/.openclaw + - ./config/openclaw:/home/$USER/.openclaw restart: unless-stopped clawshell: @@ -36,7 +36,7 @@ services: - "18790:18790" - "51121:51121" volumes: - - config-clawshell:/etc/clawshell + - ./config/clawshell:/etc/clawshell restart: unless-stopped volumes: diff --git a/src/onboard/interactive.rs b/src/onboard/interactive.rs index 0e20d26..abafc77 100644 --- a/src/onboard/interactive.rs +++ b/src/onboard/interactive.rs @@ -433,13 +433,25 @@ pub fn collect_onboard_config_tui() -> Result Result Some(MENU_ANTIGRAVITY), + (Some("oauth"), Some("antigravity"), _) if show_antigravity => Some(MENU_ANTIGRAVITY), (Some("oauth"), Some("codex"), _) | (Some("oauth"), _, _) => Some(MENU_CODEX), (_, _, Some("anthropic")) => Some(MENU_ANTHROPIC), (_, _, Some("openrouter")) => Some(MENU_OPENROUTER), @@ -459,7 +471,7 @@ pub fn collect_onboard_config_tui() -> Result Date: Sun, 1 Mar 2026 08:47:12 -0500 Subject: [PATCH 16/16] remove antigravity support --- src/config.rs | 2 +- src/main.rs | 5 - src/oauth/antigravity.rs | 1170 ------------------------------------ src/oauth/mod.rs | 7 +- src/onboard/interactive.rs | 36 +- src/onboard/types.rs | 2 +- 6 files changed, 8 insertions(+), 1214 deletions(-) delete mode 100644 src/oauth/antigravity.rs diff --git a/src/config.rs b/src/config.rs index 7344b4c..73ece6b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -110,7 +110,7 @@ pub struct KeyMapping { /// Authentication method for this key. Defaults to "static". #[serde(default)] pub auth: KeyAuthMethod, - /// Which OAuth provider supplies the token (e.g. "codex", "antigravity"). + /// 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, diff --git a/src/main.rs b/src/main.rs index 636816c..83ba5b7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -630,7 +630,6 @@ async fn build_oauth_registry( ) -> Result> { use crate::oauth::{ OAuthRegistry, TokenStorage, - antigravity::AntigravityProvider, codex::CodexProvider, }; use std::sync::Arc; @@ -647,10 +646,6 @@ async fn build_oauth_registry( let provider = CodexProvider::from_config(provider_config); registry.register(Arc::new(provider)); } - "antigravity" => { - let provider = AntigravityProvider::from_config(provider_config); - registry.register(Arc::new(provider)); - } other => { return Err(format!("Unknown OAuth provider type: '{other}'").into()); } diff --git a/src/oauth/antigravity.rs b/src/oauth/antigravity.rs deleted file mode 100644 index a3ae53b..0000000 --- a/src/oauth/antigravity.rs +++ /dev/null @@ -1,1170 +0,0 @@ -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, warn}; - -const DEFAULT_AUTH_URL: &str = "https://accounts.google.com/o/oauth2/v2/auth"; -const DEFAULT_TOKEN_URL: &str = "https://oauth2.googleapis.com/token"; -const DEFAULT_SCOPES: &[&str] = &[ - "https://www.googleapis.com/auth/cloud-platform", - "https://www.googleapis.com/auth/userinfo.email", - "https://www.googleapis.com/auth/userinfo.profile", - "https://www.googleapis.com/auth/cclog", - "https://www.googleapis.com/auth/experimentsandconfigs", -]; - -const ENDPOINT_PRODUCTION: &str = "https://cloudcode-pa.googleapis.com"; -const ENDPOINT_DAILY: &str = "https://daily-cloudcode-pa.sandbox.googleapis.com"; -const ENDPOINT_ALT: &str = "https://codeassist.googleapis.com/v1"; - -#[derive(Debug)] -pub struct AntigravityProvider { - client_id: String, - client_secret: String, - auth_url: String, - token_url: String, - scopes: Vec, - http_client: reqwest::Client, - endpoints: Vec, - default_project_id: Option, -} - -impl AntigravityProvider { - pub fn new( - client_id: Option<&str>, - auth_url: Option<&str>, - token_url: Option<&str>, - scopes: Option<&[String]>, - ) -> Self { - Self::new_with_secret(client_id, None, auth_url, token_url, scopes) - } - - pub fn new_with_secret( - client_id: Option<&str>, - client_secret: Option<&str>, - auth_url: Option<&str>, - token_url: Option<&str>, - scopes: Option<&[String]>, - ) -> Self { - // Load .env file if present (ignored if missing) - let _ = dotenvy::dotenv(); - - // Env vars take priority, then explicit constructor arguments. - // No hardcoded defaults — credentials must come from env, .env file, or config. - let resolved_client_id = std::env::var("GOOGLE_OAUTH_CLIENT_ID") - .ok() - .or_else(|| client_id.map(String::from)) - .expect("GOOGLE_OAUTH_CLIENT_ID env var or client_id argument is required"); - let resolved_client_secret = std::env::var("GOOGLE_OAUTH_CLIENT_SECRET") - .ok() - .or_else(|| client_secret.map(String::from)) - .expect("GOOGLE_OAUTH_CLIENT_SECRET env var or client_secret argument is required"); - let default_project_id = std::env::var("ANTIGRAVITY_DEFAULT_PROJECT_ID").ok(); - - Self { - client_id: resolved_client_id, - client_secret: resolved_client_secret, - 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"), - endpoints: vec![ - ENDPOINT_PRODUCTION.to_string(), - ENDPOINT_DAILY.to_string(), - ENDPOINT_ALT.to_string(), - ], - default_project_id, - } - } - - 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.as_str()), - ("client_secret", self.client_secret.as_str()), - ("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?; - let mut tokens = parse_google_token_response(&json)?; - - // Discover project ID after successful login - if let Err(e) = self.discover_project_id(&mut tokens).await { - warn!(error = %e, "Failed to discover Antigravity project ID"); - } - - Ok(tokens) - } - - async fn exchange_refresh_token( - &self, - refresh_token: &str, - ) -> Result { - let params = [ - ("grant_type", "refresh_token"), - ("client_id", self.client_id.as_str()), - ("client_secret", self.client_secret.as_str()), - ("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?; - // Google refresh responses may not include a new refresh token; - // the caller should preserve the original refresh token. - let mut tokens = parse_google_token_response(&json)?; - if tokens.refresh_token.is_none() { - tokens.refresh_token = Some(refresh_token.to_string()); - } - - // Re-discover project ID with the fresh access token. - // This also recovers from initial login discovery failures. - if let Err(e) = self.discover_project_id(&mut tokens).await { - warn!(error = %e, "Failed to discover Antigravity project ID during refresh"); - } - - Ok(tokens) - } - - async fn discover_project_id(&self, tokens: &mut OAuthTokens) -> Result<(), OAuthError> { - let url = format!("{}/v1internal:loadCodeAssist", self.endpoints[0]); - - let resp: reqwest::Response = self - .http_client - .post(&url) - .header("Authorization", format!("Bearer {}", tokens.access_token)) - .header( - "x-goog-api-client", - "google-cloud-sdk vscode_cloudshelleditor/0.1", - ) - .header( - "Client-Metadata", - r#"{"ideType":"IDE_UNSPECIFIED","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}"#, - ) - .json(&serde_json::json!({ - "metadata": { - "ideType": "IDE_UNSPECIFIED", - "platform": "PLATFORM_UNSPECIFIED", - "pluginType": "GEMINI" - } - })) - .send() - .await?; - - if !resp.status().is_success() { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - warn!( - %status, - "loadCodeAssist failed, attempting fallback to default project ID" - ); - debug!(response_body = %body, "loadCodeAssist error response"); - if let Some(ref default_id) = self.default_project_id { - tokens.extra.insert( - "project_id".to_string(), - serde_json::json!(default_id), - ); - info!(project_id = %default_id, "Using default Antigravity project ID from env"); - return Ok(()); - } - return Err(OAuthError::LoginFailed(format!( - "loadCodeAssist failed ({status}) and no ANTIGRAVITY_DEFAULT_PROJECT_ID set" - ))); - } - - let json: serde_json::Value = resp.json().await?; - - // Extract project ID — openclaw uses cloudaicompanionProject (string or {id: ...}) - let project_id = json - .get("cloudaicompanionProject") - .and_then(|v| { - v.as_str().map(String::from).or_else(|| { - v.get("id").and_then(|id| id.as_str()).map(String::from) - }) - }) - .or_else(|| { - json.get("projectId") - .and_then(|v| v.as_str()) - .map(String::from) - }); - - if let Some(pid) = project_id { - debug!(project_id = %pid, "Discovered Antigravity project ID"); - tokens - .extra - .insert("project_id".to_string(), serde_json::json!(pid)); - } else if let Some(ref default_id) = self.default_project_id { - tokens.extra.insert( - "project_id".to_string(), - serde_json::json!(default_id), - ); - info!(project_id = %default_id, "API returned no project ID, using default from env"); - } - if let Some(tier) = json.get("tier").and_then(|v| v.as_str()) { - tokens - .extra - .insert("tier".to_string(), serde_json::json!(tier)); - } - - Ok(()) - } -} - -fn parse_google_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) -} - -/// Wrap an OpenAI-format request body into Antigravity/Gemini-style format. -/// -/// Translates OpenAI chat/completions fields to Gemini generateContent fields: -/// - `messages` → `contents` + `systemInstruction` -/// - `max_tokens`/`max_completion_tokens` → `generationConfig.maxOutputTokens` -/// - `temperature`, `top_p`, `stop` → `generationConfig` -/// - `tools` (OpenAI function-calling) → Gemini `tools[].functionDeclarations` -/// - Strips OpenAI-only fields (`stream`, `stream_options`, `store`, etc.) -pub fn wrap_antigravity_request( - body: &[u8], - project_id: &str, -) -> Result, OAuthError> { - let original: serde_json::Value = serde_json::from_slice(body).map_err(|e| { - OAuthError::LoginFailed(format!("failed to parse request body as JSON: {e}")) - })?; - - let obj = original.as_object().ok_or_else(|| { - OAuthError::LoginFailed("request body is not a JSON object".to_string()) - })?; - - let raw_model = obj - .get("model") - .and_then(|v| v.as_str()) - .unwrap_or("gemini-2.0-flash"); - // Strip provider prefix (e.g. "google/gemini-2.5-flash" or - // "custom-provider/google/gemini-2.5-flash" → "gemini-2.5-flash") - let model = raw_model.rsplit('/').next().unwrap_or(raw_model); - - let mut request = serde_json::Map::new(); - - // Translate messages → contents + systemInstruction - if let Some(messages) = obj.get("messages").and_then(|v| v.as_array()) { - let mut system_parts: Vec = Vec::new(); - let mut contents: Vec = Vec::new(); - - for msg in messages { - let role = msg.get("role").and_then(|v| v.as_str()).unwrap_or(""); - let parts = message_content_to_gemini_parts(msg); - - match role { - // OpenAI/Codex "developer" has system-like semantics. - "system" | "developer" => { - system_parts.extend(parts); - } - "assistant" => { - contents.push(serde_json::json!({ - "role": "model", - "parts": parts, - })); - } - "user" => { - contents.push(serde_json::json!({ - "role": "user", - "parts": parts, - })); - } - _ => { - // Gemini contents only support "user"/"model" roles. - // Coerce tool/unknown roles to "user" to avoid INVALID_ARGUMENT. - contents.push(serde_json::json!({ - "role": "user", - "parts": parts, - })); - } - } - } - - if !system_parts.is_empty() { - request.insert( - "systemInstruction".to_string(), - serde_json::json!({ "parts": system_parts }), - ); - } - request.insert("contents".to_string(), serde_json::json!(contents)); - } - - // Translate generation parameters → generationConfig - let mut gen_config = serde_json::Map::new(); - if let Some(v) = obj.get("max_tokens").or(obj.get("max_completion_tokens")) { - gen_config.insert("maxOutputTokens".to_string(), v.clone()); - } - if let Some(v) = obj.get("temperature") { - gen_config.insert("temperature".to_string(), v.clone()); - } - if let Some(v) = obj.get("top_p") { - gen_config.insert("topP".to_string(), v.clone()); - } - if let Some(v) = obj.get("stop") { - // OpenAI: stop can be string or array; Gemini: stopSequences is always array - let sequences = if v.is_string() { - serde_json::json!([v]) - } else { - v.clone() - }; - gen_config.insert("stopSequences".to_string(), sequences); - } - if !gen_config.is_empty() { - request.insert( - "generationConfig".to_string(), - serde_json::Value::Object(gen_config), - ); - } - - // Translate tools (OpenAI function-calling → Gemini functionDeclarations) - if let Some(tools) = obj.get("tools").and_then(|v| v.as_array()) { - let mut func_decls: Vec = Vec::new(); - for tool in tools { - if tool.get("type").and_then(|v| v.as_str()) == Some("function") { - if let Some(func) = tool.get("function") { - let mut decl = serde_json::Map::new(); - if let Some(name) = func.get("name") { - decl.insert("name".to_string(), name.clone()); - } - if let Some(desc) = func.get("description") { - decl.insert("description".to_string(), desc.clone()); - } - if let Some(params) = func.get("parameters") { - decl.insert( - "parameters".to_string(), - sanitize_schema_for_gemini(params.clone()), - ); - } - func_decls.push(serde_json::Value::Object(decl)); - } - } - } - if !func_decls.is_empty() { - request.insert( - "tools".to_string(), - serde_json::json!([{ "functionDeclarations": func_decls }]), - ); - } - } - - debug!( - raw_model = raw_model, - model = model, - project = project_id, - "Antigravity request: wrapping for upstream" - ); - - let wrapped = serde_json::json!({ - "project": project_id, - "model": model, - "user_prompt_id": uuid::Uuid::new_v4().to_string(), - "request": request, - }); - - serde_json::to_vec(&wrapped) - .map_err(|e| OAuthError::LoginFailed(format!("failed to serialize wrapped body: {e}"))) -} - -/// Convert an OpenAI message's `content` field to Gemini `parts` array. -fn message_content_to_gemini_parts(msg: &serde_json::Value) -> Vec { - let Some(content) = msg.get("content") else { - return vec![]; - }; - - // String content → single text part - if let Some(text) = content.as_str() { - return vec![serde_json::json!({ "text": text })]; - } - - // Array content (multimodal) → convert each part - if let Some(parts) = content.as_array() { - return parts - .iter() - .filter_map(|part| { - match part.get("type").and_then(|v| v.as_str()) { - Some("text") => { - let text = part.get("text").and_then(|v| v.as_str()).unwrap_or(""); - Some(serde_json::json!({ "text": text })) - } - Some("image_url") => { - // Convert OpenAI image_url to Gemini inlineData or fileData - let url = part - .get("image_url") - .and_then(|v| v.get("url")) - .and_then(|v| v.as_str()) - .unwrap_or(""); - if let Some(rest) = url.strip_prefix("data:") { - // data URI → inlineData - if let Some((mime, data)) = rest.split_once(";base64,") { - return Some(serde_json::json!({ - "inlineData": { - "mimeType": mime, - "data": data, - } - })); - } - } - // URL → fileData - Some(serde_json::json!({ - "fileData": { - "fileUri": url, - } - })) - } - _ => None, - } - }) - .collect(); - } - - vec![] -} - -/// JSON Schema keywords that the Cloud Code Assist API rejects. -/// Matching openclaw's `GEMINI_UNSUPPORTED_SCHEMA_KEYWORDS`. -const GEMINI_SCHEMA_REJECTED: &[&str] = &[ - "patternProperties", - "additionalProperties", - "$schema", - "$id", - "$ref", - "$defs", - "definitions", - "examples", - "minLength", - "maxLength", - "minimum", - "maximum", - "multipleOf", - "pattern", - "format", - "minItems", - "maxItems", - "uniqueItems", - "minProperties", - "maxProperties", -]; - -/// Recursively strip JSON Schema fields that the Gemini API does not support. -fn sanitize_schema_for_gemini(mut value: serde_json::Value) -> serde_json::Value { - let Some(obj) = value.as_object_mut() else { - return value; - }; - - obj.retain(|key, _| !GEMINI_SCHEMA_REJECTED.contains(&key.as_str())); - - // Recurse into nested schemas - if let Some(props) = obj.get_mut("properties") { - if let Some(map) = props.as_object_mut() { - for v in map.values_mut() { - *v = sanitize_schema_for_gemini(v.clone()); - } - } - } - if let Some(items) = obj.get_mut("items") { - *items = sanitize_schema_for_gemini(items.clone()); - } - - value -} - -#[async_trait] -impl OAuthProvider for AntigravityProvider { - fn id(&self) -> &str { - "antigravity" - } - - fn display_name(&self) -> &str { - "Antigravity / Google (OAuth)" - } - - fn supports_headless_url(&self) -> bool { - true - } - - async fn login_browser(&self, _callback_port: u16) -> Result { - // This client ID requires a fixed redirect URI registered in Google's OAuth console. - const ANTIGRAVITY_CALLBACK_PORT: u16 = 51121; - let (verifier, challenge) = generate_pkce(); - let redirect_uri = format!("http://localhost:{ANTIGRAVITY_CALLBACK_PORT}/oauth-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={}&access_type=offline&prompt=consent", - 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 Antigravity/Google OAuth login"); - if let Err(e) = open::that(&auth_url) { - return Err(OAuthError::LoginFailed(format!( - "failed to open browser: {e}. Visit this URL manually:\n{auth_url}" - ))); - } - - let (code, received_state) = - wait_for_oauth_callback(ANTIGRAVITY_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 { - const ANTIGRAVITY_CALLBACK_PORT: u16 = 51121; - let (verifier, challenge) = generate_pkce(); - let redirect_uri = format!("http://localhost:{ANTIGRAVITY_CALLBACK_PORT}/oauth-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={}&access_type=offline&prompt=consent", - self.auth_url, - urlencoding::encode(&self.client_id), - urlencoding::encode(&redirect_uri), - urlencoding::encode(&self.scopes.join(" ")), - urlencoding::encode(&challenge), - urlencoding::encode(&state), - ); - - println!(); - println!(" Visit this URL to authenticate:"); - println!(" {auth_url}"); - println!(); - - // Start a local HTTP server to receive the OAuth callback, then wait. - let (code, received_state) = - wait_for_oauth_callback(ANTIGRAVITY_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 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()?, - ); - 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(()) - } - - async fn enrich_tokens(&self, tokens: &OAuthTokens) -> Result, OAuthError> { - if tokens.extra.contains_key("project_id") { - return Ok(None); - } - info!("Antigravity tokens missing project_id, discovering on-demand"); - let mut enriched = tokens.clone(); - self.discover_project_id(&mut enriched).await?; - Ok(Some(enriched)) - } - - 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 for Antigravity provider".to_string(), - ))?; - let wrapped = wrap_antigravity_request(body, project_id)?; - Ok(Some(wrapped)) - } - - fn upstream_url(&self, _tokens: &OAuthTokens) -> Option { - Some(self.endpoints[0].clone()) - } - - fn rewrite_request_path(&self, _path: &str) -> Option { - Some("/v1internal:streamGenerateContent?alt=sse".to_string()) - } - - fn response_format(&self, _original_path: &str) -> Option { - Some(super::ResponseFormat::GeminiSse) - } -} - -/// Parse code and state from a pasted redirect URL (headless flow). -fn parse_redirect_url(input: &str) -> Result<(String, String), OAuthError> { - let query = input - .split('?') - .nth(1) - .ok_or_else(|| OAuthError::LoginFailed("no query string in redirect URL".to_string()))?; - - 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(), - _ => {} - } - } - } - - if code.is_empty() { - return Err(OAuthError::LoginFailed( - "no authorization code found in redirect URL".to_string(), - )); - } - - Ok((code, state)) -} - -/// Wait for an OAuth callback on a local HTTP server (same as codex). -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]); - - 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)) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn test_provider() -> AntigravityProvider { - AntigravityProvider::new_with_secret( - Some("test-client-id"), - Some("test-client-secret"), - None, - None, - None, - ) - } - - #[test] - fn test_parse_google_token_response() { - let json = serde_json::json!({ - "access_token": "ya29.test", - "refresh_token": "1//test", - "expires_in": 3600, - "scope": "openid email profile", - "token_type": "Bearer" - }); - - let tokens = parse_google_token_response(&json).unwrap(); - assert_eq!(tokens.access_token, "ya29.test"); - assert_eq!(tokens.refresh_token.as_deref(), Some("1//test")); - assert!(tokens.expires_at.is_some()); - } - - #[test] - fn test_parse_google_token_response_no_refresh() { - let json = serde_json::json!({ - "access_token": "ya29.refreshed", - "expires_in": 3600 - }); - - let tokens = parse_google_token_response(&json).unwrap(); - assert_eq!(tokens.access_token, "ya29.refreshed"); - assert!(tokens.refresh_token.is_none()); - } - - #[test] - fn test_wrap_antigravity_request() { - let body = serde_json::json!({ - "model": "gemini-3-pro", - "messages": [{"role": "user", "content": "hello"}] - }); - let body_bytes = serde_json::to_vec(&body).unwrap(); - - let wrapped = wrap_antigravity_request(&body_bytes, "proj-abc-123").unwrap(); - let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); - - assert_eq!(parsed["project"], "proj-abc-123"); - assert_eq!(parsed["model"], "gemini-3-pro"); - - // messages should be translated to Gemini contents format - let contents = parsed["request"]["contents"].as_array().unwrap(); - assert_eq!(contents.len(), 1); - assert_eq!(contents[0]["role"], "user"); - assert_eq!(contents[0]["parts"][0]["text"], "hello"); - - // Raw OpenAI fields should not be present - assert!(parsed["request"].get("messages").is_none()); - } - - #[test] - fn test_wrap_antigravity_request_default_model() { - let body = serde_json::json!({ - "messages": [{"role": "user", "content": "test"}] - }); - let body_bytes = serde_json::to_vec(&body).unwrap(); - - let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); - let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); - - assert_eq!(parsed["model"], "gemini-2.0-flash"); - } - - #[test] - fn test_wrap_antigravity_request_strips_nested_provider_prefixes() { - let body = serde_json::json!({ - "model": "custom-clawshell-18790/google/gemini-2.0-flash", - "messages": [{"role": "user", "content": "hello"}] - }); - let body_bytes = serde_json::to_vec(&body).unwrap(); - - let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); - let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); - - assert_eq!(parsed["model"], "gemini-2.0-flash"); - } - - #[test] - fn test_wrap_antigravity_system_message() { - let body = serde_json::json!({ - "model": "gemini-2.0-flash", - "messages": [ - {"role": "system", "content": "You are helpful."}, - {"role": "user", "content": "hi"} - ] - }); - let body_bytes = serde_json::to_vec(&body).unwrap(); - - let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); - let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); - - // System messages become systemInstruction - assert_eq!( - parsed["request"]["systemInstruction"]["parts"][0]["text"], - "You are helpful." - ); - // Only non-system messages in contents - let contents = parsed["request"]["contents"].as_array().unwrap(); - assert_eq!(contents.len(), 1); - assert_eq!(contents[0]["role"], "user"); - } - - #[test] - fn test_wrap_antigravity_developer_message_maps_to_system_instruction() { - let body = serde_json::json!({ - "model": "gemini-2.0-flash", - "messages": [ - {"role": "developer", "content": "Follow this policy."}, - {"role": "user", "content": "hello"} - ] - }); - let body_bytes = serde_json::to_vec(&body).unwrap(); - - let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); - let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); - - assert_eq!( - parsed["request"]["systemInstruction"]["parts"][0]["text"], - "Follow this policy." - ); - let contents = parsed["request"]["contents"].as_array().unwrap(); - assert_eq!(contents.len(), 1); - assert_eq!(contents[0]["role"], "user"); - assert_eq!(contents[0]["parts"][0]["text"], "hello"); - } - - #[test] - fn test_wrap_antigravity_assistant_role() { - let body = serde_json::json!({ - "model": "gemini-2.0-flash", - "messages": [ - {"role": "user", "content": "hi"}, - {"role": "assistant", "content": "hello"} - ] - }); - let body_bytes = serde_json::to_vec(&body).unwrap(); - - let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); - let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); - - let contents = parsed["request"]["contents"].as_array().unwrap(); - assert_eq!(contents[1]["role"], "model"); - assert_eq!(contents[1]["parts"][0]["text"], "hello"); - } - - #[test] - fn test_wrap_antigravity_generation_config() { - let body = serde_json::json!({ - "model": "gemini-2.0-flash", - "messages": [{"role": "user", "content": "hi"}], - "max_completion_tokens": 1024, - "temperature": 0.7, - "top_p": 0.9, - "stop": ["\n"] - }); - let body_bytes = serde_json::to_vec(&body).unwrap(); - - let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); - let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); - - let gc = &parsed["request"]["generationConfig"]; - assert_eq!(gc["maxOutputTokens"], 1024); - assert_eq!(gc["temperature"], 0.7); - assert_eq!(gc["topP"], 0.9); - assert_eq!(gc["stopSequences"], serde_json::json!(["\n"])); - } - - #[test] - fn test_wrap_antigravity_tools() { - let body = serde_json::json!({ - "model": "gemini-2.0-flash", - "messages": [{"role": "user", "content": "hi"}], - "tools": [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get weather", - "parameters": {"type": "object", "properties": {}} - } - } - ] - }); - let body_bytes = serde_json::to_vec(&body).unwrap(); - - let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); - let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); - - let decls = &parsed["request"]["tools"][0]["functionDeclarations"]; - assert_eq!(decls[0]["name"], "get_weather"); - assert_eq!(decls[0]["description"], "Get weather"); - } - - #[test] - fn test_wrap_antigravity_strips_openai_fields() { - let body = serde_json::json!({ - "model": "gemini-2.0-flash", - "messages": [{"role": "user", "content": "hi"}], - "stream": true, - "stream_options": {"include_usage": true}, - "store": false - }); - let body_bytes = serde_json::to_vec(&body).unwrap(); - - let wrapped = wrap_antigravity_request(&body_bytes, "proj-123").unwrap(); - let parsed: serde_json::Value = serde_json::from_slice(&wrapped).unwrap(); - - assert!(parsed["request"].get("stream").is_none()); - assert!(parsed["request"].get("stream_options").is_none()); - assert!(parsed["request"].get("store").is_none()); - assert!(parsed["request"].get("messages").is_none()); - } - - #[test] - fn test_sanitize_schema_strips_unsupported_fields() { - let schema = serde_json::json!({ - "type": "object", - "properties": { - "name": { - "type": "string", - "description": "The name", - "patternProperties": {"^x-": {"type": "string"}}, - "additionalProperties": false, - "minLength": 1, - "maxLength": 100 - }, - "tags": { - "type": "array", - "items": { - "type": "string", - "$ref": "#/defs/Tag", - "default": "foo" - } - } - }, - "required": ["name"], - "$schema": "http://json-schema.org/draft-07/schema#", - "additionalProperties": false, - "$defs": {} - }); - - let sanitized = sanitize_schema_for_gemini(schema); - - // Top-level rejected fields stripped - assert!(sanitized.get("$schema").is_none()); - assert!(sanitized.get("additionalProperties").is_none()); - assert!(sanitized.get("$defs").is_none()); - // Supported fields preserved - assert_eq!(sanitized["type"], "object"); - assert_eq!(sanitized["required"], serde_json::json!(["name"])); - - // Nested property rejected fields stripped - let name_prop = &sanitized["properties"]["name"]; - assert!(name_prop.get("patternProperties").is_none()); - assert!(name_prop.get("additionalProperties").is_none()); - assert!(name_prop.get("minLength").is_none()); - assert!(name_prop.get("maxLength").is_none()); - assert_eq!(name_prop["type"], "string"); - assert_eq!(name_prop["description"], "The name"); - - // Items rejected fields stripped, but non-rejected fields preserved - let items = &sanitized["properties"]["tags"]["items"]; - assert!(items.get("$ref").is_none()); - assert_eq!(items["type"], "string"); - // "default" is NOT in the rejected list, so it's preserved - assert_eq!(items["default"], "foo"); - } - - #[test] - fn test_antigravity_provider_defaults() { - let provider = test_provider(); - assert_eq!(provider.id(), "antigravity"); - assert_eq!(provider.display_name(), "Antigravity / Google (OAuth)"); - assert!(provider.supports_headless_url()); - assert!(!provider.supports_device_code()); - } - - #[test] - fn test_inject_auth_headers() { - let provider = test_provider(); - let mut headers = HeaderMap::new(); - provider.inject_auth(&mut headers, "ya29.test").unwrap(); - - assert_eq!( - headers.get("authorization").unwrap().to_str().unwrap(), - "Bearer ya29.test" - ); - assert!(headers.get("x-goog-api-client").is_some()); - assert!(headers.get("client-metadata").is_some()); - } - - #[test] - fn test_prepare_request_body_with_project_id() { - let provider = test_provider(); - let mut tokens = OAuthTokens { - access_token: "t".to_string(), - refresh_token: None, - id_token: None, - expires_at: None, - account_id: None, - extra: BTreeMap::new(), - }; - tokens.extra.insert( - "project_id".to_string(), - serde_json::json!("proj-test"), - ); - - let body = serde_json::json!({"model": "gemini-3-pro", "messages": []}); - let result = provider - .prepare_request_body(&serde_json::to_vec(&body).unwrap(), &tokens) - .unwrap(); - assert!(result.is_some()); - - let parsed: serde_json::Value = serde_json::from_slice(&result.unwrap()).unwrap(); - assert_eq!(parsed["project"], "proj-test"); - } - - #[test] - fn test_prepare_request_body_missing_project_id() { - let provider = test_provider(); - 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": "gemini-3-pro"}); - let result = provider.prepare_request_body(&serde_json::to_vec(&body).unwrap(), &tokens); - assert!(result.is_err()); - } - - #[test] - fn test_upstream_url() { - let provider = test_provider(); - let tokens = OAuthTokens { - access_token: "t".to_string(), - refresh_token: None, - id_token: None, - expires_at: None, - account_id: None, - extra: BTreeMap::new(), - }; - - let url = provider.upstream_url(&tokens).unwrap(); - assert!(url.contains("cloudcode-pa.googleapis.com")); - assert!(!url.contains("streamGenerateContent"), "upstream_url should be base only"); - } - - #[test] - fn test_rewrite_request_path() { - let provider = test_provider(); - let rewritten = provider.rewrite_request_path("/v1/chat/completions"); - assert_eq!( - rewritten.as_deref(), - Some("/v1internal:streamGenerateContent?alt=sse") - ); - } -} diff --git a/src/oauth/mod.rs b/src/oauth/mod.rs index 4391d78..58a4cd7 100644 --- a/src/oauth/mod.rs +++ b/src/oauth/mod.rs @@ -1,7 +1,6 @@ mod storage; pub mod codex; -pub mod antigravity; pub use storage::TokenStorage; @@ -86,10 +85,10 @@ impl OAuthTokens { /// The core trait that each OAuth provider implements. #[async_trait] pub trait OAuthProvider: Send + Sync + fmt::Debug { - /// Unique identifier (e.g., "codex", "antigravity"). + /// Unique identifier (e.g., "codex"). fn id(&self) -> &str; - /// Display name (e.g., "Codex (OpenAI)", "Antigravity (Google)"). + /// Display name (e.g., "Codex (OpenAI)"). fn display_name(&self) -> &str; /// Execute browser-based OAuth login flow. @@ -105,7 +104,7 @@ pub trait OAuthProvider: Send + Sync + fmt::Debug { 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. + /// Returns None for pass-through (Codex). fn prepare_request_body( &self, _body: &[u8], diff --git a/src/onboard/interactive.rs b/src/onboard/interactive.rs index abafc77..42823db 100644 --- a/src/onboard/interactive.rs +++ b/src/onboard/interactive.rs @@ -322,13 +322,11 @@ fn mask_secret(secret: &str) -> String { /// Run the OAuth login flow for the given provider, persisting tokens. fn run_oauth_login(provider_id: &str) -> Result<(), Box> { - use crate::oauth::antigravity::AntigravityProvider; use crate::oauth::codex::CodexProvider; use crate::oauth::{OAuthProvider, TokenStorage}; let provider: Box = match provider_id { "codex" => Box::new(CodexProvider::new(None, None, None, None)), - "antigravity" => Box::new(AntigravityProvider::new(None, None, None, None)), other => return Err(format!("unknown OAuth provider: {other}").into()), }; @@ -426,32 +424,12 @@ pub fn collect_onboard_config_tui() -> Result Result Some(MENU_ANTIGRAVITY), (Some("oauth"), Some("codex"), _) | (Some("oauth"), _, _) => Some(MENU_CODEX), (_, _, Some("anthropic")) => Some(MENU_ANTHROPIC), (_, _, Some("openrouter")) => Some(MENU_OPENROUTER), @@ -471,7 +448,7 @@ pub fn collect_onboard_config_tui() -> Result Result ( - "openai".to_string(), - OnboardAuthMethod::OAuth { - provider_id: "antigravity".to_string(), - }, - ), _ => ("openai".to_string(), OnboardAuthMethod::StaticKey), }; @@ -499,7 +470,6 @@ pub fn collect_onboard_config_tui() -> Result "claude-sonnet-4-5-20250929", MENU_OPENROUTER => "openrouter/auto", MENU_CODEX => "gpt-5.2-chat-latest", - MENU_ANTIGRAVITY => "gemini-2.0-flash", _ => "gpt-5.2-chat-latest", // OpenAI default }); let model = tui::prompt_text("Enter the model name", Some(default_model))?; diff --git a/src/onboard/types.rs b/src/onboard/types.rs index ed54460..6c3933d 100644 --- a/src/onboard/types.rs +++ b/src/onboard/types.rs @@ -49,7 +49,7 @@ pub enum OnboardAuthMethod { StaticKey, /// OAuth provider supplies access tokens at runtime. OAuth { - /// Provider identifier, e.g. "codex", "antigravity". + /// Provider identifier, e.g. "codex". provider_id: String, }, }