Skip to content

Commit 5dd5cd4

Browse files
author
Mateusz
committed
fix: codex managed OAuth bootstrap, tests, and pycycle dev tooling
- Short-circuit managed OAuth load when no account files; avoid rate-limit wait on bootstrap (credentials.py). - Point config tests at src.core.services.backend_registry for patches; align Codex cleanup tests with ResponseExecutor streaming path. - Add pycycle dev dependency, [tool.pycycle], dev/scripts/run_pycycle.py, and AGENTS command row. - Include Codex quota/selector/payload updates, demo script fix, and new quota notification unit tests. Left unstaged: dev/diag_cbor_requests.py (local diagnostic scratch). Made-with: Cursor
1 parent 23401ee commit 5dd5cd4

14 files changed

Lines changed: 343 additions & 144 deletions

AGENTS.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ When working on specs, the user will invoke `/kiro:*` commands. Follow the instr
8484
| **Lint/Fix** | `./.venv/Scripts/python.exe -m ruff check --fix .` |
8585
| **Format** | `./.venv/Scripts/python.exe -m black .` |
8686
| **Boundary Type Check** | `./.venv/Scripts/python.exe dev/scripts/check_boundary_types.py` |
87+
| **Circular import scan (pycycle)** | `./.venv/Scripts/python.exe dev/scripts/run_pycycle.py` |
8788
| **Inspect CBOR wire captures** | `./.venv/Scripts/python.exe scripts/inspect_cbor_capture.py <file> --detect-issues` |
8889

8990
## Quality & Testing Standards

dev/scripts/demo_gemini_api_key_loading_fix.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ def demo_full_integration(monkeypatch) -> None:
153153
return_value=["gemini"],
154154
),
155155
patch(
156-
"src.core.config.models.backends.backend_registry.get_registered_backends",
156+
"src.core.services.backend_registry.backend_registry.get_registered_backends",
157157
return_value=["gemini"],
158158
),
159159
):
@@ -232,7 +232,7 @@ def demo_single_base_key_only(monkeypatch) -> None:
232232
return_value=["gemini"],
233233
),
234234
patch(
235-
"src.core.config.models.backends.backend_registry.get_registered_backends",
235+
"src.core.services.backend_registry.backend_registry.get_registered_backends",
236236
return_value=["gemini"],
237237
),
238238
):
@@ -271,7 +271,7 @@ def demo_single_numbered_key_only(monkeypatch) -> None:
271271
return_value=["gemini"],
272272
),
273273
patch(
274-
"src.core.config.models.backends.backend_registry.get_registered_backends",
274+
"src.core.services.backend_registry.backend_registry.get_registered_backends",
275275
return_value=["gemini"],
276276
),
277277
):

