diff --git a/README.md b/README.md index 7dbea331..ecd46a0d 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,53 @@ passed again. Smart routing uses the `task_v1` router by default. Power users ca router for a launch by setting `SMART_ROUTER_NAME`, for example `SMART_ROUTER_NAME=task_v2 ug codex --enable-smart-routing`. +### Codex shared rate limiting + +Codex requests launched by `ug codex` share a local rolling 60-second input-token budget. This +prevents several concurrent Codex sessions from independently exhausting the same Databricks +Foundation Model API quota. Budgets are isolated by workspace and model, so traffic for one model +does not consume another model's allowance. + +| Codex model | Published input tokens/minute | Local target | +|---|---:|---:| +| GPT-6 Astra | 200,000 | 180,000 | +| GPT-5.6 Sol | 2,000,000 | 1,800,000 | +| GPT-5.6 Terra | 2,000,000 | 1,800,000 | +| GPT-5.6 Luna | 2,000,000 | 1,800,000 | +| Kimi K3 | 200,000 | 180,000 | +| Qwen3.5 122B A10B | 1,000,000 | 900,000 | +| Qwen3-Next 80B A3B Instruct | 1,000,000 | 900,000 | + +The 90% target leaves headroom for estimation error and requests made outside Unity Gateway. +Unity Gateway estimates input conservatively from the uncompressed JSON request size and records +only timestamps, model keys, and estimates in `~/.ucode/codex-rate-limit-state.json`; prompts and +credentials are never stored. When capacity is unavailable, the request waits and one short notice +is written to stderr. Unknown/new models pass through the proactive token-budget check until their +published quota is added. + +Every upstream `429 Too Many Requests` response is also handled inside Unity Gateway instead of +consuming Codex's finite retry budget. Unity Gateway honors `Retry-After` when present; otherwise it +uses capped exponential backoff with jitter. The resulting cooldown is shared across local Codex +processes for the same workspace and model. If the request body is unreadable or the model is new, +the cooldown safely applies to the whole workspace. The original model and reasoning effort never +change, and the thread keeps retrying until capacity returns or the user cancels it. + +The proxy address is different for each launch, so Codex's provider must remain in the ucode profile +and launch-time overrides. During configuration, current versions remove provider values that older +ucode versions wrote to `/etc/codex/managed_config.toml`; Codex gives that file higher precedence +and would otherwise silently bypass the limiter. + +When a Codex thread switches from a model that returns visible reasoning, such as Kimi K3, to an +OpenAI reasoning model, Unity Gateway removes the nonportable `reasoning.content` field from the +outbound replay request. The saved Codex transcript is not changed. OpenAI's replayable +`encrypted_content` and all messages and tool calls remain intact. + +The limiter applies to normal, app, and smart-routed Codex launches. Caller-managed server commands +such as `codex app-server` and `codex mcp-server` keep their existing launch path. To bypass the +limiter for one launch, use `UCODE_CODEX_RATE_LIMITER=0 ug codex`; that restores direct requests and +may expose the session to 429 responses. See the current +[Databricks Foundation Model API limits](https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/limits). + To configure all tools at once: ```bash @@ -263,11 +310,12 @@ ug publish # publish it to the workspace `ug setup` walks through the agents to enable and which one bare `ug` launches, then per agent: Databricks-hosted models or an external Model Provider Service and the models to expose. Interactive -Claude Code and Codex configuration installs gateway-critical values in the OS-managed settings -scope so enterprise settings cannot silently override Unity Gateway. Non-interactive and CI runs use local -files without invoking `sudo`, and stop with an actionable error if an existing managed value -conflicts. Claude subscription relay is local-only because its loopback proxy exists only for that -session. +Claude Code configuration installs gateway-critical values in the OS-managed settings scope so +enterprise settings cannot silently override Unity Gateway. Codex keeps its provider launch-scoped +because its shared limiter uses a per-session loopback proxy. Codex configuration retires provider +values written there by older ucode versions and stops if unrelated machine policy still overrides +the launch-scoped provider. Non-interactive and CI runs never invoke `sudo`. Claude subscription +relay is also local-only because its loopback proxy exists only for that session. Claude Code is asked one model per family (opus/sonnet/haiku/fable), since it selects models by family alias; any family can be skipped. @@ -407,7 +455,7 @@ control the installation. | `~/.codex/ucode.config.toml` (or legacy `~/.codex/config.toml`) | Codex | | `~/.claude/ucode-settings.json` | Claude Code settings generated by ug | | `/etc/claude-code/managed-settings.json` (Linux) or `/Library/Application Support/ClaudeCode/managed-settings.json` (macOS) | Claude Code OS-managed settings | -| `/etc/codex/managed_config.toml` | Codex OS-managed settings | +| `/etc/codex/managed_config.toml` | Codex OS-managed settings inspected for provider conflicts; older ucode-owned values are retired | | `~/.gemini/.env` | Gemini CLI | | `~/.config/opencode/opencode.json` | OpenCode | | `~/.copilot/.env` | GitHub Copilot CLI | diff --git a/docs/os-managed-settings-design.md b/docs/os-managed-settings-design.md index 4f75e6bc..3e993243 100644 --- a/docs/os-managed-settings-design.md +++ b/docs/os-managed-settings-design.md @@ -2,19 +2,20 @@ ## Summary -Claude Code and Codex give OS-managed settings higher precedence than user settings. Previously, -ucode could write only its local configuration while an existing machine-managed file silently -overrode the gateway endpoint, authentication helper, provider headers, or model. +Claude Code and Codex give OS-managed settings higher precedence than user settings. That is useful +for Claude Code's stable gateway configuration, but a Codex provider at that scope overrides the +per-launch loopback address used for shared throttling and 429 retries. -The two stacked PRs make precedence handling deterministic: +The shared managed-file lifecycle makes precedence handling deterministic: 1. The Claude PR adds the shared managed-file lifecycle and applies it to Claude Code. -2. The Codex PR reuses that lifecycle for TOML, applies it to Codex, and removes the old optional - managed-settings path. +2. Codex reuses that lifecycle to retire provider settings older ucode versions installed and to + preserve unrelated machine policy. -After both PRs merge, interactive configuration reconciles the agent's OS-managed file by default. -Non-interactive and CI execution never elevates privileges and instead uses local settings when the -managed file is compatible. +Interactive Claude configuration reconciles its OS-managed file by default. Codex configuration +keeps its provider in `~/.codex/ucode.config.toml` plus launch-time overrides. It restores the +recorded pre-ucode managed baseline, then verifies that remaining machine policy cannot bypass that +provider. Non-interactive and CI execution never elevates privileges. ## Configuration Files @@ -23,8 +24,9 @@ managed file is compatible. | Claude Code | `~/.claude/ucode-settings.json` | Linux: `/etc/claude-code/managed-settings.json`; macOS: `/Library/Application Support/ClaudeCode/managed-settings.json` | | Codex | `~/.codex/ucode.config.toml` | `/etc/codex/managed_config.toml` | -The local file is always written. The OS-managed file is additionally reconciled during interactive -configuration, except for Claude subscription relay. +The local file is always written. Claude's OS-managed file is additionally reconciled during +interactive configuration, except for subscription relay. Codex's OS-managed file is inspected and +fingerprinted but is not populated by current ucode versions. ## Interactive Detection @@ -38,24 +40,23 @@ command shape in some flows and did not guard managed-file writes consistently. ## Behavior Matrix -| Invocation | Managed file | Behavior | +| Agent and invocation | Managed file | Behavior | | --- | --- | --- | -| Interactive | Absent | Create it from the ucode configuration after recording an absent baseline. | -| Interactive | Unrelated or partially populated | Preserve unrelated values and add or update all ucode-owned values. | -| Interactive | Conflicting | Back up the baseline, replace the conflicting ucode-owned values, and verify. | -| Interactive | Already identical | Continue without a backup, write, or `sudo` invocation. | -| Non-interactive | Absent | Use the local ucode file. Do not create the managed file. | -| Non-interactive | Ucode-owned values absent or equal | Use the local ucode file. Do not modify the managed file. | -| Non-interactive | Ucode-owned value conflicts | Stop before launching because the higher-precedence value would override ucode. | +| Claude, interactive | Absent or compatible | Record the baseline, add the stable gateway values, and verify. | +| Claude, non-interactive | Absent or compatible | Use the local ucode file without invoking `sudo`. | +| Codex, interactive | Contains tracked ucode provider values | Restore the pre-ucode baseline, preserving later external changes. | +| Codex, non-interactive | Contains tracked ucode provider values | Stop and require one interactive migration; never invoke `sudo`. | +| Codex, any | Absent or unrelated | Use the local profile and launch-scoped provider; record a compatibility fingerprint. | +| Codex, any | Selects another provider or overrides `ucode-databricks.base_url` | Stop because the managed value would bypass the launch-scoped proxy. | | Any | Invalid, unreadable, or symlinked | Stop without modifying the file because precedence cannot be established safely. | -`ucode configure`, first-time `ucode claude` or `ucode codex`, and later launches all use the same -agent-specific reconciliation path. A first-time launch from an interactive terminal can therefore -request administrator permission. A first-time non-interactive launch remains local-only. +`ucode configure`, first-time agent launches, and later launches all use the same agent-specific +compatibility path. A Codex launch requests administrator permission only when it must retire a +tracked provider installed by an older ucode version. -## Interactive Reconciliation +## Claude Interactive Reconciliation -For each agent, ucode: +For Claude Code, ucode: 1. Strictly parses the existing managed JSON or TOML document. 2. Produces the desired document by applying the same gateway overlay used for the local ucode file. @@ -66,9 +67,28 @@ For each agent, ucode: 7. Reads the installed file back and verifies its exact contents. 8. Records the last-applied snapshot, owned paths, and a launch fingerprint. -An existing managed file is reconciled even when it does not currently conflict. This ensures every -ucode-required value exists at the highest-precedence scope and avoids separate behavior for absent, -partial, and conflicting files. +An existing Claude managed file is reconciled even when it does not currently conflict. This +ensures every stable gateway value exists at the highest-precedence scope. + +## Codex Launch-Scoped Provider + +Every normal `ug codex` launch starts a loopback proxy on an ephemeral port. The launch overlay +replaces only `model_providers.ucode-databricks.base_url` with that address and disables request +compression so the proxy can identify the model and estimate input tokens. + +Codex's system `managed_config.toml` has higher precedence than both the named ucode profile and +launch overrides. Persisting the direct Databricks endpoint there therefore bypasses the proxy even +when the child command visibly receives a loopback URL. Configuration now: + +1. Uses the managed-backup manifest to restore the exact pre-ucode baseline or perform a three-way + revert when external policy changed later. +2. Preserves unrelated managed keys, including approval and model defaults. +3. Rejects a remaining external `model_provider` selection or managed + `model_providers.ucode-databricks.base_url`. +4. Fingerprints the compatible result so unchanged launches need only one `stat()` call. + +This means a bare `codex` command is not routed through Unity Gateway by current ucode versions. Use +`ug codex` when the shared limiter, gateway authentication, and automatic 429 recovery are required. ## Privileged Write Transaction @@ -215,7 +235,7 @@ administrator help. - Add Claude managed status and revert output. - Remove Claude's old managed-settings scope choice. -### PR 2: Codex +### PR 2: Codex (historical behavior) - Stack on the Claude PR and reuse the shared lifecycle with strict TOML parsing and serialization. - Make interactive Codex configuration reconcile OS-managed TOML by default. @@ -224,3 +244,10 @@ administrator help. - Keep opt-in smart-routing hooks in the local Codex config rather than adding them by default to machine-managed policy. - Remove the remaining managed-settings scope schema, resolution, setup prompt, summary, and tests. + +### Codex launch-proxy amendment + +- Retire the direct managed provider using the existing baseline and three-way restore lifecycle. +- Keep the provider in the named profile and launch overrides so every request reaches the + per-launch proxy. +- Preserve unrelated system policy and block only values that would bypass the proxy. diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 625329f3..566eaed3 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -7,14 +7,19 @@ import re import subprocess import sys +import threading import time -from collections.abc import Callable +from collections.abc import Iterator, Mapping +from contextlib import contextmanager from pathlib import Path import tomlkit from tomlkit.exceptions import ParseError +from ucode import gateway_proxy from ucode.codex_config import codex_config_args +from ucode.codex_rate_limit import SharedCodexRateLimiter +from ucode.codex_request import sanitize_reasoning_replay from ucode.config_io import ( APP_DIR, ToolSpec, @@ -32,13 +37,10 @@ from ucode.managed_files import ( OS, current_os, - managed_file_conflicts, managed_file_is_verified, managed_file_status, - managed_writes_allowed, mark_managed_file_verified, read_managed_file, - reconcile_managed_file, revert_managed_file, ) from ucode.smart_routing import v2 as smart_routing_v2 @@ -66,6 +68,8 @@ # Retained only to identify and remove state written by the legacy persisted opt-in. SMART_ROUTING_STATE_KEY = smart_routing_v2.LEGACY_STATE_KEY APP_SERVER_SMART_ROUTING_STARTING_MODEL = "gpt-5.6-luna" +CODEX_RATE_LIMITER_ENV = "UCODE_CODEX_RATE_LIMITER" +_FALSE_ENV_VALUES = frozenset({"0", "false", "no", "off"}) SPEC: ToolSpec = { "binary": "codex", @@ -363,7 +367,7 @@ def compose(base: dict) -> dict: enabled=False, ) write_toml_file(CODEX_CONFIG_PATH, doc) - _reconcile_managed_config(state, compose) + _migrate_and_check_managed_config(state) state = mark_tool_managed(state, "codex", MANAGED_KEYS) save_state(state) return state @@ -395,13 +399,18 @@ def managed_config_is_current(state: dict) -> bool: path = _managed_config_path() if path is None: return True - required_scope = "managed" if managed_writes_allowed() else None - return managed_file_is_verified(state, "codex", path, required_scope=required_scope) + return managed_file_is_verified(state, "codex", path, required_scope="local-compatible") def managed_config_status(state: dict) -> tuple[Path | None, str, str]: path = _managed_config_path() status, backup = managed_file_status(state, "codex", path, parser=_parse_managed_config) + if ( + path is not None + and status == "missing" + and managed_file_is_verified(state, "codex", path, required_scope="local-compatible") + ): + status = "compatible (local settings)" return path, status, backup @@ -414,49 +423,57 @@ def revert_managed_config() -> str: ) -def _reconcile_managed_config(state: dict, compose: Callable[[dict], dict]) -> None: - """Reconcile Codex's highest-precedence config while preserving unrelated policy.""" +def _managed_provider_conflicts(doc: dict) -> list[str]: + """Return managed values that can bypass a launch-scoped ucode provider.""" + conflicts: list[str] = [] + managed_provider = doc.get("model_provider") + if managed_provider is not None and managed_provider != CODEX_MODEL_PROVIDER_NAME: + conflicts.append("model_provider") + providers = doc.get("model_providers") + ucode_provider = ( + providers.get(CODEX_MODEL_PROVIDER_NAME) if isinstance(providers, Mapping) else None + ) + if isinstance(ucode_provider, Mapping) and "base_url" in ucode_provider: + conflicts.append(f"model_providers.{CODEX_MODEL_PROVIDER_NAME}.base_url") + return conflicts + + +def _migrate_and_check_managed_config(state: dict) -> None: + """Retire ucode's managed provider and verify that external policy will not bypass it. + + Codex gives ``/etc/codex/managed_config.toml`` precedence over profile files and launch-time + config overrides. A provider persisted there therefore bypasses the per-launch loopback proxy + used for shared throttling and 429 retries. Restore ucode's recorded pre-write baseline, then + keep only a compatibility fingerprint for unrelated machine policy. + """ path = _managed_config_path() if path is None: - print_warning_err( - "Machine-wide Codex settings aren't supported on this platform; skipped the managed " - "config." - ) return if path.is_symlink(): raise RuntimeError( f"Refusing to use Codex managed settings through symlink {path}. Replace it with a " "regular file or contact your administrator." ) + + # Releases only settings recorded as ucode-owned. The shared managed-file lifecycle restores + # the original baseline or performs a three-way revert when external policy changed later. + revert_managed_config() current_text = read_managed_file(path) try: existing = _parse_managed_config(current_text) if current_text is not None else {} except RuntimeError as exc: raise RuntimeError( - f"Cannot safely update Codex managed settings at {path}: {exc}. ucode did not modify " + f"Cannot safely inspect Codex managed settings at {path}: {exc}. ucode did not modify " "the file. Repair it or contact your administrator." ) from exc - managed_before = copy.deepcopy(existing) - desired_doc = compose(existing) - if not managed_writes_allowed(): - conflicts = managed_file_conflicts(managed_before, desired_doc, MANAGED_KEYS) - if conflicts: - raise RuntimeError( - "Codex configuration cannot be applied non-interactively because OS-managed " - f"settings at {path} override ucode values: {', '.join(conflicts)}. Run `ucode " - "configure --agent codex` from an interactive terminal or contact your " - "administrator." - ) - mark_managed_file_verified(state, "codex", path, scope="local-compatible") - return - reconcile_managed_file( - path, - tomlkit.dumps(desired_doc), - tool="codex", - display="Codex", - owned_paths=MANAGED_KEYS, - ) - mark_managed_file_verified(state, "codex", path) + conflicts = _managed_provider_conflicts(existing) + if conflicts: + raise RuntimeError( + f"Codex cannot use ucode's shared rate limiter because OS-managed settings at {path} " + f"override the launch-scoped provider: {', '.join(conflicts)}. Remove those provider " + "settings or contact your administrator." + ) + mark_managed_file_verified(state, "codex", path, scope="local-compatible") def default_model(state: dict) -> str | None: @@ -494,6 +511,94 @@ def clear_model_preferences(state: dict) -> bool: _PROFILE_REJECTED_MAX_SECONDS = 3.0 +def _codex_rate_limiter_enabled() -> bool: + return os.environ.get(CODEX_RATE_LIMITER_ENV, "").strip().lower() not in _FALSE_ENV_VALUES + + +def _server_family_subcommand(tool_args: list[str]) -> bool: + """Keep Codex server commands on their existing profile/fallback path.""" + return bool(tool_args) and tool_args[0].endswith("-server") + + +def _proxied_overlay(config: dict, proxy_base_url: str, *, provider_only: bool = False) -> dict: + """Return a launch-scoped config that routes the ucode provider via loopback.""" + providers = config.get("model_providers") + provider = providers.get(CODEX_MODEL_PROVIDER_NAME) if isinstance(providers, Mapping) else None + if not isinstance(provider, Mapping): + raise RuntimeError( + f"Cannot launch Codex through the shared limiter because {CODEX_CONFIG_PATH} " + "does not contain the ucode Databricks provider. Run `ucode configure --agents " + "codex` first." + ) + + if provider_only: + overlay = { + "model_providers": {CODEX_MODEL_PROVIDER_NAME: copy.deepcopy(provider)}, + } + else: + overlay = copy.deepcopy(config) + overlay["model_providers"][CODEX_MODEL_PROVIDER_NAME]["base_url"] = proxy_base_url + # The proxy needs the JSON body to identify the model and estimate input + # tokens. This is a launch-only override; the user's Codex config is not + # changed. Official Codex config calls this stable feature on by default. + overlay["features.enable_request_compression"] = False + return overlay + + +@contextmanager +def _codex_request_proxy(state: dict) -> Iterator[str]: + """Run a loopback Codex proxy for the lifetime of one launched process.""" + workspace = state.get("workspace") + if not isinstance(workspace, str) or not workspace: + raise RuntimeError("Codex shared rate limiting needs a configured Databricks workspace.") + + limiter = SharedCodexRateLimiter(workspace) + server, cache, client = gateway_proxy.start_proxy( + workspace, + state.get("profile"), + 0, + token_header=gateway_proxy.AUTHORIZATION_HEADER, + force_refresh_near_expiry=True, + upstream_base=f"{workspace.rstrip('/')}/ai-gateway/codex/", + request_transform=sanitize_reasoning_replay, + request_gate=limiter, + rate_limit_retry=limiter.retry_after_429, + ) + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}/v1" + finally: + cache.stop() + server.shutdown() + server.server_close() + client.close() + server_thread.join(timeout=2) + + +def _run_with_profile( + binary: str, + tool_args: list[str], + config_args: list[str] | None = None, +) -> tuple[int, bool]: + """Run Codex with the ucode profile and identify its fast parse rejection.""" + started = time.monotonic() + argv = [binary, "--profile", CODEX_PROFILE_NAME, *(config_args or []), *tool_args] + returncode = subprocess.run(argv).returncode + rejected = returncode != 0 and time.monotonic() - started < _PROFILE_REJECTED_MAX_SECONDS + return returncode, rejected + + +def _relaunch_without_profile(binary: str, tool_args: list[str]) -> None: + # Warn on stderr: app-server stdout is a JSON-RPC stream. + print_warning_err( + "ucode's `--profile` isn't accepted here (error above). Retrying " + f"without it: Codex will resolve {LEGACY_CODEX_CONFIG_PATH} and any OS-managed " + "settings instead of the ucode profile." + ) + exec_or_spawn([binary, *tool_args]) + + def launch(state: dict, tool_args: list[str]) -> None: clear_model_preferences(state) binary = SPEC["binary"] @@ -516,6 +621,20 @@ def _app_server_start_model() -> str: return codex_model_id(models[0]) return APP_SERVER_SMART_ROUTING_STARTING_MODEL + if _codex_rate_limiter_enabled(): + with _codex_request_proxy(state) as proxy_base_url: + + def render_proxied_overlay(*args, **kwargs): + return _proxied_overlay(render_overlay(*args, **kwargs), proxy_base_url) + + smart_routing_v2.launch_codex( + state, + tool_args, + binary=binary, + start_model=_app_server_start_model(), + render_overlay=render_proxied_overlay, + ) + return smart_routing_v2.launch_codex( state, tool_args, @@ -523,6 +642,7 @@ def _app_server_start_model() -> str: start_model=_app_server_start_model(), render_overlay=render_overlay, ) + return if workspace: os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) if tool_args[:1] == ["app"]: @@ -535,9 +655,25 @@ def _app_server_start_model() -> str: f"Cannot launch Codex app with the ucode profile because {CODEX_CONFIG_PATH} " "is missing or empty. Run `ucode configure --agents codex` first." ) - config_args = codex_config_args(profile_doc) - exec_or_spawn([binary, "app", *config_args, *tool_args[1:]]) - return # unreachable in production (exec replaces the process) + if not _codex_rate_limiter_enabled(): + config_args = codex_config_args(profile_doc) + exec_or_spawn([binary, "app", *config_args, *tool_args[1:]]) + return # unreachable in production (exec replaces the process) + with _codex_request_proxy(state) as proxy_base_url: + config_args = codex_config_args(_proxied_overlay(profile_doc, proxy_base_url)) + returncode = subprocess.run([binary, "app", *config_args, *tool_args[1:]]).returncode + sys.exit(returncode) + + # Codex server-family commands retain their existing launch and fallback + # behavior. Their callers own provider configuration and may keep the + # server alive after this launcher exits, so a session-local proxy would be + # the wrong lifecycle. + if _server_family_subcommand(tool_args) or not _codex_rate_limiter_enabled(): + returncode, rejected = _run_with_profile(binary, tool_args) + if rejected: + _relaunch_without_profile(binary, tool_args) + return + sys.exit(returncode) # Run codex with --profile first — the TUI and runtime subcommands # (exec/resume/mcp/...) keep ucode's Databricks routing, including any added # by future codex versions. codex rejects the global --profile on @@ -554,9 +690,12 @@ def _app_server_start_model() -> str: # inherited (no capture), so Ctrl-C reaches codex directly and the resulting # KeyboardInterrupt propagates past the retry check — quitting an interactive # session is never mistaken for a --profile rejection. - started = time.monotonic() - returncode = subprocess.run([binary, "--profile", CODEX_PROFILE_NAME, *tool_args]).returncode - if returncode != 0 and time.monotonic() - started < _PROFILE_REJECTED_MAX_SECONDS: + profile_doc = read_toml_safe(CODEX_CONFIG_PATH) + with _codex_request_proxy(state) as proxy_base_url: + launch_overlay = _proxied_overlay(profile_doc, proxy_base_url, provider_only=True) + config_args = codex_config_args(launch_overlay) + returncode, rejected = _run_with_profile(binary, tool_args, config_args) + if rejected: # Fast failure: most likely codex rejected --profile on this subcommand. # Relaunch without it, handing over the terminal. (A fast failure for # any other reason — e.g. a bad flag — just re-fails the same way here, @@ -565,12 +704,7 @@ def _app_server_start_model() -> str: # Warn on *stderr*: this path is reached by `codex app-server`, whose # stdout is a JSON-RPC stream its caller parses. Emit before handing off, # since execvp replaces this process. - print_warning_err( - "ucode's `--profile` isn't accepted here (error above). Retrying " - f"without it: Codex will resolve {LEGACY_CODEX_CONFIG_PATH} and any OS-managed " - "settings instead of the ucode profile." - ) - exec_or_spawn([binary, *tool_args]) + _relaunch_without_profile(binary, tool_args) return # unreachable in production (exec replaces the process) sys.exit(returncode) diff --git a/src/ucode/codex_config.py b/src/ucode/codex_config.py index 25070cce..c555048b 100644 --- a/src/ucode/codex_config.py +++ b/src/ucode/codex_config.py @@ -2,20 +2,32 @@ from __future__ import annotations +from collections.abc import Mapping + import tomlkit +def _inline_compatible(value: object) -> object: + """Convert parsed TOML containers into inline-table-compatible values.""" + if isinstance(value, Mapping): + return {key: _inline_compatible(entry) for key, entry in value.items()} + if isinstance(value, list): + return [_inline_compatible(entry) for entry in value] + return value + + def _toml_value(value: str | int | float | bool | list[object] | dict[str, object]) -> str: - if isinstance(value, dict): + normalized = _inline_compatible(value) + if isinstance(normalized, dict): item = tomlkit.inline_table() - item.update(value) + item.update(normalized) return item.as_string() - if isinstance(value, list) and any(isinstance(entry, dict) for entry in value): + if isinstance(normalized, list) and any(isinstance(entry, dict) for entry in normalized): wrapper = tomlkit.inline_table() - wrapper["value"] = value + wrapper["value"] = normalized rendered = wrapper.as_string() return rendered.removeprefix("{value = ").removesuffix("}") - return tomlkit.item(value).as_string() + return tomlkit.item(normalized).as_string() def codex_config_args(config: dict) -> list[str]: diff --git a/src/ucode/codex_rate_limit.py b/src/ucode/codex_rate_limit.py new file mode 100644 index 00000000..9a0d08f2 --- /dev/null +++ b/src/ucode/codex_rate_limit.py @@ -0,0 +1,460 @@ +"""Cross-process input-token limiter for Codex requests sent through FMAPI. + +Codex resends its active context on every model round-trip. Several concurrent +sessions can therefore exhaust a model's input-tokens-per-minute allowance even +when each session is new. This module reserves a conservative body-size estimate +in a rolling, per-workspace/model window shared through ``~/.ucode``. + +Request bodies are inspected in memory only. The shared state contains model +keys, timestamps, and estimated token counts; prompts and credentials are never +persisted or logged. +""" + +from __future__ import annotations + +import json +import math +import os +import random +import re +import sys +import tempfile +import threading +import time +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from email.utils import parsedate_to_datetime +from pathlib import Path +from typing import BinaryIO + +from ucode import config_io + +WINDOW_SECONDS = 60.0 +BYTES_PER_ESTIMATED_TOKEN = 3 +SAFETY_PERCENT = 90 +STATE_FILE_NAME = "codex-rate-limit-state.json" +LOCK_FILE_NAME = "codex-rate-limit.lock" +RETRY_BACKOFF_BASE_SECONDS = 2.0 +RETRY_BACKOFF_MAX_SECONDS = 300.0 +RETRY_AFTER_JITTER_PERCENT = 10 + +# Databricks Foundation Model API input-token-per-minute limits. Reservations +# use only 90% so estimation error and requests outside this ucode process have +# some headroom. Model spellings are normalized before lookup, so dotted and +# dashed UC/Codex variants share the same bucket. +PUBLISHED_INPUT_TOKENS_PER_MINUTE = { + "gpt6astra": 200_000, + "gpt56sol": 2_000_000, + "gpt56terra": 2_000_000, + "gpt56luna": 2_000_000, + "kimik3": 200_000, + "qwen35122ba10b": 1_000_000, + "qwen3next80ba3binstruct": 1_000_000, +} +DEFAULT_TARGET_LIMITS = { + model: published * SAFETY_PERCENT // 100 + for model, published in PUBLISHED_INPUT_TOKENS_PER_MINUTE.items() +} + +_PROCESS_LOCK = threading.Lock() + + +def _canonical_model_key(model: str) -> str | None: + """Return the known quota key for a Codex/UC model spelling.""" + normalized = re.sub(r"[^a-z0-9]+", "", model.strip().lower()) + for key in DEFAULT_TARGET_LIMITS: + if normalized.endswith(key): + return key + return None + + +def estimate_request(body: bytes) -> tuple[str, str, int] | None: + """Return ``(request model, quota key, estimated input tokens)``. + + Unknown models and non-JSON bodies deliberately pass through. Codex request + compression is disabled in the launch-scoped proxy config, so a normal + Responses API request reaches this function as JSON. + """ + try: + payload = json.loads(body) + except (json.JSONDecodeError, UnicodeDecodeError): + return None + if not isinstance(payload, dict): + return None + model = payload.get("model") + if not isinstance(model, str) or not model.strip(): + return None + model_key = _canonical_model_key(model) + if model_key is None: + return None + estimated_tokens = max( + 1, + (len(body) + BYTES_PER_ESTIMATED_TOKEN - 1) // BYTES_PER_ESTIMATED_TOKEN, + ) + return model, model_key, estimated_tokens + + +def _acquire_file_lock(handle: BinaryIO) -> None: + if os.name == "nt": + import msvcrt + + handle.seek(0, os.SEEK_END) + if handle.tell() == 0: + handle.write(b"\0") + handle.flush() + handle.seek(0) + msvcrt.locking(handle.fileno(), msvcrt.LK_LOCK, 1) + return + + import fcntl + + fcntl.flock(handle.fileno(), fcntl.LOCK_EX) + + +def _release_file_lock(handle: BinaryIO) -> None: + if os.name == "nt": + import msvcrt + + handle.seek(0) + msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) + return + + import fcntl + + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + + +@contextmanager +def _locked(lock_path: Path) -> Iterator[None]: + lock_path.parent.mkdir(parents=True, exist_ok=True) + # flock/locking provides process coordination. The Python lock also makes + # the behavior explicit and portable between threads in one launcher. + with _PROCESS_LOCK, lock_path.open("a+b") as handle: + _acquire_file_lock(handle) + try: + yield + finally: + _release_file_lock(handle) + + +def _empty_state() -> dict: + return {"version": 2, "buckets": {}, "cooldowns": {}} + + +def _read_state(path: Path) -> dict: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (FileNotFoundError, json.JSONDecodeError, UnicodeDecodeError): + return _empty_state() + if not isinstance(payload, dict) or not isinstance(payload.get("buckets"), dict): + return _empty_state() + return payload + + +def _valid_recent_events(raw: object, now: float) -> list[dict[str, float | int]]: + if not isinstance(raw, list): + return [] + cutoff = now - WINDOW_SECONDS + events: list[dict[str, float | int]] = [] + for event in raw: + if not isinstance(event, dict): + continue + at = event.get("at") + tokens = event.get("tokens") + if ( + isinstance(at, (int, float)) + and not isinstance(at, bool) + and isinstance(tokens, int) + and not isinstance(tokens, bool) + and tokens > 0 + and cutoff < float(at) <= now + WINDOW_SECONDS + ): + events.append({"at": float(at), "tokens": tokens}) + events.sort(key=lambda event: float(event["at"])) + return events + + +def _prune_state(state: dict, now: float) -> dict: + buckets = state.get("buckets") + clean: dict[str, list[dict[str, float | int]]] = {} + if isinstance(buckets, dict): + for key, raw_events in buckets.items(): + if not isinstance(key, str): + continue + events = _valid_recent_events(raw_events, now) + if events: + clean[key] = events + cooldowns: dict[str, float] = {} + raw_cooldowns = state.get("cooldowns") + if isinstance(raw_cooldowns, dict): + for key, raw_until in raw_cooldowns.items(): + if ( + isinstance(key, str) + and isinstance(raw_until, (int, float)) + and not isinstance(raw_until, bool) + and now < float(raw_until) + ): + cooldowns[key] = float(raw_until) + return {"version": 2, "buckets": clean, "cooldowns": cooldowns} + + +def _header(headers: Mapping[str, str], name: str) -> str | None: + """Read one response header from either a plain or case-insensitive map.""" + value = headers.get(name) + if value is not None: + return value + lowered = name.lower() + for key, candidate in headers.items(): + if key.lower() == lowered: + return candidate + return None + + +def _retry_after_seconds(headers: Mapping[str, str], now: float) -> float | None: + """Parse Retry-After in either delta-seconds or HTTP-date form.""" + value = _header(headers, "Retry-After") + if value is None: + return None + value = value.strip() + try: + seconds = float(value) + return max(0.0, seconds) if math.isfinite(seconds) else None + except ValueError: + pass + try: + parsed = parsedate_to_datetime(value) + seconds = parsed.timestamp() - now + return max(0.0, seconds) if math.isfinite(seconds) else None + except (TypeError, ValueError, OverflowError): + return None + + +def _write_state(path: Path, state: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temp_name: str | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as handle: + temp_name = handle.name + json.dump(state, handle, separators=(",", ":"), sort_keys=True) + handle.write("\n") + os.replace(temp_name, path) + temp_name = None + finally: + if temp_name is not None: + try: + Path(temp_name).unlink() + except FileNotFoundError: + pass + + +class SharedCodexRateLimiter: + """Reserve Codex request estimates in a shared rolling 60-second window.""" + + def __init__( + self, + workspace: str, + *, + state_path: Path | None = None, + lock_path: Path | None = None, + target_limits: dict[str, int] | None = None, + clock: Callable[[], float] = time.time, + sleeper: Callable[[float], None] = time.sleep, + notice: Callable[[str, float], None] | None = None, + retry_notice: Callable[[str, float, int, bool], None] | None = None, + jitter: Callable[[], float] = random.random, + ) -> None: + app_dir = config_io.APP_DIR + self.workspace = workspace.rstrip("/").lower() + self.state_path = state_path or app_dir / STATE_FILE_NAME + self.lock_path = lock_path or app_dir / LOCK_FILE_NAME + self.target_limits = dict(DEFAULT_TARGET_LIMITS if target_limits is None else target_limits) + self.clock = clock + self.sleeper = sleeper + self.notice = notice or self._default_notice + self.retry_notice = retry_notice or self._default_retry_notice + self.jitter = jitter + self._unavailable_notice_sent = False + + @staticmethod + def _default_notice(model: str, wait_seconds: float) -> None: + seconds = max(1, math.ceil(wait_seconds)) + sys.stderr.write( + f"[ucode] Pausing a {model} request for about {seconds}s to stay under the " + "shared Databricks rate limit.\n" + ) + sys.stderr.flush() + + @staticmethod + def _default_retry_notice( + model: str, + wait_seconds: float, + attempt: int, + used_retry_after: bool, + ) -> None: + seconds = max(1, math.ceil(wait_seconds)) + source = "server Retry-After" if used_retry_after else "jittered backoff" + sys.stderr.write( + f"[ucode] {model} returned 429; keeping the thread alive and retrying in " + f"about {seconds}s ({source}, attempt {attempt}).\n" + ) + sys.stderr.flush() + + def __call__(self, body: bytes) -> None: + estimate = estimate_request(body) + try: + # A 429 whose request body was unreadable creates a workspace-wide + # cooldown. Honor it even when this request is also unreadable or is + # for a model whose quota has not been added yet. + self.wait_for_cooldown(estimate[1] if estimate is not None else None) + if estimate is None: + return + model, model_key, estimated_tokens = estimate + self.wait_for_capacity(model, model_key, estimated_tokens) + except OSError as exc: + # A local permissions/filesystem problem must not make the model + # endpoint unreachable. Surface it once, then fail open. + if not self._unavailable_notice_sent: + self._unavailable_notice_sent = True + sys.stderr.write( + "[ucode] Shared Codex rate limiter is unavailable " + f"({type(exc).__name__}); sending without a local throttle.\n" + ) + sys.stderr.flush() + + def _cooldown_keys(self, model_key: str | None) -> list[str]: + keys = [f"{self.workspace}|*"] + if model_key is not None: + keys.append(f"{self.workspace}|{model_key}") + return keys + + def wait_for_cooldown(self, model_key: str | None) -> None: + """Wait for the latest applicable workspace/model 429 cooldown.""" + while True: + now = self.clock() + with _locked(self.lock_path): + state = _prune_state(_read_state(self.state_path), now) + cooldowns = state["cooldowns"] + blocked_until = max( + (float(cooldowns.get(key, 0.0)) for key in self._cooldown_keys(model_key)), + default=0.0, + ) + _write_state(self.state_path, state) + if blocked_until <= now: + return + self.sleeper(max(0.01, blocked_until - now)) + + def _fallback_retry_delay(self, attempt: int) -> float: + exponent = max(0, attempt - 1) + ceiling = min( + RETRY_BACKOFF_MAX_SECONDS, + RETRY_BACKOFF_BASE_SECONDS * (2 ** min(exponent, 30)), + ) + # Equal jitter keeps meaningful backoff while preventing many local + # Codex processes from waking on exactly the same instant. + return ceiling * (0.5 + 0.5 * min(1.0, max(0.0, self.jitter()))) + + def _schedule_cooldown(self, model_key: str | None, delay: float, now: float) -> float: + """Claim one retry slot at the tail of the shared cooldown queue.""" + bucket_key = self._cooldown_keys(model_key)[-1] + with _locked(self.lock_path): + state = _prune_state(_read_state(self.state_path), now) + cooldowns = state["cooldowns"] + scheduled_until = max(float(cooldowns.get(bucket_key, 0.0)), now) + delay + cooldowns[bucket_key] = scheduled_until + _write_state(self.state_path, state) + return scheduled_until + + def _sleep_until(self, scheduled_until: float) -> None: + while True: + remaining = scheduled_until - self.clock() + if remaining <= 0: + return + self.sleeper(max(0.01, remaining)) + + def retry_after_429( + self, + body: bytes | None, + headers: Mapping[str, str], + attempt: int, + ) -> None: + """Publish a shared cooldown, then wait before an internal 429 retry. + + The proxy calls this only after draining a 429 response. Recognized + models get an isolated bucket; unreadable/compressed bodies and future + models safely fall back to a workspace-wide bucket. No prompt, response, + header value, or credential is written to disk. + """ + estimate = estimate_request(body) if body is not None else None + model = estimate[0] if estimate is not None else "Codex" + model_key = estimate[1] if estimate is not None else None + now = self.clock() + retry_after = _retry_after_seconds(headers, now) + used_retry_after = retry_after is not None + if retry_after is None: + delay = self._fallback_retry_delay(attempt) + else: + # Positive jitter preserves the server's minimum while spreading a + # local herd. Cap only the added jitter, never Retry-After itself. + extra_ceiling = min(5.0, max(0.1, retry_after * RETRY_AFTER_JITTER_PERCENT / 100)) + delay = retry_after + extra_ceiling * min(1.0, max(0.0, self.jitter())) + delay = max(0.01, delay) + + try: + # Appending rather than overwriting gives simultaneous 429s separate + # retry slots. Each blocked handler waits for its own slot, while new + # requests observe the tail and do not jump ahead of the queue. + scheduled_until = self._schedule_cooldown(model_key, delay, now) + actual_wait = max(0.01, scheduled_until - now) + self.retry_notice(model, actual_wait, attempt, used_retry_after) + self._sleep_until(scheduled_until) + except OSError as exc: + self.retry_notice(model, delay, attempt, used_retry_after) + if not self._unavailable_notice_sent: + self._unavailable_notice_sent = True + sys.stderr.write( + "[ucode] Shared Codex rate limiter is unavailable " + f"({type(exc).__name__}); using local 429 backoff only.\n" + ) + sys.stderr.flush() + self.sleeper(delay) + + def wait_for_capacity(self, model: str, model_key: str, estimated_tokens: int) -> None: + target = self.target_limits.get(model_key) + if not isinstance(target, int) or target <= 0 or estimated_tokens <= 0: + return + + # An estimate larger than the local target can never satisfy + # total+estimate <= target. Reserve one full window instead: it proceeds + # when the bucket is empty and serializes any following request. + reservation = min(estimated_tokens, target) + bucket_key = f"{self.workspace}|{model_key}" + notice_sent = False + + while True: + now = self.clock() + with _locked(self.lock_path): + state = _prune_state(_read_state(self.state_path), now) + buckets = state["buckets"] + events = buckets.setdefault(bucket_key, []) + used = sum(int(event["tokens"]) for event in events) + if used + reservation <= target: + events.append({"at": now, "tokens": reservation}) + _write_state(self.state_path, state) + return + + oldest = min(float(event["at"]) for event in events) + wait_seconds = max(0.01, oldest + WINDOW_SECONDS - now) + _write_state(self.state_path, state) + + # Never hold the cross-process lock while sleeping or printing. + if not notice_sent: + self.notice(model, wait_seconds) + notice_sent = True + self.sleeper(wait_seconds) diff --git a/src/ucode/codex_request.py b/src/ucode/codex_request.py new file mode 100644 index 00000000..64411b0f --- /dev/null +++ b/src/ucode/codex_request.py @@ -0,0 +1,67 @@ +"""Compatibility rewrites for Codex Responses API requests. + +Some non-OpenAI FMAPI models return visible ``reasoning.content`` items. Codex +can replay those items while continuing on the same model, but OpenAI reasoning +models require replayable reasoning state in ``encrypted_content`` and reject a +non-empty ``content`` array. Keep the persisted transcript untouched and remove +only that provider-specific field from requests sent to the affected models. +""" + +from __future__ import annotations + +import json +import re + +_STRICT_REASONING_REPLAY_MODEL_KEYS = frozenset( + { + "gpt6astra", + "gpt56sol", + "gpt56terra", + "gpt56luna", + } +) + + +def _model_key(model: str) -> str: + return re.sub(r"[^a-z0-9]+", "", model.strip().lower()) + + +def _needs_strict_reasoning_replay(model: object) -> bool: + if not isinstance(model, str): + return False + normalized = _model_key(model) + return any(normalized.endswith(key) for key in _STRICT_REASONING_REPLAY_MODEL_KEYS) + + +def sanitize_reasoning_replay(body: bytes) -> bytes: + """Remove nonportable visible reasoning from strict-model request history. + + Invalid JSON, unknown models, and already-compatible requests are returned + byte-for-byte so the shared proxy remains transparent outside this narrow + Responses compatibility case. + """ + try: + payload = json.loads(body) + except (json.JSONDecodeError, UnicodeDecodeError): + return body + if not isinstance(payload, dict) or not _needs_strict_reasoning_replay(payload.get("model")): + return body + + items = payload.get("input") + if not isinstance(items, list): + return body + + changed = False + for item in items: + if ( + isinstance(item, dict) + and item.get("type") == "reasoning" + and isinstance(item.get("content"), list) + and item["content"] + ): + item.pop("content") + changed = True + + if not changed: + return body + return json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8") diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index e5cd4788..3aad0ba7 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -1,4 +1,4 @@ -"""Loopback refresh proxy for Claude gateway requests. +"""Loopback refresh proxy for coding-agent gateway requests. A relayed Model Provider Service authenticates the caller's own Anthropic subscription OAuth (which Claude Code owns in the `Authorization` header) and @@ -24,6 +24,7 @@ import threading import time import uuid +from collections.abc import Callable, Mapping from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import httpx @@ -200,6 +201,9 @@ class _ProxyHandler(BaseHTTPRequestHandler): cache: TokenCache client: httpx.Client token_header = AI_GATEWAY_TOKEN_HEADER + request_transform: Callable[[bytes], bytes] | None = None + request_gate: Callable[[bytes], None] | None = None + rate_limit_retry: Callable[[bytes | None, Mapping[str, str], int], None] | None = None def log_message(self, format: str, *args: object) -> None: return @@ -225,47 +229,61 @@ def _handle(self) -> None: path=self.path.split("?", 1)[0], ) try: - # First attempt with the current token. - headers = forwarded_request_headers(self, self.cache.token, self.token_header) - with self.client.stream(self.command, url, headers=headers, content=body) as resp: - log_proxy_diagnostic( - "upstream_headers", - request_id=diagnostic_id, - attempt=1, - status=resp.status_code, - elapsed_ms=round((time.monotonic() - started) * 1000), - ) - if resp.status_code not in (401, 403): - self._relay_response(resp, diagnostic_id=diagnostic_id, started=started) - return - # Auth rejected. Drain the (small) error body so the pooled - # connection can be reused, then fall through to one retry. - resp.read() - # A relayed 401/403 may be a stale Databricks swap token rather than a - # bad Anthropic OAuth — the two are indistinguishable from the status - # alone. Force-refresh the Databricks token and retry once. If it was the - # Anthropic layer, the retry still 401s and we relay it verbatim, so a - # genuine re-auth is triggered; a stale-Databricks 401 self-heals here - # instead of surfacing to Claude Code as a spurious Anthropic prompt. - try: - self.cache.refresh() - except RuntimeError as exc: - # Refresh failed: the Databricks OAuth session is dead (not just the - # access token) and can't be re-minted non-interactively. Surface the - # `databricks auth login` hint rather than silently relaying a bare 401, - # which otherwise reads as an Anthropic `/login` prompt and sends the - # user to the wrong re-auth. Still retry + relay with the existing token. - log_token_refresh_failure(exc) - headers = forwarded_request_headers(self, self.cache.token, self.token_header) - with self.client.stream(self.command, url, headers=headers, content=body) as resp: - log_proxy_diagnostic( - "upstream_headers", - request_id=diagnostic_id, - attempt=2, - status=resp.status_code, - elapsed_ms=round((time.monotonic() - started) * 1000), - ) - self._relay_response(resp, diagnostic_id=diagnostic_id, started=started) + if self.request_transform is not None and body is not None: + body = self.request_transform(body) + if self.request_gate is not None and body is not None: + self.request_gate(body) + auth_retried = False + rate_limit_attempt = 0 + upstream_attempt = 0 + rate_limit_retry = self.rate_limit_retry + while True: + upstream_attempt += 1 + headers = forwarded_request_headers(self, self.cache.token, self.token_header) + retry_auth = False + retry_rate_limit_headers: Mapping[str, str] | None = None + with self.client.stream(self.command, url, headers=headers, content=body) as resp: + log_proxy_diagnostic( + "upstream_headers", + request_id=diagnostic_id, + attempt=upstream_attempt, + status=resp.status_code, + elapsed_ms=round((time.monotonic() - started) * 1000), + ) + if resp.status_code in (401, 403) and not auth_retried: + # Drain the small auth error so the pooled connection can + # be reused before refreshing and retrying once. + resp.read() + retry_auth = True + elif resp.status_code == 429 and rate_limit_retry is not None: + # Keep the 429 inside the proxy so Codex never spends one + # of its own finite retries. The callback publishes and + # waits on a cross-process cooldown before we try again. + resp.read() + retry_rate_limit_headers = dict(resp.headers) + else: + self._relay_response(resp, diagnostic_id=diagnostic_id, started=started) + return + + if retry_auth: + auth_retried = True + # A relayed 401/403 may be a stale Databricks swap token rather + # than a bad Anthropic OAuth. Force-refresh once; a persistent + # auth rejection is relayed on the next loop. + try: + self.cache.refresh() + except RuntimeError as exc: + log_token_refresh_failure(exc) + continue + + if retry_rate_limit_headers is not None: + rate_limit_attempt += 1 + assert rate_limit_retry is not None + rate_limit_retry(body, retry_rate_limit_headers, rate_limit_attempt) + continue + + # Every response either relays or chooses one of the retry paths. + raise AssertionError("unreachable proxy response state") except (BrokenPipeError, ConnectionResetError): # Client closed before/while we relayed headers — routine on cancel. log_proxy_diagnostic( @@ -374,18 +392,30 @@ def start_proxy( port: int, token_header: str, force_refresh_near_expiry: bool, + *, + upstream_base: str | None = None, + request_transform: Callable[[bytes], bytes] | None = None, + request_gate: Callable[[bytes], None] | None = None, + rate_limit_retry: Callable[[bytes | None, Mapping[str, str], int], None] | None = None, ) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. Binds ``port``, falling back to a fresh OS-assigned port when it is already in use (e.g. a prior session's proxy that was killed before its teardown ran still holds the socket). The caller reads ``server.server_address[1]`` for the - actual port and points Claude Code at it. + actual port and points the coding agent at it. ``upstream_base`` defaults to + Claude's Anthropic gateway path; callers for other APIs provide their own. + ``request_transform`` and ``request_gate`` run once per downstream request, + before any upstream attempt (including retry paths). The gate receives the + transformed body. When provided, ``rate_limit_retry`` receives every drained + 429 plus the transformed request body and one-based rate-limit attempt. It + may block while applying a shared cooldown; the proxy then retries internally + instead of exposing the 429 to the coding agent's finite retry budget. Returns (server, cache, client); the caller runs the server (e.g. in a thread) and calls shutdown()/cache.stop()/client.close() on exit. """ - upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" + upstream_base = upstream_base or f"{workspace.rstrip('/')}/ai-gateway/anthropic/" cache = TokenCache( workspace, profile, @@ -399,7 +429,20 @@ def start_proxy( handler = type( "BoundProxyHandler", (_ProxyHandler,), - {"cache": cache, "client": client, "token_header": token_header}, + { + "cache": cache, + "client": client, + "token_header": token_header, + # Functions stored directly on a class are descriptors and would + # otherwise receive the handler instance as an extra argument. + "request_transform": ( + staticmethod(request_transform) if request_transform is not None else None + ), + "request_gate": staticmethod(request_gate) if request_gate is not None else None, + "rate_limit_retry": ( + staticmethod(rate_limit_retry) if rate_limit_retry is not None else None + ), + }, ) try: server = ThreadingHTTPServer(("127.0.0.1", port), handler) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 2cd11dc3..e13381c3 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -2,10 +2,11 @@ from __future__ import annotations +import contextlib import os -from pathlib import Path import pytest +import tomlkit from ucode.agents import codex from ucode.config_io import read_toml_safe @@ -566,6 +567,7 @@ def fake_run(argv, **kwargs): monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: fallbacks.append(argv)) monkeypatch.setattr(codex, "get_databricks_token", lambda workspace, profile=None: "tok") monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) + monkeypatch.setenv(codex.CODEX_RATE_LIMITER_ENV, "0") return runs, fallbacks def test_sets_oauth_token_and_runs_with_profile(self, monkeypatch): @@ -678,9 +680,159 @@ def test_fast_success_does_not_retry(self, monkeypatch): assert exc.value.code == 0 assert fallbacks == [] + def test_normal_launch_routes_provider_through_shared_limiter(self, tmp_path, monkeypatch): + profile_path = tmp_path / "ucode.config.toml" + profile_path.write_text( + 'model_provider = "ucode-databricks"\n\n' + "[model_providers.ucode-databricks]\n" + 'name = "Databricks AI Gateway"\n' + 'base_url = "https://example.databricks.com/ai-gateway/codex/v1"\n' + 'wire_api = "responses"\n', + encoding="utf-8", + ) + runs = [] + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) + monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) + monkeypatch.setattr(codex, "get_databricks_token", lambda *args: "token") + monkeypatch.setattr( + codex, + "_codex_request_proxy", + lambda state: contextlib.nullcontext("http://127.0.0.1:43210/v1"), + ) + + def run(argv, **_kwargs): + runs.append(argv) + return codex.subprocess.CompletedProcess(argv, 0) + + monkeypatch.setattr(codex.subprocess, "run", run) + monkeypatch.delenv(codex.CODEX_RATE_LIMITER_ENV, raising=False) + + with pytest.raises(SystemExit) as exc: + codex.launch({"workspace": WS}, ["exec", "hello"]) + + assert exc.value.code == 0 + argv = runs[0] + assert argv[:3] == ["codex", "--profile", "ucode"] + provider_arg = next( + arg for arg in argv if arg.startswith("model_providers.ucode-databricks=") + ) + assert 'base_url = "http://127.0.0.1:43210/v1"' in provider_arg + assert "features.enable_request_compression=false" in argv + assert argv[-2:] == ["exec", "hello"] + + def test_app_launch_routes_through_shared_limiter(self, tmp_path, monkeypatch): + profile_path = tmp_path / "ucode.config.toml" + profile_path.write_text( + 'model_provider = "ucode-databricks"\n\n' + "[model_providers.ucode-databricks]\n" + 'base_url = "https://example.databricks.com/ai-gateway/codex/v1"\n' + 'wire_api = "responses"\n', + encoding="utf-8", + ) + runs = [] + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) + monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) + monkeypatch.setattr(codex, "get_databricks_token", lambda *args: "token") + monkeypatch.setattr( + codex, + "_codex_request_proxy", + lambda state: contextlib.nullcontext("http://127.0.0.1:43210/v1"), + ) + + def run(argv, **_kwargs): + runs.append(argv) + return codex.subprocess.CompletedProcess(argv, 0) + + monkeypatch.setattr(codex.subprocess, "run", run) + monkeypatch.delenv(codex.CODEX_RATE_LIMITER_ENV, raising=False) + + with pytest.raises(SystemExit) as exc: + codex.launch({"workspace": WS}, ["app", "--new-window"]) + + assert exc.value.code == 0 + assert runs[0][:2] == ["codex", "app"] + assert "--profile" not in runs[0] + assert "features.enable_request_compression=false" in runs[0] + provider_arg = next( + arg for arg in runs[0] if arg.startswith("model_providers.ucode-databricks=") + ) + assert 'base_url = "http://127.0.0.1:43210/v1"' in provider_arg + assert runs[0][-1] == "--new-window" + + def test_server_family_keeps_existing_non_proxy_path(self, monkeypatch): + runs = [] + monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) + monkeypatch.setattr(codex, "get_databricks_token", lambda *args: "token") + monkeypatch.setattr( + codex, + "_codex_request_proxy", + lambda state: pytest.fail("server-family command must not start a session proxy"), + ) + + def run(argv, **_kwargs): + runs.append(argv) + return codex.subprocess.CompletedProcess(argv, 0) + + monkeypatch.setattr(codex.subprocess, "run", run) + monkeypatch.delenv(codex.CODEX_RATE_LIMITER_ENV, raising=False) + + with pytest.raises(SystemExit) as exc: + codex.launch({"workspace": WS}, ["app-server", "--listen", "stdio://"]) + + assert exc.value.code == 0 + assert runs == [["codex", "--profile", "ucode", "app-server", "--listen", "stdio://"]] + + def test_proxy_lifecycle_uses_codex_gateway_and_authorization(self, monkeypatch): + calls = [] + + class Server: + server_address = ("127.0.0.1", 43210) + + def serve_forever(self): + calls.append("serve") + + def shutdown(self): + calls.append("shutdown") + + def server_close(self): + calls.append("server_close") + + class Cache: + def stop(self): + calls.append("cache_stop") + + class Client: + def close(self): + calls.append("client_close") + + def start_proxy(*args, **kwargs): + calls.append((args, kwargs)) + return Server(), Cache(), Client() + + monkeypatch.setattr(codex.gateway_proxy, "start_proxy", start_proxy) + + with codex._codex_request_proxy({"workspace": WS, "profile": "test"}) as base_url: + assert base_url == "http://127.0.0.1:43210/v1" + + args, kwargs = calls[0] + assert args == (WS, "test", 0) + assert kwargs["token_header"] == codex.gateway_proxy.AUTHORIZATION_HEADER + assert kwargs["force_refresh_near_expiry"] is True + assert kwargs["upstream_base"] == f"{WS}/ai-gateway/codex/" + assert kwargs["request_transform"] is codex.sanitize_reasoning_replay + assert isinstance(kwargs["request_gate"], codex.SharedCodexRateLimiter) + assert kwargs["rate_limit_retry"] == kwargs["request_gate"].retry_after_429 + assert calls.count("serve") == 1 + assert [call for call in calls[1:] if call != "serve"] == [ + "cache_stop", + "shutdown", + "server_close", + "client_close", + ] + class TestCodexManagedConfig: - """Every normal configuration also reconciles Codex's OS-managed config.""" + """Codex must keep launch-scoped providers out of its higher-precedence managed config.""" def _patch(self, tmp_path, monkeypatch): config_path = tmp_path / ".codex" / "ucode.config.toml" @@ -689,70 +841,105 @@ def _patch(self, tmp_path, monkeypatch): monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "codex-ucode-config.backup.toml") monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") monkeypatch.setattr(codex, "save_state", lambda state: None) - monkeypatch.setattr(codex, "managed_writes_allowed", lambda: True) - # Deterministic managed path + a mocked sudo writer that writes straight to disk, so the test - # can read the TOML back and NO real sudo/`/etc` write ever happens. monkeypatch.setattr(codex, "_managed_config_path", lambda: managed_path) + return config_path, managed_path - def fake_write_managed(path, text, **kwargs): - Path(path).parent.mkdir(parents=True, exist_ok=True) - Path(path).write_text(text, encoding="utf-8") - return "written" + def test_preserves_unrelated_managed_config_without_adding_provider( + self, tmp_path, monkeypatch + ): + _, managed_path = self._patch(tmp_path, monkeypatch) + managed_path.parent.mkdir(parents=True, exist_ok=True) + original = 'model = "my-own"\napproval_policy = "on-request"\n' + managed_path.write_text(original, encoding="utf-8") + state = {"workspace": WS, "codex_models": ["gpt-5"]} + codex.write_tool_config(state) - monkeypatch.setattr(codex, "reconcile_managed_file", fake_write_managed) - return config_path, managed_path + assert managed_path.read_text(encoding="utf-8") == original + assert codex.managed_config_is_current(state) is True - def test_writes_managed_config_by_default(self, tmp_path, monkeypatch): + def test_absent_managed_config_stays_absent(self, tmp_path, monkeypatch): _, managed_path = self._patch(tmp_path, monkeypatch) state = {"workspace": WS, "codex_models": ["gpt-5"]} codex.write_tool_config(state) - doc = read_toml_safe(managed_path) - assert doc["model_provider"] == "ucode-databricks" - assert "model" not in doc - assert "ucode-databricks" in doc["model_providers"] + assert not managed_path.exists() + assert codex.managed_config_is_current(state) is True + + def test_migrates_tracked_provider_and_preserves_external_policy(self, tmp_path, monkeypatch): + from ucode import managed_files - def test_managed_config_preserves_other_keys(self, tmp_path, monkeypatch): _, managed_path = self._patch(tmp_path, monkeypatch) managed_path.parent.mkdir(parents=True, exist_ok=True) - managed_path.write_text( - 'model = "my-own"\napproval_policy = "on-request"\n', encoding="utf-8" + original = 'approval_policy = "on-request"\n[enterprise]\nkeep = true\n' + managed_path.write_text(original, encoding="utf-8") + stale = read_toml_safe(managed_path) + stale.update(codex.render_overlay(WS)) + monkeypatch.setattr(managed_files, "managed_writes_allowed", lambda: True) + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), ) + managed_files.reconcile_managed_file( + managed_path, + tomlkit.dumps(stale), + tool="codex", + display="Codex", + owned_paths=codex.MANAGED_KEYS, + ) + state = {"workspace": WS, "codex_models": ["gpt-5"]} codex.write_tool_config(state) - doc = read_toml_safe(managed_path) - # ucode removes its stale model pin, but other keys already in the managed file survive. - assert doc["approval_policy"] == "on-request" - assert "model" not in doc + assert managed_path.read_text(encoding="utf-8") == original + assert codex.managed_config_status(state)[1] == "compatible (local settings)" + + def test_migrates_tracked_provider_file_created_by_ucode(self, tmp_path, monkeypatch): + from ucode import managed_files - def test_noninteractive_uses_local_config_when_managed_config_is_compatible( - self, tmp_path, monkeypatch - ): _, managed_path = self._patch(tmp_path, monkeypatch) - monkeypatch.setattr(codex, "managed_writes_allowed", lambda: False) + monkeypatch.setattr(managed_files, "managed_writes_allowed", lambda: True) + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: ( + target.parent.mkdir(parents=True, exist_ok=True), + target.write_text(text, encoding="utf-8"), + ), + ) + monkeypatch.setattr(managed_files, "_sudo_remove", lambda target: target.unlink()) + managed_files.reconcile_managed_file( + managed_path, + tomlkit.dumps(codex.render_overlay(WS)), + tool="codex", + display="Codex", + owned_paths=codex.MANAGED_KEYS, + ) + state = {"workspace": WS, "codex_models": ["gpt-5"]} codex.write_tool_config(state) + assert not managed_path.exists() + assert codex.managed_config_status(state)[1] == "compatible (local settings)" - def test_noninteractive_preserves_unrelated_managed_config(self, tmp_path, monkeypatch): + def test_external_managed_provider_conflict_blocks_configuration(self, tmp_path, monkeypatch): _, managed_path = self._patch(tmp_path, monkeypatch) managed_path.parent.mkdir(parents=True, exist_ok=True) - original = 'approval_policy = "on-request"\n' - managed_path.write_text(original, encoding="utf-8") - monkeypatch.setattr(codex, "managed_writes_allowed", lambda: False) - - codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) + managed_path.write_text('model_provider = "enterprise"\n', encoding="utf-8") - assert managed_path.read_text(encoding="utf-8") == original + with pytest.raises(RuntimeError, match="override the launch-scoped provider"): + codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) - def test_noninteractive_fails_when_managed_config_conflicts(self, tmp_path, monkeypatch): + def test_managed_ucode_base_url_conflict_blocks_configuration(self, tmp_path, monkeypatch): _, managed_path = self._patch(tmp_path, monkeypatch) managed_path.parent.mkdir(parents=True, exist_ok=True) - managed_path.write_text('model_provider = "enterprise"\n', encoding="utf-8") - monkeypatch.setattr(codex, "managed_writes_allowed", lambda: False) + managed_path.write_text( + "[model_providers.ucode-databricks]\n" + 'base_url = "https://direct.example/ai-gateway/codex/v1"\n', + encoding="utf-8", + ) - with pytest.raises(RuntimeError, match="cannot be applied non-interactively"): + with pytest.raises(RuntimeError, match="ucode-databricks.base_url"): codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) def test_invalid_managed_toml_is_not_modified(self, tmp_path, monkeypatch): @@ -760,7 +947,7 @@ def test_invalid_managed_toml_is_not_modified(self, tmp_path, monkeypatch): managed_path.parent.mkdir(parents=True, exist_ok=True) managed_path.write_text("[invalid", encoding="utf-8") - with pytest.raises(RuntimeError, match="Cannot safely update Codex managed settings"): + with pytest.raises(RuntimeError, match="Cannot safely inspect Codex managed settings"): codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) assert managed_path.read_text(encoding="utf-8") == "[invalid" diff --git a/tests/test_codex_config.py b/tests/test_codex_config.py index 445b95fb..92c7e83a 100644 --- a/tests/test_codex_config.py +++ b/tests/test_codex_config.py @@ -1,5 +1,9 @@ from __future__ import annotations +import tomllib + +import tomlkit + from ucode.agents import codex from ucode.codex_config import codex_config_args @@ -29,3 +33,29 @@ def test_layers_provider_overrides_without_replacing_user_config(self, monkeypat assert "/ai-gateway/codex/v1" in provider_override assert 'command = "' in provider_override assert '"myprof"' in provider_override + + def test_renders_nested_tables_from_parsed_config_as_inline_tables(self): + config = tomlkit.parse( + """ +[model_providers.ucode-databricks] +name = "Databricks AI Gateway" + +[model_providers.ucode-databricks.http_headers] +User-Agent = "ucode/0.1.0" + +[model_providers.ucode-databricks.auth] +command = "ucode" +args = ["auth-token"] +""" + ) + + args = codex_config_args(config) + + assert args[0] == "--config" + key, rendered = args[1].split("=", 1) + assert key == "model_providers.ucode-databricks" + assert tomllib.loads(f"value = {rendered}")["value"] == { + "name": "Databricks AI Gateway", + "http_headers": {"User-Agent": "ucode/0.1.0"}, + "auth": {"command": "ucode", "args": ["auth-token"]}, + } diff --git a/tests/test_codex_rate_limit.py b/tests/test_codex_rate_limit.py new file mode 100644 index 00000000..3feca4cd --- /dev/null +++ b/tests/test_codex_rate_limit.py @@ -0,0 +1,319 @@ +from __future__ import annotations + +import json +import multiprocessing +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime +from pathlib import Path + +import pytest + +from ucode import codex_rate_limit + + +def _body(model: str, padding: str = "") -> bytes: + return json.dumps({"model": model, "input": padding}, separators=(",", ":")).encode() + + +def _reserve_in_process(state_path: str, lock_path: str, count: int) -> None: + limiter = codex_rate_limit.SharedCodexRateLimiter( + "https://workspace", + state_path=Path(state_path), + lock_path=Path(lock_path), + target_limits={"gpt56sol": 1_000}, + clock=lambda: 100.0, + ) + for _ in range(count): + limiter.wait_for_capacity("gpt-5.6-sol", "gpt56sol", 10) + + +@pytest.mark.parametrize( + ("model", "expected_key", "published"), + [ + ("gpt-6-astra", "gpt6astra", 200_000), + ("system.ai.gpt-5-6-sol", "gpt56sol", 2_000_000), + ("databricks-gpt-5.6-terra", "gpt56terra", 2_000_000), + ("eu/gpt-5-6-luna", "gpt56luna", 2_000_000), + ("system.ai.kimi-k3", "kimik3", 200_000), + ("databricks-qwen35-122b-a10b", "qwen35122ba10b", 1_000_000), + ( + "system.ai.qwen3-next-80b-a3b-instruct", + "qwen3next80ba3binstruct", + 1_000_000, + ), + ], +) +def test_estimates_every_known_codex_model(model, expected_key, published): + body = _body(model, "x" * 100) + + assert codex_rate_limit.estimate_request(body) == ( + model, + expected_key, + (len(body) + 2) // 3, + ) + assert codex_rate_limit.PUBLISHED_INPUT_TOKENS_PER_MINUTE[expected_key] == published + assert codex_rate_limit.DEFAULT_TARGET_LIMITS[expected_key] == published * 90 // 100 + + +@pytest.mark.parametrize( + "body", + [ + b"not json", + b"[]", + b'{"input":"hi"}', + b'{"model":"gpt-future"}', + ], +) +def test_unknown_or_unreadable_requests_pass_through(body): + assert codex_rate_limit.estimate_request(body) is None + + +def test_exact_window_boundary_releases_capacity(tmp_path): + now = [100.0] + sleeps = [] + notices = [] + + def sleep(seconds): + sleeps.append(seconds) + now[0] += seconds + + limiter = codex_rate_limit.SharedCodexRateLimiter( + "https://workspace", + state_path=tmp_path / "state.json", + lock_path=tmp_path / "state.lock", + target_limits={"gpt6astra": 10}, + clock=lambda: now[0], + sleeper=sleep, + notice=lambda model, seconds: notices.append((model, seconds)), + ) + + limiter.wait_for_capacity("gpt-6-astra", "gpt6astra", 10) + limiter.wait_for_capacity("gpt-6-astra", "gpt6astra", 1) + + assert sleeps == [60.0] + assert notices == [("gpt-6-astra", 60.0)] + state = json.loads((tmp_path / "state.json").read_text()) + assert state["buckets"]["https://workspace|gpt6astra"] == [{"at": 160.0, "tokens": 1}] + + +def test_oversized_request_uses_one_full_window_reservation(tmp_path): + limiter = codex_rate_limit.SharedCodexRateLimiter( + "https://workspace", + state_path=tmp_path / "state.json", + lock_path=tmp_path / "state.lock", + target_limits={"gpt6astra": 10}, + clock=lambda: 100.0, + sleeper=lambda _seconds: pytest.fail("first oversized request must not wait forever"), + ) + + limiter.wait_for_capacity("gpt-6-astra", "gpt6astra", 100) + + state = json.loads((tmp_path / "state.json").read_text()) + assert state["buckets"]["https://workspace|gpt6astra"] == [{"at": 100.0, "tokens": 10}] + + +def test_models_and_workspaces_have_independent_buckets(tmp_path): + common = { + "state_path": tmp_path / "state.json", + "lock_path": tmp_path / "state.lock", + "target_limits": {"gpt6astra": 10, "gpt56sol": 10}, + "clock": lambda: 100.0, + "sleeper": lambda _seconds: pytest.fail("independent buckets should not wait"), + } + first = codex_rate_limit.SharedCodexRateLimiter("https://one", **common) + second = codex_rate_limit.SharedCodexRateLimiter("https://two", **common) + + first.wait_for_capacity("gpt-6-astra", "gpt6astra", 10) + first.wait_for_capacity("gpt-5.6-sol", "gpt56sol", 10) + second.wait_for_capacity("gpt-6-astra", "gpt6astra", 10) + + state = json.loads((tmp_path / "state.json").read_text()) + assert set(state["buckets"]) == { + "https://one|gpt6astra", + "https://one|gpt56sol", + "https://two|gpt6astra", + } + + +def test_concurrent_instances_do_not_lose_reservations(tmp_path): + state_path = tmp_path / "state.json" + lock_path = tmp_path / "state.lock" + + def reserve(_index): + limiter = codex_rate_limit.SharedCodexRateLimiter( + "https://workspace", + state_path=state_path, + lock_path=lock_path, + target_limits={"gpt56sol": 1_000}, + clock=lambda: 100.0, + sleeper=lambda _seconds: pytest.fail("reservations fit in the window"), + ) + limiter.wait_for_capacity("gpt-5.6-sol", "gpt56sol", 10) + + with ThreadPoolExecutor(max_workers=12) as pool: + list(pool.map(reserve, range(40))) + + state = json.loads(state_path.read_text()) + events = state["buckets"]["https://workspace|gpt56sol"] + assert len(events) == 40 + assert sum(event["tokens"] for event in events) == 400 + + +def test_concurrent_processes_share_the_same_reservations(tmp_path): + state_path = tmp_path / "state.json" + lock_path = tmp_path / "state.lock" + context = multiprocessing.get_context("spawn") + processes = [ + context.Process( + target=_reserve_in_process, + args=(str(state_path), str(lock_path), 10), + ) + for _ in range(4) + ] + + for process in processes: + process.start() + for process in processes: + process.join(timeout=15) + assert process.exitcode == 0 + + state = json.loads(state_path.read_text()) + events = state["buckets"]["https://workspace|gpt56sol"] + assert len(events) == 40 + assert sum(event["tokens"] for event in events) == 400 + + +def test_corrupt_state_is_replaced_under_lock(tmp_path): + state_path = tmp_path / "state.json" + state_path.write_text("truncated{") + limiter = codex_rate_limit.SharedCodexRateLimiter( + "https://workspace", + state_path=state_path, + lock_path=tmp_path / "state.lock", + target_limits={"gpt56luna": 10}, + clock=lambda: 100.0, + ) + + limiter.wait_for_capacity("gpt-5.6-luna", "gpt56luna", 3) + + state = json.loads(state_path.read_text()) + assert state["buckets"]["https://workspace|gpt56luna"][0]["tokens"] == 3 + + +def test_retry_after_seconds_supports_delta_and_http_date(): + now = datetime(2026, 9, 5, 16, 0, tzinfo=UTC).timestamp() + + assert codex_rate_limit._retry_after_seconds({"Retry-After": "12"}, now) == 12 + assert ( + codex_rate_limit._retry_after_seconds({"retry-after": "Sat, 05 Sep 2026 16:00:30 GMT"}, now) + == 30 + ) + assert codex_rate_limit._retry_after_seconds({"Retry-After": "invalid"}, now) is None + + +def test_429_honors_retry_after_and_publishes_model_cooldown(tmp_path): + now = [100.0] + sleeps = [] + snapshots = [] + notices = [] + + def sleep(seconds): + sleeps.append(seconds) + snapshots.append(json.loads((tmp_path / "state.json").read_text())) + now[0] += seconds + + limiter = codex_rate_limit.SharedCodexRateLimiter( + "https://workspace", + state_path=tmp_path / "state.json", + lock_path=tmp_path / "state.lock", + target_limits={}, + clock=lambda: now[0], + sleeper=sleep, + retry_notice=lambda *args: notices.append(args), + jitter=lambda: 0.0, + ) + + limiter.retry_after_429(_body("gpt-5.6-sol"), {"Retry-After": "5"}, 1) + + assert sleeps == [5.0] + assert notices == [("gpt-5.6-sol", 5.0, 1, True)] + assert snapshots[0]["cooldowns"] == {"https://workspace|gpt56sol": 105.0} + + +def test_unknown_model_429_uses_workspace_wide_cooldown(tmp_path): + now = [100.0] + snapshots = [] + + def sleep(seconds): + snapshots.append(json.loads((tmp_path / "state.json").read_text())) + now[0] += seconds + + limiter = codex_rate_limit.SharedCodexRateLimiter( + "https://workspace", + state_path=tmp_path / "state.json", + lock_path=tmp_path / "state.lock", + target_limits={}, + clock=lambda: now[0], + sleeper=sleep, + retry_notice=lambda *_args: None, + jitter=lambda: 0.0, + ) + + limiter.retry_after_429(_body("future-model"), {"Retry-After": "3"}, 1) + + assert snapshots[0]["cooldowns"] == {"https://workspace|*": 103.0} + + +def test_model_cooldown_does_not_pause_another_model(tmp_path): + state_path = tmp_path / "state.json" + state_path.write_text( + json.dumps( + { + "version": 2, + "buckets": {}, + "cooldowns": {"https://workspace|gpt56sol": 105.0}, + } + ) + ) + sleeps = [] + limiter = codex_rate_limit.SharedCodexRateLimiter( + "https://workspace", + state_path=state_path, + lock_path=tmp_path / "state.lock", + target_limits={}, + clock=lambda: 100.0, + sleeper=sleeps.append, + ) + + limiter.wait_for_cooldown("gpt56terra") + + assert sleeps == [] + + +def test_missing_retry_after_uses_capped_jittered_exponential_backoff(tmp_path): + limiter = codex_rate_limit.SharedCodexRateLimiter( + "https://workspace", + state_path=tmp_path / "state.json", + lock_path=tmp_path / "state.lock", + jitter=lambda: 0.0, + ) + + assert [limiter._fallback_retry_delay(attempt) for attempt in (1, 2, 3, 10)] == [ + 1.0, + 2.0, + 4.0, + 150.0, + ] + + +def test_concurrent_429s_claim_separate_shared_retry_slots(tmp_path): + common = { + "state_path": tmp_path / "state.json", + "lock_path": tmp_path / "state.lock", + "target_limits": {}, + } + first = codex_rate_limit.SharedCodexRateLimiter("https://workspace", **common) + second = codex_rate_limit.SharedCodexRateLimiter("https://workspace", **common) + + assert first._schedule_cooldown("gpt56sol", 5.0, 100.0) == 105.0 + assert second._schedule_cooldown("gpt56sol", 5.0, 100.0) == 110.0 diff --git a/tests/test_codex_request.py b/tests/test_codex_request.py new file mode 100644 index 00000000..5bab836e --- /dev/null +++ b/tests/test_codex_request.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import json + +import pytest + +from ucode.codex_request import sanitize_reasoning_replay + + +def _body(model: str, items: list[dict]) -> bytes: + return json.dumps({"model": model, "input": items}, indent=2).encode() + + +@pytest.mark.parametrize( + "model", + [ + "gpt-6-astra", + "system.ai.gpt-5-6-sol", + "databricks-gpt-5.6-terra", + "eu/gpt-5-6-luna", + ], +) +def test_strips_visible_reasoning_for_strict_models(model): + reasoning = { + "type": "reasoning", + "id": "rs_provider", + "summary": [], + "content": [{"type": "reasoning_text", "text": "private"}], + "encrypted_content": None, + } + message = { + "role": "assistant", + "content": [{"type": "output_text", "text": "kept"}], + } + + result = json.loads(sanitize_reasoning_replay(_body(model, [reasoning, message]))) + + assert "content" not in result["input"][0] + assert result["input"][0]["encrypted_content"] is None + assert result["input"][1] == message + + +def test_kimi_visible_reasoning_is_unchanged(): + body = _body( + "system.ai.kimi-k3", + [ + { + "type": "reasoning", + "content": [{"type": "reasoning_text", "text": "needed by Kimi"}], + } + ], + ) + + assert sanitize_reasoning_replay(body) is body + + +@pytest.mark.parametrize( + "body", + [ + b"not json", + b"[]", + b'{"model":"gpt-future","input":[]}', + b'{"model":"gpt-5.6-sol","input":"text"}', + b'{"model":"gpt-5.6-sol","input":[{"type":"reasoning","content":[]}]}', + b'{"model":"gpt-5.6-sol","input":[{"type":"reasoning"}]}', + ], +) +def test_unrelated_or_compatible_requests_are_byte_transparent(body): + assert sanitize_reasoning_replay(body) is body diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 58e72126..543e41c8 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -1,5 +1,6 @@ from __future__ import annotations +import contextlib import json import pytest @@ -36,6 +37,7 @@ def test_rejects_unsupported_codex_version(self, monkeypatch): def test_codex_launch_dispatches_when_flag_enabled(self, monkeypatch): calls = [] monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(codex.CODEX_RATE_LIMITER_ENV, "0") monkeypatch.setattr(codex, "default_model", lambda state: "gpt-start") monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) @@ -65,6 +67,7 @@ def launch_v2(state, tool_args, **kwargs): def test_codex_launch_normalizes_cached_bootstrap_model(self, monkeypatch): calls = [] monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(codex.CODEX_RATE_LIMITER_ENV, "0") monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) monkeypatch.setattr(codex, "default_model", lambda state: None) @@ -82,6 +85,33 @@ def launch_v2(state, tool_args, **kwargs): assert calls[0]["start_model"] == "gpt-5.6-luna" + def test_codex_launch_proxies_smart_routing_app_server(self, monkeypatch): + calls = [] + monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.delenv(codex.CODEX_RATE_LIMITER_ENV, raising=False) + monkeypatch.setattr(codex, "default_model", lambda state: "gpt-5.6-sol") + monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) + monkeypatch.setattr( + codex, + "_codex_request_proxy", + lambda state: contextlib.nullcontext("http://127.0.0.1:43210/v1"), + ) + + def launch_v2(state, tool_args, **kwargs): + overlay = kwargs["render_overlay"](WS, "gpt-5.6-sol") + calls.append(overlay) + raise SystemExit(0) + + monkeypatch.setattr(v2, "launch_codex", launch_v2) + + with pytest.raises(SystemExit) as exc: + codex.launch({"workspace": WS}, ["--search"]) + + assert exc.value.code == 0 + provider = calls[0]["model_providers"]["ucode-databricks"] + assert provider["base_url"] == "http://127.0.0.1:43210/v1" + assert calls[0]["features.enable_request_compression"] is False + def test_owns_app_server_interposer_and_tui_lifecycle(self, monkeypatch): processes = [] interposer_args = {} diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 28600045..b4e888de 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -328,9 +328,11 @@ class _FakeClient: def __init__(self, responses): self._responses = list(responses) self.sent_tokens: list[str | None] = [] + self.sent_bodies: list[bytes | None] = [] def stream(self, _method, _url, headers, content): self.sent_tokens.append(headers.get(gateway_proxy.AI_GATEWAY_TOKEN_HEADER)) + self.sent_bodies.append(content) return self._responses.pop(0) @@ -392,6 +394,24 @@ def flush(self): class TestRetryOn401: + def test_request_transform_and_gate_run_once_before_auth_retry(self): + body = b'{"model":"gpt-6-astra"}' + transformed = body + b" " + transformed_bodies = [] + gated = [] + client = _FakeClient([_FakeResp(401, b"a"), _FakeResp(200, b"ok")]) + handler = _handle_handler(client, _FakeCache(), _Collect()) + handler.headers = {"Content-Length": str(len(body))} + handler.rfile = io.BytesIO(body) + handler.request_transform = lambda value: transformed_bodies.append(value) or transformed + handler.request_gate = gated.append + + handler._handle() + + assert transformed_bodies == [body] + assert gated == [transformed] + assert client.sent_bodies == [transformed, transformed] + def test_401_forces_refresh_and_retries(self): # A stale swap token yields 401; the proxy force-refreshes and retries, # this time succeeding, so Claude Code never sees the 401. @@ -437,6 +457,58 @@ def test_failed_refresh_surfaces_reauth_hint(self, capsys): assert b"401" in bytes(out.data) # the response is still relayed +class TestRetryOn429: + def test_429_is_drained_waited_and_retried_inside_proxy(self): + body = b'{"model":"future-model","input":"hello"}' + first = _FakeResp(429, b'{"error":"slow down"}', {"Retry-After": "7"}) + client = _FakeClient([first, _FakeResp(200, b"ok")]) + out = _Collect() + handler = _handle_handler(client, _FakeCache(), out) + handler.headers = {"Content-Length": str(len(body))} + handler.rfile = io.BytesIO(body) + retries = [] + handler.rate_limit_retry = lambda seen_body, headers, attempt: retries.append( + (seen_body, headers, attempt) + ) + + handler._handle() + + assert first.read_called is True + assert client.sent_bodies == [body, body] + assert retries == [(body, {"Retry-After": "7"}, 1)] + assert b"429" not in bytes(out.data) + assert b"200" in bytes(out.data) + assert b"ok" in bytes(out.data) + + def test_repeated_429s_increase_attempt_without_spending_client_retries(self): + responses = [ + _FakeResp(429, b"one"), + _FakeResp(429, b"two"), + _FakeResp(200, b"ok"), + ] + client = _FakeClient(responses) + out = _Collect() + handler = _handle_handler(client, _FakeCache(), out) + attempts = [] + handler.rate_limit_retry = lambda _body, _headers, attempt: attempts.append(attempt) + + handler._handle() + + assert attempts == [1, 2] + assert len(client.sent_bodies) == 3 + assert b"429" not in bytes(out.data) + assert b"200" in bytes(out.data) + + def test_429_is_relayed_when_no_retry_handler_is_configured(self): + client = _FakeClient([_FakeResp(429, b"limited")]) + out = _Collect() + + _handle_handler(client, _FakeCache(), out)._handle() + + assert b"429" in bytes(out.data) + assert b"limited" in bytes(out.data) + + class TestStartProxyPortFallback: def test_falls_back_to_free_port_when_cached_port_busy(self, monkeypatch): # A stale proxy from a killed session can still hold the cached port; the diff --git a/tests/test_gateway_proxy_integration.py b/tests/test_gateway_proxy_integration.py index 2809ced0..89e8a860 100644 --- a/tests/test_gateway_proxy_integration.py +++ b/tests/test_gateway_proxy_integration.py @@ -47,10 +47,15 @@ class _FakeGateway: a streaming relay is exercised, not just a single write).""" def __init__( - self, status: int = 200, chunks: list[bytes] | None = None, sse_delay: float = 0.0 + self, + status: int = 200, + chunks: list[bytes] | None = None, + sse_delay: float = 0.0, + statuses: list[int] | None = None, ): self.requests: list[_CapturedRequest] = [] self._status = status + self._statuses = list(statuses) if statuses is not None else None self._chunks = chunks if chunks is not None else [b'{"ok":true}'] self._sse_delay = sse_delay self._server: ThreadingHTTPServer | None = None @@ -63,7 +68,13 @@ def base_url(self) -> str: def start(self) -> None: captured = self.requests - status, chunks, sse_delay = self._status, self._chunks, self._sse_delay + status, statuses, chunks, sse_delay = ( + self._status, + self._statuses, + self._chunks, + self._sse_delay, + ) + status_lock = threading.Lock() class Handler(BaseHTTPRequestHandler): def _serve(self) -> None: @@ -77,7 +88,9 @@ def _serve(self) -> None: body=body, ) ) - self.send_response(status) + with status_lock: + response_status = statuses.pop(0) if statuses else status + self.send_response(response_status) self.send_header("Content-Type", "text/event-stream") self.end_headers() for chunk in chunks: @@ -142,15 +155,29 @@ def fn(_workspace, _profile, force_refresh=False): @contextlib.contextmanager -def _running_proxy(gateway: _FakeGateway, monkeypatch, token_fn=None): +def _running_proxy( + gateway: _FakeGateway, + monkeypatch, + token_fn=None, + *, + token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER, + upstream_base=None, + request_transform=None, + request_gate=None, + rate_limit_retry=None, +): """Start the real proxy pointed at `gateway`, yield its loopback URL, tear down.""" monkeypatch.setattr(gateway_proxy, "get_databricks_token", token_fn or _counting_token()) server, cache, client = gateway_proxy.start_proxy( gateway.base_url, None, 0, - token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER, + token_header=token_header, force_refresh_near_expiry=False, + upstream_base=upstream_base, + request_transform=request_transform, + request_gate=request_gate, + rate_limit_retry=rate_limit_retry, ) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() @@ -165,6 +192,54 @@ def _running_proxy(gateway: _FakeGateway, monkeypatch, token_fn=None): class TestRelayedProxyEndToEnd: + def test_codex_upstream_replaces_authorization_and_gates_body(self, make_gateway, monkeypatch): + gw = make_gateway() + gated = [] + body = b'{"model":"gpt-6-astra","input":"hello"}' + transformed = body.replace(b"hello", b"sanitized") + with _running_proxy( + gw, + monkeypatch, + _counting_token("fresh-db-token"), + token_header=gateway_proxy.AUTHORIZATION_HEADER, + upstream_base=f"{gw.base_url}/ai-gateway/codex/", + request_transform=lambda _body: transformed, + request_gate=gated.append, + ) as proxy_url: + resp = httpx.post( + f"{proxy_url}/v1/responses", + headers={"Authorization": "Bearer stale-client-token"}, + content=body, + timeout=10, + ) + + assert resp.status_code == 200 + assert gated == [transformed] + request = gw.requests[-1] + assert request.path == "/ai-gateway/codex/v1/responses" + assert request.header("Authorization") == "Bearer fresh-db-token" + assert request.body == transformed + + def test_codex_429_waits_and_retries_over_real_proxy_sockets(self, make_gateway, monkeypatch): + gw = make_gateway(statuses=[429, 200]) + retries = [] + body = b'{"model":"future-model","input":"hello"}' + with _running_proxy( + gw, + monkeypatch, + token_header=gateway_proxy.AUTHORIZATION_HEADER, + upstream_base=f"{gw.base_url}/ai-gateway/codex/", + rate_limit_retry=lambda seen_body, headers, attempt: retries.append( + (seen_body, headers, attempt) + ), + ) as proxy_url: + resp = httpx.post(f"{proxy_url}/v1/responses", content=body, timeout=10) + + assert resp.status_code == 200 + assert len(gw.requests) == 2 + assert [request.body for request in gw.requests] == [body, body] + assert [(seen_body, attempt) for seen_body, _headers, attempt in retries] == [(body, 1)] + def test_forwards_request_with_swap_header_and_passthrough(self, make_gateway, monkeypatch): # The whole relayed data-plane over real sockets: the proxy injects a fresh # swap token, passes the caller's Anthropic OAuth + the MPS routing header