dev/scripts/run_pycycle.py

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Run pycycle using settings from ``[tool.pycycle]`` in the repository ``pyproject.toml``.
4+
5+
Upstream ``pycycle --source <path>`` is unusable with current Click (the option is typed as a
6+
boolean), so this script invokes ``pycycle --here`` with the working directory set to
7+
``resolve_dir`` (typically ``src`` for this layout).
8+
9+
Usage::
10+
11+
./.venv/Scripts/python.exe dev/scripts/run_pycycle.py
12+
./.venv/Scripts/python.exe dev/scripts/run_pycycle.py --verbose
13+
"""
14+
15+
from __future__ import annotations
16+
17+
import argparse
18+
import shutil
19+
import subprocess
20+
import sys
21+
from collections.abc import Mapping
22+
from pathlib import Path
23+
from typing import Any
24+
25+
import tomli
26+
27+
28+
def _find_pyproject(start: Path) -> Path:
29+
for directory in (start, *start.parents):
30+
candidate = directory / "pyproject.toml"
31+
if candidate.is_file():
32+
return candidate
33+
msg = "pyproject.toml not found (expected an ancestor of this script)."
34+
raise SystemExit(msg)
35+
36+
37+
def _load_pycycle_config(pyproject_path: Path) -> dict[str, Any]:
38+
with pyproject_path.open("rb") as handle:
39+
data = tomli.load(handle)
40+
tool = data.get("tool") or {}
41+
if not isinstance(tool, Mapping):
42+
return {}
43+
pycycle = tool.get("pycycle") or {}
44+
if not isinstance(pycycle, Mapping):
45+
return {}
46+
return dict(pycycle)
47+
48+
49+
def _pycycle_executable() -> str:
50+
scripts_dir = Path(sys.executable).resolve().parent
51+
for name in ("pycycle.exe", "pycycle"):
52+
candidate = scripts_dir / name
53+
if candidate.is_file():
54+
return str(candidate)
55+
found = shutil.which("pycycle")
56+
if found:
57+
return found
58+
msg = (
59+
"pycycle executable not found. Install dev dependencies, for example: "
60+
"pip install -e .[dev]"
61+
)
62+
raise SystemExit(msg)
63+
64+
65+
def main() -> int:
66+
parser = argparse.ArgumentParser(description=__doc__)
67+
parser.add_argument(
68+
"--verbose",
69+
action="store_true",
70+
help="Verbose pycycle output (overrides [tool.pycycle].verbose when set).",
71+
)
72+
args = parser.parse_args()
73+
74+
script_path = Path(__file__).resolve()
75+
pyproject_path = _find_pyproject(script_path.parent)
76+
repo_root = pyproject_path.parent
77+
cfg = _load_pycycle_config(pyproject_path)
78+
79+
resolve_dir = str(cfg.get("resolve_dir") or "src")
80+
encoding = cfg.get("encoding")
81+
ignore = str(cfg.get("ignore") or "")
82+
verbose = bool(cfg.get("verbose")) or args.verbose
83+
84+
target = (repo_root / resolve_dir).resolve()
85+
try:
86+
target.relative_to(repo_root.resolve())
87+
except ValueError:
88+
msg = f"resolve_dir {resolve_dir!r} escapes repository root."
89+
raise SystemExit(msg)
90+
if not target.is_dir():
91+
msg = f"Analysis directory does not exist: {target}"
92+
raise SystemExit(msg)
93+
94+
cmd: list[str] = [_pycycle_executable(), "--here"]
95+
if encoding:
96+
cmd.extend(["--encoding", str(encoding)])
97+
stripped_ignore = ignore.strip()
98+
if stripped_ignore:
99+
cmd.extend(["--ignore", stripped_ignore])
100+
if verbose:
101+
cmd.append("--verbose")
102+
103+
completed = subprocess.run(cmd, cwd=target, check=False)
104+
return int(completed.returncode)
105+
106+
107+
if __name__ == "__main__":
108+
raise SystemExit(main())

pyproject.toml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ dev = [
9191
"respx",
9292
"dependency-injector",
9393
"vulture",
94+
"pycycle==0.0.8",
9495
"pytest-snapshot==0.9.0",
9596
"mypy==1.10.0",
9697
"hypothesis==6.112.1",
@@ -448,6 +449,16 @@ ignore_names = [
448449
"tool_name"
449450
]
450451

452+
# Circular import scan (bndr/pycycle). Upstream does not read this table; dev/scripts/run_pycycle.py does.
453+
# Note: `pycycle --source <path>` is broken with modern Click (option is inferred as boolean), so the
454+
# wrapper runs `pycycle --here` with cwd set to `resolve_dir`.
455+
[tool.pycycle]
456+
resolve_dir = "src"
457+
encoding = "utf-8"
458+
# Comma-separated directory basenames excluded during traversal (passed to pycycle --ignore).
459+
ignore = ""
460+
verbose = false
461+
451462
[tool.testmon]
452463
run_variant_expression = "','.join(sorted(str(v) for v in [sys.version_info[:2], sys.platform]))"
453464
historian = "git"

src/connectors/openai_codex/codex_quota_notifications.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@
2020

2121
_CODEX_QUOTA_NOTIFICATION_TITLE = "OpenAI Codex: Quota reached"
2222

23+
# Dedupe key component when every managed account is simultaneously unavailable:
24+
# one desktop alert per quota window, not one per account that hits 429 last.
25+
_CODEX_QUOTA_ALL_ACCOUNTS_DEDUPE_ID = "__all_managed_accounts__"
26+
2327
_EXTENDED_WINDOW = "extended (~weekly_or_plan_quota)"
2428

2529

@@ -96,7 +100,12 @@ async def maybe_notify_codex_quota_reached(
96100
until_display = until_iso if until_iso else "unknown"
97101
dedupe_until = until_iso if until_iso else "none"
98102

99-
key = (managed_account_id, quota_type, dedupe_until)
103+
dedupe_account_id = (
104+
_CODEX_QUOTA_ALL_ACCOUNTS_DEDUPE_ID
105+
if all_accounts_exhausted
106+
else managed_account_id
107+
)
108+
key = (dedupe_account_id, quota_type, dedupe_until)
100109
if key in dedupe_keys:
101110
return
102111

src/connectors/openai_codex/credentials.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -541,12 +541,21 @@ async def _load_managed_auth(self, force_reload: bool = False) -> bool:
541541
if not self._managed_enabled():
542542
return False
543543

544+
# Skip managed path when nothing is on disk — avoids selector/refresh work
545+
# and prevents long rate-limit polling loops during bootstrap (tests/CI).
546+
if not await self._managed_has_accounts():
547+
return False
548+
544549
if force_reload:
545550
await self._managed_selector.reload_accounts()
546551

547552
account = self._managed_selector.get_current_account()
548553
if account is None or account.needs_reauth:
549-
account = await self._managed_selector.get_next_account()
554+
# Initial bootstrap must not block on rate-limit recovery sleeps
555+
# (see ManagedOAuthAccountSelector.get_next_account); fall back to legacy.
556+
account = await self._managed_selector.get_next_account(
557+
wait_for_rate_limit_recovery=False,
558+
)
550559

551560
if account is None:
552561
return False
@@ -862,7 +871,10 @@ async def refresh_access_token(self) -> bool:
862871
async def _refresh_managed_access_token(self) -> bool:
863872
account = self._managed_selector.get_current_account()
864873
if account is None:
865-
account = await self._managed_selector.get_next_account()
874+
# Refresh path should also avoid blocking on rate-limit recovery sleeps.
875+
account = await self._managed_selector.get_next_account(
876+
wait_for_rate_limit_recovery=False,
877+
)
866878
if account is None:
867879
return False
868880

src/connectors/openai_codex/managed_oauth_models.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,8 @@ class ManagedOAuthConfig(BaseModel):
227227
#: when upstream sends very large ``resets_in_seconds``. Full upstream metadata is
228228
#: still stored on ``last_codex_usage_limit``.
229229
rate_limit_local_cooldown_cap_seconds: float = 1800.0
230-
#: Max idle polls (sleeps) while all accounts are rate-limited before giving up.
230+
#: Max idle polls (sleeps) while all accounts are rate-limited before giving up
231+
#: (only when ``get_next_account`` is invoked with ``wait_for_rate_limit_recovery=True``).
231232
max_rate_limit_idle_polls: int = 48
232233

233234
@classmethod

src/connectors/openai_codex/managed_oauth_selector.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -237,8 +237,17 @@ async def get_next_account(
237237
*,
238238
session_id: str | None = None,
239239
ignore_session_affinity: bool = False,
240+
wait_for_rate_limit_recovery: bool = True,
240241
) -> ManagedOAuthAccount | None:
241-
"""Return next usable account and perform proactive refresh."""
242+
"""Return next usable account and perform proactive refresh.
243+
244+
When every available account is currently rate-limited, the default
245+
behaviour is to sleep (bounded by ``max_rate_limit_wait_seconds``) and
246+
poll until one becomes eligible again or ``max_rate_limit_idle_polls``
247+
is exceeded. Callers that must not block in-flight requests (for example
248+
:meth:`rotate_on_rate_limit`) should pass ``wait_for_rate_limit_recovery=False``
249+
to return ``None`` immediately in that situation.
250+
"""
242251
await self._ensure_accounts_loaded()
243252

244253
rate_limit_idle_polls = 0
@@ -250,6 +259,8 @@ async def get_next_account(
250259
return None
251260

252261
if not eligible:
262+
if not wait_for_rate_limit_recovery:
263+
return None
253264
rate_limit_idle_polls += 1
254265
if rate_limit_idle_polls > self._max_rate_limit_idle_polls:
255266
return None
@@ -381,6 +392,7 @@ async def rotate_on_rate_limit(
381392
return await self.get_next_account(
382393
session_id=session_id,
383394
ignore_session_affinity=True,
395+
wait_for_rate_limit_recovery=False,
384396
)
385397

386398
async def rotate_on_auth_failure(
@@ -398,4 +410,5 @@ async def rotate_on_auth_failure(
398410
return await self.get_next_account(
399411
session_id=session_id,
400412
ignore_session_affinity=True,
413+
wait_for_rate_limit_recovery=False,
401414
)

src/connectors/openai_codex/payload.py

Lines changed: 6 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
from __future__ import annotations
77

8-
import json
98
import logging
109
import uuid
1110
from copy import deepcopy
@@ -665,75 +664,12 @@ def convert_dict_to_payload(
665664
@staticmethod
666665
def _sanitize_responses_input(input_value: Any) -> list[dict[str, Any] | Any]:
667666
"""Make a Responses `input` array safe for ChatGPT Codex backend.
668-
669-
- Removes `item_reference` entries (AI SDK/OpenCode server-state references)
670-
- Strips per-item `id` fields for stateless mode (`store: false`)
671-
- Removes unsupported per-item `metadata` blocks
672-
- Converts orphaned `function_call_output` entries into assistant messages
673-
to preserve context while avoiding backend validation errors
667+
668+
Note: We must NOT strip `id`, `metadata`, or `item_reference` fields, nor
669+
modify `function_call_output` entries. The Codex backend relies on the exact
670+
shape of these messages (including `id`s) to perform context (write) caching.
671+
Stripping them causes cache misses and rapid quota exhaustion.
674672
"""
675673
if not isinstance(input_value, list):
676674
return []
677-
678-
filtered: list[dict[str, Any]] = []
679-
for item in input_value:
680-
if not isinstance(item, dict):
681-
continue
682-
683-
item_type = item.get("type")
684-
if item_type == "item_reference":
685-
continue
686-
687-
item_dict = dict(item)
688-
item_dict.pop("id", None)
689-
item_dict.pop("metadata", None)
690-
filtered.append(item_dict)
691-
692-
function_call_ids: set[str] = set()
693-
for item in filtered:
694-
if item.get("type") == "function_call":
695-
call_id = item.get("call_id")
696-
if isinstance(call_id, str) and call_id:
697-
function_call_ids.add(call_id)
698-
699-
safe: list[dict[str, Any]] = []
700-
for item in filtered:
701-
if item.get("type") == "function_call_output":
702-
call_id = item.get("call_id")
703-
if (
704-
isinstance(call_id, str)
705-
and call_id
706-
and call_id not in function_call_ids
707-
):
708-
tool_name = item.get("name")
709-
if not isinstance(tool_name, str) or not tool_name:
710-
tool_name = "tool"
711-
output_val = item.get("output")
712-
if isinstance(output_val, str):
713-
output_text = output_val
714-
else:
715-
try:
716-
output_text = json.dumps(output_val)
717-
except Exception:
718-
output_text = str(output_val)
719-
720-
if len(output_text) > 16000:
721-
output_text = output_text[:16000] + "\n...[truncated]"
722-
723-
safe.append(
724-
{
725-
"type": "message",
726-
"role": "assistant",
727-
"content": [
728-
{
729-
"type": "output_text",
730-
"text": f"[Previous {tool_name} result; call_id={call_id}]: {output_text}",
731-
}
732-
],
733-
}
734-
)
735-
continue
736-
737-
safe.append(item)
738-
739-
return safe
675+
return [dict(item) if isinstance(item, dict) else item for item in input_value]

0 commit comments

Comments
 (0)