commit fc4fcbab58d00651f7c3f21bb68c3a550bd2dba8 Author: wehub-resource-sync Date: Mon Jul 13 12:22:33 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.agents/skills/codexbar/SKILL.md b/.agents/skills/codexbar/SKILL.md new file mode 100644 index 0000000..abcc6e3 --- /dev/null +++ b/.agents/skills/codexbar/SKILL.md @@ -0,0 +1,38 @@ +--- +name: codexbar +description: "CodexBar read. Provider usage, limits, credits, config health. JSON. No writes." +--- + +# CodexBar + +Read CodexBar. Never mutate config/auth. + +## Run + +```bash +skill="${CODEX_HOME:-$HOME/.codex}/skills/codexbar" +"$skill/scripts/codexbar" doctor +"$skill/scripts/codexbar" providers +"$skill/scripts/codexbar" usage +"$skill/scripts/codexbar" usage --provider codex +"$skill/scripts/codexbar" usage --all +``` + +All stdout: JSON. Upstream CodexBar shape kept. Less drift, fewer tokens. + +## Rules + +- Start `doctor` when install/config unknown. +- `usage` reads enabled providers. Prefer this. +- `usage --provider ID` reads one provider. +- `usage --all` expensive; use only when needed. +- Identities hidden by default. `--include-identities` only when user explicitly needs them. +- Secrets always hidden. +- Helper read-only: fixed allowlist only. No config writes, auth repair, enable/disable, key storage. +- Timeout means upstream stuck. Narrow provider or raise `CODEXBAR_TIMEOUT` (default 120 seconds). + +## Binary + +Auto-find: `CODEXBAR_BIN`, PATH, app bundle, Homebrew cask. If missing: open CodexBar, Preferences > Advanced > Install CLI; or set `CODEXBAR_BIN`. + +Each stdout/stderr stream capped at 1 MiB while fully drained. Timeout kills process group. diff --git a/.agents/skills/codexbar/scripts/codexbar b/.agents/skills/codexbar/scripts/codexbar new file mode 100755 index 0000000..d60c8d4 --- /dev/null +++ b/.agents/skills/codexbar/scripts/codexbar @@ -0,0 +1,321 @@ +#!/usr/bin/env python3 +"""Read-only, bounded CodexBar CLI bridge.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +import signal +import subprocess +import sys +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + + +MAX_CAPTURE_BYTES = 1024 * 1024 +DEFAULT_TIMEOUT = 120.0 +BIN_ENV = "CODEXBAR_BIN" +TIMEOUT_ENV = "CODEXBAR_TIMEOUT" +SKIP_DISCOVERY_ENV = "CODEXBAR_SKIP_DISCOVERY" + +SECRET = "" +IDENTITY = "" +EMAIL = "" + +EMAIL_RE = re.compile(r"\b[A-Z0-9._%+\-]+@[A-Z0-9.\-]+\.[A-Z]{2,}\b", re.I) +BEARER_RE = re.compile(r"(?i)\bbearer\s+[A-Z0-9._~+/=\-]+") +JWT_RE = re.compile(r"\beyJ[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\b") +LABELED_SECRET_RE = re.compile( + r"(?i)\b(authorization|api[_-]?key|access[_-]?token|refresh[_-]?token|id[_-]?token|" + r"session(?:id)?|secret|password|passwd|cookie|set-cookie)(\s*[:=]\s*)([^\s,;]+)" +) +WORD_SECRET_RE = re.compile( + r"(?i)\b(token|secret|password|passwd|cookie)\s+([A-Z0-9._~+/=\-]{6,})\b" +) +SECRET_KEY_RE = re.compile( + r"(?i)(authorization|api[_-]?key|access[_-]?token|refresh[_-]?token|id[_-]?token|" + r"session(?:id)?|secret|password|passwd|cookie)" +) +IDENTITY_KEYS = { + "accountemail", + "accountorganization", + "accountid", + "accountname", + "email", + "organization", + "userid", + "username", +} +# ProviderIdentitySnapshot.providerID is UsageProvider, not an account identifier. + + +@dataclass(frozen=True) +class Binary: + path: str + source: str + + +@dataclass(frozen=True) +class Result: + returncode: int + stdout: bytes + stderr: bytes + timed_out: bool + + +def normalized_key(value: str) -> str: + return re.sub(r"[^a-z0-9]", "", value.lower()) + + +def timeout_seconds() -> float: + try: + return max(1.0, float(os.environ.get(TIMEOUT_ENV, DEFAULT_TIMEOUT))) + except ValueError: + return DEFAULT_TIMEOUT + + +def safe_path(path: str) -> str: + home = str(Path.home()) + if path == home: + return "~" + if path.startswith(home + os.sep): + return "~" + path[len(home) :] + return re.sub(r"^/Users/[^/]+", "/Users/", path) + + +def candidates() -> list[tuple[str, Path]]: + found: list[tuple[str, Path]] = [] + override = os.environ.get(BIN_ENV) + if override: + found.append(("env", Path(override).expanduser())) + + path_hit = shutil.which("codexbar") + if path_hit: + found.append(("path", Path(path_hit))) + + if os.environ.get(SKIP_DISCOVERY_ENV) == "1": + return found + + home = Path.home() + found.extend( + [ + ("homebrew", Path("/opt/homebrew/bin/codexbar")), + ("homebrew", Path("/usr/local/bin/codexbar")), + ("app", Path("/Applications/CodexBar.app/Contents/Helpers/CodexBarCLI")), + ("app", home / "Applications/CodexBar.app/Contents/Helpers/CodexBarCLI"), + ] + ) + for root in (Path("/opt/homebrew/Caskroom/codexbar"), Path("/usr/local/Caskroom/codexbar")): + found.extend(("cask", app / "Contents/Helpers/CodexBarCLI") for app in sorted(root.glob("*/CodexBar.app"))) + return found + + +def resolve_binary() -> Binary | None: + self_path = Path(__file__).resolve() + seen: set[Path] = set() + for source, candidate in candidates(): + try: + resolved = candidate.resolve() + if resolved in seen or resolved == self_path: + continue + seen.add(resolved) + if resolved.is_file() and os.access(resolved, os.X_OK): + return Binary(str(resolved), source) + except OSError: + continue + return None + + +def drain(pipe: Any, target: bytearray) -> None: + try: + while True: + chunk = pipe.read(65536) + if not chunk: + break + room = MAX_CAPTURE_BYTES - len(target) + if room > 0: + target.extend(chunk[:room]) + finally: + pipe.close() + + +def stop_group(process: subprocess.Popen[bytes], sig: signal.Signals) -> None: + try: + os.killpg(process.pid, sig) + except ProcessLookupError: + pass + + +def run_process(argv: Sequence[str]) -> Result: + process = subprocess.Popen( + list(argv), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + start_new_session=True, + ) + assert process.stdout is not None + assert process.stderr is not None + stdout = bytearray() + stderr = bytearray() + readers = [ + threading.Thread(target=drain, args=(process.stdout, stdout), daemon=True), + threading.Thread(target=drain, args=(process.stderr, stderr), daemon=True), + ] + for reader in readers: + reader.start() + + timed_out = False + try: + process.wait(timeout=timeout_seconds()) + except subprocess.TimeoutExpired: + timed_out = True + stop_group(process, signal.SIGTERM) + try: + process.wait(timeout=2) + except subprocess.TimeoutExpired: + stop_group(process, signal.SIGKILL) + process.wait() + + for reader in readers: + reader.join(timeout=3) + return Result(124 if timed_out else process.returncode, bytes(stdout), bytes(stderr), timed_out) + + +def decode(value: bytes) -> str: + return value.decode("utf-8", errors="replace") + + +def sanitize_text(value: str, include_identities: bool) -> str: + value = JWT_RE.sub(SECRET, value) + value = BEARER_RE.sub("Bearer " + SECRET, value) + value = LABELED_SECRET_RE.sub(lambda match: match.group(1) + match.group(2) + SECRET, value) + value = WORD_SECRET_RE.sub(lambda match: match.group(1) + " " + SECRET, value) + if not include_identities: + value = EMAIL_RE.sub(EMAIL, value) + return value + + +def sanitize(value: Any, include_identities: bool, key: str | None = None) -> Any: + if key and SECRET_KEY_RE.search(key): + return SECRET + if key and normalized_key(key) in IDENTITY_KEYS and not include_identities and value is not None: + return EMAIL if "email" in normalized_key(key) else IDENTITY + if isinstance(value, dict): + return {str(child_key): sanitize(child, include_identities, str(child_key)) for child_key, child in value.items()} + if isinstance(value, list): + return [sanitize(child, include_identities) for child in value] + if isinstance(value, str): + return sanitize_text(value, include_identities) + return value + + +def print_json(value: Any) -> None: + print(json.dumps(value, indent=2, sort_keys=True)) + + +def print_error(kind: str, message: str, *, binary: Binary | None = None) -> None: + payload: dict[str, Any] = {"error": {"kind": kind, "message": sanitize_text(message, False)}} + if binary: + payload["binary"] = {"path": safe_path(binary.path), "source": binary.source} + print_json(payload) + + +def emit_stderr(result: Result, include_identities: bool) -> None: + text = sanitize_text(decode(result.stderr), include_identities).strip() + if text: + print(text, file=sys.stderr) + + +def read_json(binary: Binary, argv: Sequence[str], include_identities: bool) -> int: + result = run_process([binary.path, *argv]) + emit_stderr(result, include_identities) + if result.timed_out: + print_error("timeout", f"CodexBar exceeded {timeout_seconds():g}s and was stopped.", binary=binary) + return 124 + try: + payload = json.loads(decode(result.stdout)) + except json.JSONDecodeError as error: + preview = sanitize_text(decode(result.stdout[:400]), False) + print_error("invalid_json", f"CodexBar JSON failed: {error.msg}. stdout={preview!r}", binary=binary) + return result.returncode or 1 + print_json(sanitize(payload, include_identities)) + return result.returncode + + +def doctor(binary: Binary, include_identities: bool) -> int: + version = run_process([binary.path, "--version"]) + if version.timed_out: + print_error("timeout", f"CodexBar version exceeded {timeout_seconds():g}s.", binary=binary) + return 124 + validation = run_process([binary.path, "config", "validate", "--format", "json", "--json-only"]) + emit_stderr(version, include_identities) + emit_stderr(validation, include_identities) + if validation.timed_out: + print_error("timeout", f"CodexBar config validation exceeded {timeout_seconds():g}s.", binary=binary) + return 124 + try: + issues = json.loads(decode(validation.stdout)) + except json.JSONDecodeError as error: + print_error("invalid_json", f"CodexBar config validation failed: {error.msg}.", binary=binary) + return validation.returncode or 1 + print_json( + { + "binary": {"path": safe_path(binary.path), "source": binary.source}, + "configIssues": sanitize(issues, include_identities), + "version": sanitize_text((decode(version.stdout) or decode(version.stderr)).strip(), include_identities), + } + ) + return version.returncode or validation.returncode + + +def parser() -> argparse.ArgumentParser: + root = argparse.ArgumentParser(prog="codexbar", description="CodexBar read. JSON out. No writes.") + commands = root.add_subparsers(dest="command", required=True) + for name in ("doctor", "providers"): + command = commands.add_parser(name) + command.add_argument("--include-identities", action="store_true") + usage = commands.add_parser("usage") + scope = usage.add_mutually_exclusive_group() + scope.add_argument("--all", action="store_true") + scope.add_argument("--provider") + usage.add_argument("--include-identities", action="store_true") + return root + + +def main(argv: Sequence[str] | None = None) -> int: + args = parser().parse_args(argv) + binary = resolve_binary() + if not binary: + print_error("missing", "CodexBar CLI not found. Install CLI in CodexBar Advanced settings or set CODEXBAR_BIN.") + return 1 + try: + if args.command == "doctor": + return doctor(binary, args.include_identities) + if args.command == "providers": + return read_json( + binary, + ["config", "providers", "--format", "json", "--json-only"], + args.include_identities, + ) + usage_args = ["usage", "--format", "json", "--json-only"] + if args.all: + usage_args.extend(["--provider", "all"]) + elif args.provider: + usage_args.extend(["--provider", args.provider]) + return read_json(binary, usage_args, args.include_identities) + except OSError as error: + print_error("launch", str(error), binary=binary) + return 1 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except KeyboardInterrupt: + print_error("interrupted", "Interrupted.") + raise SystemExit(130) diff --git a/.agents/skills/codexbar/scripts/test_codexbar.py b/.agents/skills/codexbar/scripts/test_codexbar.py new file mode 100644 index 0000000..ceb9dc7 --- /dev/null +++ b/.agents/skills/codexbar/scripts/test_codexbar.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +import json +import os +import runpy +import subprocess +import sys +import tempfile +import textwrap +import time +import unittest +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("codexbar") +MODULE = runpy.run_path(str(SCRIPT), run_name="codexbar_skill") +MAX_CAPTURE_BYTES = MODULE["MAX_CAPTURE_BYTES"] +run_process = MODULE["run_process"] + +FAKE = textwrap.dedent( + """\ + #!/usr/bin/env python3 + import json + import os + import subprocess + import sys + import time + + argv = sys.argv[1:] + log = os.environ.get("FAKE_LOG") + if log: + with open(log, "a", encoding="utf-8") as handle: + handle.write(json.dumps(argv) + "\\n") + + if argv == ["--version"]: + key = "VERSION" + elif argv[:2] == ["config", "validate"]: + key = "VALIDATE" + elif argv[:2] == ["config", "providers"]: + key = "PROVIDERS" + elif argv and argv[0] == "usage": + key = "USAGE" + else: + key = "OTHER" + + if os.environ.get("FAKE_SPAWN_CHILD") == "1": + marker = os.environ["FAKE_MARKER"] + subprocess.Popen([ + sys.executable, + "-c", + f"import pathlib,time; time.sleep(2); pathlib.Path({marker!r}).write_text('alive')", + ]) + time.sleep(5) + + delay = os.environ.get(f"FAKE_{key}_DELAY") + if delay: + time.sleep(float(delay)) + stdout = os.environ.get(f"FAKE_{key}_STDOUT", "") + stderr = os.environ.get(f"FAKE_{key}_STDERR", "") + sys.stdout.write(stdout) + sys.stderr.write(stderr) + raise SystemExit(int(os.environ.get(f"FAKE_{key}_EXIT", "0"))) + """ +) + + +def make_env(root: Path, *, install: bool = True) -> tuple[dict[str, str], Path]: + binary = root / "CodexBar.app" / "Contents" / "Helpers" / "CodexBarCLI" + binary.parent.mkdir(parents=True) + binary.write_text(FAKE, encoding="utf-8") + binary.chmod(0o755) + env = os.environ.copy() + env["CODEXBAR_SKIP_DISCOVERY"] = "1" + env["CODEXBAR_TIMEOUT"] = "3" + env["FAKE_LOG"] = str(root / "argv.log") + env["FAKE_VERSION_STDOUT"] = "CodexBar 1.2.3\n" + env["FAKE_VALIDATE_STDOUT"] = "[]" + env["FAKE_PROVIDERS_STDOUT"] = json.dumps( + [{"provider": "codex", "displayName": "Codex", "enabled": True}] + ) + env["FAKE_USAGE_STDOUT"] = json.dumps( + [ + { + "provider": "codex", + "source": "oauth", + "usage": { + "accountEmail": "alice@example.com", + "accountOrganization": "Example Org", + "identity": {"providerID": "codex", "accountID": "acct-123"}, + "primary": {"usedPercent": 42, "windowMinutes": 300}, + }, + } + ] + ) + if install: + env["CODEXBAR_BIN"] = str(binary) + else: + env.pop("CODEXBAR_BIN", None) + env["PATH"] = "/usr/bin:/bin:/usr/sbin:/sbin" + return env, binary + + +def helper(*args: str, env: dict[str, str], cwd: Path | None = None) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(SCRIPT), *args], + capture_output=True, + text=True, + check=False, + env=env, + cwd=cwd, + ) + + +class CodexBarSkillTests(unittest.TestCase): + def test_help_is_small_and_read_only(self) -> None: + result = helper("--help", env=os.environ.copy()) + self.assertEqual(result.returncode, 0) + self.assertIn("CodexBar read. JSON out. No writes.", result.stdout) + self.assertNotIn("enable", result.stdout) + self.assertNotIn("set-api-key", result.stdout) + + def test_missing_binary_is_json(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + env, _ = make_env(Path(tmp), install=False) + result = helper("doctor", env=env) + self.assertEqual(result.returncode, 1) + self.assertEqual(json.loads(result.stdout)["error"]["kind"], "missing") + + def test_doctor_reports_version_and_raw_validation_shape(self) -> None: + with tempfile.TemporaryDirectory(dir=Path.home()) as tmp: + env, binary = make_env(Path(tmp)) + result = helper("doctor", env=env) + self.assertEqual(result.returncode, 0) + payload = json.loads(result.stdout) + self.assertEqual(payload["version"], "CodexBar 1.2.3") + self.assertEqual(payload["configIssues"], []) + self.assertEqual(payload["binary"]["path"], "~" + str(binary)[len(str(Path.home())) :]) + + def test_providers_passes_upstream_json_without_second_schema(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + env, _ = make_env(Path(tmp)) + result = helper("providers", env=env) + self.assertEqual(result.returncode, 0) + self.assertEqual( + json.loads(result.stdout), + [{"provider": "codex", "displayName": "Codex", "enabled": True}], + ) + + def test_usage_defaults_to_enabled_and_runs_from_any_cwd(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + env, _ = make_env(root) + elsewhere = root / "elsewhere" + elsewhere.mkdir() + result = helper("usage", env=env, cwd=elsewhere) + calls = [json.loads(line) for line in (root / "argv.log").read_text().splitlines()] + self.assertEqual(result.returncode, 0) + self.assertEqual(calls, [["usage", "--format", "json", "--json-only"]]) + + def test_usage_scope_maps_to_upstream_cli(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + env, _ = make_env(root) + self.assertEqual(helper("usage", "--all", env=env).returncode, 0) + self.assertEqual(helper("usage", "--provider", "zai", env=env).returncode, 0) + calls = [json.loads(line) for line in (root / "argv.log").read_text().splitlines()] + self.assertEqual(calls[0][-2:], ["--provider", "all"]) + self.assertEqual(calls[1][-2:], ["--provider", "zai"]) + + def test_usage_hides_identity_but_preserves_provider_id(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + env, _ = make_env(Path(tmp)) + env["FAKE_USAGE_STDERR"] = ( + "Authorization: Bearer secret-token alice@example.com; token sk-live-prose-secret\n" + ) + result = helper("usage", env=env) + payload = json.loads(result.stdout)[0] + self.assertEqual(payload["provider"], "codex") + self.assertEqual(payload["usage"]["identity"]["providerID"], "codex") + self.assertEqual(payload["usage"]["accountEmail"], "") + self.assertEqual(payload["usage"]["accountOrganization"], "") + self.assertEqual(payload["usage"]["identity"]["accountID"], "") + self.assertNotIn("secret-token", result.stderr) + self.assertNotIn("sk-live-prose-secret", result.stderr) + self.assertIn("", result.stderr) + + def test_include_identities_never_exposes_secret_keys(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + env, _ = make_env(Path(tmp)) + payload = json.loads(env["FAKE_USAGE_STDOUT"]) + payload[0]["usage"]["apiKey"] = "super-secret" + env["FAKE_USAGE_STDOUT"] = json.dumps(payload) + result = helper("usage", "--include-identities", env=env) + usage = json.loads(result.stdout)[0]["usage"] + self.assertEqual(usage["accountEmail"], "alice@example.com") + self.assertEqual(usage["apiKey"], "") + + def test_each_stream_is_capped_while_fully_drained(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + writer = root / "writer" + writer.write_text( + "#!/usr/bin/env python3\n" + "import sys\n" + "sys.stdout.buffer.write(b'o' * 2000000)\n" + "sys.stderr.buffer.write(b'e' * 2000000)\n", + encoding="utf-8", + ) + writer.chmod(0o755) + result = run_process([str(writer)]) + self.assertEqual(result.returncode, 0) + self.assertEqual(len(result.stdout), MAX_CAPTURE_BYTES) + self.assertEqual(len(result.stderr), MAX_CAPTURE_BYTES) + + def test_timeout_kills_process_group(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + env, _ = make_env(root) + marker = root / "child-survived" + env["CODEXBAR_TIMEOUT"] = "1" + env["FAKE_SPAWN_CHILD"] = "1" + env["FAKE_MARKER"] = str(marker) + result = helper("usage", env=env) + time.sleep(2.2) + self.assertEqual(result.returncode, 124) + self.assertEqual(json.loads(result.stdout)["error"]["kind"], "timeout") + self.assertFalse(marker.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/.agents/skills/qa-test/SKILL.md b/.agents/skills/qa-test/SKILL.md new file mode 100644 index 0000000..81c752a --- /dev/null +++ b/.agents/skills/qa-test/SKILL.md @@ -0,0 +1,118 @@ +--- +name: qa-test +description: "CodexBar live QA/e2e testing: run provider usage matrix checks, validate real app config, use Peekaboo for menu proof, use Browser Use/official docs for API spec or logged-in dashboard checks, and handle 1Password credentials safely." +--- + +# CodexBar Live QA + +Use for live provider testing, release smoke tests, menu verification, or debugging “provider works/fails” reports. + +## Rules + +- Work from the CodexBar repo checkout. +- Use the packaged CLI first: `CodexBar.app/Contents/Helpers/CodexBarCLI`. +- Do not use `CodexBar.app/Contents/MacOS/codexbar`; that is the app binary and may appear to hang as a CLI. +- Never run broad `env`, `set`, or secret regex dumps. +- Use `$one-password` for secrets: all `op` commands inside one persistent tmux session, service account first, no raw secret output. +- Treat browser-cookie/keychain flows as prompt-risky. Prefer CLI/API-token checks and `KeychainNoUIQuery`-safe tests unless the user explicitly requested live UI. +- For current API behavior, browse official provider docs only. + +## CLI Matrix + +Run the bundled script: + +```bash +.agents/skills/qa-test/scripts/live_provider_matrix.sh --enabled +``` + +Useful modes: + +```bash +.agents/skills/qa-test/scripts/live_provider_matrix.sh --provider all +.agents/skills/qa-test/scripts/live_provider_matrix.sh --providers openai,zai,deepseek +.agents/skills/qa-test/scripts/live_provider_matrix.sh --default +``` + +Interpretation: + +- `--enabled` asks `CodexBarCLI config providers` for enabled providers, honoring `CODEXBAR_CONFIG` and default toggles. +- `--default` runs the app-facing default command with no provider override. +- `--provider all` forces every registered provider and is expected to fail for providers without sessions/keys. +- A green app config needs `--enabled` and `--default` clean; `--provider all` is a discovery/triage tool. + +## Config QA + +Validate config: + +```bash +CodexBar.app/Contents/Helpers/CodexBarCLI config validate +stat -f '%Lp %N' "$HOME/.codexbar/config.json" +``` + +Redact config shape: + +```bash +jq '(.providers // []) |= map(.apiKey = (if .apiKey then "" else .apiKey end) | + .secretKey = (if .secretKey then "" else .secretKey end) | + .cookieHeader = (if .cookieHeader then "" else .cookieHeader end) | + (if .id == "stepfun" and has("region") then .region = "" else . end) | + .tokenAccounts = (if .tokenAccounts then (.tokenAccounts | .accounts = (.accounts | map(.token = ""))) else .tokenAccounts end))' \ + "$HOME/.codexbar/config.json" +``` + +Before editing config, make a backup: + +```bash +cp "$HOME/.codexbar/config.json" "$HOME/.codexbar/config.pre-qa-$(date +%Y%m%d%H%M%S).json" +chmod 600 "$HOME/.codexbar"/config.pre-qa-*.json +``` + +## Live Menu QA + +Use Peekaboo after CLI checks: + +```bash +pkill -x CodexBar || pkill -f 'CodexBar.app/Contents/MacOS/CodexBar' || true +open -n "$PWD/CodexBar.app" +peekaboo menu list-all --json | rg -i 'codexbar' +peekaboo menu click-extra --title codexbar-merged --json +screencapture -x /tmp/codexbar-live-menu.png +``` + +Crop top-right menu if needed: + +```bash +sips --cropToHeightWidth 900 340 --cropOffset 20 2650 /tmp/codexbar-live-menu.png \ + --out /tmp/codexbar-live-menu-crop.png >/dev/null +``` + +Verify visually with `view_image`. Confirm provider tabs/rows match enabled config and no failing provider dominates the first screen. + +## Browser Use + +Use `$browser-use` only when a logged-in dashboard, API key page, or provider docs need browser/profile state. + +Existing Chrome path: + +```bash +mcporter call chrome-devtools.list_pages --args '{}' --output text +mcporter call chrome-devtools.navigate_page --args '{"url":"https://provider.example"}' --output text +mcporter call chrome-devtools.take_snapshot --args '{}' --output text +``` + +If Browser Use is unavailable, say so and use web search for public official docs; do not substitute isolated Playwright for login/profile-dependent pages. + +## Fix Triage + +- Missing auth/session: configure key/session if available; otherwise leave provider disabled or report blocked auth. +- Wrong provider API/spec: inspect official docs, then patch fetcher/settings/tests. +- Provider key exists but live API rejects it: keep key stored if useful, disable provider if the menu would show a persistent error. +- User-facing behavior changes need `CHANGELOG.md`. +- Code fixes need focused tests, `make check`, `$autoreview`, and live CLI proof before landing. + +## Known CodexBar QA Notes + +- OpenAI Admin API key is the useful usage provider key. Project `OPENAI_API_KEY` values can fail legacy credit-balance fallback with 403. +- Deepgram usage requires a key/project with Management API permissions; transcription-only keys can return 403. +- Groq usage uses the Prometheus metrics API, not ordinary inference endpoints. +- MiniMax pay-as-you-go API keys and Token Plan/Coding Plan keys are different; wrong key kind can leave usage unavailable. diff --git a/.agents/skills/qa-test/agents/openai.yaml b/.agents/skills/qa-test/agents/openai.yaml new file mode 100644 index 0000000..3bf7a7b --- /dev/null +++ b/.agents/skills/qa-test/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "CodexBar QA Test" + short_description: "Run live CodexBar CLI and menu QA safely." + default_prompt: "Run CodexBar live QA with CLI, Peekaboo, browser docs, and 1Password-safe credential checks." diff --git a/.agents/skills/qa-test/references/api-specs.md b/.agents/skills/qa-test/references/api-specs.md new file mode 100644 index 0000000..8bb7298 --- /dev/null +++ b/.agents/skills/qa-test/references/api-specs.md @@ -0,0 +1,10 @@ +# API Spec Pointers + +Use current official docs for provider API behavior. Prefer these searches/pages before patching fetchers: + +- MiniMax: `https://platform.minimax.io/docs/llms.txt`; key types differ between pay-as-you-go API keys and Token Plan/Coding Plan keys. +- Deepgram: `https://developers.deepgram.com/llms.txt`; usage/project APIs require Management permissions and project-scoped keys. +- Groq: `https://console.groq.com/docs/prometheus-metrics`; usage metrics use `https://api.groq.com/v1/metrics/prometheus`. +- LLM Proxy/LiteLLM: `https://docs.litellm.ai/`; CodexBar expects an LLM-API-Key-Proxy compatible `/v1/quota-stats` endpoint plus base URL. + +When citing docs in a user-facing answer, browse the current page and include source links. diff --git a/.agents/skills/qa-test/scripts/live_provider_matrix.sh b/.agents/skills/qa-test/scripts/live_provider_matrix.sh new file mode 100755 index 0000000..0c8240c --- /dev/null +++ b/.agents/skills/qa-test/scripts/live_provider_matrix.sh @@ -0,0 +1,183 @@ +#!/usr/bin/env bash +set -u + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" +CLI="${CODEXBAR_CLI:-$ROOT/CodexBar.app/Contents/Helpers/CodexBarCLI}" +TIMEOUT_BIN="${TIMEOUT_BIN:-$(command -v gtimeout || command -v timeout || true)}" +WEB_TIMEOUT="${CODEXBAR_QA_WEB_TIMEOUT:-12}" +CASE_TIMEOUT="${CODEXBAR_QA_CASE_TIMEOUT:-60}" + +usage() { + cat <<'USAGE' +Usage: + live_provider_matrix.sh --enabled + live_provider_matrix.sh --default + live_provider_matrix.sh --provider all + live_provider_matrix.sh --providers openai,zai,deepseek + +Environment: + CODEXBAR_CLI=/path/to/CodexBarCLI + CODEXBAR_CONFIG=/path/to/config.json + CODEXBAR_QA_WEB_TIMEOUT=12 + CODEXBAR_QA_CASE_TIMEOUT=60 +USAGE +} + +if [[ ! -x "$CLI" ]]; then + echo "missing CodexBarCLI at $CLI" >&2 + exit 2 +fi +if [[ -z "$TIMEOUT_BIN" ]]; then + echo "missing timeout command (install coreutils for gtimeout)" >&2 + exit 2 +fi +if ! command -v node >/dev/null 2>&1; then + echo "missing node" >&2 + exit 2 +fi + +mode="${1:-}" +shift || true + +providers=() +case "$mode" in + --enabled) + provider_status="$(mktemp)" + provider_err="$(mktemp)" + provider_list="$(mktemp)" + if ! "$CLI" config providers --format json --json-only >"$provider_status" 2>"$provider_err"; then + rm -f "$provider_status" "$provider_err" "$provider_list" + echo "failed to list providers via CodexBarCLI config providers" >&2 + exit 2 + fi + if ! node - "$provider_status" >"$provider_list" <<'NODE'; then +const fs = require("fs"); +const path = process.argv[2]; +const raw = fs.readFileSync(path, "utf8").trim(); +const payload = JSON.parse(raw); +if (!Array.isArray(payload)) { + throw new Error("config providers output is not an array"); +} +for (const item of payload) { + if (item && item.enabled === true && typeof item.provider === "string" && item.provider) { + console.log(item.provider); + } +} +NODE + rm -f "$provider_status" "$provider_err" "$provider_list" + echo "failed to parse CodexBarCLI config providers output" >&2 + exit 2 + fi + while IFS= read -r provider; do + [[ -n "$provider" ]] && providers+=("$provider") + done <"$provider_list" + rm -f "$provider_status" "$provider_err" "$provider_list" + if [[ "${#providers[@]}" -eq 0 ]]; then + echo "no enabled providers found via CodexBarCLI config providers" >&2 + exit 2 + fi + ;; + --default) + providers=("__default__") + ;; + --provider) + if [[ -z "${1:-}" ]]; then + echo "missing provider" >&2 + exit 2 + fi + providers=("${1:-}") + ;; + --providers) + if [[ -z "${1:-}" ]]; then + echo "missing providers" >&2 + exit 2 + fi + IFS=',' read -r -a providers <<< "${1:-}" + ;; + -h|--help|"") + usage + exit 0 + ;; + *) + echo "unknown mode: $mode" >&2 + usage >&2 + exit 2 + ;; +esac + +redact_node=' +const redact = s => String(s || "") + .replace(/[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+/g, "") + .replace(/sk-[A-Za-z0-9_-]{12,}/g, "sk-REDACTED") + .replace(/gsk_[A-Za-z0-9_-]{12,}/g, "gsk_REDACTED") + .replace(/[A-Za-z0-9_-]{32,}/g, m => /[A-Za-z]/.test(m) && /[0-9]/.test(m) ? "" : m); +' + +run_one() { + local name="$1" + shift + local out err start end elapsed st node_status + out="$(mktemp)" + err="$(mktemp)" + start="$(date +%s)" + "$TIMEOUT_BIN" "$CASE_TIMEOUT" "$CLI" usage "$@" --format json --json-only --web-timeout "$WEB_TIMEOUT" >"$out" 2>"$err" + st=$? + end="$(date +%s)" + elapsed=$((end - start)) + node - "$name" "$st" "$elapsed" "$out" "$err" <&2 + exit 2 +fi +exit "$overall" diff --git a/.agents/skills/release-codexbar/SKILL.md b/.agents/skills/release-codexbar/SKILL.md new file mode 100644 index 0000000..2d823e8 --- /dev/null +++ b/.agents/skills/release-codexbar/SKILL.md @@ -0,0 +1,141 @@ +--- +name: release-codexbar +description: "CodexBar release: versioning, notarization, appcast, Homebrew, post-release bump." +--- + +# CodexBar Release + +Use for releasing signed/notarized macOS apps, especially repos with Sparkle appcasts and Homebrew casks. + +## Start + +1. Work from the app repo unless asked otherwise. +2. Check repo state, current version, latest tag/release, and release docs/scripts. +3. Confirm `CHANGELOG.md` is complete, user-facing, deduped, and dated for the release. +4. Prefer the repo release script; patch small script/test blockers instead of bypassing the release path. +5. Never print key material. Keep 1Password references and local key paths as references only. +6. Load `$release-private` if it exists before resolving Peter-owned credential locators. + +## Key Material + +Use `$one-password` for secret handling. `op` only in tmux/persistent shell; no broad `env`, `set`, `export -p`, or secret scans. + +Known App Store Connect shape: + +- fields: `private_key_p8`, `key_id`, `issuer_id` +- keep all three fields from the same 1Password item; do not mix with stale values from `~/.profile` +- resolve Peter-owned item refs from `$release-private` + +Known Sparkle key: + +- resolve the private key file from `$release-private` +- pass as `SPARKLE_PRIVATE_KEY_FILE` + +Safe env file pattern: + +```text +APP_STORE_CONNECT_API_KEY_P8=<1Password ref from release-private> +APP_STORE_CONNECT_KEY_ID=<1Password ref from release-private> +APP_STORE_CONNECT_ISSUER_ID=<1Password ref from release-private> +SPARKLE_PRIVATE_KEY_FILE= +``` + +Run with `op run --account my.1password.com --env-file -- `. + /// WebKit `document.body.innerText` often does not include the email, so we parse it from HTML. + public static func parseSignedInEmailFromClientBootstrap(html: String) -> String? { + guard let data = self.clientBootstrapJSONData(fromHTML: html) else { return nil } + guard let json = try? JSONSerialization.jsonObject(with: data, options: []) else { return nil } + + // Fast path: common structure. + if let dict = json as? [String: Any] { + if let session = dict["session"] as? [String: Any], + let user = session["user"] as? [String: Any], + let email = user["email"] as? String, + email.contains("@") + { + return email.trimmingCharacters(in: .whitespacesAndNewlines) + } + + if let user = dict["user"] as? [String: Any], + let email = user["email"] as? String, + email.contains("@") + { + return email.trimmingCharacters(in: .whitespacesAndNewlines) + } + } + + // Fallback: BFS scan for an email key/value. + var queue: [Any] = [json] + var seen = 0 + while !queue.isEmpty, seen < 4000 { + let cur = queue.removeFirst() + seen += 1 + if let dict = cur as? [String: Any] { + for (k, v) in dict { + if k.lowercased() == "email", let email = v as? String, email.contains("@") { + return email.trimmingCharacters(in: .whitespacesAndNewlines) + } + queue.append(v) + } + } else if let arr = cur as? [Any] { + queue.append(contentsOf: arr) + } + } + return nil + } + + /// Extracts the auth status from `client-bootstrap`, if present. + /// Expected values include `logged_in` and `logged_out`. + public static func parseAuthStatusFromClientBootstrap(html: String) -> String? { + guard let data = self.clientBootstrapJSONData(fromHTML: html) else { return nil } + guard let json = try? JSONSerialization.jsonObject(with: data, options: []) else { return nil } + guard let dict = json as? [String: Any] else { return nil } + if let authStatus = dict["authStatus"] as? String, !authStatus.isEmpty { + return authStatus.trimmingCharacters(in: .whitespacesAndNewlines) + } + return nil + } + + public static func parseCodeReviewRemainingPercent(bodyText: String) -> Double? { + let cleaned = bodyText.replacingOccurrences(of: "\r", with: "\n") + for regex in self.codeReviewRegexes { + let range = NSRange(cleaned.startIndex..= 2, + let r = Range(match.range(at: 1), in: cleaned) + else { continue } + if let val = Double(cleaned[r]) { return min(100, max(0, val)) } + } + return nil + } + + public static func parseCreditsRemaining(bodyText: String) -> Double? { + let cleaned = bodyText.replacingOccurrences(of: "\r", with: "\n") + let patterns = [ + #"credits\s*remaining[^0-9]*([0-9][0-9.,]*)"#, + #"remaining\s*credits[^0-9]*([0-9][0-9.,]*)"#, + #"credit\s*balance[^0-9]*([0-9][0-9.,]*)"#, + ] + for pattern in patterns { + if let val = TextParsing.firstNumber(pattern: pattern, text: cleaned) { return val } + } + return nil + } + + public static func parseRateLimits( + bodyText: String, + now: Date = .init()) -> (primary: RateWindow?, secondary: RateWindow?) + { + let cleaned = bodyText.replacingOccurrences(of: "\r", with: "\n") + let lines = cleaned + .split(whereSeparator: \.isNewline) + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + .filter { !$0.isEmpty } + + let primary = self.parseRateWindow( + lines: lines, + match: self.isFiveHourLimitLine, + windowMinutes: 5 * 60, + now: now) + let secondary = self.parseRateWindow( + lines: lines, + match: self.isWeeklyLimitLine, + windowMinutes: 7 * 24 * 60, + now: now) + return (primary, secondary) + } + + public static func parseCodeReviewLimit(bodyText: String, now: Date = .init()) -> RateWindow? { + let cleaned = bodyText.replacingOccurrences(of: "\r", with: "\n") + let lines = cleaned + .split(whereSeparator: \.isNewline) + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + .filter { !$0.isEmpty } + + return self.parseRateWindow( + lines: lines, + match: self.isCodeReviewLimitLine, + windowMinutes: nil, + now: now) + } + + public static func parsePlanFromHTML(html: String) -> String? { + if let data = self.clientBootstrapJSONData(fromHTML: html), + let plan = self.findPlan(in: data) + { + return plan + } + if let data = self.nextDataJSONData(fromHTML: html), + let plan = self.findPlan(in: data) + { + return plan + } + return nil + } + + public static func parseCreditEvents(rows: [[String]]) -> [CreditEvent] { + let formatter = self.creditDateFormatter() + + return rows.compactMap { row in + guard row.count >= 3 else { return nil } + let dateString = row[0] + let service = row[1].trimmingCharacters(in: .whitespacesAndNewlines) + let amountString = row[2] + guard let date = formatter.date(from: dateString) else { return nil } + let creditsUsed = Self.parseCreditsUsed(amountString) + return CreditEvent(date: date, service: service, creditsUsed: creditsUsed) + } + .sorted { $0.date > $1.date } + } + + private static func parseCreditsUsed(_ text: String) -> Double { + guard let raw = self.firstNumberToken(in: text) else { return 0 } + let token = raw + .replacingOccurrences(of: "\u{00A0}", with: "") + .replacingOccurrences(of: "\u{202F}", with: "") + .replacingOccurrences(of: " ", with: "") + let hasComma = token.contains(",") + let hasDot = token.contains(".") + if hasComma, hasDot { + return TextParsing.firstNumber(pattern: #"([0-9][0-9.,\s\p{Zs}]*)"#, text: token) ?? 0 + } + if hasComma { + if self.usesLocalizedDecimalCommaCreditLabel(text) { + return Double(token.replacingOccurrences(of: ",", with: ".")) ?? 0 + } + if token.range(of: #"^\d{1,3}(,\d{3})+$"#, options: .regularExpression) != nil { + return Double(token.replacingOccurrences(of: ",", with: "")) ?? 0 + } + return Double(token.replacingOccurrences(of: ",", with: ".")) ?? 0 + } + return Double(token) ?? 0 + } + + private static func usesLocalizedDecimalCommaCreditLabel(_ text: String) -> Bool { + text + .lowercased() + .contains("crédit") + } + + private static func firstNumberToken(in text: String) -> String? { + guard let regex = try? NSRegularExpression( + pattern: #"([0-9][0-9.,\s\p{Zs}]*)"#, + options: []) + else { + return nil + } + let range = NSRange(text.startIndex..= 2, + let tokenRange = Range(match.range(at: 1), in: text) + else { + return nil + } + return String(text[tokenRange]) + } + + // MARK: - Private + + private static let codeReviewRegexes: [NSRegularExpression] = { + let patterns = [ + #"Code\s*review[^0-9%]*([0-9]{1,3})%\s*remaining"#, + #"Core\s*review[^0-9%]*([0-9]{1,3})%\s*remaining"#, + ] + return patterns.compactMap { pattern in + try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) + } + }() + + private static let creditDateFormatterKey = "OpenAIDashboardParser.creditDateFormatter" + private static let clientBootstrapNeedle = Data("id=\"client-bootstrap\"".utf8) + private static let nextDataNeedle = Data("id=\"__NEXT_DATA__\"".utf8) + private static let scriptCloseNeedle = Data("".utf8) + + private static func creditDateFormatter() -> DateFormatter { + let threadDict = Thread.current.threadDictionary + if let cached = threadDict[self.creditDateFormatterKey] as? DateFormatter { + return cached + } + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.dateFormat = "MMM d, yyyy" + threadDict[self.creditDateFormatterKey] = formatter + return formatter + } + + private static func clientBootstrapJSONData(fromHTML html: String) -> Data? { + let data = Data(html.utf8) + guard let idRange = data.range(of: self.clientBootstrapNeedle) else { return nil } + + guard let openTagEnd = data[idRange.upperBound...].firstIndex(of: UInt8(ascii: ">")) else { return nil } + let contentStart = data.index(after: openTagEnd) + guard let closeRange = data.range( + of: self.scriptCloseNeedle, + options: [], + in: contentStart.. Data? { + let data = Data(html.utf8) + guard let idRange = data.range(of: self.nextDataNeedle) else { return nil } + + guard let openTagEnd = data[idRange.upperBound...].firstIndex(of: UInt8(ascii: ">")) else { return nil } + let contentStart = data.index(after: openTagEnd) + guard let closeRange = data.range( + of: self.scriptCloseNeedle, + options: [], + in: contentStart.. Data { + guard !data.isEmpty else { return data } + var start = data.startIndex + var end = data.endIndex + + while start < end, data[start].isASCIIWhitespace { + start = data.index(after: start) + } + while end > start { + let prev = data.index(before: end) + if data[prev].isASCIIWhitespace { + end = prev + } else { + break + } + } + return data.subdata(in: start.. Bool, + windowMinutes: Int?, + now: Date) -> RateWindow? + { + for idx in lines.indices where match(lines[idx]) { + let end = min(lines.count - 1, idx + 5) + let windowLines = Array(lines[idx...end]) + + var percentValue: Double? + var isRemaining = true + for line in windowLines { + if let percent = self.parsePercent(from: line) { + percentValue = percent.value + isRemaining = percent.isRemaining + break + } + } + + guard let percentValue else { continue } + let usedPercent = isRemaining ? max(0, min(100, 100 - percentValue)) : max(0, min(100, percentValue)) + + let resetLine = windowLines.first { $0.localizedCaseInsensitiveContains("reset") } + let resetDescription = resetLine?.trimmingCharacters(in: .whitespacesAndNewlines) + let resetsAt = resetLine.flatMap { self.parseResetDate(from: $0, now: now) } + let fallbackDescription = resetsAt.map { UsageFormatter.resetDescription(from: $0) } + + return RateWindow( + usedPercent: usedPercent, + windowMinutes: windowMinutes, + resetsAt: resetsAt, + resetDescription: resetDescription ?? fallbackDescription) + } + return nil + } + + private static func parsePercent(from line: String) -> (value: Double, isRemaining: Bool)? { + guard let percent = TextParsing.firstNumber(pattern: #"([0-9]{1,3})\s*%"#, text: line) else { return nil } + let lower = line.lowercased() + let isRemaining = lower.contains("remaining") || lower.contains("left") + let isUsed = lower.contains("used") || lower.contains("spent") || lower.contains("consumed") + if isUsed { return (percent, false) } + if isRemaining { return (percent, true) } + return (percent, true) + } + + private static func isFiveHourLimitLine(_ line: String) -> Bool { + let lower = line.lowercased() + if lower.contains("5h") { return true } + if lower.range(of: #"\b5\s*h\b"#, options: .regularExpression) != nil { return true } + if lower.contains("5-hour") { return true } + if lower.contains("5 hour") { return true } + return false + } + + private static func isWeeklyLimitLine(_ line: String) -> Bool { + let lower = line.lowercased() + if lower.contains("weekly") { return true } + if lower.contains("7-day") { return true } + if lower.contains("7 day") { return true } + if lower.contains("7d") { return true } + if lower.range(of: #"\b7\s*d\b"#, options: .regularExpression) != nil { return true } + return false + } + + private static func isCodeReviewLimitLine(_ line: String) -> Bool { + let lower = line.lowercased() + guard lower.contains("code review") || lower.contains("core review") else { return false } + if lower.contains("github code review") { return false } + return true + } + + private static func parseResetDate(from line: String, now: Date) -> Date? { + var raw = line.trimmingCharacters(in: .whitespacesAndNewlines) + raw = raw.replacingOccurrences(of: #"(?i)^resets?:?\s*"#, with: "", options: .regularExpression) + raw = raw.replacingOccurrences(of: " at ", with: " ", options: .caseInsensitive) + raw = raw.replacingOccurrences(of: " on ", with: " ", options: .caseInsensitive) + raw = raw.replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespacesAndNewlines) + + let calendar = Calendar(identifier: .gregorian) + let monthDayFormatter = DateFormatter() + monthDayFormatter.locale = Locale(identifier: "en_US_POSIX") + monthDayFormatter.timeZone = TimeZone.current + monthDayFormatter.dateFormat = "MMM d" + + var candidate = raw + let lower = candidate.lowercased() + var usedRelativeDay = false + + if lower.contains("today") { + usedRelativeDay = true + let dateText = monthDayFormatter.string(from: now) + candidate = candidate.replacingOccurrences(of: "today", with: dateText, options: .caseInsensitive) + } else if lower.contains("tomorrow") { + usedRelativeDay = true + if let tomorrow = calendar.date(byAdding: .day, value: 1, to: now) { + let dateText = monthDayFormatter.string(from: tomorrow) + candidate = candidate.replacingOccurrences(of: "tomorrow", with: dateText, options: .caseInsensitive) + } + } + + if let weekdayMatch = self.weekdayMatch(in: candidate) { + usedRelativeDay = true + let target = self.nextWeekdayDate(weekday: weekdayMatch.weekday, now: now, calendar: calendar) + let dateText = monthDayFormatter.string(from: target) + candidate = candidate.replacingOccurrences( + of: weekdayMatch.matched, + with: dateText, + options: .caseInsensitive) + } + + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = TimeZone.current + formatter.defaultDate = now + + let formats = [ + "MMM d h:mma", + "MMM d, h:mma", + "MMM d h:mm a", + "MMM d, h:mm a", + "MMM d HH:mm", + "MMM d, HH:mm", + "MMM d", + "M/d h:mma", + "M/d h:mm a", + "M/d/yyyy h:mm a", + "M/d/yy h:mm a", + "M/d", + "yyyy-MM-dd HH:mm", + "yyyy-MM-dd h:mm a", + "yyyy-MM-dd", + ] + + for format in formats { + formatter.dateFormat = format + if let date = formatter.date(from: candidate) { + if usedRelativeDay, date < now { + if lower.contains("today"), + let bumped = calendar.date(byAdding: .day, value: 1, to: date) + { + return bumped + } + if let bumped = calendar.date(byAdding: .day, value: 7, to: date) { + return bumped + } + } + return date + } + } + return nil + } + + private struct WeekdayMatch { + let matched: String + let weekday: Int + } + + private static func weekdayMatch(in text: String) -> WeekdayMatch? { + // The optional "(day)?" suffix requires each alternative to be long enough that + // adding "day" reproduces the full weekday name. "wed"+"day"="wedday" and + // "sat"+"day"="satday", so the longer prefixes "wednes" and "satur" are needed + // to cover the full "Wednesday" and "Saturday" forms — mirroring how the existing + // "tue|tues" and "thu|thur|thurs" alternatives layer up to "tuesday"/"thursday". + let pattern = #"\b(mon|tue|tues|wed|wednes|thu|thur|thurs|fri|sat|satur|sun)(day)?\b"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) else { return nil } + let range = NSRange(text.startIndex.. Date { + let currentWeekday = calendar.component(.weekday, from: now) + var delta = weekday - currentWeekday + if delta < 0 { delta += 7 } + guard let next = calendar.date(byAdding: .day, value: delta, to: calendar.startOfDay(for: now)) else { + return now + } + return next + } + + private static func findPlan(in data: Data) -> String? { + guard let json = try? JSONSerialization.jsonObject(with: data, options: []) else { return nil } + return self.findPlan(in: json) + } + + private static func findPlan(in json: Any) -> String? { + var queue: [Any] = [json] + var seen = 0 + while !queue.isEmpty, seen < 6000 { + let cur = queue.removeFirst() + seen += 1 + if let dict = cur as? [String: Any] { + for (k, v) in dict { + if let plan = self.planCandidate(forKey: k, value: v) { return plan } + queue.append(v) + } + } else if let arr = cur as? [Any] { + queue.append(contentsOf: arr) + } + } + return nil + } + + private static func planCandidate(forKey key: String, value: Any) -> String? { + guard self.isPlanKey(key) else { return nil } + if let str = value as? String { + return self.normalizePlanValue(str) + } + if let dict = value as? [String: Any] { + if let name = dict["name"] as? String, let plan = self.normalizePlanValue(name) { return plan } + if let display = dict["displayName"] as? String, let plan = self.normalizePlanValue(display) { return plan } + if let tier = dict["tier"] as? String, let plan = self.normalizePlanValue(tier) { return plan } + } + return nil + } + + private static func isPlanKey(_ key: String) -> Bool { + let lower = key.lowercased() + return lower.contains("plan") || lower.contains("tier") || lower.contains("subscription") + } + + private static func normalizePlanValue(_ value: String) -> String? { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + let lower = trimmed.lowercased() + let allowed = [ + "free", + "plus", + "pro", + "team", + "enterprise", + "business", + "edu", + "education", + "gov", + "premium", + "essential", + ] + guard allowed.contains(where: { lower.contains($0) }) else { return nil } + return CodexPlanFormatting.displayName(trimmed) ?? UsageFormatter.cleanPlanName(trimmed) + } +} + +extension UInt8 { + fileprivate var isASCIIWhitespace: Bool { + switch self { + case 9, 10, 13, 32: true + default: false + } + } +} diff --git a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardScrapeScript.swift b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardScrapeScript.swift new file mode 100644 index 0000000..b45d573 --- /dev/null +++ b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardScrapeScript.swift @@ -0,0 +1,907 @@ +#if os(macOS) +let openAIDashboardScrapeScript = """ + + (() => { + const textOf = el => { + const raw = el && (el.innerText || el.textContent) ? String(el.innerText || el.textContent) : ''; + return raw.trim(); + }; + const parseHexColor = (color) => { + if (!color) return null; + const c = String(color).trim().toLowerCase(); + if (c.startsWith('#')) { + if (c.length === 4) { + return '#' + c[1] + c[1] + c[2] + c[2] + c[3] + c[3]; + } + if (c.length === 7) return c; + return c; + } + const m = c.match(/^rgba?\\(([^)]+)\\)$/); + if (m) { + const parts = m[1].split(',').map(x => parseFloat(x.trim())).filter(x => Number.isFinite(x)); + if (parts.length >= 3) { + const r = Math.max(0, Math.min(255, Math.round(parts[0]))); + const g = Math.max(0, Math.min(255, Math.round(parts[1]))); + const b = Math.max(0, Math.min(255, Math.round(parts[2]))); + const toHex = n => n.toString(16).padStart(2, '0'); + return '#' + toHex(r) + toHex(g) + toHex(b); + } + } + return c; + }; + const reactPropsOf = (el) => { + if (!el) return null; + try { + const keys = Object.keys(el); + const propsKey = keys.find(k => k.startsWith('__reactProps$')); + if (propsKey) return el[propsKey] || null; + const fiberKey = keys.find(k => k.startsWith('__reactFiber$')); + if (fiberKey) { + const fiber = el[fiberKey]; + return (fiber && (fiber.memoizedProps || fiber.pendingProps)) || null; + } + } catch {} + return null; + }; + const reactFiberOf = (el) => { + if (!el) return null; + try { + const keys = Object.keys(el); + const fiberKey = keys.find(k => k.startsWith('__reactFiber$')); + return fiberKey ? (el[fiberKey] || null) : null; + } catch { + return null; + } + }; + const nestedBarMetaOf = (root) => { + if (!root || typeof root !== 'object') return null; + const queue = [root]; + const seen = typeof WeakSet !== 'undefined' ? new WeakSet() : null; + let steps = 0; + while (queue.length && steps < 250) { + const cur = queue.shift(); + steps++; + if (!cur || typeof cur !== 'object') continue; + if (seen) { + if (seen.has(cur)) continue; + seen.add(cur); + } + if (cur.payload && (cur.dataKey || cur.name || cur.value !== undefined)) return cur; + const values = Array.isArray(cur) ? cur : Object.values(cur); + for (const v of values) { + if (v && typeof v === 'object') queue.push(v); + } + } + return null; + }; + const barMetaFromElement = (el) => { + const direct = reactPropsOf(el); + if (direct && direct.payload && (direct.dataKey || direct.name || direct.value !== undefined)) return direct; + + const fiber = reactFiberOf(el); + if (fiber) { + let cur = fiber; + for (let i = 0; i < 10 && cur; i++) { + const props = (cur.memoizedProps || cur.pendingProps) || null; + if (props && props.payload && (props.dataKey || props.name || props.value !== undefined)) return props; + const nested = props ? nestedBarMetaOf(props) : null; + if (nested) return nested; + cur = cur.return || null; + } + } + + if (direct) { + const nested = nestedBarMetaOf(direct); + if (nested) return nested; + } + return null; + }; + const normalizeHref = (raw) => { + if (!raw) return null; + const href = String(raw).trim(); + if (!href) return null; + if (href.startsWith('http://') || href.startsWith('https://')) return href; + if (href.startsWith('//')) return window.location.protocol + href; + if (href.startsWith('/')) return window.location.origin + href; + return window.location.origin + '/' + href; + }; + const isLikelyCreditsURL = (raw) => { + if (!raw) return false; + try { + const url = new URL(raw, window.location.origin); + if (!url.host || !url.host.includes('chatgpt.com')) return false; + const path = url.pathname.toLowerCase(); + return ( + path.includes('settings') || + path.includes('usage') || + path.includes('billing') || + path.includes('credits') + ); + } catch { + return false; + } + }; + const purchaseTextMatches = (text) => { + const lower = String(text || '').trim().toLowerCase(); + if (!lower) return false; + if (lower.includes('add more')) return true; + if (!lower.includes('credit')) return false; + return ( + lower.includes('buy') || + lower.includes('add') || + lower.includes('purchase') || + lower.includes('top up') || + lower.includes('top-up') + ); + }; + const elementLabel = (el) => { + if (!el) return ''; + return ( + textOf(el) || + el.getAttribute('aria-label') || + el.getAttribute('title') || + '' + ); + }; + const urlFromProps = (props) => { + if (!props || typeof props !== 'object') return null; + const candidates = [ + props.href, + props.to, + props.url, + props.link, + props.destination, + props.navigateTo + ]; + for (const candidate of candidates) { + if (typeof candidate === 'string' && candidate.trim()) { + return normalizeHref(candidate); + } + } + return null; + }; + const purchaseURLFromElement = (el) => { + if (!el) return null; + const isAnchor = el.tagName && el.tagName.toLowerCase() === 'a'; + const anchor = isAnchor ? el : (el.closest ? el.closest('a') : null); + const anchorHref = anchor ? anchor.getAttribute('href') : null; + const dataHref = el.getAttribute + ? (el.getAttribute('data-href') || + el.getAttribute('data-url') || + el.getAttribute('data-link') || + el.getAttribute('data-destination')) + : null; + const propHref = urlFromProps(reactPropsOf(el)) || urlFromProps(reactPropsOf(anchor)); + const normalized = normalizeHref(anchorHref || dataHref || propHref); + return normalized && isLikelyCreditsURL(normalized) ? normalized : null; + }; + const cleanPlanName = (raw) => String(raw || '') + .replace(/\\b(claude|codex|account|plan)\\b/gi, ' ') + .replace(/_/g, ' ') + .replace(/-/g, ' ') + .replace(/\\s+/g, ' ') + .trim(); + const codexPlanDisplayName = (raw) => { + const trimmed = String(raw || '').trim(); + if (!trimmed) return null; + const lower = trimmed.toLowerCase(); + const exact = { + pro: 'Pro 20x', + prolite: 'Pro 5x', + 'pro_lite': 'Pro 5x', + 'pro-lite': 'Pro 5x', + 'pro lite': 'Pro 5x' + }; + if (exact[lower]) return exact[lower]; + const cleaned = cleanPlanName(trimmed); + if (!cleaned) return trimmed; + if (exact[cleaned.toLowerCase()]) return exact[cleaned.toLowerCase()]; + return cleaned.split(' ') + .filter(Boolean) + .map(word => { + const wordLower = word.toLowerCase(); + if (wordLower === 'cbp' || wordLower === 'k12') return wordLower.toUpperCase(); + if (word === word.toUpperCase() && /[a-z]/i.test(word)) return word; + return word.charAt(0).toUpperCase() + word.slice(1); + }) + .join(' ') || cleaned; + }; + const normalizePlanValue = (value) => { + const trimmed = String(value || '').trim(); + if (!trimmed) return null; + const lower = trimmed.toLowerCase(); + const allowed = [ + 'free', + 'plus', + 'pro', + 'team', + 'enterprise', + 'business', + 'edu', + 'education', + 'gov', + 'premium', + 'essential' + ]; + if (!allowed.some(token => lower.includes(token))) return null; + return codexPlanDisplayName(trimmed) || cleanPlanName(trimmed); + }; + const planCandidate = (key, value) => { + const lower = String(key || '').toLowerCase(); + if (!lower.includes('plan') && !lower.includes('tier') && !lower.includes('subscription')) return null; + if (typeof value === 'string') return normalizePlanValue(value); + if (value && typeof value === 'object' && !Array.isArray(value)) { + return normalizePlanValue(value.name) || + normalizePlanValue(value.displayName) || + normalizePlanValue(value.tier); + } + return null; + }; + const findPlan = (root) => { + if (!root || typeof root !== 'object') return null; + const queue = [root]; + const seenObjects = typeof WeakSet !== 'undefined' ? new WeakSet() : null; + let index = 0; + let seen = 0; + while (index < queue.length && seen < 6000) { + const cur = queue[index++]; + seen++; + if (!cur || typeof cur !== 'object') continue; + if (seenObjects) { + if (seenObjects.has(cur)) continue; + seenObjects.add(cur); + } + if (Array.isArray(cur)) { + for (const v of cur) { + if (v && typeof v === 'object') queue.push(v); + } + continue; + } + for (const [k, v] of Object.entries(cur)) { + const plan = planCandidate(k, v); + if (plan) return plan; + if (v && typeof v === 'object') queue.push(v); + } + } + return null; + }; + const parseJSONScript = (id) => { + try { + const node = document.getElementById(id); + const raw = node && node.textContent ? String(node.textContent) : ''; + return raw ? JSON.parse(raw) : null; + } catch { + return null; + } + }; + const pickLikelyPurchaseButton = (buttons) => { + if (!buttons || buttons.length === 0) return null; + const labeled = buttons.find(btn => { + const label = elementLabel(btn); + if (purchaseTextMatches(label)) return true; + const aria = String(btn.getAttribute('aria-label') || '').toLowerCase(); + return aria.includes('credit') || aria.includes('buy') || aria.includes('add'); + }); + return labeled || buttons[0]; + }; + const findCreditsPurchaseButton = () => { + const nodes = Array.from(document.querySelectorAll('h1,h2,h3,div,span,p')); + const labelMatch = nodes.find(node => { + const lower = textOf(node).toLowerCase(); + return lower === 'credits remaining' || (lower.includes('credits') && lower.includes('remaining')); + }); + if (!labelMatch) return null; + let cur = labelMatch; + for (let i = 0; i < 6 && cur; i++) { + const buttons = Array.from(cur.querySelectorAll('button, a')); + const picked = pickLikelyPurchaseButton(buttons); + if (picked) return picked; + cur = cur.parentElement; + } + return null; + }; + const dayKeyFromPayload = (payload) => { + if (!payload || typeof payload !== 'object') return null; + const localDayKeyForDate = (date) => { + const year = date.getFullYear(); + const month = String(date.getMonth() + 1).padStart(2, '0'); + const day = String(date.getDate()).padStart(2, '0'); + return `${year}-${month}-${day}`; + }; + const keys = ['day', 'date', 'name', 'label', 'x', 'time', 'timestamp']; + for (const k of keys) { + const v = payload[k]; + if (typeof v === 'string') { + const s = v.trim(); + if (/^\\d{4}-\\d{2}-\\d{2}$/.test(s)) return s; + const iso = s.match(/^(\\d{4}-\\d{2}-\\d{2})/); + if (iso) return iso[1]; + } + if (typeof v === 'number' && Number.isFinite(v) && (k === 'timestamp' || k === 'time' || k === 'x')) { + try { + const d = new Date(v); + if (!isNaN(d.getTime())) return localDayKeyForDate(d); + } catch {} + } + } + return null; + }; + const isSkillUsageServiceKey = (raw) => { + const key = raw === null || raw === undefined ? '' : String(raw).trim().toLowerCase(); + return key.startsWith('skillusage:'); + }; + const displayNameForUsageServiceKey = (raw) => { + const key = raw === null || raw === undefined ? '' : String(raw).trim(); + if (!key) return key; + if (isSkillUsageServiceKey(key)) return null; + if (key.toUpperCase() === key && key.length <= 6) return key; + const lower = key.toLowerCase(); + if (lower === 'cli') return 'CLI'; + if (lower.includes('github') && lower.includes('review')) return 'GitHub Code Review'; + const words = lower.replace(/[_-]+/g, ' ').split(' ').filter(Boolean); + return words.map(w => w.length <= 2 ? w.toUpperCase() : w.charAt(0).toUpperCase() + w.slice(1)).join(' '); + }; + const isLikelyCodexUsageService = (raw) => { + const service = raw === null || raw === undefined ? '' : String(raw).trim().toLowerCase(); + return ( + service === 'cli' || + service === 'desktop' || + service === 'desktop app' || + service === 'vscode' || + service === 'vs code' || + service === 'unknown' || + (service.includes('github') && service.includes('review')) + ); + }; + const usageChartRootForPath = (path) => { + if (!path || !path.closest) return null; + return ( + path.closest('.recharts-wrapper') || + path.closest('svg.recharts-surface') || + path.closest('section') || + path.parentElement || + null + ); + }; + const uniqueUsageChartRoots = (paths) => { + const roots = []; + for (const path of paths) { + const root = usageChartRootForPath(path); + if (root && !roots.includes(root)) roots.push(root); + } + return roots; + }; + const usageBreakdownTitleScore = (title) => { + const lower = String(title || '').trim().toLowerCase().replace(/\\s+/g, ' '); + if (!lower) return 0; + if (lower === 'usage breakdown') return 1000000; + if (lower.includes('usage breakdown')) return 900000; + if (lower === 'personal usage') return 800000; + if (lower.includes('threads') || + lower.includes('turns') || + lower.includes('client') || + lower.includes('skill') || + lower.includes('invocation')) return -1000000; + return 0; + }; + const titleLikeElements = (scope) => { + try { + return Array.from(scope.querySelectorAll('h1,h2,h3,[role=\"heading\"],div,span,p')) + .filter(el => { + const title = textOf(el); + const lower = title.toLowerCase(); + const tag = el.tagName ? el.tagName.toLowerCase() : ''; + const isHeading = tag === 'h1' || + tag === 'h2' || + tag === 'h3' || + String(el.getAttribute('role') || '').toLowerCase() === 'heading'; + return title.length > 0 && + title.length <= 80 && + ( + isHeading || + usageBreakdownTitleScore(title) !== 0 || + lower.includes('usage breakdown') || + lower.includes('threads') || + lower.includes('turns') || + lower.includes('client') || + lower.includes('skill') || + lower.includes('invocation') + ); + }); + } catch { + return []; + } + }; + const titleNodePrecedesRoot = (titleNode, root) => { + if (!titleNode || titleNode === root || root.contains(titleNode) || titleNode.contains(root)) return false; + const relation = titleNode.compareDocumentPosition(root); + return Boolean(relation & Node.DOCUMENT_POSITION_FOLLOWING); + }; + const nearestScoredChartTitleInScope = (scope, root) => { + let best = null; + for (const titleNode of titleLikeElements(scope)) { + if (!titleNodePrecedesRoot(titleNode, root)) continue; + const title = textOf(titleNode); + const score = usageBreakdownTitleScore(title); + if (score === 0) continue; + if (!best || score >= best.score) best = { title, score }; + } + return best ? best.title : ''; + }; + const chartTitleBoundaryForRoot = (root) => { + if (!root) return null; + try { + return root.closest('section,[role=\"region\"],article') || root.parentElement || null; + } catch { + return root.parentElement || null; + } + }; + const nearestTitleTextInScope = (scope, root) => { + if (!scope) return ''; + let nearest = null; + for (const titleNode of titleLikeElements(scope)) { + if (titleNodePrecedesRoot(titleNode, root)) nearest = titleNode; + } + return textOf(nearest); + }; + const nearestChartTitleTextForRoot = (root) => { + if (!root) return ''; + try { + const boundary = chartTitleBoundaryForRoot(root) || root.parentElement || null; + let ancestor = root.parentElement || null; + for (let i = 0; i < 8 && ancestor; i++) { + const scoredTitle = nearestScoredChartTitleInScope(ancestor, root); + if (scoredTitle) return scoredTitle; + if (ancestor === boundary) break; + ancestor = ancestor.parentElement || null; + } + + return nearestTitleTextInScope(boundary, root); + } catch { + return ''; + } + }; + const legendMapForUsageChartRoot = (root) => { + const legendMap = {}; + const scopes = [ + root, + root && root.parentElement, + root && root.closest ? root.closest('section') : null + ].filter(Boolean); + for (const scope of scopes) { + try { + const legendItems = Array.from(scope.querySelectorAll('div[title]')); + for (const item of legendItems) { + const title = item.getAttribute('title') ? String(item.getAttribute('title')).trim() : ''; + const square = item.querySelector('div[style*=\"background-color\"]'); + const color = (square && square.style && square.style.backgroundColor) + ? square.style.backgroundColor + : null; + const hex = parseHexColor(color); + if (title && hex) legendMap[hex] = title; + } + } catch {} + if (Object.keys(legendMap).length > 0) break; + } + return legendMap; + }; + const parseUsageBreakdownFromChartPaths = (paths, legendMap) => { + const totalsByDay = {}; // day -> service -> value + const addValue = (day, service, value) => { + if (!day || !service || isSkillUsageServiceKey(service)) return false; + if (typeof value !== 'number' || !Number.isFinite(value) || value <= 0) return false; + if (!totalsByDay[day]) totalsByDay[day] = {}; + totalsByDay[day][service] = (totalsByDay[day][service] || 0) + value; + return true; + }; + let pointCount = 0; + for (const path of paths) { + const meta = barMetaFromElement(path) || barMetaFromElement(path.parentElement) || null; + if (!meta) continue; + + const payload = meta.payload || null; + const day = dayKeyFromPayload(payload); + if (!day) continue; + + const valuesObj = (payload && payload.values && typeof payload.values === 'object') ? payload.values : null; + if (valuesObj) { + for (const [k, v] of Object.entries(valuesObj)) { + const service = displayNameForUsageServiceKey(k); + if (addValue(day, service, v)) pointCount++; + } + continue; + } + + let value = null; + if (typeof meta.value === 'number' && Number.isFinite(meta.value)) value = meta.value; + if (value === null && typeof meta.value === 'string') { + const v = parseFloat(meta.value.replace(/,/g, '')); + if (Number.isFinite(v)) value = v; + } + if (value === null) continue; + + const fill = parseHexColor(meta.fill || path.getAttribute('fill')); + const service = + (fill && legendMap[fill]) || + (typeof meta.name === 'string' && meta.name) || + null; + if (addValue(day, service, value)) pointCount++; + } + + const dayKeys = Object.keys(totalsByDay) + .filter(day => Object.keys(totalsByDay[day] || {}).length > 0) + .sort((a, b) => b.localeCompare(a)) + .slice(0, 30); + const breakdown = dayKeys.map(day => { + const servicesMap = totalsByDay[day] || {}; + const services = Object.keys(servicesMap).map(service => ({ + service, + creditsUsed: servicesMap[service] + })).sort((a, b) => { + if (a.creditsUsed === b.creditsUsed) return a.service.localeCompare(b.service); + return b.creditsUsed - a.creditsUsed; + }); + const totalCreditsUsed = services.reduce((sum, s) => sum + (Number(s.creditsUsed) || 0), 0); + return { day, services, totalCreditsUsed }; + }); + const services = Array.from(new Set(breakdown.flatMap(day => day.services.map(service => service.service)))); + const totalCreditsUsed = breakdown.reduce((sum, day) => sum + (Number(day.totalCreditsUsed) || 0), 0); + const likelyCodexServiceCount = services.filter(isLikelyCodexUsageService).length; + return { + breakdown, + pointCount, + services, + totalCreditsUsed, + likelyCodexServiceCount, + score: likelyCodexServiceCount * 1000 + services.length * 100 + pointCount + totalCreditsUsed / 1000 + }; + }; + const usageBreakdownJSON = (() => { + try { + if (window.__codexbarUsageBreakdownJSON) return window.__codexbarUsageBreakdownJSON; + + const paths = Array.from(document.querySelectorAll('g.recharts-bar-rectangle path.recharts-rectangle')); + let debug = { + pathCount: paths.length, + chartCount: 0, + eligibleCandidateCount: 0, + selectedCandidateTitle: null, + candidateSummaries: [], + sampleReactKeys: null, + sampleMetaKeys: null, + samplePayloadKeys: null, + sampleValuesKeys: null, + sampleDayKey: null + }; + try { + const sample = paths[0] || null; + if (sample) { + const names = Object.getOwnPropertyNames(sample); + debug.sampleReactKeys = names.filter(k => k.includes('react')).slice(0, 10); + const metaSample = barMetaFromElement(sample) || barMetaFromElement(sample.parentElement) || null; + if (metaSample) { + debug.sampleMetaKeys = Object.keys(metaSample).slice(0, 12); + const payload = metaSample.payload || null; + if (payload && typeof payload === 'object') { + debug.samplePayloadKeys = Object.keys(payload).slice(0, 12); + debug.sampleDayKey = dayKeyFromPayload(payload); + const values = payload.values || null; + if (values && typeof values === 'object') { + debug.sampleValuesKeys = Object.keys(values).slice(0, 12); + } + } + } + } + } catch {} + + const roots = uniqueUsageChartRoots(paths); + debug.chartCount = roots.length; + const candidates = roots.map(root => { + const chartPaths = paths.filter(path => usageChartRootForPath(path) === root); + const title = nearestChartTitleTextForRoot(root); + const titleScore = usageBreakdownTitleScore(title); + const parsed = parseUsageBreakdownFromChartPaths(chartPaths, legendMapForUsageChartRoot(root)); + return { + root, + title, + titleScore, + pathCount: chartPaths.length, + ...parsed, + score: titleScore + parsed.score + }; + }).filter(candidate => candidate.breakdown.length > 0); + const rejectedTitleCandidates = candidates.filter(candidate => candidate.titleScore < 0); + const titledCandidates = candidates.filter(candidate => candidate.titleScore > 0); + const unknownTitleCandidates = candidates.filter(candidate => candidate.titleScore === 0); + const eligibleCandidates = titledCandidates; + eligibleCandidates.sort((a, b) => b.score - a.score); + debug.eligibleCandidateCount = eligibleCandidates.length; + debug.selectedCandidateTitle = eligibleCandidates[0] ? eligibleCandidates[0].title : null; + if (eligibleCandidates.length === 0 && candidates.length > 0) { + if (unknownTitleCandidates.length > 0) { + debug.error = 'No English usage breakdown chart title found. Candidate titles: ' + + candidates.map(candidate => candidate.title || 'Untitled chart').join(', '); + } else if (rejectedTitleCandidates.length > 0) { + debug.error = 'Only non-usage chart candidates found: ' + + rejectedTitleCandidates.map(candidate => candidate.title || 'Untitled chart').join(', '); + } + } + debug.candidateSummaries = candidates.slice(0, 6).map(candidate => ({ + title: candidate.title, + titleScore: candidate.titleScore, + pathCount: candidate.pathCount, + dayCount: candidate.breakdown.length, + pointCount: candidate.pointCount, + serviceCount: candidate.services.length, + likelyCodexServiceCount: candidate.likelyCodexServiceCount, + services: candidate.services.slice(0, 8) + })); + + const breakdown = eligibleCandidates[0] ? eligibleCandidates[0].breakdown : []; + const json = (breakdown.length > 0) ? JSON.stringify(breakdown) : null; + window.__codexbarUsageBreakdownJSON = json; + window.__codexbarUsageBreakdownDebug = json ? null : JSON.stringify(debug); + return json; + } catch { + return null; + } + })(); + const usageBreakdownDebug = (() => { + try { + return window.__codexbarUsageBreakdownDebug || null; + } catch { + return null; + } + })(); + const usageBreakdownError = (() => { + try { + if (!usageBreakdownDebug) return null; + const parsed = JSON.parse(usageBreakdownDebug); + return parsed && parsed.error ? String(parsed.error) : null; + } catch { + return null; + } + })(); + const bodyText = document.body ? String(document.body.innerText || '').trim() : ''; + const href = window.location ? String(window.location.href || '') : ''; + const workspacePicker = bodyText.includes('Select a workspace'); + const title = document.title ? String(document.title || '') : ''; + const cloudflareInterstitial = + title.toLowerCase().includes('just a moment') || + bodyText.toLowerCase().includes('checking your browser') || + bodyText.toLowerCase().includes('cloudflare'); + const authSelector = [ + 'input[type="email"]', + 'input[type="password"]', + 'input[name="username"]' + ].join(', '); + const hasAuthInputs = !!document.querySelector(authSelector); + const lower = bodyText.toLowerCase(); + const loginCTA = + lower.includes('sign in') || + lower.includes('log in') || + lower.includes('continue with google') || + lower.includes('continue with apple') || + lower.includes('continue with microsoft'); + const loginRequired = + href.includes('/auth/') || + href.includes('/login') || + (hasAuthInputs && loginCTA) || + (!hasAuthInputs && loginCTA && href.includes('chatgpt.com')); + const scrollY = (typeof window.scrollY === 'number') ? window.scrollY : 0; + const scrollHeight = document.documentElement ? (document.documentElement.scrollHeight || 0) : 0; + const viewportHeight = (typeof window.innerHeight === 'number') ? window.innerHeight : 0; + + let creditsHeaderPresent = false; + let creditsHeaderInViewport = false; + let didScrollToCredits = false; + let rows = []; + try { + const looksLikeCreditsEventRow = (cells) => { + if (!cells || cells.length < 3) return false; + const first = String(cells[0] || ''); + const amount = String(cells[2] || ''); + return /\\d{4}|\\d{1,2}[\\/.\\-]\\d{1,2}/.test(first) && /\\d/.test(amount); + }; + const allTableRows = () => Array.from(document.querySelectorAll('tbody tr')).map(tr => { + const cells = Array.from(tr.querySelectorAll('td')).map(td => textOf(td)); + return cells; + }).filter(looksLikeCreditsEventRow); + const headings = Array.from(document.querySelectorAll('h1,h2,h3')); + const header = headings.find(h => textOf(h).toLowerCase() === 'credits usage history'); + if (header) { + creditsHeaderPresent = true; + const rect = header.getBoundingClientRect(); + creditsHeaderInViewport = rect.top >= 0 && rect.top <= viewportHeight; + + // Only scrape rows from the *credits usage history* table. The page can contain other tables, + // and treating any as credits history can prevent our scroll-to-load logic from running. + const container = header.closest('section') || header.parentElement || document; + const table = container.querySelector('table') || null; + const scope = table || container; + rows = Array.from(scope.querySelectorAll('tbody tr')).map(tr => { + const cells = Array.from(tr.querySelectorAll('td')).map(td => textOf(td)); + return cells; + }).filter(r => r.length >= 3); + if (rows.length === 0) { + rows = allTableRows(); + } + if (rows.length === 0 && !window.__codexbarDidScrollToCredits) { + window.__codexbarDidScrollToCredits = true; + // If the table is virtualized/lazy-loaded, we need to scroll to trigger rendering even if the + // header is already in view. + header.scrollIntoView({ block: 'start', inline: 'nearest' }); + if (creditsHeaderInViewport) { + window.scrollBy(0, Math.max(220, viewportHeight * 0.6)); + } + didScrollToCredits = true; + } + } else if (rows.length === 0 && !window.__codexbarDidScrollToCredits && scrollHeight > viewportHeight * 1.5) { + rows = allTableRows(); + if (rows.length > 0) { + creditsHeaderPresent = true; + creditsHeaderInViewport = true; + } + } + if (rows.length === 0 && !window.__codexbarDidScrollToCredits && scrollHeight > viewportHeight * 1.5) { + // The credits history section often isn't part of the DOM until you scroll down. Nudge the page + // once so subsequent scrapes can find the header and rows. + window.__codexbarDidScrollToCredits = true; + window.scrollTo(0, Math.max(0, scrollHeight - viewportHeight - 40)); + didScrollToCredits = true; + } + } catch {} + + let creditsPurchaseURL = null; + try { + const creditsButton = findCreditsPurchaseButton(); + if (creditsButton) { + const url = purchaseURLFromElement(creditsButton); + if (url) creditsPurchaseURL = url; + } + const candidates = Array.from(document.querySelectorAll('a, button')); + for (const node of candidates) { + const label = elementLabel(node); + if (!purchaseTextMatches(label)) continue; + const url = purchaseURLFromElement(node); + if (url) { + creditsPurchaseURL = url; + break; + } + } + if (!creditsPurchaseURL) { + const anchors = Array.from(document.querySelectorAll('a[href]')); + for (const anchor of anchors) { + const label = elementLabel(anchor); + const href = anchor.getAttribute('href') || ''; + const hrefLooksRelevant = /credits|billing/i.test(href); + if (!hrefLooksRelevant && !purchaseTextMatches(label)) continue; + const url = normalizeHref(href); + if (url) { + creditsPurchaseURL = url; + break; + } + } + } + } catch {} + + let signedInEmail = null; + let authStatus = null; + let accountPlan = null; + try { + const next = window.__NEXT_DATA__ || null; + const props = (next && next.props && next.props.pageProps) ? next.props.pageProps : null; + const userEmail = (props && props.user) ? props.user.email : null; + const sessionEmail = (props && props.session && props.session.user) ? props.session.user.email : null; + signedInEmail = userEmail || sessionEmail || null; + } catch {} + + const clientBootstrap = parseJSONScript('client-bootstrap'); + if (clientBootstrap) { + try { + authStatus = typeof clientBootstrap.authStatus === 'string' ? clientBootstrap.authStatus : null; + if (!signedInEmail) { + const session = clientBootstrap.session || null; + const user = (session && session.user) || clientBootstrap.user || null; + const email = user && typeof user.email === 'string' ? user.email : null; + if (email && email.includes('@')) signedInEmail = email; + } + if (!accountPlan) accountPlan = findPlan(clientBootstrap); + } catch {} + } + if (!accountPlan) { + try { + accountPlan = findPlan(window.__NEXT_DATA__ || parseJSONScript('__NEXT_DATA__')); + } catch {} + } + + if (!signedInEmail) { + try { + const obj = parseJSONScript('__NEXT_DATA__'); + if (obj) { + const queue = [obj]; + let seen = 0; + while (queue.length && seen < 2000 && !signedInEmail) { + const cur = queue.shift(); + seen++; + if (!cur) continue; + if (typeof cur === 'string') { + if (cur.includes('@')) signedInEmail = cur; + continue; + } + if (typeof cur !== 'object') continue; + for (const [k, v] of Object.entries(cur)) { + if (signedInEmail) break; + if (k === 'email' && typeof v === 'string' && v.includes('@')) { + signedInEmail = v; + break; + } + if (typeof v === 'object' && v) queue.push(v); + } + } + } + } catch {} + } + + if (!signedInEmail) { + try { + const emailRe = /[A-Z0-9._%+-]+@[A-Z0-9.-]+\\.[A-Z]{2,}/ig; + const found = (bodyText.match(emailRe) || []).map(x => String(x).trim().toLowerCase()); + const unique = Array.from(new Set(found)); + if (unique.length === 1) { + signedInEmail = unique[0]; + } else if (unique.length > 1) { + signedInEmail = unique[0]; + } + } catch {} + } + + if (!signedInEmail) { + // Last resort: open the account menu so the email becomes part of the DOM text. + try { + const emailRe = /[A-Z0-9._%+-]+@[A-Z0-9.-]+\\.[A-Z]{2,}/ig; + const hasMenu = Boolean(document.querySelector('[role="menu"]')); + if (!hasMenu) { + const button = + document.querySelector('button[aria-haspopup="menu"]') || + document.querySelector('button[aria-expanded]'); + if (button && !button.disabled) { + button.click(); + } + } + const afterText = document.body ? String(document.body.innerText || '').trim() : ''; + const found = (afterText.match(emailRe) || []).map(x => String(x).trim().toLowerCase()); + const unique = Array.from(new Set(found)); + if (unique.length === 1) { + signedInEmail = unique[0]; + } else if (unique.length > 1) { + signedInEmail = unique[0]; + } + } catch {} + } + + return { + loginRequired, + workspacePicker, + cloudflareInterstitial, + href, + bodyText, + signedInEmail, + authStatus, + accountPlan, + creditsPurchaseURL, + rows, + usageBreakdownJSON, + usageBreakdownDebug, + usageBreakdownError, + scrollY, + scrollHeight, + viewportHeight, + creditsHeaderPresent, + creditsHeaderInViewport, + didScrollToCredits + }; + })(); + +""" +#endif diff --git a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebViewCache.swift b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebViewCache.swift new file mode 100644 index 0000000..c406e99 --- /dev/null +++ b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebViewCache.swift @@ -0,0 +1,782 @@ +#if os(macOS) +import AppKit +import Foundation +import WebKit + +struct OpenAIDashboardWebViewLease { + let webView: WKWebView + let log: (String) -> Void + let setPreserveLoadedPageOnRelease: (Bool) -> Void + let release: () -> Void +} + +@MainActor +final class OpenAIDashboardWebViewCache { + static let shared = OpenAIDashboardWebViewCache() + fileprivate static let log = CodexBarLog.logger(LogCategories.openAIWebview) + + private final class ReleaseState { + var preserveLoadedPageOnRelease: Bool + + init(preserveLoadedPageOnRelease: Bool) { + self.preserveLoadedPageOnRelease = preserveLoadedPageOnRelease + } + } + + private struct AcquireOptions { + let allowTimeoutRetry: Bool + let preserveLoadedPageOnRelease: Bool + } + + @MainActor + private final class NavigationCancellationState { + private weak var webView: WKWebView? + private var delegate: NavigationDelegate? + private var isCancelled = false + + func install(webView: WKWebView, delegate: NavigationDelegate) { + self.webView = webView + self.delegate = delegate + if self.isCancelled { + self.cancel() + } + } + + func cancel() { + self.isCancelled = true + guard let webView, let delegate else { return } + delegate.cancel() + if webView.codexNavigationDelegate === delegate { + webView.stopLoading() + webView.navigationDelegate = nil + webView.codexNavigationDelegate = nil + } + self.delegate = nil + self.webView = nil + } + } + + private final class Entry { + let webView: WKWebView + let host: OffscreenWebViewHost + var lastUsedAt: Date + var isBusy: Bool + var preservedPageExpiresAt: Date? + var preservedPageExpiryWorkItem: DispatchWorkItem? + + init( + webView: WKWebView, + host: OffscreenWebViewHost, + lastUsedAt: Date, + isBusy: Bool, + preservedPageExpiresAt: Date? = nil) + { + self.webView = webView + self.host = host + self.lastUsedAt = lastUsedAt + self.isBusy = isBusy + self.preservedPageExpiresAt = preservedPageExpiresAt + } + + func armPreservedPage(until expiry: Date) { + self.preservedPageExpiresAt = expiry + } + + func setPreservedPageExpiryWorkItem(_ workItem: DispatchWorkItem?) { + self.preservedPageExpiryWorkItem?.cancel() + self.preservedPageExpiryWorkItem = workItem + } + + func clearPreservedPage() { + self.preservedPageExpiresAt = nil + self.preservedPageExpiryWorkItem?.cancel() + self.preservedPageExpiryWorkItem = nil + } + + func consumePreservedPageReuseIfAvailable(now: Date) -> Bool { + guard let preservedPageExpiresAt else { return false } + self.preservedPageExpiresAt = nil + self.preservedPageExpiryWorkItem?.cancel() + self.preservedPageExpiryWorkItem = nil + return preservedPageExpiresAt > now + } + + func hasExpiredPreservedPage(now: Date) -> Bool { + guard let preservedPageExpiresAt else { return false } + return preservedPageExpiresAt <= now + } + } + + private var entries: [ObjectIdentifier: Entry] = [:] + /// Keep the WebView alive only long enough for immediate retries/menu reopens. + /// Long-lived hidden ChatGPT tabs still consume noticeable energy on some setups. + private let idleTimeout: TimeInterval + private var idlePruneWorkItem: DispatchWorkItem? + private var idlePruneGeneration = 0 + #if DEBUG + private(set) var idlePruneDeadlineForTesting: Date? + #endif + /// Reuse the validated analytics page only for the immediate next handoff. + private let preservedPageHandoffTimeout: TimeInterval = 5 + private let blankURL = URL(string: "about:blank")! + private let idlePageClearScript = """ + (() => { + try { + document.documentElement.innerHTML = ''; + return true; + } catch { + return false; + } + })(); + """ + private let reusablePageResetScript = """ + (() => { + try { + delete window.__codexbarDidScrollToCredits; + delete window.__codexbarUsageBreakdownJSON; + delete window.__codexbarUsageBreakdownDebug; + return true; + } catch { + return false; + } + })(); + """ + private let preferredLanguageScript = """ + (() => { + const define = (target, name, value) => { + try { + Object.defineProperty(target, name, { + get: () => value, + configurable: true + }); + } catch {} + }; + define(Navigator.prototype, 'language', 'en-US'); + define(Navigator.prototype, 'languages', ['en-US', 'en']); + define(navigator, 'language', 'en-US'); + define(navigator, 'languages', ['en-US', 'en']); + })(); + """ + + init(idleTimeout: TimeInterval = 60) { + self.idleTimeout = idleTimeout + } + + nonisolated static func remainingNavigationTimeout( + until deadline: Date, + now: Date = Date()) throws -> TimeInterval + { + let remaining = deadline.timeIntervalSince(now) + guard remaining > 0 else { throw URLError(.timedOut) } + return remaining + } + + private func releaseCachedEntry(_ entry: Entry, preserveLoadedPage: Bool) { + entry.isBusy = false + let now = Date() + entry.lastUsedAt = now + self.updatePreservedPageState(for: entry, preserveLoadedPage: preserveLoadedPage) + self.prepareCachedWebViewForIdle( + entry.webView, + host: entry.host, + preserveLoadedPage: preserveLoadedPage) + self.prune(now: now) + self.scheduleNextIdlePrune(now: now) + } + + private func releaseNewEntry(_ entry: Entry, webView: WKWebView, preserveLoadedPage: Bool) { + entry.isBusy = false + let now = Date() + entry.lastUsedAt = now + self.updatePreservedPageState(for: entry, preserveLoadedPage: preserveLoadedPage) + self.prepareCachedWebViewForIdle( + webView, + host: entry.host, + preserveLoadedPage: preserveLoadedPage) + self.prune(now: now) + self.scheduleNextIdlePrune(now: now) + } + + // MARK: - Testing support + + #if DEBUG + /// Number of cached WebView entries (for testing). + var entryCount: Int { + self.entries.count + } + + /// Check if a WebView is cached for the given data store (for testing). + func hasCachedEntry(for websiteDataStore: WKWebsiteDataStore) -> Bool { + let key = ObjectIdentifier(websiteDataStore) + return self.entries[key] != nil + } + + /// Force prune with a custom "now" timestamp (for testing idle timeout). + func pruneForTesting(now: Date) { + self.prune(now: now) + } + + var idleTimeoutForTesting: TimeInterval { + self.idleTimeout + } + + var preservedPageHandoffTimeoutForTesting: TimeInterval { + self.preservedPageHandoffTimeout + } + + func hasPreservedPageForTesting(for websiteDataStore: WKWebsiteDataStore) -> Bool { + let key = ObjectIdentifier(websiteDataStore) + return self.entries[key]?.preservedPageExpiresAt != nil + } + + func markPreservedPageForTesting( + websiteDataStore: WKWebsiteDataStore, + expiresAt: Date = .init().addingTimeInterval(5)) + { + let key = ObjectIdentifier(websiteDataStore) + guard let entry = self.entries[key] else { return } + entry.armPreservedPage(until: expiresAt) + self.schedulePreservedPageExpiry(for: key, entry: entry, expiresAt: expiresAt) + } + + func consumePreservedPageForTesting(websiteDataStore: WKWebsiteDataStore, now: Date = Date()) -> Bool { + let key = ObjectIdentifier(websiteDataStore) + guard let entry = self.entries[key] else { return false } + return entry.consumePreservedPageReuseIfAvailable(now: now) + } + + /// Seed a cached entry without navigating a real page (for test stability). + @discardableResult + func cacheEntryForTesting( + websiteDataStore: WKWebsiteDataStore, + lastUsedAt: Date = Date(), + isBusy: Bool = false) -> WKWebView + { + let key = ObjectIdentifier(websiteDataStore) + if let existing = self.entries.removeValue(forKey: key) { + existing.host.close() + } + + let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore) + let entry = Entry(webView: webView, host: host, lastUsedAt: lastUsedAt, isBusy: isBusy) + self.entries[key] = entry + if isBusy { + host.show() + } else { + host.hide() + } + return webView + } + + /// Clear all cached entries (for test isolation). + func clearAllForTesting() { + self.cancelIdlePrune() + for (_, entry) in self.entries { + entry.clearPreservedPage() + entry.host.close() + } + self.entries.removeAll() + } + + func resetReusablePageStateForTesting(_ webView: WKWebView) async -> Bool { + await self.resetReusablePageState(webView) + } + #endif + + func acquire( + websiteDataStore: WKWebsiteDataStore, + usageURL: URL, + logger: ((String) -> Void)?, + navigationTimeout: TimeInterval = 15, + allowTimeoutRetry: Bool = true, + preserveLoadedPageOnRelease: Bool = false) async throws -> OpenAIDashboardWebViewLease + { + let deadline = Date().addingTimeInterval(max(navigationTimeout, 0.01)) + return try await self.acquire( + websiteDataStore: websiteDataStore, + usageURL: usageURL, + logger: logger, + deadline: deadline, + options: .init( + allowTimeoutRetry: allowTimeoutRetry, + preserveLoadedPageOnRelease: preserveLoadedPageOnRelease)) + } + + private func acquire( + websiteDataStore: WKWebsiteDataStore, + usageURL: URL, + logger: ((String) -> Void)?, + deadline: Date, + options: AcquireOptions) async throws -> OpenAIDashboardWebViewLease + { + let now = Date() + self.prune(now: now) + + let log: (String) -> Void = { message in + logger?("[webview] \(message)") + } + let key = ObjectIdentifier(websiteDataStore) + let remainingTimeout = try Self.remainingNavigationTimeout(until: deadline, now: now) + + if let entry = self.entries[key] { + if entry.isBusy { + log("Cached WebView busy; using a temporary WebView.") + let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore) + host.show() + do { + try await self.prepareWebView( + webView, + usageURL: usageURL, + timeout: remainingTimeout, + canReuseLoadedPage: false) + } catch { + if options.allowTimeoutRetry, Self.isPrepareTimeout(error) { + host.close() + log("Temporary OpenAI WebView timed out; retrying with a fresh WebView.") + return try await self.acquireTemporaryWebView( + websiteDataStore: websiteDataStore, + usageURL: usageURL, + log: log, + deadline: deadline) + } + host.close() + throw error + } + return OpenAIDashboardWebViewLease( + webView: webView, + log: log, + setPreserveLoadedPageOnRelease: { _ in }, + release: { host.close() }) + } + + entry.isBusy = true + entry.lastUsedAt = now + let canReuseLoadedPage = entry.consumePreservedPageReuseIfAvailable(now: now) + let releaseState = ReleaseState(preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease) + entry.host.show() + do { + try await self.prepareWebView( + entry.webView, + usageURL: usageURL, + timeout: remainingTimeout, + canReuseLoadedPage: canReuseLoadedPage) + } catch { + if options.allowTimeoutRetry, Self.isPrepareTimeout(error) { + entry.isBusy = false + entry.lastUsedAt = Date() + entry.clearPreservedPage() + entry.host.close() + self.entries.removeValue(forKey: key) + log("Cached OpenAI WebView timed out; recreating it.") + return try await self.acquire( + websiteDataStore: websiteDataStore, + usageURL: usageURL, + logger: logger, + deadline: deadline, + options: .init( + allowTimeoutRetry: false, + preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease)) + } + entry.isBusy = false + entry.lastUsedAt = Date() + entry.clearPreservedPage() + entry.host.close() + self.entries.removeValue(forKey: key) + Self.log.warning("OpenAI webview prepare failed") + throw error + } + + return OpenAIDashboardWebViewLease( + webView: entry.webView, + log: log, + setPreserveLoadedPageOnRelease: { preserveLoadedPageOnRelease in + releaseState.preserveLoadedPageOnRelease = preserveLoadedPageOnRelease + }, + release: { [weak self, weak entry] in + guard let self, let entry else { return } + self.releaseCachedEntry( + entry, + preserveLoadedPage: releaseState.preserveLoadedPageOnRelease) + }) + } + + let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore) + let entry = Entry(webView: webView, host: host, lastUsedAt: now, isBusy: true) + self.entries[key] = entry + host.show() + let releaseState = ReleaseState(preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease) + + do { + try await self.prepareWebView( + webView, + usageURL: usageURL, + timeout: remainingTimeout, + canReuseLoadedPage: false) + } catch { + if options.allowTimeoutRetry, Self.isPrepareTimeout(error) { + self.entries.removeValue(forKey: key) + host.close() + log("OpenAI WebView timed out during prepare; retrying once.") + return try await self.acquire( + websiteDataStore: websiteDataStore, + usageURL: usageURL, + logger: logger, + deadline: deadline, + options: .init( + allowTimeoutRetry: false, + preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease)) + } + self.entries.removeValue(forKey: key) + host.close() + Self.log.warning("OpenAI webview prepare failed") + throw error + } + + return OpenAIDashboardWebViewLease( + webView: webView, + log: log, + setPreserveLoadedPageOnRelease: { preserveLoadedPageOnRelease in + releaseState.preserveLoadedPageOnRelease = preserveLoadedPageOnRelease + }, + release: { [weak self, weak entry] in + guard let self, let entry else { return } + self.releaseNewEntry( + entry, + webView: webView, + preserveLoadedPage: releaseState.preserveLoadedPageOnRelease) + }) + } + + func evict(websiteDataStore: WKWebsiteDataStore) { + let key = ObjectIdentifier(websiteDataStore) + guard let entry = self.entries.removeValue(forKey: key) else { return } + entry.clearPreservedPage() + Self.log.debug("OpenAI webview evicted") + entry.host.close() + self.scheduleNextIdlePrune() + } + + func evictAll() { + self.cancelIdlePrune() + let existing = self.entries + self.entries.removeAll() + for (_, entry) in existing { + entry.clearPreservedPage() + entry.host.close() + } + if !existing.isEmpty { + Self.log.debug("OpenAI webview evicted all") + } + } + + func evictIdle() { + let idleEntries = self.entries.filter { _, entry in + !entry.isBusy + } + guard !idleEntries.isEmpty else { return } + + for (key, entry) in idleEntries { + entry.clearPreservedPage() + entry.host.close() + self.entries.removeValue(forKey: key) + } + Self.log.debug("OpenAI idle webviews evicted", metadata: ["count": "\(idleEntries.count)"]) + self.scheduleNextIdlePrune() + } + + private func prepareCachedWebViewForIdle( + _ webView: WKWebView, + host: OffscreenWebViewHost, + preserveLoadedPage: Bool) + { + webView.navigationDelegate = nil + webView.codexNavigationDelegate = nil + if preserveLoadedPage { + host.hide() + return + } + + // Detach the heavyweight ChatGPT SPA as soon as a scrape completes. Keeping the WebView object around + // still helps with immediate reuse, but letting chatgpt.com remain the active document is too expensive. + webView.stopLoading() + webView.evaluateJavaScript(self.idlePageClearScript, completionHandler: nil) + _ = webView.load(URLRequest(url: self.blankURL)) + host.hide() + } + + /// Schedule against the oldest idle entry so later releases cannot postpone its eviction. + private func scheduleNextIdlePrune(now: Date = Date()) { + self.cancelIdlePrune() + + guard let nextExpiry = self.entries.values + .filter({ !$0.isBusy }) + .map({ $0.lastUsedAt.addingTimeInterval(self.idleTimeout) }) + .min() + else { return } + + let generation = self.idlePruneGeneration + let workItem = DispatchWorkItem { [weak self] in + MainActor.assumeIsolated { + guard let self, self.idlePruneGeneration == generation else { return } + self.idlePruneWorkItem = nil + #if DEBUG + self.idlePruneDeadlineForTesting = nil + #endif + let pruneTime = Date() + self.prune(now: pruneTime) + self.scheduleNextIdlePrune(now: pruneTime) + } + } + self.idlePruneWorkItem = workItem + #if DEBUG + self.idlePruneDeadlineForTesting = nextExpiry + #endif + let delay = max(0, nextExpiry.timeIntervalSince(now)) + 0.01 + DispatchQueue.main.asyncAfter(deadline: .now() + delay, execute: workItem) + } + + private func cancelIdlePrune() { + self.idlePruneGeneration &+= 1 + self.idlePruneWorkItem?.cancel() + self.idlePruneWorkItem = nil + #if DEBUG + self.idlePruneDeadlineForTesting = nil + #endif + } + + private func prune(now: Date) { + for entry in self.entries.values where !entry.isBusy && entry.hasExpiredPreservedPage(now: now) { + entry.clearPreservedPage() + self.prepareCachedWebViewForIdle( + entry.webView, + host: entry.host, + preserveLoadedPage: false) + Self.log.debug("OpenAI webview preserved page expired") + } + + let expired = self.entries.filter { _, entry in + !entry.isBusy && now.timeIntervalSince(entry.lastUsedAt) >= self.idleTimeout + } + for (key, entry) in expired { + entry.host.close() + self.entries.removeValue(forKey: key) + Self.log.debug("OpenAI webview pruned") + } + } + + private func makeWebView(websiteDataStore: WKWebsiteDataStore) -> (WKWebView, OffscreenWebViewHost) { + let config = WKWebViewConfiguration() + config.websiteDataStore = websiteDataStore + let userContentController = WKUserContentController() + userContentController.addUserScript(WKUserScript( + source: self.preferredLanguageScript, + injectionTime: .atDocumentStart, + forMainFrameOnly: false)) + config.userContentController = userContentController + if #available(macOS 14.0, *) { + config.preferences.inactiveSchedulingPolicy = .suspend + } + + let webView = WKWebView(frame: .zero, configuration: config) + let host = OffscreenWebViewHost(webView: webView) + return (webView, host) + } + + private func prepareWebView( + _ webView: WKWebView, + usageURL: URL, + timeout: TimeInterval, + canReuseLoadedPage: Bool) async throws + { + #if DEBUG + if usageURL.absoluteString == "about:blank" { + _ = webView.loadHTMLString("", baseURL: nil) + return + } + #endif + + if canReuseLoadedPage, + let currentURL = webView.url?.absoluteString, + OpenAIDashboardFetcher.isUsageRoute(currentURL) + { + if await self.resetReusablePageState(webView) { + return + } + + Self.log.debug("OpenAI preserved page reset failed; reloading usage URL") + } + + try Task.checkCancellation() + let cancellationState = NavigationCancellationState() + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { (cont: CheckedContinuation) in + let delegate = NavigationDelegate { result in + cont.resume(with: result) + } + webView.navigationDelegate = delegate + webView.codexNavigationDelegate = delegate + cancellationState.install(webView: webView, delegate: delegate) + delegate.armTimeout(seconds: timeout) + _ = webView.load(OpenAIDashboardFetcher.usageURLRequest(url: usageURL)) + if Task.isCancelled { + cancellationState.cancel() + } + } + } onCancel: { + Task { @MainActor in + cancellationState.cancel() + } + } + } + + private func acquireTemporaryWebView( + websiteDataStore: WKWebsiteDataStore, + usageURL: URL, + log: @escaping (String) -> Void, + deadline: Date) async throws -> OpenAIDashboardWebViewLease + { + let remainingTimeout = try Self.remainingNavigationTimeout(until: deadline) + let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore) + host.show() + do { + try await self.prepareWebView( + webView, + usageURL: usageURL, + timeout: remainingTimeout, + canReuseLoadedPage: false) + } catch { + host.close() + throw error + } + return OpenAIDashboardWebViewLease( + webView: webView, + log: log, + setPreserveLoadedPageOnRelease: { _ in }, + release: { host.close() }) + } + + private static func isPrepareTimeout(_ error: Error) -> Bool { + let nsError = error as NSError + return nsError.domain == NSURLErrorDomain && nsError.code == NSURLErrorTimedOut + } + + private func updatePreservedPageState(for entry: Entry, preserveLoadedPage: Bool) { + if preserveLoadedPage { + let expiresAt = Date().addingTimeInterval(self.preservedPageHandoffTimeout) + entry.armPreservedPage(until: expiresAt) + if let key = self.entries.first(where: { $0.value === entry })?.key { + self.schedulePreservedPageExpiry(for: key, entry: entry, expiresAt: expiresAt) + } + } else { + entry.clearPreservedPage() + } + } + + private func schedulePreservedPageExpiry( + for key: ObjectIdentifier, + entry: Entry, + expiresAt: Date) + { + let delay = max(0, expiresAt.timeIntervalSinceNow) + let workItem = DispatchWorkItem { [weak self] in + MainActor.assumeIsolated { + self?.expirePreservedPageIfNeeded(for: key, expectedExpiry: expiresAt) + } + } + entry.setPreservedPageExpiryWorkItem(workItem) + DispatchQueue.main.asyncAfter(deadline: .now() + delay, execute: workItem) + } + + private func expirePreservedPageIfNeeded(for key: ObjectIdentifier, expectedExpiry: Date) { + guard let entry = self.entries[key], + !entry.isBusy, + let preservedPageExpiresAt = entry.preservedPageExpiresAt, + preservedPageExpiresAt == expectedExpiry, + preservedPageExpiresAt <= Date() + else { + return + } + + entry.clearPreservedPage() + self.prepareCachedWebViewForIdle( + entry.webView, + host: entry.host, + preserveLoadedPage: false) + Self.log.debug("OpenAI webview preserved page expired") + self.prune(now: Date()) + } + + private func resetReusablePageState(_ webView: WKWebView) async -> Bool { + do { + let any = try await webView.evaluateJavaScript(self.reusablePageResetScript) + return (any as? Bool) ?? true + } catch { + return false + } + } +} + +@MainActor +private final class OffscreenWebViewHost { + private let window: NSWindow + private weak var webView: WKWebView? + + init(webView: WKWebView) { + // WebKit throttles timers/RAF aggressively when a WKWebView is not considered "visible". + // The Codex usage page uses streaming SSR + client hydration; if RAF is throttled, the + // dashboard never becomes part of the visible DOM and `document.body.innerText` stays tiny. + // + // Keep a transparent (mouse-ignoring) window technically "on-screen" while scraping, but + // place it almost entirely off-screen so we never ghost-render dashboard UI over the desktop. + let visibleFrame = NSScreen.main?.visibleFrame ?? NSRect(x: 0, y: 0, width: 900, height: 700) + let frame = OpenAIDashboardFetcher.offscreenHostWindowFrame(for: visibleFrame) + let window = NSWindow( + contentRect: frame, + styleMask: [.borderless], + backing: .buffered, + defer: false) + window.isReleasedWhenClosed = false + window.backgroundColor = .clear + window.isOpaque = false + // Keep it effectively invisible, but non-zero alpha so WebKit treats it as "visible" and doesn't + // stall hydration (we've observed a head-only HTML shell for minutes at alpha=0). + window.alphaValue = OpenAIDashboardFetcher.offscreenHostAlphaValue() + window.hasShadow = false + window.ignoresMouseEvents = true + window.level = .floating + window.collectionBehavior = [.canJoinAllSpaces, .fullScreenAuxiliary] + window.isExcludedFromWindowsMenu = true + window.contentView = webView + + self.window = window + self.webView = webView + } + + func show() { + OpenAIDashboardWebViewCache.log.debug("OpenAI webview show") + self.window.alphaValue = OpenAIDashboardFetcher.offscreenHostAlphaValue() + self.window.orderFrontRegardless() + } + + func hide() { + // Set alpha to 0 so WebKit recognizes the page as inactive and applies + // its scheduling policy (throttle/suspend), reducing CPU when idle. + OpenAIDashboardWebViewCache.log.debug("OpenAI webview hide") + self.window.alphaValue = 0.0 + self.window.orderOut(nil) + } + + func close() { + OpenAIDashboardWebViewCache.log.debug("OpenAI webview close") + WebKitTeardown.scheduleCleanup( + owner: self, + window: self.window, + webView: self.webView, + closeWindow: { [window] in + window.orderOut(nil) + window.close() + }) + } +} + +#endif diff --git a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebsiteDataStore.swift b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebsiteDataStore.swift new file mode 100644 index 0000000..a868996 --- /dev/null +++ b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebsiteDataStore.swift @@ -0,0 +1,118 @@ +#if os(macOS) +import CryptoKit +import Foundation +import WebKit + +/// Per-account persistent `WKWebsiteDataStore` for the OpenAI dashboard scrape. +/// +/// Why: `WKWebsiteDataStore.default()` is a single shared cookie jar. If the user switches Codex accounts, +/// we want to keep multiple signed-in dashboard sessions around (one per email) without clearing cookies. +/// +/// Implementation detail: macOS 14+ supports `WKWebsiteDataStore.dataStore(forIdentifier:)`, which creates +/// persistent isolated stores keyed by an identifier. We derive a stable UUID from the email and optional +/// Codex source scope so distinct profiles with the same email never share a cookie store. +/// +/// Important: We cache the `WKWebsiteDataStore` instances so the same object is returned for the same +/// account email. This ensures `OpenAIDashboardWebViewCache` can use object identity for cache lookups. +@MainActor +public enum OpenAIDashboardWebsiteDataStore { + /// Cached data store instances keyed by normalized email and optional Codex source scope. + /// Using the same instance ensures stable object identity for WebView cache lookups. + private static var cachedStores: [String: WKWebsiteDataStore] = [:] + + public static func store( + forAccountEmail email: String?, + scope: CookieHeaderCache.Scope? = nil) -> WKWebsiteDataStore + { + guard let normalized = normalizeEmail(email) else { return .default() } + let storageKey = self.storageKey(normalizedEmail: normalized, scope: scope) + + // Return cached instance if available to maintain stable object identity + if let cached = cachedStores[storageKey] { + return cached + } + + let id = Self.identifier(forStorageKey: storageKey) + let store = WKWebsiteDataStore(forIdentifier: id) + self.cachedStores[storageKey] = store + return store + } + + /// Clears the persistent cookie store for a single account email. + /// + /// Note: this does *not* impact other accounts, and is safe to use when the stored session is "stuck" + /// or signed in to a different account than expected. + public static func clearStore( + forAccountEmail email: String?, + scope: CookieHeaderCache.Scope? = nil) async + { + // Clear only ChatGPT/OpenAI domain data for the per-account store. + // Avoid deleting the entire persistent store (WebKit requires all WKWebViews using it to be released). + let store = self.store(forAccountEmail: email, scope: scope) + await withCheckedContinuation { cont in + store.fetchDataRecords(ofTypes: WKWebsiteDataStore.allWebsiteDataTypes()) { records in + let filtered = records.filter { record in + let name = record.displayName.lowercased() + return name.contains("chatgpt.com") || name.contains("openai.com") + } + store.removeData(ofTypes: WKWebsiteDataStore.allWebsiteDataTypes(), for: filtered) { + cont.resume() + } + } + } + + // Remove from cache so a fresh instance is created on next access + if let normalized = normalizeEmail(email) { + self.cachedStores.removeValue(forKey: self.storageKey(normalizedEmail: normalized, scope: scope)) + } + } + + #if DEBUG + /// Clear all cached store instances (for test isolation). + public static func clearCacheForTesting() { + self.cachedStores.removeAll() + } + + #endif + + // MARK: - Private + + private static func normalizeEmail(_ email: String?) -> String? { + guard let raw = email?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { return nil } + return raw.lowercased() + } + + private static func storageKey(normalizedEmail: String, scope: CookieHeaderCache.Scope?) -> String { + guard let scope else { return normalizedEmail } + return "\(normalizedEmail)|\(scope.isolationIdentifier)" + } + + private static func identifier(forStorageKey storageKey: String) -> UUID { + let digest = SHA256.hash(data: Data(storageKey.utf8)) + var bytes = Array(digest.prefix(16)) + + // Make it a well-formed UUID (v4 + RFC4122 variant) while staying deterministic. + bytes[6] = (bytes[6] & 0x0F) | 0x40 + bytes[8] = (bytes[8] & 0x3F) | 0x80 + + let uuidBytes: uuid_t = ( + bytes[0], + bytes[1], + bytes[2], + bytes[3], + bytes[4], + bytes[5], + bytes[6], + bytes[7], + bytes[8], + bytes[9], + bytes[10], + bytes[11], + bytes[12], + bytes[13], + bytes[14], + bytes[15]) + return UUID(uuid: uuidBytes) + } +} +#endif diff --git a/Sources/CodexBarCore/PathEnvironment.swift b/Sources/CodexBarCore/PathEnvironment.swift new file mode 100644 index 0000000..631ac04 --- /dev/null +++ b/Sources/CodexBarCore/PathEnvironment.swift @@ -0,0 +1,1228 @@ +import Foundation +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif +#if os(macOS) +import Security +#endif + +public enum PathPurpose: Hashable, Sendable { + case rpc + case tty + case nodeTooling +} + +public struct PathDebugSnapshot: Equatable, Sendable { + public let codexBinary: String? + public let claudeBinary: String? + public let geminiBinary: String? + public let effectivePATH: String + public let loginShellPATH: String? + + public static let empty = PathDebugSnapshot( + codexBinary: nil, + claudeBinary: nil, + geminiBinary: nil, + effectivePATH: "", + loginShellPATH: nil) + + public init( + codexBinary: String?, + claudeBinary: String?, + geminiBinary: String? = nil, + effectivePATH: String, + loginShellPATH: String?) + { + self.codexBinary = codexBinary + self.claudeBinary = claudeBinary + self.geminiBinary = geminiBinary + self.effectivePATH = effectivePATH + self.loginShellPATH = loginShellPATH + } +} + +public enum BinaryLocator { + /// Test-only override so parallel Gemini suites can point at fake binaries + /// without mutating process-wide `GEMINI_CLI_PATH`. + @TaskLocal public static var geminiBinaryPathOverrideForTesting: String? + + public static func resolveClaudeBinary( + env: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current, + commandV: (String, String?, TimeInterval, FileManager) -> String? = ShellCommandLocator.commandV, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String? = ShellCommandLocator + .resolveAlias, + fileManager: FileManager = .default, + home: String = NSHomeDirectory()) -> String? + { + self.resolveBinary( + name: "claude", + overrideKey: "CLAUDE_CLI_PATH", + env: env, + loginPATH: loginPATH, + commandV: commandV, + aliasResolver: aliasResolver, + wellKnownPaths: self.claudeWellKnownPaths(home: home), + fileManager: fileManager, + home: home) + } + + /// Well-known installation paths for the Claude CLI binary. + /// Covers Anthropic's native installer (`~/.local/bin`), the `claude migrate-installer` + /// self-updating location (`~/.claude/local`), the legacy per-user installer + /// (`~/.claude/bin`), Homebrew, and the macOS Terminal installer (cmux.app). + static func claudeWellKnownPaths(home: String) -> [String] { + [ + "\(home)/.local/bin/claude", + "\(home)/.claude/local/claude", + "\(home)/.claude/bin/claude", + "/opt/homebrew/bin/claude", + "/usr/local/bin/claude", + "/Applications/cmux.app/Contents/Resources/bin/claude", + ] + } + + public static func resolveAntigravityBinary( + env: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current, + commandV: (String, String?, TimeInterval, FileManager) -> String? = ShellCommandLocator.commandV, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String? = ShellCommandLocator + .resolveAlias, + fileManager: FileManager = .default, + home: String = NSHomeDirectory()) -> String? + { + self.resolveBinary( + name: "agy", + overrideKey: "ANTIGRAVITY_CLI_PATH", + env: env, + loginPATH: loginPATH, + commandV: commandV, + aliasResolver: aliasResolver, + wellKnownPaths: [ + "\(home)/.local/bin/agy", + "/opt/homebrew/bin/agy", + "/usr/local/bin/agy", + ], + fileManager: fileManager, + home: home) + } + + public static func resolveCodexBinary( + env: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current, + commandV: (String, String?, TimeInterval, FileManager) -> String? = ShellCommandLocator.commandV, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String? = ShellCommandLocator + .resolveAlias, + launchCandidateFilter: (String, FileManager) -> Bool = CodexLaunchPreflight.isLaunchCandidateAllowed, + fileManager: FileManager = .default, + home: String = NSHomeDirectory()) -> String? + { + self.resolveBinary( + name: "codex", + overrideKey: "CODEX_CLI_PATH", + env: env, + loginPATH: loginPATH, + commandV: commandV, + aliasResolver: aliasResolver, + wellKnownPaths: self.codexWellKnownPaths(home: home), + launchCandidateFilter: launchCandidateFilter, + fileManager: fileManager, + home: home) + } + + /// Well-known installation paths for the signed Codex CLI bundled with current and legacy desktop apps. + /// Keep these after PATH lookups, but use them as a safe fallback when a PATH shim is blocked. + static func codexWellKnownPaths(home: String) -> [String] { + #if os(macOS) + [ + "\(home)/Applications/ChatGPT.app/Contents/Resources/codex", + "\(home)/Applications/Codex.app/Contents/Resources/codex", + "/Applications/ChatGPT.app/Contents/Resources/codex", + "/Applications/Codex.app/Contents/Resources/codex", + ] + #else + [] + #endif + } + + public static func resolveGeminiBinary( + env: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current, + commandV: (String, String?, TimeInterval, FileManager) -> String? = ShellCommandLocator.commandV, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String? = ShellCommandLocator + .resolveAlias, + fileManager: FileManager = .default, + home: String = NSHomeDirectory()) -> String? + { + if let override = self.geminiBinaryPathOverrideForTesting, + fileManager.isExecutableFile(atPath: override) + { + return override + } + return self.resolveBinary( + name: "gemini", + overrideKey: "GEMINI_CLI_PATH", + env: env, + loginPATH: loginPATH, + commandV: commandV, + aliasResolver: aliasResolver, + fileManager: fileManager, + home: home) + } + + public static func resolveGrokBinary( + env: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current, + commandV: (String, String?, TimeInterval, FileManager) -> String? = ShellCommandLocator.commandV, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String? = ShellCommandLocator + .resolveAlias, + fileManager: FileManager = .default, + home: String = NSHomeDirectory()) -> String? + { + self.resolveBinary( + name: "grok", + overrideKey: "GROK_CLI_PATH", + env: env, + loginPATH: loginPATH, + commandV: commandV, + aliasResolver: aliasResolver, + wellKnownPaths: self.grokWellKnownPaths(home: home), + fileManager: fileManager, + home: home) + } + + public static func resolveAmpBinary( + env: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current, + commandV: (String, String?, TimeInterval, FileManager) -> String? = ShellCommandLocator.commandV, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String? = ShellCommandLocator + .resolveAlias, + fileManager: FileManager = .default, + home: String = NSHomeDirectory()) -> String? + { + self.resolveBinary( + name: "amp", + overrideKey: "AMP_CLI_PATH", + env: env, + loginPATH: loginPATH, + commandV: commandV, + aliasResolver: aliasResolver, + wellKnownPaths: self.ampWellKnownPaths(home: home), + fileManager: fileManager, + home: home) + } + + static func ampWellKnownPaths(home: String) -> [String] { + [ + "\(home)/.local/bin/amp", + "\(home)/.amp/bin/amp", + "/opt/homebrew/bin/amp", + "/usr/local/bin/amp", + ] + } + + /// Well-known install locations for the Grok Build CLI binary. + /// Covers the installer's default (`~/.grok/bin/grok`) and the symlinks it sometimes + /// creates into `~/.local/bin` and `/usr/local/bin`. + static func grokWellKnownPaths(home: String) -> [String] { + [ + "\(home)/.grok/bin/grok", + "\(home)/.local/bin/grok", + "/usr/local/bin/grok", + "/opt/homebrew/bin/grok", + ] + } + + public static func resolveAWSBinary( + env: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current, + commandV: (String, String?, TimeInterval, FileManager) -> String? = ShellCommandLocator.commandV, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String? = ShellCommandLocator + .resolveAlias, + fileManager: FileManager = .default, + home: String = NSHomeDirectory()) -> String? + { + self.resolveBinary( + name: "aws", + overrideKey: "AWS_CLI_PATH", + env: env, + loginPATH: loginPATH, + commandV: commandV, + aliasResolver: aliasResolver, + wellKnownPaths: self.awsWellKnownPaths(home: home), + fileManager: fileManager, + home: home) + } + + /// Well-known install locations for the AWS CLI v2 (`aws`). + /// Covers Homebrew (Apple Silicon + Intel) and the per-user pip/uv install path. + static func awsWellKnownPaths(home: String) -> [String] { + [ + "/opt/homebrew/bin/aws", + "/usr/local/bin/aws", + "\(home)/.local/bin/aws", + ] + } + + public static func resolveAuggieBinary( + env: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current, + commandV: (String, String?, TimeInterval, FileManager) -> String? = ShellCommandLocator.commandV, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String? = ShellCommandLocator + .resolveAlias, + fileManager: FileManager = .default, + home: String = NSHomeDirectory()) -> String? + { + self.resolveBinary( + name: "auggie", + overrideKey: "AUGGIE_CLI_PATH", + env: env, + loginPATH: loginPATH, + commandV: commandV, + aliasResolver: aliasResolver, + fileManager: fileManager, + home: home) + } + + // swiftlint:disable function_parameter_count + private static func resolveBinary( + name: String, + overrideKey: String, + env: [String: String], + loginPATH: [String]?, + commandV: (String, String?, TimeInterval, FileManager) -> String?, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String?, + wellKnownPaths: [String] = [], + launchCandidateFilter: (String, FileManager) -> Bool = { _, _ in true }, + fileManager: FileManager, + home: String) -> String? + { + // swiftlint:enable function_parameter_count + // 1) Explicit override + if let override = env[overrideKey], fileManager.isExecutableFile(atPath: override) { + return override + } + + // 2) Login-shell PATH (captured once per launch) + if let loginPATH, + let pathHit = self.find( + name, + in: loginPATH, + fileManager: fileManager, + launchCandidateFilter: launchCandidateFilter) + { + return pathHit + } + + // 3) Existing PATH + if let existingPATH = env["PATH"], + let pathHit = self.find( + name, + in: existingPATH.split(separator: ":").map(String.init), + fileManager: fileManager, + launchCandidateFilter: launchCandidateFilter) + { + return pathHit + } + + // 4) Well-known installation paths (e.g. Homebrew, cmux.app bundle, ~/.claude/bin). + // Prefer these before shell probing to avoid running interactive shell init for common installs. + for candidate in wellKnownPaths + where fileManager.isExecutableFile(atPath: candidate) && launchCandidateFilter(candidate, fileManager) + { + return candidate + } + + // 5) Interactive login shell lookup (captures nvm/fnm/mise paths from .zshrc/.bashrc) + if let shellHit = commandV(name, env["SHELL"], 2.0, fileManager), + fileManager.isExecutableFile(atPath: shellHit), + launchCandidateFilter(shellHit, fileManager) + { + return shellHit + } + + // 5b) Alias fallback (login shell); only attempt after all standard lookups fail. + if let aliasHit = aliasResolver(name, env["SHELL"], 2.0, fileManager, home), + fileManager.isExecutableFile(atPath: aliasHit), + launchCandidateFilter(aliasHit, fileManager) + { + return aliasHit + } + + // 6) Minimal fallback + let fallback = ["/usr/bin", "/bin", "/usr/sbin", "/sbin"] + if let pathHit = self.find( + name, + in: fallback, + fileManager: fileManager, + launchCandidateFilter: launchCandidateFilter) + { + return pathHit + } + + return nil + } + + private static func find( + _ binary: String, + in paths: [String], + fileManager: FileManager, + launchCandidateFilter: (String, FileManager) -> Bool = { _, _ in true }) -> String? + { + for path in paths where !path.isEmpty { + let candidate = "\(path.hasSuffix("/") ? String(path.dropLast()) : path)/\(binary)" + if fileManager.isExecutableFile(atPath: candidate), launchCandidateFilter(candidate, fileManager) { + return candidate + } + } + return nil + } +} + +public enum CodexLaunchPreflight { + struct GatekeeperAssessment { + let output: String + let exitStatus: Int32 + } + + public static func isLaunchCandidateAllowed(path: String, fileManager: FileManager = .default) -> Bool { + #if os(macOS) + self.isLaunchCandidateAllowed( + path: path, + fileManager: fileManager, + hasExtendedAttribute: self.hasExtendedAttribute, + spctlAssessment: { self.spctlAssessment(path: $0) }, + appSignatureIsTrusted: self.isExpectedOpenAIAppSignature, + isMachOExecutable: self.isMachOExecutable) + #else + _ = path + _ = fileManager + return true + #endif + } + + #if os(macOS) + // Keep each security boundary injectable so preflight tests never inspect or launch host binaries. + // swiftlint:disable:next function_parameter_count + static func isLaunchCandidateAllowed( + path: String, + fileManager: FileManager, + hasExtendedAttribute: (String, String) -> Bool, + spctlAssessment: (String) -> GatekeeperAssessment?, + appSignatureIsTrusted: (String) -> Bool, + isMachOExecutable: (String) -> Bool) -> Bool + { + let realPath = URL(fileURLWithPath: path).resolvingSymlinksInPath().path + let sourceAppBundlePath = self.containingAppBundlePath(for: path) + let resolvedAppBundlePath = self.containingAppBundlePath(for: realPath) + let appBundlePath: String? + if let sourceAppBundlePath { + let resolvedSourceBundle = URL(fileURLWithPath: sourceAppBundlePath).resolvingSymlinksInPath().path + guard resolvedAppBundlePath == resolvedSourceBundle else { return false } + appBundlePath = resolvedSourceBundle + } else { + appBundlePath = resolvedAppBundlePath + } + let appBundlePaths = [sourceAppBundlePath, appBundlePath].compactMap(\.self) + let pathsToCheck = [path, realPath] + appBundlePaths + self + .nativeCodexExecutableCandidates( + for: realPath, + fileManager: fileManager) + + for candidate in Set(pathsToCheck) where hasExtendedAttribute(candidate, "com.apple.malware") { + return false + } + + let hasQuarantine = Set(pathsToCheck).contains { hasExtendedAttribute($0, "com.apple.quarantine") } + if let appBundlePath { + guard appSignatureIsTrusted(appBundlePath), + let assessment = spctlAssessment(appBundlePath), + assessment.exitStatus == 0, + self.isAcceptedAssessment(assessment.output, path: appBundlePath), + !self.isExplicitlyBlockedAssessment(assessment.output, path: appBundlePath) + else { + return false + } + return true + } + + guard let native = pathsToCheck.first(where: isMachOExecutable) else { + return !hasQuarantine + } + + guard let assessment = spctlAssessment(native) + else { + return !hasQuarantine + } + + return !self.isExplicitlyBlockedAssessment(assessment.output, path: native) + } + + private static func containingAppBundlePath(for path: String) -> String? { + var candidate = URL(fileURLWithPath: path).standardizedFileURL + while candidate.path != "/" { + if candidate.pathExtension.caseInsensitiveCompare("app") == .orderedSame { + return candidate.path + } + let parent = candidate.deletingLastPathComponent() + guard parent.path != candidate.path else { return nil } + candidate = parent + } + return nil + } + + private static func nativeCodexExecutableCandidates(for path: String, fileManager: FileManager) -> [String] { + let url = URL(fileURLWithPath: path) + guard url.lastPathComponent == "codex.js" else { return [] } + + let packageRoot = url.deletingLastPathComponent().deletingLastPathComponent() + return self.npmNativeCodexCandidates(packageRoot: packageRoot) + .map(\.path) + .filter { fileManager.isExecutableFile(atPath: $0) } + } + + private static func npmNativeCodexCandidates(packageRoot: URL) -> [URL] { + guard let target = self.darwinCodexTarget else { return [] } + let optionalPackage = packageRoot + .appendingPathComponent("node_modules") + .appendingPathComponent("@openai") + .appendingPathComponent(target.packageName) + + return [ + optionalPackage, + packageRoot, + ].map { + $0.appendingPathComponent("vendor") + .appendingPathComponent(target.triple) + .appendingPathComponent("codex") + .appendingPathComponent("codex") + } + } + + private static var darwinCodexTarget: (packageName: String, triple: String)? { + #if arch(arm64) + ("codex-darwin-arm64", "aarch64-apple-darwin") + #elseif arch(x86_64) + ("codex-darwin-x64", "x86_64-apple-darwin") + #else + nil + #endif + } + + private static func hasExtendedAttribute(path: String, name: String) -> Bool { + path.withCString { pathPointer in + name.withCString { namePointer in + getxattr(pathPointer, namePointer, nil, 0, 0, 0) >= 0 + } + } + } + + private static func isMachOExecutable(path: String) -> Bool { + guard let handle = try? FileHandle(forReadingFrom: URL(fileURLWithPath: path)) else { return false } + defer { try? handle.close() } + + guard let data = try? handle.read(upToCount: 4), data.count == 4 else { return false } + let bytes = [UInt8](data) + return bytes == [0xFE, 0xED, 0xFA, 0xCE] || + bytes == [0xCE, 0xFA, 0xED, 0xFE] || + bytes == [0xFE, 0xED, 0xFA, 0xCF] || + bytes == [0xCF, 0xFA, 0xED, 0xFE] || + bytes == [0xCA, 0xFE, 0xBA, 0xBE] || + bytes == [0xCA, 0xFE, 0xBA, 0xBF] + } + + private static func spctlAssessment(path: String, timeout: TimeInterval = 2.0) -> GatekeeperAssessment? { + let spctlPath = "/usr/sbin/spctl" + guard FileManager.default.isExecutableFile(atPath: spctlPath) else { return nil } + + let process = Process() + process.executableURL = URL(fileURLWithPath: spctlPath) + process.arguments = ["--assess", "--type", "execute", "--verbose=4", path] + + let output = Pipe() + process.standardOutput = output + process.standardError = output + + let finished = DispatchSemaphore(value: 0) + process.terminationHandler = { _ in finished.signal() } + + do { + try process.run() + } catch { + return nil + } + + if finished.wait(timeout: .now() + timeout) != .success { + process.terminate() + return nil + } + + let data = output.fileHandleForReading.readDataToEndOfFile() + guard let text = String(data: data, encoding: .utf8) else { return nil } + return GatekeeperAssessment(output: text, exitStatus: process.terminationStatus) + } + + private static func isExpectedOpenAIAppSignature(path: String) -> Bool { + let requirementText = + "identifier \"com.openai.codex\" and anchor apple generic " + + "and certificate leaf[subject.OU] = \"2DC432GLL2\"" + var staticCode: SecStaticCode? + guard SecStaticCodeCreateWithPath( + URL(fileURLWithPath: path) as CFURL, + SecCSFlags(), + &staticCode) == errSecSuccess, + let staticCode + else { + return false + } + + var requirement: SecRequirement? + guard SecRequirementCreateWithString( + requirementText as CFString, + SecCSFlags(), + &requirement) == errSecSuccess + else { + return false + } + + // Pin the publisher and bundle identity here; the Gatekeeper assessment below performs full bundle validation. + let validationFlags = SecCSFlags(rawValue: kSecCSBasicValidateOnly) + return SecStaticCodeCheckValidity(staticCode, validationFlags, requirement) == errSecSuccess + } + + private static func isAcceptedAssessment(_ assessment: String, path: String) -> Bool { + self.assessmentDiagnosticText(assessment, path: path) + .split(whereSeparator: \.isNewline) + .first? + .trimmingCharacters(in: .whitespacesAndNewlines) + .localizedCaseInsensitiveCompare("accepted") == .orderedSame + } + + private static func isExplicitlyBlockedAssessment(_ assessment: String, path: String) -> Bool { + let lower = self.assessmentDiagnosticText(assessment, path: path).lowercased() + if lower.contains("denied") || + lower.contains("cssmerr_tp_cert_revoked") || + lower.contains("revoked") || + lower.contains("malware") || + lower.contains("quarantine") + { + return true + } + if lower.contains("rejected") { + return !lower.contains("code is valid but does not seem to be an app") + } + return false + } + + private static func assessmentDiagnosticText(_ assessment: String, path: String) -> String { + assessment + .split(whereSeparator: \.isNewline) + .enumerated() + .compactMap { offset, line -> String? in + var text = line.trimmingCharacters(in: .whitespacesAndNewlines) + if offset == 0, text.hasPrefix("\(path):") { + text = String(text.dropFirst(path.count + 1)) + .trimmingCharacters(in: .whitespacesAndNewlines) + } + let lower = text.lowercased() + guard !lower.hasPrefix("source="), !lower.hasPrefix("origin=") else { + return nil + } + return text + } + .joined(separator: "\n") + } + #endif +} + +public enum ShellCommandLocator { + #if canImport(Darwin) + private static let shellSpawnFlags = Int16(POSIX_SPAWN_SETSID) + #else + private static let shellSpawnFlags: Int16 = 0x80 // glibc/musl POSIX_SPAWN_SETSID. + #endif + + static func test_runShellCommand( + shell: String, + arguments: [String], + timeout: TimeInterval) -> Data? + { + self.runShellCommand(shell: shell, arguments: arguments, timeout: timeout) + } + + static var test_shellSpawnFlags: Int16 { + self.shellSpawnFlags + } + + public static func commandV( + _ tool: String, + _ shell: String?, + _ timeout: TimeInterval, + _ fileManager: FileManager) -> String? + { + let text = self.runShellCapture(shell, timeout, "command -v \(tool)")? + .trimmingCharacters(in: .whitespacesAndNewlines) + guard let text, !text.isEmpty else { return nil } + + let lines = text.split(separator: "\n").map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + for line in lines.reversed() where line.hasPrefix("/") { + let path = line + if fileManager.isExecutableFile(atPath: path) { + return path + } + } + + return nil + } + + public static func resolveAlias( + _ tool: String, + _ shell: String?, + _ timeout: TimeInterval, + _ fileManager: FileManager, + _ home: String) -> String? + { + let command = "alias \(tool) 2>/dev/null; type -a \(tool) 2>/dev/null" + guard let text = self.runShellCapture(shell, timeout, command) else { return nil } + let lines = text.split(separator: "\n").map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + + if let aliasPath = self.parseAliasPath(lines, tool: tool, home: home, fileManager: fileManager) { + return aliasPath + } + + for line in lines { + if let path = self.extractPathCandidate(line: line, tool: tool, home: home), + fileManager.isExecutableFile(atPath: path) + { + return path + } + } + + return nil + } + + /// Thread-safe buffer for collecting pipe output from a readability handler. + private final class CapturedData: @unchecked Sendable { + private let lock = NSLock() + private var data = Data() + + func append(_ other: Data) { + self.lock.lock() + self.data.append(other) + self.lock.unlock() + } + + func drain() -> Data { + self.lock.lock() + let result = self.data + self.lock.unlock() + return result + } + } + + /// Idempotent one-shot flag — `fire()` returns true exactly once. + /// Used to make `DispatchGroup.leave()` safe to attempt from multiple paths. + private final class OnceFlag: @unchecked Sendable { + private let lock = NSLock() + private var fired = false + + func fire() -> Bool { + self.lock.lock() + defer { self.lock.unlock() } + if self.fired { + return false + } + self.fired = true + return true + } + } + + // swiftlint:disable cyclomatic_complexity function_body_length + /// Runs a shell command, draining both stdout and stderr concurrently so that + /// verbose shell init scripts (oh-my-zsh, nvm, pyenv, etc.) cannot deadlock on + /// a full pipe buffer. The child is launched via `posix_spawn` with + /// `POSIX_SPAWN_SETSID` so it cannot take ownership of the caller's controlling + /// terminal. The new session also makes the child its own process-group leader + /// before `exec`, which guarantees that subsequent `kill(-pgid, ...)` calls reach + /// background helpers spawned by shell init, both after timeout and normal exit. + fileprivate static func runShellCommand( + shell: String, + arguments: [String], + timeout: TimeInterval) -> Data? + { + // Pipes for stdout/stderr. stdin is redirected from /dev/null in the child + // via posix_spawn_file_actions_addopen below. + var stdoutFds: (read: Int32, write: Int32) = (-1, -1) + var stderrFds: (read: Int32, write: Int32) = (-1, -1) + guard withUnsafeMutablePointer(to: &stdoutFds, { + $0.withMemoryRebound(to: Int32.self, capacity: 2) { pipe($0) == 0 } + }) else { return nil } + guard withUnsafeMutablePointer(to: &stderrFds, { + $0.withMemoryRebound(to: Int32.self, capacity: 2) { pipe($0) == 0 } + }) else { + close(stdoutFds.read); close(stdoutFds.write) + return nil + } + + // Build file actions: redirect stdin from /dev/null, dup pipe write ends to + // fds 1 and 2, and close every pipe fd in the child. The init pattern + // differs between platforms because the typedef is an opaque pointer on + // Darwin and a struct on Linux C modules. + #if canImport(Darwin) + var fileActions: posix_spawn_file_actions_t? + #else + var fileActions = posix_spawn_file_actions_t() + #endif + guard posix_spawn_file_actions_init(&fileActions) == 0 else { + close(stdoutFds.read); close(stdoutFds.write) + close(stderrFds.read); close(stderrFds.write) + return nil + } + defer { posix_spawn_file_actions_destroy(&fileActions) } + posix_spawn_file_actions_addopen(&fileActions, 0, "/dev/null", O_RDONLY, 0) + posix_spawn_file_actions_adddup2(&fileActions, stdoutFds.write, 1) + posix_spawn_file_actions_adddup2(&fileActions, stderrFds.write, 2) + posix_spawn_file_actions_addclose(&fileActions, stdoutFds.read) + posix_spawn_file_actions_addclose(&fileActions, stdoutFds.write) + posix_spawn_file_actions_addclose(&fileActions, stderrFds.read) + posix_spawn_file_actions_addclose(&fileActions, stderrFds.write) + + // Build attributes: detach the child into a new session before exec. This + // prevents interactive shell startup from changing the caller's foreground + // process group while retaining a stable process group for cleanup. + #if canImport(Darwin) + var attr: posix_spawnattr_t? + #else + var attr = posix_spawnattr_t() + #endif + guard posix_spawnattr_init(&attr) == 0 else { + close(stdoutFds.read); close(stdoutFds.write) + close(stderrFds.read); close(stderrFds.write) + return nil + } + defer { posix_spawnattr_destroy(&attr) } + guard posix_spawnattr_setflags(&attr, self.shellSpawnFlags) == 0 else { + close(stdoutFds.read); close(stdoutFds.write) + close(stderrFds.read); close(stderrFds.write) + return nil + } + + // Build argv (argv[0] is conventionally the executable path). + var cArgs: [UnsafeMutablePointer?] = [] + cArgs.append(strdup(shell)) + for arg in arguments { + cArgs.append(strdup(arg)) + } + cArgs.append(nil) + defer { + for p in cArgs { + if let p { + free(p) + } + } + } + + // Inherit the parent environment. Build a NULL-terminated `KEY=VALUE` + // array since `extern char **environ` isn't directly visible from Swift. + var cEnv: [UnsafeMutablePointer?] = [] + for (key, value) in ProcessInfo.processInfo.environment { + cEnv.append(strdup("\(key)=\(value)")) + } + cEnv.append(nil) + defer { + for p in cEnv { + if let p { + free(p) + } + } + } + + var pid: pid_t = 0 + let spawnResult = shell.withCString { execPath in + posix_spawn(&pid, execPath, &fileActions, &attr, cArgs, cEnv) + } + + // Close the write ends in the parent so EOF will arrive on the read ends + // once every descendant in the process group also closes them. + close(stdoutFds.write) + close(stderrFds.write) + + guard spawnResult == 0 else { + close(stdoutFds.read); close(stderrFds.read) + return nil + } + + // POSIX_SPAWN_SETSID guarantees the child's session ID and pgid equal its pid. + let pgid: pid_t = pid + + // Track EOF on each pipe so we can wait for full drain instead of sleeping. + // The readability handler fires with empty data when every writer end is + // closed (i.e. the child *and* any inheriting background helpers are gone). + let drainGroup = DispatchGroup() + drainGroup.enter() + drainGroup.enter() + let stdoutDone = OnceFlag() + let stderrDone = OnceFlag() + + let stdoutCollector = CapturedData() + let stdoutHandle = FileHandle(fileDescriptor: stdoutFds.read, closeOnDealloc: true) + stdoutHandle.readabilityHandler = { handle in + let data = handle.availableData + if data.isEmpty { + handle.readabilityHandler = nil + if stdoutDone.fire() { + drainGroup.leave() + } + } else { + stdoutCollector.append(data) + } + } + + let stderrHandle = FileHandle(fileDescriptor: stderrFds.read, closeOnDealloc: true) + stderrHandle.readabilityHandler = { handle in + let data = handle.availableData + if data.isEmpty { + handle.readabilityHandler = nil + if stderrDone.fire() { + drainGroup.leave() + } + } + } + + // Reap the child on a background queue and signal a semaphore on exit. + let exitSemaphore = DispatchSemaphore(value: 0) + let waitPid = pid + DispatchQueue.global(qos: .userInitiated).async { + var status: Int32 = 0 + while waitpid(waitPid, &status, 0) == -1, errno == EINTR { + // retry + } + exitSemaphore.signal() + } + + let finishedInTime = exitSemaphore.wait(timeout: .now() + timeout) == .success + + if !finishedInTime { + kill(-pgid, SIGTERM) + kill(pid, SIGTERM) + if exitSemaphore.wait(timeout: .now() + 0.4) != .success { + kill(-pgid, SIGKILL) + kill(pid, SIGKILL) + _ = exitSemaphore.wait(timeout: .now() + 1.0) + } + stdoutHandle.readabilityHandler = nil + stderrHandle.readabilityHandler = nil + if stdoutDone.fire() { + drainGroup.leave() + } + if stderrDone.fire() { + drainGroup.leave() + } + return nil + } + + // Normal completion — clean up any background children spawned by shell init. + // Without this, helpers that inherited stdout/stderr keep the pipe write ends + // open and we never see EOF on the read ends. + kill(-pgid, SIGTERM) + + // Wait for both pipes to deliver EOF so no buffered bytes are lost. + // Bounded so a stuck handler can't hang the caller indefinitely. + if drainGroup.wait(timeout: .now() + 0.4) != .success { + kill(-pgid, SIGKILL) + } + if drainGroup.wait(timeout: .now() + 0.6) != .success { + stdoutHandle.readabilityHandler = nil + stderrHandle.readabilityHandler = nil + if stdoutDone.fire() { + drainGroup.leave() + } + if stderrDone.fire() { + drainGroup.leave() + } + } + return stdoutCollector.drain() + } + + // swiftlint:enable cyclomatic_complexity function_body_length + + private static func runShellCapture(_ shell: String?, _ timeout: TimeInterval, _ command: String) -> String? { + let shellPath = (shell?.isEmpty == false) ? shell! : "/bin/zsh" + let isCI = ["1", "true"].contains(ProcessInfo.processInfo.environment["CI"]?.lowercased()) + // Interactive login shell to pick up PATH mutations from shell init (nvm/fnm/mise). + // CI runners can have shell init hooks that emit missing CLI errors; avoid them in CI. + let args = isCI ? ["-c", command] : ["-l", "-i", "-c", command] + guard let data = runShellCommand(shell: shellPath, arguments: args, timeout: timeout) else { + return nil + } + return String(data: data, encoding: .utf8) + } + + private static func parseAliasPath( + _ lines: [String], + tool: String, + home: String, + fileManager: FileManager) -> String? + { + for line in lines { + if line.hasPrefix("alias \(tool)=") { + let value = line.replacingOccurrences(of: "alias \(tool)=", with: "") + if let path = self.extractAliasExpansion(value, home: home), + fileManager.isExecutableFile(atPath: path) + { + return path + } + } + if line.lowercased().contains("aliased to") { + if let range = line.range(of: "aliased to") { + let value = line[range.upperBound...].trimmingCharacters(in: .whitespacesAndNewlines) + if let path = self.extractAliasExpansion(String(value), home: home), + fileManager.isExecutableFile(atPath: path) + { + return path + } + } + } + } + return nil + } + + private static func extractAliasExpansion(_ raw: String, home: String) -> String? { + let trimmed = raw.trimmingCharacters(in: CharacterSet(charactersIn: " \t\"'`")) + guard !trimmed.isEmpty else { return nil } + let parts = trimmed.split(separator: " ").map(String.init) + guard let first = parts.first else { return nil } + return self.expandPath(first, home: home) + } + + private static func extractPathCandidate(line: String, tool: String, home: String) -> String? { + let tokens = line.split(whereSeparator: { $0 == " " || $0 == "\t" }).map(String.init) + for token in tokens { + let candidate = self.expandPath(token, home: home) + if candidate.hasPrefix("/"), + URL(fileURLWithPath: candidate).lastPathComponent == tool + { + return candidate + } + } + return nil + } + + private static func expandPath(_ raw: String, home: String) -> String { + if raw == "~" { + return home + } + if raw.hasPrefix("~/") { + return home + String(raw.dropFirst()) + } + return raw + } +} + +public enum PathBuilder { + public static func effectivePATH( + purposes _: Set, + env: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current, + home _: String = NSHomeDirectory()) -> String + { + var parts: [String] = [] + + if let loginPATH, !loginPATH.isEmpty { + parts.append(contentsOf: loginPATH) + } + + if let existing = env["PATH"], !existing.isEmpty { + parts.append(contentsOf: existing.split(separator: ":").map(String.init)) + } + + if parts.isEmpty { + parts.append(contentsOf: ["/usr/bin", "/bin", "/usr/sbin", "/sbin"]) + } + + var seen = Set() + let deduped = parts.compactMap { part -> String? in + guard !part.isEmpty else { return nil } + if seen.insert(part).inserted { + return part + } + return nil + } + + return deduped.joined(separator: ":") + } + + public static func debugSnapshot( + purposes: Set, + env: [String: String] = ProcessInfo.processInfo.environment, + home: String = NSHomeDirectory()) -> PathDebugSnapshot + { + let login = LoginShellPathCache.shared.current + let effective = self.effectivePATH( + purposes: purposes, + env: env, + loginPATH: login, + home: home) + let codex = BinaryLocator.resolveCodexBinary(env: env, loginPATH: login, home: home) + let claude = BinaryLocator.resolveClaudeBinary(env: env, loginPATH: login, home: home) + let gemini = BinaryLocator.resolveGeminiBinary(env: env, loginPATH: login, home: home) + let loginString = login?.joined(separator: ":") + return PathDebugSnapshot( + codexBinary: codex, + claudeBinary: claude, + geminiBinary: gemini, + effectivePATH: effective, + loginShellPATH: loginString) + } + + public static func debugSnapshotAsync( + purposes: Set, + env: [String: String] = ProcessInfo.processInfo.environment, + home: String = NSHomeDirectory()) async -> PathDebugSnapshot + { + await Task.detached(priority: .userInitiated) { + self.debugSnapshot(purposes: purposes, env: env, home: home) + }.value + } +} + +enum LoginShellPathCapturer { + static let defaultTimeout: TimeInterval = 6.0 + + static func capture( + shell: String? = ProcessInfo.processInfo.environment["SHELL"], + timeout: TimeInterval = Self.defaultTimeout) -> [String]? + { + let shellPath = (shell?.isEmpty == false) ? shell! : "/bin/zsh" + let isCI = ["1", "true"].contains(ProcessInfo.processInfo.environment["CI"]?.lowercased()) + let marker = "__CODEXBAR_PATH__" + // Skip interactive login shells in CI to avoid noisy init hooks. + let args = isCI + ? ["-c", "printf '\(marker)%s\(marker)' \"$PATH\""] + : ["-l", "-i", "-c", "printf '\(marker)%s\(marker)' \"$PATH\""] + guard let data = ShellCommandLocator.runShellCommand( + shell: shellPath, + arguments: args, + timeout: timeout), + let raw = String(data: data, encoding: .utf8), + !raw.isEmpty + else { return nil } + + let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines) + let extracted = if let start = trimmed.range(of: marker), + let end = trimmed.range(of: marker, options: .backwards), + start.upperBound <= end.lowerBound + { + String(trimmed[start.upperBound.. [String]? + private var captured: [String]? + private var isCapturing = false + private var callbacks: [([String]?) -> Void] = [] + + init(capture: @escaping @Sendable (String?, TimeInterval) -> [String]? = LoginShellPathCapturer.capture) { + self.capture = capture + } + + public var current: [String]? { + self.lock.lock() + let value = self.captured + self.lock.unlock() + return value + } + + public func captureOnce( + shell: String? = ProcessInfo.processInfo.environment["SHELL"], + timeout: TimeInterval = 6.0, + onFinish: (([String]?) -> Void)? = nil) + { + self.lock.lock() + if let captured { + self.lock.unlock() + onFinish?(captured) + return + } + + if let onFinish { + self.callbacks.append(onFinish) + } + + if self.isCapturing { + self.lock.unlock() + return + } + + self.isCapturing = true + self.lock.unlock() + + let capture = self.capture + DispatchQueue.global(qos: .utility).async { [weak self] in + let result = capture(shell, timeout) + guard let self else { return } + + self.lock.lock() + self.captured = result + self.isCapturing = false + let callbacks = self.callbacks + self.callbacks.removeAll() + self.lock.unlock() + + callbacks.forEach { $0(result) } + } + } + + public func currentOrCapture( + shell: String? = ProcessInfo.processInfo.environment["SHELL"], + timeout: TimeInterval = 6.0) -> [String]? + { + self.lock.lock() + if let captured { + self.lock.unlock() + return captured + } + + if self.isCapturing { + let semaphore = DispatchSemaphore(value: 0) + var callbackResult: [String]? + self.callbacks.append { result in + callbackResult = result + semaphore.signal() + } + self.lock.unlock() + let deadline = DispatchTime.now() + timeout + _ = semaphore.wait(timeout: deadline) + return callbackResult ?? self.current + } + + self.isCapturing = true + self.lock.unlock() + + let result = self.capture(shell, timeout) + self.lock.lock() + self.captured = result + self.isCapturing = false + let callbacks = self.callbacks + self.callbacks.removeAll() + self.lock.unlock() + + callbacks.forEach { $0(result) } + return result + } +} diff --git a/Sources/CodexBarCore/PiSessionCostCache.swift b/Sources/CodexBarCore/PiSessionCostCache.swift new file mode 100644 index 0000000..355be9c --- /dev/null +++ b/Sources/CodexBarCore/PiSessionCostCache.swift @@ -0,0 +1,97 @@ +import Foundation + +enum PiSessionCostCacheIO { + /// Artifact schema version. Pricing changes are tracked separately by `pricingKey`. + private static let artifactVersion = 5 + + private static func defaultCacheRoot() -> URL { + let root = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask).first! + return root.appendingPathComponent("CodexBar", isDirectory: true) + } + + static func cacheFileURL(cacheRoot: URL? = nil) -> URL { + let root = cacheRoot ?? self.defaultCacheRoot() + return root + .appendingPathComponent("cost-usage", isDirectory: true) + .appendingPathComponent("pi-sessions-v\(Self.artifactVersion).json", isDirectory: false) + } + + static func load(cacheRoot: URL? = nil) -> PiSessionCostCache { + let url = self.cacheFileURL(cacheRoot: cacheRoot) + guard let data = try? Data(contentsOf: url), + let decoded = try? JSONDecoder().decode(PiSessionCostCache.self, from: data), + decoded.version == Self.artifactVersion + else { + return PiSessionCostCache(version: Self.artifactVersion) + } + return decoded + } + + static func save(cache: PiSessionCostCache, cacheRoot: URL? = nil) { + let url = self.cacheFileURL(cacheRoot: cacheRoot) + let dir = url.deletingLastPathComponent() + try? FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + + let tmp = dir.appendingPathComponent(".tmp-\(UUID().uuidString).json", isDirectory: false) + let data = (try? JSONEncoder().encode(cache)) ?? Data() + do { + try data.write(to: tmp, options: [.atomic]) + if FileManager.default.fileExists(atPath: url.path) { + _ = try FileManager.default.replaceItemAt(url, withItemAt: tmp) + } else { + try FileManager.default.moveItem(at: tmp, to: url) + } + } catch { + try? FileManager.default.removeItem(at: tmp) + } + } +} + +struct PiSessionCostCache: Codable { + var version: Int + var lastScanUnixMs: Int64 = 0 + var scanSinceKey: String? + var scanUntilKey: String? + var pricingKey: String? + var daysByProvider: [String: [String: [String: PiPackedUsage]]] = [:] + var files: [String: PiSessionFileUsage] = [:] + + init(version: Int = 5) { + self.version = version + } +} + +struct PiSessionFileUsage: Codable { + var mtimeUnixMs: Int64 + var size: Int64 + var parsedBytes: Int64 + var lastModelContext: PiModelContext? + var contributions: [String: [String: [String: PiPackedUsage]]] +} + +struct PiModelContext: Codable, Equatable { + var providerRawValue: String + var modelName: String +} + +struct PiPackedUsage: Codable, Equatable { + var inputTokens: Int = 0 + var cacheReadTokens: Int = 0 + var cacheWriteTokens: Int = 0 + var outputTokens: Int = 0 + var totalTokens: Int = 0 + var costNanos: Int64 = 0 + var costSampleCount: Int = 0 + var usageSampleCount: Int? + + var isZero: Bool { + self.inputTokens == 0 + && self.cacheReadTokens == 0 + && self.cacheWriteTokens == 0 + && self.outputTokens == 0 + && self.totalTokens == 0 + && self.costNanos == 0 + && self.costSampleCount == 0 + && (self.usageSampleCount ?? 0) == 0 + } +} diff --git a/Sources/CodexBarCore/PiSessionCostScanner.swift b/Sources/CodexBarCore/PiSessionCostScanner.swift new file mode 100644 index 0000000..27b3568 --- /dev/null +++ b/Sources/CodexBarCore/PiSessionCostScanner.swift @@ -0,0 +1,943 @@ +import Foundation + +private final class PiSessionISO8601FormatterBox: @unchecked Sendable { + let lock = NSLock() + let withFractional: ISO8601DateFormatter = { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + return formatter + }() + + let plain: ISO8601DateFormatter = { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime] + return formatter + }() +} + +enum PiSessionCostScanner { + struct Options { + var piSessionsRoot: URL? + var cacheRoot: URL? + var refreshMinIntervalSeconds: TimeInterval = 60 + var forceRescan: Bool = false + + init( + piSessionsRoot: URL? = nil, + cacheRoot: URL? = nil, + refreshMinIntervalSeconds: TimeInterval = 60, + forceRescan: Bool = false) + { + self.piSessionsRoot = piSessionsRoot + self.cacheRoot = cacheRoot + self.refreshMinIntervalSeconds = refreshMinIntervalSeconds + self.forceRescan = forceRescan + } + } + + private struct ParseResult { + let contributions: [String: [String: [String: PiPackedUsage]]] + let parsedBytes: Int64 + let lastModelContext: PiModelContext? + } + + private struct AssistantIdentity { + let provider: UsageProvider + let modelName: String + } + + private struct ModelsDevPricingContext { + let catalog: ModelsDevCatalog? + let cacheRoot: URL? + let pricingKey: String + } + + private struct ScanContext { + let range: CostUsageScanner.CostUsageDayRange + let forceRescan: Bool + let pricingContext: ModelsDevPricingContext + let checkCancellation: CostUsageScanner.CancellationCheck? + } + + private static let costScale = 1_000_000_000.0 + /// Bump for Pi-only cost formula changes not represented by the parser or pricing fingerprints. + private static let costFormulaVersion = 1 + private static let maxLineBytes = 16 * 1024 * 1024 + private static let maxSafeRoundedInt = Double(Int.max) - 1 + private static let sessionStartFilenameRegex = try? NSRegularExpression( + pattern: "^(\\d{4}-\\d{2}-\\d{2})T(\\d{2})-(\\d{2})-(\\d{2})-(\\d{3})Z_") + private static let isoFormatterBox = PiSessionISO8601FormatterBox() + + static func loadDailyReport( + provider: UsageProvider, + since: Date, + until: Date, + now: Date = Date(), + options: Options = Options()) -> CostUsageDailyReport + { + ( + try? self.loadDailyReportCancellable( + provider: provider, + since: since, + until: until, + now: now, + options: options, + checkCancellation: nil)) ?? CostUsageDailyReport(data: [], summary: nil) + } + + static func loadDailyReportCancellable( + provider: UsageProvider, + since: Date, + until: Date, + now: Date = Date(), + options: Options = Options(), + checkCancellation: CostUsageScanner.CancellationCheck?) throws -> CostUsageDailyReport + { + guard provider == .codex || provider == .claude else { + return CostUsageDailyReport(data: [], summary: nil) + } + + let range = CostUsageScanner.CostUsageDayRange(since: since, until: until) + var cache = PiSessionCostCacheIO.load(cacheRoot: options.cacheRoot) + let nowMs = Int64(now.timeIntervalSince1970 * 1000) + let refreshMs = Int64(max(0, options.refreshMinIntervalSeconds) * 1000) + let pricingContext = self.pricingContext(now: now, cacheRoot: options.cacheRoot) + let windowExpanded = self.requestedWindowExpandsCache(range: range, cache: cache) + let pricingChanged = cache.pricingKey != pricingContext.pricingKey + let shouldRefresh = options.forceRescan + || windowExpanded + || pricingChanged + || refreshMs == 0 + || cache.lastScanUnixMs == 0 + || nowMs - cache.lastScanUnixMs > refreshMs + + if shouldRefresh { + try checkCancellation?() + let root = self.defaultPiSessionsRoot(options: options) + let startCutoff = self.dateFromDayKey(range.scanSinceKey) ?? since + let files = self.listPiSessionFiles(root: root, startCutoffLocal: startCutoff) + let filePathsInScan = Set(files.map(\.path)) + + for fileURL in files { + try self.scanPiSessionFile( + fileURL: fileURL, + cache: &cache, + context: ScanContext( + range: range, + forceRescan: options.forceRescan || windowExpanded || pricingChanged, + pricingContext: pricingContext, + checkCancellation: checkCancellation)) + } + try checkCancellation?() + + for key in cache.files.keys where !filePathsInScan.contains(key) { + if let old = cache.files[key] { + self.applyContributions( + daysByProvider: &cache.daysByProvider, + contributions: old.contributions, + sign: -1) + } + cache.files.removeValue(forKey: key) + } + + cache.scanSinceKey = range.scanSinceKey + cache.scanUntilKey = range.scanUntilKey + cache.pricingKey = pricingContext.pricingKey + cache.lastScanUnixMs = nowMs + try checkCancellation?() + PiSessionCostCacheIO.save(cache: cache, cacheRoot: options.cacheRoot) + } + + return self.buildReport( + provider: provider, + cache: cache, + range: range, + pricingContext: pricingContext) + } + + struct CachedDailyReportResult { + let report: CostUsageDailyReport + let lastScanAt: Date? + } + + static func loadCachedDailyReport( + provider: UsageProvider, + since: Date, + until: Date, + now: Date = Date(), + cacheRoot: URL? = nil) -> CostUsageDailyReport? + { + self.loadCachedDailyReportResult( + provider: provider, + since: since, + until: until, + now: now, + cacheRoot: cacheRoot)?.report + } + + static func loadCachedDailyReportResult( + provider: UsageProvider, + since: Date, + until: Date, + now: Date = Date(), + cacheRoot: URL? = nil) -> CachedDailyReportResult? + { + guard provider == .codex || provider == .claude else { return nil } + + let range = CostUsageScanner.CostUsageDayRange(since: since, until: until) + let cache = PiSessionCostCacheIO.load(cacheRoot: cacheRoot) + guard !cache.daysByProvider.isEmpty else { return nil } + guard !self.requestedWindowExpandsCache(range: range, cache: cache) else { return nil } + + let pricingContext = self.pricingContext(now: now, cacheRoot: cacheRoot) + guard cache.pricingKey == pricingContext.pricingKey else { return nil } + let report = self.buildReport( + provider: provider, + cache: cache, + range: range, + pricingContext: pricingContext) + guard !report.data.isEmpty else { return nil } + let lastScanAt = cache.lastScanUnixMs > 0 + ? Date(timeIntervalSince1970: TimeInterval(cache.lastScanUnixMs) / 1000) + : nil + return CachedDailyReportResult(report: report, lastScanAt: lastScanAt) + } + + private static func pricingContext(now: Date, cacheRoot: URL?) -> ModelsDevPricingContext { + let modelsDevArtifact = ModelsDevCache.load(now: now, cacheRoot: cacheRoot).artifact + return ModelsDevPricingContext( + catalog: modelsDevArtifact?.catalog, + cacheRoot: cacheRoot, + pricingKey: CostUsagePricingKey.codex( + modelsDevArtifact: modelsDevArtifact, + formulaVersion: Self.costFormulaVersion, + parserHash: CodexParserHash.value, + modelsDevProviderIDs: ["anthropic", "openai"])) + } + + private static func requestedWindowExpandsCache( + range: CostUsageScanner.CostUsageDayRange, + cache: PiSessionCostCache) -> Bool + { + guard let cachedSince = cache.scanSinceKey, + let cachedUntil = cache.scanUntilKey + else { + return true + } + + if range.scanSinceKey < cachedSince { + return true + } + if range.scanUntilKey > cachedUntil { + return true + } + return false + } + + private static func defaultPiSessionsRoot(options: Options) -> URL { + if let override = options.piSessionsRoot { return override } + return FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent(".pi", isDirectory: true) + .appendingPathComponent("agent", isDirectory: true) + .appendingPathComponent("sessions", isDirectory: true) + } + + private static func listPiSessionFiles(root: URL, startCutoffLocal: Date) -> [URL] { + guard FileManager.default.fileExists(atPath: root.path) else { return [] } + + let keys: Set = [.isRegularFileKey, .contentModificationDateKey] + guard let enumerator = FileManager.default.enumerator( + at: root, + includingPropertiesForKeys: Array(keys), + options: [.skipsHiddenFiles]) + else { + return [] + } + + var output: [URL] = [] + while let item = enumerator.nextObject() as? URL { + guard item.pathExtension.lowercased() == "jsonl" else { continue } + let values = try? item.resourceValues(forKeys: keys) + guard values?.isRegularFile == true else { continue } + + let startedAt = self.parseSessionStartFromFilename(item.lastPathComponent) + let modifiedAt = values?.contentModificationDate + if self + .shouldIncludeFile(startedAt: startedAt, modifiedAt: modifiedAt, startCutoffLocal: startCutoffLocal) + { + output.append(item) + } + } + + return output.sorted(by: { $0.path < $1.path }) + } + + private static func shouldIncludeFile( + startedAt: Date?, + modifiedAt: Date?, + startCutoffLocal: Date) -> Bool + { + if let modifiedAt, self.localMidnight(modifiedAt) >= startCutoffLocal { + return true + } + if let startedAt, self.localMidnight(startedAt) >= startCutoffLocal { + return true + } + return false + } + + private static func scanPiSessionFile( + fileURL: URL, + cache: inout PiSessionCostCache, + context: ScanContext) + throws + { + try context.checkCancellation?() + let path = fileURL.path + let attrs = (try? FileManager.default.attributesOfItem(atPath: path)) ?? [:] + let mtime = (attrs[.modificationDate] as? Date)?.timeIntervalSince1970 ?? 0 + let size = (attrs[.size] as? NSNumber)?.int64Value ?? 0 + let mtimeMs = Int64(mtime * 1000) + + func storeFileUsage(_ usage: PiSessionFileUsage) { + cache.files[path] = usage + } + + let cached = cache.files[path] + if !context.forceRescan, + let cached, + cached.mtimeUnixMs == mtimeMs, + cached.size == size + { + return + } + + if !context.forceRescan, + let cached, + size > cached.size, + cached.parsedBytes > 0, + cached.parsedBytes <= size + { + let delta = try self.parsePiSessionFile( + fileURL: fileURL, + range: context.range, + startOffset: cached.parsedBytes, + initialModelContext: cached.lastModelContext, + pricingContext: context.pricingContext, + checkCancellation: context.checkCancellation) + if !delta.contributions.isEmpty { + self.applyContributions( + daysByProvider: &cache.daysByProvider, + contributions: delta.contributions, + sign: 1) + } + let merged = self.mergedContributions(existing: cached.contributions, delta: delta.contributions) + storeFileUsage(PiSessionFileUsage( + mtimeUnixMs: mtimeMs, + size: size, + parsedBytes: delta.parsedBytes, + lastModelContext: delta.lastModelContext, + contributions: merged)) + return + } + + if let cached { + self.applyContributions( + daysByProvider: &cache.daysByProvider, + contributions: cached.contributions, + sign: -1) + } + + let parsed = try self.parsePiSessionFile( + fileURL: fileURL, + range: context.range, + pricingContext: context.pricingContext, + checkCancellation: context.checkCancellation) + if !parsed.contributions.isEmpty { + self.applyContributions(daysByProvider: &cache.daysByProvider, contributions: parsed.contributions, sign: 1) + } + + storeFileUsage(PiSessionFileUsage( + mtimeUnixMs: mtimeMs, + size: size, + parsedBytes: parsed.parsedBytes, + lastModelContext: parsed.lastModelContext, + contributions: parsed.contributions)) + } + + private static func parsePiSessionFile( + fileURL: URL, + range: CostUsageScanner.CostUsageDayRange, + startOffset: Int64 = 0, + initialModelContext: PiModelContext? = nil, + pricingContext: ModelsDevPricingContext? = nil, + checkCancellation: CostUsageScanner.CancellationCheck? = nil) throws -> ParseResult + { + var currentModelContext = initialModelContext + var contributions: [String: [String: [String: PiPackedUsage]]] = [:] + + func add(provider: UsageProvider, dayKey: String, modelName: String, usage: PiPackedUsage) { + guard !usage.isZero else { return } + guard CostUsageScanner.CostUsageDayRange.isInRange( + dayKey: dayKey, + since: range.scanSinceKey, + until: range.scanUntilKey) + else { + return + } + + let providerKey = provider.rawValue + var providerDays = contributions[providerKey] ?? [:] + var dayModels = providerDays[dayKey] ?? [:] + let merged = self.addPacked(a: dayModels[modelName] ?? PiPackedUsage(), b: usage, sign: 1) + if merged.isZero { + dayModels.removeValue(forKey: modelName) + } else { + dayModels[modelName] = merged + } + if dayModels.isEmpty { + providerDays.removeValue(forKey: dayKey) + } else { + providerDays[dayKey] = dayModels + } + if providerDays.isEmpty { + contributions.removeValue(forKey: providerKey) + } else { + contributions[providerKey] = providerDays + } + } + + let parsedBytes: Int64 + do { + parsedBytes = try CostUsageJsonl.scan( + fileURL: fileURL, + offset: startOffset, + maxLineBytes: Self.maxLineBytes, + prefixBytes: Self.maxLineBytes, + checkCancellation: checkCancellation, + onLine: { line in + guard !line.bytes.isEmpty, !line.wasTruncated else { return } + autoreleasepool { + guard let object = (try? JSONSerialization.jsonObject(with: line.bytes)) as? [String: Any] + else { return } + guard let type = object["type"] as? String else { return } + + if type == "model_change" { + currentModelContext = self.modelContext(from: object) + return + } + + guard type == "message", let message = object["message"] as? [String: Any] else { return } + guard (message["role"] as? String) == "assistant" else { return } + + let identity = self.resolveAssistantIdentity( + entry: object, + message: message, + fallback: currentModelContext) + guard let identity else { return } + guard let date = self.timestampDate(entry: object, message: message) else { return } + let dayKey = CostUsageScanner.CostUsageDayRange.dayKey(from: date) + let usage = self.extractUsage( + provider: identity.provider, + modelName: identity.modelName, + message: message, + pricingDate: date, + pricingContext: pricingContext) + add(provider: identity.provider, dayKey: dayKey, modelName: identity.modelName, usage: usage) + } + }) + } catch is CancellationError { + throw CancellationError() + } catch { + parsedBytes = startOffset + } + + return ParseResult( + contributions: contributions, + parsedBytes: parsedBytes, + lastModelContext: currentModelContext) + } + + private static func modelContext(from object: [String: Any]) -> PiModelContext? { + guard let providerText = object["provider"] as? String, + let provider = self.mappedProvider(fromPiProvider: providerText) + else { + return nil + } + let rawModelName = (object["modelId"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + guard let modelName = self.normalizeModelName(rawModelName, provider: provider) else { return nil } + return PiModelContext(providerRawValue: provider.rawValue, modelName: modelName) + } + + private static func resolveAssistantIdentity( + entry: [String: Any], + message: [String: Any], + fallback: PiModelContext?) -> AssistantIdentity? + { + let explicitProviderText = self.extractProviderText(entry: entry, message: message) + let explicitProvider = explicitProviderText.flatMap(self.mappedProvider(fromPiProvider:)) + let explicitModelText = self.extractModelText(entry: entry, message: message) + + if explicitProviderText != nil, explicitProvider == nil { + return nil + } + + if let explicitProvider, + let explicitModelText, + let explicitModel = self.normalizeModelName(explicitModelText, provider: explicitProvider) + { + return AssistantIdentity(provider: explicitProvider, modelName: explicitModel) + } + + if let explicitProvider, + let fallback, + fallback.providerRawValue == explicitProvider.rawValue + { + return AssistantIdentity(provider: explicitProvider, modelName: fallback.modelName) + } + + if explicitProviderText == nil, + let explicitModelText, + let fallbackProvider = fallback.flatMap({ UsageProvider(rawValue: $0.providerRawValue) }), + let explicitModel = self.normalizeModelName(explicitModelText, provider: fallbackProvider) + { + return AssistantIdentity(provider: fallbackProvider, modelName: explicitModel) + } + + if explicitProviderText == nil, + let fallback, + let provider = UsageProvider(rawValue: fallback.providerRawValue) + { + return AssistantIdentity(provider: provider, modelName: fallback.modelName) + } + + return nil + } + + private static func extractProviderText(entry: [String: Any], message: [String: Any]) -> String? { + if let provider = (message["provider"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines), + !provider.isEmpty + { + return provider + } + if let provider = (entry["provider"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines), + !provider.isEmpty + { + return provider + } + return nil + } + + private static func extractModelText(entry: [String: Any], message: [String: Any]) -> String? { + for value in [message["model"], entry["model"], message["modelId"], entry["modelId"]] { + if let model = (value as? String)?.trimmingCharacters(in: .whitespacesAndNewlines), !model.isEmpty { + return model + } + } + return nil + } + + private static func normalizeModelName(_ raw: String, provider: UsageProvider) -> String? { + let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + return switch provider { + case .codex: + CostUsagePricing.normalizeCodexModel(trimmed) + case .claude: + CostUsagePricing.normalizeClaudeModel(trimmed) + default: + trimmed + } + } + + private static func timestampDate(entry: [String: Any], message: [String: Any]) -> Date? { + self.parseTimestampValue(message["timestamp"]) + ?? self.parseTimestampValue(entry["timestamp"]) + } + + private static func parseTimestampValue(_ value: Any?) -> Date? { + if let number = value as? NSNumber { + let raw = number.doubleValue + guard raw.isFinite else { return nil } + if raw > 1_000_000_000_000 { + return Date(timeIntervalSince1970: raw / 1000) + } + return Date(timeIntervalSince1970: raw) + } + + if let string = value as? String { + if let numeric = Double(string), numeric.isFinite { + if numeric > 1_000_000_000_000 { + return Date(timeIntervalSince1970: numeric / 1000) + } + return Date(timeIntervalSince1970: numeric) + } + return self.parseISO(string) + } + + return nil + } + + private static func extractUsage( + provider: UsageProvider, + modelName: String, + message: [String: Any], + pricingDate: Date? = nil, + pricingContext: ModelsDevPricingContext? = nil) -> PiPackedUsage + { + let usage = (message["usage"] as? [String: Any]) ?? [:] + let input = self.readNonNegativeInt( + usage["input"] + ?? usage["inputTokens"] + ?? usage["input_tokens"] + ?? usage["promptTokens"] + ?? usage["prompt_tokens"]) + let cacheRead = self.readNonNegativeInt( + usage["cacheRead"] + ?? usage["cacheReadTokens"] + ?? usage["cache_read"] + ?? usage["cache_read_tokens"] + ?? usage["cacheReadInputTokens"] + ?? usage["cache_read_input_tokens"]) + let cacheWrite = self.readNonNegativeInt( + usage["cacheWrite"] + ?? usage["cacheWriteTokens"] + ?? usage["cache_write"] + ?? usage["cache_write_tokens"] + ?? usage["cacheCreationTokens"] + ?? usage["cache_creation_tokens"] + ?? usage["cacheCreationInputTokens"] + ?? usage["cache_creation_input_tokens"]) + let output = self.readNonNegativeInt( + usage["output"] + ?? usage["outputTokens"] + ?? usage["output_tokens"] + ?? usage["completionTokens"] + ?? usage["completion_tokens"]) + + let directTotal = self.readNonNegativeInt( + usage["totalTokens"] + ?? usage["total_tokens"] + ?? usage["tokenCount"] + ?? usage["token_count"] + ?? usage["tokens"]) + let derivedTotal = input + cacheRead + cacheWrite + output + let totalTokens = max(directTotal, derivedTotal) + + let rawUsage = PiPackedUsage( + inputTokens: input, + cacheReadTokens: cacheRead, + cacheWriteTokens: cacheWrite, + outputTokens: output, + totalTokens: totalTokens) + // Pi JSONL does not record Anthropic cache retention, so use Pi's persisted default tariff. + let costUSD = self.computedCostUSD( + provider: provider, + modelName: modelName, + usage: rawUsage, + pricingDate: pricingDate, + pricingContext: pricingContext) + let costNanos = costUSD.map { Int64(($0 * self.costScale).rounded()) } ?? 0 + + return PiPackedUsage( + inputTokens: rawUsage.inputTokens, + cacheReadTokens: rawUsage.cacheReadTokens, + cacheWriteTokens: rawUsage.cacheWriteTokens, + outputTokens: rawUsage.outputTokens, + totalTokens: rawUsage.totalTokens, + costNanos: costNanos, + costSampleCount: costUSD == nil ? 0 : 1, + usageSampleCount: 1) + } + + private static func computedCostUSD( + provider: UsageProvider, + modelName: String, + usage: PiPackedUsage, + pricingDate: Date? = nil, + pricingContext: ModelsDevPricingContext? = nil) -> Double? + { + switch provider { + case .codex: + // Pi records input, cache reads, and cache writes as disjoint counts. Codex pricing + // expects cached/write tokens to be subsets of total input, so reconstruct that total + // here and pass writes separately (1.25x input for GPT-5.6 when rates are known). + CostUsagePricing.codexCostUSD( + model: modelName, + inputTokens: usage.inputTokens + usage.cacheReadTokens + usage.cacheWriteTokens, + cachedInputTokens: usage.cacheReadTokens, + outputTokens: usage.outputTokens, + cacheWriteInputTokens: usage.cacheWriteTokens, + modelsDevCatalog: pricingContext?.catalog, + modelsDevCacheRoot: pricingContext?.cacheRoot) + case .claude: + CostUsagePricing.claudeCostUSD( + model: modelName, + inputTokens: usage.inputTokens, + cacheReadInputTokens: usage.cacheReadTokens, + cacheCreationInputTokens: usage.cacheWriteTokens, + outputTokens: usage.outputTokens, + pricingDate: pricingDate, + modelsDevCatalog: pricingContext?.catalog, + modelsDevCacheRoot: pricingContext?.cacheRoot) + default: + nil + } + } + + private static func readNonNegativeInt(_ value: Any?) -> Int { + if let number = value as? NSNumber { + let numeric = number.doubleValue + guard numeric.isFinite, numeric >= 0, numeric <= self.maxSafeRoundedInt else { return 0 } + return Int(numeric.rounded()) + } + if let string = value as? String, + let numeric = Double(string), + numeric.isFinite, + numeric >= 0, + numeric <= self.maxSafeRoundedInt + { + return Int(numeric.rounded()) + } + return 0 + } +} + +extension PiSessionCostScanner { + private static func mappedProvider(fromPiProvider provider: String) -> UsageProvider? { + switch provider.lowercased() { + case "openai-codex": + .codex + case "anthropic": + .claude + default: + nil + } + } + + private static func buildReport( + provider: UsageProvider, + cache: PiSessionCostCache, + range: CostUsageScanner.CostUsageDayRange, + pricingContext: ModelsDevPricingContext? = nil) -> CostUsageDailyReport + { + guard let providerDays = cache.daysByProvider[provider.rawValue] else { + return CostUsageDailyReport(data: [], summary: nil) + } + + let dayKeys = providerDays.keys.sorted().filter { + CostUsageScanner.CostUsageDayRange.isInRange(dayKey: $0, since: range.sinceKey, until: range.untilKey) + } + + var entries: [CostUsageDailyReport.Entry] = [] + var totalInput = 0 + var totalOutput = 0 + var totalCacheRead = 0 + var totalCacheWrite = 0 + var totalTokens = 0 + var totalCostNanos: Int64 = 0 + var totalCostSamples = 0 + + for dayKey in dayKeys { + guard let models = providerDays[dayKey] else { continue } + let modelNames = models.keys.sorted() + + var dayInput = 0 + var dayOutput = 0 + var dayCacheRead = 0 + var dayCacheWrite = 0 + var dayTotalTokens = 0 + var dayCostNanos: Int64 = 0 + var dayCostSamples = 0 + var breakdown: [CostUsageDailyReport.ModelBreakdown] = [] + + for modelName in modelNames { + let packed = models[modelName] ?? PiPackedUsage() + let modelTotalTokens = max( + packed.totalTokens, + packed.inputTokens + packed.cacheReadTokens + packed.cacheWriteTokens + packed.outputTokens) + let currentPricingCost = self.computedCostUSD( + provider: provider, + modelName: modelName, + usage: packed, + pricingContext: pricingContext) + let usageSampleCount = packed.usageSampleCount + let hasCompleteCachedCost = (usageSampleCount ?? 0) > 0 + && packed.costSampleCount == usageSampleCount + // Cached costs are accumulated per message, which preserves Claude long-context threshold boundaries. + let costNanos = hasCompleteCachedCost + ? packed.costNanos + : currentPricingCost.map { Int64(($0 * self.costScale).rounded()) } + breakdown.append(CostUsageDailyReport.ModelBreakdown( + modelName: modelName, + costUSD: costNanos.map { Double($0) / Self.costScale }, + totalTokens: modelTotalTokens > 0 ? modelTotalTokens : nil)) + dayInput += packed.inputTokens + dayOutput += packed.outputTokens + dayCacheRead += packed.cacheReadTokens + dayCacheWrite += packed.cacheWriteTokens + dayTotalTokens += modelTotalTokens + if let costNanos { + dayCostNanos += costNanos + dayCostSamples += 1 + } + } + + let sortedBreakdown = self.sortedModelBreakdowns(breakdown) + entries.append(CostUsageDailyReport.Entry( + date: dayKey, + inputTokens: dayInput > 0 ? dayInput : nil, + outputTokens: dayOutput > 0 ? dayOutput : nil, + cacheReadTokens: dayCacheRead > 0 ? dayCacheRead : nil, + cacheCreationTokens: dayCacheWrite > 0 ? dayCacheWrite : nil, + totalTokens: dayTotalTokens > 0 ? dayTotalTokens : nil, + costUSD: dayCostSamples > 0 ? Double(dayCostNanos) / Self.costScale : nil, + modelsUsed: modelNames, + modelBreakdowns: sortedBreakdown)) + + totalInput += dayInput + totalOutput += dayOutput + totalCacheRead += dayCacheRead + totalCacheWrite += dayCacheWrite + totalTokens += dayTotalTokens + totalCostNanos += dayCostNanos + totalCostSamples += dayCostSamples + } + + guard !entries.isEmpty else { return CostUsageDailyReport(data: [], summary: nil) } + return CostUsageDailyReport( + data: entries, + summary: CostUsageDailyReport.Summary( + totalInputTokens: totalInput > 0 ? totalInput : nil, + totalOutputTokens: totalOutput > 0 ? totalOutput : nil, + cacheReadTokens: totalCacheRead > 0 ? totalCacheRead : nil, + cacheCreationTokens: totalCacheWrite > 0 ? totalCacheWrite : nil, + totalTokens: totalTokens > 0 ? totalTokens : nil, + totalCostUSD: totalCostSamples > 0 ? Double(totalCostNanos) / Self.costScale : nil)) + } + + private static func mergedContributions( + existing: [String: [String: [String: PiPackedUsage]]], + delta: [String: [String: [String: PiPackedUsage]]]) -> [String: [String: [String: PiPackedUsage]]] + { + var merged = existing + self.applyContributions(daysByProvider: &merged, contributions: delta, sign: 1) + return merged + } + + private static func applyContributions( + daysByProvider: inout [String: [String: [String: PiPackedUsage]]], + contributions: [String: [String: [String: PiPackedUsage]]], + sign: Int) + { + for (providerKey, providerDays) in contributions { + var mergedProviderDays = daysByProvider[providerKey] ?? [:] + for (dayKey, dayModels) in providerDays { + var mergedDayModels = mergedProviderDays[dayKey] ?? [:] + for (modelName, packed) in dayModels { + let updated = self.addPacked( + a: mergedDayModels[modelName] ?? PiPackedUsage(), + b: packed, + sign: sign) + if updated.isZero { + mergedDayModels.removeValue(forKey: modelName) + } else { + mergedDayModels[modelName] = updated + } + } + if mergedDayModels.isEmpty { + mergedProviderDays.removeValue(forKey: dayKey) + } else { + mergedProviderDays[dayKey] = mergedDayModels + } + } + if mergedProviderDays.isEmpty { + daysByProvider.removeValue(forKey: providerKey) + } else { + daysByProvider[providerKey] = mergedProviderDays + } + } + } + + private static func addPacked(a: PiPackedUsage, b: PiPackedUsage, sign: Int) -> PiPackedUsage { + let aUsageSampleCount = a.usageSampleCount ?? (a.isZero ? 0 : nil) + let bUsageSampleCount = b.usageSampleCount ?? (b.isZero ? 0 : nil) + let usageSampleCount: Int? = if let aCount = aUsageSampleCount, let bCount = bUsageSampleCount { + max(0, aCount + sign * bCount) + } else { + nil + } + + return PiPackedUsage( + inputTokens: max(0, a.inputTokens + sign * b.inputTokens), + cacheReadTokens: max(0, a.cacheReadTokens + sign * b.cacheReadTokens), + cacheWriteTokens: max(0, a.cacheWriteTokens + sign * b.cacheWriteTokens), + outputTokens: max(0, a.outputTokens + sign * b.outputTokens), + totalTokens: max(0, a.totalTokens + sign * b.totalTokens), + costNanos: max(0, a.costNanos + Int64(sign) * b.costNanos), + costSampleCount: max(0, a.costSampleCount + sign * b.costSampleCount), + usageSampleCount: usageSampleCount) + } + + private static func parseSessionStartFromFilename(_ filename: String) -> Date? { + guard let regex = self.sessionStartFilenameRegex else { return nil } + let range = NSRange(filename.startIndex.. Date? { + self.isoFormatterBox.lock.lock() + defer { self.isoFormatterBox.lock.unlock() } + return self.isoFormatterBox.withFractional.date(from: text) + ?? self.isoFormatterBox.plain.date(from: text) + } + + private static func localMidnight(_ date: Date) -> Date { + let components = Calendar.current.dateComponents([.year, .month, .day], from: date) + return Calendar.current.date(from: components) ?? date + } + + private static func dateFromDayKey(_ key: String) -> Date? { + let parts = key.split(separator: "-") + guard parts.count == 3, + let year = Int(parts[0]), + let month = Int(parts[1]), + let day = Int(parts[2]) else { return nil } + + var components = DateComponents() + components.calendar = Calendar.current + components.timeZone = TimeZone.current + components.year = year + components.month = month + components.day = day + components.hour = 0 + return components.date + } + + private static func sortedModelBreakdowns(_ breakdowns: [CostUsageDailyReport.ModelBreakdown]) + -> [CostUsageDailyReport.ModelBreakdown] + { + breakdowns.sorted { lhs, rhs in + let lhsCost = lhs.costUSD ?? -1 + let rhsCost = rhs.costUSD ?? -1 + if lhsCost != rhsCost { + return lhsCost > rhsCost + } + + let lhsTokens = lhs.totalTokens ?? -1 + let rhsTokens = rhs.totalTokens ?? -1 + if lhsTokens != rhsTokens { + return lhsTokens > rhsTokens + } + + return lhs.modelName > rhs.modelName + } + } +} diff --git a/Sources/CodexBarCore/ProviderAccountSnapshot.swift b/Sources/CodexBarCore/ProviderAccountSnapshot.swift new file mode 100644 index 0000000..6a569c1 --- /dev/null +++ b/Sources/CodexBarCore/ProviderAccountSnapshot.swift @@ -0,0 +1,55 @@ +import Foundation + +/// Stable identity for one account row surfaced by a multi-account source adapter. +/// +/// `source` names the adapter (for example `claude-swap`) and `opaqueID` is the +/// source-issued identifier (for example a numeric slot). Identity never derives +/// from emails or credential material, per +/// `docs/claude-multi-account-and-status-items.md`. +public struct ProviderAccountIdentity: Hashable, Sendable { + public let source: String + public let opaqueID: String + + public init(source: String, opaqueID: String) { + self.source = source + self.opaqueID = opaqueID + } +} + +/// Provider-neutral projection of one account's usage, consumed by menus (and, +/// later, per-account status items) without teaching UI code about any specific +/// credential source. +public struct ProviderAccountUsageSnapshot: Identifiable, Sendable { + public let id: ProviderAccountIdentity + public let provider: UsageProvider + /// Display-only label (may contain personal data such as an email); UI is + /// responsible for privacy redaction. Never logged or persisted. + public let displayLabel: String + public let isActive: Bool + /// Whether the source can make this inactive account the provider's active account. + /// Activation remains source-owned; CodexBar never handles credential material. + public let canActivate: Bool + public let snapshot: UsageSnapshot? + public let error: String? + public let sourceLabel: String? + + public init( + id: ProviderAccountIdentity, + provider: UsageProvider, + displayLabel: String, + isActive: Bool, + canActivate: Bool = false, + snapshot: UsageSnapshot?, + error: String?, + sourceLabel: String?) + { + self.id = id + self.provider = provider + self.displayLabel = displayLabel + self.isActive = isActive + self.canActivate = canActivate + self.snapshot = snapshot + self.error = error + self.sourceLabel = sourceLabel + } +} diff --git a/Sources/CodexBarCore/ProviderCostSnapshot.swift b/Sources/CodexBarCore/ProviderCostSnapshot.swift new file mode 100644 index 0000000..5d838d5 --- /dev/null +++ b/Sources/CodexBarCore/ProviderCostSnapshot.swift @@ -0,0 +1,38 @@ +import Foundation + +/// Provider-specific spend/budget snapshot (e.g. Claude "Extra usage" monthly spend vs limit). +public struct ProviderCostSnapshot: Equatable, Codable, Sendable { + public let used: Double + public let limit: Double + public let currencyCode: String + /// Human-friendly period label (e.g. "Monthly"). Optional; some providers don't expose a period. + public let period: String? + /// Optional renewal/reset timestamp for the period. + public let resetsAt: Date? + /// Optional amount restored on the next regeneration tick for providers with rolling credit recovery. + public let nextRegenAmount: Double? + /// This account's own contribution when `used`/`limit` describe a shared/pooled budget + /// (e.g. Cursor team on-demand pool). nil when the budget is already personal. + public let personalUsed: Double? + public let updatedAt: Date + + public init( + used: Double, + limit: Double, + currencyCode: String, + period: String? = nil, + resetsAt: Date? = nil, + nextRegenAmount: Double? = nil, + personalUsed: Double? = nil, + updatedAt: Date) + { + self.used = used + self.limit = limit + self.currencyCode = currencyCode + self.period = period + self.resetsAt = resetsAt + self.nextRegenAmount = nextRegenAmount + self.personalUsed = personalUsed + self.updatedAt = updatedAt + } +} diff --git a/Sources/CodexBarCore/ProviderEndpointOverrideValidator.swift b/Sources/CodexBarCore/ProviderEndpointOverrideValidator.swift new file mode 100644 index 0000000..6b4ee4a --- /dev/null +++ b/Sources/CodexBarCore/ProviderEndpointOverrideValidator.swift @@ -0,0 +1,170 @@ +import Foundation + +enum ProviderEndpointOverrideError: LocalizedError, Equatable { + case minimax(String) + case alibabaCodingPlan(String) + + var errorDescription: String? { + switch self { + case let .minimax(key): + "MiniMax endpoint override \(key) is not allowed. " + + "Use an HTTPS endpoint without user info or encoded host tricks. " + + "If MINIMAX_REQUIRE_PROVIDER_ENDPOINT_OVERRIDES=true is set, the endpoint must also be MiniMax-owned." + case let .alibabaCodingPlan(key): + "Alibaba Coding Plan endpoint override \(key) is not allowed. " + + "Use an HTTPS endpoint without user info or encoded host tricks. " + + "If ALIBABA_CODING_PLAN_REQUIRE_PROVIDER_ENDPOINT_OVERRIDES=true is set, " + + "the endpoint must also be Alibaba-owned." + } + } +} + +struct ProviderEndpointOverrideValidator { + enum HostPolicy { + case allowAnyHTTPSHost + case providerOwnedOnly + } + + private let allowedHosts: Set + private let allowedDomainSuffixes: Set + + init(allowedHosts: [String] = [], allowedDomainSuffixes: [String] = []) { + self.allowedHosts = Set(allowedHosts.map { $0.lowercased() }) + self.allowedDomainSuffixes = Set(allowedDomainSuffixes.map { $0.lowercased() }) + } + + func validatedHost(_ raw: String?, policy: HostPolicy = .allowAnyHTTPSHost) -> String? { + guard let raw, + let url = self.url(from: raw), + let host = self.validatedDecodedHost(for: url, policy: policy) + else { return nil } + return self.hostAuthority(host: host, port: url.port) + } + + func validatedURL(_ raw: String?, policy: HostPolicy = .allowAnyHTTPSHost) -> URL? { + guard let raw, + let url = self.url(from: raw), + self.validatedDecodedHost(for: url, policy: policy) != nil + else { return nil } + return url + } + + func validatedURLAllowingLoopbackHTTP(_ raw: String?) -> URL? { + guard let raw, + Self.hasExplicitURLScheme(raw), + let url = URL(string: raw), + let scheme = url.scheme?.lowercased(), + scheme == "https" || scheme == "http", + url.user == nil, + url.password == nil, + let host = self.validatedDecodedHost(for: url, policy: .allowAnyHTTPSHost), + scheme == "https" || Self.isLoopbackHost(host) + else { return nil } + return url + } + + static func normalizedHTTPSURL(from raw: String) -> URL? { + let url = if Self.hasExplicitURLScheme(raw) { + URL(string: raw) + } else { + URL(string: "https://\(raw)") + } + guard let url else { return nil } + guard let scheme = url.scheme?.lowercased(), scheme == "https" else { return nil } + guard url.user == nil, url.password == nil else { return nil } + guard let decodedHost = url.host(percentEncoded: false)?.lowercased(), + !decodedHost.isEmpty, + !decodedHost.contains("%"), + decodedHost.rangeOfCharacter(from: .whitespacesAndNewlines) == nil, + decodedHost.rangeOfCharacter(from: .controlCharacters) == nil, + let encodedHost = url.host(percentEncoded: true)?.lowercased(), + Self.hostHasNoEncodedDelimiters(encodedHost, decodedHost: decodedHost, url: url) + else { return nil } + return url + } + + private func url(from raw: String) -> URL? { + Self.normalizedHTTPSURL(from: raw) + } + + private static func hasExplicitURLScheme(_ raw: String) -> Bool { + guard let colonIndex = raw.firstIndex(of: ":") else { return false } + if raw[colonIndex...].hasPrefix("://") { return true } + + if let authorityEnd = raw.firstIndex(where: { ["/", "?", "#"].contains($0) }), + colonIndex > authorityEnd + { + return false + } + + let afterColon = raw.index(after: colonIndex) + guard afterColon < raw.endIndex else { return true } + let portEnd = raw[afterColon...].firstIndex { Set(["/", "?", "#"]).contains($0) } ?? raw.endIndex + let suffix = raw[afterColon.. Bool { + if host == "localhost" || host == "::1" { return true } + let octets = host.split(separator: ".", omittingEmptySubsequences: false) + guard octets.count == 4, + let first = UInt8(octets[0]), + octets.dropFirst().allSatisfy({ UInt8($0) != nil }) + else { return false } + return first == 127 + } + + private func hostAuthority(host: String, port: Int?) -> String { + let authorityHost = host.contains(":") ? "[\(host)]" : host + guard let port else { return authorityHost } + return "\(authorityHost):\(port)" + } + + private func validatedDecodedHost(for url: URL, policy: HostPolicy) -> String? { + guard let decodedHost = url.host(percentEncoded: false)?.lowercased(), + !decodedHost.isEmpty, + !decodedHost.contains("%"), + decodedHost.rangeOfCharacter(from: .whitespacesAndNewlines) == nil, + decodedHost.rangeOfCharacter(from: .controlCharacters) == nil, + let encodedHost = url.host(percentEncoded: true)?.lowercased(), + Self.hostHasNoEncodedDelimiters(encodedHost, decodedHost: decodedHost, url: url) + else { return nil } + + switch policy { + case .allowAnyHTTPSHost: + return decodedHost + case .providerOwnedOnly: + let isAllowedHost = self.allowedHosts.contains(decodedHost) + let isAllowedSuffix = self.allowedDomainSuffixes.contains { suffix in + decodedHost == suffix || decodedHost.hasSuffix(".\(suffix)") + } + guard isAllowedHost || isAllowedSuffix else { return nil } + return decodedHost + } + } + + private static func hostHasNoEncodedDelimiters(_ encodedHost: String, decodedHost: String, url: URL) -> Bool { + if decodedHost.contains(":") { + guard encodedHost == decodedHost, + let componentHost = URLComponents(url: url, resolvingAgainstBaseURL: false)?.host, + componentHost.hasPrefix("["), + componentHost.hasSuffix("]") + else { return false } + + let address = componentHost.dropFirst().dropLast() + return !address.isEmpty && address.allSatisfy { $0.isHexDigit || $0 == ":" || $0 == "." } + } + + let decodedDelimiters = CharacterSet(charactersIn: "/\\?#@:") + guard decodedHost.rangeOfCharacter(from: decodedDelimiters) == nil else { return false } + + let encodedDelimiters = ["%2f", "%5c", "%3f", "%23", "%40", "%3a"] + return !encodedDelimiters.contains { encodedHost.contains($0) } + } +} diff --git a/Sources/CodexBarCore/ProviderHTTPClient.swift b/Sources/CodexBarCore/ProviderHTTPClient.swift new file mode 100644 index 0000000..68b8b38 --- /dev/null +++ b/Sources/CodexBarCore/ProviderHTTPClient.swift @@ -0,0 +1,253 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public protocol ProviderHTTPTransport: Sendable { + func data(for request: URLRequest) async throws -> (Data, URLResponse) +} + +#if !os(Linux) +extension URLSession: ProviderHTTPTransport {} +#endif + +extension URLSession { + public func response(for request: URLRequest) async throws -> ProviderHTTPResponse { + let (data, response) = try await self.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw URLError(.badServerResponse) + } + return ProviderHTTPResponse(data: data, response: httpResponse) + } +} + +public struct ProviderHTTPResponse: Sendable { + public let data: Data + public let response: HTTPURLResponse + + public init(data: Data, response: HTTPURLResponse) { + self.data = data + self.response = response + } + + public var statusCode: Int { + self.response.statusCode + } +} + +public struct ProviderHTTPRetryPolicy: Sendable { + public let maxRetries: Int + public let retryableStatusCodes: Set + public let retryableURLErrorCodes: Set + public let retryableMethods: Set + public let baseDelaySeconds: TimeInterval + public let maxDelaySeconds: TimeInterval + + public init( + maxRetries: Int, + retryableStatusCodes: Set = [408, 429, 500, 502, 503, 504], + retryableURLErrorCodes: Set = [ + .timedOut, + .networkConnectionLost, + .cannotConnectToHost, + .cannotFindHost, + .dnsLookupFailed, + ], + retryableMethods: Set = ["GET", "HEAD", "OPTIONS"], + baseDelaySeconds: TimeInterval = 1, + maxDelaySeconds: TimeInterval = 10) + { + self.maxRetries = max(0, maxRetries) + self.retryableStatusCodes = retryableStatusCodes + self.retryableURLErrorCodes = retryableURLErrorCodes + self.retryableMethods = retryableMethods + self.baseDelaySeconds = max(0, baseDelaySeconds) + self.maxDelaySeconds = max(0, maxDelaySeconds) + } + + public static let disabled = ProviderHTTPRetryPolicy( + maxRetries: 0, + retryableStatusCodes: [], + retryableURLErrorCodes: [], + baseDelaySeconds: 0, + maxDelaySeconds: 0) + + public static let transientIdempotent = ProviderHTTPRetryPolicy(maxRetries: 1) + + func shouldRetry(request: URLRequest, attempt: Int, statusCode: Int) -> Bool { + self.canRetry(request: request, attempt: attempt) + && self.retryableStatusCodes.contains(statusCode) + } + + func shouldRetry(request: URLRequest, attempt: Int, error: Error) -> Bool { + guard self.canRetry(request: request, attempt: attempt) else { return false } + guard let urlError = error as? URLError else { return false } + return self.retryableURLErrorCodes.contains(urlError.code) + } + + func delaySeconds(attempt: Int, response: HTTPURLResponse?) -> TimeInterval { + if let retryAfter = response?.value(forHTTPHeaderField: "Retry-After"), + let seconds = TimeInterval(retryAfter.trimmingCharacters(in: .whitespacesAndNewlines)), + seconds >= 0 + { + return min(seconds, self.maxDelaySeconds) + } + + guard self.baseDelaySeconds > 0 else { return 0 } + let multiplier = pow(2, Double(max(0, attempt))) + return min(self.baseDelaySeconds * multiplier, self.maxDelaySeconds) + } + + private func canRetry(request: URLRequest, attempt: Int) -> Bool { + guard attempt < self.maxRetries else { return false } + let method = request.httpMethod?.uppercased() ?? "GET" + return self.retryableMethods.contains(method) + } +} + +public struct ProviderHTTPTransportHandler: ProviderHTTPTransport { + private let handler: @Sendable (URLRequest) async throws -> (Data, URLResponse) + + public init(_ handler: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) { + self.handler = handler + } + + public func data(for request: URLRequest) async throws -> (Data, URLResponse) { + try await self.handler(request) + } +} + +extension ProviderHTTPTransport { + public func response(for request: URLRequest) async throws -> ProviderHTTPResponse { + try await self.response(for: request, retryPolicy: .disabled) + } + + public func response( + for request: URLRequest, + retryPolicy: ProviderHTTPRetryPolicy) async throws -> ProviderHTTPResponse + { + var attempt = 0 + + while true { + do { + let (data, response) = try await self.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw URLError(.badServerResponse) + } + let providerResponse = ProviderHTTPResponse(data: data, response: httpResponse) + guard retryPolicy.shouldRetry( + request: request, + attempt: attempt, + statusCode: providerResponse.statusCode) + else { + return providerResponse + } + try await Self.sleepBeforeRetry(policy: retryPolicy, attempt: attempt, response: httpResponse) + attempt += 1 + } catch { + guard retryPolicy.shouldRetry(request: request, attempt: attempt, error: error) else { + throw error + } + try await Self.sleepBeforeRetry(policy: retryPolicy, attempt: attempt, response: nil) + attempt += 1 + } + } + } + + private static func sleepBeforeRetry( + policy: ProviderHTTPRetryPolicy, + attempt: Int, + response: HTTPURLResponse?) async throws + { + let delay = policy.delaySeconds(attempt: attempt, response: response) + guard delay > 0 else { return } + try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) + } +} + +public final class ProviderHTTPClient: ProviderHTTPTransport, @unchecked Sendable { + public static let shared = ProviderHTTPClient(session: ProviderHTTPClient.sharedSession()) + + private let session: URLSession + + public init(session: URLSession? = nil) { + self.session = session ?? Self.redirectGuardedSession() + } + + static func defaultConfiguration() -> URLSessionConfiguration { + let configuration = URLSessionConfiguration.default + configuration.timeoutIntervalForRequest = 30 + configuration.timeoutIntervalForResource = 90 + #if !os(Linux) + configuration.waitsForConnectivity = false + #endif + return configuration + } + + private static func sharedSession() -> URLSession { + if self.isRunningTests { + // XCTest URLProtocol.registerClass stubs only intercept URLSession.shared on macOS. + return .shared + } + return self.redirectGuardedSession() + } + + static func redirectGuardedSession( + configuration: URLSessionConfiguration = ProviderHTTPClient.defaultConfiguration()) -> URLSession + { + URLSession( + configuration: configuration, + delegate: ProviderHTTPRedirectGuardDelegate(), + delegateQueue: nil) + } + + private static var isRunningTests: Bool { + let environment = ProcessInfo.processInfo.environment + if environment["XCTestConfigurationFilePath"] != nil || environment["XCTestBundlePath"] != nil { + return true + } + if ProcessInfo.processInfo.processName.lowercased().contains("xctest") { + return true + } + return CommandLine.arguments.contains { $0.lowercased().contains(".xctest") } + } + + public func data(for request: URLRequest) async throws -> (Data, URLResponse) { + try await self.session.data(for: request) + } +} + +final class ProviderHTTPRedirectGuardDelegate: NSObject, URLSessionTaskDelegate, @unchecked Sendable { + func urlSession( + _: URLSession, + task: URLSessionTask, + willPerformHTTPRedirection _: HTTPURLResponse, + newRequest request: URLRequest, + completionHandler: @escaping @Sendable (URLRequest?) -> Void) + { + completionHandler(Self.guardedRedirectRequest(originalURL: task.originalRequest?.url, redirectRequest: request)) + } + + static func guardedRedirectRequest(originalURL: URL?, redirectRequest request: URLRequest) -> URLRequest? { + guard let originalURL, let redirectedURL = request.url else { return nil } + guard originalURL.scheme?.caseInsensitiveCompare("https") == .orderedSame else { return nil } + guard redirectedURL.scheme?.caseInsensitiveCompare("https") == .orderedSame else { return nil } + guard self.isSameOrigin(originalURL, redirectedURL) else { return nil } + return request + } + + private static func isSameOrigin(_ lhs: URL, _ rhs: URL) -> Bool { + lhs.scheme?.lowercased() == rhs.scheme?.lowercased() + && lhs.host?.lowercased() == rhs.host?.lowercased() + && self.normalizedPort(lhs) == self.normalizedPort(rhs) + } + + private static func normalizedPort(_ url: URL) -> Int? { + if let port = url.port { return port } + switch url.scheme?.lowercased() { + case "http": return 80 + case "https": return 443 + default: return nil + } + } +} diff --git a/Sources/CodexBarCore/ProviderStorageFootprint.swift b/Sources/CodexBarCore/ProviderStorageFootprint.swift new file mode 100644 index 0000000..75feb7a --- /dev/null +++ b/Sources/CodexBarCore/ProviderStorageFootprint.swift @@ -0,0 +1,543 @@ +import Foundation + +public struct ProviderStorageFootprint: Sendable, Equatable { + public struct Component: Sendable, Equatable, Identifiable { + public let id: String + public let path: String + public let totalBytes: Int64 + + public init(path: String, totalBytes: Int64) { + self.id = path + self.path = path + self.totalBytes = totalBytes + } + + public var name: String { + let url = URL(fileURLWithPath: self.path) + let last = url.lastPathComponent + if last.isEmpty { return self.path } + return last + } + } + + public let provider: UsageProvider + public let totalBytes: Int64 + public let paths: [String] + public let missingPaths: [String] + public let unreadablePaths: [String] + public let components: [Component] + public let updatedAt: Date + + public init( + provider: UsageProvider, + totalBytes: Int64, + paths: [String], + missingPaths: [String], + unreadablePaths: [String], + components: [Component] = [], + updatedAt: Date) + { + self.provider = provider + self.totalBytes = totalBytes + self.paths = paths + self.missingPaths = missingPaths + self.unreadablePaths = unreadablePaths + self.components = components + self.updatedAt = updatedAt + } + + public var hasLocalData: Bool { + self.totalBytes > 0 + } + + /// Value equality that ignores `updatedAt`. Two scans of identical on-disk data differ only by + /// their scan timestamp, so callers use this to avoid re-publishing observable state (and the + /// menu-invalidation churn that follows) when nothing the user sees has actually changed. + public func hasSameContents(as other: ProviderStorageFootprint) -> Bool { + self.provider == other.provider && + self.totalBytes == other.totalBytes && + self.paths == other.paths && + self.missingPaths == other.missingPaths && + self.unreadablePaths == other.unreadablePaths && + self.components == other.components + } + + public var cleanupRecommendations: [ProviderStorageRecommendation] { + ProviderStorageRecommendation.recommendations(for: self) + } + + public func replacingProvider(_ provider: UsageProvider) -> ProviderStorageFootprint { + ProviderStorageFootprint( + provider: provider, + totalBytes: self.totalBytes, + paths: self.paths, + missingPaths: self.missingPaths, + unreadablePaths: self.unreadablePaths, + components: self.components, + updatedAt: self.updatedAt) + } +} + +public struct ProviderStorageRecommendation: Sendable, Equatable, Identifiable { + public enum RiskLevel: String, Sendable { + case informational + case manualCleanup + } + + public let id: String + public let provider: UsageProvider + public let path: String + public let bytes: Int64 + public let title: String + public let riskLevel: RiskLevel + public let consequence: String + public let sortPriority: Int + + public init( + provider: UsageProvider, + path: String, + bytes: Int64, + title: String, + riskLevel: RiskLevel, + consequence: String, + sortPriority: Int) + { + self.id = path + self.provider = provider + self.path = path + self.bytes = bytes + self.title = title + self.riskLevel = riskLevel + self.consequence = consequence + self.sortPriority = sortPriority + } + + public static func recommendations(for footprint: ProviderStorageFootprint) -> [ProviderStorageRecommendation] { + let candidates: [ProviderStorageRecommendation] = footprint.components.compactMap { component in + switch footprint.provider { + case .claude: + self.claudeRecommendation(for: component) + case .codex: + self.codexRecommendation(for: component, roots: footprint.paths) + default: + nil + } + } + + return candidates.sorted { lhs, rhs in + if lhs.sortPriority == rhs.sortPriority { + if lhs.bytes == rhs.bytes { + return lhs.path.localizedCaseInsensitiveCompare(rhs.path) == .orderedAscending + } + return lhs.bytes > rhs.bytes + } + return lhs.sortPriority < rhs.sortPriority + } + } + + private static func claudeRecommendation( + for component: ProviderStorageFootprint.Component) + -> ProviderStorageRecommendation? + { + switch component.name { + case "projects": + self.make( + provider: .claude, + component: component, + title: "Manual cleanup: past sessions", + consequence: "Clearing removes past resume, continue, and rewind history.", + priority: 10) + case "file-history": + self.make( + provider: .claude, + component: component, + title: "Manual cleanup: file checkpoints", + consequence: "Clearing removes checkpoint restore data for previous edits.", + priority: 20) + case "plans": + self.make( + provider: .claude, + component: component, + title: "Manual cleanup: saved plans", + consequence: "Clearing removes old plan-mode files.", + priority: 30) + case "debug": + self.make( + provider: .claude, + component: component, + title: "Manual cleanup: debug logs", + consequence: "Clearing removes past debug logs.", + priority: 40) + case "paste-cache", "image-cache": + self.make( + provider: .claude, + component: component, + title: "Manual cleanup: attachment cache", + consequence: "Clearing removes cached large pastes or attached images.", + priority: 50) + case "session-env": + self.make( + provider: .claude, + component: component, + title: "Manual cleanup: session metadata", + consequence: "Clearing removes per-session environment metadata.", + priority: 60) + case "shell-snapshots": + self.make( + provider: .claude, + component: component, + title: "Manual cleanup: shell snapshots", + consequence: "Clearing removes leftover runtime shell snapshot files.", + priority: 70) + case "todos": + self.make( + provider: .claude, + component: component, + title: "Manual cleanup: legacy todos", + consequence: "Clearing removes legacy per-session task lists.", + priority: 80) + default: + nil + } + } + + private static func codexRecommendation( + for component: ProviderStorageFootprint.Component, + roots: [String]) + -> ProviderStorageRecommendation? + { + guard self.path(component.path, isContainedIn: roots) else { return nil } + + return switch component.name { + case "sessions": + self.make( + provider: .codex, + component: component, + title: "Manual cleanup: sessions", + consequence: "Clearing removes past Codex session history.", + priority: 10) + case "archived_sessions": + self.make( + provider: .codex, + component: component, + title: "Manual cleanup: archived sessions", + consequence: "Clearing removes archived Codex session history.", + priority: 20) + case "cache", "caches", "Cache", "Caches": + self.make( + provider: .codex, + component: component, + title: "Manual cleanup: cache", + consequence: "Clearing removes provider-owned cached data.", + priority: 30) + case "log", "logs", "debug": + self.make( + provider: .codex, + component: component, + title: "Manual cleanup: logs", + consequence: "Clearing removes local diagnostic logs.", + priority: 40) + case let name where name.hasPrefix("logs_") && name.hasSuffix(".sqlite"): + self.make( + provider: .codex, + component: component, + title: "Manual cleanup: logs", + consequence: "Clearing removes local diagnostic logs.", + priority: 40) + case "file-history": + self.make( + provider: .codex, + component: component, + title: "Manual cleanup: file history", + consequence: "Clearing removes local edit checkpoint history.", + priority: 50) + case "paste-cache", "image-cache", "session-env", "shell-snapshots", "shell_snapshots", "tmp", "temp", ".tmp": + self.make( + provider: .codex, + component: component, + title: "Manual cleanup: temporary data", + consequence: "Clearing removes local temporary provider data.", + priority: 60) + default: + nil + } + } + + private static func make( + provider: UsageProvider, + component: ProviderStorageFootprint.Component, + title: String, + consequence: String, + priority: Int) + -> ProviderStorageRecommendation + { + ProviderStorageRecommendation( + provider: provider, + path: component.path, + bytes: component.totalBytes, + title: title, + riskLevel: .manualCleanup, + consequence: consequence, + sortPriority: priority) + } + + private static func path(_ path: String, isContainedIn roots: [String]) -> Bool { + let standardizedPath = URL(fileURLWithPath: path).standardizedFileURL.path + return roots.contains { root in + let standardizedRoot = URL(fileURLWithPath: root, isDirectory: true).standardizedFileURL.path + return standardizedPath == standardizedRoot || standardizedPath.hasPrefix(standardizedRoot + "/") + } + } +} + +public enum ProviderStoragePathCatalog { + public static func candidatePaths( + for provider: UsageProvider, + environment: [String: String], + managedCodexAccounts: [ManagedCodexAccount] = [], + fileManager: FileManager = .default) + -> [String] + { + let home = fileManager.homeDirectoryForCurrentUser + + func homePath(_ relativePath: String) -> String { + home.appendingPathComponent(relativePath, isDirectory: true).path + } + + let candidates: [String] = switch provider { + case .codex: + [CodexHomeScope.ambientHomeURL(env: environment, fileManager: fileManager).path] + + managedCodexAccounts.map(\.managedHomePath) + case .claude: + [ + homePath(".claude"), + homePath(".config/claude"), + home + .appendingPathComponent("Library/Application Support/CodexBar/ClaudeProbe", isDirectory: true) + .path, + ] + case .gemini: + [ + homePath(".gemini"), + homePath(".config/gemini"), + ] + case .opencode, .opencodego: + [ + homePath(".config/opencode"), + ] + case .copilot: + [ + homePath(".config/github-copilot"), + ] + case .cursor: + [ + homePath("Library/Application Support/Cursor"), + homePath("Library/Application Support/Caches/cursor-updater"), + homePath(".cursor"), + homePath("Library/Caches/Cursor"), + homePath("Library/Caches/com.todesktop.230313mzl4w4u92"), + homePath("Library/Caches/com.todesktop.230313mzl4w4u92.ShipIt"), + homePath("Library/Caches/cursor-compile-cache"), + homePath("Library/HTTPStorages/com.todesktop.230313mzl4w4u92"), + ] + default: + [] + } + + return Self.uniqueStandardizedPaths(candidates) + } + + private static func uniqueStandardizedPaths(_ paths: [String]) -> [String] { + var seen: Set = [] + var result: [String] = [] + for path in paths { + let trimmed = path.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { continue } + let standardized = URL(fileURLWithPath: trimmed, isDirectory: true).standardizedFileURL.path + guard seen.insert(standardized).inserted else { continue } + result.append(standardized) + } + return result + } +} + +public struct ProviderStorageScanner: @unchecked Sendable { + private struct DirectoryScanResult { + var bytes: Int64 = 0 + var unreadablePaths: [String] = [] + var componentBytes: [String: Int64] = [:] + } + + private let fileManager: FileManager + + public init(fileManager: FileManager = .default) { + self.fileManager = fileManager + } + + public func scan( + provider: UsageProvider, + candidatePaths: [String], + now: Date = Date()) + -> ProviderStorageFootprint + { + var totalBytes: Int64 = 0 + var existingPaths: [String] = [] + var missingPaths: [String] = [] + var unreadablePaths: [String] = [] + var components: [ProviderStorageFootprint.Component] = [] + + for path in candidatePaths { + if Task.isCancelled { break } + var isDirectory: ObjCBool = false + guard self.fileManager.fileExists(atPath: path, isDirectory: &isDirectory) else { + missingPaths.append(path) + continue + } + + existingPaths.append(path) + let url = URL(fileURLWithPath: path, isDirectory: isDirectory.boolValue) + if self.isSymbolicLink(at: url) { + continue + } + if isDirectory.boolValue { + let result = self.scanDirectory(at: url) + if Task.isCancelled { break } + totalBytes += result.bytes + unreadablePaths.append(contentsOf: result.unreadablePaths) + components.append(contentsOf: result.componentBytes.map { + ProviderStorageFootprint.Component(path: $0.key, totalBytes: $0.value) + }) + } else { + let result = self.sizeOfFile(at: url) + totalBytes += result.bytes + unreadablePaths.append(contentsOf: result.unreadablePaths) + if result.bytes > 0 { + components.append(.init(path: url.path, totalBytes: result.bytes)) + } + } + } + + return ProviderStorageFootprint( + provider: provider, + totalBytes: totalBytes, + paths: existingPaths, + missingPaths: missingPaths, + unreadablePaths: unreadablePaths, + components: components.sorted { lhs, rhs in + if lhs.totalBytes == rhs.totalBytes { + return lhs.path.localizedCaseInsensitiveCompare(rhs.path) == .orderedAscending + } + return lhs.totalBytes > rhs.totalBytes + }, + updatedAt: now) + } + + private func isSymbolicLink(at url: URL) -> Bool { + (try? url.resourceValues(forKeys: [.isSymbolicLinkKey]).isSymbolicLink) == true + } + + private func sizeOfFile(at url: URL) -> (bytes: Int64, unreadablePaths: [String]) { + if Task.isCancelled { return (0, []) } + let keys: Set = [ + .isRegularFileKey, + .isSymbolicLinkKey, + .fileSizeKey, + ] + + guard let values = try? url.resourceValues(forKeys: keys) else { + return (0, [url.path]) + } + + if values.isSymbolicLink == true { + return (0, []) + } + + if values.isRegularFile == true { + return (Int64(values.fileSize ?? 0), []) + } + + return (0, []) + } + + private func scanDirectory(at url: URL) -> DirectoryScanResult { + if Task.isCancelled { return DirectoryScanResult() } + let keys: Set = [ + .isDirectoryKey, + .isRegularFileKey, + .isSymbolicLinkKey, + .fileSizeKey, + ] + + let unreadableCollector = ProviderStorageUnreadablePathCollector() + guard let enumerator = self.fileManager.enumerator( + at: url, + includingPropertiesForKeys: Array(keys), + options: [.skipsPackageDescendants], + errorHandler: { url, _ in + unreadableCollector.append(url.path) + return true + }) + else { + return DirectoryScanResult(unreadablePaths: [url.path]) + } + + var result = DirectoryScanResult() + let rootPath = url.standardizedFileURL.path + for case let itemURL as URL in enumerator { + if Task.isCancelled { + enumerator.skipDescendants() + break + } + guard let itemValues = try? itemURL.resourceValues(forKeys: keys) else { + unreadableCollector.append(itemURL.path) + continue + } + if itemValues.isSymbolicLink == true { + if itemValues.isDirectory == true { + enumerator.skipDescendants() + } + continue + } + if itemValues.isRegularFile == true { + let bytes = Int64(itemValues.fileSize ?? 0) + result.bytes += bytes + if bytes > 0, let componentPath = self.topLevelComponentPath(for: itemURL, rootPath: rootPath) { + result.componentBytes[componentPath, default: 0] += bytes + } + } + } + result.unreadablePaths = unreadableCollector.paths + return result + } + + private func topLevelComponentPath(for url: URL, rootPath: String) -> String? { + let itemPath = url.standardizedFileURL.path + let pathPrefix = rootPath.hasSuffix("/") ? rootPath : "\(rootPath)/" + guard itemPath.hasPrefix(pathPrefix) else { return nil } + let suffix = itemPath.dropFirst(pathPrefix.count) + let relative = suffix.drop { $0 == "/" } + guard let first = relative.split(separator: "/", maxSplits: 1, omittingEmptySubsequences: true).first else { + return nil + } + return URL(fileURLWithPath: rootPath, isDirectory: true) + .appendingPathComponent(String(first)) + .path + } +} + +private final class ProviderStorageUnreadablePathCollector: @unchecked Sendable { + private let lock = NSLock() + private var storage: [String] = [] + + var paths: [String] { + self.lock.lock() + defer { self.lock.unlock() } + return self.storage + } + + func append(_ path: String) { + self.lock.lock() + defer { self.lock.unlock() } + self.storage.append(path) + } +} diff --git a/Sources/CodexBarCore/Providers/APITokenFetchStrategy.swift b/Sources/CodexBarCore/Providers/APITokenFetchStrategy.swift new file mode 100644 index 0000000..8e33789 --- /dev/null +++ b/Sources/CodexBarCore/Providers/APITokenFetchStrategy.swift @@ -0,0 +1,72 @@ +import Foundation + +public struct APITokenFetchStrategy: ProviderFetchStrategy { + public typealias TokenResolver = @Sendable ([String: String]) -> String? + public typealias MissingCredentialsError = @Sendable () -> (any Error & Sendable) + public typealias UsageLoader = @Sendable (String, ProviderFetchContext) async throws -> UsageSnapshot + + public let id: String + public let kind: ProviderFetchKind = .apiToken + + private let sourceLabel: String + private let reportsMissingCredentials: Bool + private let resolveToken: TokenResolver + private let missingCredentialsError: MissingCredentialsError + private let loadUsage: UsageLoader + + public init( + id: String, + sourceLabel: String = "api", + reportsMissingCredentials: Bool = false, + resolveToken: @escaping TokenResolver, + missingCredentialsError: @escaping MissingCredentialsError, + loadUsage: @escaping UsageLoader) + { + self.id = id + self.sourceLabel = sourceLabel + self.reportsMissingCredentials = reportsMissingCredentials + self.resolveToken = resolveToken + self.missingCredentialsError = missingCredentialsError + self.loadUsage = loadUsage + } + + public func isAvailable(_ context: ProviderFetchContext) async -> Bool { + self.reportsMissingCredentials || self.resolveToken(context.env) != nil + } + + public func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let token = self.resolveToken(context.env) else { + throw self.missingCredentialsError() + } + return try await self.makeResult( + usage: self.loadUsage(token, context), + sourceLabel: self.sourceLabel) + } + + public func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} + +extension ProviderFetchPlan { + public static func apiToken( + strategyID: String, + sourceLabel: String = "api", + reportsMissingCredentials: Bool = false, + resolveToken: @escaping APITokenFetchStrategy.TokenResolver, + missingCredentialsError: @escaping APITokenFetchStrategy.MissingCredentialsError, + loadUsage: @escaping APITokenFetchStrategy.UsageLoader) -> ProviderFetchPlan + { + ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + [APITokenFetchStrategy( + id: strategyID, + sourceLabel: sourceLabel, + reportsMissingCredentials: reportsMissingCredentials, + resolveToken: resolveToken, + missingCredentialsError: missingCredentialsError, + loadUsage: loadUsage)] + })) + } +} diff --git a/Sources/CodexBarCore/Providers/Abacus/AbacusCookieImporter.swift b/Sources/CodexBarCore/Providers/Abacus/AbacusCookieImporter.swift new file mode 100644 index 0000000..c32c8bb --- /dev/null +++ b/Sources/CodexBarCore/Providers/Abacus/AbacusCookieImporter.swift @@ -0,0 +1,134 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit + +// MARK: - Abacus Cookie Importer + +public enum AbacusCookieImporter { + private static let log = CodexBarLog.logger(LogCategories.abacusCookie) + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["abacus.ai", "apps.abacus.ai"] + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.abacus]?.browserCookieOrder ?? Browser.defaultImportOrder + + /// Exact cookie names known to carry Abacus session state. + /// CSRF tokens are deliberately excluded — they are present in anonymous + /// jars and do not indicate an authenticated session. + private static let knownSessionCookieNames: Set = [ + "sessionid", "session_id", "session_token", + "auth_token", "access_token", + ] + + /// Substrings that indicate a session or auth cookie (applied only when + /// no exact-name match is found). Deliberately excludes overly broad + /// patterns like "id" and "token" that match analytics/CSRF cookies. + private static let sessionCookieSubstrings = ["session", "auth", "sid", "jwt"] + + /// Cookie name prefixes that indicate a non-session cookie even when a + /// substring match would otherwise accept it (e.g. "csrftoken"). + private static let excludedCookiePrefixes = ["csrf", "_ga", "_gid", "tracking", "analytics"] + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + } + + /// Returns all candidate sessions across browsers/profiles, ordered by + /// import priority. Callers should try each in turn so that a stale + /// session in the first source doesn't block a valid one further down. + /// + /// Defaults to Chrome-only per AGENTS.md guideline. Pass an empty + /// `preferredBrowsers` list to fall back to the full descriptor-defined + /// import order (Safari, Firefox, etc.) when Chrome has no cookies. + public static func importSessions( + browserDetection: BrowserDetection = BrowserDetection(), + preferredBrowsers: [Browser] = [.chrome], + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + var candidates: [SessionInfo] = [] + let installedBrowsers = preferredBrowsers.isEmpty + ? self.cookieImportOrder.cookieImportCandidates(using: browserDetection) + : preferredBrowsers.cookieImportCandidates(using: browserDetection) + + for browserSource in installedBrowsers { + do { + let query = BrowserCookieQuery(domains: self.cookieDomains) + let sources = try Self.cookieClient.codexBarRecords( + matching: query, + in: browserSource, + logger: { msg in self.emit(msg, logger: logger) }) + for source in sources where !source.records.isEmpty { + let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin) + guard !httpCookies.isEmpty else { continue } + + guard Self.containsSessionCookie(httpCookies) else { + self.emit( + "Skipping \(source.label): no session cookie found", + logger: logger) + continue + } + + self.emit( + "Found \(httpCookies.count) session cookies in \(source.label)", + logger: logger) + candidates.append(SessionInfo(cookies: httpCookies, sourceLabel: source.label)) + } + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + self.emit( + "\(browserSource.displayName) cookie import failed: \(error.localizedDescription)", + logger: logger) + } + } + + guard !candidates.isEmpty else { + throw AbacusUsageError.noSessionCookie + } + return candidates + } + + /// Cheap check for whether any browser has an Abacus session cookie, + /// used by the fetch strategy's `isAvailable()`. + public static func hasSession( + browserDetection: BrowserDetection = BrowserDetection(), + preferredBrowsers: [Browser] = [.chrome], + logger: ((String) -> Void)? = nil) -> Bool + { + do { + return try !self.importSessions( + browserDetection: browserDetection, + preferredBrowsers: preferredBrowsers, + logger: logger).isEmpty + } catch { + return false + } + } + + /// Returns `true` if the cookie set contains at least one cookie whose name + /// indicates session or authentication state. Checks exact known names + /// first, then falls back to conservative substring matching. + private static func containsSessionCookie(_ cookies: [HTTPCookie]) -> Bool { + cookies.contains { cookie in + let lower = cookie.name.lowercased() + if self.knownSessionCookieNames.contains(lower) { return true } + if self.excludedCookiePrefixes.contains(where: { lower.hasPrefix($0) }) { return false } + return self.sessionCookieSubstrings.contains { lower.contains($0) } + } + } + + private static func emit(_ message: String, logger: ((String) -> Void)?) { + logger?("[abacus-cookie] \(message)") + self.log.debug(message) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Abacus/AbacusProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Abacus/AbacusProviderDescriptor.swift new file mode 100644 index 0000000..2afe1dd --- /dev/null +++ b/Sources/CodexBarCore/Providers/Abacus/AbacusProviderDescriptor.swift @@ -0,0 +1,91 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit +#endif + +public enum AbacusProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .abacus, + metadata: ProviderMetadata( + id: .abacus, + displayName: "Abacus AI", + sessionLabel: "Credits", + weeklyLabel: "Weekly", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Abacus AI usage", + cliName: "abacusai", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder, + dashboardURL: "https://apps.abacus.ai/chatllm/admin/compute-points-usage", + statusPageURL: nil, + statusLinkURL: nil), + branding: ProviderBranding( + iconStyle: .abacus, + iconResourceName: "ProviderIcon-abacus", + color: ProviderColor(red: 56 / 255, green: 189 / 255, blue: 248 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Abacus AI cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + [AbacusWebFetchStrategy()] + })), + cli: ProviderCLIConfig( + name: "abacusai", + aliases: ["abacus-ai"], + versionDetector: nil)) + } +} + +// MARK: - Fetch Strategy + +struct AbacusWebFetchStrategy: ProviderFetchStrategy { + let id: String = "abacus.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + context.settings?.abacus?.cookieSource != .off + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let manual: String? + if context.settings?.abacus?.cookieSource == .manual { + guard let header = Self.manualCookieHeader(from: context) else { + throw AbacusUsageError.noSessionCookie + } + manual = header + } else { + manual = nil + } + let logger: ((String) -> Void)? = context.verbose + ? { msg in CodexBarLog.logger(LogCategories.abacusUsage).verbose(msg) } + : nil + let snap = try await AbacusUsageFetcher.fetchUsage( + cookieHeaderOverride: manual, + browserDetection: context.browserDetection, + timeout: context.webTimeout, + logger: logger) + return self.makeResult( + usage: snap.toUsageSnapshot(), + sourceLabel: "web") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func manualCookieHeader(from context: ProviderFetchContext) -> String? { + guard context.settings?.abacus?.cookieSource == .manual else { return nil } + return CookieHeaderNormalizer.normalize(context.settings?.abacus?.manualCookieHeader) + } +} diff --git a/Sources/CodexBarCore/Providers/Abacus/AbacusUsageError.swift b/Sources/CodexBarCore/Providers/Abacus/AbacusUsageError.swift new file mode 100644 index 0000000..0f5c6cc --- /dev/null +++ b/Sources/CodexBarCore/Providers/Abacus/AbacusUsageError.swift @@ -0,0 +1,67 @@ +import Foundation + +// MARK: - Abacus Usage Error + +public enum AbacusUsageError: LocalizedError, Sendable, Equatable { + case noSessionCookie + case sessionExpired + case networkError(String) + case parseFailed(String) + case unauthorized + + /// Whether this error indicates an authentication/session problem that + /// should trigger cache eviction. + public var isRecoverable: Bool { + switch self { + case .unauthorized, .sessionExpired: true + default: false + } + } + + public var isAuthRelated: Bool { + self.isRecoverable + } + + /// Whether browser-import scanning should continue to later sessions after + /// this failure. Imported sessions can differ by profile/browser, so we keep + /// scanning on per-session fetch failures and surface the first one only if + /// every candidate is exhausted. + var shouldTryNextImportedSession: Bool { + switch self { + case .unauthorized, .sessionExpired, .networkError, .parseFailed: true + case .noSessionCookie: false + } + } + + /// Whether a cached cookie header should be evicted before falling back to + /// a fresh browser import. Parse/auth failures usually indicate that the + /// cached session is stale or no longer accepted. + var shouldClearCachedCookie: Bool { + switch self { + case .unauthorized, .sessionExpired, .parseFailed: true + case .networkError, .noSessionCookie: false + } + } + + public var errorDescription: String? { + switch self { + case .noSessionCookie: + "No Abacus AI session found. Please log in to apps.abacus.ai in your browser " + + "or paste a Cookie header in manual mode." + case .sessionExpired: + "Abacus AI session expired. Please log in again." + case let .networkError(msg): + "Abacus AI API error: \(msg)" + case let .parseFailed(msg): + "Could not parse Abacus AI usage: \(msg)" + case .unauthorized: + "Unauthorized. Please log in to Abacus AI." + } + } +} + +#if !os(macOS) +extension AbacusUsageError { + public static let notSupported = AbacusUsageError.networkError("Abacus AI is only supported on macOS.") +} +#endif diff --git a/Sources/CodexBarCore/Providers/Abacus/AbacusUsageFetcher.swift b/Sources/CodexBarCore/Providers/Abacus/AbacusUsageFetcher.swift new file mode 100644 index 0000000..1c449d4 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Abacus/AbacusUsageFetcher.swift @@ -0,0 +1,336 @@ +import Foundation + +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +#if os(macOS) +import SweetCookieKit + +// MARK: - Abacus Usage Fetcher + +public enum AbacusUsageFetcher { + private struct BrowserFetchRequest { + let browserDetection: BrowserDetection + let preferredBrowsers: [Browser] + let label: String + let timeout: TimeInterval + let logger: ((String) -> Void)? + } + + /// Parsed JSON dictionaries are treated as immutable snapshots here and are + /// only moved between sibling fetch tasks before being consumed locally. + private struct JSONDictionaryBox: @unchecked Sendable { + let value: [String: Any] + } + + private static let log = CodexBarLog.logger(LogCategories.abacusUsage) + private static let computePointsURL = + URL(string: "https://apps.abacus.ai/api/_getOrganizationComputePoints")! + private static let billingInfoURL = + URL(string: "https://apps.abacus.ai/api/_getBillingInfo")! + + public static func fetchUsage( + cookieHeaderOverride: String? = nil, + browserDetection: BrowserDetection = BrowserDetection(), + timeout: TimeInterval = 15.0, + logger: ((String) -> Void)? = nil) async throws -> AbacusUsageSnapshot + { + // Manual cookie header — no fallback, errors propagate directly + if let override = CookieHeaderNormalizer.normalize(cookieHeaderOverride) { + self.emit("Using manual cookie header", logger: logger) + return try await self.fetchWithCookieHeader(override, timeout: timeout, logger: logger) + } + + // Cached cookie header — fall back to a fresh browser import when the + // cached session is rejected or looks stale. + if let cached = CookieHeaderCache.load(provider: .abacus), + !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + self.emit("Using cached cookie header from \(cached.sourceLabel)", logger: logger) + do { + return try await self.fetchWithCookieHeader( + cached.cookieHeader, timeout: timeout, logger: logger) + } catch let error as AbacusUsageError where error.shouldTryNextImportedSession { + if error.shouldClearCachedCookie { + CookieHeaderCache.clear(provider: .abacus) + self.emit( + "Cached cookie failed (\(error.localizedDescription)); cleared, trying fresh import", + logger: logger) + } else { + self.emit( + "Cached cookie failed (\(error.localizedDescription)); trying fresh import", + logger: logger) + } + } + } + + // Fresh browser import — try Chrome first (AGENTS.md default), then broaden + // to all browsers if Chrome has no sessions OR if every imported Chrome + // session is exhausted without a successful fetch. + var lastError: AbacusUsageError = .noSessionCookie + if let snapshot = try await self.tryFetchFromBrowsers( + BrowserFetchRequest( + browserDetection: browserDetection, + preferredBrowsers: [.chrome], + label: "Chrome", + timeout: timeout, + logger: logger), + lastError: &lastError) + { + return snapshot + } + + self.emit("Chrome sessions exhausted; falling back to all browsers", logger: logger) + if let snapshot = try await self.tryFetchFromBrowsers( + BrowserFetchRequest( + browserDetection: browserDetection, + preferredBrowsers: [], + label: "all browsers", + timeout: timeout, + logger: logger), + lastError: &lastError) + { + return snapshot + } + + throw lastError + } + + /// Tries to import sessions from `preferredBrowsers` and fetch usage. Returns + /// the snapshot on success, nil if no sessions were available or every + /// imported session was exhausted without success. + private static func tryFetchFromBrowsers( + _ request: BrowserFetchRequest, + lastError: inout AbacusUsageError) async throws -> AbacusUsageSnapshot? + { + let sessions: [AbacusCookieImporter.SessionInfo] + do { + sessions = try AbacusCookieImporter.importSessions( + browserDetection: request.browserDetection, + preferredBrowsers: request.preferredBrowsers, + logger: request.logger) + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + self.emit( + "\(request.label) cookie import failed: \(error.localizedDescription)", + logger: request.logger) + return nil + } + + for session in sessions { + self.emit("Trying cookies from \(session.sourceLabel)", logger: request.logger) + do { + let snapshot = try await self.fetchWithCookieHeader( + session.cookieHeader, + timeout: request.timeout, + logger: request.logger) + CookieHeaderCache.store( + provider: .abacus, + cookieHeader: session.cookieHeader, + sourceLabel: session.sourceLabel) + return snapshot + } catch let error as AbacusUsageError where error.shouldTryNextImportedSession { + self.emit( + "\(session.sourceLabel): \(error.localizedDescription), trying next source", + logger: request.logger) + lastError = error + continue + } catch { + self.emit( + "\(session.sourceLabel): \(error.localizedDescription), trying next source", + logger: request.logger) + lastError = .networkError(error.localizedDescription) + continue + } + } + return nil + } + + // MARK: - API Requests + + private static func fetchWithCookieHeader( + _ cookieHeader: String, + timeout: TimeInterval, + logger: ((String) -> Void)? = nil) async throws -> AbacusUsageSnapshot + { + enum FetchPart: Sendable { + case computePoints(JSONDictionaryBox) + case billingInfoSuccess(JSONDictionaryBox) + case billingInfoFailure(String) + } + + // Fetch compute points (required, full timeout) and billing info + // (optional, shorter budget) concurrently. Billing is bounded so a + // slow/flaky billing endpoint can't delay credit rendering. + let billingBudget = min(timeout, 5.0) + + var computePointsResult: [String: Any]? + var billingInfoResult: [String: Any] = [:] + + try await withThrowingTaskGroup(of: FetchPart.self) { group in + group.addTask { + let result = try await self.fetchJSON( + url: self.computePointsURL, + method: "GET", + cookieHeader: cookieHeader, + timeout: timeout) + return .computePoints(JSONDictionaryBox(value: result)) + } + group.addTask { + do { + let result = try await self.fetchJSON( + url: self.billingInfoURL, + method: "POST", + cookieHeader: cookieHeader, + timeout: billingBudget) + return .billingInfoSuccess(JSONDictionaryBox(value: result)) + } catch { + return .billingInfoFailure(error.localizedDescription) + } + } + + while let result = try await group.next() { + switch result { + case let .computePoints(value): + computePointsResult = value.value + case let .billingInfoSuccess(value): + billingInfoResult = value.value + case let .billingInfoFailure(message): + self.emit( + "Billing info fetch failed: \(message); credits shown without plan/reset", + logger: logger) + } + } + } + + guard let computePointsResult else { + throw AbacusUsageError.networkError("Abacus compute points fetch did not complete") + } + + return try self.parseResults(computePoints: computePointsResult, billingInfo: billingInfoResult) + } + + private static func fetchJSON( + url: URL, method: String, cookieHeader: String, timeout: TimeInterval) async throws -> [String: Any] + { + var request = URLRequest(url: url) + request.httpMethod = method + request.timeoutInterval = timeout + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + if method == "POST" { + request.httpBody = Data("{}".utf8) + } + + let response = try await ProviderHTTPClient.shared.response(for: request) + let data = response.data + let statusCode = response.statusCode + + if statusCode == 401 || statusCode == 403 { + throw AbacusUsageError.unauthorized + } + + guard statusCode == 200 else { + let body = String(data: data.prefix(200), encoding: .utf8) ?? "" + throw AbacusUsageError.networkError("HTTP \(statusCode): \(body)") + } + + let parsed: Any + do { + parsed = try JSONSerialization.jsonObject(with: data) + } catch { + let preview = String(data: data.prefix(200), encoding: .utf8) ?? "" + throw AbacusUsageError.parseFailed( + "\(url.lastPathComponent): \(error.localizedDescription) — preview: \(preview)") + } + + guard let root = parsed as? [String: Any] else { + throw AbacusUsageError.parseFailed("\(url.lastPathComponent): top-level JSON is not a dictionary") + } + + guard root["success"] as? Bool == true, + let result = root["result"] as? [String: Any] + else { + let errorMsg = (root["error"] as? String ?? "Unknown error").lowercased() + if errorMsg.contains("expired") || errorMsg.contains("session") + || errorMsg.contains("login") || errorMsg.contains("authenticate") + || errorMsg.contains("unauthorized") || errorMsg.contains("unauthenticated") + || errorMsg.contains("forbidden") + { + throw AbacusUsageError.unauthorized + } + throw AbacusUsageError.parseFailed("\(url.lastPathComponent): \(errorMsg)") + } + + return result + } + + // MARK: - Parsing + + private static func parseResults( + computePoints: [String: Any], billingInfo: [String: Any]) throws -> AbacusUsageSnapshot + { + let totalCredits = self.double(from: computePoints["totalComputePoints"]) + let creditsLeft = self.double(from: computePoints["computePointsLeft"]) + + guard let totalCredits, let creditsLeft else { + let keys = computePoints.keys.sorted().joined(separator: ", ") + throw AbacusUsageError.parseFailed( + "Missing credit fields in compute points response. Keys: [\(keys)]") + } + + let creditsUsed = totalCredits - creditsLeft + + let nextBillingDate = billingInfo["nextBillingDate"] as? String + let currentTier = billingInfo["currentTier"] as? String + let resetsAt = self.parseDate(nextBillingDate) + + return AbacusUsageSnapshot( + creditsUsed: creditsUsed, + creditsTotal: totalCredits, + resetsAt: resetsAt, + planName: currentTier) + } + + private static func double(from value: Any?) -> Double? { + if let d = value as? Double { return d } + if let i = value as? Int { return Double(i) } + if let n = value as? NSNumber { return n.doubleValue } + return nil + } + + private static func parseDate(_ isoString: String?) -> Date? { + guard let isoString else { return nil } + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = formatter.date(from: isoString) { return date } + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: isoString) + } + + // MARK: - Logging + + private static func emit(_ message: String, logger: ((String) -> Void)?) { + logger?("[abacus] \(message)") + self.log.debug(message) + } +} + +#else + +// MARK: - Abacus (Unsupported) + +public enum AbacusUsageFetcher { + public static func fetchUsage( + cookieHeaderOverride _: String? = nil, + browserDetection _: BrowserDetection = BrowserDetection(), + timeout _: TimeInterval = 15.0, + logger _: ((String) -> Void)? = nil) async throws -> AbacusUsageSnapshot + { + throw AbacusUsageError.notSupported + } +} + +#endif diff --git a/Sources/CodexBarCore/Providers/Abacus/AbacusUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Abacus/AbacusUsageSnapshot.swift new file mode 100644 index 0000000..0186844 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Abacus/AbacusUsageSnapshot.swift @@ -0,0 +1,77 @@ +import Foundation + +// MARK: - Abacus Usage Snapshot + +public struct AbacusUsageSnapshot: Sendable { + public let creditsUsed: Double? + public let creditsTotal: Double? + public let resetsAt: Date? + public let planName: String? + + public init( + creditsUsed: Double? = nil, + creditsTotal: Double? = nil, + resetsAt: Date? = nil, + planName: String? = nil) + { + self.creditsUsed = creditsUsed + self.creditsTotal = creditsTotal + self.resetsAt = resetsAt + self.planName = planName + } + + public func toUsageSnapshot() -> UsageSnapshot { + let percentUsed: Double = if let used = self.creditsUsed, let total = self.creditsTotal, total > 0 { + (used / total) * 100.0 + } else { + 0 + } + + let resetDesc: String? = if let used = self.creditsUsed, let total = self.creditsTotal { + "\(Self.formatCredits(used)) / \(Self.formatCredits(total)) credits" + } else { + nil + } + + // Derive window from actual billing cycle when possible. + // Assume the cycle started one calendar month before resetsAt. + let windowMinutes: Int = if let resetDate = self.resetsAt, + let cycleStart = Calendar.current.date(byAdding: .month, value: -1, to: resetDate) + { + max(1, Int(resetDate.timeIntervalSince(cycleStart) / 60)) + } else { + 30 * 24 * 60 + } + + let primary = RateWindow( + usedPercent: percentUsed, + windowMinutes: windowMinutes, + resetsAt: self.resetsAt, + resetDescription: resetDesc) + + let identity = ProviderIdentitySnapshot( + providerID: .abacus, + accountEmail: nil, + accountOrganization: nil, + loginMethod: self.planName) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: Date(), + identity: identity) + } + + // MARK: - Formatting + + /// Thread-safe credit formatting — allocates per call to avoid shared mutable state. + private static func formatCredits(_ value: Double) -> String { + let formatter = NumberFormatter() + formatter.numberStyle = .decimal + formatter.locale = Locale(identifier: "en_US") + formatter.maximumFractionDigits = value >= 1000 ? 0 : 1 + return formatter.string(from: NSNumber(value: value)) ?? String(format: "%.0f", value) + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanAPIRegion.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanAPIRegion.swift new file mode 100644 index 0000000..d86585a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanAPIRegion.swift @@ -0,0 +1,133 @@ +import Foundation + +public enum AlibabaCodingPlanAPIRegion: String, CaseIterable, Sendable { + case international = "intl" + case chinaMainland = "cn" + + private static let endpointPath = "data/api.json" + + public var displayName: String { + switch self { + case .international: + "International (modelstudio.console.alibabacloud.com)" + case .chinaMainland: + "China mainland (bailian.console.aliyun.com)" + } + } + + public var gatewayBaseURLString: String { + switch self { + case .international: + "https://modelstudio.console.alibabacloud.com" + case .chinaMainland: + "https://bailian.console.aliyun.com" + } + } + + public var dashboardURL: URL { + switch self { + case .international: + URL( + string: "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=coding-plan#/efm/coding_plan")! + case .chinaMainland: + URL(string: "https://bailian.console.aliyun.com/cn-beijing/?tab=model#/efm/coding_plan")! + } + } + + public var consoleDomain: String { + switch self { + case .international: + "modelstudio.console.alibabacloud.com" + case .chinaMainland: + "bailian.console.aliyun.com" + } + } + + public var consoleSite: String { + switch self { + case .international: + "MODELSTUDIO_ALIBABACLOUD" + case .chinaMainland: + "BAILIAN_ALIYUN" + } + } + + public var consoleRefererURL: URL { + switch self { + case .international: + URL(string: "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=coding-plan")! + case .chinaMainland: + URL(string: "https://bailian.console.aliyun.com/cn-beijing/?tab=model")! + } + } + + public var quotaURL: URL { + var components = URLComponents(string: self.gatewayBaseURLString)! + components.path = "/" + Self.endpointPath + components.queryItems = [ + URLQueryItem( + name: "action", + value: "zeldaEasy.broadscope-bailian.codingPlan.queryCodingPlanInstanceInfoV2"), + URLQueryItem(name: "product", value: "broadscope-bailian"), + URLQueryItem(name: "api", value: "queryCodingPlanInstanceInfoV2"), + URLQueryItem(name: "currentRegionId", value: self.currentRegionID), + ] + return components.url! + } + + public var consoleRPCBaseURLString: String { + switch self { + case .international: + "https://bailian-singapore-cs.alibabacloud.com" + case .chinaMainland: + "https://bailian-cs.console.aliyun.com" + } + } + + public var consoleRPCURL: URL { + var components = URLComponents(string: self.consoleRPCBaseURLString)! + components.path = "/" + Self.endpointPath + components.queryItems = [ + URLQueryItem(name: "action", value: self.consoleRPCAction), + URLQueryItem(name: "product", value: self.consoleRPCProduct), + URLQueryItem(name: "api", value: self.consoleQuotaAPIName), + URLQueryItem(name: "_v", value: "undefined"), + ] + return components.url! + } + + public var consoleRPCAction: String { + switch self { + case .international: + "IntlBroadScopeAspnGateway" + case .chinaMainland: + "BroadScopeAspnGateway" + } + } + + public var consoleRPCProduct: String { + "sfm_bailian" + } + + public var consoleQuotaAPIName: String { + "zeldaEasy.broadscope-bailian.codingPlan.queryCodingPlanInstanceInfoV2" + } + + public var commodityCode: String { + switch self { + case .international: + "sfm_codingplan_public_intl" + case .chinaMainland: + "sfm_codingplan_public_cn" + } + } + + public var currentRegionID: String { + switch self { + case .international: + "ap-southeast-1" + case .chinaMainland: + "cn-beijing" + } + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanCookieImporter.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanCookieImporter.swift new file mode 100644 index 0000000..1eb1e1f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanCookieImporter.swift @@ -0,0 +1,510 @@ +import Foundation + +#if os(macOS) +import CommonCrypto +import Security +import SQLite3 +import SweetCookieKit + +private let alibabaCookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.alibaba]?.browserCookieOrder ?? Browser.defaultImportOrder + +public enum AlibabaCodingPlanCookieImporter { + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = [ + "bailian-singapore-cs.alibabacloud.com", + "bailian-cs.console.aliyun.com", + "bailian-beijing-cs.aliyuncs.com", + "modelstudio.console.alibabacloud.com", + "bailian.console.aliyun.com", + "free.aliyun.com", + "account.aliyun.com", + "signin.aliyun.com", + "passport.alibabacloud.com", + "console.alibabacloud.com", + "console.aliyun.com", + "alibabacloud.com", + "aliyun.com", + ] + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String { + var byName: [String: HTTPCookie] = [:] + byName.reserveCapacity(self.cookies.count) + + for cookie in self.cookies { + if let expiry = cookie.expiresDate, expiry < Date() { + continue + } + guard !cookie.value.isEmpty else { continue } + if let existing = byName[cookie.name] { + let existingExpiry = existing.expiresDate ?? .distantPast + let candidateExpiry = cookie.expiresDate ?? .distantPast + if candidateExpiry >= existingExpiry { + byName[cookie.name] = cookie + } + } else { + byName[cookie.name] = cookie + } + } + + return byName.keys.sorted().compactMap { name in + guard let cookie = byName[name] else { return nil } + return "\(cookie.name)=\(cookie.value)" + }.joined(separator: "; ") + } + } + + nonisolated(unsafe) static var importSessionOverrideForTesting: + ((BrowserDetection, ((String) -> Void)?) throws -> SessionInfo)? + + public static func importSession( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) throws -> SessionInfo + { + if let override = self.importSessionOverrideForTesting { + return try override(browserDetection, logger) + } + let log: (String) -> Void = { msg in logger?("[alibaba-cookie] \(msg)") } + var accessDeniedHints: [String] = [] + var failureDetails: [String] = [] + let installedBrowsers = self.cookieImportCandidates(browserDetection: browserDetection) + log("Cookie import candidates: \(installedBrowsers.map(\.displayName).joined(separator: ", "))") + + for browserSource in installedBrowsers { + do { + log("Checking \(browserSource.displayName)") + let query = BrowserCookieQuery(domains: self.cookieDomains) + let sources = try Self.cookieClient.codexBarRecords( + matching: query, + in: browserSource, + logger: log) + if sources.isEmpty { + log("No matching cookie records in \(browserSource.displayName)") + if let fallbackSession = try Self.importChromiumFallbackSession( + browser: browserSource, + logger: log) + { + return fallbackSession + } + } + for source in sources where !source.records.isEmpty { + let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin) + if self.isAuthenticatedSession(cookies: httpCookies) { + log("Found \(httpCookies.count) Alibaba cookies in \(source.label)") + return SessionInfo(cookies: httpCookies, sourceLabel: source.label) + } + let cookieNames = Set(httpCookies.map(\.name)) + let hasTicket = cookieNames.contains("login_aliyunid_ticket") + let hasAccount = + cookieNames.contains("login_aliyunid_pk") || + cookieNames.contains("login_current_pk") || + cookieNames.contains("login_aliyunid") + log("Skipping \(source.label): missing auth cookies (ticket=\(hasTicket), account=\(hasAccount))") + } + if let fallbackSession = try Self.importChromiumFallbackSession(browser: browserSource, logger: log) { + return fallbackSession + } + } catch let error as BrowserCookieError { + BrowserCookieAccessGate.recordIfNeeded(error) + if let hint = error.accessDeniedHint { + accessDeniedHints.append(hint) + } + failureDetails.append("\(browserSource.displayName): \(error.localizedDescription)") + log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)") + } catch { + failureDetails.append("\(browserSource.displayName): \(error.localizedDescription)") + log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)") + } + } + + let details = (Array(Set(accessDeniedHints)).sorted() + Array(Set(failureDetails)).sorted()) + .joined(separator: " ") + throw AlibabaCodingPlanSettingsError.missingCookie(details: details.isEmpty ? nil : details) + } + + private static func isAuthenticatedSession(cookies: [HTTPCookie]) -> Bool { + guard !cookies.isEmpty else { return false } + let names = Set(cookies.map(\.name)) + let hasTicket = names.contains("login_aliyunid_ticket") + let hasAccount = + names.contains("login_aliyunid_pk") || + names.contains("login_current_pk") || + names.contains("login_aliyunid") + return hasTicket && hasAccount + } + + public static func hasSession( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) -> Bool + { + do { + _ = try self.importSession(browserDetection: browserDetection, logger: logger) + return true + } catch { + return false + } + } + + private static func importChromiumFallbackSession( + browser: Browser, + logger: ((String) -> Void)? = nil) throws -> SessionInfo? + { + guard browser.usesChromiumProfileStore else { return nil } + return try AlibabaChromiumCookieFallbackImporter.importSession( + browser: browser, + domains: self.cookieDomains, + logger: logger) + } + + static func cookieImportCandidates( + browserDetection: BrowserDetection, + importOrder: BrowserCookieImportOrder = alibabaCookieImportOrder) -> [Browser] + { + importOrder.cookieImportCandidates(using: browserDetection) + } + + static func matchesCookieDomain(_ domain: String, patterns: [String] = Self.cookieDomains) -> Bool { + let normalized = self.normalizeCookieDomain(domain) + return patterns.contains { pattern in + let normalizedPattern = self.normalizeCookieDomain(pattern) + return normalized == normalizedPattern || normalized.hasSuffix(".\(normalizedPattern)") + } + } + + static func normalizeCookieDomain(_ domain: String) -> String { + let trimmed = domain.trimmingCharacters(in: .whitespacesAndNewlines) + let normalized = trimmed.hasPrefix(".") ? String(trimmed.dropFirst()) : trimmed + return normalized.lowercased() + } +} + +enum AlibabaChromiumCookieFallbackImporter { + private struct ChromiumCookieRecord { + let domain: String + let name: String + let path: String + let value: String + let expires: Date? + let isSecure: Bool + } + + enum ImportError: LocalizedError { + case keyUnavailable(browser: Browser) + case keychainDenied(browser: Browser) + case sqliteFailed(label: String, details: String) + + var errorDescription: String? { + switch self { + case let .keyUnavailable(browser): + "\(browser.displayName) Safe Storage key not found." + case let .keychainDenied(browser): + "macOS Keychain denied access to \(browser.displayName) Safe Storage." + case let .sqliteFailed(label, details): + "\(label) cookie fallback failed: \(details)" + } + } + } + + static func importSession( + browser: Browser, + domains: [String], + cookieClient: BrowserCookieClient = BrowserCookieClient(), + logger: ((String) -> Void)? = nil) throws -> AlibabaCodingPlanCookieImporter.SessionInfo? + { + let stores = try cookieClient.codexBarStores(for: browser).filter { $0.databaseURL != nil } + guard !stores.isEmpty else { return nil } + + logger?("[alibaba-cookie] Trying \(browser.displayName) Chromium fallback") + let keys = try self.derivedKeys(for: browser) + for store in stores { + let cookies = try self.loadCookies(from: store, domains: domains, keys: keys) + guard !cookies.isEmpty else { continue } + if self.isAuthenticatedSession(cookies) { + logger?("[alibaba-cookie] Found \(cookies.count) Alibaba cookies via \(store.label) fallback") + return AlibabaCodingPlanCookieImporter.SessionInfo(cookies: cookies, sourceLabel: store.label) + } + } + return nil + } + + private static func isAuthenticatedSession(_ cookies: [HTTPCookie]) -> Bool { + let names = Set(cookies.map(\.name)) + let hasTicket = names.contains("login_aliyunid_ticket") + let hasAccount = + names.contains("login_aliyunid_pk") || + names.contains("login_current_pk") || + names.contains("login_aliyunid") + return hasTicket && hasAccount + } + + private static func loadCookies( + from store: BrowserCookieStore, + domains: [String], + keys: [Data]) throws -> [HTTPCookie] + { + guard let sourceDB = store.databaseURL else { return [] } + let records = try self.readCookiesFromLockedDB( + sourceDB: sourceDB, + domains: domains, + keys: keys, + label: store.label) + return records.compactMap(self.makeCookie) + } + + private static func readCookiesFromLockedDB( + sourceDB: URL, + domains: [String], + keys: [Data], + label: String) throws -> [ChromiumCookieRecord] + { + let tempDir = FileManager.default.temporaryDirectory + .appendingPathComponent("alibaba-chromium-cookies-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true) + + let copiedDB = tempDir.appendingPathComponent("Cookies") + try FileManager.default.copyItem(at: sourceDB, to: copiedDB) + for suffix in ["-wal", "-shm"] { + let src = URL(fileURLWithPath: sourceDB.path + suffix) + if FileManager.default.fileExists(atPath: src.path) { + let dst = URL(fileURLWithPath: copiedDB.path + suffix) + try? FileManager.default.copyItem(at: src, to: dst) + } + } + defer { try? FileManager.default.removeItem(at: tempDir) } + + return try self.readCookies(fromDB: copiedDB.path, domains: domains, keys: keys, label: label) + } + + private static func readCookies( + fromDB path: String, + domains: [String], + keys: [Data], + label: String) throws -> [ChromiumCookieRecord] + { + var db: OpaquePointer? + guard sqlite3_open_v2(path, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { + throw ImportError.sqliteFailed(label: label, details: String(cString: sqlite3_errmsg(db))) + } + defer { sqlite3_close(db) } + + let sql = "SELECT host_key, name, path, expires_utc, is_secure, value, encrypted_value FROM cookies" + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, sql, -1, &stmt, nil) == SQLITE_OK else { + throw ImportError.sqliteFailed(label: label, details: String(cString: sqlite3_errmsg(db))) + } + defer { sqlite3_finalize(stmt) } + + var records: [ChromiumCookieRecord] = [] + while sqlite3_step(stmt) == SQLITE_ROW { + guard let hostKey = self.readText(stmt, index: 0), self.matches(domain: hostKey, patterns: domains) else { + continue + } + guard let name = self.readText(stmt, index: 1), let path = self.readText(stmt, index: 2) else { + continue + } + + let value: String? = if let plain = self.readText(stmt, index: 5), !plain.isEmpty { + plain + } else if let encrypted = self.readBlob(stmt, index: 6) { + self.decrypt(encrypted, usingAnyOf: keys) + } else { + nil + } + guard let value, !value.isEmpty else { continue } + + records.append(ChromiumCookieRecord( + domain: AlibabaCodingPlanCookieImporter.normalizeCookieDomain(hostKey), + name: name, + path: path, + value: value, + expires: self.chromiumExpiry(sqlite3_column_int64(stmt, 3)), + isSecure: sqlite3_column_int(stmt, 4) != 0)) + } + + return records.filter { record in + guard let expires = record.expires else { return true } + return expires >= Date() + } + } + + private static func derivedKeys(for browser: Browser) throws -> [Data] { + var keys: [Data] = [] + var sawDenied = false + + for label in browser.safeStorageLabels { + switch KeychainAccessPreflight.checkGenericPassword(service: label.service, account: label.account) { + case .interactionRequired: + sawDenied = true + continue + case .allowed, .notFound, .failure: + break + } + + if let password = self.safeStoragePassword(service: label.service, account: label.account) { + keys.append(self.deriveKey(from: password)) + } + } + + if !keys.isEmpty { + return keys + } + if sawDenied { + throw ImportError.keychainDenied(browser: browser) + } + throw ImportError.keyUnavailable(browser: browser) + } + + private static func safeStoragePassword(service: String, account: String) -> String? { + // The preflight classifies prompt-requiring items as .interactionRequired, but its + // .notFound (gate disabled) and .failure outcomes still reach this read. Honor the + // access gate and keep the read strictly non-interactive so it can never prompt. + guard !KeychainAccessGate.isDisabled else { return nil } + var query: [String: Any] = [ + kSecClass as String: kSecClassGenericPassword, + kSecAttrService as String: service, + kSecAttrAccount as String: account, + kSecMatchLimit as String: kSecMatchLimitOne, + kSecReturnData as String: true, + ] + KeychainNoUIQuery.apply(to: &query) + + var result: AnyObject? + let status = KeychainSecurity.copyMatching(query as CFDictionary, &result) + guard status == errSecSuccess, let data = result as? Data else { return nil } + return String(data: data, encoding: .utf8) + } + + private static func deriveKey(from password: String) -> Data { + let salt = Data("saltysalt".utf8) + var key = Data(count: kCCKeySizeAES128) + let keyLength = key.count + _ = key.withUnsafeMutableBytes { keyBytes in + password.utf8CString.withUnsafeBytes { passBytes in + salt.withUnsafeBytes { saltBytes in + CCKeyDerivationPBKDF( + CCPBKDFAlgorithm(kCCPBKDF2), + passBytes.bindMemory(to: Int8.self).baseAddress, + passBytes.count - 1, + saltBytes.bindMemory(to: UInt8.self).baseAddress, + salt.count, + CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA1), + 1003, + keyBytes.bindMemory(to: UInt8.self).baseAddress, + keyLength) + } + } + } + return key + } + + private static func decrypt(_ encryptedValue: Data, usingAnyOf keys: [Data]) -> String? { + for key in keys { + if let value = self.decrypt(encryptedValue, key: key) { + return value + } + } + return nil + } + + private static func decrypt(_ encryptedValue: Data, key: Data) -> String? { + guard encryptedValue.count > 3 else { return nil } + let prefix = String(data: encryptedValue.prefix(3), encoding: .utf8) + guard prefix == "v10" else { return nil } + + let payload = Data(encryptedValue.dropFirst(3)) + let iv = Data(repeating: 0x20, count: kCCBlockSizeAES128) + var outLength = 0 + var out = Data(count: payload.count + kCCBlockSizeAES128) + let outCapacity = out.count + + let status = out.withUnsafeMutableBytes { outBytes in + payload.withUnsafeBytes { payloadBytes in + key.withUnsafeBytes { keyBytes in + iv.withUnsafeBytes { ivBytes in + CCCrypt( + CCOperation(kCCDecrypt), + CCAlgorithm(kCCAlgorithmAES), + CCOptions(kCCOptionPKCS7Padding), + keyBytes.baseAddress, + key.count, + ivBytes.baseAddress, + payloadBytes.baseAddress, + payload.count, + outBytes.baseAddress, + outCapacity, + &outLength) + } + } + } + } + + guard status == kCCSuccess else { return nil } + out.count = outLength + + if let value = String(data: out, encoding: .utf8), !value.isEmpty { + return value + } + if out.count > 32 { + let trimmed = out.dropFirst(32) + if let value = String(data: trimmed, encoding: .utf8), !value.isEmpty { + return value + } + } + return nil + } + + private static func makeCookie(from record: ChromiumCookieRecord) -> HTTPCookie? { + var properties: [HTTPCookiePropertyKey: Any] = [ + .domain: record.domain, + .path: record.path, + .name: record.name, + .value: record.value, + ] + if record.isSecure { + properties[.secure] = true + } + if let expires = record.expires { + properties[.expires] = expires + } + return HTTPCookie(properties: properties) + } + + private static func readText(_ stmt: OpaquePointer?, index: Int32) -> String? { + guard sqlite3_column_type(stmt, index) != SQLITE_NULL, + let value = sqlite3_column_text(stmt, index) + else { + return nil + } + return String(cString: value) + } + + private static func readBlob(_ stmt: OpaquePointer?, index: Int32) -> Data? { + guard sqlite3_column_type(stmt, index) != SQLITE_NULL, + let bytes = sqlite3_column_blob(stmt, index) + else { + return nil + } + return Data(bytes: bytes, count: Int(sqlite3_column_bytes(stmt, index))) + } + + private static func matches(domain: String, patterns: [String]) -> Bool { + AlibabaCodingPlanCookieImporter.matchesCookieDomain(domain, patterns: patterns) + } + + private static func chromiumExpiry(_ expiresUTC: Int64) -> Date? { + guard expiresUTC > 0 else { return nil } + let seconds = (Double(expiresUTC) / 1_000_000.0) - 11_644_473_600.0 + guard seconds > 0 else { return nil } + return Date(timeIntervalSince1970: seconds) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanProviderDescriptor.swift new file mode 100644 index 0000000..3958664 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanProviderDescriptor.swift @@ -0,0 +1,266 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit +#endif + +public enum AlibabaCodingPlanProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + #if os(macOS) + let browserOrder: BrowserCookieImportOrder = [ + .chrome, + .chromeBeta, + .brave, + .edge, + .arc, + .firefox, + .safari, + ] + #else + let browserOrder: BrowserCookieImportOrder? = nil + #endif + + return ProviderDescriptor( + id: .alibaba, + metadata: ProviderMetadata( + id: .alibaba, + displayName: "Alibaba", + sessionLabel: "5-hour", + weeklyLabel: "Weekly", + opusLabel: "Monthly", + supportsOpus: true, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Alibaba usage", + cliName: "alibaba-coding-plan", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: browserOrder, + dashboardURL: AlibabaCodingPlanAPIRegion.international.dashboardURL.absoluteString, + statusPageURL: nil, + statusLinkURL: "https://status.aliyun.com"), + branding: ProviderBranding( + iconStyle: .alibaba, + iconResourceName: "ProviderIcon-alibaba", + color: ProviderColor(red: 1.0, green: 106 / 255, blue: 0)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Alibaba Coding Plan cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), + cli: ProviderCLIConfig( + name: "alibaba-coding-plan", + aliases: ["alibaba", "bailian"], + versionDetector: nil)) + } + + private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] { + switch context.sourceMode { + case .web: + return [AlibabaCodingPlanWebFetchStrategy()] + case .api: + return [AlibabaCodingPlanAPIFetchStrategy()] + case .cli, .oauth: + return [] + case .auto: + break + } + + if context.settings?.alibaba?.cookieSource == .off { + return [AlibabaCodingPlanAPIFetchStrategy()] + } + + return [AlibabaCodingPlanWebFetchStrategy(), AlibabaCodingPlanAPIFetchStrategy()] + } +} + +struct AlibabaCodingPlanWebFetchStrategy: ProviderFetchStrategy { + let id: String = "alibaba-coding-plan.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + guard context.settings?.alibaba?.cookieSource != .off else { return false } + + if AlibabaCodingPlanSettingsReader.cookieHeader(environment: context.env) != nil { + return true + } + + if let settings = context.settings?.alibaba, + settings.cookieSource == .manual + { + return CookieHeaderNormalizer.normalize(settings.manualCookieHeader) != nil + } + + #if os(macOS) + if let cached = CookieHeaderCache.load(provider: .alibaba), + !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + return true + } + if AlibabaCodingPlanCookieImporter.hasSession(browserDetection: context.browserDetection) { + return true + } + return false + #else + return false + #endif + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let cookieSource = context.settings?.alibaba?.cookieSource ?? .auto + let cookieHeader = try Self.resolveCookieHeader(context: context, allowCached: true) + do { + let region = context.settings?.alibaba?.apiRegion ?? .international + let usage = try await AlibabaCodingPlanUsageFetcher.fetchUsage( + cookieHeader: cookieHeader, + region: region, + environment: context.env) + return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web") + } catch let error as AlibabaCodingPlanUsageError + where (error == .invalidCredentials || error == .loginRequired) && cookieSource != .manual + { + #if os(macOS) + CookieHeaderCache.clear(provider: .alibaba) + let refreshedHeader = try Self.resolveCookieHeader(context: context, allowCached: false) + let region = context.settings?.alibaba?.apiRegion ?? .international + let usage = try await AlibabaCodingPlanUsageFetcher.fetchUsage( + cookieHeader: refreshedHeader, + region: region, + environment: context.env) + return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web") + #else + throw AlibabaCodingPlanUsageError.invalidCredentials + #endif + } + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + guard context.sourceMode == .auto else { return false } + + if let urlError = error as? URLError { + switch urlError.code { + case .timedOut, + .cannotFindHost, + .cannotConnectToHost, + .networkConnectionLost, + .dnsLookupFailed, + .secureConnectionFailed, + .serverCertificateHasBadDate, + .serverCertificateUntrusted, + .serverCertificateHasUnknownRoot, + .serverCertificateNotYetValid, + .clientCertificateRejected, + .clientCertificateRequired, + .cannotLoadFromNetwork, + .internationalRoamingOff, + .callIsActive, + .dataNotAllowed, + .requestBodyStreamExhausted, + .resourceUnavailable, + .notConnectedToInternet: + return true + default: + break + } + } + + if let settingsError = error as? AlibabaCodingPlanSettingsError { + switch settingsError { + case .missingCookie, .invalidCookie: + return true + case .missingToken: + return false + } + } + + guard let alibabaError = error as? AlibabaCodingPlanUsageError else { return false } + switch alibabaError { + case .loginRequired: + return true + case .invalidCredentials: + return true + case .apiKeyUnavailableInRegion: + return false + case let .apiError(message): + return message.contains("HTTP 404") || message.contains("HTTP 403") + case .networkError: + return true + case .parseFailed: + return false + } + } + + static func resolveCookieHeader(context: ProviderFetchContext, allowCached: Bool) throws -> String { + if let settings = context.settings?.alibaba, + settings.cookieSource == .manual + { + guard let header = CookieHeaderNormalizer.normalize(settings.manualCookieHeader) else { + throw AlibabaCodingPlanSettingsError.invalidCookie + } + return header + } + + if let envCookie = AlibabaCodingPlanSettingsReader.cookieHeader(environment: context.env), + let normalized = CookieHeaderNormalizer.normalize(envCookie) + { + return normalized + } + + #if os(macOS) + if allowCached, + let cached = CookieHeaderCache.load(provider: .alibaba), + !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + return cached.cookieHeader + } + + do { + let session = try AlibabaCodingPlanCookieImporter.importSession(browserDetection: context.browserDetection) + CookieHeaderCache.store( + provider: .alibaba, + cookieHeader: session.cookieHeader, + sourceLabel: session.sourceLabel) + return session.cookieHeader + } catch { + throw error + } + #else + throw AlibabaCodingPlanSettingsError.missingCookie() + #endif + } +} + +struct AlibabaCodingPlanAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "alibaba-coding-plan.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolveToken(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveToken(environment: context.env) else { + throw AlibabaCodingPlanSettingsError.missingToken + } + let region = context.settings?.alibaba?.apiRegion ?? .international + let usage = try await AlibabaCodingPlanUsageFetcher.fetchUsage( + apiKey: apiKey, + region: region, + environment: context.env) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.alibabaToken(environment: environment) + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanSettingsReader.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanSettingsReader.swift new file mode 100644 index 0000000..c082477 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanSettingsReader.swift @@ -0,0 +1,118 @@ +import Foundation + +public struct AlibabaCodingPlanSettingsReader: Sendable { + private static let endpointValidator = ProviderEndpointOverrideValidator( + allowedHosts: [ + "modelstudio.console.alibabacloud.com", + "bailian-singapore-cs.alibabacloud.com", + "bailian.console.aliyun.com", + "bailian-cs.console.aliyun.com", + "bailian-beijing-cs.aliyuncs.com", + ]) + + public static let apiTokenKey = "ALIBABA_CODING_PLAN_API_KEY" + public static let qwenAPITokenKey = "ALIBABA_QWEN_API_KEY" + public static let dashScopeAPITokenKey = "DASHSCOPE_API_KEY" + public static let apiTokenEnvironmentKeys = [ + Self.apiTokenKey, + Self.qwenAPITokenKey, + Self.dashScopeAPITokenKey, + ] + public static let cookieHeaderKey = "ALIBABA_CODING_PLAN_COOKIE" + public static let hostKey = "ALIBABA_CODING_PLAN_HOST" + public static let quotaURLKey = "ALIBABA_CODING_PLAN_QUOTA_URL" + public static let requireProviderEndpointOverridesKey = "ALIBABA_CODING_PLAN_REQUIRE_PROVIDER_ENDPOINT_OVERRIDES" + private static let endpointOverrideKeys = [ + Self.hostKey, + Self.quotaURLKey, + ] + + public static func apiToken( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + for key in self.apiTokenEnvironmentKeys { + if let token = self.cleaned(environment[key]) { return token } + } + return nil + } + + public static func hostOverride( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.endpointValidator.validatedHost( + self.cleaned(environment[self.hostKey]), + policy: self.endpointOverrideHostPolicy(environment: environment)) + } + + public static func rejectedEndpointOverrideKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + let policy = self.endpointOverrideHostPolicy(environment: environment) + return self.endpointOverrideKeys.first { key in + guard let value = self.cleaned(environment[key]) else { return false } + if key == Self.hostKey { + return self.endpointValidator.validatedHost(value, policy: policy) == nil + } + return self.endpointValidator.validatedURL(value, policy: policy) == nil + } + } + + public static func cookieHeader( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.cookieHeaderKey]) + } + + public static func quotaURL( + environment: [String: String] = ProcessInfo.processInfo.environment) -> URL? + { + self.endpointValidator.validatedURL( + self.cleaned(environment[self.quotaURLKey]), + policy: self.endpointOverrideHostPolicy(environment: environment)) + } + + static func endpointOverrideHostPolicy(environment: [String: String]) -> ProviderEndpointOverrideValidator + .HostPolicy { + guard let value = self.cleaned(environment[self.requireProviderEndpointOverridesKey])?.lowercased(), + ["1", "true", "yes", "on"].contains(value) + else { return .allowAnyHTTPSHost } + return .providerOwnedOnly + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} + +public enum AlibabaCodingPlanSettingsError: LocalizedError, Sendable { + case missingToken + case missingCookie(details: String? = nil) + case invalidCookie + + public var errorDescription: String? { + switch self { + case .missingToken: + return "Alibaba Coding Plan API key not found. " + + "Set apiKey in ~/.codexbar/config.json, ALIBABA_CODING_PLAN_API_KEY, " + + "ALIBABA_QWEN_API_KEY, or DASHSCOPE_API_KEY." + case let .missingCookie(details): + let base = "No Alibaba Coding Plan session cookies found in browsers. " + + "If you use Safari, enable Full Disk Access for CodexBar/Terminal or paste a manual Cookie header." + guard let details, !details.isEmpty else { return base } + return "\(base) \(details)" + case .invalidCookie: + return "Alibaba Coding Plan cookie header is invalid." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanUsageFetcher.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanUsageFetcher.swift new file mode 100644 index 0000000..2eb7918 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanUsageFetcher.swift @@ -0,0 +1,1121 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +// swiftlint:disable type_body_length +public struct AlibabaCodingPlanUsageFetcher: Sendable { + private static let log = CodexBarLog.logger("alibaba-coding-plan") + private static let browserLikeUserAgent = + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36" + private static let safariLikeUserAgent = + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/605.1.15 (KHTML, like Gecko) Version/26.3 Safari/605.1.15" + + enum AuthMode { + case apiKey + case webSession + } + + public static func fetchUsage( + apiKey: String, + region: AlibabaCodingPlanAPIRegion = .international, + environment: [String: String] = ProcessInfo.processInfo.environment, + now: Date = Date(), + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> AlibabaCodingPlanUsageSnapshot + { + let cleanedKey = apiKey.trimmingCharacters(in: .whitespacesAndNewlines) + guard !cleanedKey.isEmpty else { + throw AlibabaCodingPlanUsageError.invalidCredentials + } + if let rejectedKey = AlibabaCodingPlanSettingsReader.rejectedEndpointOverrideKey(environment: environment) { + throw ProviderEndpointOverrideError.alibabaCodingPlan(rejectedKey) + } + + if region != .international { + return try await self.fetchUsageOnce( + apiKey: cleanedKey, + region: region, + environment: environment, + now: now, + transport: transport) + } + + do { + return try await self.fetchUsageOnce( + apiKey: cleanedKey, + region: .international, + environment: environment, + now: now, + transport: transport) + } catch let error as AlibabaCodingPlanUsageError { + guard error.shouldRetryOnAlternateRegion else { throw error } + Self.log.debug("Alibaba Coding Plan request failed on intl host; retrying cn host") + return try await self.fetchUsageOnce( + apiKey: cleanedKey, + region: .chinaMainland, + environment: environment, + now: now, + transport: transport) + } + } + + public static func fetchUsage( + cookieHeader: String, + region: AlibabaCodingPlanAPIRegion = .international, + environment: [String: String] = ProcessInfo.processInfo.environment, + now: Date = Date(), + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> AlibabaCodingPlanUsageSnapshot + { + guard let normalizedCookie = CookieHeaderNormalizer.normalize(cookieHeader) else { + throw AlibabaCodingPlanSettingsError.invalidCookie + } + if let rejectedKey = AlibabaCodingPlanSettingsReader.rejectedEndpointOverrideKey(environment: environment) { + throw ProviderEndpointOverrideError.alibabaCodingPlan(rejectedKey) + } + + if region != .international { + return try await self.fetchUsageOnce( + cookieHeader: normalizedCookie, + region: region, + environment: environment, + now: now, + transport: transport) + } + + do { + return try await self.fetchUsageOnce( + cookieHeader: normalizedCookie, + region: .international, + environment: environment, + now: now, + transport: transport) + } catch let error as AlibabaCodingPlanUsageError { + guard error.shouldRetryOnAlternateRegion else { throw error } + Self.log.debug("Alibaba Coding Plan cookie request failed on intl host; retrying cn host") + return try await self.fetchUsageOnce( + cookieHeader: normalizedCookie, + region: .chinaMainland, + environment: environment, + now: now, + transport: transport) + } + } + + private static func fetchUsageOnce( + apiKey: String, + region: AlibabaCodingPlanAPIRegion, + environment: [String: String], + now: Date, + transport: any ProviderHTTPTransport) async throws -> AlibabaCodingPlanUsageSnapshot + { + let url = self.resolveQuotaURL(region: region, environment: environment) + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.httpBody = self.queryCodingPlanAPIRequestBody(region: region) + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue(apiKey, forHTTPHeaderField: "x-api-key") + request.setValue(apiKey, forHTTPHeaderField: "X-DashScope-API-Key") + request.setValue(Self.browserLikeUserAgent, forHTTPHeaderField: "User-Agent") + request.setValue(region.gatewayBaseURLString, forHTTPHeaderField: "Origin") + request.setValue(region.dashboardURL.absoluteString, forHTTPHeaderField: "Referer") + + let response = try await transport.response(for: request) + let data = response.data + guard response.statusCode == 200 else { + if response.statusCode == 401 || response.statusCode == 403 { + throw AlibabaCodingPlanUsageError.invalidCredentials + } + let body = String(data: data, encoding: .utf8) ?? "" + Self.log.error("Alibaba Coding Plan returned \(response.statusCode): \(body)") + throw AlibabaCodingPlanUsageError.apiError("HTTP \(response.statusCode)") + } + + return try self.parseUsageSnapshot(from: data, now: now, authMode: .apiKey) + } + + private static func fetchUsageOnce( + cookieHeader: String, + region: AlibabaCodingPlanAPIRegion, + environment: [String: String], + now: Date, + transport: any ProviderHTTPTransport) async throws -> AlibabaCodingPlanUsageSnapshot + { + let url = self.resolveConsoleQuotaURL(region: region, environment: environment) + let secToken = try await self.resolveConsoleSECToken( + cookieHeader: cookieHeader, + region: region, + environment: environment, + transport: transport) + let anonymousID = self.extractCookieValue(name: "cna", from: cookieHeader) + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.httpBody = self.queryCodingPlanConsoleRequestBody( + region: region, + secToken: secToken, + anonymousID: anonymousID) + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.setValue("*/*", forHTTPHeaderField: "Accept") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + if let csrf = self.extractCookieValue(name: "login_aliyunid_csrf", from: cookieHeader) ?? + self.extractCookieValue(name: "csrf", from: cookieHeader) + { + request.setValue(csrf, forHTTPHeaderField: "x-xsrf-token") + request.setValue(csrf, forHTTPHeaderField: "x-csrf-token") + } + request.setValue("XMLHttpRequest", forHTTPHeaderField: "X-Requested-With") + request.setValue(Self.browserLikeUserAgent, forHTTPHeaderField: "User-Agent") + request.setValue(region.gatewayBaseURLString, forHTTPHeaderField: "Origin") + request.setValue(region.consoleRefererURL.absoluteString, forHTTPHeaderField: "Referer") + + let response = try await transport.response(for: request) + let data = response.data + guard response.statusCode == 200 else { + if response.statusCode == 401 || response.statusCode == 403 { + throw AlibabaCodingPlanUsageError.loginRequired + } + let body = String(data: data, encoding: .utf8) ?? "" + Self.log.error("Alibaba Coding Plan returned \(response.statusCode): \(body)") + throw AlibabaCodingPlanUsageError.apiError("HTTP \(response.statusCode)") + } + + return try self.parseUsageSnapshot(from: data, now: now, authMode: .webSession) + } + + private static func queryCodingPlanAPIRequestBody(region: AlibabaCodingPlanAPIRegion) -> Data { + let payload: [String: Any] = [ + "queryCodingPlanInstanceInfoRequest": [ + "commodityCode": region.commodityCode, + ], + ] + return (try? JSONSerialization.data(withJSONObject: payload, options: [])) ?? Data("{}".utf8) + } + + private static func queryCodingPlanConsoleRequestBody( + region: AlibabaCodingPlanAPIRegion, + secToken: String, + anonymousID: String?) -> Data + { + let traceID = UUID().uuidString.lowercased() + var cornerstoneParam: [String: Any] = [ + "feTraceId": traceID, + "feURL": region.dashboardURL.absoluteString, + "protocol": "V2", + "console": "ONE_CONSOLE", + "productCode": "p_efm", + "domain": region.consoleDomain, + "consoleSite": region.consoleSite, + "userNickName": "", + "userPrincipalName": "", + "xsp_lang": "en-US", + ] + if let anonymousID, !anonymousID.isEmpty { + cornerstoneParam["X-Anonymous-Id"] = anonymousID + } + + let paramsObject: [String: Any] = [ + "Api": region.consoleQuotaAPIName, + "V": "1.0", + "Data": [ + "queryCodingPlanInstanceInfoRequest": [ + "commodityCode": region.commodityCode, + "onlyLatestOne": true, + ], + "cornerstoneParam": cornerstoneParam, + ], + ] + + guard let paramsData = try? JSONSerialization.data(withJSONObject: paramsObject, options: []), + let paramsString = String(data: paramsData, encoding: .utf8) + else { + return Data() + } + + var components = URLComponents() + components.queryItems = [ + URLQueryItem(name: "params", value: paramsString), + URLQueryItem(name: "region", value: region.currentRegionID), + URLQueryItem(name: "sec_token", value: secToken), + ] + return Data((components.percentEncodedQuery ?? "").utf8) + } + + static func resolveConsoleQuotaURL( + region: AlibabaCodingPlanAPIRegion, + environment: [String: String]) -> URL + { + if let override = AlibabaCodingPlanSettingsReader.quotaURL(environment: environment) { + return override + } + if let host = AlibabaCodingPlanSettingsReader.hostOverride(environment: environment), + let hostURL = self.consoleURL(from: host, region: region) + { + return hostURL + } + return region.consoleRPCURL + } + + static func resolveQuotaURL( + region: AlibabaCodingPlanAPIRegion, + environment: [String: String]) -> URL + { + if let override = AlibabaCodingPlanSettingsReader.quotaURL(environment: environment) { + return override + } + if let host = AlibabaCodingPlanSettingsReader.hostOverride(environment: environment), + let hostURL = self.url(from: host, region: region) + { + return hostURL + } + return region.quotaURL + } + + static func resolveConsoleDashboardURL( + region: AlibabaCodingPlanAPIRegion, + environment: [String: String]) -> URL + { + if let override = AlibabaCodingPlanSettingsReader.hostOverride(environment: environment), + let hostURL = self.dashboardURL(from: override, region: region) + { + return hostURL + } + return region.dashboardURL + } + + static func url(from rawHost: String, region: AlibabaCodingPlanAPIRegion) -> URL? { + let cleaned = AlibabaCodingPlanSettingsReader.cleaned(rawHost) + guard let cleaned else { return nil } + + let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: cleaned) + guard let base else { return nil } + + var components = URLComponents(url: base, resolvingAgainstBaseURL: false) + components?.path = "/data/api.json" + components?.queryItems = [ + URLQueryItem( + name: "action", + value: "zeldaEasy.broadscope-bailian.codingPlan.queryCodingPlanInstanceInfoV2"), + URLQueryItem(name: "product", value: "broadscope-bailian"), + URLQueryItem(name: "api", value: "queryCodingPlanInstanceInfoV2"), + URLQueryItem(name: "currentRegionId", value: region.currentRegionID), + ] + return components?.url + } + + static func consoleURL(from rawHost: String, region: AlibabaCodingPlanAPIRegion) -> URL? { + let cleaned = AlibabaCodingPlanSettingsReader.cleaned(rawHost) + guard let cleaned else { return nil } + + let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: cleaned) + guard let base else { return nil } + + var components = URLComponents(url: base, resolvingAgainstBaseURL: false) + components?.path = "/data/api.json" + components?.queryItems = [ + URLQueryItem(name: "action", value: region.consoleRPCAction), + URLQueryItem(name: "product", value: region.consoleRPCProduct), + URLQueryItem(name: "api", value: region.consoleQuotaAPIName), + URLQueryItem(name: "_v", value: "undefined"), + ] + return components?.url + } + + private static func resolveConsoleSECToken( + cookieHeader: String, + region: AlibabaCodingPlanAPIRegion, + environment: [String: String], + transport: any ProviderHTTPTransport) async throws -> String + { + let cookieSECToken = self.extractCookieValue(name: "sec_token", from: cookieHeader) + + let dashboardURL = self.resolveConsoleDashboardURL(region: region, environment: environment) + + var request = URLRequest(url: dashboardURL) + request.httpMethod = "GET" + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue(Self.safariLikeUserAgent, forHTTPHeaderField: "User-Agent") + request.setValue( + "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", + forHTTPHeaderField: "Accept") + + if let response = try? await transport.response(for: request), + response.statusCode == 200, + let html = String(data: response.data, encoding: .utf8), + let token = self.extractConsoleSECToken(from: html), + !token.isEmpty + { + return token + } + + if let token = try? await self.fetchSECTokenFromUserInfo( + cookieHeader: cookieHeader, + region: region, + environment: environment, + transport: transport) + { + return token + } + + if let cookieSECToken, !cookieSECToken.isEmpty { + return cookieSECToken + } + + throw AlibabaCodingPlanUsageError.loginRequired + } + + static func dashboardURL(from rawHost: String, region: AlibabaCodingPlanAPIRegion) -> URL? { + let cleaned = AlibabaCodingPlanSettingsReader.cleaned(rawHost) + guard let cleaned else { return nil } + + let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: cleaned) + guard let base else { return nil } + + guard var components = URLComponents(url: base, resolvingAgainstBaseURL: false), + let dashboardComponents = URLComponents( + url: region.dashboardURL, + resolvingAgainstBaseURL: false) + else { + return nil + } + + components.path = dashboardComponents.path + components.percentEncodedQuery = dashboardComponents.percentEncodedQuery + components.fragment = dashboardComponents.fragment + return components.url + } + + private static func fetchSECTokenFromUserInfo( + cookieHeader: String, + region: AlibabaCodingPlanAPIRegion, + environment: [String: String], + transport: any ProviderHTTPTransport) async throws -> String? + { + let gatewayBaseURL = self.resolveConsoleGatewayBaseURL(region: region, environment: environment) + let userInfoURL = gatewayBaseURL.appendingPathComponent("tool/user/info.json") + var request = URLRequest(url: userInfoURL) + request.httpMethod = "GET" + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue(Self.safariLikeUserAgent, forHTTPHeaderField: "User-Agent") + request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept") + let referer = gatewayBaseURL.absoluteString.hasSuffix("/") ? gatewayBaseURL.absoluteString : gatewayBaseURL + .absoluteString + "/" + request.setValue(referer, forHTTPHeaderField: "Referer") + + let response = try await transport.response(for: request) + guard response.statusCode == 200 else { + return nil + } + + let object = try JSONSerialization.jsonObject(with: response.data, options: []) + let expanded = self.expandedJSON(object) + return self.findFirstString(forKeys: ["secToken", "sec_token"], in: expanded) + } + + private static func resolveConsoleGatewayBaseURL( + region: AlibabaCodingPlanAPIRegion, + environment: [String: String]) -> URL + { + if let host = AlibabaCodingPlanSettingsReader.hostOverride(environment: environment), + let hostURL = self.baseURL(from: host) + { + return hostURL + } + return URL(string: region.gatewayBaseURLString)! + } + + private static func baseURL(from rawHost: String) -> URL? { + let cleaned = AlibabaCodingPlanSettingsReader.cleaned(rawHost) + guard let cleaned else { return nil } + + let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: cleaned) + guard let base else { return nil } + + guard var components = URLComponents(url: base, resolvingAgainstBaseURL: false) else { + return nil + } + components.path = "" + components.query = nil + components.fragment = nil + return components.url + } + + static func parseUsageSnapshot( + from data: Data, + now: Date = Date(), + authMode: AuthMode = .webSession) throws -> AlibabaCodingPlanUsageSnapshot + { + guard !data.isEmpty else { + throw AlibabaCodingPlanUsageError.parseFailed("Empty response body") + } + + let object = try JSONSerialization.jsonObject(with: data, options: []) + let expanded = self.expandedJSON(object) + guard let dictionary = expanded as? [String: Any] else { + throw AlibabaCodingPlanUsageError.parseFailed("Unexpected payload") + } + + if let statusCode = self.findFirstInt(forKeys: ["statusCode", "status_code", "code"], in: dictionary), + statusCode != 0, + statusCode != 200 + { + let message = self.findFirstString( + forKeys: ["statusMessage", "status_msg", "message", "msg"], + in: dictionary) + ?? "status code \(statusCode)" + let lower = message.lowercased() + if statusCode == 401 || statusCode == 403 || lower.contains("api key") || lower.contains("unauthorized") { + throw AlibabaCodingPlanUsageError.invalidCredentials + } + throw AlibabaCodingPlanUsageError.apiError(message) + } + + if let codeText = self.findFirstString(forKeys: ["code", "status", "statusCode"], in: dictionary) { + let normalizedCode = codeText.lowercased() + if normalizedCode.contains("needlogin") || normalizedCode.contains("login") { + if authMode == .apiKey { + throw AlibabaCodingPlanUsageError.apiKeyUnavailableInRegion + } + throw AlibabaCodingPlanUsageError.loginRequired + } + } + if let messageText = self.findFirstString(forKeys: ["message", "msg", "statusMessage"], in: dictionary) { + let normalizedMessage = messageText.lowercased() + if normalizedMessage.contains("log in") || normalizedMessage.contains("login") { + if authMode == .apiKey { + throw AlibabaCodingPlanUsageError.apiKeyUnavailableInRegion + } + throw AlibabaCodingPlanUsageError.loginRequired + } + if authMode == .apiKey, + normalizedMessage.contains("console session") || + normalizedMessage.contains("api key mode may be unavailable") + { + throw AlibabaCodingPlanUsageError.apiKeyUnavailableInRegion + } + } + + let instanceInfo = self.findActiveInstanceInfo(in: dictionary, now: now) + let instanceCount = self.findFirstArray( + forKeys: ["codingPlanInstanceInfos", "coding_plan_instance_infos"], + in: dictionary)? + .compactMap { $0 as? [String: Any] } + .count ?? 0 + let shouldScopeQuotaToSelectedInstance = + instanceCount > 1 && + (instanceInfo.map { self.activeSignalScore(in: $0, now: now) > 0 } ?? false) + let quota = if shouldScopeQuotaToSelectedInstance, let instanceInfo { + self.findQuotaInfo(in: instanceInfo) + } else { + self.findQuotaInfo(in: instanceInfo ?? [:]) ?? self.findQuotaInfo(in: dictionary) + } + guard let quota else { + if let fallback = self.parsePlanVisibleActiveFallback( + payload: dictionary, + instanceInfo: instanceInfo, + now: now) + { + return fallback + } + let diagnostics = self.payloadDiagnostics(payload: dictionary) + Self.log.error("Alibaba coding plan quota payload missing expected fields: \(diagnostics)") + throw AlibabaCodingPlanUsageError.parseFailed("Missing coding plan quota data (\(diagnostics))") + } + + let planName = self.findPlanName(in: instanceInfo ?? [:]) ?? self.findPlanName(in: dictionary) + + let snapshot = AlibabaCodingPlanUsageSnapshot( + planName: planName, + fiveHourUsedQuota: self.anyInt(for: ["per5HourUsedQuota", "perFiveHourUsedQuota"], in: quota), + fiveHourTotalQuota: self.anyInt(for: ["per5HourTotalQuota", "perFiveHourTotalQuota"], in: quota), + fiveHourNextRefreshTime: self.anyDate( + for: ["per5HourQuotaNextRefreshTime", "perFiveHourQuotaNextRefreshTime"], + in: quota), + weeklyUsedQuota: self.anyInt(for: ["perWeekUsedQuota"], in: quota), + weeklyTotalQuota: self.anyInt(for: ["perWeekTotalQuota"], in: quota), + weeklyNextRefreshTime: self.anyDate(for: ["perWeekQuotaNextRefreshTime"], in: quota), + monthlyUsedQuota: self.anyInt(for: ["perBillMonthUsedQuota", "perMonthUsedQuota"], in: quota), + monthlyTotalQuota: self.anyInt(for: ["perBillMonthTotalQuota", "perMonthTotalQuota"], in: quota), + monthlyNextRefreshTime: self.anyDate( + for: ["perBillMonthQuotaNextRefreshTime", "perMonthQuotaNextRefreshTime"], + in: quota), + updatedAt: now) + + if snapshot.fiveHourTotalQuota == nil, + snapshot.weeklyTotalQuota == nil, + snapshot.monthlyTotalQuota == nil + { + if let fallback = self.parsePlanVisibleActiveFallback( + payload: dictionary, + instanceInfo: instanceInfo, + now: now) + { + return fallback + } + let diagnostics = self.payloadDiagnostics(payload: dictionary) + Self.log.error("Alibaba coding plan payload had no usable windows: \(diagnostics)") + throw AlibabaCodingPlanUsageError.parseFailed("No quota windows found in payload (\(diagnostics))") + } + + return snapshot + } + + private static func findPlanName(in payload: [String: Any]) -> String? { + if let infos = self.findFirstArray( + forKeys: ["codingPlanInstanceInfos", "coding_plan_instance_infos"], + in: payload) + { + for item in infos { + guard let info = item as? [String: Any] else { continue } + let candidates = [ + self.anyString(for: ["planName", "plan_name"], in: info), + self.anyString(for: ["instanceName", "instance_name"], in: info), + self.anyString(for: ["packageName", "package_name"], in: info), + ] + for candidate in candidates { + if let candidate, !candidate.isEmpty { + return candidate + } + } + } + } + return self.findFirstString(forKeys: ["planName", "plan_name", "packageName", "package_name"], in: payload) + } + + private static func findActiveInstanceInfo(in payload: [String: Any], now: Date = Date()) -> [String: Any]? { + guard let infos = self.findFirstArray( + forKeys: ["codingPlanInstanceInfos", "coding_plan_instance_infos"], + in: payload) + else { + return nil + } + + var first: [String: Any]? + var best: [String: Any]? + var bestScore = Int.min + for item in infos { + guard let info = item as? [String: Any] else { continue } + first = first ?? info + let score = self.activeSignalScore(in: info, now: now) + if score > bestScore { + best = info + bestScore = score + } + } + if bestScore > 0 { + return best + } + return first + } + + private static func parsePlanVisibleActiveFallback( + payload: [String: Any], + instanceInfo: [String: Any]?, + now: Date) -> AlibabaCodingPlanUsageSnapshot? + { + let source = instanceInfo ?? payload + guard self.hasPositiveActivePlanSignal(in: source, payload: payload, now: now), + let planName = self.findPlanName(in: source) ?? self.findPlanName(in: payload) + else { + return nil + } + + return AlibabaCodingPlanUsageSnapshot( + planName: planName, + fiveHourUsedQuota: nil, + fiveHourTotalQuota: nil, + fiveHourNextRefreshTime: nil, + weeklyUsedQuota: nil, + weeklyTotalQuota: nil, + weeklyNextRefreshTime: nil, + monthlyUsedQuota: nil, + monthlyTotalQuota: nil, + monthlyNextRefreshTime: nil, + updatedAt: now) + } + + private static func hasPositiveActivePlanSignal( + in source: [String: Any], + payload: [String: Any], + now: Date) -> Bool + { + if self.containsPlanInstances(in: payload) { + return self.activeSignalScore(in: source, now: now) > 0 + } + return self.activeSignalScore(in: source, now: now) > 0 || self.activeSignalScore(in: payload, now: now) > 0 + } + + private static func containsPlanInstances(in payload: [String: Any]) -> Bool { + guard let infos = self.findFirstArray( + forKeys: ["codingPlanInstanceInfos", "coding_plan_instance_infos"], + in: payload) + else { + return false + } + return infos.contains { $0 is [String: Any] } + } + + private static func activeSignalScore(in source: [String: Any], now: Date) -> Int { + if let status = self.anyString(for: ["status", "instanceStatus"], in: source)?.uppercased() { + if ["VALID", "ACTIVE"].contains(status) { + return 3 + } + if ["EXPIRED", "INVALID", "INACTIVE", "DISABLED", "TERMINATED", "STOPPED"].contains(status) { + return -1 + } + } + + if let isActive = self.anyBool(for: ["isActive", "active"], in: source) { + return isActive ? 3 : -1 + } + + if let expiry = self.anyDate(for: ["endTime", "periodEndTime", "expireTime", "expirationTime"], in: source), + expiry.timeIntervalSince(now) > 0 + { + return 1 + } + + return 0 + } + + private static func payloadDiagnostics(payload: [String: Any]) -> String { + let topKeys = payload.keys.sorted() + let dataDict = self.findFirstDictionary(forKeys: ["data", "successResponse", "success_response"], in: payload) + let dataKeys = dataDict?.keys.sorted() ?? [] + let instanceInfo = self.findActiveInstanceInfo(in: payload) + let instanceKeys = instanceInfo?.keys.sorted() ?? [] + let hasQuota = self.findQuotaInfo(in: payload) != nil + let planUsage = + self.anyString(for: ["planUsage", "usageRate", "usage", "usageValue"], in: instanceInfo ?? [:]) ?? + self.findFirstString(forKeys: ["planUsage", "usageRate", "usage", "usageValue"], in: payload) + let compactPlanUsage = planUsage?.replacingOccurrences(of: "\n", with: " ") ?? "none" + let status = self.anyString(for: ["status", "instanceStatus"], in: instanceInfo ?? [:]) ?? "none" + + return "topKeys=\(topKeys.joined(separator: ",")) dataKeys=\(dataKeys.joined(separator: ",")) " + + "instanceKeys=\(instanceKeys.joined(separator: ",")) hasQuota=\(hasQuota ? "1" : "0") " + + "status=\(status) planUsage=\(compactPlanUsage)" + } + + private static func findQuotaInfo(in payload: [String: Any]) -> [String: Any]? { + if let direct = self.findFirstDictionary( + forKeys: ["codingPlanQuotaInfo", "coding_plan_quota_info"], + in: payload) + { + return direct + } + return self.findFirstDictionary( + matchingAnyKey: [ + "per5HourUsedQuota", + "per5HourTotalQuota", + "perWeekUsedQuota", + "perWeekTotalQuota", + "perBillMonthUsedQuota", + "perBillMonthTotalQuota", + ], + in: payload) + } + + private static func findFirstDictionary(forKeys keys: [String], in value: Any) -> [String: Any]? { + guard let dict = value as? [String: Any] else { return nil } + for key in keys { + if let nested = dict[key] as? [String: Any] { + return nested + } + } + for nestedValue in dict.values { + if let nested = self.findFirstDictionary(forKeys: keys, in: nestedValue) { + return nested + } + } + if let array = value as? [Any] { + for item in array { + if let nested = self.findFirstDictionary(forKeys: keys, in: item) { + return nested + } + } + } + return nil + } + + private static func findFirstDictionary(matchingAnyKey keys: [String], in value: Any) -> [String: Any]? { + if let dict = value as? [String: Any] { + if keys.contains(where: { dict[$0] != nil }) { + return dict + } + for nestedValue in dict.values { + if let nested = self.findFirstDictionary(matchingAnyKey: keys, in: nestedValue) { + return nested + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let nested = self.findFirstDictionary(matchingAnyKey: keys, in: item) { + return nested + } + } + } + return nil + } + + private static func findFirstArray(forKeys keys: [String], in value: Any) -> [Any]? { + guard let dict = value as? [String: Any] else { + if let array = value as? [Any] { + for item in array { + if let found = self.findFirstArray(forKeys: keys, in: item) { + return found + } + } + } + return nil + } + for key in keys { + if let array = dict[key] as? [Any] { + return array + } + } + for nested in dict.values { + if let found = self.findFirstArray(forKeys: keys, in: nested) { + return found + } + } + return nil + } + + private static func findFirstInt(forKeys keys: [String], in value: Any) -> Int? { + if let dict = value as? [String: Any] { + for key in keys { + if let parsed = self.parseInt(dict[key]) { + return parsed + } + } + for nested in dict.values { + if let parsed = self.findFirstInt(forKeys: keys, in: nested) { + return parsed + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let parsed = self.findFirstInt(forKeys: keys, in: item) { + return parsed + } + } + } + return nil + } + + private static func findFirstString(forKeys keys: [String], in value: Any) -> String? { + if let dict = value as? [String: Any] { + for key in keys { + if let parsed = self.parseString(dict[key]) { + return parsed + } + } + for nested in dict.values { + if let parsed = self.findFirstString(forKeys: keys, in: nested) { + return parsed + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let parsed = self.findFirstString(forKeys: keys, in: item) { + return parsed + } + } + } + return nil + } + + private static func findFirstDate(forKeys keys: [String], in value: Any) -> Date? { + if let dict = value as? [String: Any] { + for key in keys { + if let parsed = self.parseDate(dict[key]) { + return parsed + } + } + for nested in dict.values { + if let parsed = self.findFirstDate(forKeys: keys, in: nested) { + return parsed + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let parsed = self.findFirstDate(forKeys: keys, in: item) { + return parsed + } + } + } + return nil + } + + private static func expandedJSON(_ value: Any) -> Any { + if let dict = value as? [String: Any] { + var expanded: [String: Any] = [:] + expanded.reserveCapacity(dict.count) + for (key, nested) in dict { + expanded[key] = self.expandedJSON(nested) + } + return expanded + } + if let array = value as? [Any] { + return array.map { self.expandedJSON($0) } + } + if let string = value as? String, + let data = string.data(using: .utf8), + let nested = try? JSONSerialization.jsonObject(with: data, options: []), + nested is [String: Any] || nested is [Any] + { + return self.expandedJSON(nested) + } + return value + } + + private static func anyInt(for keys: [String], in dict: [String: Any]) -> Int? { + for key in keys { + if let value = self.parseInt(dict[key]) { + return value + } + } + return nil + } + + private static func anyString(for keys: [String], in dict: [String: Any]) -> String? { + for key in keys { + if let value = self.parseString(dict[key]) { + return value + } + } + return nil + } + + private static func anyDate(for keys: [String], in dict: [String: Any]) -> Date? { + for key in keys { + if let value = self.parseDate(dict[key]) { + return value + } + } + return nil + } + + private static func anyPercent(for keys: [String], in dict: [String: Any]) -> Double? { + for key in keys { + if let value = self.parsePercent(dict[key]) { + return value + } + } + return nil + } + + private static func anyBool(for keys: [String], in dict: [String: Any]) -> Bool? { + for key in keys { + if let value = self.parseBool(dict[key]) { + return value + } + } + return nil + } + + private static func findFirstPercent(forKeys keys: [String], in value: Any) -> Double? { + if let dict = value as? [String: Any] { + for key in keys { + if let parsed = self.parsePercent(dict[key]) { + return parsed + } + } + for nested in dict.values { + if let parsed = self.findFirstPercent(forKeys: keys, in: nested) { + return parsed + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let parsed = self.findFirstPercent(forKeys: keys, in: item) { + return parsed + } + } + } + return nil + } + + private static func parseDate(_ raw: Any?) -> Date? { + if let intValue = self.parseInt(raw) { + if intValue > 1_000_000_000_000 { + return Date(timeIntervalSince1970: TimeInterval(intValue) / 1000) + } + if intValue > 1_000_000_000 { + return Date(timeIntervalSince1970: TimeInterval(intValue)) + } + } + if let string = self.parseString(raw) { + let formatter = ISO8601DateFormatter() + if let date = formatter.date(from: string) { + return date + } + let dateFormatter = DateFormatter() + dateFormatter.locale = Locale(identifier: "en_US_POSIX") + for format in ["yyyy-MM-dd HH:mm", "yyyy-MM-dd HH:mm:ss"] { + dateFormatter.dateFormat = format + if let date = dateFormatter.date(from: string) { + return date + } + } + } + return nil + } + + private static func parseInt(_ raw: Any?) -> Int? { + if let value = raw as? Int { return value } + if let value = raw as? Int64 { return Int(value) } + if let value = raw as? Double { return Int(value) } + if let value = raw as? NSNumber { return value.intValue } + if let value = raw as? String { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + return Int(trimmed) + } + return nil + } + + private static func parseString(_ raw: Any?) -> String? { + guard let value = raw as? String else { return nil } + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + } + + private static func parsePercent(_ raw: Any?) -> Double? { + if let intValue = self.parseInt(raw) { + return max(0, min(Double(intValue), 100)) + } + guard let rawString = self.parseString(raw) else { return nil } + let cleaned = rawString + .replacingOccurrences(of: "%", with: "") + .trimmingCharacters(in: .whitespacesAndNewlines) + guard let parsed = Double(cleaned) else { return nil } + return max(0, min(parsed, 100)) + } + + private static func parseBool(_ raw: Any?) -> Bool? { + if let value = raw as? Bool { + return value + } + if let number = raw as? NSNumber { + return number.boolValue + } + guard let string = self.parseString(raw)?.lowercased() else { return nil } + switch string { + case "true", "1", "yes", "active", "valid": + return true + case "false", "0", "no", "inactive", "invalid", "expired": + return false + default: + return nil + } + } + + private static func matchFirstGroup(pattern: String, in text: String) -> String? { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { + return nil + } + let range = NSRange(text.startIndex.. 1, + let valueRange = Range(match.range(at: 1), in: text) + else { + return nil + } + let value = text[valueRange].trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : String(value) + } + + private static func extractConsoleSECToken(from html: String) -> String? { + let patterns = [ + #"SEC_TOKEN\s*:\s*\"([^\"]+)\""#, + #"SEC_TOKEN\s*:\s*'([^']+)'"#, + #"secToken\s*:\s*\"([^\"]+)\""#, + #"sec_token\s*:\s*\"([^\"]+)\""#, + #"sec_token\s*:\s*'([^']+)'"#, + #"\"SEC_TOKEN\"\s*:\s*\"([^\"]+)\""#, + #"\"sec_token\"\s*:\s*\"([^\"]+)\""#, + ] + + for pattern in patterns { + if let token = self.matchFirstGroup(pattern: pattern, in: html), !token.isEmpty { + return token + } + } + return nil + } + + private static func extractCookieValue(name: String, from cookieHeader: String) -> String? { + let segments = cookieHeader.split(separator: ";") + for segment in segments { + let part = segment.trimmingCharacters(in: .whitespacesAndNewlines) + guard let index = part.firstIndex(of: "=") else { continue } + let key = String(part[.. Providers -> Alibaba if the Alibaba console exposes a " + + "compatible session, or switch regions if quota API access is available there." + case let .networkError(message): + "Alibaba Coding Plan network error: \(message)" + case let .apiError(message): + "Alibaba Coding Plan API error: \(message)" + case let .parseFailed(message): + "Failed to parse Alibaba Coding Plan response: \(message)" + } + } +} + +// swiftlint:enable type_body_length diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanUsageSnapshot.swift new file mode 100644 index 0000000..2385b9b --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaCodingPlanUsageSnapshot.swift @@ -0,0 +1,126 @@ +import Foundation + +public struct AlibabaCodingPlanUsageSnapshot: Sendable { + public let planName: String? + public let fiveHourUsedQuota: Int? + public let fiveHourTotalQuota: Int? + public let fiveHourNextRefreshTime: Date? + public let weeklyUsedQuota: Int? + public let weeklyTotalQuota: Int? + public let weeklyNextRefreshTime: Date? + public let monthlyUsedQuota: Int? + public let monthlyTotalQuota: Int? + public let monthlyNextRefreshTime: Date? + public let updatedAt: Date + + public init( + planName: String?, + fiveHourUsedQuota: Int?, + fiveHourTotalQuota: Int?, + fiveHourNextRefreshTime: Date?, + weeklyUsedQuota: Int?, + weeklyTotalQuota: Int?, + weeklyNextRefreshTime: Date?, + monthlyUsedQuota: Int?, + monthlyTotalQuota: Int?, + monthlyNextRefreshTime: Date?, + updatedAt: Date) + { + self.planName = planName + self.fiveHourUsedQuota = fiveHourUsedQuota + self.fiveHourTotalQuota = fiveHourTotalQuota + self.fiveHourNextRefreshTime = fiveHourNextRefreshTime + self.weeklyUsedQuota = weeklyUsedQuota + self.weeklyTotalQuota = weeklyTotalQuota + self.weeklyNextRefreshTime = weeklyNextRefreshTime + self.monthlyUsedQuota = monthlyUsedQuota + self.monthlyTotalQuota = monthlyTotalQuota + self.monthlyNextRefreshTime = monthlyNextRefreshTime + self.updatedAt = updatedAt + } +} + +extension AlibabaCodingPlanUsageSnapshot { + public func toUsageSnapshot() -> UsageSnapshot { + let primaryPercent = Self.usedPercent(used: self.fiveHourUsedQuota, total: self.fiveHourTotalQuota) + let secondaryPercent = Self.usedPercent(used: self.weeklyUsedQuota, total: self.weeklyTotalQuota) + let tertiaryPercent = Self.usedPercent(used: self.monthlyUsedQuota, total: self.monthlyTotalQuota) + + let primary: RateWindow? = if let primaryPercent { + RateWindow( + usedPercent: primaryPercent, + windowMinutes: 5 * 60, + resetsAt: Self.normalizedFiveHourReset( + self.fiveHourNextRefreshTime, + updatedAt: self.updatedAt), + resetDescription: Self.usageDetail(used: self.fiveHourUsedQuota, total: self.fiveHourTotalQuota)) + } else { + nil + } + + let secondary: RateWindow? = if let secondaryPercent { + RateWindow( + usedPercent: secondaryPercent, + windowMinutes: 7 * 24 * 60, + resetsAt: self.weeklyNextRefreshTime, + resetDescription: Self.usageDetail(used: self.weeklyUsedQuota, total: self.weeklyTotalQuota)) + } else { + nil + } + + let tertiary: RateWindow? = if let tertiaryPercent { + RateWindow( + usedPercent: tertiaryPercent, + windowMinutes: 30 * 24 * 60, + resetsAt: self.monthlyNextRefreshTime, + resetDescription: Self.usageDetail(used: self.monthlyUsedQuota, total: self.monthlyTotalQuota)) + } else { + nil + } + + let loginMethod = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines) + let identity = ProviderIdentitySnapshot( + providerID: .alibaba, + accountEmail: nil, + accountOrganization: nil, + loginMethod: loginMethod) + + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: tertiary, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } + + private static func usedPercent(used: Int?, total: Int?) -> Double? { + guard let used, let total, total > 0 else { return nil } + let normalizedUsed = max(0, min(used, total)) + return Double(normalizedUsed) / Double(total) * 100 + } + + private static func limitDescription(total: Int?, label: String) -> String? { + guard let total, total > 0 else { return nil } + return "\(total) requests / \(label)" + } + + private static func usageDetail(used: Int?, total: Int?) -> String? { + guard let used, let total, total > 0 else { return nil } + return "\(used) / \(total) used" + } + + private static func normalizedFiveHourReset(_ raw: Date?, updatedAt: Date) -> Date? { + guard let raw else { return nil } + if raw.timeIntervalSince(updatedAt) >= 60 { + return raw + } + + let shifted = raw.addingTimeInterval(TimeInterval(5 * 60 * 60)) + if shifted.timeIntervalSince(updatedAt) >= 60 { + return shifted + } + + return updatedAt.addingTimeInterval(TimeInterval(5 * 60 * 60)) + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanAPIRegion.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanAPIRegion.swift new file mode 100644 index 0000000..ae7453f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanAPIRegion.swift @@ -0,0 +1,60 @@ +import Foundation + +public enum AlibabaTokenPlanAPIRegion: String, CaseIterable, Sendable { + case international = "intl" + case chinaMainland = "cn" + + public var displayName: String { + switch self { + case .international: + "International (modelstudio.console.alibabacloud.com)" + case .chinaMainland: + "China mainland (bailian.console.aliyun.com)" + } + } + + public var gatewayBaseURLString: String { + switch self { + case .international: + "https://modelstudio.console.alibabacloud.com" + case .chinaMainland: + "https://bailian.console.aliyun.com" + } + } + + public var dashboardOriginURLString: String { + self.gatewayBaseURLString + } + + public var dashboardURL: URL { + switch self { + case .international: + URL( + string: "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=plan#/efm/subscription/token-plan")! + case .chinaMainland: + URL(string: "https://bailian.console.aliyun.com/cn-beijing?tab=plan#/efm/subscription/token-plan")! + } + } + + public var currentRegionID: String { + switch self { + case .international: + "ap-southeast-1" + case .chinaMainland: + "cn-beijing" + } + } + + public var tokenPlanProductCode: String { + switch self { + case .international: + "sfm_tokenplanteams_dp_intl" + case .chinaMainland: + "sfm_tokenplanteams_dp_cn" + } + } + + public var cookieCacheScope: CookieHeaderCache.Scope { + .providerVariant("alibaba-token-plan.\(self.rawValue)") + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanCookieHeader.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanCookieHeader.swift new file mode 100644 index 0000000..14a99af --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanCookieHeader.swift @@ -0,0 +1,163 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +struct AlibabaTokenPlanCookieHeaders { + private static let cachedAPIHeaderName = "__codexbar_alibaba_token_plan_api" + private static let cachedDashboardHeaderName = "__codexbar_alibaba_token_plan_dashboard" + + let apiCookieHeader: String + let dashboardCookieHeader: String + + init(apiCookieHeader: String, dashboardCookieHeader: String) { + self.apiCookieHeader = apiCookieHeader + self.dashboardCookieHeader = dashboardCookieHeader + } + + init?(singleHeader raw: String?) { + guard let normalized = CookieHeaderNormalizer.normalize(raw) else { return nil } + self.apiCookieHeader = normalized + self.dashboardCookieHeader = normalized + } + + init?(cachedHeader raw: String?) { + var valuesByName: [String: String] = [:] + for pair in CookieHeaderNormalizer.pairs(from: raw ?? "") { + valuesByName[pair.name] = pair.value + } + if let encodedAPI = valuesByName[Self.cachedAPIHeaderName], + let encodedDashboard = valuesByName[Self.cachedDashboardHeaderName], + let apiHeader = Self.decodeCachedHeader(encodedAPI), + let dashboardHeader = Self.decodeCachedHeader(encodedDashboard), + let normalizedAPI = CookieHeaderNormalizer.normalize(apiHeader), + let normalizedDashboard = CookieHeaderNormalizer.normalize(dashboardHeader) + { + self.init(apiCookieHeader: normalizedAPI, dashboardCookieHeader: normalizedDashboard) + return + } + + self.init(singleHeader: raw) + } + + var cacheCookieHeader: String { + [ + "\(Self.cachedAPIHeaderName)=\(Self.encodeCachedHeader(self.apiCookieHeader))", + "\(Self.cachedDashboardHeaderName)=\(Self.encodeCachedHeader(self.dashboardCookieHeader))", + ].joined(separator: "; ") + } + + var apiCookieNames: [String] { + Self.cookieNames(from: self.apiCookieHeader) + } + + var dashboardCookieNames: [String] { + Self.cookieNames(from: self.dashboardCookieHeader) + } + + func hasCookie(named name: String) -> Bool { + Self.cookieNames(from: self.apiCookieHeader).contains(name) || + Self.cookieNames(from: self.dashboardCookieHeader).contains(name) + } + + private static func cookieNames(from header: String) -> [String] { + CookieHeaderNormalizer.pairs(from: header) + .map(\.name) + .filter { !$0.isEmpty } + .uniquedSorted() + } + + private static func encodeCachedHeader(_ header: String) -> String { + Data(header.utf8).base64EncodedString() + } + + private static func decodeCachedHeader(_ encoded: String) -> String? { + guard let data = Data(base64Encoded: encoded) else { return nil } + return String(data: data, encoding: .utf8) + } +} + +enum AlibabaTokenPlanCookieHeader { + static func headers( + from cookies: [HTTPCookie], + region: AlibabaTokenPlanAPIRegion = .international, + environment: [String: String] = ProcessInfo.processInfo.environment) -> AlibabaTokenPlanCookieHeaders? + { + guard let apiHeader = self.header( + from: cookies, + targetURL: AlibabaTokenPlanUsageFetcher.resolveQuotaURL(region: region, environment: environment)), + let dashboardHeader = self.header( + from: cookies, + targetURL: AlibabaTokenPlanUsageFetcher.dashboardURL(region: region, environment: environment)) + else { + return nil + } + return AlibabaTokenPlanCookieHeaders(apiCookieHeader: apiHeader, dashboardCookieHeader: dashboardHeader) + } + + static func header(from cookies: [HTTPCookie], targetURL: URL) -> String? { + var byName: [String: HTTPCookie] = [:] + for cookie in cookies { + guard !cookie.name.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { continue } + guard !cookie.value.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { continue } + if let expiry = cookie.expiresDate, expiry < Date() { continue } + guard self.matchesRequestURL(cookie: cookie, url: targetURL) else { continue } + + if let existing = byName[cookie.name] { + if self.cookieSortKey(for: cookie) >= self.cookieSortKey(for: existing) { + byName[cookie.name] = cookie + } + } else { + byName[cookie.name] = cookie + } + } + + guard !byName.isEmpty else { return nil } + return byName.keys.sorted().compactMap { name in + guard let cookie = byName[name] else { return nil } + return "\(cookie.name)=\(cookie.value)" + }.joined(separator: "; ") + } + + private static func matchesRequestURL(cookie: HTTPCookie, url: URL) -> Bool { + guard let host = url.host?.lowercased() else { return false } + let normalizedDomain = cookie.domain.lowercased().trimmingCharacters(in: CharacterSet(charactersIn: ".")) + guard !normalizedDomain.isEmpty else { return false } + guard host == normalizedDomain || host.hasSuffix(".\(normalizedDomain)") else { return false } + + let cookiePath = cookie.path.isEmpty ? "/" : cookie.path + let requestPath = url.path.isEmpty ? "/" : url.path + if requestPath == cookiePath { + return true + } + guard requestPath.hasPrefix(cookiePath) else { return false } + guard cookiePath != "/" else { return true } + if cookiePath.hasSuffix("/") { + return true + } + guard + let boundaryIndex = requestPath.index( + requestPath.startIndex, + offsetBy: cookiePath.count, + limitedBy: requestPath.endIndex), + boundaryIndex < requestPath.endIndex + else { + return true + } + return requestPath[boundaryIndex] == "/" + } + + private static func cookieSortKey(for cookie: HTTPCookie) -> (Int, Int, Date) { + let pathLength = cookie.path.count + let normalizedDomain = cookie.domain.lowercased().trimmingCharacters(in: CharacterSet(charactersIn: ".")) + let domainLength = normalizedDomain.count + let expiry = cookie.expiresDate ?? .distantPast + return (pathLength, domainLength, expiry) + } +} + +extension [String] { + fileprivate func uniquedSorted() -> [String] { + Array(Set(self)).sorted() + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanProviderDescriptor.swift new file mode 100644 index 0000000..6b5012a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanProviderDescriptor.swift @@ -0,0 +1,289 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit +#endif + +public enum AlibabaTokenPlanProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + #if os(macOS) + let browserOrder: BrowserCookieImportOrder = [ + .chrome, + .chromeBeta, + .brave, + .edge, + .arc, + .firefox, + .safari, + ] + #else + let browserOrder: BrowserCookieImportOrder? = nil + #endif + + return ProviderDescriptor( + id: .alibabatokenplan, + metadata: ProviderMetadata( + id: .alibabatokenplan, + displayName: "Alibaba Token Plan", + sessionLabel: "Credits", + weeklyLabel: "Usage", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Alibaba Token Plan usage", + cliName: "alibaba-token-plan", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: browserOrder, + dashboardURL: AlibabaTokenPlanUsageFetcher.dashboardURL.absoluteString, + statusPageURL: nil, + statusLinkURL: "https://status.aliyun.com"), + branding: ProviderBranding( + iconStyle: .alibaba, + iconResourceName: "ProviderIcon-alibaba", + color: ProviderColor(red: 1.0, green: 106 / 255, blue: 0)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Alibaba Token Plan cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), + cli: ProviderCLIConfig( + name: "alibaba-token-plan", + aliases: ["alibaba-token", "bailian-token-plan"], + versionDetector: nil)) + } + + private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] { + guard context.settings?.alibabaTokenPlan?.cookieSource != .off else { return [] } + switch context.sourceMode { + case .auto, .web: + return [AlibabaTokenPlanWebFetchStrategy()] + case .api, .cli, .oauth: + return [] + } + } +} + +struct AlibabaTokenPlanWebFetchStrategy: ProviderFetchStrategy { + private static let log = CodexBarLog.logger("alibaba-token-plan") + + let id: String = "alibaba-token-plan.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + guard context.settings?.alibabaTokenPlan?.cookieSource != .off else { return false } + let region = context.settings?.alibabaTokenPlan?.apiRegion ?? .international + + if AlibabaTokenPlanSettingsReader.cookieHeader(environment: context.env) != nil { + return true + } + + if let settings = context.settings?.alibabaTokenPlan, + settings.cookieSource == .manual + { + return CookieHeaderNormalizer.normalize(settings.manualCookieHeader) != nil + } + + #if os(macOS) + if let cached = Self.cachedCookieEntry(region: region), + !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + return true + } + return true + #else + return false + #endif + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let cookieSource = context.settings?.alibabaTokenPlan?.cookieSource ?? .auto + let region = context.settings?.alibabaTokenPlan?.apiRegion ?? .international + let cookieHeaders = try Self.resolveCookieHeaders(context: context, allowCached: true, region: region) + do { + let usage = try await AlibabaTokenPlanUsageFetcher.fetchUsage( + apiCookieHeader: cookieHeaders.apiCookieHeader, + dashboardCookieHeader: cookieHeaders.dashboardCookieHeader, + region: region, + environment: context.env) + return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web") + } catch let error as AlibabaTokenPlanUsageError + where error.isCredentialFailure && cookieSource != .manual + { + #if os(macOS) + CookieHeaderCache.clear(provider: .alibabatokenplan, scope: region.cookieCacheScope) + let refreshedHeaders = try Self.resolveCookieHeaders(context: context, allowCached: false, region: region) + let usage = try await AlibabaTokenPlanUsageFetcher.fetchUsage( + apiCookieHeader: refreshedHeaders.apiCookieHeader, + dashboardCookieHeader: refreshedHeaders.dashboardCookieHeader, + region: region, + environment: context.env) + return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "web") + #else + throw error + #endif + } + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + static func resolveCookieHeader(context: ProviderFetchContext, allowCached: Bool) throws -> String { + try self.resolveCookieHeaders(context: context, allowCached: allowCached, region: .international) + .apiCookieHeader + } + + static func resolveCookieHeaders( + context: ProviderFetchContext, + allowCached: Bool, + region: AlibabaTokenPlanAPIRegion = .international) throws -> AlibabaTokenPlanCookieHeaders + { + if let settings = context.settings?.alibabaTokenPlan, + settings.cookieSource == .manual + { + guard let headers = AlibabaTokenPlanCookieHeaders(singleHeader: settings.manualCookieHeader) else { + self.log.warning("Alibaba Token Plan manual cookie header is invalid") + throw AlibabaTokenPlanSettingsError.invalidCookie + } + Self.log.info( + "Alibaba Token Plan using manual cookie header", + metadata: [ + "apiCookieNames": headers.apiCookieNames.joined(separator: ","), + "dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","), + "hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0", + ]) + return headers + } + + if let envCookie = AlibabaTokenPlanSettingsReader.cookieHeader(environment: context.env), + let headers = AlibabaTokenPlanCookieHeaders(singleHeader: envCookie) + { + Self.log.info( + "Alibaba Token Plan using environment cookie header", + metadata: [ + "apiCookieNames": headers.apiCookieNames.joined(separator: ","), + "dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","), + "hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0", + ]) + return headers + } + + #if os(macOS) + if allowCached, + let cached = Self.cachedCookieEntry(region: region), + let headers = AlibabaTokenPlanCookieHeaders(cachedHeader: cached.cookieHeader) + { + Self.log.info( + "Alibaba Token Plan using cached browser cookie header", + metadata: [ + "source": cached.sourceLabel, + "apiCookieNames": headers.apiCookieNames.joined(separator: ","), + "dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","), + "hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0", + ]) + return headers + } + + do { + var importLog: [String] = [] + let session = try AlibabaCodingPlanCookieImporter.importSession( + browserDetection: context.browserDetection, + logger: { importLog.append($0) }) + let rawCookieNames = session.cookies.map(\.name).filter { !$0.isEmpty }.uniquedSorted() + guard let headers = AlibabaTokenPlanCookieHeader.headers( + from: session.cookies, + region: region, + environment: context.env) + else { + Self.log.warning( + "Alibaba Token Plan browser cookie header was empty", + metadata: [ + "source": session.sourceLabel, + "rawCookieNames": rawCookieNames.joined(separator: ","), + ]) + throw AlibabaTokenPlanSettingsError.missingCookie( + details: "No Alibaba Token Plan browser cookies were available after import.") + } + CookieHeaderCache.store( + provider: .alibabatokenplan, + scope: region.cookieCacheScope, + cookieHeader: headers.cacheCookieHeader, + sourceLabel: session.sourceLabel) + Self.log.info( + "Alibaba Token Plan imported browser cookies", + metadata: [ + "source": session.sourceLabel, + "rawCookieNames": rawCookieNames.joined(separator: ","), + "apiCookieNames": headers.apiCookieNames.joined(separator: ","), + "dashboardCookieNames": headers.dashboardCookieNames.joined(separator: ","), + "hasSecToken": headers.hasCookie(named: "sec_token") ? "1" : "0", + "importLogLines": "\(importLog.count)", + ]) + return headers + } catch { + Self.log.warning( + "Alibaba Token Plan cookie resolution failed", + metadata: ["error": error.localizedDescription]) + throw AlibabaTokenPlanSettingsError.missingCookie(details: Self.missingCookieDetails(from: error)) + } + #else + throw AlibabaTokenPlanSettingsError.missingCookie() + #endif + } + + #if os(macOS) + /// The former unscoped cache only ever represented the China gateway. Never expose it to + /// International requests; migrate it into the China scope after a successful scoped write. + private static func cachedCookieEntry(region: AlibabaTokenPlanAPIRegion) -> CookieHeaderCache.Entry? { + if let scoped = CookieHeaderCache.load(provider: .alibabatokenplan, scope: region.cookieCacheScope) { + return scoped + } + guard region == .chinaMainland, + let legacy = CookieHeaderCache.load(provider: .alibabatokenplan) + else { return nil } + + CookieHeaderCache.store( + provider: .alibabatokenplan, + scope: region.cookieCacheScope, + cookieHeader: legacy.cookieHeader, + sourceLabel: legacy.sourceLabel, + now: legacy.storedAt) + if let migrated = CookieHeaderCache.load(provider: .alibabatokenplan, scope: region.cookieCacheScope) { + CookieHeaderCache.clear(provider: .alibabatokenplan) + return migrated + } + return legacy + } + #endif + + private static func missingCookieDetails(from error: Error) -> String? { + if case let AlibabaCodingPlanSettingsError.missingCookie(details) = error { + return details + } + let message = error.localizedDescription.trimmingCharacters(in: .whitespacesAndNewlines) + return message.isEmpty ? nil : message + } +} + +extension [String] { + fileprivate func uniquedSorted() -> [String] { + Array(Set(self)).sorted() + } +} + +extension AlibabaTokenPlanUsageError { + fileprivate var isCredentialFailure: Bool { + switch self { + case .loginRequired, .invalidCredentials: + true + case .apiError, .networkError, .parseFailed: + false + } + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanSettingsReader.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanSettingsReader.swift new file mode 100644 index 0000000..1074d6c --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanSettingsReader.swift @@ -0,0 +1,67 @@ +import Foundation + +public struct AlibabaTokenPlanSettingsReader: Sendable { + public static let cookieHeaderKey = "ALIBABA_TOKEN_PLAN_COOKIE" + public static let hostKey = "ALIBABA_TOKEN_PLAN_HOST" + public static let quotaURLKey = "ALIBABA_TOKEN_PLAN_QUOTA_URL" + + private static let endpointValidator = ProviderEndpointOverrideValidator( + allowedHosts: [ + "modelstudio.console.alibabacloud.com", + "bailian.console.aliyun.com", + ]) + + public static func cookieHeader( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.cookieHeaderKey]) + } + + public static func hostOverride( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.endpointValidator.validatedHost( + self.cleaned(environment[self.hostKey]), + policy: .allowAnyHTTPSHost) + } + + public static func quotaURL( + environment: [String: String] = ProcessInfo.processInfo.environment) -> URL? + { + self.endpointValidator.validatedURL( + self.cleaned(environment[self.quotaURLKey]), + policy: .allowAnyHTTPSHost) + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} + +public enum AlibabaTokenPlanSettingsError: LocalizedError, Sendable { + case missingCookie(details: String? = nil) + case invalidCookie + + public var errorDescription: String? { + switch self { + case let .missingCookie(details): + let base = "No Alibaba Token Plan session cookies found in browsers. " + + "Sign in to Model Studio/Bailian in Chrome, " + + "allow CodexBar to access Chrome Safe Storage in Keychain Access, " + + "or paste a manual Cookie header." + guard let details, !details.isEmpty else { return base } + return "\(base) \(details)" + case .invalidCookie: + return "Alibaba Token Plan cookie header is invalid." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanUsageFetcher.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanUsageFetcher.swift new file mode 100644 index 0000000..27aee5a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanUsageFetcher.swift @@ -0,0 +1,997 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum AlibabaTokenPlanUsageError: LocalizedError, Sendable, Equatable { + case loginRequired + case invalidCredentials + case apiError(String) + case networkError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .loginRequired: + "Alibaba Token Plan login required." + case .invalidCredentials: + "Alibaba Token Plan credentials are invalid." + case let .apiError(message): + "Alibaba Token Plan API error: \(message)" + case let .networkError(message): + "Alibaba Token Plan network error: \(message)" + case let .parseFailed(message): + "Could not parse Alibaba Token Plan usage: \(message)" + } + } +} + +// swiftlint:disable:next type_body_length +public struct AlibabaTokenPlanUsageFetcher: Sendable { + private static let log = CodexBarLog.logger("alibaba-token-plan") + private static let bssServiceCode = "BssOpenAPI-V3" + private static let subscriptionSummaryAction = "GetSubscriptionSummary" + private static let browserLikeUserAgent = + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36" + private static let safariLikeUserAgent = + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/605.1.15 (KHTML, like Gecko) Version/26.3 Safari/605.1.15" + + public static var dashboardURL: URL { + Self.dashboardURL(region: .international, environment: ProcessInfo.processInfo.environment) + } + + public static func dashboardURL( + region: AlibabaTokenPlanAPIRegion, + environment: [String: String] = ProcessInfo.processInfo.environment) -> URL + { + if let host = AlibabaTokenPlanSettingsReader.hostOverride(environment: environment), + let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: host), + var components = URLComponents(url: base, resolvingAgainstBaseURL: false), + let dashboardComponents = URLComponents(url: region.dashboardURL, resolvingAgainstBaseURL: false) + { + components.path = dashboardComponents.path + components.percentEncodedQuery = dashboardComponents.percentEncodedQuery + components.fragment = dashboardComponents.fragment + return components.url ?? region.dashboardURL + } + return region.dashboardURL + } + + public static func fetchUsage( + cookieHeader: String, + region: AlibabaTokenPlanAPIRegion = .international, + environment: [String: String] = ProcessInfo.processInfo.environment, + now: Date = Date()) async throws -> AlibabaTokenPlanUsageSnapshot + { + guard let headers = AlibabaTokenPlanCookieHeaders(singleHeader: cookieHeader) else { + throw AlibabaTokenPlanSettingsError.invalidCookie + } + return try await self.fetchUsage( + apiCookieHeader: headers.apiCookieHeader, + dashboardCookieHeader: headers.dashboardCookieHeader, + region: region, + environment: environment, + now: now) + } + + static func fetchUsage( + apiCookieHeader: String, + dashboardCookieHeader: String, + region: AlibabaTokenPlanAPIRegion = .international, + environment: [String: String] = ProcessInfo.processInfo.environment, + now: Date = Date(), + session overrideSession: URLSession? = nil) async throws -> AlibabaTokenPlanUsageSnapshot + { + guard let normalizedAPIHeader = CookieHeaderNormalizer.normalize(apiCookieHeader), + let normalizedDashboardHeader = CookieHeaderNormalizer.normalize(dashboardCookieHeader) + else { + throw AlibabaTokenPlanSettingsError.invalidCookie + } + + let url = self.resolveQuotaURL(region: region, environment: environment) + let apiRedirectDiagnostics = RedirectDiagnostics(cookieHeader: normalizedAPIHeader) + let dashboardRedirectDiagnostics: RedirectDiagnostics? + let apiSession: URLSession + let dashboardSession: URLSession + if let overrideSession { + apiSession = overrideSession + dashboardSession = overrideSession + dashboardRedirectDiagnostics = nil + } else { + let dashboardDiagnostics = RedirectDiagnostics(cookieHeader: normalizedDashboardHeader) + apiSession = URLSession( + configuration: .default, + delegate: apiRedirectDiagnostics, + delegateQueue: nil) + dashboardSession = URLSession( + configuration: .default, + delegate: dashboardDiagnostics, + delegateQueue: nil) + dashboardRedirectDiagnostics = dashboardDiagnostics + } + defer { + if overrideSession == nil { + apiSession.invalidateAndCancel() + dashboardSession.invalidateAndCancel() + } + } + let secToken = await self.resolveSECToken( + dashboardCookieHeader: normalizedDashboardHeader, + apiCookieHeader: normalizedAPIHeader, + region: region, + environment: environment, + session: dashboardSession) + Self.log.info( + "Fetching Alibaba Token Plan usage", + metadata: [ + "apiHost": url.host ?? "unknown", + "region": region.rawValue, + "apiCookieNames": self.cookieNamesDescription(from: normalizedAPIHeader), + "dashboardCookieNames": self.cookieNamesDescription(from: normalizedDashboardHeader), + "hasCSRF": self.hasCSRF(in: normalizedAPIHeader) ? "1" : "0", + "secTokenSource": secToken == nil ? "missing" : "resolved", + ]) + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.timeoutInterval = 20 + request.httpBody = self.subscriptionSummaryRequestBody(region: region, secToken: secToken) + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.setValue("*/*", forHTTPHeaderField: "Accept") + request.setValue(normalizedAPIHeader, forHTTPHeaderField: "Cookie") + if let csrf = self.extractCookieValue(name: "login_aliyunid_csrf", from: normalizedAPIHeader) ?? + self.extractCookieValue(name: "csrf", from: normalizedAPIHeader) + { + request.setValue(csrf, forHTTPHeaderField: "x-xsrf-token") + request.setValue(csrf, forHTTPHeaderField: "x-csrf-token") + } + request.setValue("XMLHttpRequest", forHTTPHeaderField: "X-Requested-With") + request.setValue(Self.browserLikeUserAgent, forHTTPHeaderField: "User-Agent") + request.setValue(region.dashboardOriginURLString, forHTTPHeaderField: "Origin") + request.setValue( + Self.dashboardURL(region: region, environment: environment).absoluteString, + forHTTPHeaderField: "Referer") + + let data: Data + let response: URLResponse + do { + (data, response) = try await apiSession.data(for: request) + } catch { + Self.log.error( + "Alibaba Token Plan request failed", + metadata: [ + "apiHost": url.host ?? "unknown", + "error": error.localizedDescription, + ]) + throw AlibabaTokenPlanUsageError.networkError(error.localizedDescription) + } + if let dashboardRedirectDiagnostics, !dashboardRedirectDiagnostics.redirects.isEmpty { + Self.log.info( + "Alibaba Token Plan dashboard redirects", + metadata: [ + "count": "\(dashboardRedirectDiagnostics.redirects.count)", + "items": dashboardRedirectDiagnostics.redirects.joined(separator: " | "), + ]) + } + if !apiRedirectDiagnostics.redirects.isEmpty { + Self.log.info( + "Alibaba Token Plan redirects", + metadata: [ + "count": "\(apiRedirectDiagnostics.redirects.count)", + "items": apiRedirectDiagnostics.redirects.joined(separator: " | "), + ]) + } + guard let httpResponse = response as? HTTPURLResponse else { + Self.log.error("Alibaba Token Plan response was not HTTP") + throw AlibabaTokenPlanUsageError.networkError("Invalid response") + } + Self.log.info( + "Alibaba Token Plan HTTP response", + metadata: [ + "status": "\(httpResponse.statusCode)", + "bodyBytes": "\(data.count)", + "contentType": httpResponse.value(forHTTPHeaderField: "Content-Type") ?? "none", + ]) + guard httpResponse.statusCode == 200 else { + if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 { + throw AlibabaTokenPlanUsageError.loginRequired + } + Self.log.error("Alibaba Token Plan returned HTTP \(httpResponse.statusCode)") + throw AlibabaTokenPlanUsageError.apiError("HTTP \(httpResponse.statusCode)") + } + + return try self.parseUsageSnapshot(from: data, now: now) + } + + static func resolveQuotaURL( + region: AlibabaTokenPlanAPIRegion, + environment: [String: String]) -> URL + { + if let override = AlibabaTokenPlanSettingsReader.quotaURL(environment: environment) { + return override + } + if let host = AlibabaTokenPlanSettingsReader.hostOverride(environment: environment), + let hostURL = self.quotaURL(from: host) + { + return hostURL + } + return self.defaultQuotaURL(region: region) + } + + static var defaultQuotaURL: URL { + self.defaultQuotaURL(region: .international) + } + + static func defaultQuotaURL(region: AlibabaTokenPlanAPIRegion) -> URL { + var components = URLComponents(string: region.gatewayBaseURLString)! + components.path = "/data/api.json" + components.queryItems = [ + URLQueryItem(name: "action", value: Self.subscriptionSummaryAction), + URLQueryItem(name: "product", value: Self.bssServiceCode), + URLQueryItem(name: "_tag", value: ""), + ] + return components.url! + } + + static func parseUsageSnapshot(from data: Data, now: Date = Date()) throws -> AlibabaTokenPlanUsageSnapshot { + guard !data.isEmpty else { + throw AlibabaTokenPlanUsageError.parseFailed("Empty response body") + } + + let object: Any + do { + object = try JSONSerialization.jsonObject(with: data, options: []) + } catch { + if self.isLikelyLoginHTML(data) { + throw AlibabaTokenPlanUsageError.loginRequired + } + throw AlibabaTokenPlanUsageError.parseFailed("Invalid JSON response") + } + let expanded = self.expandedJSON(object) + guard let dictionary = expanded as? [String: Any] else { + throw AlibabaTokenPlanUsageError.parseFailed("Unexpected payload") + } + + try self.throwIfErrorPayload(dictionary) + + let summary = self.findSubscriptionSummary(in: dictionary) ?? dictionary + let total = self.anyDouble(for: Self.totalQuotaKeys, in: summary) + let remaining = self.anyDouble(for: Self.remainingQuotaKeys, in: summary) + let used = self.anyDouble(for: Self.usedQuotaKeys, in: summary) ?? + total.flatMap { total in remaining.map { max(0, total - $0) } } + let resetsAt = self.findResetDate(in: summary) ?? self.findResetDate(in: dictionary) + let totalCount = self.anyDouble(for: Self.subscriptionCountKeys, in: summary) + let planName = self.findPlanName(in: summary) ?? ((totalCount ?? 0) > 0 || total != nil ? "TOKEN PLAN" : nil) + + if planName == nil, total == nil, used == nil, remaining == nil, totalCount == nil { + let diagnostics = self.payloadDiagnostics(payload: dictionary) + Self.log.error("Alibaba Token Plan payload missing expected fields: \(diagnostics)") + throw AlibabaTokenPlanUsageError.parseFailed("Missing token plan data (\(diagnostics))") + } + + return AlibabaTokenPlanUsageSnapshot( + planName: planName, + usedQuota: used, + totalQuota: total, + remainingQuota: remaining, + resetsAt: resetsAt, + updatedAt: now) + } + + private static func subscriptionSummaryRequestBody(region: AlibabaTokenPlanAPIRegion, secToken: String?) -> Data { + let paramsObject = ["ProductCode": region.tokenPlanProductCode] + guard let paramsData = try? JSONSerialization.data(withJSONObject: paramsObject, options: []), + let paramsString = String(data: paramsData, encoding: .utf8) + else { + return Data() + } + + var components = URLComponents() + var queryItems = [ + URLQueryItem(name: "product", value: Self.bssServiceCode), + URLQueryItem(name: "action", value: Self.subscriptionSummaryAction), + URLQueryItem(name: "params", value: paramsString), + URLQueryItem(name: "region", value: region.currentRegionID), + ] + if let secToken, !secToken.isEmpty { + queryItems.append(URLQueryItem(name: "sec_token", value: secToken)) + } + components.queryItems = queryItems + return Data((components.percentEncodedQuery ?? "").utf8) + } + + private static func resolveSECToken( + dashboardCookieHeader: String, + apiCookieHeader: String, + region: AlibabaTokenPlanAPIRegion, + environment: [String: String], + session: URLSession) async -> String? + { + let cookieSECToken = self.extractCookieValue(name: "sec_token", from: dashboardCookieHeader) ?? + self.extractCookieValue(name: "sec_token", from: apiCookieHeader) + var request = URLRequest(url: self.dashboardURL(region: region, environment: environment)) + request.httpMethod = "GET" + request.timeoutInterval = 10 + request.setValue(dashboardCookieHeader, forHTTPHeaderField: "Cookie") + request.setValue(Self.safariLikeUserAgent, forHTTPHeaderField: "User-Agent") + request.setValue( + "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", + forHTTPHeaderField: "Accept") + + if let (data, response) = try? await session.data(for: request), + let httpResponse = response as? HTTPURLResponse, + httpResponse.statusCode == 200, + let html = String(data: data, encoding: .utf8), + let token = self.extractSECToken(from: html) + { + Self.log.info( + "Resolved Alibaba Token Plan sec_token from dashboard HTML", + metadata: [ + "dashboardHost": request.url?.host ?? "unknown", + "htmlBytes": "\(data.count)", + ]) + return token + } + + if let token = await self.fetchSECTokenFromUserInfo( + cookieHeader: dashboardCookieHeader, + region: region, + environment: environment, + session: session) + { + return token + } + + if let cookieSECToken, !cookieSECToken.isEmpty { + Self.log.info("Resolved Alibaba Token Plan sec_token from cookies") + return cookieSECToken + } + + Self.log.info( + "Alibaba Token Plan sec_token missing; continuing with cookie-only request", + metadata: [ + "dashboardCookieNames": self.cookieNamesDescription(from: dashboardCookieHeader), + "apiCookieNames": self.cookieNamesDescription(from: apiCookieHeader), + ]) + return nil + } + + private static func fetchSECTokenFromUserInfo( + cookieHeader: String, + region: AlibabaTokenPlanAPIRegion, + environment: [String: String], + session: URLSession) async -> String? + { + let baseURL = self.consoleBaseURL(region: region, environment: environment) + let userInfoURL = baseURL.appendingPathComponent("tool/user/info.json") + var request = URLRequest(url: userInfoURL) + request.httpMethod = "GET" + request.timeoutInterval = 10 + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue(Self.safariLikeUserAgent, forHTTPHeaderField: "User-Agent") + request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept") + let referer = baseURL.absoluteString.hasSuffix("/") ? baseURL.absoluteString : "\(baseURL.absoluteString)/" + request.setValue(referer, forHTTPHeaderField: "Referer") + + guard let (data, response) = try? await session.data(for: request), + let httpResponse = response as? HTTPURLResponse, + httpResponse.statusCode == 200, + let object = try? JSONSerialization.jsonObject(with: data, options: []) + else { + return nil + } + + let expanded = self.expandedJSON(object) + guard let token = self.findFirstString(forKeys: ["secToken", "sec_token"], in: expanded), + !token.isEmpty + else { + return nil + } + + Self.log.info( + "Resolved Alibaba Token Plan sec_token from user info", + metadata: [ + "userInfoHost": userInfoURL.host ?? "unknown", + "bodyBytes": "\(data.count)", + ]) + return token + } + + private static func consoleBaseURL( + region: AlibabaTokenPlanAPIRegion, + environment: [String: String]) -> URL + { + let dashboard = self.dashboardURL(region: region, environment: environment) + var components = URLComponents() + components.scheme = dashboard.scheme + components.host = dashboard.host + components.port = dashboard.port + return components.url ?? URL(string: region.dashboardOriginURLString)! + } + + private static func quotaURL(from rawHost: String) -> URL? { + let cleaned = AlibabaTokenPlanSettingsReader.cleaned(rawHost) + guard let cleaned else { return nil } + guard let base = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: cleaned) else { return nil } + var components = URLComponents(url: base, resolvingAgainstBaseURL: false) + let defaultComponents = URLComponents( + url: Self.defaultQuotaURL(region: .international), + resolvingAgainstBaseURL: false) + components?.path = "/data/api.json" + components?.queryItems = defaultComponents?.queryItems + return components?.url + } + + private final class RedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable { + private let cookieHeader: String + var redirects: [String] = [] + + init(cookieHeader: String) { + self.cookieHeader = cookieHeader + } + + func urlSession( + _: URLSession, + task _: URLSessionTask, + willPerformHTTPRedirection response: HTTPURLResponse, + newRequest request: URLRequest, + completionHandler: @escaping (URLRequest?) -> Void) + { + let from = AlibabaTokenPlanUsageFetcher.redactedURLDescription(response.url) + let to = AlibabaTokenPlanUsageFetcher.redactedURLDescription(request.url) + self.redirects.append("\(response.statusCode) \(from) -> \(to)") + + completionHandler(AlibabaTokenPlanUsageFetcher.redirectedRequest( + response: response, + request: request, + cookieHeader: self.cookieHeader)) + } + } + + private static func redactedURLDescription(_ url: URL?) -> String { + guard let url else { return "unknown" } + var components = URLComponents(url: url, resolvingAgainstBaseURL: false) + components?.query = nil + components?.fragment = nil + return components?.string ?? "\(url.scheme ?? "unknown")://\(url.host ?? "unknown")" + } + + static func redirectedRequest( + response: HTTPURLResponse, + request: URLRequest, + cookieHeader: String) -> URLRequest? + { + guard request.url?.scheme?.lowercased() == "https" else { + return nil + } + + var updated = request + if self.shouldForwardRedirectCookies(from: response.url, to: request.url) { + updated.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + } else { + updated.setValue(nil, forHTTPHeaderField: "Cookie") + } + return updated + } + + private static func shouldForwardRedirectCookies(from sourceURL: URL?, to targetURL: URL?) -> Bool { + guard let sourceHost = sourceURL?.host?.lowercased(), + let targetHost = targetURL?.host?.lowercased() + else { + return false + } + return sourceHost == targetHost + } + + private static func throwIfErrorPayload(_ dictionary: [String: Any]) throws { + if self.parseBool(dictionary["successResponse"]) == false { + if let statusCode = self.findFirstInt(forKeys: ["statusCode", "status_code", "code"], in: dictionary), + statusCode == 401 || statusCode == 403 + { + throw AlibabaTokenPlanUsageError.invalidCredentials + } + let code = self.findFirstString(forKeys: ["code", "status", "statusCode"], in: dictionary) + let message = self.findFirstString(forKeys: ["message", "msg", "statusMessage"], in: dictionary) ?? + code ?? + "request was not successful" + if self.isLoginOrTokenError(code: code, message: message) { + throw AlibabaTokenPlanUsageError.loginRequired + } + throw AlibabaTokenPlanUsageError.apiError(message) + } + + if self.findBoolValues(forKeys: ["Success", "success"], in: dictionary).contains(false) { + let code = self.findFirstString(forKeys: ["Code", "code"], in: dictionary) + let message = self.findFirstString(forKeys: ["Message", "message", "msg", "Code", "code"], in: dictionary) + ?? "request was not successful" + if self.isLoginOrTokenError(code: code, message: message) { + throw AlibabaTokenPlanUsageError.loginRequired + } + throw AlibabaTokenPlanUsageError.apiError(message) + } + + if let statusCode = self.findFirstInt(forKeys: ["statusCode", "status_code", "code"], in: dictionary), + statusCode != 0, + statusCode != 200 + { + let message = self.findFirstString( + forKeys: ["statusMessage", "status_msg", "message", "msg"], + in: dictionary) + ?? "status code \(statusCode)" + if statusCode == 401 || statusCode == 403 { + throw AlibabaTokenPlanUsageError.invalidCredentials + } + throw AlibabaTokenPlanUsageError.apiError(message) + } + + let codeText = self.findFirstString(forKeys: ["code", "status", "statusCode"], in: dictionary)?.lowercased() + let messageText = self.findFirstString(forKeys: ["message", "msg", "statusMessage"], in: dictionary)? + .lowercased() + if self.isLoginOrTokenError(code: codeText, message: messageText) { + throw AlibabaTokenPlanUsageError.loginRequired + } + } + + private static func isLoginOrTokenError(code: String?, message: String?) -> Bool { + let combined = [code, message] + .compactMap { $0?.lowercased() } + .joined(separator: " ") + return combined.contains("needlogin") || + combined.contains("login") || + combined.contains("postonlyortokenerror") || + combined.contains("tokenerror") || + combined.contains("request has expired") || + combined.contains("refresh page") || + combined.contains("请求已经过期") + } + + private static let planNameKeys = [ + "planName", + "plan_name", + "packageName", + "package_name", + "commodityName", + "commodity_name", + "instanceName", + "instance_name", + "displayName", + "display_name", + "ProductName", + "productName", + "name", + "title", + "planType", + "plan_type", + ] + private static let usedQuotaKeys = [ + "usedQuota", + "used_quota", + "usedCredits", + "usedCredit", + "consumedCredits", + "usage", + "used", + "usedAmount", + "consumeAmount", + "usedValue", + "UsedValue", + "consumedValue", + "ConsumedValue", + ] + private static let totalQuotaKeys = [ + "totalQuota", + "total_quota", + "totalCredits", + "totalCredit", + "quota", + "creditLimit", + "creditsTotal", + "monthlyTotalQuota", + "amount", + "totalValue", + "TotalValue", + ] + private static let remainingQuotaKeys = [ + "remainingQuota", + "remainQuota", + "remainingCredits", + "remainingCredit", + "availableCredits", + "balance", + "remaining", + "availableAmount", + "remainAmount", + "totalSurplusValue", + "TotalSurplusValue", + "surplusValue", + "SurplusValue", + ] + private static let subscriptionCountKeys = [ + "totalCount", + "TotalCount", + "subscriptionTotalNumber", + "SubscriptionTotalNumber", + ] + private static let resetDateKeys = [ + "nextRefreshTime", + "resetTime", + "periodEndTime", + "billingCycleEnd", + "billCycleEndTime", + "expireTime", + "expirationTime", + "endTime", + "validEndTime", + "instanceEndTime", + "nearestExpireDate", + "NearestExpireDate", + ] + + private static func findSubscriptionSummary(in payload: [String: Any]) -> [String: Any]? { + if let data = self.findFirstDictionary( + forKeys: ["Data", "data", "successResponse", "success_response"], + in: payload), + self.containsSubscriptionSummaryFields(data) + { + return data + } + return self.findFirstDictionary( + matchingAnyKey: Self.usedQuotaKeys + Self.totalQuotaKeys + Self.remainingQuotaKeys + + Self.subscriptionCountKeys, + in: payload) + } + + private static func containsSubscriptionSummaryFields(_ payload: [String: Any]) -> Bool { + let keys = self.usedQuotaKeys + self.totalQuotaKeys + self.remainingQuotaKeys + self.subscriptionCountKeys + return keys.contains { payload[$0] != nil } + } + + private static func findPlanName(in payload: [String: Any]) -> String? { + self.anyString(for: self.planNameKeys, in: payload) ?? + self.findFirstString(forKeys: self.planNameKeys, in: payload) + } + + private static func findResetDate(in payload: [String: Any]) -> Date? { + self.anyDate(for: self.resetDateKeys, in: payload) ?? + self.findFirstDate(forKeys: self.resetDateKeys, in: payload) + } + + private static func payloadDiagnostics(payload: [String: Any]) -> String { + let topKeys = payload.keys.sorted() + let dataDict = self.findFirstDictionary( + forKeys: ["Data", "data", "successResponse", "success_response"], + in: payload) + let dataKeys = dataDict?.keys.sorted() ?? [] + return "topKeys=\(topKeys.joined(separator: ",")) dataKeys=\(dataKeys.joined(separator: ","))" + } + + private static func isLikelyLoginHTML(_ data: Data) -> Bool { + guard let text = String(data: data, encoding: .utf8)?.lowercased() else { return false } + return text.contains(" [String: Any]? { + if let dict = value as? [String: Any] { + for key in keys { + if let nested = dict[key] as? [String: Any] { + return nested + } + } + for nestedValue in dict.values { + if let nested = self.findFirstDictionary(forKeys: keys, in: nestedValue) { + return nested + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let nested = self.findFirstDictionary(forKeys: keys, in: item) { + return nested + } + } + } + return nil + } + + private static func findFirstDictionary(matchingAnyKey keys: [String], in value: Any) -> [String: Any]? { + if let dict = value as? [String: Any] { + if keys.contains(where: { dict[$0] != nil }) { + return dict + } + for nestedValue in dict.values { + if let nested = self.findFirstDictionary(matchingAnyKey: keys, in: nestedValue) { + return nested + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let nested = self.findFirstDictionary(matchingAnyKey: keys, in: item) { + return nested + } + } + } + return nil + } + + private static func findFirstString(forKeys keys: [String], in value: Any) -> String? { + if let dict = value as? [String: Any] { + for key in keys { + if let parsed = self.parseString(dict[key]) { + return parsed + } + } + for nestedValue in dict.values { + if let parsed = self.findFirstString(forKeys: keys, in: nestedValue) { + return parsed + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let parsed = self.findFirstString(forKeys: keys, in: item) { + return parsed + } + } + } + return nil + } + + private static func findBoolValues(forKeys keys: [String], in value: Any) -> [Bool] { + if let dict = value as? [String: Any] { + let directValues = keys.compactMap { self.parseBool(dict[$0]) } + let nestedValues = dict.values.flatMap { self.findBoolValues(forKeys: keys, in: $0) } + return directValues + nestedValues + } + if let array = value as? [Any] { + return array.flatMap { self.findBoolValues(forKeys: keys, in: $0) } + } + return [] + } + + private static func findFirstInt(forKeys keys: [String], in value: Any) -> Int? { + if let dict = value as? [String: Any] { + for key in keys { + if let parsed = self.parseInt(dict[key]) { + return parsed + } + } + for nestedValue in dict.values { + if let parsed = self.findFirstInt(forKeys: keys, in: nestedValue) { + return parsed + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let parsed = self.findFirstInt(forKeys: keys, in: item) { + return parsed + } + } + } + return nil + } + + private static func findFirstDate(forKeys keys: [String], in value: Any) -> Date? { + if let dict = value as? [String: Any] { + for key in keys { + if let parsed = self.parseDate(dict[key]) { + return parsed + } + } + for nestedValue in dict.values { + if let parsed = self.findFirstDate(forKeys: keys, in: nestedValue) { + return parsed + } + } + return nil + } + if let array = value as? [Any] { + for item in array { + if let parsed = self.findFirstDate(forKeys: keys, in: item) { + return parsed + } + } + } + return nil + } + + private static func expandedJSON(_ value: Any) -> Any { + if let dict = value as? [String: Any] { + var expanded: [String: Any] = [:] + expanded.reserveCapacity(dict.count) + for (key, nested) in dict { + expanded[key] = self.expandedJSON(nested) + } + return expanded + } + if let array = value as? [Any] { + return array.map { self.expandedJSON($0) } + } + if let string = value as? String, + let data = string.data(using: .utf8), + let nested = try? JSONSerialization.jsonObject(with: data, options: []), + nested is [String: Any] || nested is [Any] + { + return self.expandedJSON(nested) + } + return value + } + + private static func anyString(for keys: [String], in dict: [String: Any]) -> String? { + for key in keys { + if let value = self.parseString(dict[key]) { + return value + } + } + return nil + } + + private static func anyDouble(for keys: [String], in dict: [String: Any]) -> Double? { + for key in keys { + if let value = self.parseDouble(dict[key]) { + return value + } + } + return nil + } + + private static func anyDate(for keys: [String], in dict: [String: Any]) -> Date? { + for key in keys { + if let value = self.parseDate(dict[key]) { + return value + } + } + return nil + } + + private static func anyBool(for keys: [String], in dict: [String: Any]) -> Bool? { + for key in keys { + if let value = self.parseBool(dict[key]) { + return value + } + } + return nil + } + + private static func parseInt(_ raw: Any?) -> Int? { + if let value = raw as? Int { return value } + if let value = raw as? Int64 { return Int(value) } + if let value = raw as? Double { return Int(value) } + if let value = raw as? NSNumber { return value.intValue } + if let value = self.parseString(raw) { return Int(value) } + return nil + } + + private static func parseDouble(_ raw: Any?) -> Double? { + if let value = raw as? Double { return value } + if let value = raw as? Int { return Double(value) } + if let value = raw as? Int64 { return Double(value) } + if let value = raw as? NSNumber { return value.doubleValue } + if let value = self.parseString(raw) { + let cleaned = value.replacingOccurrences(of: ",", with: "") + return Double(cleaned) + } + return nil + } + + private static func parseString(_ raw: Any?) -> String? { + guard let value = raw as? String else { return nil } + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + } + + private static func parseDate(_ raw: Any?) -> Date? { + if let intValue = self.parseInt(raw) { + if intValue > 1_000_000_000_000 { + return Date(timeIntervalSince1970: TimeInterval(intValue) / 1000) + } + if intValue > 1_000_000_000 { + return Date(timeIntervalSince1970: TimeInterval(intValue)) + } + } + if let string = self.parseString(raw) { + let formatter = ISO8601DateFormatter() + if let date = formatter.date(from: string) { + return date + } + let dateFormatter = DateFormatter() + dateFormatter.locale = Locale(identifier: "en_US_POSIX") + for format in ["yyyy-MM-dd", "yyyy-MM-dd HH:mm", "yyyy-MM-dd HH:mm:ss"] { + dateFormatter.dateFormat = format + if let date = dateFormatter.date(from: string) { + return date + } + } + } + return nil + } + + private static func parseBool(_ raw: Any?) -> Bool? { + if let value = raw as? Bool { return value } + if let number = raw as? NSNumber { return number.boolValue } + guard let string = self.parseString(raw)?.lowercased() else { return nil } + switch string { + case "true", "1", "yes", "active", "valid", "normal": + return true + case "false", "0", "no", "inactive", "invalid", "expired": + return false + default: + return nil + } + } + + private static func extractCookieValue(name: String, from cookieHeader: String) -> String? { + cookieHeader + .split(separator: ";") + .compactMap { part -> (String, String)? in + let pieces = part.split(separator: "=", maxSplits: 1).map { + $0.trimmingCharacters(in: .whitespacesAndNewlines) + } + guard pieces.count == 2 else { return nil } + return (pieces[0], pieces[1]) + } + .first { $0.0 == name }? + .1 + } + + private static func hasCSRF(in cookieHeader: String) -> Bool { + self.extractCookieValue(name: "login_aliyunid_csrf", from: cookieHeader) != nil || + self.extractCookieValue(name: "csrf", from: cookieHeader) != nil + } + + static func cookieNames(from cookieHeader: String) -> [String] { + CookieHeaderNormalizer.pairs(from: cookieHeader) + .map(\.name) + .filter { !$0.isEmpty } + .uniquedSorted() + } + + static func cookieNamesDescription(from cookieHeader: String) -> String { + let names = self.cookieNames(from: cookieHeader) + return names.isEmpty ? "none" : names.joined(separator: ",") + } + + private static func extractSECToken(from html: String) -> String? { + let patterns = [ + #""secToken"\s*:\s*"([^"]+)""#, + #""sec_token"\s*:\s*"([^"]+)""#, + #"secToken['"]?\s*[:=]\s*['"]([^'"]+)['"]"#, + #"sec_token['"]?\s*[:=]\s*['"]([^'"]+)['"]"#, + ] + for pattern in patterns { + if let token = self.matchFirstGroup(pattern: pattern, in: html), !token.isEmpty { + return token + } + } + return nil + } + + private static func matchFirstGroup(pattern: String, in text: String) -> String? { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { + return nil + } + let range = NSRange(text.startIndex.. 1, + let valueRange = Range(match.range(at: 1), in: text) + else { + return nil + } + let value = text[valueRange].trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : String(value) + } +} + +extension [String] { + fileprivate func uniquedSorted() -> [String] { + Array(Set(self)).sorted() + } +} diff --git a/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanUsageSnapshot.swift new file mode 100644 index 0000000..36cee2a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Alibaba/AlibabaTokenPlanUsageSnapshot.swift @@ -0,0 +1,96 @@ +import Foundation + +public struct AlibabaTokenPlanUsageSnapshot: Sendable { + public let planName: String? + public let usedQuota: Double? + public let totalQuota: Double? + public let remainingQuota: Double? + public let resetsAt: Date? + public let updatedAt: Date + + public init( + planName: String?, + usedQuota: Double?, + totalQuota: Double?, + remainingQuota: Double?, + resetsAt: Date?, + updatedAt: Date) + { + self.planName = planName + self.usedQuota = usedQuota + self.totalQuota = totalQuota + self.remainingQuota = remainingQuota + self.resetsAt = resetsAt + self.updatedAt = updatedAt + } +} + +extension AlibabaTokenPlanUsageSnapshot { + public func toUsageSnapshot() -> UsageSnapshot { + let primary: RateWindow? = Self.usedPercent( + used: self.usedQuota, + total: self.totalQuota, + remaining: self.remainingQuota).map { + RateWindow( + usedPercent: $0, + windowMinutes: 30 * 24 * 60, + resetsAt: self.resetsAt, + resetDescription: Self.quotaDetail( + used: self.usedQuota, + total: self.totalQuota, + remaining: self.remainingQuota)) + } + + let planName = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines) + let loginMethod = (planName?.isEmpty ?? true) ? nil : planName + let identity = ProviderIdentitySnapshot( + providerID: .alibabatokenplan, + accountEmail: nil, + accountOrganization: nil, + loginMethod: loginMethod) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } + + private static func usedPercent(used: Double?, total: Double?, remaining: Double?) -> Double? { + guard let total, total > 0 else { return nil } + let usedValue: Double? = if let used { + used + } else if let remaining { + total - remaining + } else { + nil + } + guard let usedValue else { return nil } + let normalizedUsed = max(0, min(usedValue, total)) + return normalizedUsed / total * 100 + } + + private static func quotaDetail(used: Double?, total: Double?, remaining: Double?) -> String? { + if let used, let total, total > 0 { + return "\(self.format(used)) / \(self.format(total)) credits used" + } + if let remaining, let total, total > 0 { + return "\(Self.format(remaining)) / \(Self.format(total)) credits left" + } + if let remaining { + return "\(Self.format(remaining)) credits left" + } + return nil + } + + private static func format(_ value: Double) -> String { + let formatter = NumberFormatter() + formatter.numberStyle = .decimal + formatter.usesGroupingSeparator = true + formatter.maximumFractionDigits = value.rounded() == value ? 0 : 2 + formatter.locale = Locale(identifier: "en_US_POSIX") + return formatter.string(from: NSNumber(value: value)) ?? String(format: "%.2f", value) + } +} diff --git a/Sources/CodexBarCore/Providers/Amp/AmpCLIProbe.swift b/Sources/CodexBarCore/Providers/Amp/AmpCLIProbe.swift new file mode 100644 index 0000000..ab995dd --- /dev/null +++ b/Sources/CodexBarCore/Providers/Amp/AmpCLIProbe.swift @@ -0,0 +1,46 @@ +import Foundation + +public struct AmpCLIProbe: Sendable { + private static let commandTimeout: TimeInterval = 15 + private let arguments: [String] + + public init() { + self.arguments = ["usage"] + } + + init(arguments: [String]) { + self.arguments = arguments + } + + public func fetch( + environment: [String: String] = ProcessInfo.processInfo.environment, + now: Date = Date()) async throws -> AmpUsageSnapshot + { + let loginPATH = LoginShellPathCache.shared.current + guard let executable = BinaryLocator.resolveAmpBinary(env: environment, loginPATH: loginPATH) else { + throw SubprocessRunnerError.binaryNotFound("amp") + } + + var commandEnvironment = environment + commandEnvironment["NO_COLOR"] = "1" + commandEnvironment["PATH"] = PathBuilder.effectivePATH( + purposes: [.tty, .nodeTooling], + env: environment, + loginPATH: loginPATH) + + let result = try await SubprocessRunner.run( + binary: executable, + arguments: self.arguments, + environment: commandEnvironment, + timeout: Self.commandTimeout, + standardInput: FileHandle.nullDevice, + label: "amp-usage") + let output = result.stdout.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + ? result.stderr + : result.stdout + guard !output.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw AmpUsageError.parseFailed("The Amp CLI returned no usage data.") + } + return try AmpUsageParser.parse(displayText: output, now: now) + } +} diff --git a/Sources/CodexBarCore/Providers/Amp/AmpProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Amp/AmpProviderDescriptor.swift new file mode 100644 index 0000000..91ecef4 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Amp/AmpProviderDescriptor.swift @@ -0,0 +1,169 @@ +import Foundation + +public enum AmpProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .amp, + metadata: ProviderMetadata( + id: .amp, + displayName: "Amp", + sessionLabel: "Amp Free", + weeklyLabel: "Balance", + opusLabel: nil, + supportsOpus: false, + supportsCredits: true, + creditsHint: "Individual and workspace credit balances from Amp.", + toggleTitle: "Show Amp usage", + cliName: "amp", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder, + dashboardURL: "https://ampcode.com/settings/usage", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .amp, + iconResourceName: "ProviderIcon-amp", + color: ProviderColor(red: 220 / 255, green: 38 / 255, blue: 38 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Amp cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api, .web, .cli], + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), + cli: ProviderCLIConfig( + name: "amp", + versionDetector: nil)) + } + + private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] { + switch context.sourceMode { + case .auto: + [AmpCLIFetchStrategy(), AmpAPIFetchStrategy(), AmpStatusFetchStrategy()] + case .cli: + [AmpCLIFetchStrategy()] + case .api: + [AmpAPIFetchStrategy()] + case .web: + [AmpStatusFetchStrategy()] + case .oauth: + [] + } + } +} + +struct AmpCLIFetchStrategy: ProviderFetchStrategy { + let id: String = "amp.cli" + let kind: ProviderFetchKind = .cli + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + BinaryLocator.resolveAmpBinary( + env: context.env, + loginPATH: LoginShellPathCache.shared.current) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let snapshot = try await AmpCLIProbe().fetch(environment: context.env) + return self.makeResult( + usage: snapshot.toUsageSnapshot(now: snapshot.updatedAt), + sourceLabel: "cli") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + if error is CancellationError || (error as? URLError)?.code == .cancelled { + return false + } + return context.sourceMode == .auto + } +} + +struct AmpAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "amp.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + _ = context + return true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let token = ProviderTokenResolver.ampToken(environment: context.env) else { + throw AmpUsageError.missingAPIToken + } + let logger: ((String) -> Void)? = context.verbose + ? { msg in CodexBarLog.logger(LogCategories.amp).verbose(msg) } + : nil + let snapshot = try await AmpUsageFetcher(browserDetection: context.browserDetection) + .fetch(apiToken: token, logger: logger) + return self.makeResult( + usage: snapshot.toUsageSnapshot(now: snapshot.updatedAt), + sourceLabel: "api") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + guard context.sourceMode == .auto else { return false } + return !(error is CancellationError) && (error as? URLError)?.code != .cancelled + } +} + +struct AmpStatusFetchStrategy: ProviderFetchStrategy { + let id: String = "amp.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + guard context.sourceMode.usesWeb else { return false } + #if os(macOS) + let canImportBrowserCookies = true + #else + let canImportBrowserCookies = false + #endif + return Self.canUseWebFallback( + settings: context.settings?.amp, + canImportBrowserCookies: canImportBrowserCookies) + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let fetcher = AmpUsageFetcher(browserDetection: context.browserDetection) + let manual = Self.manualCookieHeader(from: context) + let logger: ((String) -> Void)? = context.verbose + ? { msg in CodexBarLog.logger(LogCategories.amp).verbose(msg) } + : nil + let snap = try await fetcher.fetch(cookieHeaderOverride: manual, logger: logger) + return self.makeResult( + usage: snap.toUsageSnapshot(now: snap.updatedAt), + sourceLabel: "web") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + static func canUseWebFallback( + settings: ProviderSettingsSnapshot.AmpProviderSettings?, + canImportBrowserCookies: Bool) -> Bool + { + guard let settings else { return canImportBrowserCookies } + switch settings.cookieSource { + case .auto: + return canImportBrowserCookies + case .manual: + return Self.manualCookieHeader(from: settings.manualCookieHeader) != nil + case .off: + return false + } + } + + private static func manualCookieHeader(from context: ProviderFetchContext) -> String? { + guard let settings = context.settings?.amp, settings.cookieSource == .manual else { return nil } + return Self.manualCookieHeader(from: settings.manualCookieHeader) + } + + private static func manualCookieHeader(from rawHeader: String?) -> String? { + guard let header = CookieHeaderNormalizer.normalize(rawHeader), + CookieHeaderNormalizer.pairs(from: header).contains(where: { $0.name == "session" }) + else { return nil } + return header + } +} diff --git a/Sources/CodexBarCore/Providers/Amp/AmpSettingsReader.swift b/Sources/CodexBarCore/Providers/Amp/AmpSettingsReader.swift new file mode 100644 index 0000000..5997e35 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Amp/AmpSettingsReader.swift @@ -0,0 +1,22 @@ +import Foundation + +public enum AmpSettingsReader { + public static let apiTokenKey = "AMP_API_KEY" + + public static func apiToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.cleaned(environment[self.apiTokenKey]) + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} diff --git a/Sources/CodexBarCore/Providers/Amp/AmpUsageFetcher.swift b/Sources/CodexBarCore/Providers/Amp/AmpUsageFetcher.swift new file mode 100644 index 0000000..17d9c29 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Amp/AmpUsageFetcher.swift @@ -0,0 +1,491 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +#if os(macOS) +import SweetCookieKit +#endif + +public enum AmpUsageError: LocalizedError, Sendable { + case notLoggedIn + case invalidCredentials + case missingAPIToken + case invalidAPIToken + case parseFailed(String) + case networkError(String) + case noSessionCookie + + public var errorDescription: String? { + switch self { + case .notLoggedIn: + "Not logged in to Amp. Please log in via ampcode.com." + case .invalidCredentials: + "Amp session cookie expired. Please log in again." + case .missingAPIToken: + "Amp access token not configured. Set AMP_API_KEY or add it in Settings." + case .invalidAPIToken: + "Amp access token is invalid or expired." + case let .parseFailed(message): + "Could not parse Amp usage: \(message)" + case let .networkError(message): + "Amp request failed: \(message)" + case .noSessionCookie: + "No Amp session cookie found. Please log in to ampcode.com in your browser." + } + } +} + +#if os(macOS) +private let ampCookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.amp]?.browserCookieOrder ?? Browser.defaultImportOrder + +public enum AmpCookieImporter { + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["ampcode.com", "www.ampcode.com"] + private static let sessionCookieNames: Set = [ + "session", + ] + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + } + + public static func importSession( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) throws -> SessionInfo + { + let log: (String) -> Void = { msg in logger?("[amp-cookie] \(msg)") } + + let installed = ampCookieImportOrder.cookieImportCandidates(using: browserDetection) + for browserSource in installed { + do { + let query = BrowserCookieQuery(domains: self.cookieDomains) + let sources = try Self.cookieClient.codexBarRecords( + matching: query, + in: browserSource, + logger: log) + for source in sources where !source.records.isEmpty { + let cookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin) + guard !cookies.isEmpty else { continue } + let names = cookies.map(\.name).joined(separator: ", ") + log("\(source.label) cookies: \(names)") + let sessionCookies = cookies.filter { Self.sessionCookieNames.contains($0.name) } + if !sessionCookies.isEmpty { + log("Found Amp session cookie in \(source.label)") + return SessionInfo(cookies: sessionCookies, sourceLabel: source.label) + } + log("\(source.label) cookies found, but no Amp session cookie present") + log("Expected one of: \(Self.sessionCookieNames.joined(separator: ", "))") + } + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)") + } + } + + throw AmpUsageError.noSessionCookie + } +} +#endif + +public struct AmpUsageFetcher: Sendable { + private static let settingsURL = URL(string: "https://ampcode.com/settings")! + static let usageURL = URL(string: "https://ampcode.com/api/internal?userDisplayBalanceInfo")! + @MainActor private static var recentDumps: [String] = [] + + public let browserDetection: BrowserDetection + + public init(browserDetection: BrowserDetection) { + self.browserDetection = browserDetection + } + + public func fetch( + cookieHeaderOverride: String? = nil, + logger: ((String) -> Void)? = nil, + now: Date = Date()) async throws -> AmpUsageSnapshot + { + let log: (String) -> Void = { msg in logger?("[amp] \(msg)") } + let cookieHeader = try await self.resolveCookieHeader(override: cookieHeaderOverride, logger: log) + + if let logger { + let names = self.cookieNames(from: cookieHeader) + if !names.isEmpty { + logger("[amp] Cookie names: \(names.joined(separator: ", "))") + } + let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: logger) + do { + let (html, responseInfo) = try await self.fetchLegacyHTMLWithDiagnostics( + cookieHeader: cookieHeader, + diagnostics: diagnostics) + self.logDiagnostics(responseInfo: responseInfo, diagnostics: diagnostics, logger: logger) + return try AmpUsageParser.parse(html: html, now: now) + } catch { + self.logDiagnostics(responseInfo: nil, diagnostics: diagnostics, logger: logger) + logger("[amp] Fetch failed: \(error.localizedDescription)") + throw error + } + } + + let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: nil) + let (html, _) = try await self.fetchLegacyHTMLWithDiagnostics( + cookieHeader: cookieHeader, + diagnostics: diagnostics) + return try AmpUsageParser.parse(html: html, now: now) + } + + public func fetch( + apiToken: String, + logger: ((String) -> Void)? = nil, + now: Date = Date()) async throws -> AmpUsageSnapshot + { + guard let token = AmpSettingsReader.cleaned(apiToken) else { + throw AmpUsageError.missingAPIToken + } + let request = try Self.makeUsageAPIRequest(apiToken: token) + let diagnostics = APIRedirectDiagnostics(logger: logger) + let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil) + let httpResponse = try await session.response(for: request) + logger?("[amp] API response: \(httpResponse.statusCode) " + + "\(httpResponse.response.url?.absoluteString ?? "unknown")") + try Self.validateAPIResponse(httpResponse) + return try Self.parseUsageAPIResponse(httpResponse.data, now: now) + } + + public func debugRawProbe(cookieHeaderOverride: String? = nil) async -> String { + let stamp = ISO8601DateFormatter().string(from: Date()) + var lines: [String] = [] + lines.append("=== Amp Debug Probe @ \(stamp) ===") + lines.append("") + + do { + let cookieHeader = try await self.resolveCookieHeader( + override: cookieHeaderOverride, + logger: { msg in lines.append("[cookie] \(msg)") }) + let diagnostics = RedirectDiagnostics(cookieHeader: cookieHeader, logger: nil) + let cookieNames = CookieHeaderNormalizer.pairs(from: cookieHeader).map(\.name) + lines.append("Cookie names: \(cookieNames.joined(separator: ", "))") + + let (html, responseInfo) = try await self.fetchLegacyHTMLWithDiagnostics( + cookieHeader: cookieHeader, + diagnostics: diagnostics) + let snapshot = try AmpUsageParser.parse(html: html) + + lines.append("") + lines.append("Fetch Success") + lines.append("Status: \(responseInfo.statusCode) \(responseInfo.url)") + + if !diagnostics.redirects.isEmpty { + lines.append("") + lines.append("Redirects:") + for entry in diagnostics.redirects { + lines.append(" \(entry)") + } + } + + lines.append("") + lines.append("Amp Free:") + lines.append(" quota=\(snapshot.freeQuota?.description ?? "nil")") + lines.append(" used=\(snapshot.freeUsed?.description ?? "nil")") + lines.append(" hourlyReplenishment=\(snapshot.hourlyReplenishment?.description ?? "nil")") + lines.append(" windowHours=\(snapshot.windowHours?.description ?? "nil")") + lines.append(" individualCredits=\(snapshot.individualCredits?.description ?? "nil")") + for workspace in snapshot.workspaceBalances { + lines.append(" workspace[\(workspace.name)]=\(workspace.remaining)") + } + + let output = lines.joined(separator: "\n") + Task { @MainActor in Self.recordDump(output) } + return output + } catch { + lines.append("") + lines.append("Probe Failed: \(error.localizedDescription)") + let output = lines.joined(separator: "\n") + Task { @MainActor in Self.recordDump(output) } + return output + } + } + + public static func latestDumps() async -> String { + await MainActor.run { + let result = Self.recentDumps.joined(separator: "\n\n---\n\n") + return result.isEmpty ? "No Amp probe dumps captured yet." : result + } + } + + private func resolveCookieHeader( + override: String?, + logger: ((String) -> Void)?) async throws -> String + { + if let override = CookieHeaderNormalizer.normalize(override) { + if let sessionHeader = self.sessionCookieHeader(from: override) { + logger?("[amp] Using manual session cookie") + return sessionHeader + } + throw AmpUsageError.noSessionCookie + } + #if os(macOS) + let session = try AmpCookieImporter.importSession(browserDetection: self.browserDetection, logger: logger) + logger?("[amp] Using cookies from \(session.sourceLabel)") + return session.cookieHeader + #else + throw AmpUsageError.noSessionCookie + #endif + } + + static func makeUsageAPIRequest(apiToken: String) throws -> URLRequest { + var request = URLRequest(url: Self.usageURL) + request.httpMethod = "POST" + request.httpBody = try JSONSerialization.data(withJSONObject: [ + "method": "userDisplayBalanceInfo", + "params": [:], + ]) + request.setValue("Bearer \(apiToken)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "accept") + request.setValue("application/json", forHTTPHeaderField: "content-type") + return request + } + + private func fetchLegacyHTMLWithDiagnostics( + cookieHeader: String, + diagnostics: RedirectDiagnostics) async throws -> (String, ResponseInfo) + { + var request = URLRequest(url: Self.settingsURL) + request.httpMethod = "GET" + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue( + "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", + forHTTPHeaderField: "accept") + Self.applyBrowserHeaders(to: &request) + + let session = URLSession(configuration: .ephemeral, delegate: diagnostics, delegateQueue: nil) + let httpResponse = try await session.response(for: request) + let responseInfo = ResponseInfo( + statusCode: httpResponse.statusCode, + url: httpResponse.response.url?.absoluteString ?? "unknown") + try Self.validateBrowserResponse(response: httpResponse, diagnostics: diagnostics) + + let html = String(data: httpResponse.data, encoding: .utf8) ?? "" + return (html, responseInfo) + } + + static func parseUsageAPIResponse(_ data: Data, now: Date = Date()) throws -> AmpUsageSnapshot { + let response: UsageAPIResponse + do { + response = try JSONDecoder().decode(UsageAPIResponse.self, from: data) + } catch { + throw AmpUsageError.parseFailed("Invalid Amp usage API response.") + } + + guard response.ok else { + if response.error?.code == "auth-required" { + throw AmpUsageError.invalidAPIToken + } + throw AmpUsageError.networkError(response.error?.message ?? "Amp usage API returned an error.") + } + guard let displayText = response.result?.displayText, !displayText.isEmpty else { + throw AmpUsageError.parseFailed("Missing Amp usage display text.") + } + return try AmpUsageParser.parse(displayText: displayText, now: now) + } + + private static func applyBrowserHeaders(to request: inout URLRequest) { + request.setValue( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36", + forHTTPHeaderField: "user-agent") + request.setValue("en-US,en;q=0.9", forHTTPHeaderField: "accept-language") + request.setValue("https://ampcode.com", forHTTPHeaderField: "origin") + request.setValue(self.settingsURL.absoluteString, forHTTPHeaderField: "referer") + } + + private static func validateBrowserResponse( + response: ProviderHTTPResponse, + diagnostics: RedirectDiagnostics) throws + { + guard response.statusCode == 200 else { + if response.statusCode == 401 || response.statusCode == 403 || diagnostics.detectedLoginRedirect { + throw AmpUsageError.invalidCredentials + } + throw AmpUsageError.networkError("HTTP \(response.statusCode)") + } + } + + private static func validateAPIResponse(_ response: ProviderHTTPResponse) throws { + guard response.statusCode == 200 else { + if response.statusCode == 401 || response.statusCode == 403 { + throw AmpUsageError.invalidAPIToken + } + throw AmpUsageError.networkError("HTTP \(response.statusCode)") + } + } + + @MainActor private static func recordDump(_ text: String) { + if self.recentDumps.count >= 5 { self.recentDumps.removeFirst() } + self.recentDumps.append(text) + } + + private final class RedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable { + private let cookieHeader: String + private let logger: ((String) -> Void)? + var redirects: [String] = [] + private(set) var detectedLoginRedirect = false + + init(cookieHeader: String, logger: ((String) -> Void)?) { + self.cookieHeader = cookieHeader + self.logger = logger + } + + func urlSession( + _: URLSession, + task: URLSessionTask, + willPerformHTTPRedirection response: HTTPURLResponse, + newRequest request: URLRequest, + completionHandler: @escaping (URLRequest?) -> Void) + { + let from = response.url?.absoluteString ?? "unknown" + let to = request.url?.absoluteString ?? "unknown" + self.redirects.append("\(response.statusCode) \(from) -> \(to)") + + if let toURL = request.url, AmpUsageFetcher.isLoginRedirect(toURL) { + if let logger { + logger("[amp] Detected login redirect, aborting (invalid session)") + } + self.detectedLoginRedirect = true + completionHandler(nil) + return + } + + var updated = request + if AmpUsageFetcher.shouldAttachCookie(to: request.url), !self.cookieHeader.isEmpty { + updated.setValue(self.cookieHeader, forHTTPHeaderField: "Cookie") + } else { + updated.setValue(nil, forHTTPHeaderField: "Cookie") + } + if let referer = response.url?.absoluteString { + updated.setValue(referer, forHTTPHeaderField: "referer") + } + if let logger { + logger("[amp] Redirect \(response.statusCode) \(from) -> \(to)") + } + completionHandler(updated) + } + } + + /// Amp's balance RPC should not redirect. Refusing redirects guarantees the bearer token cannot cross hosts. + private final class APIRedirectDiagnostics: NSObject, URLSessionTaskDelegate, @unchecked Sendable { + private let logger: ((String) -> Void)? + + init(logger: ((String) -> Void)?) { + self.logger = logger + } + + func urlSession( + _: URLSession, + task _: URLSessionTask, + willPerformHTTPRedirection response: HTTPURLResponse, + newRequest request: URLRequest, + completionHandler: @escaping (URLRequest?) -> Void) + { + let from = response.url?.absoluteString ?? "unknown" + let to = request.url?.absoluteString ?? "unknown" + self.logger?("[amp] API redirect blocked: \(response.statusCode) \(from) -> \(to)") + completionHandler(nil) + } + } + + private struct ResponseInfo { + let statusCode: Int + let url: String + } + + private struct UsageAPIResponse: Decodable { + let ok: Bool + let result: Result? + let error: APIError? + + struct Result: Decodable { + let displayText: String + } + + struct APIError: Decodable { + let code: String? + let message: String? + } + } + + private func logDiagnostics( + responseInfo: ResponseInfo?, + diagnostics: RedirectDiagnostics, + logger: (String) -> Void) + { + if let responseInfo { + logger("[amp] Response: \(responseInfo.statusCode) \(responseInfo.url)") + } + if !diagnostics.redirects.isEmpty { + logger("[amp] Redirects:") + for entry in diagnostics.redirects { + logger("[amp] \(entry)") + } + } + } + + private func cookieNames(from header: String) -> [String] { + header.split(separator: ";").compactMap { part in + let trimmed = part.trimmingCharacters(in: .whitespacesAndNewlines) + guard let idx = trimmed.firstIndex(of: "=") else { return nil } + let name = trimmed[.. String? { + let pairs = CookieHeaderNormalizer.pairs(from: header) + let sessionPairs = pairs.filter { $0.name == "session" } + guard !sessionPairs.isEmpty else { return nil } + return sessionPairs.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + + static func shouldAttachCookie(to url: URL?) -> Bool { + guard url?.scheme?.lowercased() == "https" else { return false } + return self.isAmpHost(url) + } + + private static func isAmpHost(_ url: URL?) -> Bool { + guard let host = url?.host?.lowercased() else { return false } + if host == "ampcode.com" || host == "www.ampcode.com" { return true } + return host.hasSuffix(".ampcode.com") + } + + static func isLoginRedirect(_ url: URL) -> Bool { + guard self.isAmpHost(url) else { return false } + if url.host?.lowercased() == "auth.ampcode.com" { return true } + + let path = url.path.lowercased() + let components = path.split(separator: "/").map(String.init) + if components.contains("login") { return true } + if components.contains("signin") { return true } + if components.contains("sign-in") { return true } + + // Amp currently redirects to /auth/sign-in?returnTo=... when session is invalid. Keep this slightly broader + // than one exact path so we keep working if Amp changes auth routes. + if components.contains("auth") { + let query = url.query?.lowercased() ?? "" + if query.contains("returnto=") { return true } + if query.contains("redirect=") { return true } + if query.contains("redirectto=") { return true } + } + + return false + } +} diff --git a/Sources/CodexBarCore/Providers/Amp/AmpUsageParser.swift b/Sources/CodexBarCore/Providers/Amp/AmpUsageParser.swift new file mode 100644 index 0000000..90ac8fc --- /dev/null +++ b/Sources/CodexBarCore/Providers/Amp/AmpUsageParser.swift @@ -0,0 +1,221 @@ +import Foundation + +enum AmpUsageParser { + static func parse(html: String, now: Date = Date()) throws -> AmpUsageSnapshot { + guard let usage = self.parseFreeTierUsage(html) else { + if self.looksSignedOut(html) { + throw AmpUsageError.notLoggedIn + } + throw AmpUsageError.parseFailed("Missing Amp Free usage data.") + } + + return AmpUsageSnapshot( + freeQuota: usage.quota, + freeUsed: usage.used, + hourlyReplenishment: usage.hourlyReplenishment, + windowHours: usage.windowHours, + updatedAt: now) + } + + static func parse(displayText: String, now: Date = Date()) throws -> AmpUsageSnapshot { + let text = TextParsing.stripANSICodes(displayText) + let identityPattern = #"(?im)^\s*Signed in as\s+([^\s(]+)(?:\s+\(([^\r\n)]+)\))?\s*$"# + let identity = self.captures(in: text, pattern: identityPattern) + if identity == nil, self.looksSignedOut(text) { + throw AmpUsageError.notLoggedIn + } + + let amountPattern = #"([0-9][0-9,]*(?:\.[0-9]+)?)"# + let freePattern = #"(?im)^\s*Amp Free:\s*\$?"# + amountPattern + + #"\s*/\s*\$?"# + amountPattern + + #"\s+remaining(?:\s*\(replenishes\s*\+\$?"# + amountPattern + #"\s*/\s*hour\))?"# + let freePercentPattern = #"(?im)^\s*Amp Free:\s*"# + amountPattern + + #"\s*%\s+remaining(?:\s+today)?(?:\s*\(resets\s+daily\))?"# + let creditsPattern = #"(?im)^\s*Individual credits:\s*\$?"# + amountPattern + #"\s+remaining"# + let individualCredits = self.captures(in: text, pattern: creditsPattern)?.first + .flatMap(self.number(from:)) + let workspacePattern = #"(?im)^\s*Workspace\s+(.+?):\s*\$?"# + amountPattern + #"\s+remaining"# + let workspaceBalances: [AmpWorkspaceBalance] = self.allCaptures( + in: text, + pattern: workspacePattern).compactMap { captures -> AmpWorkspaceBalance? in + guard captures.count == 2, + let name = self.nonEmpty(captures[0]), + let remaining = self.number(from: captures[1]) + else { return nil } + return AmpWorkspaceBalance(name: name, remaining: remaining) + } + let freeUsage: FreeTierUsage? = { + guard let free = self.captures(in: text, pattern: freePattern), + let remaining = self.number(from: free[0]), + let quota = self.number(from: free[1]) + else { return nil } + let hourlyReplenishment = self.number(from: free[2]) ?? 0 + let windowHours = hourlyReplenishment > 0 + ? max(1, (quota / hourlyReplenishment).rounded()) + : nil + return FreeTierUsage( + quota: quota, + used: max(0, quota - remaining), + hourlyReplenishment: hourlyReplenishment, + windowHours: windowHours, + resetDescription: nil) + }() + let freePercentUsage: FreeTierUsage? = { + guard let remainingText = self.captures(in: text, pattern: freePercentPattern)?.first, + let remaining = self.number(from: remainingText) + else { return nil } + let clampedRemaining = min(100, max(0, remaining)) + return FreeTierUsage( + quota: 100, + used: 100 - clampedRemaining, + hourlyReplenishment: 0, + windowHours: 24, + resetDescription: "resets daily") + }() + let resolvedFreeUsage = freeUsage ?? freePercentUsage + guard resolvedFreeUsage != nil || individualCredits != nil || !workspaceBalances.isEmpty else { + throw AmpUsageError.parseFailed("Missing Amp usage data.") + } + + return AmpUsageSnapshot( + freeQuota: resolvedFreeUsage?.quota, + freeUsed: resolvedFreeUsage?.used, + hourlyReplenishment: resolvedFreeUsage?.hourlyReplenishment, + windowHours: resolvedFreeUsage?.windowHours, + individualCredits: individualCredits, + workspaceBalances: workspaceBalances, + accountEmail: self.nonEmpty(identity?[0]), + accountOrganization: self.nonEmpty(identity?[1]), + updatedAt: now, + freeResetDescription: resolvedFreeUsage?.resetDescription) + } + + private struct FreeTierUsage { + let quota: Double + let used: Double + let hourlyReplenishment: Double + let windowHours: Double? + let resetDescription: String? + } + + private static func parseFreeTierUsage(_ html: String) -> FreeTierUsage? { + let tokens = ["freeTierUsage", "getFreeTierUsage"] + for token in tokens { + if let object = self.extractObject(named: token, in: html), + let usage = self.parseFreeTierUsageObject(object) + { + return usage + } + } + return nil + } + + private static func parseFreeTierUsageObject(_ object: String) -> FreeTierUsage? { + guard let quota = self.number(for: "quota", in: object), + let used = self.number(for: "used", in: object), + let hourly = self.number(for: "hourlyReplenishment", in: object) + else { return nil } + + let windowHours = self.number(for: "windowHours", in: object) + return FreeTierUsage( + quota: quota, + used: used, + hourlyReplenishment: hourly, + windowHours: windowHours, + resetDescription: nil) + } + + private static func extractObject(named token: String, in text: String) -> String? { + guard let tokenRange = text.range(of: token) else { return nil } + guard let braceIndex = text[tokenRange.upperBound...].firstIndex(of: "{") else { return nil } + + var depth = 0 + var inString = false + var isEscaped = false + var index = braceIndex + + while index < text.endIndex { + let char = text[index] + if inString { + if isEscaped { + isEscaped = false + } else if char == "\\" { + isEscaped = true + } else if char == "\"" { + inString = false + } + } else { + if char == "\"" { + inString = true + } else if char == "{" { + depth += 1 + } else if char == "}" { + depth -= 1 + if depth == 0 { + return String(text[braceIndex...index]) + } + } + } + index = text.index(after: index) + } + + return nil + } + + private static func number(for key: String, in text: String) -> Double? { + let pattern = "\\b\(NSRegularExpression.escapedPattern(for: key))\\b\\s*:\\s*([0-9]+(?:\\.[0-9]+)?)" + guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil } + let range = NSRange(text.startIndex.. 1, + let valueRange = Range(match.range(at: 1), in: text) + else { return nil } + return Double(text[valueRange]) + } + + private static func captures(in text: String, pattern: String) -> [String]? { + guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil } + let range = NSRange(text.startIndex.. [[String]] { + guard let regex = try? NSRegularExpression(pattern: pattern) else { return [] } + let range = NSRange(text.startIndex.. [String] { + (1.. Double? { + Double(text.replacingOccurrences(of: ",", with: "")) + } + + private static func nonEmpty(_ text: String?) -> String? { + guard let text, !text.isEmpty else { return nil } + return text + } + + private static func looksSignedOut(_ html: String) -> Bool { + let lower = html.lowercased() + if lower.contains("sign in") || lower.contains("log in") || lower.contains("login") { + return true + } + if lower.contains("/login") || lower.contains("ampcode.com/login") { + return true + } + return false + } +} diff --git a/Sources/CodexBarCore/Providers/Amp/AmpUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Amp/AmpUsageSnapshot.swift new file mode 100644 index 0000000..dfa670d --- /dev/null +++ b/Sources/CodexBarCore/Providers/Amp/AmpUsageSnapshot.swift @@ -0,0 +1,109 @@ +import Foundation + +public struct AmpWorkspaceBalance: Codable, Equatable, Sendable { + public let name: String + public let remaining: Double + + public init(name: String, remaining: Double) { + self.name = name + self.remaining = remaining + } +} + +public struct AmpUsageDetails: Codable, Equatable, Sendable { + public let individualCredits: Double? + public let workspaceBalances: [AmpWorkspaceBalance] + + public init(individualCredits: Double?, workspaceBalances: [AmpWorkspaceBalance]) { + self.individualCredits = individualCredits + self.workspaceBalances = workspaceBalances + } +} + +public struct AmpUsageSnapshot: Sendable { + public let freeQuota: Double? + public let freeUsed: Double? + public let hourlyReplenishment: Double? + public let windowHours: Double? + public let individualCredits: Double? + public let workspaceBalances: [AmpWorkspaceBalance] + public let accountEmail: String? + public let accountOrganization: String? + public let updatedAt: Date + public let freeResetDescription: String? + + public init( + freeQuota: Double?, + freeUsed: Double?, + hourlyReplenishment: Double?, + windowHours: Double?, + individualCredits: Double? = nil, + workspaceBalances: [AmpWorkspaceBalance] = [], + accountEmail: String? = nil, + accountOrganization: String? = nil, + updatedAt: Date, + freeResetDescription: String? = nil) + { + self.freeQuota = freeQuota + self.freeUsed = freeUsed + self.hourlyReplenishment = hourlyReplenishment + self.windowHours = windowHours + self.individualCredits = individualCredits + self.workspaceBalances = workspaceBalances + self.accountEmail = accountEmail + self.accountOrganization = accountOrganization + self.updatedAt = updatedAt + self.freeResetDescription = freeResetDescription + } +} + +extension AmpUsageSnapshot { + public func toUsageSnapshot(now: Date = Date()) -> UsageSnapshot { + let primary: RateWindow? = if let freeQuota, let freeUsed { + { + let quota = max(0, freeQuota) + let used = max(0, freeUsed) + let percent = quota > 0 ? min(100, (used / quota) * 100) : 0 + let windowMinutes: Int? = if let hours = self.windowHours, hours > 0 { + Int((hours * 60).rounded()) + } else { + nil + } + let resetsAt: Date? = { + guard quota > 0, let hourlyReplenishment, hourlyReplenishment > 0 else { return nil } + return now.addingTimeInterval(max(0, used / hourlyReplenishment * 3600)) + }() + return RateWindow( + usedPercent: percent, + windowMinutes: windowMinutes, + resetsAt: resetsAt, + resetDescription: self.freeResetDescription) + }() + } else { + nil + } + + let identity = ProviderIdentitySnapshot( + providerID: .amp, + accountEmail: self.accountEmail, + accountOrganization: self.accountOrganization, + loginMethod: primary == nil ? "Amp" : "Amp Free") + + let ampUsage: AmpUsageDetails? = if self.individualCredits != nil || !self.workspaceBalances.isEmpty { + AmpUsageDetails( + individualCredits: self.individualCredits, + workspaceBalances: self.workspaceBalances) + } else { + nil + } + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + ampUsage: ampUsage, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityCLISession.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityCLISession.swift new file mode 100644 index 0000000..d86d929 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityCLISession.swift @@ -0,0 +1,1297 @@ +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif +import Foundation + +// MARK: - Antigravity CLI Process Abstractions + +protocol AntigravityCLIProcessHandle: AnyObject, Sendable { + var pid: pid_t { get } + var isRunning: Bool { get } + var processGroup: pid_t? { get } + + func assignProcessGroup() -> pid_t? + func sendExit() throws + func closePTY() + func terminateRoot() + func killRoot() + func descendantPIDs() -> [pid_t] + func terminateTree(signal: Int32, knownDescendants: [pid_t]) + func killDescendants(_ descendants: [pid_t]) + func drainOutput() -> Data +} + +protocol AntigravityCLIProcessLaunching: Sendable { + func launch(binary: String) throws -> any AntigravityCLIProcessHandle +} + +enum AntigravityCLIAuthenticationPrompt { + static let evidence = Data("Select login method:".utf8) + private static let promptPattern = #"select\s+login\s+method\s*:?"# + + static func contains(_ output: Data) -> Bool { + // `agy` briefly prints "You are currently not signed in" before it + // auto-refreshes an existing login. The actual blocking state is the + // interactive login-method prompt. The exact prompt is a CLI-owned TUI + // string, so keep matching tolerant to casing and whitespace changes. + if output.range(of: self.evidence) != nil { + return true + } + let asciiBytes = output.map { $0 < 0x80 ? $0 : 0x20 } + let text = String(bytes: asciiBytes, encoding: .utf8) ?? "" + return text.range( + of: self.promptPattern, + options: [.regularExpression, .caseInsensitive]) != nil + } +} + +struct AntigravityCLIProcessIdentity: Equatable { + let executablePath: String + let startEpoch: TimeInterval +} + +protocol AntigravityCLIProcessIdentityProviding: Sendable { + func identity(for pid: pid_t) -> AntigravityCLIProcessIdentity? +} + +struct AntigravityCLISessionRecord: Codable, Equatable { + let pid: pid_t + let requestedBinaryPath: String + let executablePath: String + let startEpoch: TimeInterval + let processGroup: pid_t? + let ownerPID: pid_t? + let ownerExecutablePath: String? + let ownerStartEpoch: TimeInterval? + + init( + pid: pid_t, + requestedBinaryPath: String, + executablePath: String, + startEpoch: TimeInterval, + processGroup: pid_t?, + ownerPID: pid_t? = nil, + ownerExecutablePath: String? = nil, + ownerStartEpoch: TimeInterval? = nil) + { + self.pid = pid + self.requestedBinaryPath = requestedBinaryPath + self.executablePath = executablePath + self.startEpoch = startEpoch + self.processGroup = processGroup + self.ownerPID = ownerPID + self.ownerExecutablePath = ownerExecutablePath + self.ownerStartEpoch = ownerStartEpoch + } +} + +protocol AntigravityCLISessionRecordStoring: Sendable { + func load() throws -> [AntigravityCLISessionRecord] + func save(_ record: AntigravityCLISessionRecord) throws + func remove(_ record: AntigravityCLISessionRecord) throws +} + +protocol AntigravityCLISessionLaunchLocking: Sendable { + func withLock(_ operation: () throws -> T) throws -> T +} + +// MARK: - AntigravityCLISession + +/// Manages a bounded background ``agy`` process whose embedded localhost server +/// provides the same ``GetUserStatus`` endpoint as the desktop Antigravity app's +/// ``language_server``. The CLI is kept alive in a PTY so its daemon stays bound +/// to a local port - this lets CodexBar read Claude + Gemini quotas even when +/// the desktop Antigravity app is closed. +/// +/// The session intentionally does not scrape TUI output. It only launches and +/// keeps the process reachable for HTTPS probing, drains discarded PTY output so +/// the CLI cannot block on a full terminal buffer, and bounds the warm lifetime +/// with an idle timer so CodexBar does not run an IDE backend forever. +actor AntigravityCLISession { + static let shared = AntigravityCLISession() + private static let log = CodexBarLog.logger(LogCategories.antigravity) + + enum ResetCause: Int { + case deferred + case oneShotCLI + case unhealthy + case authentication + + var message: String { + switch self { + case .deferred: "deferred reset" + case .oneShotCLI: "one-shot CLI fetch" + case .unhealthy: "unhealthy CLI HTTPS session" + case .authentication: "authentication required" + } + } + } + + private struct LaunchOutcome { + let pid: pid_t + let rejectedProcess: (any AntigravityCLIProcessHandle)? + let rejectionMessage: String? + let holdsLaunchReservation: Bool + } + + struct Dependencies { + var launcher: any AntigravityCLIProcessLaunching + var identityProvider: any AntigravityCLIProcessIdentityProviding + var recordStore: any AntigravityCLISessionRecordStoring + var launchLock: any AntigravityCLISessionLaunchLocking + var beginAppShutdownTrackedLaunch: @Sendable () -> Bool + var endAppShutdownTrackedLaunch: @Sendable () -> Void + var registerForAppShutdown: @Sendable (pid_t, String) -> Bool + var updateAppShutdownProcessGroup: @Sendable (pid_t, pid_t?) -> Void + var unregisterForAppShutdown: @Sendable (pid_t) -> Void + var descendantPIDs: @Sendable (pid_t) -> [pid_t] + var terminateProcessTree: @Sendable (pid_t, pid_t?, Int32, [pid_t]) -> Void + var currentProcessID: @Sendable () -> pid_t + var now: @Sendable () -> Date + var sleep: @Sendable (UInt64) async throws -> Void + var idleWindow: TimeInterval + var failureRelaunchThreshold: Int + var terminationGracePeriod: TimeInterval + + static func live() -> Self { + Self( + launcher: AntigravityPTYProcessLauncher(), + identityProvider: AntigravityProcessIdentityProvider(), + recordStore: AntigravityFileCLISessionRecordStore(), + launchLock: AntigravityFileCLISessionLaunchLock(), + beginAppShutdownTrackedLaunch: { + TTYCommandRunner.beginActiveProcessLaunchForAppShutdown() + }, + endAppShutdownTrackedLaunch: { + TTYCommandRunner.endActiveProcessLaunchForAppShutdown() + }, + registerForAppShutdown: { pid, binary in + TTYCommandRunner.registerActiveProcessForAppShutdown(pid: pid, binary: binary) + }, + updateAppShutdownProcessGroup: { pid, group in + TTYCommandRunner.updateActiveProcessGroupForAppShutdown(pid: pid, processGroup: group) + }, + unregisterForAppShutdown: { pid in + TTYCommandRunner.unregisterActiveProcessForAppShutdown(pid: pid) + }, + descendantPIDs: { pid in + TTYProcessTreeTerminator.descendantPIDs(of: pid) + }, + terminateProcessTree: { pid, group, signal, knownDescendants in + TTYProcessTreeTerminator.terminateProcessTree( + rootPID: pid, + processGroup: group, + signal: signal, + knownDescendants: knownDescendants) + }, + currentProcessID: getpid, + now: Date.init, + sleep: { nanoseconds in + try await Task.sleep(nanoseconds: nanoseconds) + }, + idleWindow: 180, + failureRelaunchThreshold: 2, + terminationGracePeriod: 1) + } + } + + // MARK: State + + private let dependencies: Dependencies + private var process: (any AntigravityCLIProcessHandle)? + private var binaryPath: String? + private var sessionIdleWindow: TimeInterval + private var activeProbeCount = 0 + private var activeSessionProbeCount = 0 + private var resetRequestedWhenIdle = false + private var hardResetRequestedWhenIdle = false + private var pendingResetCause: ResetCause? + private var idleTask: Task? + private var sessionGeneration: UInt64 = 0 + private var consecutiveProbeFailures = 0 + private var persistedProcessIdentity: AntigravityCLIProcessIdentity? + private var recentOutput = Data() + private var authenticationPromptObserved = false + private var lastStopReason: String? + private var lifecycleOperationInProgress = false + private var lifecycleWaiters: [CheckedContinuation] = [] + private var exclusiveProbeWaiters: [CheckedContinuation] = [] + + init(dependencies: Dependencies = .live()) { + self.dependencies = dependencies + self.sessionIdleWindow = dependencies.idleWindow + } + + /// The pid of the running ``agy`` process, exposed so callers can discover + /// its listening ports via `lsof`. + var pid: pid_t? { + guard let proc = self.process, proc.isRunning else { return nil } + return proc.pid + } + + /// Whether the managed process is alive and matches ``binaryPath``. + var isRunning: Bool { + guard let proc = self.process, proc.isRunning, self.binaryPath != nil else { return false } + return true + } + + var failureCountForTesting: Int { + self.consecutiveProbeFailures + } + + var idleWindowForTesting: TimeInterval { + self.sessionIdleWindow + } + + var activeProbeCountForTesting: Int { + self.activeProbeCount + } + + var lastStopReasonForTesting: String? { + self.lastStopReason + } + + // MARK: Lifecycle + + /// Mark a probe as active and ensure a warm ``agy`` is running on the given binary path. + /// + /// Callers must balance this with ``finishProbe(success:resetAfterFetch:)`` so + /// idle/reset cleanup cannot kill the process while its ports are being probed. + /// If previous probes repeatedly failed while the process stayed alive, this + /// force-relaunches instead of reusing a wedged HTTPS server forever. + func beginProbe(binary: String, idleWindow: TimeInterval? = nil) async throws -> pid_t { + self.activeProbeCount += 1 + if let idleWindow, idleWindow > 0 { + self.sessionIdleWindow = max(self.dependencies.idleWindow, idleWindow) + } + self.cancelIdleTimer() + do { + return try await self.withLifecycleOperation { + let pid = try await self.ensureStartedLocked(binary: binary) + self.activeSessionProbeCount += 1 + return pid + } + } catch { + self.activeProbeCount = max(0, self.activeProbeCount - 1) + self.notifyExclusiveProbeWaitersIfNeeded() + if self.activeProbeCount == 0, self.resetRequestedWhenIdle { + await self.withLifecycleOperation { + guard self.activeProbeCount == 0 else { + self.resetRequestedWhenIdle = true + return + } + let forceTerminate = self.hardResetRequestedWhenIdle || self.consecutiveProbeFailures > 0 + self.resetRequestedWhenIdle = false + self.hardResetRequestedWhenIdle = false + await self.stopCurrentSessionLocked( + reason: "deferred reset after failed begin", + clearRecord: true, + graceful: !forceTerminate) + } + } + throw error + } + } + + /// Record probe completion and either keep the session warm for the bounded + /// idle window or tear it down immediately for one-shot CLI invocations. + func finishProbe(success: Bool, resetAfterFetch: Bool, forceTerminate: Bool = false) async { + if success { + self.consecutiveProbeFailures = 0 + } else { + self.consecutiveProbeFailures += 1 + } + + self.activeProbeCount = max(0, self.activeProbeCount - 1) + self.activeSessionProbeCount = max(0, self.activeSessionProbeCount - 1) + self.notifyExclusiveProbeWaitersIfNeeded() + let shouldForceStopUnhealthy = !success && + self.consecutiveProbeFailures >= max(1, self.dependencies.failureRelaunchThreshold) + let shouldReset = forceTerminate || resetAfterFetch || self.resetRequestedWhenIdle || shouldForceStopUnhealthy + let currentResetCause = Self.resetCause( + authenticationRequired: forceTerminate, + resetAfterFetch: resetAfterFetch, + shouldForceStopUnhealthy: shouldForceStopUnhealthy) + + guard self.activeProbeCount == 0 else { + if shouldReset { + self.resetRequestedWhenIdle = true + self.pendingResetCause = Self.strongestResetCause(self.pendingResetCause, currentResetCause) + } + if shouldReset, !success || forceTerminate { + self.hardResetRequestedWhenIdle = true + } + return + } + + if shouldReset { + let shouldForceTerminate = !success || forceTerminate || self.hardResetRequestedWhenIdle + let resetCause = Self.strongestResetCause(self.pendingResetCause, currentResetCause) + await self.withLifecycleOperation { + guard self.activeProbeCount == 0 else { + self.resetRequestedWhenIdle = true + self.pendingResetCause = Self.strongestResetCause(self.pendingResetCause, resetCause) + if shouldForceTerminate { + self.hardResetRequestedWhenIdle = true + } + return + } + self.resetRequestedWhenIdle = false + self.hardResetRequestedWhenIdle = false + self.pendingResetCause = nil + await self.stopCurrentSessionLocked( + reason: resetCause.message, + clearRecord: true, + graceful: !shouldForceTerminate) + } + } else { + self.armIdleTimer() + } + } + + static func resetCause( + authenticationRequired: Bool, + resetAfterFetch: Bool, + shouldForceStopUnhealthy: Bool) -> ResetCause + { + if authenticationRequired { + .authentication + } else if shouldForceStopUnhealthy { + .unhealthy + } else if resetAfterFetch { + .oneShotCLI + } else { + .deferred + } + } + + static func strongestResetCause(_ existing: ResetCause?, _ incoming: ResetCause) -> ResetCause { + guard let existing else { return incoming } + return existing.rawValue >= incoming.rawValue ? existing : incoming + } + + /// Ensure a warm ``agy`` is running on the given binary path. + /// + /// - If the process is already alive with the same binary, this returns immediately. + /// - If the process died, the binary changed, or repeated probes failed, it tears down the old one first. + /// - Returns the process identifier for port discovery. + func ensureStarted(binary: String) async throws -> pid_t { + try await self.withLifecycleOperation { + try await self.ensureStartedLocked(binary: binary) + } + } + + /// Request teardown. If a probe is in flight, cleanup is deferred until the + /// matching ``finishProbe(success:resetAfterFetch:)`` call. + func reset() async { + self.cancelIdleTimer() + guard self.activeProbeCount == 0 else { + self.resetRequestedWhenIdle = true + return + } + await self.withLifecycleOperation { + guard self.activeProbeCount == 0 else { + self.resetRequestedWhenIdle = true + return + } + self.resetRequestedWhenIdle = false + let forceTerminate = self.hardResetRequestedWhenIdle || self.consecutiveProbeFailures > 0 + self.hardResetRequestedWhenIdle = false + await self.stopCurrentSessionLocked( + reason: "manual reset", + clearRecord: true, + graceful: !forceTerminate) + } + } + + /// Drain PTY output into one rolling buffer shared by concurrent probes. + func drainOutput() -> Data { + var searchableOutput = self.recentOutput + if let output = self.process?.drainOutput(), !output.isEmpty { + searchableOutput.append(output) + self.recentOutput = Data(searchableOutput.suffix(4096)) + } + + if !self.authenticationPromptObserved, + AntigravityCLIAuthenticationPrompt.contains(searchableOutput) + { + self.authenticationPromptObserved = true + } + if self.authenticationPromptObserved, + !AntigravityCLIAuthenticationPrompt.contains(searchableOutput) + { + searchableOutput.append(AntigravityCLIAuthenticationPrompt.evidence) + } + return searchableOutput + } + + // MARK: Errors + + enum SessionError: LocalizedError { + case launchFailed(String) + + var errorDescription: String? { + switch self { + case let .launchFailed(msg): "Failed to launch Antigravity CLI session: \(msg)" + } + } + } + + // MARK: Private + + private func withLifecycleOperation(_ operation: () async throws -> T) async throws -> T { + await self.acquireLifecycleOperation() + defer { self.releaseLifecycleOperation() } + return try await operation() + } + + private func withLifecycleOperation(_ operation: () async -> Void) async { + await self.acquireLifecycleOperation() + defer { self.releaseLifecycleOperation() } + await operation() + } + + private func acquireLifecycleOperation() async { + guard self.lifecycleOperationInProgress else { + self.lifecycleOperationInProgress = true + return + } + await withCheckedContinuation { continuation in + self.lifecycleWaiters.append(continuation) + } + } + + private func releaseLifecycleOperation() { + guard !self.lifecycleWaiters.isEmpty else { + self.lifecycleOperationInProgress = false + return + } + let next = self.lifecycleWaiters.removeFirst() + next.resume() + } + + private func waitForExclusiveProbeIfNeeded() async { + while self.activeSessionProbeCount > 0 { + await withCheckedContinuation { continuation in + self.exclusiveProbeWaiters.append(continuation) + } + } + } + + private func notifyExclusiveProbeWaitersIfNeeded() { + guard self.activeSessionProbeCount == 0, !self.exclusiveProbeWaiters.isEmpty else { return } + let waiters = self.exclusiveProbeWaiters + self.exclusiveProbeWaiters.removeAll() + for waiter in waiters { + waiter.resume() + } + } + + private func ensureStartedLocked(binary: String) async throws -> pid_t { + while true { + if let proc = self.process, + proc.isRunning, + self.binaryPath == binary, + !self.resetRequestedWhenIdle, + !self.hardResetRequestedWhenIdle, + self.consecutiveProbeFailures < max(1, self.dependencies.failureRelaunchThreshold) + { + try? self.dependencies.launchLock.withLock { + self.reapRecordedSessionsIfNeeded() + } + self.persistCurrentRecordIfNeeded(proc, binary: binary) + Self.log.debug("Antigravity CLI session reused", metadata: ["pid": "\(proc.pid)"]) + return proc.pid + } + + if self.process != nil { + if self.activeSessionProbeCount > 0 { + await self.waitForExclusiveProbeIfNeeded() + continue + } + + let reason = self.consecutiveProbeFailures >= max(1, self.dependencies.failureRelaunchThreshold) + ? "relaunching unhealthy session" + : "replacing stale session" + let forceTerminate = self.hardResetRequestedWhenIdle || self.consecutiveProbeFailures > 0 + self.resetRequestedWhenIdle = false + self.hardResetRequestedWhenIdle = false + await self.stopCurrentSessionLocked( + reason: reason, + clearRecord: true, + graceful: !forceTerminate) + } + + let binaryName = URL(fileURLWithPath: binary).lastPathComponent + let launch: (Bool) throws -> LaunchOutcome = { canPersistRecord in + if canPersistRecord { + self.prepareRecordStoreForLaunch() + } + guard self.dependencies.beginAppShutdownTrackedLaunch() else { + throw SessionError.launchFailed("App shutdown in progress") + } + let launched: any AntigravityCLIProcessHandle + do { + launched = try self.dependencies.launcher.launch(binary: binary) + } catch { + self.dependencies.endAppShutdownTrackedLaunch() + throw error + } + let launchedPID = launched.pid + guard self.dependencies.registerForAppShutdown(launchedPID, binaryName) else { + return LaunchOutcome( + pid: launchedPID, + rejectedProcess: launched, + rejectionMessage: "App shutdown in progress", + holdsLaunchReservation: true) + } + + let processGroup = launched.processGroup ?? launched.assignProcessGroup() + self.dependencies.updateAppShutdownProcessGroup(launchedPID, processGroup) + + self.process = launched + self.binaryPath = binary + self.recentOutput.removeAll(keepingCapacity: true) + self.authenticationPromptObserved = false + self.consecutiveProbeFailures = 0 + self.sessionGeneration &+= 1 + if canPersistRecord { + _ = self.persistRecord(pid: launchedPID, binary: binary, processGroup: processGroup) + } else { + self.persistedProcessIdentity = nil + } + self.dependencies.endAppShutdownTrackedLaunch() + return LaunchOutcome( + pid: launchedPID, + rejectedProcess: nil, + rejectionMessage: nil, + holdsLaunchReservation: false) + } + + var lockedLaunch: Result? + var lockFailure: Error? + do { + try self.dependencies.launchLock.withLock { + lockedLaunch = Result { try launch(true) } + } + } catch { + lockFailure = error + } + + let outcome: LaunchOutcome + if let lockFailure { + Self.log.warning( + "Antigravity CLI session coordination unavailable", + metadata: ["error": lockFailure.localizedDescription]) + outcome = try launch(false) + } else if let lockedLaunch { + outcome = try lockedLaunch.get() + } else { + throw SessionError.launchFailed("CLI session launch did not complete") + } + if let rejectedProcess = outcome.rejectedProcess { + await self.terminateLaunchedProcess(rejectedProcess, graceful: false) + if outcome.holdsLaunchReservation { + self.dependencies.endAppShutdownTrackedLaunch() + } + throw SessionError.launchFailed(outcome.rejectionMessage ?? "App shutdown in progress") + } + + Self.log.debug( + "Antigravity CLI session started", + metadata: [ + "binary": binaryName, + "pid": "\(outcome.pid)", + ]) + return outcome.pid + } + } + + private func prepareRecordStoreForLaunch() { + guard self.process == nil || self.persistedProcessIdentity == nil else { return } + self.reapRecordedSessionsIfNeeded() + } + + private func persistCurrentRecordIfNeeded(_ proc: any AntigravityCLIProcessHandle, binary: String) { + guard self.persistedProcessIdentity == nil else { return } + try? self.dependencies.launchLock.withLock { + self.prepareRecordStoreForLaunch() + _ = self.persistRecord(pid: proc.pid, binary: binary, processGroup: proc.processGroup) + } + } + + private func cancelIdleTimer() { + self.idleTask?.cancel() + self.idleTask = nil + } + + private func armIdleTimer() { + guard self.process != nil, self.sessionIdleWindow > 0 else { return } + self.cancelIdleTimer() + let generation = self.sessionGeneration + let nanoseconds = Self.nanoseconds(from: self.sessionIdleWindow) + let sleep = self.dependencies.sleep + self.idleTask = Task { [weak self] in + do { + try await sleep(nanoseconds) + await self?.stopIfIdle(generation: generation) + } catch { + // Cancellation is the normal path when a refresh reuses the warm session. + } + } + } + + private func stopIfIdle(generation: UInt64) async { + await self.withLifecycleOperation { + guard generation == self.sessionGeneration else { return } + guard self.activeProbeCount == 0 else { + self.armIdleTimer() + return + } + let forceTerminate = self.hardResetRequestedWhenIdle || self.consecutiveProbeFailures > 0 + self.resetRequestedWhenIdle = false + self.hardResetRequestedWhenIdle = false + await self.stopCurrentSessionLocked( + reason: "idle timeout", + clearRecord: true, + graceful: !forceTerminate) + } + } + + private func stopCurrentSessionLocked(reason: String, clearRecord: Bool, graceful: Bool = true) async { + self.cancelIdleTimer() + let effectiveReason = self.pendingResetCause?.message ?? reason + self.pendingResetCause = nil + self.lastStopReason = effectiveReason + guard let proc = self.process else { + if clearRecord { + try? self.dependencies.launchLock.withLock { + self.reapRecordedSessionsIfNeeded() + } + } + self.sessionIdleWindow = self.dependencies.idleWindow + self.recentOutput.removeAll(keepingCapacity: true) + self.authenticationPromptObserved = false + return + } + + let pid = proc.pid + let identity = self.persistedProcessIdentity + Self.log.debug("Antigravity CLI session stopping", metadata: ["pid": "\(pid)", "reason": "\(effectiveReason)"]) + + self.process = nil + self.binaryPath = nil + self.persistedProcessIdentity = nil + self.sessionIdleWindow = self.dependencies.idleWindow + self.recentOutput.removeAll(keepingCapacity: true) + self.authenticationPromptObserved = false + self.sessionGeneration &+= 1 + + await self.terminateLaunchedProcess(proc, graceful: graceful) + self.dependencies.unregisterForAppShutdown(pid) + if clearRecord { + self.removeRecordIfMatches(pid: pid, identity: identity) + } + } + + private func terminateLaunchedProcess( + _ proc: any AntigravityCLIProcessHandle, + graceful: Bool = true) async + { + if graceful { + try? proc.sendExit() + } + proc.closePTY() + + let descendants = proc.descendantPIDs() + if proc.isRunning { + proc.terminateRoot() + } + proc.terminateTree(signal: SIGTERM, knownDescendants: descendants) + + let gracePeriod = self.dependencies.terminationGracePeriod + if gracePeriod > 0 { + let deadline = self.dependencies.now().addingTimeInterval(gracePeriod) + while proc.isRunning, self.dependencies.now() < deadline { + try? await self.dependencies.sleep(100_000_000) + } + } + + if proc.isRunning { + proc.terminateTree(signal: SIGKILL, knownDescendants: descendants) + await self.waitUntilProcessExits(proc, timeout: 1) + } else { + proc.killDescendants(descendants) + } + } + + private func waitUntilProcessExits(_ proc: any AntigravityCLIProcessHandle, timeout: TimeInterval) async { + let deadline = self.dependencies.now().addingTimeInterval(timeout) + while proc.isRunning, self.dependencies.now() < deadline { + try? await self.dependencies.sleep(50_000_000) + } + _ = proc.isRunning + } + + @discardableResult + private func persistRecord(pid: pid_t, binary: String, processGroup: pid_t?) -> Bool { + guard let identity = self.dependencies.identityProvider.identity(for: pid) else { + self.persistedProcessIdentity = nil + return false + } + let ownerPID = self.dependencies.currentProcessID() + let ownerIdentity = self.dependencies.identityProvider.identity(for: ownerPID) + let record = AntigravityCLISessionRecord( + pid: pid, + requestedBinaryPath: binary, + executablePath: identity.executablePath, + startEpoch: identity.startEpoch, + processGroup: processGroup, + ownerPID: ownerPID, + ownerExecutablePath: ownerIdentity?.executablePath, + ownerStartEpoch: ownerIdentity?.startEpoch) + do { + try self.dependencies.recordStore.save(record) + self.persistedProcessIdentity = identity + return true + } catch { + self.persistedProcessIdentity = nil + return false + } + } + + private func reapRecordedSessionsIfNeeded() { + guard let records = try? self.dependencies.recordStore.load() else { return } + for record in records { + guard let liveIdentity = self.dependencies.identityProvider.identity(for: record.pid), + liveIdentity.executablePath == record.executablePath, + abs(liveIdentity.startEpoch - record.startEpoch) < 0.001 + else { + try? self.dependencies.recordStore.remove(record) + continue + } + if let proc = self.process, proc.pid == record.pid { + self.persistedProcessIdentity = liveIdentity + continue + } + if let ownerPID = record.ownerPID, + ownerPID != self.dependencies.currentProcessID(), + let ownerExecutablePath = record.ownerExecutablePath, + let ownerStartEpoch = record.ownerStartEpoch, + let liveOwnerIdentity = self.dependencies.identityProvider.identity(for: ownerPID), + liveOwnerIdentity.executablePath == ownerExecutablePath, + abs(liveOwnerIdentity.startEpoch - ownerStartEpoch) < 0.001 + { + continue + } + + let knownDescendants = self.dependencies.descendantPIDs(record.pid) + Self.log.debug("Reaping stale Antigravity CLI session", metadata: ["pid": "\(record.pid)"]) + self.dependencies.terminateProcessTree(record.pid, record.processGroup, SIGTERM, knownDescendants) + self.dependencies.terminateProcessTree(record.pid, record.processGroup, SIGKILL, knownDescendants) + try? self.dependencies.recordStore.remove(record) + } + } + + private func removeRecordIfMatches(pid: pid_t, identity: AntigravityCLIProcessIdentity?) { + try? self.dependencies.launchLock.withLock { + guard let identity, let records = try? self.dependencies.recordStore.load() else { return } + for record in records + where record.pid == pid && + record.executablePath == identity.executablePath && + abs(record.startEpoch - identity.startEpoch) < 0.001 + { + try? self.dependencies.recordStore.remove(record) + } + } + } + + private static func nanoseconds(from interval: TimeInterval) -> UInt64 { + guard interval > 0 else { return 0 } + guard interval.isFinite else { return UInt64.max } + let nanoseconds = interval * 1_000_000_000 + guard nanoseconds < TimeInterval(UInt64.max) else { return UInt64.max } + return UInt64(nanoseconds) + } +} + +// MARK: - Production Process Implementation + +struct AntigravityPTYProcessLauncher: AntigravityCLIProcessLaunching { + static func defaultSignalsForSpawn() -> sigset_t { + var signals = sigset_t() + sigemptyset(&signals) + sigaddset(&signals, SIGINT) + sigaddset(&signals, SIGTERM) + sigaddset(&signals, SIGHUP) + return signals + } + + static func spawnWithTextBusyRetry( + maxAttempts: Int = 3, + retryDelay: TimeInterval = 0.01, + spawn: () -> Int32) -> Int32 + { + var result = spawn() + guard maxAttempts > 1 else { return result } + + for _ in 1.. 0 { + Thread.sleep(forTimeInterval: retryDelay) + } + result = spawn() + } + return result + } + + func launch(binary: String) throws -> any AntigravityCLIProcessHandle { + try self.launch(binary: binary, arguments: []) + } + + func launch(binary: String, arguments: [String]) throws -> any AntigravityCLIProcessHandle { + var primaryFD: Int32 = -1 + var secondaryFD: Int32 = -1 + var win = winsize(ws_row: 50, ws_col: 160, ws_xpixel: 0, ws_ypixel: 0) + guard openpty(&primaryFD, &secondaryFD, nil, nil, &win) == 0 else { + throw AntigravityCLISession.SessionError.launchFailed("openpty failed") + } + _ = fcntl(primaryFD, F_SETFL, O_NONBLOCK) + + let primaryHandle = FileHandle(fileDescriptor: primaryFD, closeOnDealloc: true) + let secondaryHandle = FileHandle(fileDescriptor: secondaryFD, closeOnDealloc: true) + + #if canImport(Darwin) + var fileActions: posix_spawn_file_actions_t? + #else + var fileActions = posix_spawn_file_actions_t() + #endif + guard posix_spawn_file_actions_init(&fileActions) == 0 else { + try? primaryHandle.close() + try? secondaryHandle.close() + throw AntigravityCLISession.SessionError.launchFailed("posix_spawn_file_actions_init failed") + } + defer { posix_spawn_file_actions_destroy(&fileActions) } + + posix_spawn_file_actions_adddup2(&fileActions, secondaryFD, 0) + posix_spawn_file_actions_adddup2(&fileActions, secondaryFD, 1) + posix_spawn_file_actions_adddup2(&fileActions, secondaryFD, 2) + posix_spawn_file_actions_addclose(&fileActions, primaryFD) + posix_spawn_file_actions_addclose(&fileActions, secondaryFD) + + let homeDirectory = NSHomeDirectory() + _ = homeDirectory.withCString { path in + posix_spawn_file_actions_addchdir_np(&fileActions, path) + } + #if canImport(Glibc) || canImport(Musl) + do { + try PosixSpawnFileActionsCloseFrom.addCloseFrom(&fileActions, startingAt: 3) + } catch { + try? primaryHandle.close() + try? secondaryHandle.close() + throw AntigravityCLISession.SessionError.launchFailed(error.localizedDescription) + } + #endif + + #if canImport(Darwin) + var attr: posix_spawnattr_t? + #else + var attr = posix_spawnattr_t() + #endif + guard posix_spawnattr_init(&attr) == 0 else { + try? primaryHandle.close() + try? secondaryHandle.close() + throw AntigravityCLISession.SessionError.launchFailed("posix_spawnattr_init failed") + } + defer { posix_spawnattr_destroy(&attr) } + var defaultSignals = Self.defaultSignalsForSpawn() + posix_spawnattr_setsigdefault(&attr, &defaultSignals) + #if canImport(Darwin) + let spawnFlags = POSIX_SPAWN_SETPGROUP | POSIX_SPAWN_SETSIGDEF | POSIX_SPAWN_CLOEXEC_DEFAULT + #else + let spawnFlags = POSIX_SPAWN_SETPGROUP | POSIX_SPAWN_SETSIGDEF + #endif + posix_spawnattr_setflags(&attr, Int16(spawnFlags)) + posix_spawnattr_setpgroup(&attr, 0) + + var env = TTYCommandRunner.enrichedEnvironment() + env["PWD"] = NSHomeDirectory() + env["TERM"] = "xterm-256color" + + var cArgs = ([binary] + arguments).map { strdup($0) as UnsafeMutablePointer? } + cArgs.append(nil) + defer { + for arg in cArgs { + if let arg { + free(arg) + } + } + } + + var cEnv: [UnsafeMutablePointer?] = env.map { key, value in + strdup("\(key)=\(value)") + } + cEnv.append(nil) + defer { + for entry in cEnv { + if let entry { + free(entry) + } + } + } + + var pid: pid_t = 0 + let spawnResult = Self.spawnWithTextBusyRetry { + binary.withCString { execPath in + posix_spawn(&pid, execPath, &fileActions, &attr, cArgs, cEnv) + } + } + guard spawnResult == 0 else { + try? primaryHandle.close() + try? secondaryHandle.close() + throw AntigravityCLISession.SessionError.launchFailed(String(cString: strerror(spawnResult))) + } + + return AntigravitySpawnedPTYProcessHandle( + pid: pid, + processGroup: pid, + primaryFD: primaryFD, + primaryHandle: primaryHandle, + secondaryHandle: secondaryHandle) + } +} + +final class AntigravitySpawnedPTYProcessHandle: AntigravityCLIProcessHandle, @unchecked Sendable { + private let lock = NSLock() + private let processPID: pid_t + private let processGroupID: pid_t + private let primaryFD: Int32 + private let primaryHandle: FileHandle + private let secondaryHandle: FileHandle + private var reaped = false + + init( + pid: pid_t, + processGroup: pid_t, + primaryFD: Int32, + primaryHandle: FileHandle, + secondaryHandle: FileHandle) + { + self.processPID = pid + self.processGroupID = processGroup + self.primaryFD = primaryFD + self.primaryHandle = primaryHandle + self.secondaryHandle = secondaryHandle + } + + var pid: pid_t { + self.processPID + } + + var isRunning: Bool { + self.lock.lock() + if self.reaped { + self.lock.unlock() + return false + } + self.lock.unlock() + + var status: Int32 = 0 + let result = waitpid(self.processPID, &status, WNOHANG) + + self.lock.lock() + defer { self.lock.unlock() } + switch result { + case 0: + return true + case self.processPID: + self.reaped = true + return false + case -1 where errno == ECHILD: + self.reaped = true + return false + default: + return kill(self.processPID, 0) == 0 || errno == EPERM + } + } + + var processGroup: pid_t? { + self.processGroupID + } + + func assignProcessGroup() -> pid_t? { + self.processGroupID + } + + func sendExit() throws { + try self.writeAllToPrimary(Data("/exit\r".utf8)) + } + + func closePTY() { + try? self.primaryHandle.close() + try? self.secondaryHandle.close() + } + + func terminateRoot() { + kill(self.processPID, SIGTERM) + } + + func killRoot() { + kill(self.processPID, SIGKILL) + } + + func descendantPIDs() -> [pid_t] { + TTYProcessTreeTerminator.descendantPIDs(of: self.processPID) + } + + func terminateTree(signal: Int32, knownDescendants: [pid_t]) { + TTYProcessTreeTerminator.terminateProcessTree( + rootPID: self.processPID, + processGroup: self.processGroupID, + signal: signal, + knownDescendants: knownDescendants) + } + + func killDescendants(_ descendants: [pid_t]) { + for pid in descendants where pid > 0 { + kill(pid, SIGKILL) + } + } + + func drainOutput() -> Data { + var tmp = [UInt8](repeating: 0, count: 8192) + var output: [UInt8] = [] + for _ in 0..<64 { + let n = read(self.primaryFD, &tmp, tmp.count) + if n > 0 { + output.append(contentsOf: tmp.prefix(n)) + continue + } + break + } + return Data(output) + } + + private func writeAllToPrimary(_ data: Data) throws { + data.withUnsafeBytes { rawBytes in + guard let baseAddress = rawBytes.baseAddress else { return } + var offset = 0 + var retries = 0 + while offset < rawBytes.count { + let written = write(self.primaryFD, baseAddress.advanced(by: offset), rawBytes.count - offset) + if written > 0 { + offset += written + retries = 0 + continue + } + if written == 0 { break } + + let err = errno + if err == EINTR || err == EAGAIN || err == EWOULDBLOCK { + retries += 1 + if retries > 200 { return } + usleep(5000) + continue + } + return + } + } + } +} + +// MARK: - Production Stale Session Identity + Storage + +struct AntigravityProcessIdentityProvider: AntigravityCLIProcessIdentityProviding { + static var currentUserID: UInt32 { + UInt32(getuid()) + } + + func ownerUserID(for pid: pid_t) -> UInt32? { + #if canImport(Darwin) + var info = proc_bsdinfo() + let size = proc_pidinfo( + pid, + PROC_PIDTBSDINFO, + 0, + &info, + Int32(MemoryLayout.stride)) + guard size == Int32(MemoryLayout.stride) else { return nil } + return info.pbi_uid + #else + guard let status = try? String(contentsOfFile: "/proc/\(pid)/status", encoding: .utf8), + let uidLine = status.split(separator: "\n").first(where: { $0.hasPrefix("Uid:") }), + let owner = uidLine.split(whereSeparator: \.isWhitespace).dropFirst().first, + let userID = UInt32(owner) + else { + return nil + } + return userID + #endif + } + + func identity(for pid: pid_t) -> AntigravityCLIProcessIdentity? { + #if canImport(Darwin) + var pathBuffer = [CChar](repeating: 0, count: 4096) + let pathLength = proc_pidpath(pid, &pathBuffer, UInt32(pathBuffer.count)) + guard pathLength > 0 else { return nil } + let executablePath = pathBuffer.withUnsafeBufferPointer { buffer -> String? in + let rawBytes = UnsafeRawBufferPointer(start: buffer.baseAddress, count: Int(pathLength)) + return String(bytes: rawBytes.prefix { $0 != 0 }, encoding: .utf8) + } + guard let executablePath, !executablePath.isEmpty else { return nil } + + var info = proc_bsdinfo() + let size = proc_pidinfo( + pid, + PROC_PIDTBSDINFO, + 0, + &info, + Int32(MemoryLayout.stride)) + guard size == Int32(MemoryLayout.stride) else { return nil } + let startEpoch = TimeInterval(info.pbi_start_tvsec) + (TimeInterval(info.pbi_start_tvusec) / 1_000_000) + return AntigravityCLIProcessIdentity(executablePath: executablePath, startEpoch: startEpoch) + #else + let procDirectory = "/proc/\(pid)" + guard let executablePath = try? FileManager.default.destinationOfSymbolicLink( + atPath: "\(procDirectory)/exe"), + let stat = try? String(contentsOfFile: "\(procDirectory)/stat", encoding: .utf8), + let closeParen = stat.lastIndex(of: ")") + else { + return nil + } + let fields = stat[stat.index(after: closeParen)...].split(whereSeparator: \.isWhitespace) + let clockTicksPerSecond = sysconf(Int32(_SC_CLK_TCK)) + let systemStat = try? String(contentsOfFile: "/proc/stat", encoding: .utf8) + let bootEpoch = systemStat? + .split(separator: "\n") + .first { $0.hasPrefix("btime ") }? + .split(whereSeparator: \.isWhitespace) + .dropFirst() + .first + .flatMap { TimeInterval($0) } + guard fields.count > 19, + let startTicks = TimeInterval(fields[19]), + clockTicksPerSecond > 0, + let bootEpoch + else { + return nil + } + let startEpoch = bootEpoch + (startTicks / TimeInterval(clockTicksPerSecond)) + return AntigravityCLIProcessIdentity(executablePath: executablePath, startEpoch: startEpoch) + #endif + } +} + +final class AntigravityFileCLISessionRecordStore: AntigravityCLISessionRecordStoring, @unchecked Sendable { + private let fileURL: URL + private let fileManager: FileManager + + init( + fileURL: URL = URL(fileURLWithPath: NSHomeDirectory(), isDirectory: true) + .appendingPathComponent(".codexbar", isDirectory: true) + .appendingPathComponent("antigravity", isDirectory: true) + .appendingPathComponent("agy-session.json"), + fileManager: FileManager = .default) + { + self.fileURL = fileURL + self.fileManager = fileManager + } + + func load() throws -> [AntigravityCLISessionRecord] { + guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return [] } + let data = try Data(contentsOf: self.fileURL) + let decoder = JSONDecoder() + if let records = try? decoder.decode([AntigravityCLISessionRecord].self, from: data) { + return records + } + return try [decoder.decode(AntigravityCLISessionRecord.self, from: data)] + } + + func save(_ record: AntigravityCLISessionRecord) throws { + var records = (try? self.load()) ?? [] + records.removeAll { Self.sameOwner($0, record) } + records.append(record) + try self.write(records) + } + + func remove(_ record: AntigravityCLISessionRecord) throws { + var records = try self.load() + records.removeAll { + $0.pid == record.pid && + $0.executablePath == record.executablePath && + abs($0.startEpoch - record.startEpoch) < 0.001 + } + if records.isEmpty { + guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return } + try self.fileManager.removeItem(at: self.fileURL) + } else { + try self.write(records) + } + } + + private func write(_ records: [AntigravityCLISessionRecord]) throws { + let directory = self.fileURL.deletingLastPathComponent() + try self.fileManager.createDirectory(at: directory, withIntermediateDirectories: true) + let data = try JSONEncoder().encode(records) + try data.write(to: self.fileURL, options: [.atomic]) + } + + private static func sameOwner( + _ lhs: AntigravityCLISessionRecord, + _ rhs: AntigravityCLISessionRecord) -> Bool + { + if let lhsOwnerPID = lhs.ownerPID, + let rhsOwnerPID = rhs.ownerPID, + let lhsOwnerPath = lhs.ownerExecutablePath, + let rhsOwnerPath = rhs.ownerExecutablePath, + let lhsOwnerStart = lhs.ownerStartEpoch, + let rhsOwnerStart = rhs.ownerStartEpoch + { + return lhsOwnerPID == rhsOwnerPID && + lhsOwnerPath == rhsOwnerPath && + abs(lhsOwnerStart - rhsOwnerStart) < 0.001 + } + return lhs.pid == rhs.pid && + lhs.executablePath == rhs.executablePath && + abs(lhs.startEpoch - rhs.startEpoch) < 0.001 + } +} + +final class AntigravityFileCLISessionLaunchLock: AntigravityCLISessionLaunchLocking, @unchecked Sendable { + private let fileURL: URL + private let fileManager: FileManager + + init( + fileURL: URL = URL(fileURLWithPath: NSHomeDirectory(), isDirectory: true) + .appendingPathComponent(".codexbar", isDirectory: true) + .appendingPathComponent("antigravity", isDirectory: true) + .appendingPathComponent("agy-session.lock"), + fileManager: FileManager = .default) + { + self.fileURL = fileURL + self.fileManager = fileManager + } + + func withLock(_ operation: () throws -> T) throws -> T { + let directory = self.fileURL.deletingLastPathComponent() + try self.fileManager.createDirectory(at: directory, withIntermediateDirectories: true) + let fd = open(self.fileURL.path, O_CREAT | O_RDWR | O_CLOEXEC, S_IRUSR | S_IWUSR) + guard fd >= 0 else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + defer { + _ = flock(fd, LOCK_UN) + close(fd) + } + + while flock(fd, LOCK_EX) != 0 { + guard errno == EINTR else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + } + return try operation() + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityLocalSnapshotSelection.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityLocalSnapshotSelection.swift new file mode 100644 index 0000000..83785a1 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityLocalSnapshotSelection.swift @@ -0,0 +1,37 @@ +import Foundation + +extension AntigravityStatusProbe { + static func preferredLocalSnapshot( + _ snapshots: [AntigravityStatusSnapshot], + matchingAccountEmail expectedAccountEmail: String?) -> AntigravityStatusSnapshot? + { + var candidates = snapshots + if let expected = self.normalizedAccountEmail(expectedAccountEmail) { + let matches = snapshots.filter { + guard let found = self.normalizedAccountEmail($0.accountEmail) else { return false } + return found.caseInsensitiveCompare(expected) == .orderedSame + } + if !matches.isEmpty { + candidates = matches + } + } + + var bestSnapshot: AntigravityStatusSnapshot? + var bestScore = Int.min + for snapshot in candidates { + let score = self.localSnapshotScore(snapshot) + if score > bestScore { + bestSnapshot = snapshot + bestScore = score + } + } + return bestSnapshot + } + + private static func normalizedAccountEmail(_ email: String?) -> String? { + guard let trimmed = email?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityLocalhostSession.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityLocalhostSession.swift new file mode 100644 index 0000000..5a372ea --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityLocalhostSession.swift @@ -0,0 +1,109 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +enum LocalhostTrustPolicy { + static func shouldAcceptServerTrust( + host: String, + authenticationMethod: String, + hasServerTrust: Bool) -> Bool + { + #if !os(Linux) + guard authenticationMethod == NSURLAuthenticationMethodServerTrust else { return false } + #endif + let normalizedHost = host.lowercased() + guard normalizedHost == "127.0.0.1" || normalizedHost == "localhost" else { return false } + return hasServerTrust + } +} + +final class LocalhostSessionDelegate: NSObject { + func data(for request: URLRequest, session: URLSession) async throws -> (Data, URLResponse) { + let state = LocalhostSessionTaskState() + return try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + let task = session.dataTask(with: request) { data, response, error in + if let error { + continuation.resume(throwing: error) + return + } + guard let data, let response else { + continuation.resume(throwing: AntigravityStatusProbeError.apiError("Invalid response")) + return + } + continuation.resume(returning: (data, response)) + } + state.setTask(task) + task.resume() + } + } onCancel: { + state.cancel() + } + } + + private func challengeResult(_ challenge: URLAuthenticationChallenge) -> ( + disposition: URLSession.AuthChallengeDisposition, + credential: URLCredential?) + { + #if os(Linux) + return (.performDefaultHandling, nil) + #else + let protectionSpace = challenge.protectionSpace + let trust = protectionSpace.serverTrust + guard LocalhostTrustPolicy.shouldAcceptServerTrust( + host: protectionSpace.host, + authenticationMethod: protectionSpace.authenticationMethod, + hasServerTrust: trust != nil), + let trust + else { + return (.performDefaultHandling, nil) + } + return (.useCredential, URLCredential(trust: trust)) + #endif + } +} + +extension LocalhostSessionDelegate: URLSessionDelegate { + func urlSession( + _ session: URLSession, + didReceive challenge: URLAuthenticationChallenge) async -> (URLSession.AuthChallengeDisposition, URLCredential?) + { + self.challengeResult(challenge) + } +} + +extension LocalhostSessionDelegate: URLSessionTaskDelegate { + func urlSession( + _ session: URLSession, + task: URLSessionTask, + didReceive challenge: URLAuthenticationChallenge) async -> (URLSession.AuthChallengeDisposition, URLCredential?) + { + self.challengeResult(challenge) + } +} + +private final class LocalhostSessionTaskState: @unchecked Sendable { + private let lock = NSLock() + private var task: URLSessionDataTask? + private var isCancelled = false + + func setTask(_ task: URLSessionDataTask) { + self.lock.lock() + self.task = task + let shouldCancel = self.isCancelled + self.lock.unlock() + + if shouldCancel { + task.cancel() + } + } + + func cancel() { + self.lock.lock() + self.isCancelled = true + let task = self.task + self.lock.unlock() + task?.cancel() + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuthCredentialsStore.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuthCredentialsStore.swift new file mode 100644 index 0000000..06819e0 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityOAuthCredentialsStore.swift @@ -0,0 +1,516 @@ +import Foundation + +public struct AntigravityOAuthCredentials: Codable, Sendable, Equatable { + public var accessToken: String? + public var refreshToken: String? + public var expiryDateMilliseconds: Double? + public var idToken: String? + public var email: String? + public var projectID: String? + public var clientID: String? + public var clientSecret: String? + + public init( + accessToken: String?, + refreshToken: String?, + expiryDate: Date?, + idToken: String? = nil, + email: String? = nil, + projectID: String? = nil, + clientID: String? = nil, + clientSecret: String? = nil) + { + self.accessToken = accessToken + self.refreshToken = refreshToken + self.expiryDateMilliseconds = expiryDate.map { $0.timeIntervalSince1970 * 1000 } + self.idToken = idToken + self.email = email + self.projectID = projectID + self.clientID = clientID + self.clientSecret = clientSecret + } + + public var expiryDate: Date? { + guard let expiryDateMilliseconds else { return nil } + return Date(timeIntervalSince1970: expiryDateMilliseconds / 1000) + } + + /// Email of the Google account these credentials authenticate, preferring the + /// signed `id_token` claim (what the remote OAuth fetcher reports) and falling + /// back to the stored `email` field. Used to verify that an ambient local/CLI + /// Antigravity snapshot belongs to the account the user explicitly selected. + public var resolvedAccountEmail: String? { + Self.email(fromIDToken: self.idToken) ?? self.email?.trimmedNonEmptyEmail + } + + static func email(fromIDToken idToken: String?) -> String? { + guard let idToken else { return nil } + let parts = idToken.components(separatedBy: ".") + guard parts.count >= 2 else { return nil } + var payload = parts[1] + .replacingOccurrences(of: "-", with: "+") + .replacingOccurrences(of: "_", with: "/") + let remainder = payload.count % 4 + if remainder > 0 { + payload += String(repeating: "=", count: 4 - remainder) + } + guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return nil + } + return (json["email"] as? String)?.trimmedNonEmptyEmail + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.accessToken = + try container.decodeIfPresent(String.self, forKey: .accessTokenSnake) + ?? container.decodeIfPresent(String.self, forKey: .accessTokenCamel) + self.refreshToken = + try container.decodeIfPresent(String.self, forKey: .refreshTokenSnake) + ?? container.decodeIfPresent(String.self, forKey: .refreshTokenCamel) + self.idToken = + try container.decodeIfPresent(String.self, forKey: .idTokenSnake) + ?? container.decodeIfPresent(String.self, forKey: .idTokenCamel) + self.email = try container.decodeIfPresent(String.self, forKey: .email) + self.projectID = + try container.decodeIfPresent(String.self, forKey: .projectIDSnake) + ?? container.decodeIfPresent(String.self, forKey: .projectIDCamel) + self.clientID = + try container.decodeIfPresent(String.self, forKey: .clientIDSnake) + ?? container.decodeIfPresent(String.self, forKey: .clientIDCamel) + self.clientSecret = + try container.decodeIfPresent(String.self, forKey: .clientSecretSnake) + ?? container.decodeIfPresent(String.self, forKey: .clientSecretCamel) + + if let expiryDateMilliseconds = try container.decodeIfPresent(Double.self, forKey: .expiryDateSnake) + ?? container.decodeIfPresent(Double.self, forKey: .expiresAtCamel) + { + self.expiryDateMilliseconds = expiryDateMilliseconds + } else if let expiryDateMilliseconds = try container.decodeIfPresent(Int.self, forKey: .expiryDateSnake) + ?? container.decodeIfPresent(Int.self, forKey: .expiresAtCamel) + { + self.expiryDateMilliseconds = Double(expiryDateMilliseconds) + } else { + self.expiryDateMilliseconds = nil + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encodeIfPresent(self.accessToken, forKey: .accessTokenSnake) + try container.encodeIfPresent(self.refreshToken, forKey: .refreshTokenSnake) + try container.encodeIfPresent(self.expiryDateMilliseconds, forKey: .expiryDateSnake) + try container.encodeIfPresent(self.idToken, forKey: .idTokenSnake) + try container.encodeIfPresent(self.email, forKey: .email) + try container.encodeIfPresent(self.projectID, forKey: .projectIDSnake) + try container.encodeIfPresent(self.clientID, forKey: .clientIDSnake) + try container.encodeIfPresent(self.clientSecret, forKey: .clientSecretSnake) + } + + enum CodingKeys: String, CodingKey { + case accessTokenSnake = "access_token" + case accessTokenCamel = "accessToken" + case refreshTokenSnake = "refresh_token" + case refreshTokenCamel = "refreshToken" + case expiryDateSnake = "expiry_date" + case expiresAtCamel = "expiresAt" + case idTokenSnake = "id_token" + case idTokenCamel = "idToken" + case email + case projectIDSnake = "project_id" + case projectIDCamel = "projectId" + case clientIDSnake = "client_id" + case clientIDCamel = "clientId" + case clientSecretSnake = "client_secret" + case clientSecretCamel = "clientSecret" + } +} + +public struct AntigravityOAuthClient: Sendable, Equatable { + public let clientID: String + public let clientSecret: String + + public init(clientID: String, clientSecret: String) { + self.clientID = clientID + self.clientSecret = clientSecret + } +} + +public enum AntigravityOAuthConfig { + public static var configuredClientID: String? { + let value = ProcessInfo.processInfo.environment["ANTIGRAVITY_OAUTH_CLIENT_ID"] + return value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty + } + + public static var configuredClientSecret: String? { + let value = ProcessInfo.processInfo.environment["ANTIGRAVITY_OAUTH_CLIENT_SECRET"] + return value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty + } + + public static let authURL = URL(string: "https://accounts.google.com/o/oauth2/v2/auth")! + public static let tokenURL = URL(string: "https://oauth2.googleapis.com/token")! + public static let userInfoURL = URL(string: "https://www.googleapis.com/oauth2/v2/userinfo")! + public static let scopes = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/userinfo.email", + ] + + public static let missingCredentialsMessage = + """ + Antigravity OAuth client is not configured. Install Antigravity.app or set \ + ANTIGRAVITY_OAUTH_CLIENT_ID and ANTIGRAVITY_OAUTH_CLIENT_SECRET before logging in. + """ + + public static func resolvedClient() -> AntigravityOAuthClient? { + if let client = environmentClient() { + return client + } + return Self.discoverClientFromInstalledApp() + } + + private static func environmentClient() -> AntigravityOAuthClient? { + guard let clientID = configuredClientID, + let clientSecret = configuredClientSecret + else { + return nil + } + return AntigravityOAuthClient(clientID: clientID, clientSecret: clientSecret) + } + + static func discoverClientFromInstalledApp( + applicationRoots: [URL]? = nil, + fileManager: FileManager = .default) -> AntigravityOAuthClient? + { + for url in self.candidateOAuthClientArtifactURLs( + applicationRoots: applicationRoots, + fileManager: fileManager) + where fileManager.fileExists(atPath: url.path) + { + guard let data = try? Data(contentsOf: url), + let client = Self.parseClient(fromInstalledArtifactData: data) + else { + continue + } + return client + } + return nil + } + + static func candidateOAuthClientArtifactURLs( + applicationRoots: [URL]? = nil, + fileManager: FileManager = .default) -> [URL] + { + let roots = [ + URL(fileURLWithPath: "/Applications", isDirectory: true), + fileManager.homeDirectoryForCurrentUser.appendingPathComponent("Applications", isDirectory: true), + ] + let applicationRoots = applicationRoots ?? roots + let appBundleURLs = self.candidateAntigravityAppBundleURLs( + applicationRoots: applicationRoots, + fileManager: fileManager) + let relativePaths = [ + "Contents/Resources/app/extensions/antigravity/bin/language_server_macos_arm", + "Contents/Resources/app/extensions/antigravity/bin/language_server_macos_x64", + "Contents/Resources/app/extensions/antigravity/bin/language_server_macos", + "Contents/Resources/app/out/main.js", + "Contents/Resources/bin/language_server", + "Contents/Resources/bin/language_server_macos", + ] + return appBundleURLs.flatMap { bundleURL in + relativePaths.map { bundleURL.appendingPathComponent($0) } + } + } + + private static func candidateAntigravityAppBundleURLs( + applicationRoots: [URL], + fileManager: FileManager) -> [URL] + { + var urls: [URL] = [] + + for root in applicationRoots { + urls.append(root.appendingPathComponent("Antigravity.app", isDirectory: true)) + + let appURLs = (try? fileManager.contentsOfDirectory( + at: root, + includingPropertiesForKeys: [.isDirectoryKey], + options: [.skipsHiddenFiles])) ?? [] + for appURL in appURLs where appURL.pathExtension == "app" { + guard self.isAntigravityAppBundle(appURL) else { continue } + urls.append(appURL) + } + } + + var seen = Set() + return urls.filter { url in + let key = url.standardizedFileURL.path + guard !seen.contains(key) else { return false } + seen.insert(key) + return true + } + } + + private static func isAntigravityAppBundle(_ url: URL) -> Bool { + switch Bundle(url: url)?.bundleIdentifier { + case "com.google.antigravity", "com.google.antigravity-ide": + true + default: + false + } + } + + static func parseClient(fromInstalledArtifactData data: Data) -> AntigravityOAuthClient? { + if let content = String(data: data, encoding: .utf8), + let client = parseClient(fromInstalledArtifactText: content) + { + return client + } + + let clientIDs = Self.clientIDs(in: data) + let clientSecrets = Self.clientSecrets(in: data) + guard let client = Self.preferredBinaryClient( + clientIDs: clientIDs, + clientSecrets: clientSecrets) + else { + return nil + } + + return client + } + + static func parseClient(fromInstalledArtifactText content: String) -> AntigravityOAuthClient? { + let marker = "vs/platform/cloudCode/common/oauthClient.js" + let searchStart = content.range(of: marker)?.lowerBound ?? content.startIndex + let searchEnd = content.index(searchStart, offsetBy: 4000, limitedBy: content.endIndex) ?? content.endIndex + let haystack = String(content[searchStart.. [String] { + let suffix = Data(".apps.googleusercontent.com".utf8) + var searchRange = data.startIndex.. data.startIndex { + let previous = data.index(before: start) + guard Self.isOAuthClientIDPrefixByte(data[previous]) else { break } + start = previous + } + + let candidateData = Data(data[start.. [String] { + let prefix = Data("GOCSPX-".utf8) + let secretLength = 35 + var searchRange = data.startIndex.. AntigravityOAuthClient? + { + guard !clientIDs.isEmpty, + !clientSecrets.isEmpty + else { + return nil + } + + if clientSecrets.count == 1, clientIDs.count > 1 { + return AntigravityOAuthClient(clientID: clientIDs[clientIDs.count - 1], clientSecret: clientSecrets[0]) + } + + let clientSecret: String = if clientSecrets.count == clientIDs.count, clientSecrets.count > 1 { + // Antigravity 2's language_server binary stores the secret table before the client id table. + clientSecrets[clientSecrets.count - 1] + } else { + clientSecrets[0] + } + + return AntigravityOAuthClient(clientID: clientIDs[0], clientSecret: clientSecret) + } + + private static func isOAuthClientIDPrefixByte(_ byte: UInt8) -> Bool { + (byte >= 48 && byte <= 57) + || (byte >= 65 && byte <= 90) + || (byte >= 97 && byte <= 122) + || byte == 45 + || byte == 95 + } + + private static func isOAuthClientSecretByte(_ byte: UInt8) -> Bool { + self.isOAuthClientIDPrefixByte(byte) + } + + private static func firstMatch(pattern: String, in text: String) -> String? { + guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil } + let range = NSRange(text.startIndex.. [String] { + var seen = Set() + return values.filter { value in + guard !seen.contains(value) else { return false } + seen.insert(value) + return true + } + } +} + +public struct AntigravityOAuthCredentialsStore: @unchecked Sendable { + public static let environmentCredentialsKey = "ANTIGRAVITY_OAUTH_CREDENTIALS_JSON" + private static let fileLock = NSLock() + + public let fileURL: URL + private let fileManager: FileManager + + public init(fileURL: URL = Self.defaultURL(), fileManager: FileManager = .default) { + self.fileURL = fileURL + self.fileManager = fileManager + } + + public func load() throws -> AntigravityOAuthCredentials? { + try Self.fileLock.withLock { + try self.loadUnlocked() + } + } + + private func loadUnlocked() throws -> AntigravityOAuthCredentials? { + guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return nil } + let data = try Data(contentsOf: self.fileURL) + return try JSONDecoder().decode(AntigravityOAuthCredentials.self, from: data) + } + + public func save(_ credentials: AntigravityOAuthCredentials) throws { + try Self.fileLock.withLock { + let data = try JSONEncoder.antigravityCredentials.encode(credentials) + let directory = self.fileURL.deletingLastPathComponent() + if !self.fileManager.fileExists(atPath: directory.path) { + try self.fileManager.createDirectory(at: directory, withIntermediateDirectories: true) + } + try data.write(to: self.fileURL, options: [.atomic]) + try self.applySecurePermissionsIfNeeded() + } + } + + public func deleteIfPresent() throws { + try Self.fileLock.withLock { + try self.deleteIfPresentUnlocked() + } + } + + public func deleteIfPresent( + matching predicate: (AntigravityOAuthCredentials) -> Bool) throws + { + try Self.fileLock.withLock { + guard let credentials = try self.loadUnlocked(), predicate(credentials) else { return } + try self.deleteIfPresentUnlocked() + } + } + + private func deleteIfPresentUnlocked() throws { + guard self.fileManager.fileExists(atPath: self.fileURL.path) else { return } + try self.fileManager.removeItem(at: self.fileURL) + } + + public static func defaultDirectoryURL(home: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL { + home + .appendingPathComponent(".codexbar", isDirectory: true) + .appendingPathComponent("antigravity", isDirectory: true) + } + + public static func defaultURL(home: URL = FileManager.default.homeDirectoryForCurrentUser) -> URL { + self.defaultDirectoryURL(home: home) + .appendingPathComponent("oauth_creds.json") + } + + public static func tokenAccountValue(for credentials: AntigravityOAuthCredentials) throws -> String { + let data = try JSONEncoder.antigravityCredentials.encode(credentials) + guard let value = String(data: data, encoding: .utf8) else { + throw CocoaError(.coderInvalidValue) + } + return value + } + + public static func credentials(fromTokenAccountValue value: String) -> AntigravityOAuthCredentials? { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty, let data = trimmed.data(using: .utf8) else { return nil } + return try? JSONDecoder().decode(AntigravityOAuthCredentials.self, from: data) + } + + private func applySecurePermissionsIfNeeded() throws { + #if os(macOS) || os(Linux) + try self.fileManager.setAttributes([ + .posixPermissions: NSNumber(value: Int16(0o600)), + ], ofItemAtPath: self.fileURL.path) + #endif + } +} + +extension JSONEncoder { + fileprivate static let antigravityCredentials: JSONEncoder = { + let encoder = JSONEncoder() + encoder.outputFormatting = [.prettyPrinted, .sortedKeys] + return encoder + }() +} + +extension String { + fileprivate var nilIfEmpty: String? { + self.isEmpty ? nil : self + } + + fileprivate var trimmedNonEmptyEmail: String? { + let trimmed = self.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift new file mode 100644 index 0000000..10b89cd --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityProviderDescriptor.swift @@ -0,0 +1,673 @@ +import Foundation + +public enum AntigravityProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .antigravity, + metadata: ProviderMetadata( + id: .antigravity, + displayName: "Antigravity", + sessionLabel: "Gemini Models", + weeklyLabel: "Claude and GPT", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Antigravity usage (experimental)", + cliName: "antigravity", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + dashboardURL: nil, + statusPageURL: nil, + statusLinkURL: "https://www.google.com/appsstatus/dashboard/products/npdyhgECDJ6tB66MxXyo/history", + statusWorkspaceProductID: "npdyhgECDJ6tB66MxXyo"), + branding: ProviderBranding( + iconStyle: .antigravity, + iconResourceName: "ProviderIcon-antigravity", + color: ProviderColor(red: 96 / 255, green: 186 / 255, blue: 126 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Antigravity cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .cli, .oauth], + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), + cli: ProviderCLIConfig( + name: "antigravity", + versionDetector: nil)) + } + + private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] { + let app = AntigravityStatusFetchStrategy(source: .app) + let cli = AntigravityCLIHTTPSFetchStrategy() + let ide = AntigravityStatusFetchStrategy(source: .ide) + let oauth = AntigravityOAuthFetchStrategy() + switch context.sourceMode { + case .cli: + return [app, cli, ide] + case .oauth: + return [oauth] + case .auto: + if context.selectedTokenAccountID != nil || + context.env[AntigravityOAuthCredentialsStore.environmentCredentialsKey] != nil || + self.hasSharedOAuthCredentials(context: context) + { + return [app, cli, ide, oauth] + } + return [app, cli, ide] + case .web, .api: + return [] + } + } + + private static func hasSharedOAuthCredentials(context: ProviderFetchContext) -> Bool { + let homeURL = context.env["HOME"] + .flatMap { $0.isEmpty ? nil : URL(fileURLWithPath: $0, isDirectory: true) } + ?? FileManager.default.homeDirectoryForCurrentUser + let fileURL = AntigravityOAuthCredentialsStore.defaultURL(home: homeURL) + return FileManager.default.fileExists(atPath: fileURL.path) + } +} + +struct AntigravityStatusFetchStrategy: ProviderFetchStrategy { + enum Source { + case app + case ide + + var id: String { + switch self { + case .app: "antigravity.app-local" + case .ide: "antigravity.ide-local" + } + } + + var processScope: AntigravityStatusProbe.ProcessScope { + switch self { + case .app: .appOnly + case .ide: .ideOnly + } + } + + var sourceLabel: String { + switch self { + case .app: "app" + case .ide: "ide" + } + } + } + + let source: Source + var id: String { + self.source.id + } + + let kind: ProviderFetchKind = .localProbe + + init(source: Source = .app) { + self.source = source + } + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let probe = AntigravityStatusProbe(processScope: self.source.processScope) + let selectedAccountEmail: String? = if context.sourceMode == .auto, context.selectedTokenAccountID != nil { + AntigravitySelectedAccountGuard.selectedAccountEmail(context: context) + } else { + nil + } + let snap = try await probe.fetch(matchingAccountEmail: selectedAccountEmail) + let usage = try snap.toUsageSnapshot() + try AntigravitySelectedAccountGuard.validate(usage, context: context) + return self.makeResult( + usage: usage, + sourceLabel: self.source.sourceLabel) + } + + func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool { + context.sourceMode == .auto || context.sourceMode == .cli + } +} + +/// When the Antigravity 2.0 app is closed or unavailable, this strategy spawns +/// or reuses ``agy`` and talks to the localhost server embedded in that CLI +/// process. ``agy`` is an interactive REPL, not a query command, so CodexBar +/// never scrapes TUI output here; it only keeps the process alive long enough +/// for the server to answer quota endpoints. +struct AntigravityCLIHTTPSFetchStrategy: ProviderFetchStrategy { + static let sourceLabel = "cli" + let id: String = "antigravity.cli-https" + let kind: ProviderFetchKind = .cli + private static let log = CodexBarLog.logger(LogCategories.antigravity) + + struct SnapshotWaitDependencies { + let pollIntervalNanoseconds: UInt64 + let listeningPorts: @Sendable (Int, TimeInterval) async throws -> [Int] + let drainOutput: @Sendable () async -> Data + let fetchSnapshot: @Sendable ([Int]) async throws -> AntigravityStatusSnapshot + let now: @Sendable () -> Date + + init( + pollIntervalNanoseconds: UInt64, + listeningPorts: @escaping @Sendable (Int, TimeInterval) async throws -> [Int], + drainOutput: @escaping @Sendable () async -> Data, + fetchSnapshot: @escaping @Sendable ([Int]) async throws -> AntigravityStatusSnapshot, + now: @escaping @Sendable () -> Date = Date.init) + { + self.pollIntervalNanoseconds = pollIntervalNanoseconds + self.listeningPorts = listeningPorts + self.drainOutput = drainOutput + self.fetchSnapshot = fetchSnapshot + self.now = now + } + } + + /// Seams for discovering and reusing an already-running ``agy`` CLI language + /// server, so a fresh spawn (and its multi-second ``GetUserStatus`` warm-up) + /// can be skipped when a warm server is already present. + struct WarmAgyDependencies { + let processInfos: @Sendable (TimeInterval) async throws -> [AntigravityStatusProbe.ProcessInfoResult] + let listeningPorts: @Sendable (Int, TimeInterval) async throws -> [Int] + let fetchSnapshot: @Sendable ([Int], TimeInterval) async throws -> AntigravityStatusSnapshot + let processOwnerUserID: @Sendable (Int) -> UInt32? + let currentUserID: @Sendable () -> UInt32 + /// The pid of an ``agy`` that CodexBar itself spawned and manages through + /// ``AntigravityCLISession`` (if any). Such a process must NOT be reused + /// through the warm path: doing so bypasses `beginProbe`/`finishProbe`, so + /// the idle timer is never cancelled/extended and `stopIfIdle` could tear + /// the managed session down mid-poll. Externally owned `agy` (an IDE, a + /// long-lived `agy`, or another CodexBar host) has no such accounting. + let ownedPID: @Sendable () async -> Int? + let now: @Sendable () -> Date + + init( + processInfos: @escaping @Sendable (TimeInterval) async throws + -> [AntigravityStatusProbe.ProcessInfoResult], + listeningPorts: @escaping @Sendable (Int, TimeInterval) async throws -> [Int], + fetchSnapshot: @escaping @Sendable ([Int], TimeInterval) async throws -> AntigravityStatusSnapshot, + processOwnerUserID: @escaping @Sendable (Int) -> UInt32? = { _ in 0 }, + currentUserID: @escaping @Sendable () -> UInt32 = { 0 }, + ownedPID: @escaping @Sendable () async -> Int? = { nil }, + now: @escaping @Sendable () -> Date = Date.init) + { + self.processInfos = processInfos + self.listeningPorts = listeningPorts + self.fetchSnapshot = fetchSnapshot + self.processOwnerUserID = processOwnerUserID + self.currentUserID = currentUserID + self.ownedPID = ownedPID + self.now = now + } + } + + /// Discover an already-running, authenticated ``agy`` CLI language server and + /// reuse its listening ports instead of spawning a fresh process. + /// + /// One-shot CLI invocations otherwise spawn a brand-new ``agy`` on every + /// call; a fresh server binds its port quickly but ``GetUserStatus`` returns + /// transient initialization failures for a few seconds, so the readiness + /// deadline is occasionally missed. When a warm CLI server is already up, we + /// can talk to it immediately — it needs no CSRF token (``cliHTTPS``). + /// + /// Returns the snapshot from the first warm server that answers with + /// parseable usage for the requested account, or `nil` when none is found or + /// none answers — in which case the caller falls back to the existing spawn + /// path unchanged. + static func tryWarmAgyFetch( + timeout: TimeInterval, + expectedBinaryPath: String? = nil, + expectedAccountEmail: String? = nil, + dependencies: WarmAgyDependencies) async throws -> AntigravityStatusSnapshot? + { + try Task.checkCancellation() + let deadline = dependencies.now().addingTimeInterval(timeout) + guard let discoveryTimeout = Self.remainingWarmProbeTime(deadline: deadline, now: dependencies.now) else { + return nil + } + let processInfos: [AntigravityStatusProbe.ProcessInfoResult] + do { + processInfos = try await dependencies.processInfos(discoveryTimeout) + } catch let error as CancellationError { + throw error + } catch { + return nil + } + try Task.checkCancellation() + let ownedPID = await dependencies.ownedPID() + try Task.checkCancellation() + let currentUserID = dependencies.currentUserID() + // Only the CLI's language server needs no CSRF token; the IDE/app servers + // require one and must not be reused through this token-less fast path. + // Also exclude any `agy` CodexBar itself spawned and manages: reusing it + // here would bypass session lifecycle accounting (see `ownedPID`). + let cliProcesses = processInfos.filter { info in + info.pid != ownedPID && + dependencies.processOwnerUserID(info.pid) == currentUserID && + AntigravityStatusProbe.antigravityProcessKind(info.commandLine) == .cli + } + guard !cliProcesses.isEmpty else { return nil } + + for info in cliProcesses { + if let expectedBinaryPath { + guard Self.commandLine(info.commandLine, matchesBinaryPath: expectedBinaryPath) + else { + continue + } + } + guard let portTimeout = Self.remainingWarmProbeTime(deadline: deadline, now: dependencies.now) else { + return nil + } + let ports: [Int] + do { + ports = try await dependencies.listeningPorts(info.pid, portTimeout) + } catch let error as CancellationError { + throw error + } catch { + continue + } + try Task.checkCancellation() + guard !ports.isEmpty else { continue } + guard let fetchTimeout = Self.remainingWarmProbeTime(deadline: deadline, now: dependencies.now) else { + return nil + } + let snapshot: AntigravityStatusSnapshot + do { + snapshot = try await dependencies.fetchSnapshot(ports, fetchTimeout) + } catch let error as CancellationError { + throw error + } catch { + continue + } + try Task.checkCancellation() + guard (try? snapshot.toUsageSnapshot()) != nil, + AntigravitySelectedAccountGuard.matches( + snapshotAccountEmail: snapshot.accountEmail, + expectedAccountEmail: expectedAccountEmail) + else { + continue + } + Self.log.debug("Antigravity CLI HTTPS reusing warm agy", metadata: [ + "pid": "\(info.pid)", + "ports": ports.map(String.init).joined(separator: ","), + ]) + return snapshot + } + try Task.checkCancellation() + return nil + } + + private static func remainingWarmProbeTime( + deadline: Date, + now: @Sendable () -> Date) -> TimeInterval? + { + let remaining = deadline.timeIntervalSince(now()) + return remaining > 0 ? remaining : nil + } + + private static func commandLine(_ commandLine: String, matchesBinaryPath binaryPath: String) -> Bool { + let candidates = [ + URL(fileURLWithPath: binaryPath).standardizedFileURL.path, + URL(fileURLWithPath: binaryPath).resolvingSymlinksInPath().standardizedFileURL.path, + ] + return candidates.contains { candidate in + commandLine == candidate || commandLine.hasPrefix("\(candidate) ") + } + } + + /// Production wiring for ``tryWarmAgyFetch``: list processes via `ps`, find + /// listening ports via `lsof`, and probe the token-less CLI HTTPS endpoint. + static func liveWarmAgyDependencies() -> WarmAgyDependencies { + WarmAgyDependencies( + processInfos: { timeout in + // A missing-CSRF/notRunning throw means no reusable server; the + // caller maps any throw to "no warm agy" and spawns instead. + try await AntigravityStatusProbe.detectProcessInfos( + timeout: timeout, + scope: .ideAndCLI) + }, + listeningPorts: { pid, timeout in + try await AntigravityStatusProbe.listeningPorts(pid: pid, timeout: timeout) + }, + fetchSnapshot: { ports, timeout in + let deadline = Date().addingTimeInterval(timeout) + return try await AntigravityStatusProbe(timeout: timeout) + .fetchFromPorts(ports, deadline: deadline) + }, + processOwnerUserID: { pid in + AntigravityProcessIdentityProvider().ownerUserID(for: pid_t(pid)) + }, + currentUserID: { + AntigravityProcessIdentityProvider.currentUserID + }, + ownedPID: { + // The pid of the `agy` CodexBar manages through the shared + // session, so the warm scan never reuses our own process. + await AntigravityCLISession.shared.pid.map(Int.init) + }) + } + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + BinaryLocator.resolveAntigravityBinary(env: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let binary = BinaryLocator.resolveAntigravityBinary(env: context.env) else { + throw AntigravityStatusProbeError.notRunning + } + let expectedAccountEmail: String? = if context.sourceMode == .auto, + context.selectedTokenAccountID != nil + { + AntigravitySelectedAccountGuard.selectedAccountEmail(context: context) + } else { + nil + } + let result = try await self.fetchUsingWarmSession( + binary: binary, + idleWindow: context.persistentCLISessionIdleWindow, + resetAfterFetch: Self.shouldResetSessionAfterFetch(context), + expectedAccountEmail: expectedAccountEmail) + try AntigravitySelectedAccountGuard.validate(result.usage, context: context) + return result + } + + private func fetchUsingWarmSession( + binary: String, + idleWindow: TimeInterval?, + resetAfterFetch: Bool, + expectedAccountEmail: String?) async throws -> ProviderFetchResult + { + try await self.fetchUsingWarmSession( + binary: binary, + idleWindow: idleWindow, + resetAfterFetch: resetAfterFetch, + expectedAccountEmail: expectedAccountEmail, + warmDependencies: Self.liveWarmAgyDependencies(), + spawnFetch: { binary, idleWindow, resetAfterFetch in + try await self.fetchBySpawning( + binary: binary, + idleWindow: idleWindow, + resetAfterFetch: resetAfterFetch) + }) + } + + /// Testable core of the CLI fetch: try the warm-reuse fast path first, then + /// fall back to spawning. The `spawnFetch` seam lets tests assert the spawn + /// path is skipped when a warm server is reused. + func fetchUsingWarmSession( + binary: String, + idleWindow: TimeInterval?, + resetAfterFetch: Bool, + expectedAccountEmail: String? = nil, + warmDependencies: WarmAgyDependencies, + spawnFetch: @Sendable (String, TimeInterval?, Bool) async throws -> ProviderFetchResult) + async throws -> ProviderFetchResult + { + // Fast path: reuse an already-running, authenticated `agy` CLI server if + // one is present, avoiding a fresh spawn and its multi-second warm-up. + // When none is found (or none answers), fall through to the spawn path. + // + // The warm path deliberately does NOT touch `AntigravityCLISession` + // (`beginProbe`/`finishProbe`): a discovered `agy` is owned by another + // process (an IDE, a long-lived `agy`, or another CodexBar host), so + // CodexBar must not manage its lifecycle, idle timeout, or + // `resetAfterFetch` teardown. Those apply only to processes CodexBar + // itself spawns on the fallback path below. + // Long-lived hosts already keep a managed session warm. Restrict external + // process reuse to one-shot CLI calls so app/server lifecycle accounting + // stays entirely inside AntigravityCLISession. + if resetAfterFetch, let warmSnapshot = try await Self.tryWarmAgyFetch( + timeout: 2.0, + expectedBinaryPath: binary, + expectedAccountEmail: expectedAccountEmail, + dependencies: warmDependencies) + { + // `tryWarmAgyFetch` only returns a snapshot whose `toUsageSnapshot()` + // already succeeded, so this conversion must not silently fail. + let warmUsage = try warmSnapshot.toUsageSnapshot() + return self.makeResult( + usage: warmUsage, + sourceLabel: Self.sourceLabel) + } + + try Task.checkCancellation() + return try await spawnFetch(binary, idleWindow, resetAfterFetch) + } + + /// Spawn (or reuse CodexBar's own warm) `agy` session and wait for the CLI + /// HTTPS endpoint to report ready. This is the original behavior, unchanged. + private func fetchBySpawning( + binary: String, + idleWindow: TimeInterval?, + resetAfterFetch: Bool) async throws -> ProviderFetchResult + { + let session = AntigravityCLISession.shared + let pid = try await session.beginProbe(binary: binary, idleWindow: idleWindow) + let deadline = Date().addingTimeInterval(5.0) + let snap: AntigravityStatusSnapshot + let usage: UsageSnapshot + do { + snap = try await Self.waitForSnapshot( + pid: pid, + deadline: deadline, + dependencies: SnapshotWaitDependencies( + pollIntervalNanoseconds: 200_000_000, + listeningPorts: { pid, timeout in + try await AntigravityStatusProbe.listeningPorts(pid: pid, timeout: timeout) + }, + drainOutput: { + await session.drainOutput() + }, + fetchSnapshot: { ports in + let timeout = min(2.0, max(0.2, deadline.timeIntervalSinceNow)) + return try await AntigravityStatusProbe(timeout: timeout) + .fetchFromPorts(ports, deadline: deadline) + })) + usage = try snap.toUsageSnapshot() + await session.finishProbe(success: true, resetAfterFetch: resetAfterFetch) + } catch { + let authenticationRequired = (error as? AntigravityStatusProbeError) == .authenticationRequired + await session.finishProbe( + success: false, + resetAfterFetch: resetAfterFetch || authenticationRequired, + forceTerminate: authenticationRequired) + throw error + } + + return self.makeResult( + usage: usage, + sourceLabel: Self.sourceLabel) + } + + static func shouldResetSessionAfterFetch(_ context: ProviderFetchContext) -> Bool { + // Long-lived hosts (the app, `codexbar serve`) keep the warm `agy` + // session between fetches; only one-shot CLI invocations reset it. + context.runtime == .cli && !context.persistsCLISessions + } + + /// Waits for real API readiness, not just socket readiness. Fresh ``agy`` + /// processes bind ports quickly, but ``GetUserStatus`` can return transient + /// initialization failures for a few seconds after the port appears. + static func waitForSnapshot( + pid: pid_t, + deadline: Date, + dependencies: SnapshotWaitDependencies) async throws -> AntigravityStatusSnapshot + { + var lastFetchError: Error? + while dependencies.now() < deadline { + try await Self.checkAuthenticationPrompt(dependencies) + let remaining = deadline.timeIntervalSince(dependencies.now()) + let portProbeTimeout = min(2.0, max(0.2, remaining)) + let ports: [Int] + do { + ports = try await dependencies.listeningPorts(Int(pid), portProbeTimeout) + } catch { + guard Self.isNoListeningPortsError(error) else { + try await Self.checkAuthenticationPrompt(dependencies) + throw error + } + ports = [] + } + if !ports.isEmpty { + var readySnapshot: AntigravityStatusSnapshot? + do { + let snapshot = try await dependencies.fetchSnapshot(ports) + _ = try snapshot.toUsageSnapshot() + readySnapshot = snapshot + } catch { + try await Self.checkAuthenticationPrompt(dependencies) + lastFetchError = error + Self.log.debug("Antigravity CLI HTTPS endpoint not ready", metadata: [ + "pid": "\(pid)", + "ports": ports.map(String.init).joined(separator: ","), + "error": error.localizedDescription, + ]) + } + if let readySnapshot { + try await Self.checkAuthenticationPrompt(dependencies) + return readySnapshot + } + } + + let remainingNanoseconds = UInt64( + max(0, deadline.timeIntervalSince(dependencies.now())) * 1_000_000_000) + guard remainingNanoseconds > 0 else { break } + let sleepNanoseconds = min(dependencies.pollIntervalNanoseconds, remainingNanoseconds) + if sleepNanoseconds > 0 { + try await Task.sleep(nanoseconds: sleepNanoseconds) + } + } + + try await Self.checkAuthenticationPrompt(dependencies) + if let lastFetchError { + throw lastFetchError + } + Self.log.warning("Antigravity CLI HTTPS: no ports found for pid \(pid)") + throw AntigravityStatusProbeError.portDetectionFailed( + "Antigravity CLI started but no listening ports found") + } + + static func containsAuthenticationPrompt(_ output: Data) -> Bool { + AntigravityCLIAuthenticationPrompt.contains(output) + } + + private static func checkAuthenticationPrompt(_ dependencies: SnapshotWaitDependencies) async throws { + let terminalOutput = await dependencies.drainOutput() + if Self.containsAuthenticationPrompt(terminalOutput) { + throw AntigravityStatusProbeError.authenticationRequired + } + } + + private static func isNoListeningPortsError(_ error: Error) -> Bool { + if case let AntigravityStatusProbeError.portDetectionFailed(message) = error { + return message == "no listening ports found" + } + if case let SubprocessRunnerError.nonZeroExit(code, stderr) = error { + return code == 1 && stderr.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + } + return false + } + + func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool { + context.sourceMode == .auto || context.sourceMode == .cli + } +} + +struct AntigravityOAuthFetchStrategy: ProviderFetchStrategy { + let id: String = "antigravity.oauth" + let kind: ProviderFetchKind = .oauth + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + static func usageSnapshot( + from snapshot: AntigravityStatusSnapshot, + updatedAt: Date = Date()) throws -> UsageSnapshot + { + if snapshot.modelQuotas.isEmpty { + return UsageSnapshot( + primary: nil, + secondary: nil, + tertiary: nil, + updatedAt: updatedAt, + identity: ProviderIdentitySnapshot( + providerID: .antigravity, + accountEmail: snapshot.accountEmail, + accountOrganization: nil, + loginMethod: snapshot.accountPlan)) + } + return try snapshot.toUsageSnapshot() + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let fetcher = AntigravityRemoteUsageFetcher( + environment: context.env, + credentialsUpdateHandler: { credentials in + guard let accountID = context.selectedTokenAccountID, + let updater = context.tokenAccountTokenUpdater + else { + return + } + let token = try AntigravityOAuthCredentialsStore.tokenAccountValue(for: credentials) + await updater(.antigravity, accountID, token) + }) + let snapshot = try await fetcher.fetch() + let usage = try Self.usageSnapshot(from: snapshot) + return self.makeResult( + usage: usage, + sourceLabel: "oauth") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} + +/// Guards ambient Antigravity snapshots against the explicitly selected account. +/// +/// The local desktop probe and the ``agy`` CLI HTTPS server report whichever +/// Antigravity account is signed into the local session. When the user has +/// selected a specific saved Google account, an ambient probe can return a +/// *different* account's quota. Only the OAuth strategy is account-scoped (it +/// fetches with the selected account's injected credentials), so in ``auto`` +/// mode we reject a snapshot whose identity does not match the selected account +/// and let the pipeline fall through to OAuth. Explicit ``cli``/``oauth`` source +/// modes stay authoritative and are never second-guessed here. +enum AntigravitySelectedAccountGuard { + static func matches(snapshotAccountEmail: String?, expectedAccountEmail: String?) -> Bool { + guard let expected = self.normalizedEmail(expectedAccountEmail) else { return true } + guard let found = self.normalizedEmail(snapshotAccountEmail) else { return false } + return found.caseInsensitiveCompare(expected) == .orderedSame + } + + static func validate(_ usage: UsageSnapshot, context: ProviderFetchContext) throws { + guard context.sourceMode == .auto, context.selectedTokenAccountID != nil else { return } + let expected = self.selectedAccountEmail(context: context) + let found = self.normalizedEmail(usage.identity?.accountEmail) + guard let expected, let found, found.caseInsensitiveCompare(expected) == .orderedSame else { + throw AntigravityStatusProbeError.accountMismatch(expected: expected, found: found) + } + } + + /// Email of the selected token account, read from the same injected + /// credentials the OAuth strategy would use (`ANTIGRAVITY_OAUTH_CREDENTIALS_JSON`). + static func selectedAccountEmail(context: ProviderFetchContext) -> String? { + guard let value = context.env[AntigravityOAuthCredentialsStore.environmentCredentialsKey], + let credentials = AntigravityOAuthCredentialsStore.credentials(fromTokenAccountValue: value) + else { + return nil + } + return credentials.resolvedAccountEmail + } + + private static func normalizedEmail(_ email: String?) -> String? { + guard let trimmed = email?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityQuotaSummaryParser.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityQuotaSummaryParser.swift new file mode 100644 index 0000000..9928f76 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityQuotaSummaryParser.swift @@ -0,0 +1,173 @@ +import Foundation + +// swiftformat:disable:next redundantSendable +struct AntigravityQuotaSummary: Sendable, Equatable { + let description: String? + let groups: [AntigravityQuotaSummaryGroup] +} + +// swiftformat:disable:next redundantSendable +struct AntigravityQuotaSummaryGroup: Sendable, Equatable { + let displayName: String + let description: String? + let buckets: [AntigravityQuotaSummaryBucket] +} + +// swiftformat:disable:next redundantSendable +struct AntigravityQuotaSummaryBucket: Sendable, Equatable { + let bucketId: String + let displayName: String + let remainingFraction: Double? + let resetTime: Date? + let resetDescription: String? + let disabled: Bool + + init( + bucketId: String, + displayName: String, + remainingFraction: Double?, + resetTime: Date? = nil, + resetDescription: String?, + disabled: Bool) + { + self.bucketId = bucketId + self.displayName = displayName + self.remainingFraction = remainingFraction + self.resetTime = resetTime + self.resetDescription = resetDescription + self.disabled = disabled + } +} + +extension AntigravityStatusProbe { + static func parseQuotaSummaryResponse(_ data: Data) throws -> AntigravityStatusSnapshot { + let decoder = JSONDecoder() + let response = try decoder.decode(QuotaSummaryResponse.self, from: data) + if let invalid = Self.invalidCode(response.code) { + throw AntigravityStatusProbeError.apiError(invalid) + } + let payload = response.response ?? response.summary ?? response.rootPayload + guard let payload else { + throw AntigravityStatusProbeError.parseFailed("Missing quota summary") + } + let groups = payload.groups.compactMap(self.quotaSummaryGroup(from:)) + guard !groups.isEmpty else { + throw AntigravityStatusProbeError.parseFailed("Missing quota groups") + } + return AntigravityStatusSnapshot( + quotaSummary: AntigravityQuotaSummary( + description: payload.description, + groups: groups), + accountEmail: nil, + accountPlan: nil, + source: .local) + } + + private static func quotaSummaryGroup(from payload: QuotaSummaryGroupPayload) -> AntigravityQuotaSummaryGroup? { + let displayName = payload.displayName?.trimmingCharacters(in: .whitespacesAndNewlines) + let buckets = (payload.buckets ?? []).compactMap(self.quotaSummaryBucket(from:)) + guard !buckets.isEmpty else { return nil } + return AntigravityQuotaSummaryGroup( + displayName: self.nonEmpty(displayName) ?? "Quota", + description: payload.description, + buckets: buckets) + } + + private static func quotaSummaryBucket(from payload: QuotaSummaryBucketPayload) -> AntigravityQuotaSummaryBucket? { + let bucketId = payload.bucketId?.trimmingCharacters(in: .whitespacesAndNewlines) + let displayName = payload.displayName?.trimmingCharacters(in: .whitespacesAndNewlines) + guard let resolvedBucketId = bucketId, !resolvedBucketId.isEmpty else { return nil } + let resetTime = payload.resetTime.flatMap { Self.parseDate($0) } + return AntigravityQuotaSummaryBucket( + bucketId: resolvedBucketId, + displayName: self.nonEmpty(displayName) ?? resolvedBucketId, + remainingFraction: payload.resolvedRemainingFraction, + resetTime: resetTime, + resetDescription: payload.description, + disabled: payload.disabled ?? false) + } + + private static func nonEmpty(_ value: String?) -> String? { + guard let value, !value.isEmpty else { return nil } + return value + } +} + +private struct QuotaSummaryResponse: Decodable { + let code: CodeValue? + let message: String? + let response: QuotaSummaryPayload? + let summary: QuotaSummaryPayload? + let description: String? + let groups: [QuotaSummaryGroupPayload]? + + var rootPayload: QuotaSummaryPayload? { + guard let groups else { return nil } + return QuotaSummaryPayload(description: self.description, groups: groups) + } +} + +private struct QuotaSummaryPayload: Decodable { + let description: String? + let groups: [QuotaSummaryGroupPayload] + + init(description: String?, groups: [QuotaSummaryGroupPayload]) { + self.description = description + self.groups = groups + } + + private enum CodingKeys: String, CodingKey { + case description + case groups + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.description = try container.decodeIfPresent(String.self, forKey: .description) + self.groups = try container.decodeIfPresent([QuotaSummaryGroupPayload].self, forKey: .groups) ?? [] + } +} + +private struct QuotaSummaryGroupPayload: Decodable { + let displayName: String? + let description: String? + let buckets: [QuotaSummaryBucketPayload]? +} + +private struct QuotaSummaryBucketPayload: Decodable { + let bucketId: String? + let displayName: String? + let description: String? + let disabled: Bool? + let remainingFraction: Double? + let remaining: QuotaSummaryRemainingPayload? + let resetTime: String? + + var resolvedRemainingFraction: Double? { + self.remainingFraction ?? self.remaining?.remainingFraction + } +} + +private struct QuotaSummaryRemainingPayload: Decodable { + let remainingFraction: Double? + + private enum CodingKeys: String, CodingKey { + case remainingFraction + case oneofCase = "case" + case value + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + if let remainingFraction = try container.decodeIfPresent(Double.self, forKey: .remainingFraction) { + self.remainingFraction = remainingFraction + return + } + let oneofCase = try container.decodeIfPresent(String.self, forKey: .oneofCase) + if oneofCase == "remainingFraction" { + self.remainingFraction = try container.decodeIfPresent(Double.self, forKey: .value) + } else { + self.remainingFraction = nil + } + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityRemoteUsageFetcher.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityRemoteUsageFetcher.swift new file mode 100644 index 0000000..536e1e0 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityRemoteUsageFetcher.swift @@ -0,0 +1,777 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum AntigravityRemoteFetchError: LocalizedError, Sendable, Equatable { + case notLoggedIn + case permissionDenied(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .notLoggedIn: + "Antigravity Google auth not found. Use Antigravity login to authenticate." + case let .permissionDenied(message): + "Antigravity remote API permission denied: \(message)" + case let .apiError(message): + "Antigravity remote API error: \(message)" + case let .parseFailed(message): + "Could not parse Antigravity remote usage: \(message)" + } + } +} + +public struct AntigravityRemoteUsageFetcher: Sendable { + public var timeout: TimeInterval = 10.0 + public var homeDirectory: String + public var environment: [String: String] + public var dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse) + public var oauthClientResolver: @Sendable () -> AntigravityOAuthClient? + public var credentialsUpdateHandler: @Sendable (AntigravityOAuthCredentials) async throws -> Void + + private static let log = CodexBarLog.logger(LogCategories.antigravity) + private static let userAgent = "antigravity" + private static let baseURL = "https://cloudcode-pa.googleapis.com" + private static let loadCodeAssistEndpoint = "\(baseURL)/v1internal:loadCodeAssist" + private static let onboardUserEndpoint = "\(baseURL)/v1internal:onboardUser" + private static let fetchAvailableModelsEndpoint = "\(baseURL)/v1internal:fetchAvailableModels" + private static let retrieveUserQuotaEndpoint = "\(baseURL)/v1internal:retrieveUserQuota" + private static let refreshSafetyWindow: TimeInterval = 60 + + private struct FetchContext { + let timeout: TimeInterval + let store: AntigravityOAuthCredentialsStore? + let dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse) + let oauthClientResolver: @Sendable () -> AntigravityOAuthClient? + let credentialsUpdateHandler: @Sendable (AntigravityOAuthCredentials) async throws -> Void + + func persistCredentials(_ credentials: AntigravityOAuthCredentials) async throws { + if let store { + try store.save(credentials) + } else { + try await self.credentialsUpdateHandler(credentials) + } + } + } + + public init( + timeout: TimeInterval = 10.0, + homeDirectory: String = NSHomeDirectory(), + environment: [String: String] = ProcessInfo.processInfo.environment, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse) = { request in + try await ProviderHTTPClient.shared.data(for: request) + }, + oauthClientResolver: @escaping @Sendable () -> AntigravityOAuthClient? = { + AntigravityOAuthConfig.resolvedClient() + }, + credentialsUpdateHandler: @escaping @Sendable (AntigravityOAuthCredentials) async throws -> Void = { _ in }) + { + self.timeout = timeout + self.homeDirectory = homeDirectory + self.environment = environment + self.dataLoader = dataLoader + self.oauthClientResolver = oauthClientResolver + self.credentialsUpdateHandler = credentialsUpdateHandler + } + + public func fetch() async throws -> AntigravityStatusSnapshot { + let source = try Self.resolveCredentialSource(homeDirectory: self.homeDirectory, environment: self.environment) + guard let credentials = source.credentials else { + throw AntigravityRemoteFetchError.notLoggedIn + } + return try await self.fetchSnapshot( + using: credentials, + store: source.store) + } + + private func fetchSnapshot( + using initialCredentials: AntigravityOAuthCredentials, + store: AntigravityOAuthCredentialsStore?) async throws + -> AntigravityStatusSnapshot + { + guard let storedAccessToken = initialCredentials.accessToken?.trimmedNonEmpty else { + throw AntigravityRemoteFetchError.notLoggedIn + } + + var credentials = initialCredentials + var accessToken = storedAccessToken + let context = FetchContext( + timeout: self.timeout, + store: store, + dataLoader: self.dataLoader, + oauthClientResolver: self.oauthClientResolver, + credentialsUpdateHandler: self.credentialsUpdateHandler) + if Self.shouldRefresh(expiryDate: credentials.expiryDate, now: Date()) { + guard let refreshToken = credentials.refreshToken?.trimmedNonEmpty else { + throw AntigravityRemoteFetchError.notLoggedIn + } + let refreshed = try await Self.refreshAccessToken( + credentials: credentials, + refreshToken: refreshToken, + context: context) + accessToken = refreshed.accessToken + credentials = refreshed.credentials + if let store { + credentials = try store.load() ?? credentials + } + credentials.accessToken = credentials.accessToken?.trimmedNonEmpty ?? accessToken + } + + let claims = Self.extractClaims(from: credentials) + let codeAssist = try await Self.loadCodeAssist( + accessToken: accessToken, + timeout: self.timeout, + dataLoader: self.dataLoader) + let projectId = try await Self.resolveProjectID( + accessToken: accessToken, + storedProjectID: credentials.projectID?.trimmedNonEmpty, + initialResponse: codeAssist, + context: context) + if let projectId, credentials.projectID?.trimmedNonEmpty != projectId { + credentials.projectID = projectId + do { + try await context.persistCredentials(credentials) + } catch { + Self.log.warning("Could not persist Antigravity project ID: \(error.localizedDescription)") + } + } + let models = try await Self.fetchModelQuotas( + accessToken: accessToken, + projectId: projectId, + timeout: self.timeout, + dataLoader: self.dataLoader) + + return AntigravityStatusSnapshot( + modelQuotas: models, + accountEmail: claims.email, + accountPlan: Self.resolvePlan(response: codeAssist, claims: claims), + source: .remote) + } + + private static func shouldRefresh(expiryDate: Date?, now: Date) -> Bool { + guard let expiryDate else { return false } + return expiryDate.timeIntervalSince(now) <= Self.refreshSafetyWindow + } + + private static func loadCodeAssist( + accessToken: String, + timeout: TimeInterval, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> CodeAssistResponse + { + let body = [ + "metadata": [ + "ideType": "ANTIGRAVITY", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI", + ], + ] + return try await Self.sendRequest( + endpoint: Self.loadCodeAssistEndpoint, + accessToken: accessToken, + body: body, + timeout: timeout, + dataLoader: dataLoader) + } + + private static func fetchAvailableModels( + accessToken: String, + projectId: String?, + timeout: TimeInterval, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> FetchAvailableModelsResponse + { + let body: [String: Any] = if let projectId = projectId?.trimmedNonEmpty { + ["project": projectId] + } else { + [:] + } + return try await Self.sendRequest( + endpoint: Self.fetchAvailableModelsEndpoint, + accessToken: accessToken, + body: body, + timeout: timeout, + dataLoader: dataLoader) + } + + private static func fetchModelQuotas( + accessToken: String, + projectId: String?, + timeout: TimeInterval, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> [AntigravityModelQuota] + { + do { + let response = try await Self.fetchAvailableModels( + accessToken: accessToken, + projectId: projectId, + timeout: timeout, + dataLoader: dataLoader) + let modelQuotas = try Self.parseModelQuotas(response) + if Self.shouldVerifyFullRemoteQuotas(modelQuotas) { + let quotaBuckets = try await Self.fetchQuotaBucketsIfPermitted( + accessToken: accessToken, + projectId: projectId, + timeout: timeout, + dataLoader: dataLoader) + guard let quotaBuckets, Self.hasQuotaFractionData(quotaBuckets) else { + return [] + } + return Self.mergeVerifiedQuotas(modelQuotas: modelQuotas, verifiedQuotas: quotaBuckets) + } + return modelQuotas + } catch let error as AntigravityRemoteFetchError { + guard case .permissionDenied = error else { + throw error + } + Self.log.info("Falling back to retrieveUserQuota for Antigravity remote usage") + return try await Self.fetchQuotaBucketsIfPermitted( + accessToken: accessToken, + projectId: projectId, + timeout: timeout, + dataLoader: dataLoader) ?? [] + } + } + + private static func mergeVerifiedQuotas( + modelQuotas: [AntigravityModelQuota], + verifiedQuotas: [AntigravityModelQuota]) -> [AntigravityModelQuota] + { + var verifiedByModelID = Dictionary( + uniqueKeysWithValues: verifiedQuotas.map { (Self.quotaKey($0), $0) }) + var merged = modelQuotas.compactMap { modelQuota -> AntigravityModelQuota? in + guard let verifiedQuota = verifiedByModelID.removeValue(forKey: Self.quotaKey(modelQuota)) else { + return nil + } + let resetTime = verifiedQuota.resetTime ?? modelQuota.resetTime + return AntigravityModelQuota( + label: modelQuota.label, + modelId: modelQuota.modelId, + remainingFraction: verifiedQuota.remainingFraction ?? modelQuota.remainingFraction, + resetTime: resetTime, + resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) }) + } + let unmatchedVerifiedQuotas = verifiedByModelID.values + .filter { $0.remainingFraction != nil } + .sorted { lhs, rhs in + lhs.modelId.localizedCaseInsensitiveCompare(rhs.modelId) == .orderedAscending + } + merged.append(contentsOf: unmatchedVerifiedQuotas) + return merged + } + + private static func quotaKey(_ quota: AntigravityModelQuota) -> String { + quota.modelId.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } + + private static func shouldVerifyFullRemoteQuotas(_ quotas: [AntigravityModelQuota]) -> Bool { + guard !quotas.isEmpty else { return false } + return quotas.allSatisfy { quota in + guard let remaining = quota.remainingFraction else { return false } + return remaining >= 0.999 + } + } + + private static func hasQuotaFractionData(_ quotas: [AntigravityModelQuota]) -> Bool { + quotas.contains { quota in + quota.remainingFraction != nil + } + } + + private static func fetchQuotaBucketsIfPermitted( + accessToken: String, + projectId: String?, + timeout: TimeInterval, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> [AntigravityModelQuota]? + { + do { + let response = try await Self.retrieveUserQuota( + accessToken: accessToken, + projectId: projectId, + timeout: timeout, + dataLoader: dataLoader) + return try Self.parseQuotaBuckets(response) + } catch let quotaError as AntigravityRemoteFetchError { + guard case .permissionDenied = quotaError else { + throw quotaError + } + Self.log.info("Antigravity remote quota endpoint is not permitted for this account") + return nil + } + } + + private static func retrieveUserQuota( + accessToken: String, + projectId: String?, + timeout: TimeInterval, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> RetrieveUserQuotaResponse + { + let body: [String: Any] = if let projectId = projectId?.trimmedNonEmpty { + ["project": projectId] + } else { + [:] + } + return try await Self.sendRequest( + endpoint: Self.retrieveUserQuotaEndpoint, + accessToken: accessToken, + body: body, + timeout: timeout, + dataLoader: dataLoader) + } + + private static func resolveProjectID( + accessToken: String, + storedProjectID: String?, + initialResponse: CodeAssistResponse, + context: FetchContext) async throws + -> String? + { + if let storedProjectID { + return storedProjectID + } + + if let projectID = initialResponse.projectID { + return projectID + } + + guard let tierID = Self.pickOnboardTier(from: initialResponse) else { + return nil + } + + let onboardBody: [String: Any] = [ + "tierId": tierID, + "metadata": [ + "ideType": "ANTIGRAVITY", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI", + ], + ] + + do { + let onboardResponse: OnboardResponse = try await Self.sendRequest( + endpoint: Self.onboardUserEndpoint, + accessToken: accessToken, + body: onboardBody, + timeout: context.timeout, + dataLoader: context.dataLoader) + if let projectID = onboardResponse.projectID { + return projectID + } + } catch { + Self.log.warning("Antigravity onboarding request failed", metadata: [ + "error": "\(error.localizedDescription)", + ]) + } + + for _ in 0..<5 { + try? await Task.sleep(for: .milliseconds(2000)) + let refreshed = try await Self.loadCodeAssist( + accessToken: accessToken, + timeout: context.timeout, + dataLoader: context.dataLoader) + if let projectID = refreshed.projectID { + return projectID + } + } + + return nil + } + + private static func sendRequest( + endpoint: String, + accessToken: String, + body: [String: Any], + timeout: TimeInterval, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> Response + { + guard let url = URL(string: endpoint) else { + throw AntigravityRemoteFetchError.apiError("Invalid endpoint URL") + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.timeoutInterval = timeout + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue(Self.userAgent, forHTTPHeaderField: "User-Agent") + request.httpBody = try JSONSerialization.data(withJSONObject: body) + + let httpResponse = try await ProviderHTTPTransportHandler(dataLoader).response(for: request) + + switch httpResponse.statusCode { + case 200: + break + case 401: + throw AntigravityRemoteFetchError.notLoggedIn + case 403: + let message = String(data: httpResponse.data, encoding: .utf8)?.trimmedNonEmpty ?? "HTTP 403" + throw AntigravityRemoteFetchError.permissionDenied(message) + default: + let message = String(data: httpResponse.data, encoding: .utf8)?.trimmedNonEmpty + ?? "HTTP \(httpResponse.statusCode)" + throw AntigravityRemoteFetchError.apiError("HTTP \(httpResponse.statusCode): \(message)") + } + + do { + return try JSONDecoder().decode(Response.self, from: httpResponse.data) + } catch { + throw AntigravityRemoteFetchError.parseFailed(error.localizedDescription) + } + } + + private static func parseModelQuotas(_ response: FetchAvailableModelsResponse) throws -> [AntigravityModelQuota] { + let models = response.models ?? [:] + return models.compactMap { modelID, model in + guard let quotaInfo = model.quotaInfo else { return nil } + let resetTime = quotaInfo.resetTime.flatMap(Self.parseResetTime(_:)) + let label = model.displayName?.trimmedNonEmpty + ?? model.label?.trimmedNonEmpty + ?? modelID + return AntigravityModelQuota( + label: label, + modelId: modelID, + remainingFraction: quotaInfo.remainingFraction, + resetTime: resetTime, + resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) }) + } + } + + private static func parseQuotaBuckets(_ response: RetrieveUserQuotaResponse) throws -> [AntigravityModelQuota] { + guard let buckets = response.buckets else { + throw AntigravityRemoteFetchError.parseFailed("No quota buckets in response") + } + guard !buckets.isEmpty else { return [] } + + var modelQuotaMap: [String: (fraction: Double?, resetTime: String?)] = [:] + for bucket in buckets { + guard let modelID = bucket.modelId?.trimmedNonEmpty else { continue } + let next = (bucket.remainingFraction, bucket.resetTime) + if let existing = modelQuotaMap[modelID] { + let existingValue = existing.fraction ?? Double.greatestFiniteMagnitude + let nextValue = next.0 ?? Double.greatestFiniteMagnitude + if nextValue < existingValue { + modelQuotaMap[modelID] = next + } + } else { + modelQuotaMap[modelID] = next + } + } + + return modelQuotaMap.keys.sorted().compactMap { modelID in + guard let info = modelQuotaMap[modelID] else { return nil } + let resetTime = info.resetTime.flatMap(Self.parseResetTime(_:)) + return AntigravityModelQuota( + label: modelID, + modelId: modelID, + remainingFraction: info.fraction, + resetTime: resetTime, + resetDescription: resetTime.map { UsageFormatter.resetDescription(from: $0) }) + } + } + + private static func resolvePlan(response: CodeAssistResponse, claims: TokenClaims) -> String? { + if let planType = response.planInfo?.planType?.trimmedNonEmpty { + return planType + } + + switch (response.currentTier?.id?.trimmedNonEmpty, claims.hostedDomain) { + case ("standard-tier", _): + return "Paid" + case ("free-tier", .some): + return "Workspace" + case ("free-tier", .none): + return "Free" + case ("legacy-tier", _): + return "Legacy" + default: + return response.currentTier?.name?.trimmedNonEmpty + } + } + + private static func pickOnboardTier(from response: CodeAssistResponse) -> String? { + if let defaultTier = response.allowedTiers? + .first(where: { $0.isDefault == true && $0.id?.trimmedNonEmpty != nil })?.id?.trimmedNonEmpty + { + return defaultTier + } + if let firstTier = response.allowedTiers? + .first(where: { $0.id?.trimmedNonEmpty != nil })?.id?.trimmedNonEmpty + { + return firstTier + } + if let paidTier = response.paidTier?.id?.trimmedNonEmpty { + return paidTier + } + if let currentTier = response.currentTier?.id?.trimmedNonEmpty { + return currentTier + } + return nil + } + + private static func parseResetTime(_ value: String) -> Date? { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = formatter.date(from: value) { + return date + } + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: value) + } + + private static func credentialsStore(homeDirectory: String) -> AntigravityOAuthCredentialsStore { + let homeURL = URL(fileURLWithPath: homeDirectory, isDirectory: true) + return AntigravityOAuthCredentialsStore(fileURL: AntigravityOAuthCredentialsStore.defaultURL(home: homeURL)) + } + + private static func resolveCredentialSource( + homeDirectory: String, + environment: [String: String]) throws -> ( + credentials: AntigravityOAuthCredentials?, + store: AntigravityOAuthCredentialsStore?) + { + let primaryStore = Self.credentialsStore(homeDirectory: homeDirectory) + if let tokenValue = environment[AntigravityOAuthCredentialsStore.environmentCredentialsKey] { + guard let credentials = AntigravityOAuthCredentialsStore.credentials(fromTokenAccountValue: tokenValue) + else { + throw AntigravityRemoteFetchError.parseFailed("Could not decode selected account credentials.") + } + return (credentials, nil) + } + return try (primaryStore.load(), primaryStore) + } + + private struct RefreshResult { + let accessToken: String + let credentials: AntigravityOAuthCredentials + } + + private static func refreshAccessToken( + credentials: AntigravityOAuthCredentials, + refreshToken: String, + context: FetchContext) async throws + -> RefreshResult + { + let oauthClient = try Self.refreshOAuthClient( + from: credentials, + oauthClientResolver: context.oauthClientResolver) + + var request = URLRequest(url: AntigravityOAuthConfig.tokenURL) + request.httpMethod = "POST" + request.timeoutInterval = context.timeout + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.httpBody = Self.formBody([ + "client_id": oauthClient.clientID, + "client_secret": oauthClient.clientSecret, + "refresh_token": refreshToken, + "grant_type": "refresh_token", + ]) + + let httpResponse = try await ProviderHTTPTransportHandler(context.dataLoader).response(for: request) + guard httpResponse.statusCode == 200 else { + throw AntigravityRemoteFetchError.notLoggedIn + } + guard let json = try JSONSerialization.jsonObject(with: httpResponse.data) as? [String: Any], + let accessToken = json["access_token"] as? String + else { + throw AntigravityRemoteFetchError.parseFailed("Could not parse refresh response") + } + + let updatedCredentials = Self.updatedCredentials(credentials, refreshResponse: json) + try await context.persistCredentials(updatedCredentials) + return RefreshResult(accessToken: accessToken, credentials: updatedCredentials) + } + + private static func refreshOAuthClient( + from credentials: AntigravityOAuthCredentials, + oauthClientResolver: @escaping @Sendable () -> AntigravityOAuthClient?) throws + -> AntigravityOAuthClient + { + if let clientID = credentials.clientID?.trimmedNonEmpty, + let clientSecret = credentials.clientSecret?.trimmedNonEmpty + { + return AntigravityOAuthClient(clientID: clientID, clientSecret: clientSecret) + } + + guard let client = oauthClientResolver() else { + throw AntigravityRemoteFetchError.apiError(AntigravityOAuthConfig.missingCredentialsMessage) + } + return client + } + + private static func updatedCredentials( + _ credentials: AntigravityOAuthCredentials, + refreshResponse: [String: Any]) -> AntigravityOAuthCredentials + { + var credentials = credentials + if let accessToken = refreshResponse["access_token"] as? String { + credentials.accessToken = accessToken + } + if let expiresIn = refreshResponse["expires_in"] as? Double { + credentials.expiryDateMilliseconds = (Date().timeIntervalSince1970 + expiresIn) * 1000 + } + if let expiresIn = refreshResponse["expires_in"] as? Int { + credentials.expiryDateMilliseconds = (Date().timeIntervalSince1970 + Double(expiresIn)) * 1000 + } + if let idToken = refreshResponse["id_token"] as? String { + credentials.idToken = idToken + } + return credentials + } + + private static func formBody(_ values: [String: String]) -> Data? { + var components = URLComponents() + components.queryItems = values.map { key, value in + URLQueryItem(name: key, value: value) + } + return components.query?.data(using: .utf8) + } + + private struct TokenClaims { + let email: String? + let hostedDomain: String? + } + + private static func extractClaims(from credentials: AntigravityOAuthCredentials) -> TokenClaims { + let tokenClaims = Self.extractClaimsFromToken(credentials.idToken) + return TokenClaims( + email: tokenClaims.email ?? credentials.email?.trimmedNonEmpty, + hostedDomain: tokenClaims.hostedDomain) + } + + private static func extractClaimsFromToken(_ idToken: String?) -> TokenClaims { + guard let idToken else { + return TokenClaims(email: nil, hostedDomain: nil) + } + + let parts = idToken.components(separatedBy: ".") + guard parts.count >= 2 else { + return TokenClaims(email: nil, hostedDomain: nil) + } + + var payload = parts[1] + .replacingOccurrences(of: "-", with: "+") + .replacingOccurrences(of: "_", with: "/") + let remainder = payload.count % 4 + if remainder > 0 { + payload += String(repeating: "=", count: 4 - remainder) + } + + guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return TokenClaims(email: nil, hostedDomain: nil) + } + + return TokenClaims( + email: (json["email"] as? String)?.trimmedNonEmpty, + hostedDomain: (json["hd"] as? String)?.trimmedNonEmpty) + } +} + +extension String { + fileprivate var trimmedNonEmpty: String? { + let trimmed = self.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + } +} + +private struct ProjectReference: Decodable { + let value: String? + + init(from decoder: Decoder) throws { + let single = try decoder.singleValueContainer() + if let stringValue = try? single.decode(String.self) { + self.value = stringValue + return + } + + let keyed = try decoder.container(keyedBy: CodingKeys.self) + self.value = try keyed.decodeIfPresent(String.self, forKey: .id) + ?? keyed.decodeIfPresent(String.self, forKey: .projectID) + } + + private enum CodingKeys: String, CodingKey { + case id + case projectID = "projectId" + } +} + +private struct CodeAssistResponse: Decodable { + let planInfo: CodeAssistPlanInfo? + let currentTier: TierInfo? + let paidTier: TierInfo? + let allowedTiers: [AllowedTier]? + let cloudaicompanionProject: ProjectReference? + + var projectID: String? { + self.cloudaicompanionProject?.value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty + } +} + +private struct CodeAssistPlanInfo: Decodable { + let planType: String? +} + +private struct TierInfo: Decodable { + let id: String? + let name: String? +} + +private struct AllowedTier: Decodable { + let id: String? + let isDefault: Bool? +} + +private struct OnboardResponse: Decodable { + let response: OnboardInnerResponse? + + var projectID: String? { + self.response?.cloudaicompanionProject?.value?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty + } +} + +private struct OnboardInnerResponse: Decodable { + let cloudaicompanionProject: ProjectReference? +} + +private struct FetchAvailableModelsResponse: Decodable { + let models: [String: AntigravityRemoteModel]? +} + +private struct RetrieveUserQuotaResponse: Decodable { + let buckets: [RetrieveUserQuotaBucket]? +} + +private struct RetrieveUserQuotaBucket: Decodable { + let modelId: String? + let remainingFraction: Double? + let resetTime: String? +} + +private struct AntigravityRemoteModel: Decodable { + let displayName: String? + let label: String? + let quotaInfo: AntigravityRemoteQuotaInfo? +} + +private struct AntigravityRemoteQuotaInfo: Decodable { + let remainingFraction: Double? + let resetTime: String? +} + +extension String? { + fileprivate var trimmedNonEmpty: String? { + self?.trimmingCharacters(in: .whitespacesAndNewlines).nilIfEmpty + } +} + +extension String { + fileprivate var nilIfEmpty: String? { + self.isEmpty ? nil : self + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+Deadline.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+Deadline.swift new file mode 100644 index 0000000..da86b63 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+Deadline.swift @@ -0,0 +1,128 @@ +import Foundation + +extension AntigravityStatusProbe { + private static let processProbeLog = CodexBarLog.logger(LogCategories.antigravity) + + private enum ProcessSnapshotFetchFailure { + case antigravity(AntigravityStatusProbeError) + case url(URLError) + case cancellation + case other(String) + + init(_ error: Error) { + if let error = error as? AntigravityStatusProbeError { + self = .antigravity(error) + } else if let error = error as? URLError, error.code == .cancelled { + self = .cancellation + } else if let error = error as? URLError { + self = .url(error) + } else if error is CancellationError { + self = .cancellation + } else { + self = .other(error.localizedDescription) + } + } + + var error: Error { + switch self { + case let .antigravity(error): + error + case let .url(error): + error + case .cancellation: + CancellationError() + case let .other(message): + AntigravityStatusProbeError.apiError(message) + } + } + } + + private enum ProcessSnapshotFetchOutcome { + case success(index: Int, snapshot: AntigravityStatusSnapshot) + case failure(index: Int, pid: Int, failure: ProcessSnapshotFetchFailure) + } + + static func fetchProcessSnapshots( + processInfos: [ProcessInfoResult], + fetch: @escaping @Sendable (ProcessInfoResult) async throws -> AntigravityStatusSnapshot) + async throws -> (snapshots: [AntigravityStatusSnapshot], lastError: Error?) + { + let outcomes = await withTaskGroup(of: ProcessSnapshotFetchOutcome.self) { group in + for (index, processInfo) in processInfos.enumerated() { + group.addTask { + do { + return try await .success(index: index, snapshot: fetch(processInfo)) + } catch { + return .failure(index: index, pid: processInfo.pid, failure: .init(error)) + } + } + } + + var ordered = [ProcessSnapshotFetchOutcome?](repeating: nil, count: processInfos.count) + for await outcome in group { + switch outcome { + case let .success(index, _), let .failure(index, _, _): + ordered[index] = outcome + } + } + return ordered.compactMap(\.self) + } + + var snapshots: [AntigravityStatusSnapshot] = [] + var lastError: Error? + for outcome in outcomes { + switch outcome { + case let .success(_, snapshot): + snapshots.append(snapshot) + case let .failure(_, pid, failure): + if case .cancellation = failure { + throw CancellationError() + } + let error = failure.error + lastError = error + Self.processProbeLog.debug("Antigravity local process probe failed", metadata: [ + "pid": "\(pid)", + "error": error.localizedDescription, + ]) + } + } + return (snapshots, lastError) + } + + static func fetch( + processInfo: ProcessInfoResult, + timeout: TimeInterval, + deadline: Date) async throws -> AntigravityStatusSnapshot + { + guard let portTimeout = timeoutForNextAttempt(timeout: timeout, deadline: deadline) else { + throw AntigravityStatusProbeError.timedOut + } + let ports = try await Self.listeningPorts(pid: processInfo.pid, timeout: portTimeout) + let endpoint = try await Self.resolveWorkingEndpoint( + candidateEndpoints: Self.connectionCandidates( + listeningPorts: ports, + languageServerCSRFToken: processInfo.csrfToken, + extensionServerPort: processInfo.extensionPort, + extensionServerCSRFToken: processInfo.extensionServerCSRFToken), + timeout: timeout, + deadline: deadline) + let context = RequestContext( + endpoints: Self.requestEndpoints( + resolvedEndpoint: endpoint, + listeningPorts: ports, + languageServerCSRFToken: processInfo.csrfToken, + extensionServerPort: processInfo.extensionPort, + extensionServerCSRFToken: processInfo.extensionServerCSRFToken), + timeout: timeout, + deadline: deadline) + + return try await Self.fetchSnapshot(context: context) + } + + static func timeoutForNextAttempt(timeout: TimeInterval, deadline: Date?) -> TimeInterval? { + guard let deadline else { return timeout } + let remaining = deadline.timeIntervalSinceNow + guard remaining > 0 else { return nil } + return min(timeout, remaining) + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+LocalEndpoints.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+LocalEndpoints.swift new file mode 100644 index 0000000..fdb4ff6 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+LocalEndpoints.swift @@ -0,0 +1,23 @@ +extension AntigravityStatusProbe { + static func cliEndpoints(ports: [Int]) -> [AntigravityConnectionEndpoint] { + ports.flatMap { port in + self.localProbeSchemes.map { scheme in + AntigravityConnectionEndpoint( + scheme: scheme, + port: port, + csrfToken: "", + source: .cliHTTPS) + } + } + } + + static var localProbeSchemes: [String] { + #if os(Linux) + // FoundationNetworking cannot trust Antigravity's self-signed TLS cert. + // Requests remain pinned to 127.0.0.1 in makeRequest. + ["https", "http"] + #else + ["https"] + #endif + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+PortDetection.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+PortDetection.swift new file mode 100644 index 0000000..7510912 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+PortDetection.swift @@ -0,0 +1,137 @@ +import Foundation + +/// Parses Linux `/proc//net/tcp{,6}` output to recover the listening ports +/// owned by a process. The parsing is platform-independent for focused tests. +enum ProcNetTCPListeningPortParser { + /// The `st` column value for a socket in the LISTEN state. + private static let listenState = "0A" + + /// Extracts the socket inode from a `/proc//fd` symlink destination such + /// as `socket:[12345]`. Returns nil for non-socket descriptors. + static func socketInode(fromLink destination: String) -> String? { + let prefix = "socket:[" + guard destination.hasPrefix(prefix), destination.hasSuffix("]") else { return nil } + let inode = destination.dropFirst(prefix.count).dropLast() + return inode.isEmpty ? nil : String(inode) + } + + /// Returns the local ports of LISTEN sockets whose inode is in `socketInodes`. + /// + /// `content` is the raw text of a process-scoped `tcp` or `tcp6` table. Each + /// row encodes the local endpoint as `ADDRESS:PORT` (for example, + /// `0100007F:1F90` uses port 8080) and the owning socket inode in column ten. + static func listeningPorts(_ content: String, socketInodes: Set) -> Set { + var ports: Set = [] + for line in content.split(separator: "\n") { + let columns = line.split(separator: " ", omittingEmptySubsequences: true) + // Columns: sl local_address rem_address st ... uid timeout inode + guard columns.count > 9, + columns[3] == self.listenState, + socketInodes.contains(String(columns[9])) + else { continue } + let localAddress = columns[1] + guard let separator = localAddress.lastIndex(of: ":"), + let port = Int(localAddress[localAddress.index(after: separator)...], radix: 16), + (0...Int(UInt16.max)).contains(port) + else { continue } + ports.insert(port) + } + return ports + } +} + +extension AntigravityStatusProbe { + /// Resolves the TCP ports the process `pid` is listening on. Uses `lsof` when + /// present (the common denominator across macOS and Linux) and falls back to + /// the kernel's `/proc` interface on Linux hosts without `lsof`. + static func listeningPorts(pid: Int, timeout: TimeInterval) async throws -> [Int] { + let lsof = ["/usr/sbin/lsof", "/usr/bin/lsof"].first(where: { + FileManager.default.isExecutableFile(atPath: $0) + }) + + if let lsof { + return try await Self.lsofListeningPorts(lsof: lsof, pid: pid, timeout: timeout) + } + + #if os(Linux) + // `lsof` is frequently absent on minimal Linux hosts. Fall back to the + // kernel's /proc interface, mirroring the /proc//cwd fallback that + // LocalAgentSessionScanner.cwdByPID already uses when lsof is missing. + let ports = Self.procListeningPorts(pid: pid) + if ports.isEmpty { + throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found") + } + return ports + #else + throw AntigravityStatusProbeError.portDetectionFailed("lsof not available") + #endif + } + + private static func lsofListeningPorts( + lsof: String, + pid: Int, + timeout: TimeInterval) async throws -> [Int] + { + let env = ProcessInfo.processInfo.environment + let result: SubprocessResult + do { + result = try await SubprocessRunner.run( + binary: lsof, + arguments: ["-nP", "-iTCP", "-sTCP:LISTEN", "-a", "-p", String(pid)], + environment: env, + timeout: timeout, + label: "antigravity-lsof") + } catch let SubprocessRunnerError.nonZeroExit(code, stderr) + where code == 1 && stderr.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found") + } + let ports = Self.parseListeningPorts(result.stdout) + if ports.isEmpty { + throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found") + } + return ports + } + + private static func parseListeningPorts(_ output: String) -> [Int] { + guard let regex = try? NSRegularExpression(pattern: #":(\d+)\s+\(LISTEN\)"#) else { return [] } + let range = NSRange(output.startIndex.. = [] + regex.enumerateMatches(in: output, options: [], range: range) { match, _, _ in + guard let match, + let range = Range(match.range(at: 1), in: output), + let value = Int(output[range]) else { return } + ports.insert(value) + } + return ports.sorted() + } + + /// Recovers the listening ports owned by `pid` by matching its open socket + /// inodes against the TCP tables from the same process/network namespace. + static func procListeningPorts(pid: Int, procRoot: String = "/proc") -> [Int] { + let processRoot = "\(procRoot)/\(pid)" + let inodes = Self.socketInodes(processRoot: processRoot) + guard !inodes.isEmpty else { return [] } + var ports: Set = [] + for path in ["\(processRoot)/net/tcp", "\(processRoot)/net/tcp6"] { + guard let content = try? String(contentsOfFile: path, encoding: .utf8) else { continue } + ports.formUnion(ProcNetTCPListeningPortParser.listeningPorts(content, socketInodes: inodes)) + } + return ports.sorted() + } + + /// Collects the socket inodes referenced by the process's open descriptors. + private static func socketInodes(processRoot: String) -> Set { + let fdDirectory = "\(processRoot)/fd" + guard let entries = try? FileManager.default.contentsOfDirectory(atPath: fdDirectory) else { return [] } + var inodes: Set = [] + for entry in entries { + guard let destination = try? FileManager.default.destinationOfSymbolicLink( + atPath: "\(fdDirectory)/\(entry)"), + let inode = ProcNetTCPListeningPortParser.socketInode(fromLink: destination) + else { continue } + inodes.insert(inode) + } + return inodes + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+ResponseModels.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+ResponseModels.swift new file mode 100644 index 0000000..275b9d3 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe+ResponseModels.swift @@ -0,0 +1,111 @@ +import Foundation + +struct UserStatusResponse: Decodable { + let code: CodeValue? + let message: String? + let userStatus: UserStatus? +} + +struct CommandModelConfigResponse: Decodable { + let code: CodeValue? + let message: String? + let clientModelConfigs: [ModelConfig]? +} + +struct UserStatus: Decodable { + let email: String? + let planStatus: PlanStatus? + let cascadeModelConfigData: ModelConfigData? + let userTier: UserTier? +} + +struct UserTier: Decodable { + let id: String? + let name: String? + let description: String? + + var preferredName: String? { + guard let value = self.name?.trimmingCharacters(in: .whitespacesAndNewlines) else { return nil } + return value.isEmpty ? nil : value + } +} + +struct PlanStatus: Decodable { + let planInfo: PlanInfo? +} + +struct PlanInfo: Decodable { + let planName: String? + let planDisplayName: String? + let displayName: String? + let productName: String? + let planShortName: String? + + var preferredName: String? { + let candidates = [ + self.planDisplayName, + self.displayName, + self.productName, + self.planName, + self.planShortName, + ] + for candidate in candidates { + guard let value = candidate?.trimmingCharacters(in: .whitespacesAndNewlines) else { continue } + if !value.isEmpty { return value } + } + return nil + } +} + +struct ModelConfigData: Decodable { + let clientModelConfigs: [ModelConfig]? +} + +struct ModelConfig: Decodable { + let label: String + let modelOrAlias: ModelAlias + let quotaInfo: QuotaInfo? +} + +struct ModelAlias: Decodable { + let model: String +} + +struct QuotaInfo: Decodable { + let remainingFraction: Double? + let resetTime: String? +} + +enum CodeValue: Decodable { + case int(Int) + case string(String) + + var isOK: Bool { + switch self { + case let .int(value): + value == 0 + case let .string(value): + value.lowercased() == "ok" || value.lowercased() == "success" || value == "0" + } + } + + var rawValue: String { + switch self { + case let .int(value): "\(value)" + case let .string(value): value + } + } + + init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if let value = try? container.decode(Int.self) { + self = .int(value) + return + } + if let value = try? container.decode(String.self) { + self = .string(value) + return + } + throw DecodingError.dataCorruptedError(in: container, debugDescription: "Unsupported code type") + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe.swift new file mode 100644 index 0000000..2093f14 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityStatusProbe.swift @@ -0,0 +1,1642 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct AntigravityModelQuota: Sendable { + public let label: String + public let modelId: String + public let remainingFraction: Double? + public let resetTime: Date? + public let resetDescription: String? + + public init( + label: String, + modelId: String, + remainingFraction: Double?, + resetTime: Date?, + resetDescription: String?) + { + self.label = label + self.modelId = modelId + self.remainingFraction = remainingFraction + self.resetTime = resetTime + self.resetDescription = resetDescription + } + + public var remainingPercent: Double { + guard let remainingFraction else { return 0 } + return max(0, min(100, remainingFraction * 100)) + } +} + +private enum AntigravityModelFamily { + case claude + case gpt + case geminiPro + case geminiFlash + case unknown +} + +private enum AntigravityUsagePool: Hashable { + case gemini + case claudeGPT + + var id: String { + switch self { + case .gemini: "antigravity-gemini" + case .claudeGPT: "antigravity-claude-gpt" + } + } + + var title: String { + switch self { + case .gemini: "Gemini Models" + case .claudeGPT: "Claude and GPT models" + } + } + + var sortRank: Int { + switch self { + case .gemini: 0 + case .claudeGPT: 1 + } + } +} + +private struct AntigravityModelVersion: Comparable { + let major: Int + let minor: Int + + static func < (lhs: AntigravityModelVersion, rhs: AntigravityModelVersion) -> Bool { + if lhs.major != rhs.major { return lhs.major < rhs.major } + return lhs.minor < rhs.minor + } +} + +private struct AntigravityNormalizedModel { + let quota: AntigravityModelQuota + let family: AntigravityModelFamily + let selectionPriority: Int? + let isImage: Bool + let isLite: Bool + let isAutocomplete: Bool + let version: AntigravityModelVersion? + let tier: Int +} + +public enum AntigravityModelQuotaSource: Sendable { + case local + case remote +} + +public struct AntigravityStatusSnapshot: Sendable { + public let modelQuotas: [AntigravityModelQuota] + public let accountEmail: String? + public let accountPlan: String? + public let source: AntigravityModelQuotaSource + let quotaSummary: AntigravityQuotaSummary? + + public init( + modelQuotas: [AntigravityModelQuota], + accountEmail: String?, + accountPlan: String?, + source: AntigravityModelQuotaSource = .remote) + { + self.modelQuotas = modelQuotas + self.accountEmail = accountEmail + self.accountPlan = accountPlan + self.source = source + self.quotaSummary = nil + } + + init( + quotaSummary: AntigravityQuotaSummary, + accountEmail: String?, + accountPlan: String?, + source: AntigravityModelQuotaSource = .local) + { + self.modelQuotas = [] + self.accountEmail = accountEmail + self.accountPlan = accountPlan + self.source = source + self.quotaSummary = quotaSummary + } + + public func toUsageSnapshot() throws -> UsageSnapshot { + if let quotaSummary { + return try Self.usageSnapshot( + from: quotaSummary, + accountEmail: self.accountEmail, + accountPlan: self.accountPlan) + } + + guard !self.modelQuotas.isEmpty else { + throw AntigravityStatusProbeError.parseFailed("No quota models available") + } + + let normalized = Self.normalizedModels(self.modelQuotas) + let summaryCandidates = normalized.filter(Self.isSummaryCandidate) + let primaryQuota = Self.representative(for: .gemini, in: summaryCandidates) + let secondaryQuota = Self.representative(for: .claudeGPT, in: summaryCandidates) + let fallbackQuota: AntigravityModelQuota? = if primaryQuota == nil, secondaryQuota == nil { + switch self.source { + case .local: + Self.fallbackRepresentative(in: normalized.filter { + $0.family == .unknown && + Self.isSelectableTextModel($0) && + $0.quota.remainingFraction != nil + }) + case .remote: + nil + } + } else { + nil + } + + let primary = primaryQuota.map(Self.rateWindow(for:)) + let secondary = secondaryQuota.map(Self.rateWindow(for:)) + let extraWindows = Self.extraRateWindows( + from: normalized, + summaryCandidates: summaryCandidates, + compactFallbackModelID: fallbackQuota?.modelId, + representedPools: Set([ + primaryQuota.map { _ in AntigravityUsagePool.gemini }, + secondaryQuota.map { _ in AntigravityUsagePool.claudeGPT }, + ].compactMap(\.self))) + + let identity = ProviderIdentitySnapshot( + providerID: .antigravity, + accountEmail: self.accountEmail, + accountOrganization: nil, + loginMethod: self.accountPlan) + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: nil, + extraRateWindows: extraWindows.isEmpty ? nil : extraWindows, + updatedAt: Date(), + identity: identity) + } + + func withIdentity(from snapshot: AntigravityStatusSnapshot?) -> AntigravityStatusSnapshot { + guard let snapshot else { return self } + let accountEmail = snapshot.accountEmail ?? self.accountEmail + let accountPlan = snapshot.accountPlan ?? self.accountPlan + if let quotaSummary { + return AntigravityStatusSnapshot( + quotaSummary: quotaSummary, + accountEmail: accountEmail, + accountPlan: accountPlan, + source: self.source) + } + return AntigravityStatusSnapshot( + modelQuotas: self.modelQuotas, + accountEmail: accountEmail, + accountPlan: accountPlan, + source: self.source) + } + + private static func usageSnapshot( + from quotaSummary: AntigravityQuotaSummary, + accountEmail: String?, + accountPlan: String?) throws -> UsageSnapshot + { + let namedWindows = Self.quotaSummaryWindows(from: quotaSummary) + guard !namedWindows.isEmpty else { + throw AntigravityStatusProbeError.parseFailed("No quota buckets available") + } + + let primary = Self.quotaSummaryRepresentative( + matching: { $0.lowercased().contains("gemini") }, + in: namedWindows) + let secondary = Self.quotaSummaryRepresentative( + matching: { $0.lowercased().contains("claude") || $0.lowercased().contains("gpt") }, + in: namedWindows) + + let identity = ProviderIdentitySnapshot( + providerID: .antigravity, + accountEmail: accountEmail, + accountOrganization: nil, + loginMethod: accountPlan) + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: nil, + extraRateWindows: namedWindows, + updatedAt: Date(), + identity: identity) + } + + private static func quotaSummaryRepresentative( + matching predicate: (String) -> Bool, + in windows: [NamedRateWindow]) -> RateWindow? + { + windows + .filter { $0.usageKnown && predicate($0.title) } + .max { lhs, rhs in + if lhs.window.usedPercent != rhs.window.usedPercent { + return lhs.window.usedPercent < rhs.window.usedPercent + } + return lhs.title.localizedCaseInsensitiveCompare(rhs.title) == .orderedDescending + }? + .window + } + + private static func quotaSummaryWindows(from quotaSummary: AntigravityQuotaSummary) -> [NamedRateWindow] { + let sortedGroups = quotaSummary.groups.enumerated().sorted { lhs, rhs in + let lhsRank = Self.quotaGroupSortRank(lhs.element) + let rhsRank = Self.quotaGroupSortRank(rhs.element) + if lhsRank != rhsRank { + return lhsRank < rhsRank + } + return lhs.offset < rhs.offset + }.map(\.element) + + return sortedGroups.flatMap { group in + let groupTitle = Self.displayTitle(forQuotaGroup: group) + let sortedBuckets = group.buckets.enumerated().sorted { lhs, rhs in + let lhsRank = Self.quotaBucketSortRank(lhs.element) + let rhsRank = Self.quotaBucketSortRank(rhs.element) + if lhsRank != rhsRank { + return lhsRank < rhsRank + } + return lhs.offset < rhs.offset + }.map(\.element) + + return sortedBuckets.map { bucket in + let bucketTitle = Self.displayTitle(forQuotaBucket: bucket) + let remainingPercent = Self.remainingPercent(from: bucket.remainingFraction) + let usedPercent = remainingPercent.map { 100 - $0 } ?? 0 + let window = RateWindow( + usedPercent: usedPercent, + windowMinutes: Self.windowMinutes(forQuotaBucket: bucket), + resetsAt: bucket.resetTime, + resetDescription: bucket.resetDescription) + return NamedRateWindow( + id: Self.quotaSummaryWindowID(for: bucket), + title: "\(groupTitle) \(bucketTitle)", + window: window, + usageKnown: !bucket.disabled && bucket.remainingFraction != nil) + } + } + } + + static func isQuotaSummaryWindowID(_ id: String) -> Bool { + id.hasPrefix(self.quotaSummaryWindowIDPrefix) + } + + private static let quotaSummaryWindowIDPrefix = "antigravity-quota-summary-" + + private static func quotaSummaryWindowID(for bucket: AntigravityQuotaSummaryBucket) -> String { + self.quotaSummaryWindowIDPrefix + bucket.bucketId + } + + private static func remainingPercent(from remainingFraction: Double?) -> Double? { + guard let remainingFraction else { return nil } + return max(0, min(100, remainingFraction * 100)) + } + + private static func displayTitle(forQuotaGroup group: AntigravityQuotaSummaryGroup) -> String { + let title = group.displayName.trimmingCharacters(in: .whitespacesAndNewlines) + let lowercasedTitle = title.lowercased() + if lowercasedTitle.contains("gemini") { + return "Gemini" + } + if lowercasedTitle.contains("claude") || lowercasedTitle.contains("gpt") { + return "Claude/GPT" + } + return title.isEmpty ? "Quota" : title + } + + private static func displayTitle(forQuotaBucket bucket: AntigravityQuotaSummaryBucket) -> String { + switch self.quotaBucketKind(for: bucket) { + case .session: + "5-hour" + case .weekly: + "weekly" + case .other: + bucket.displayName + } + } + + private static func windowMinutes(forQuotaBucket bucket: AntigravityQuotaSummaryBucket) -> Int? { + switch self.quotaBucketKind(for: bucket) { + case .session: + 300 + case .weekly: + 10080 + case .other: + nil + } + } + + private enum QuotaBucketKind { + case session + case weekly + case other + } + + private static func quotaGroupSortRank(_ group: AntigravityQuotaSummaryGroup) -> Int { + let title = group.displayName.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + if title.contains("gemini") { + return 0 + } + if title.contains("claude") || title.contains("gpt") { + return 1 + } + return 2 + } + + private static func quotaBucketSortRank(_ bucket: AntigravityQuotaSummaryBucket) -> Int { + switch self.quotaBucketKind(for: bucket) { + case .session: + 0 + case .weekly: + 1 + case .other: + 2 + } + } + + private static func quotaBucketKind(for bucket: AntigravityQuotaSummaryBucket) -> QuotaBucketKind { + let combined = "\(bucket.bucketId) \(bucket.displayName)".lowercased() + if combined.contains("5h") || combined.contains("5-hour") || combined.contains("five hour") { + return .session + } + if combined.contains("weekly") { + return .weekly + } + return .other + } + + static func quotaDisplayLabel(_ quota: AntigravityModelQuota) -> String { + let trimmed = quota.label.trimmingCharacters(in: .whitespacesAndNewlines) + guard trimmed.isEmpty || trimmed == quota.modelId else { return quota.label } + return Self.humanizedModelID(quota.modelId) + } + + static func humanizedModelID(_ modelId: String) -> String { + let parts = modelId.split(separator: "-").map(String.init) + var words: [String] = [] + var index = 0 + + while index < parts.count { + let part = parts[index] + if Self.isSingleDigit(part), index + 1 < parts.count, Self.isSingleDigit(parts[index + 1]) { + var versionParts = [part] + index += 1 + while index < parts.count, Self.isSingleDigit(parts[index]) { + versionParts.append(parts[index]) + index += 1 + } + words.append(versionParts.joined(separator: ".")) + continue + } + + if part.allSatisfy({ $0.isNumber || $0 == "." }) { + words.append(part) + } else if Self.modelLabelAcronyms.contains(part.lowercased()) { + words.append(part.uppercased()) + } else { + words.append(part.prefix(1).uppercased() + part.dropFirst()) + } + index += 1 + } + + return words.joined(separator: " ") + } + + private static let modelLabelAcronyms: Set = ["ai", "api", "gpt", "oss"] + + private static func isSingleDigit(_ value: String) -> Bool { + value.count == 1 && value.allSatisfy(\.isNumber) + } + + private static func rateWindow(for quota: AntigravityModelQuota) -> RateWindow { + RateWindow( + usedPercent: 100 - quota.remainingPercent, + windowMinutes: nil, + resetsAt: quota.resetTime, + resetDescription: quota.resetDescription) + } + + private static func modelOrderPrecedes( + _ lhs: AntigravityNormalizedModel, + _ rhs: AntigravityNormalizedModel) -> Bool + { + // 1. Family rank: claude=0, geminiPro=1, geminiFlash=2, unknown=3 + let lhsFamilyRank = Self.familyRank(lhs.family) + let rhsFamilyRank = Self.familyRank(rhs.family) + if lhsFamilyRank != rhsFamilyRank { + return lhsFamilyRank < rhsFamilyRank + } + + // 2. Version descending (newer first); nil version sorts after non-nil + switch (lhs.version, rhs.version) { + case let (.some(lv), .some(rv)): + if lv != rv { + return lv > rv + } + case (.some, .none): + return true + case (.none, .some): + return false + case (.none, .none): + break + } + + // 3. Tier ascending: High(0) < Medium(1) < Low(2) + if lhs.tier != rhs.tier { + return lhs.tier < rhs.tier + } + + // 4. Label tiebreaker + return lhs.quota.label.localizedCaseInsensitiveCompare(rhs.quota.label) == .orderedAscending + } + + private static func familyRank(_ family: AntigravityModelFamily) -> Int { + switch family { + case .claude: 0 + case .gpt: 1 + case .geminiPro: 2 + case .geminiFlash: 3 + case .unknown: 4 + } + } + + private static func isSummaryCandidate(_ model: AntigravityNormalizedModel) -> Bool { + self.usagePool(for: model) != nil && self.isSelectableTextModel(model) + } + + private static func isSelectableTextModel(_ model: AntigravityNormalizedModel) -> Bool { + !model.isLite && !model.isAutocomplete && !model.isImage + } + + private static func normalizedModels(_ models: [AntigravityModelQuota]) -> [AntigravityNormalizedModel] { + models.map { self.normalizeModel($0) } + } + + private static func normalizeModel(_ quota: AntigravityModelQuota) -> AntigravityNormalizedModel { + let modelId = quota.modelId.lowercased() + let label = quota.label.lowercased() + let family = Self.family(forModelID: modelId, label: label) + + let isLite = modelId.contains("lite") || label.contains("lite") + let isAutocomplete = modelId.contains("autocomplete") || label.contains("autocomplete") || modelId + .hasPrefix("tab_") + let isImage = modelId.contains("image") || label.contains("image") + let isSelectableTextModel = !isLite && !isAutocomplete && !isImage + let isLowPriorityGeminiPro = modelId.contains("pro-low") + || (label.contains("pro") && label.contains("low")) + + let selectionPriority: Int? = switch family { + case .claude, .gpt: + 0 + case .geminiPro: + if isLowPriorityGeminiPro, isSelectableTextModel { + 0 + } else if isSelectableTextModel { + 1 + } else { + nil + } + case .geminiFlash: + isSelectableTextModel ? 0 : nil + case .unknown: + nil + } + + let version = Self.parseVersion(from: label) + let tier = Self.parseTier(from: label, modelId: modelId) + + return AntigravityNormalizedModel( + quota: quota, + family: family, + selectionPriority: selectionPriority, + isImage: isImage, + isLite: isLite, + isAutocomplete: isAutocomplete, + version: version, + tier: tier) + } + + private static func parseVersion(from label: String) -> AntigravityModelVersion? { + // Accept either "." or "-" between major and minor so a raw model id used as the + // label when displayName is missing (e.g. "gemini-3-1-pro-low") still parses 3.1. + guard let regex = try? NSRegularExpression(pattern: #"(\d+)(?:[.\-](\d+))?"#) else { return nil } + let nsLabel = label as NSString + let range = NSRange(location: 0, length: nsLabel.length) + guard let match = regex.firstMatch(in: label, options: [], range: range) else { return nil } + let majorRange = Range(match.range(at: 1), in: label) + guard let majorRange, let major = Int(label[majorRange]) else { return nil } + let minor: Int = if match.range(at: 2).location != NSNotFound, + let minorRange = Range(match.range(at: 2), in: label), + let parsed = Int(label[minorRange]) + { + parsed + } else { + 0 + } + return AntigravityModelVersion(major: major, minor: minor) + } + + private static func parseTier(from label: String, modelId: String) -> Int { + let combined = label + " " + modelId + if combined.contains("high") { return 0 } + if combined.contains("medium") { return 1 } + if combined.contains("low") { return 2 } + return 1 + } + + private static func representative( + for pool: AntigravityUsagePool, + in models: [AntigravityNormalizedModel]) -> AntigravityModelQuota? + { + let candidates = models.filter { + Self.usagePool(for: $0) == pool && $0.quota.remainingFraction != nil + } + guard !candidates.isEmpty else { return nil } + return candidates.min { lhs, rhs in + if lhs.quota.remainingPercent != rhs.quota.remainingPercent { + return lhs.quota.remainingPercent < rhs.quota.remainingPercent + } + switch (lhs.quota.resetTime, rhs.quota.resetTime) { + case let (.some(left), .some(right)) where left != right: + return left < right + case (.some, .none): + return true + case (.none, .some): + return false + default: + return lhs.quota.label.localizedCaseInsensitiveCompare(rhs.quota.label) == .orderedAscending + } + }?.quota + } + + private static func fallbackRepresentative(in models: [AntigravityNormalizedModel]) -> AntigravityModelQuota? { + models.min { lhs, rhs in + if lhs.quota.remainingPercent != rhs.quota.remainingPercent { + return lhs.quota.remainingPercent < rhs.quota.remainingPercent + } + return lhs.quota.label.localizedCaseInsensitiveCompare(rhs.quota.label) == .orderedAscending + }?.quota + } + + private static func extraRateWindows( + from models: [AntigravityNormalizedModel], + summaryCandidates: [AntigravityNormalizedModel], + compactFallbackModelID: String?, + representedPools: Set) -> [NamedRateWindow] + { + let resetOnlyPoolWindows = [AntigravityUsagePool.gemini, .claudeGPT].compactMap { pool -> NamedRateWindow? in + guard !representedPools.contains(pool) else { return nil } + let candidates = summaryCandidates.filter { Self.usagePool(for: $0) == pool } + guard let resetOnly = candidates.first(where: { model in + model.quota.remainingFraction == nil && + (model.quota.resetTime != nil || model.quota.resetDescription != nil) + }) else { + return nil + } + return NamedRateWindow( + id: pool.id, + title: pool.title, + window: Self.rateWindow(for: resetOnly.quota), + usageKnown: false) + } + + let distinctWindows = models + .filter { + $0.quota.modelId == compactFallbackModelID || Self.shouldShowDistinctExtraWindow($0) + } + .sorted(by: Self.modelOrderPrecedes) + .map { m in + NamedRateWindow( + id: m.quota.modelId == compactFallbackModelID + ? Self.compactFallbackWindowID(modelID: m.quota.modelId) + : m.quota.modelId, + title: Self.quotaDisplayLabel(m.quota), + window: Self.rateWindow(for: m.quota), + usageKnown: m.quota.remainingFraction != nil) + } + + return resetOnlyPoolWindows.sorted { lhs, rhs in + guard let lhsPool = Self.pool(forExtraWindowID: lhs.id), + let rhsPool = Self.pool(forExtraWindowID: rhs.id) + else { + return lhs.title.localizedCaseInsensitiveCompare(rhs.title) == .orderedAscending + } + return lhsPool.sortRank < rhsPool.sortRank + } + distinctWindows + } + + private static func compactFallbackWindowID(modelID: String) -> String { + "antigravity-compact-fallback-\(modelID)" + } + + private static func shouldShowDistinctExtraWindow(_ model: AntigravityNormalizedModel) -> Bool { + guard !self.isSummaryCandidate(model) else { return false } + if model.quota.remainingFraction == nil { + return model.quota.resetTime != nil || model.quota.resetDescription != nil + } + return model.quota.remainingPercent < 99.9 + } + + private static func pool(forExtraWindowID id: String) -> AntigravityUsagePool? { + switch id { + case AntigravityUsagePool.gemini.id: .gemini + case AntigravityUsagePool.claudeGPT.id: .claudeGPT + default: nil + } + } + + private static func usagePool(for model: AntigravityNormalizedModel) -> AntigravityUsagePool? { + switch model.family { + case .geminiPro, .geminiFlash: + .gemini + case .claude, .gpt: + .claudeGPT + case .unknown: + nil + } + } + + private static func family(forModelID modelId: String, label: String) -> AntigravityModelFamily { + let modelIDFamily = Self.family(from: modelId) + if modelIDFamily != .unknown { + return modelIDFamily + } + return Self.family(from: label) + } + + private static func family(from text: String) -> AntigravityModelFamily { + if text.contains("claude") { + return .claude + } + if text.contains("gpt") || text.contains("openai") { + return .gpt + } + if text.contains("gemini"), text.contains("pro") { + return .geminiPro + } + if text.contains("gemini"), text.contains("flash") { + return .geminiFlash + } + return .unknown + } +} + +public struct AntigravityPlanInfoSummary: Sendable, Codable, Equatable { + public let planName: String? + public let planDisplayName: String? + public let displayName: String? + public let productName: String? + public let planShortName: String? +} + +public enum AntigravityStatusProbeError: LocalizedError, Sendable, Equatable { + case notRunning + case missingCSRFToken + case portDetectionFailed(String) + case apiError(String) + case parseFailed(String) + case timedOut + case authenticationRequired + case accountMismatch(expected: String?, found: String?) + + public var errorDescription: String? { + switch self { + case .notRunning: + "Antigravity language server not detected. Launch Antigravity and retry." + case .missingCSRFToken: + "Antigravity CSRF token not found. Restart Antigravity and retry." + case let .portDetectionFailed(message): + Self.portDetectionDescription(message) + case let .apiError(message): + Self.apiErrorDescription(message) + case let .parseFailed(message): + "Could not parse Antigravity quota: \(message)" + case .timedOut: + "Antigravity quota request timed out." + case .authenticationRequired: + "Antigravity CLI is signed out. Run agy in a terminal to sign in, then retry." + case let .accountMismatch(expected, found): + Self.accountMismatchDescription(expected: expected, found: found) + } + } + + private static func accountMismatchDescription(expected: String?, found: String?) -> String { + let selected = expected ?? "the selected account" + if let found { + return "Antigravity local session is signed in as \(found), not \(selected); " + + "using the selected account's OAuth data instead." + } + return "Antigravity local session did not report an account matching \(selected); " + + "using the selected account's OAuth data instead." + } + + private static func portDetectionDescription(_ message: String) -> String { + switch message { + case "lsof not available": + "Antigravity port detection needs lsof. Install it, then retry." + case "no listening ports found": + "Antigravity is running but not exposing ports yet. Try again in a few seconds." + default: + "Antigravity port detection failed: \(message)" + } + } + + private static func apiErrorDescription(_ message: String) -> String { + if message.contains("HTTP 401") || message.contains("HTTP 403") { + return "Antigravity session expired. Restart Antigravity and retry." + } + return "Antigravity API error: \(message)" + } +} + +public struct AntigravityStatusProbe: Sendable { + /// Which local Antigravity processes the probe may attach to. + public enum ProcessScope: Sendable { + /// Match Antigravity app, Antigravity IDE, and the `agy` CLI language server. + case ideAndCLI + /// Match only the Antigravity 2.0 app language server. + case appOnly + /// Match only the Antigravity IDE extension language server. + case ideOnly + } + + public var timeout: TimeInterval = 8.0 + public var processScope: ProcessScope = .ideAndCLI + + private static let getUserStatusPath = "/exa.language_server_pb.LanguageServerService/GetUserStatus" + private static let commandModelConfigPath = + "/exa.language_server_pb.LanguageServerService/GetCommandModelConfigs" + private static let quotaSummaryPath = + "/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary" + private static let unleashPath = "/exa.language_server_pb.LanguageServerService/GetUnleashData" + private static let log = CodexBarLog.logger(LogCategories.antigravity) + + public init(timeout: TimeInterval = 8.0, processScope: ProcessScope = .ideAndCLI) { + self.timeout = timeout + self.processScope = processScope + } + + public func fetch(matchingAccountEmail expectedAccountEmail: String? = nil) async throws + -> AntigravityStatusSnapshot + { + let deadline = Date().addingTimeInterval(self.timeout) + let processInfos = try await Self.detectProcessInfos(timeout: self.timeout, scope: self.processScope) + let result = try await Self.fetchProcessSnapshots(processInfos: processInfos) { processInfo in + try await Self.fetch( + processInfo: processInfo, + timeout: self.timeout, + deadline: deadline) + } + + if let bestSnapshot = Self.preferredLocalSnapshot( + result.snapshots, + matchingAccountEmail: expectedAccountEmail) + { + return bestSnapshot + } + throw result.lastError ?? AntigravityStatusProbeError.notRunning + } + + public func fetchPlanInfoSummary() async throws -> AntigravityPlanInfoSummary? { + let processInfo = try await Self.detectProcessInfo(timeout: self.timeout) + let ports = try await Self.listeningPorts(pid: processInfo.pid, timeout: self.timeout) + let endpoint = try await Self.resolveWorkingEndpoint( + candidateEndpoints: Self.connectionCandidates( + listeningPorts: ports, + languageServerCSRFToken: processInfo.csrfToken, + extensionServerPort: processInfo.extensionPort, + extensionServerCSRFToken: processInfo.extensionServerCSRFToken), + timeout: self.timeout) + return try await Self.makeParsedRequest( + payload: RequestPayload( + path: Self.getUserStatusPath, + body: Self.defaultRequestBody()), + context: RequestContext( + endpoints: Self.requestEndpoints( + resolvedEndpoint: endpoint, + listeningPorts: ports, + languageServerCSRFToken: processInfo.csrfToken, + extensionServerPort: processInfo.extensionPort, + extensionServerCSRFToken: processInfo.extensionServerCSRFToken), + timeout: self.timeout), + parse: Self.parsePlanInfoSummary) + } + + static func localSnapshotScore(_ snapshot: AntigravityStatusSnapshot) -> Int { + var score = 0 + if let quotaSummary = snapshot.quotaSummary { + let buckets = quotaSummary.groups.flatMap(\.buckets) + let knownBuckets = buckets.count(where: { !$0.disabled && $0.remainingFraction != nil }) + score += 1000 + score += quotaSummary.groups.count * 10 + score += buckets.count + score += knownBuckets * 20 + } else { + let knownRows = snapshot.modelQuotas.count(where: { $0.remainingFraction != nil }) + score += snapshot.modelQuotas.count + score += knownRows * 10 + } + if snapshot.accountEmail != nil { + score += 2 + } + if snapshot.accountPlan != nil { + score += 1 + } + return score + } + + public static func isRunning(timeout: TimeInterval = 4.0) async -> Bool { + await (try? self.detectProcessInfo(timeout: timeout)) != nil + } + + public static func detectVersion(timeout: TimeInterval = 4.0) async -> String? { + let running = await Self.isRunning(timeout: timeout) + return running ? "running" : nil + } + + // MARK: - CLI Local Fetch + + /// Fetch usage data from a known set of local ports (discovered via + /// ``AntigravityCLISession``'s ``pid``), without requiring a running + /// ``language_server`` process or CSRF token. + /// + /// The ``agy`` CLI exposes the same ``GetUserStatus`` gRPC-web endpoint as + /// the desktop ``language_server``. Unlike the desktop endpoint, it does + /// not require a CSRF token header. + public func fetchFromPorts(_ ports: [Int], deadline: Date? = nil) async throws -> AntigravityStatusSnapshot { + guard !ports.isEmpty else { + throw AntigravityStatusProbeError.portDetectionFailed("no listening ports found") + } + let endpoints = Self.cliEndpoints(ports: ports) + let context = RequestContext(endpoints: endpoints, timeout: self.timeout, deadline: deadline) + return try await Self.fetchSnapshot(context: context) + } + + // MARK: - Parsing + + public static func parseUserStatusResponse(_ data: Data) throws -> AntigravityStatusSnapshot { + let decoder = JSONDecoder() + let response = try decoder.decode(UserStatusResponse.self, from: data) + if let invalid = Self.invalidCode(response.code) { + throw AntigravityStatusProbeError.apiError(invalid) + } + guard let userStatus = response.userStatus else { + throw AntigravityStatusProbeError.parseFailed("Missing userStatus") + } + + let modelConfigs = userStatus.cascadeModelConfigData?.clientModelConfigs ?? [] + let models = modelConfigs.compactMap(Self.quotaFromConfig(_:)) + let email = userStatus.email + // Prefer userTier.name (actual subscription tier) over planInfo (shows "Pro" for Ultra users) + let planName = userStatus.userTier?.preferredName ?? userStatus.planStatus?.planInfo?.preferredName + + return AntigravityStatusSnapshot( + modelQuotas: models, + accountEmail: email, + accountPlan: planName, + source: .local) + } + + static func parsePlanInfoSummary(_ data: Data) throws -> AntigravityPlanInfoSummary? { + let decoder = JSONDecoder() + let response = try decoder.decode(UserStatusResponse.self, from: data) + if let invalid = Self.invalidCode(response.code) { + throw AntigravityStatusProbeError.apiError(invalid) + } + guard let userStatus = response.userStatus else { + throw AntigravityStatusProbeError.parseFailed("Missing userStatus") + } + guard let planInfo = userStatus.planStatus?.planInfo else { return nil } + return AntigravityPlanInfoSummary( + planName: planInfo.planName, + planDisplayName: planInfo.planDisplayName, + displayName: planInfo.displayName, + productName: planInfo.productName, + planShortName: planInfo.planShortName) + } + + static func parseCommandModelResponse(_ data: Data) throws -> AntigravityStatusSnapshot { + let decoder = JSONDecoder() + let response = try decoder.decode(CommandModelConfigResponse.self, from: data) + if let invalid = Self.invalidCode(response.code) { + throw AntigravityStatusProbeError.apiError(invalid) + } + let modelConfigs = response.clientModelConfigs ?? [] + let models = modelConfigs.compactMap(Self.quotaFromConfig(_:)) + return AntigravityStatusSnapshot(modelQuotas: models, accountEmail: nil, accountPlan: nil, source: .local) + } + + private static func quotaFromConfig(_ config: ModelConfig) -> AntigravityModelQuota? { + guard let quota = config.quotaInfo else { return nil } + let reset = quota.resetTime.flatMap { Self.parseDate($0) } + return AntigravityModelQuota( + label: config.label, + modelId: config.modelOrAlias.model, + remainingFraction: quota.remainingFraction, + resetTime: reset, + resetDescription: nil) + } + + static func invalidCode(_ code: CodeValue?) -> String? { + guard let code else { return nil } + if code.isOK { return nil } + return "\(code.rawValue)" + } + + static func parseDate(_ value: String) -> Date? { + if let date = ISO8601DateFormatter().date(from: value) { + return date + } + if let seconds = Double(value) { + return Date(timeIntervalSince1970: seconds) + } + return nil + } + + // MARK: - Port detection + + struct ProcessInfoResult { + let pid: Int + let extensionPort: Int? + let extensionServerCSRFToken: String? + let csrfToken: String + let commandLine: String + } + + struct AntigravityConnectionEndpoint: Equatable { + enum Source: String { + case languageServer = "language-server" + case extensionServer = "extension-server" + case cliHTTPS = "cli-https" + } + + let scheme: String + let port: Int + let csrfToken: String + let source: Source + /// Whether this endpoint needs a CSRF token header. + /// The CLI HTTPS endpoint (``Source/cliHTTPS``) speaks the same HTTP API + /// but does not require a CSRF token. + var requiresCSRFToken: Bool { + switch self.source { + case .languageServer, .extensionServer: true + case .cliHTTPS: false + } + } + + func matchesRequestTarget(_ other: Self) -> Bool { + self.scheme == other.scheme && self.port == other.port && self.csrfToken == other.csrfToken + } + } + + private static func detectProcessInfo( + timeout: TimeInterval, + scope: ProcessScope = .ideAndCLI) async throws -> ProcessInfoResult + { + let processInfos = try await self.detectProcessInfos(timeout: timeout, scope: scope) + guard let first = processInfos.first else { + throw AntigravityStatusProbeError.notRunning + } + return first + } + + /// Internal (not private): called by AntigravityCLIHTTPSFetchStrategy + /// .liveWarmAgyDependencies to discover an already-running `agy` for reuse. + static func detectProcessInfos( + timeout: TimeInterval, + scope: ProcessScope = .ideAndCLI) async throws -> [ProcessInfoResult] + { + let env = ProcessInfo.processInfo.environment + let result = try await SubprocessRunner.run( + binary: "/bin/ps", + arguments: ["-ax", "-o", "pid=,command="], + environment: env, + timeout: timeout, + label: "antigravity-ps") + + return try Self.processInfos(fromProcessListOutput: result.stdout, scope: scope) + } + + static func processInfo( + fromProcessListOutput output: String, + scope: ProcessScope = .ideAndCLI) throws -> ProcessInfoResult + { + let processInfos = try self.processInfos(fromProcessListOutput: output, scope: scope) + guard let first = processInfos.first else { + throw AntigravityStatusProbeError.notRunning + } + return first + } + + static func processInfos( + fromProcessListOutput output: String, + scope: ProcessScope = .ideAndCLI) throws -> [ProcessInfoResult] + { + let lines = output.split(separator: "\n") + var sawTokenlessIDE = false + var results: [ProcessInfoResult] = [] + for line in lines { + let text = String(line) + guard let match = Self.matchProcessLine(text) else { continue } + guard let kind = Self.antigravityProcessKind(match.command) else { continue } + if !Self.processKind(kind, matches: scope) { continue } + // The IDE language server authenticates local requests with a + // `--csrf_token` and must keep requiring it: skip a tokenless IDE + // or app match so a later valid server can still be found (and surface + // `missingCSRFToken` if none is). The CLI's language server exposes + // no token flag and needs none, so an empty token is allowed there. + guard let token = Self.resolvedCSRFToken(forKind: kind, command: match.command) else { + sawTokenlessIDE = true + continue + } + let port = Self.extractPort("--extension_server_port", from: match.command) + let extensionServerCSRFToken = Self.extractFlag("--extension_server_csrf_token", from: match.command) + results.append(ProcessInfoResult( + pid: match.pid, + extensionPort: port, + extensionServerCSRFToken: extensionServerCSRFToken, + csrfToken: token, + commandLine: match.command)) + } + + if !results.isEmpty { + return results + } + if sawTokenlessIDE { + throw AntigravityStatusProbeError.missingCSRFToken + } + throw AntigravityStatusProbeError.notRunning + } + + private struct ProcessLineMatch { + let pid: Int + let command: String + } + + private static func matchProcessLine(_ line: String) -> ProcessLineMatch? { + let trimmed = line.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + let parts = trimmed.split(separator: " ", maxSplits: 1, omittingEmptySubsequences: true) + guard parts.count == 2, let pid = Int(parts[0]) else { return nil } + return ProcessLineMatch(pid: pid, command: String(parts[1])) + } + + enum AntigravityProcessKind: Equatable { + /// Antigravity 2.0 app language server. Requires a `--csrf_token`. + case app + /// Antigravity IDE extension language server. Requires a `--csrf_token`. + case ide + /// CLI language server (`agy` / `antigravity-cli`). Needs no CSRF token. + case cli + } + + static func isAntigravityLanguageServerCommandLine(_ command: String) -> Bool { + self.antigravityProcessKind(command) != nil + } + + /// Classify a process command line as the Antigravity app language server, + /// the Antigravity IDE language server, the Antigravity CLI language server, + /// or neither. Desktop language servers take precedence so their CSRF-token + /// requirement is preserved. + static func antigravityProcessKind(_ command: String) -> AntigravityProcessKind? { + let lower = command.lowercased() + if Self.isLanguageServerCommandLine(lower), Self.isAntigravityCommandLine(lower) { + return Self.isAntigravityIDECommandLine(lower) ? .ide : .app + } + if Self.isAntigravityCLICommandLine(lower) { + return .cli + } + return nil + } + + private static func processKind(_ kind: AntigravityProcessKind, matches scope: ProcessScope) -> Bool { + switch scope { + case .ideAndCLI: + true + case .appOnly: + kind == .app + case .ideOnly: + kind == .ide + } + } + + /// Resolve the CSRF token to use for a matched process, or `nil` when the + /// match must be skipped. IDE matches keep requiring `--csrf_token` + /// (tokenless IDE matches are skipped). CLI matches accept an empty token + /// because the CLI's language server requires none. + static func resolvedCSRFToken(forKind kind: AntigravityProcessKind, command: String) -> String? { + if let token = extractFlag("--csrf_token", from: command) { + return token + } + switch kind { + case .app, .ide: return nil + case .cli: return "" + } + } + + private static func isLanguageServerCommandLine(_ lowerCommand: String) -> Bool { + let pattern = #"(^|[/\\])language(?:_|-)server(?:[_-][a-z0-9]+)*(?:\.exe)?(\s|$)"# + return lowerCommand.range(of: pattern, options: .regularExpression) != nil + } + + /// The Antigravity CLI (`agy` / `antigravity-cli`) hosts the same language + /// server locally as the IDE, but launches it without a `--csrf_token` flag + /// and under a different process name. Match it so usage can be probed when + /// only the CLI is running. + private static func isAntigravityCLICommandLine(_ lowerCommand: String) -> Bool { + let cliPathPattern = #"(^|[/\\])(antigravity-cli|antigravity_cli)([\s/\\]|$)"# + if lowerCommand.range(of: cliPathPattern, options: .regularExpression) != nil { + return true + } + let agyPattern = #"(^|[/\\])agy(\s|$)"# + return lowerCommand.range(of: agyPattern, options: .regularExpression) != nil + } + + private static func isAntigravityCommandLine(_ command: String) -> Bool { + if command.contains("--app_data_dir") && command.contains("antigravity") { return true } + if command.contains("antigravity.app/") || command.contains("antigravity.app\\") { return true } + if command.contains("antigravity ide.app/") || command.contains("antigravity ide.app\\") { return true } + if command.contains("/antigravity/") || command.contains("\\antigravity\\") { return true } + return false + } + + private static func isAntigravityIDECommandLine(_ lowerCommand: String) -> Bool { + [ + "antigravity ide.app/", + "antigravity ide.app\\", + "--app_data_dir antigravity-ide", + "--app_data_dir=antigravity-ide", + "/extensions/antigravity/bin/language_server", + "\\extensions\\antigravity\\bin\\language_server", + ].contains { lowerCommand.contains($0) } + } + + private static func extractFlag(_ flag: String, from command: String) -> String? { + let pattern = "\(NSRegularExpression.escapedPattern(for: flag))[=\\s]+([^\\s]+)" + guard let regex = try? NSRegularExpression(pattern: pattern, options: .caseInsensitive) else { return nil } + let range = NSRange(command.startIndex.. Int? { + guard let raw = extractFlag(flag, from: command) else { return nil } + return Int(raw) + } + + static func connectionCandidates( + listeningPorts: [Int], + languageServerCSRFToken: String, + extensionServerPort: Int?, + extensionServerCSRFToken: String?) -> [AntigravityConnectionEndpoint] + { + var endpoints = Self.languageServerEndpoints( + listeningPorts: listeningPorts, + languageServerCSRFToken: languageServerCSRFToken) + + for endpoint in Self.extensionServerEndpoints( + extensionServerPort: extensionServerPort, + languageServerCSRFToken: languageServerCSRFToken, + extensionServerCSRFToken: extensionServerCSRFToken) + { + guard !endpoints.contains(where: { $0.matchesRequestTarget(endpoint) }) else { continue } + endpoints.append(endpoint) + } + + return endpoints + } + + static func requestEndpoints( + resolvedEndpoint: AntigravityConnectionEndpoint, + listeningPorts: [Int], + languageServerCSRFToken: String, + extensionServerPort: Int?, + extensionServerCSRFToken: String?) -> [AntigravityConnectionEndpoint] + { + var endpoints = [resolvedEndpoint] + + if resolvedEndpoint.source == .extensionServer { + Self.appendUniqueRequestTargets( + from: Self.extensionServerEndpoints( + extensionServerPort: extensionServerPort, + languageServerCSRFToken: languageServerCSRFToken, + extensionServerCSRFToken: extensionServerCSRFToken), + to: &endpoints) + Self.appendUniqueRequestTargets( + from: Self.languageServerEndpoints( + listeningPorts: listeningPorts, + languageServerCSRFToken: languageServerCSRFToken), + to: &endpoints) + } else { + Self.appendUniqueRequestTargets( + from: Self.languageServerEndpoints( + listeningPorts: listeningPorts, + languageServerCSRFToken: languageServerCSRFToken), + to: &endpoints) + Self.appendUniqueRequestTargets( + from: Self.extensionServerEndpoints( + extensionServerPort: extensionServerPort, + languageServerCSRFToken: languageServerCSRFToken, + extensionServerCSRFToken: extensionServerCSRFToken), + to: &endpoints) + } + + return endpoints + } + + private static func languageServerEndpoints( + listeningPorts: [Int], + languageServerCSRFToken: String) -> [AntigravityConnectionEndpoint] + { + listeningPorts.flatMap { port in + self.localProbeSchemes.map { scheme in + AntigravityConnectionEndpoint( + scheme: scheme, + port: port, + csrfToken: languageServerCSRFToken, + source: .languageServer) + } + } + } + + private static func extensionServerEndpoints( + extensionServerPort: Int?, + languageServerCSRFToken: String, + extensionServerCSRFToken: String?) -> [AntigravityConnectionEndpoint] + { + guard let extensionServerPort else { return [] } + + var endpoints: [AntigravityConnectionEndpoint] = [] + if let extensionServerCSRFToken { + endpoints.append( + AntigravityConnectionEndpoint( + scheme: "http", + port: extensionServerPort, + csrfToken: extensionServerCSRFToken, + source: .extensionServer)) + } + + if extensionServerCSRFToken != languageServerCSRFToken { + endpoints.append( + AntigravityConnectionEndpoint( + scheme: "http", + port: extensionServerPort, + csrfToken: languageServerCSRFToken, + source: .extensionServer)) + } + + return endpoints + } + + private static func appendUniqueRequestTargets( + from candidates: [AntigravityConnectionEndpoint], + to endpoints: inout [AntigravityConnectionEndpoint]) + { + for endpoint in candidates { + guard !endpoints.contains(where: { $0.matchesRequestTarget(endpoint) }) else { continue } + endpoints.append(endpoint) + } + } + + static func resolveWorkingEndpoint( + candidateEndpoints: [AntigravityConnectionEndpoint], + timeout: TimeInterval, + deadline: Date? = nil, + testConnectivity: @escaping @Sendable (AntigravityConnectionEndpoint, TimeInterval) async -> Bool = Self + .testEndpointConnectivity) async throws -> AntigravityConnectionEndpoint + { + for (index, endpoint) in candidateEndpoints.enumerated() { + let remainingAttemptCount = candidateEndpoints.count - index + guard let attemptTimeout = timeoutForEndpointAttempt( + timeout: timeout, + deadline: deadline, + remainingAttemptCount: remainingAttemptCount) + else { + throw AntigravityStatusProbeError.timedOut + } + let ok = await testConnectivity(endpoint, attemptTimeout) + if ok { return endpoint } + } + if let fallback = fallbackProbeEndpoint(candidateEndpoints) { + self.log.debug("Port probe fell back to best-effort endpoint", metadata: [ + "source": fallback.source.rawValue, + "scheme": fallback.scheme, + "port": "\(fallback.port)", + ]) + return fallback + } + throw AntigravityStatusProbeError.portDetectionFailed("no working API port found") + } + + private static func timeoutForEndpointAttempt( + timeout: TimeInterval, + deadline: Date?, + remainingAttemptCount: Int) -> TimeInterval? + { + guard let deadline else { return timeout } + let remaining = deadline.timeIntervalSinceNow + guard remaining > 0 else { return nil } + return min(timeout, remaining / Double(max(1, remainingAttemptCount))) + } + + static func fallbackProbePort(ports: [Int], extensionPort: Int?) -> Int? { + if let nonExtension = ports.first(where: { $0 != extensionPort }) { + return nonExtension + } + if let extensionPort { + return extensionPort + } + return ports.first + } + + static func isReachableProbeError(_ error: Error) -> Bool { + guard case let AntigravityStatusProbeError.apiError(message) = error else { return false } + return message.hasPrefix("HTTP ") + } + + private static func fallbackProbeEndpoint( + _ endpoints: [AntigravityConnectionEndpoint]) -> AntigravityConnectionEndpoint? + { + if let languageServerEndpoint = endpoints.first(where: { $0.source == .languageServer }) { + return languageServerEndpoint + } + return endpoints.first + } + + private static func testEndpointConnectivity( + _ endpoint: AntigravityConnectionEndpoint, + timeout: TimeInterval) async -> Bool + { + do { + _ = try await self.makeRequest( + payload: RequestPayload( + path: self.unleashPath, + body: self.unleashRequestBody()), + context: RequestContext(endpoints: [endpoint], timeout: timeout)) + return true + } catch { + if self.isReachableProbeError(error) { + self.log.debug("Port probe received HTTP response; treating endpoint as reachable", metadata: [ + "source": endpoint.source.rawValue, + "scheme": endpoint.scheme, + "port": "\(endpoint.port)", + "error": error.localizedDescription, + ]) + return true + } + self.log.debug("Port probe failed", metadata: [ + "source": endpoint.source.rawValue, + "scheme": endpoint.scheme, + "port": "\(endpoint.port)", + "error": error.localizedDescription, + ]) + return false + } + } + + // MARK: - HTTP + + struct RequestPayload { + let path: String + let body: [String: Any] + } + + struct RequestContext { + let endpoints: [AntigravityConnectionEndpoint] + let timeout: TimeInterval + let deadline: Date? + + init(endpoints: [AntigravityConnectionEndpoint], timeout: TimeInterval, deadline: Date? = nil) { + self.endpoints = endpoints + self.timeout = timeout + self.deadline = deadline + } + + func timeoutForNextAttempt() -> TimeInterval? { + AntigravityStatusProbe.timeoutForNextAttempt(timeout: self.timeout, deadline: self.deadline) + } + } + + private static func defaultRequestBody() -> [String: Any] { + [ + "metadata": [ + "ideName": "antigravity", + "extensionName": "antigravity", + "ideVersion": "unknown", + "locale": "en", + ], + ] + } + + private static func unleashRequestBody() -> [String: Any] { + [ + "context": [ + "properties": [ + "devMode": "false", + "extensionVersion": "unknown", + "hasAnthropicModelAccess": "true", + "ide": "antigravity", + "ideVersion": "unknown", + "installationId": "codexbar", + "language": "UNSPECIFIED", + "os": "macos", + "requestedModelId": "MODEL_UNSPECIFIED", + ], + ], + ] + } + + static func fetchSnapshot( + context: RequestContext, + send: @escaping @Sendable (RequestPayload, AntigravityConnectionEndpoint, TimeInterval) async throws -> Data = + sendRequest) async throws -> AntigravityStatusSnapshot + { + do { + let quotaSummary = try await self.makeParsedRequest( + payload: RequestPayload( + path: self.quotaSummaryPath, + body: ["forceRefresh": true]), + context: self.quotaSummaryRequestContext(from: context), + send: send, + parse: self.parseQuotaSummaryResponse) + guard quotaSummary.quotaSummary?.groups.contains(where: { group in + group.buckets.contains { !$0.disabled && $0.remainingFraction != nil } + }) == true else { + throw AntigravityStatusProbeError.parseFailed("Quota summary has no usable quota buckets") + } + let identity = try? await self.makeParsedRequest( + payload: RequestPayload( + path: self.getUserStatusPath, + body: self.defaultRequestBody()), + context: self.identityRequestContext(from: context), + send: send, + parse: self.parseUserStatusResponse) + return quotaSummary.withIdentity(from: identity) + } catch { + self.log.debug("Antigravity quota summary unavailable; falling back to model quotas", metadata: [ + "error": error.localizedDescription, + ]) + } + + do { + return try await self.makeParsedRequest( + payload: RequestPayload( + path: self.getUserStatusPath, + body: self.defaultRequestBody()), + context: self.legacyUserStatusRequestContext(from: context), + send: send, + parse: self.parseUserStatusResponse) + } catch { + return try await self.makeParsedRequest( + payload: RequestPayload( + path: self.commandModelConfigPath, + body: self.defaultRequestBody()), + context: context, + send: send, + parse: self.parseCommandModelResponse) + } + } + + private static func legacyUserStatusRequestContext(from context: RequestContext) -> RequestContext { + guard let deadline = context.deadline else { return context } + let remaining = max(0, deadline.timeIntervalSinceNow) + let userStatusBudget = remaining / 2 + return RequestContext( + endpoints: context.endpoints, + timeout: min(context.timeout, userStatusBudget), + deadline: Date().addingTimeInterval(userStatusBudget)) + } + + private static func quotaSummaryRequestContext(from context: RequestContext) -> RequestContext { + guard let deadline = context.deadline else { return context } + let remaining = max(0, deadline.timeIntervalSinceNow) + let quotaSummaryBudget = remaining / 2 + return RequestContext( + endpoints: context.endpoints, + timeout: min(context.timeout, quotaSummaryBudget), + deadline: Date().addingTimeInterval(quotaSummaryBudget)) + } + + private static func identityRequestContext(from context: RequestContext) -> RequestContext { + RequestContext( + endpoints: context.endpoints, + timeout: min(context.timeout, 1), + deadline: context.deadline) + } + + private static func makeRequest( + payload: RequestPayload, + context: RequestContext) async throws -> Data + { + try await self.sendRequest(payload: payload, context: context) + } + + static func makeParsedRequest( + payload: RequestPayload, + context: RequestContext, + send: @escaping @Sendable (RequestPayload, AntigravityConnectionEndpoint, TimeInterval) async throws -> Data = + sendRequest, + parse: @escaping @Sendable (Data) throws -> T) async throws -> T + { + var lastError: Error? + + for endpoint in context.endpoints { + guard let timeout = context.timeoutForNextAttempt() else { + lastError = lastError ?? AntigravityStatusProbeError.timedOut + break + } + do { + let data = try await send(payload, endpoint, timeout) + return try parse(data) + } catch { + lastError = error + Self.log.debug("Antigravity request/parse attempt failed", metadata: [ + "path": payload.path, + "source": endpoint.source.rawValue, + "scheme": endpoint.scheme, + "port": "\(endpoint.port)", + "error": error.localizedDescription, + ]) + } + } + + throw lastError ?? AntigravityStatusProbeError.apiError("Invalid response") + } + + private static func sendRequest( + payload: RequestPayload, + context: RequestContext) async throws -> Data + { + var lastError: Error? + + for endpoint in context.endpoints { + guard let timeout = context.timeoutForNextAttempt() else { + lastError = lastError ?? AntigravityStatusProbeError.timedOut + break + } + do { + return try await Self.sendRequest(payload: payload, endpoint: endpoint, timeout: timeout) + } catch { + lastError = error + Self.log.debug("Antigravity request attempt failed", metadata: [ + "path": payload.path, + "source": endpoint.source.rawValue, + "scheme": endpoint.scheme, + "port": "\(endpoint.port)", + "error": error.localizedDescription, + ]) + } + } + + throw lastError ?? AntigravityStatusProbeError.apiError("Invalid URL") + } + + private static func sendRequest( + payload: RequestPayload, + endpoint: AntigravityConnectionEndpoint, + timeout: TimeInterval) async throws -> Data + { + guard let url = URL(string: "\(endpoint.scheme)://127.0.0.1:\(endpoint.port)\(payload.path)") else { + throw AntigravityStatusProbeError.apiError("Invalid URL") + } + + let body = try JSONSerialization.data(withJSONObject: payload.body, options: []) + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.httpBody = body + request.timeoutInterval = timeout + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue(String(body.count), forHTTPHeaderField: "Content-Length") + request.setValue("1", forHTTPHeaderField: "Connect-Protocol-Version") + if endpoint.requiresCSRFToken { + request.setValue(endpoint.csrfToken, forHTTPHeaderField: "X-Codeium-Csrf-Token") + } + + let config = URLSessionConfiguration.ephemeral + config.timeoutIntervalForRequest = timeout + config.timeoutIntervalForResource = timeout + #if !os(Linux) + config.waitsForConnectivity = false + #endif + + let delegate = LocalhostSessionDelegate() + let session = URLSession(configuration: config, delegate: delegate, delegateQueue: nil) + defer { session.invalidateAndCancel() } + + let (data, response) = try await delegate.data(for: request, session: session) + guard let http = response as? HTTPURLResponse else { + throw AntigravityStatusProbeError.apiError("Invalid response") + } + guard http.statusCode == 200 else { + let message = String(data: data, encoding: .utf8) ?? "" + throw AntigravityStatusProbeError.apiError("HTTP \(http.statusCode): \(message)") + } + return data + } +} diff --git a/Sources/CodexBarCore/Providers/Antigravity/AntigravityUsageDataSource.swift b/Sources/CodexBarCore/Providers/Antigravity/AntigravityUsageDataSource.swift new file mode 100644 index 0000000..60136c6 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Antigravity/AntigravityUsageDataSource.swift @@ -0,0 +1,30 @@ +import Foundation + +public enum AntigravityUsageDataSource: String, CaseIterable, Identifiable, Sendable { + case auto + case oauth + case cli + + public var id: String { + self.rawValue + } + + public var displayName: String { + switch self { + case .auto: "Auto" + case .oauth: "Google OAuth" + case .cli: "Local API / agy CLI" + } + } + + public var sourceLabel: String { + switch self { + case .auto: + "auto" + case .oauth: + "oauth" + case .cli: + "cli" + } + } +} diff --git a/Sources/CodexBarCore/Providers/Augment/AuggieCLIProbe.swift b/Sources/CodexBarCore/Providers/Augment/AuggieCLIProbe.swift new file mode 100644 index 0000000..ddbb6de --- /dev/null +++ b/Sources/CodexBarCore/Providers/Augment/AuggieCLIProbe.swift @@ -0,0 +1,191 @@ +import Foundation + +#if os(macOS) + +/// Fetches Augment usage via `auggie account status` CLI command +public struct AuggieCLIProbe: Sendable { + private static let log = CodexBarLog.logger(LogCategories.auggieCLI) + + public init() {} + + public func fetch() async throws -> AugmentStatusSnapshot { + let output = try await self.runAuggieAccountStatus() + return try self.parse(output) + } + + /// Timeout for the `auggie account status` command. + private static let commandTimeout: TimeInterval = 15 + + private func runAuggieAccountStatus() async throws -> String { + let env = ProcessInfo.processInfo.environment + let loginPATH = LoginShellPathCache.shared.current + let executable = BinaryLocator.resolveAuggieBinary(env: env, loginPATH: loginPATH) ?? "auggie" + + var pathEnv = env + pathEnv["PATH"] = PathBuilder.effectivePATH( + purposes: [.tty, .nodeTooling], + env: env, + loginPATH: loginPATH) + + let result = try await SubprocessRunner.run( + binary: executable, + arguments: ["account", "status"], + environment: pathEnv, + timeout: Self.commandTimeout, + label: "auggie-account-status") + + let output = result.stdout + let errorOutput = result.stderr + + guard !output.isEmpty else { + if !errorOutput.isEmpty { + Self.log.error("Auggie stderr: \(errorOutput)") + } + throw AuggieCLIError.noOutput + } + + // Check for auth errors + if output.contains("Authentication failed") || output.contains("auggie login") { + throw AuggieCLIError.notAuthenticated + } + + return output + } + + func parse(_ output: String) throws -> AugmentStatusSnapshot { + // Legacy output: + // Max Plan 450,000 credits / month + // 11,657 remaining · 953,170 / 964,827 credits used + // 2 days remaining in this billing cycle (ends 1/8/2026) + // + // Current output (2026+): + // 319,054 credits remaining Max Plan + // 450,000 credits / month + // 9 days remaining in this billing cycle (ends 6/9/2026) + + var maxCredits: Int? + var remaining: Int? + var used: Int? + var total: Int? + var billingCycleEnd: Date? + + for line in output.split(separator: "\n") { + let trimmed = line.trimmingCharacters(in: .whitespaces) + + if trimmed.contains("credits / month") { + if let match = trimmed.range(of: #"([\d,]+)\s+credits\s*/\s*month"#, options: .regularExpression) { + let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "") + .replacingOccurrences(of: " credits", with: "") + .replacingOccurrences(of: " / month", with: "") + maxCredits = Int(numberStr) + total = total ?? Int(numberStr) + } + } else if trimmed.contains("Max Plan"), trimmed.contains("credits"), !trimmed.contains("remaining") { + if let match = trimmed.range(of: #"([\d,]+)\s+credits"#, options: .regularExpression) { + let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "") + .replacingOccurrences(of: " credits", with: "") + maxCredits = Int(numberStr) + } + } + + if trimmed.contains("credits remaining"), !trimmed.contains("billing cycle") { + if let match = trimmed.range(of: #"([\d,]+)\s+credits\s+remaining"#, options: .regularExpression) { + let numberStr = String(trimmed[match]).replacingOccurrences(of: ",", with: "") + .replacingOccurrences(of: " credits", with: "") + .replacingOccurrences(of: " remaining", with: "") + remaining = Int(numberStr) + } + } + + // Parse "11,657 remaining · 953,170 / 964,827 credits used" + if trimmed.contains("remaining"), trimmed.contains("credits used") { + if let remMatch = trimmed.range(of: #"([\d,]+)\s+remaining"#, options: .regularExpression) { + let numStr = String(trimmed[remMatch]) + .replacingOccurrences(of: ",", with: "") + .replacingOccurrences(of: " remaining", with: "") + remaining = Int(numStr) + } + + if let usedMatch = trimmed.range( + of: #"([\d,]+)\s*/\s*([\d,]+)\s+credits used"#, + options: .regularExpression) + { + let parts = String(trimmed[usedMatch]) + .replacingOccurrences(of: " credits used", with: "") + .split(separator: "/") + if parts.count == 2 { + used = Int(parts[0].replacingOccurrences(of: ",", with: "") + .trimmingCharacters(in: .whitespaces)) + total = Int(parts[1].replacingOccurrences(of: ",", with: "") + .trimmingCharacters(in: .whitespaces)) + } + } + } + + if trimmed.contains("billing cycle"), trimmed.contains("ends") { + if let dateMatch = trimmed.range(of: #"ends\s+([\d/]+)"#, options: .regularExpression) { + let dateStr = String(trimmed[dateMatch]) + .replacingOccurrences(of: "ends", with: "") + .trimmingCharacters(in: .whitespaces) + + let formatter = DateFormatter() + formatter.dateFormat = "M/d/yyyy" + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = TimeZone.current + billingCycleEnd = formatter.date(from: dateStr) + } + } + } + + guard let finalRemaining = remaining else { + Self.log.error("Failed to parse auggie output: \(output)") + throw AuggieCLIError.parseError("Could not extract credits from output") + } + + let finalTotal = total ?? maxCredits + guard let finalTotal else { + Self.log.error("Failed to parse auggie output: \(output)") + throw AuggieCLIError.parseError("Could not extract credits from output") + } + + let finalUsed = used ?? max(0, finalTotal - finalRemaining) + + return AugmentStatusSnapshot( + creditsRemaining: Double(finalRemaining), + creditsUsed: Double(finalUsed), + creditsLimit: Double(finalTotal), + billingCycleEnd: billingCycleEnd, + accountEmail: nil, + accountPlan: maxCredits.map { "\($0.formatted()) credits/month" }, + rawJSON: nil) + } +} + +#else + +public struct AuggieCLIProbe: Sendable { + public init() {} + + public func fetch() async throws -> AugmentStatusSnapshot { + throw AugmentStatusProbeError.notSupported + } +} + +#endif + +public enum AuggieCLIError: LocalizedError { + case noOutput + case notAuthenticated + case parseError(String) + + public var errorDescription: String? { + switch self { + case .noOutput: + "Auggie CLI returned no output" + case .notAuthenticated: + "Not authenticated. Run 'auggie login' to authenticate." + case let .parseError(msg): + "Failed to parse auggie output: \(msg)" + } + } +} diff --git a/Sources/CodexBarCore/Providers/Augment/AugmentProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Augment/AugmentProviderDescriptor.swift new file mode 100644 index 0000000..529bd39 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Augment/AugmentProviderDescriptor.swift @@ -0,0 +1,133 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit +#endif + +public enum AugmentProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + #if os(macOS) + // Custom browser order that includes Chrome Beta and other variants + // to support users running beta/canary versions + let browserOrder: BrowserCookieImportOrder = [ + .safari, + .chrome, + .chromeBeta, // Added for Chrome Beta support + .chromeCanary, // Added for Chrome Canary support + .edge, + .edgeBeta, + .brave, + .arc, + .dia, + .arcBeta, + .firefox, + ] + #else + let browserOrder: BrowserCookieImportOrder? = nil + #endif + + return ProviderDescriptor( + id: .augment, + metadata: ProviderMetadata( + id: .augment, + displayName: "Augment", + sessionLabel: "Credits", + weeklyLabel: "Usage", + opusLabel: nil, + supportsOpus: false, + supportsCredits: true, + creditsHint: "Augment Code credits for AI-powered coding assistance.", + toggleTitle: "Show Augment usage", + cliName: "augment", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: browserOrder, + dashboardURL: "https://app.augmentcode.com/account/subscription", + statusPageURL: "https://status.augmentcode.com"), + branding: ProviderBranding( + iconStyle: .augment, + iconResourceName: "ProviderIcon-augment", + color: ProviderColor(red: 99 / 255, green: 102 / 255, blue: 241 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Augment cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .cli], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + var strategies: [any ProviderFetchStrategy] = [] + // Try CLI first (no browser prompts!) + strategies.append(AugmentCLIFetchStrategy()) + // Fallback to web (browser cookies) + strategies.append(AugmentStatusFetchStrategy()) + return strategies + })), + cli: ProviderCLIConfig( + name: "augment", + versionDetector: nil)) + } +} + +struct AugmentCLIFetchStrategy: ProviderFetchStrategy { + let id: String = "augment.cli" + let kind: ProviderFetchKind = .cli + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + // Check if auggie CLI is installed + let env = ProcessInfo.processInfo.environment + let loginPATH = LoginShellPathCache.shared.current + return BinaryLocator.resolveAuggieBinary(env: env, loginPATH: loginPATH) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let probe = AuggieCLIProbe() + let snap = try await probe.fetch() + return self.makeResult( + usage: snap.toUsageSnapshot(), + sourceLabel: "cli") + } + + func shouldFallback(on error: Error, context _: ProviderFetchContext) -> Bool { + // Fallback to web if CLI fails (not authenticated, etc.) + if let cliError = error as? AuggieCLIError { + switch cliError { + case .notAuthenticated, .noOutput, .parseError: + return true + } + } + return true + } +} + +struct AugmentStatusFetchStrategy: ProviderFetchStrategy { + let id: String = "augment.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + guard context.settings?.augment?.cookieSource != .off else { return false } + return true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let probe = AugmentStatusProbe() + let manual = Self.manualCookieHeader(from: context) + let logger: ((String) -> Void)? = context.verbose + ? { msg in CodexBarLog.logger(LogCategories.augment).verbose(msg) } + : nil + let snap = try await probe.fetch(cookieHeaderOverride: manual, logger: logger) + return self.makeResult( + usage: snap.toUsageSnapshot(), + sourceLabel: "web") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func manualCookieHeader(from context: ProviderFetchContext) -> String? { + guard context.settings?.augment?.cookieSource == .manual else { return nil } + return CookieHeaderNormalizer.normalize(context.settings?.augment?.manualCookieHeader) + } +} diff --git a/Sources/CodexBarCore/Providers/Augment/AugmentSessionKeepalive.swift b/Sources/CodexBarCore/Providers/Augment/AugmentSessionKeepalive.swift new file mode 100644 index 0000000..7ffa2ae --- /dev/null +++ b/Sources/CodexBarCore/Providers/Augment/AugmentSessionKeepalive.swift @@ -0,0 +1,491 @@ +import Foundation + +#if os(macOS) +import AppKit +import UserNotifications + +/// Manages automatic session keepalive for Augment to prevent cookie expiration. +/// +/// This actor monitors cookie expiration and proactively refreshes the session +/// before cookies expire, ensuring uninterrupted access to Augment APIs. +@MainActor +public final class AugmentSessionKeepalive { + // MARK: - Configuration + + /// How often to check if session needs refresh (default: 1 minute for faster recovery) + private let checkInterval: TimeInterval = 60 + + /// Refresh session this many seconds before cookie expiration (default: 5 minutes) + private let refreshBufferSeconds: TimeInterval = 300 + + /// Minimum time between refresh attempts (default: 1 minute for faster recovery) + private let minRefreshInterval: TimeInterval = 60 + + /// Maximum time to wait for session refresh (default: 30 seconds) + private let refreshTimeout: TimeInterval = 30 + + // MARK: - State + + private var timerTask: Task? + private var lastRefreshAttempt: Date? + private var lastSuccessfulRefresh: Date? + private var isRefreshing = false + private let logger: ((String) -> Void)? + private var onSessionRecovered: (() async -> Void)? + + /// Track consecutive failures to stop retrying after too many failures + private var consecutiveFailures = 0 + private let maxConsecutiveFailures = 3 // Stop after 3 failures + private var hasGivenUp = false + + // MARK: - Initialization + + public init(logger: ((String) -> Void)? = nil, onSessionRecovered: (() async -> Void)? = nil) { + self.logger = logger + self.onSessionRecovered = onSessionRecovered + } + + deinit { + self.timerTask?.cancel() + } + + // MARK: - Public API + + /// Start the automatic session keepalive timer + public func start() { + guard self.timerTask == nil else { + self.log("Keepalive already running") + return + } + + self.log("🚀 Starting Augment session keepalive") + self.log( + " - Check interval: \(Int(self.checkInterval))s " + + "(\(Self.durationDescription(seconds: self.checkInterval)))") + self.log( + " - Refresh buffer: \(Int(self.refreshBufferSeconds))s " + + "(\(Self.durationDescription(seconds: self.refreshBufferSeconds)) before expiry)") + self.log( + " - Min refresh interval: \(Int(self.minRefreshInterval))s " + + "(\(Self.durationDescription(seconds: self.minRefreshInterval)))") + + self.timerTask = Task.detached(priority: .utility) { [weak self] in + while !Task.isCancelled { + try? await Task.sleep(for: .seconds(self?.checkInterval ?? 300)) + await self?.checkAndRefreshIfNeeded() + } + } + + self.log("✅ Keepalive timer started successfully") + } + + /// Stop the automatic session keepalive timer + public func stop() { + self.log("Stopping Augment session keepalive") + self.timerTask?.cancel() + self.timerTask = nil + } + + /// Manually trigger a session refresh (bypasses rate limiting) + public func forceRefresh() async { + self.log("Force refresh requested") + await self.performRefresh(forced: true) + } + + // MARK: - Private Implementation + + private func checkAndRefreshIfNeeded() async { + guard !self.isRefreshing else { + self.log("Refresh already in progress, skipping check") + return + } + + // Stop trying if we've given up + if self.hasGivenUp { + self.log("⏸️ Keepalive has given up after \(self.maxConsecutiveFailures) consecutive failures") + self.log(" User must manually log in to Augment and click 'Refresh Session'") + return + } + + // Rate limit: don't refresh too frequently + if let lastAttempt = self.lastRefreshAttempt { + let timeSinceLastAttempt = Date().timeIntervalSince(lastAttempt) + if timeSinceLastAttempt < self.minRefreshInterval { + self.log( + "Skipping refresh (last attempt \(Int(timeSinceLastAttempt))s ago, " + + "min interval: \(Int(self.minRefreshInterval))s)") + return + } + } + + // Check if cookies are about to expire + let shouldRefresh = await self.shouldRefreshSession() + if shouldRefresh { + await self.performRefresh(forced: false) + } + } + + private func shouldRefreshSession() async -> Bool { + do { + let session = try AugmentCookieImporter.importSession(logger: self.logger) + + self.log("📊 Cookie Status Check:") + self.log(" Total cookies: \(session.cookies.count)") + self.log(" Source: \(session.sourceLabel)") + + // Log each cookie's expiration status + for cookie in session.cookies { + if let expiry = cookie.expiresDate { + let timeUntil = expiry.timeIntervalSinceNow + let status = timeUntil > 0 ? "expires in \(Int(timeUntil))s" : "EXPIRED \(Int(-timeUntil))s ago" + self.log(" - \(cookie.name): \(status)") + } else { + self.log(" - \(cookie.name): session cookie (no expiry)") + } + } + + // Find the earliest expiration date among session cookies + let expirationDates = session.cookies.compactMap(\.expiresDate) + + guard !expirationDates.isEmpty else { + // Session cookies (no expiration) - refresh periodically + self.log(" All cookies are session cookies (no expiration dates)") + if let lastRefresh = self.lastSuccessfulRefresh { + let timeSinceRefresh = Date().timeIntervalSince(lastRefresh) + // Refresh every 30 minutes for session cookies + if timeSinceRefresh > 1800 { + self.log(" ⚠️ Need periodic refresh (\(Int(timeSinceRefresh))s since last refresh)") + return true + } else { + self.log(" ✅ Recently refreshed (\(Int(timeSinceRefresh))s ago)") + return false + } + } else { + // Never refreshed - do it now + self.log(" ⚠️ Never refreshed - doing initial refresh") + return true + } + } + + let earliestExpiration = expirationDates.min()! + let timeUntilExpiration = earliestExpiration.timeIntervalSinceNow + let expiringCookie = session.cookies.first { $0.expiresDate == earliestExpiration } + + if timeUntilExpiration < self.refreshBufferSeconds { + self.log(" ⚠️ REFRESH NEEDED:") + self.log(" Earliest expiring cookie: \(expiringCookie?.name ?? "unknown")") + self.log(" Time until expiration: \(Int(timeUntilExpiration))s") + self.log(" Refresh threshold: \(Int(self.refreshBufferSeconds))s") + return true + } else { + self.log(" ✅ Session healthy:") + self.log(" Earliest expiring cookie: \(expiringCookie?.name ?? "unknown")") + self.log(" Time until expiration: \(Int(timeUntilExpiration))s") + return false + } + } catch { + self.log("✗ Failed to check session: \(error.localizedDescription)") + return false + } + } + + private func performRefresh(forced: Bool) async { + self.isRefreshing = true + self.lastRefreshAttempt = Date() + defer { self.isRefreshing = false } + + self.log(forced ? "Performing forced session refresh..." : "Performing automatic session refresh...") + + // If this is a forced refresh (user clicked "Refresh Session"), reset failure tracking + if forced { + self.consecutiveFailures = 0 + self.hasGivenUp = false + self.log("🔄 Manual refresh - resetting failure tracking") + } + + do { + // Step 1: Ping the session endpoint to trigger cookie refresh + let refreshed = try await self.pingSessionEndpoint() + + if refreshed { + // Step 2: Re-import cookies from browser + try await Task.sleep(for: .seconds(1)) // Brief delay for browser to update cookies + let newSession = try AugmentCookieImporter.importSession(logger: self.logger) + + await AugmentSessionStore.shared.setCookies(newSession.cookies) + CookieHeaderCache.store( + provider: .augment, + cookieHeader: newSession.cookieHeader, + sourceLabel: newSession.sourceLabel) + + self.log( + "✅ Session refresh successful - imported \(newSession.cookies.count) cookies " + + "from \(newSession.sourceLabel)") + self.lastSuccessfulRefresh = Date() + + // Reset failure tracking on success + self.consecutiveFailures = 0 + self.hasGivenUp = false + + if let callback = self.onSessionRecovered { + self.log("🔄 Triggering usage refresh after session refresh") + await callback() + } + } else { + self.log("⚠️ Session refresh returned no new cookies") + self.consecutiveFailures += 1 + self.checkIfShouldGiveUp() + } + } catch AugmentSessionKeepaliveError.sessionExpired { + self.log("🔐 Session expired - attempting automatic recovery...") + self.consecutiveFailures += 1 + + if self.consecutiveFailures >= self.maxConsecutiveFailures { + self.log("❌ Too many consecutive failures (\(self.consecutiveFailures)) - giving up") + self.log(" User must manually log in to Augment and click 'Refresh Session'") + self.hasGivenUp = true + self.notifyUserLoginRequired() + } else { + await self.attemptSessionRecovery() + } + } catch { + self.log("✗ Session refresh failed: \(error.localizedDescription)") + self.consecutiveFailures += 1 + self.checkIfShouldGiveUp() + } + } + + private func checkIfShouldGiveUp() { + if self.consecutiveFailures >= self.maxConsecutiveFailures { + self.log("❌ Too many consecutive failures (\(self.consecutiveFailures)) - giving up") + self.log(" User must manually log in to Augment and click 'Refresh Session'") + self.hasGivenUp = true + self.notifyUserLoginRequired() + } + } + + /// Attempt to recover from an expired session by triggering browser re-authentication + private func attemptSessionRecovery() async { + self.log("🔄 Attempting automatic session recovery...") + self.log(" Strategy: Open Augment dashboard to trigger browser re-auth") + + #if os(macOS) + // Open the Augment dashboard in the default browser + // This will trigger the browser to re-authenticate if the user is still logged in + if let url = URL(string: "https://app.augmentcode.com") { + _ = await MainActor.run { + NSWorkspace.shared.open(url) + } + self.log(" ✅ Opened Augment dashboard in browser") + self.log(" ⏳ Waiting 5 seconds for browser to re-authenticate...") + + // Wait for browser to potentially re-authenticate + try? await Task.sleep(for: .seconds(5)) + + // Try to import cookies again + do { + let newSession = try AugmentCookieImporter.importSession(logger: self.logger) + self.log(" ✅ Session recovery successful - imported \(newSession.cookies.count) cookies") + self.lastSuccessfulRefresh = Date() + + // Verify the session is actually valid by pinging the API + let isValid = try await self.pingSessionEndpoint() + if isValid { + self.log(" ✅ Session verified - recovery complete!") + // Notify UsageStore to refresh Augment usage + if let callback = self.onSessionRecovered { + self.log(" 🔄 Triggering usage refresh after successful recovery") + await callback() + } + } else { + self.log(" ⚠️ Session imported but not yet valid - may need manual login") + self.notifyUserLoginRequired() + } + } catch { + self.log(" ✗ Session recovery failed: \(error.localizedDescription)") + self.log(" ℹ️ User needs to manually log in to Augment") + self.notifyUserLoginRequired() + } + } + #else + self.log(" ✗ Automatic recovery not supported on this platform") + #endif + } + + /// Notify the user that they need to log in to Augment + private func notifyUserLoginRequired() { + #if os(macOS) + self.log("📢 Sending notification: Augment session expired") + + Task { + let center = UNUserNotificationCenter.current() + + // Request authorization if needed + do { + let granted = try await center.requestAuthorization(options: [.alert, .sound]) + guard granted else { + self.log("⚠️ Notification permission denied") + return + } + } catch { + self.log("✗ Failed to request notification permission: \(error)") + return + } + + // Create notification content + let content = UNMutableNotificationContent() + content.title = "Augment Session Expired" + content.body = "Please log in to app.augmentcode.com to restore your session." + content.sound = .default + + // Create trigger (deliver immediately) + let request = UNNotificationRequest( + identifier: "augment-session-expired-\(UUID().uuidString)", + content: content, + trigger: nil) + + // Deliver notification + do { + try await center.add(request) + self.log("✅ Notification delivered successfully") + } catch { + self.log("✗ Failed to deliver notification: \(error)") + } + } + #endif + } + + /// Ping Augment's session endpoint to trigger cookie refresh + private func pingSessionEndpoint() async throws -> Bool { + // Try to get current cookies first + let currentSession = try? AugmentCookieImporter.importSession(logger: self.logger) + guard let cookieHeader = currentSession?.cookieHeader else { + self.log("No cookies available for session ping") + return false + } + + self.log("🔄 Attempting session refresh...") + self.log(" Cookies being sent: \(cookieHeader.prefix(100))...") + + // Try multiple endpoints - Augment might use different auth patterns + let endpoints = [ + "https://app.augmentcode.com/api/auth/session", // NextAuth pattern + "https://app.augmentcode.com/api/session", // Alternative + "https://app.augmentcode.com/api/user", // User endpoint + ] + + var receivedUnauthorized = false + + for (index, urlString) in endpoints.enumerated() { + self.log(" Trying endpoint \(index + 1)/\(endpoints.count): \(urlString)") + + guard let sessionURL = URL(string: urlString) else { continue } + var request = URLRequest(url: sessionURL) + request.timeoutInterval = self.refreshTimeout + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("https://app.augmentcode.com", forHTTPHeaderField: "Origin") + request.setValue("https://app.augmentcode.com", forHTTPHeaderField: "Referer") + + do { + let (data, response) = try await ProviderHTTPClient.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + self.log(" ✗ Invalid response type") + continue + } + + self.log(" Response: HTTP \(httpResponse.statusCode)") + + // Log Set-Cookie headers if present + if let setCookies = httpResponse.allHeaderFields["Set-Cookie"] as? String { + self.log(" Set-Cookie headers received: \(setCookies.prefix(100))...") + } + + if httpResponse.statusCode == 200 { + // Check if we got a valid session response + if let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] { + self.log(" JSON response keys: \(json.keys.joined(separator: ", "))") + + if json["user"] != nil || json["email"] != nil || json["session"] != nil { + self.log(" ✅ Valid session data found!") + return true + } else { + self.log(" ⚠️ 200 OK but no session data in response") + // Try next endpoint + continue + } + } else { + self.log(" ⚠️ 200 OK but response is not JSON") + if let responseText = String(data: data, encoding: .utf8) { + self.log(" Response text: \(responseText.prefix(200))...") + } + continue + } + } else if httpResponse.statusCode == 401 { + self.log(" ✗ 401 Unauthorized - session expired") + receivedUnauthorized = true + // Don't throw immediately - try all endpoints first + continue + } else if httpResponse.statusCode == 404 { + self.log(" ✗ 404 Not Found - trying next endpoint") + continue + } else { + self.log(" ✗ HTTP \(httpResponse.statusCode) - trying next endpoint") + continue + } + } catch { + self.log(" ✗ Request failed: \(error.localizedDescription)") + continue + } + } + + // If we got 401 from all endpoints, the session is definitely expired + if receivedUnauthorized { + self.log("⚠️ All endpoints returned 401 - session is expired") + throw AugmentSessionKeepaliveError.sessionExpired + } + + self.log("⚠️ All session endpoints failed or returned no valid data") + return false + } + + private static let log = CodexBarLog.logger(LogCategories.augmentKeepalive) + + private static func durationDescription(seconds: TimeInterval) -> String { + let totalSeconds = max(0, Int(seconds.rounded())) + if totalSeconds >= 60, totalSeconds % 60 == 0 { + let minutes = totalSeconds / 60 + return "\(minutes) minute\(minutes == 1 ? "" : "s")" + } + return "\(totalSeconds) second\(totalSeconds == 1 ? "" : "s")" + } + + private func log(_ message: String) { + let timestamp = Date().formatted(date: .omitted, time: .standard) + let fullMessage = "[\(timestamp)] [AugmentKeepalive] \(message)" + self.logger?(fullMessage) + Self.log.debug(fullMessage) + } +} + +// MARK: - Errors + +public enum AugmentSessionKeepaliveError: LocalizedError, Sendable { + case invalidResponse + case sessionExpired + case networkError(String) + + public var errorDescription: String? { + switch self { + case .invalidResponse: + "Invalid response from session endpoint" + case .sessionExpired: + "Session has expired" + case let .networkError(message): + "Network error: \(message)" + } + } +} + +#endif diff --git a/Sources/CodexBarCore/Providers/Augment/AugmentStatusProbe.swift b/Sources/CodexBarCore/Providers/Augment/AugmentStatusProbe.swift new file mode 100644 index 0000000..52ae33a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Augment/AugmentStatusProbe.swift @@ -0,0 +1,698 @@ +import Foundation +import SweetCookieKit + +#if os(macOS) + +private let augmentCookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.augment]?.browserCookieOrder ?? Browser.defaultImportOrder + +// MARK: - Augment Cookie Importer + +/// Imports Augment session cookies from browser cookies. +public enum AugmentCookieImporter { + private static let cookieClient = BrowserCookieClient() + /// Auth0 session cookies used by Augment + /// NOTE: This list may not be exhaustive. If authentication fails with cookies present, + /// check debug logs for cookie names and report them. + private static let sessionCookieNames: Set = [ + "session", // Augment auth session (auth.augmentcode.com) + "_session", // Legacy session cookie (app.augmentcode.com) + "web_rpc_proxy_session", // Augment RPC proxy session + "auth0", // Auth0 session + "auth0.is.authenticated", // Auth0 authentication flag + "a0.spajs.txs", // Auth0 SPA transaction state + "__Secure-next-auth.session-token", // NextAuth secure session + "next-auth.session-token", // NextAuth session + "__Secure-authjs.session-token", // AuthJS secure session + "__Host-authjs.csrf-token", // AuthJS CSRF token + "authjs.session-token", // AuthJS session + ] + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + + /// Returns cookie header filtered for a specific target URL + public func cookieHeader(for url: URL) -> String { + guard let host = url.host else { return "" } + + let matchingCookies = self.cookies.filter { cookie in + let domain = cookie.domain + + // Handle wildcard domains (e.g., ".augmentcode.com") + if domain.hasPrefix(".") { + let baseDomain = String(domain.dropFirst()) + return host == baseDomain || host.hasSuffix(".\(baseDomain)") + } + + // Exact match or subdomain match + return host == domain || host.hasSuffix(".\(domain)") + } + + return matchingCookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + } + + /// Attempts to import Augment cookies using the standard browser import order. + public static func importSession(logger: ((String) -> Void)? = nil) throws -> SessionInfo { + let log: (String) -> Void = { msg in logger?("[augment-cookie] \(msg)") } + + let cookieDomains = ["augmentcode.com", "app.augmentcode.com"] + for browserSource in augmentCookieImportOrder { + do { + let query = BrowserCookieQuery(domains: cookieDomains) + let sources = try Self.cookieClient.codexBarRecords( + matching: query, + in: browserSource, + logger: log) + for source in sources where !source.records.isEmpty { + let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin) + + // Log all cookies found for debugging + let cookieNames = httpCookies.map(\.name).joined(separator: ", ") + log("Found \(httpCookies.count) cookies in \(source.label): \(cookieNames)") + + // Check if we have any session cookies + let matchingCookies = httpCookies.filter { Self.sessionCookieNames.contains($0.name) } + if !matchingCookies.isEmpty { + let matchingCookieNames = matchingCookies.map(\.name).joined(separator: ", ") + log("✓ Found Augment session cookies in \(source.label): \(matchingCookieNames)") + return SessionInfo(cookies: httpCookies, sourceLabel: source.label) + } else if !httpCookies.isEmpty { + // Log unrecognized cookies to help discover missing session cookie names + log( + "⚠️ \(source.label) has \(httpCookies.count) cookies but none match known session cookies") + log(" Cookie names found: \(cookieNames)") + log(" If you're logged into Augment, please report these cookie names") + log(" to help improve detection") + } + } + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)") + } + } + + throw AugmentStatusProbeError.noSessionCookie + } + + /// Check if Augment session cookies are available + public static func hasSession(logger: ((String) -> Void)? = nil) -> Bool { + do { + let session = try self.importSession(logger: logger) + return !session.cookies.isEmpty + } catch { + return false + } + } +} + +// MARK: - Augment API Models + +public struct AugmentCreditsResponse: Codable, Sendable { + public let usageUnitsRemaining: Double? + public let usageUnitsConsumedThisBillingCycle: Double? + public let usageUnitsAvailable: Double? + public let usageBalanceStatus: String? + + /// Computed properties for compatibility with existing code + public var credits: Double? { + self.usageUnitsRemaining + } + + public var creditsUsed: Double? { + self.usageUnitsConsumedThisBillingCycle + } + + public var creditsLimit: Double? { + if let available = self.usageUnitsAvailable, available > 0 { + return available + } + + guard let remaining = self.usageUnitsRemaining, + let consumed = self.usageUnitsConsumedThisBillingCycle + else { + return nil + } + return remaining + consumed + } + + private enum CodingKeys: String, CodingKey { + case usageUnitsRemaining + case usageUnitsConsumedThisBillingCycle + case usageUnitsAvailable + case usageBalanceStatus + } +} + +public struct AugmentSubscriptionResponse: Codable, Sendable { + public let planName: String? + public let billingPeriodEnd: String? + public let email: String? + public let organization: String? + + private enum CodingKeys: String, CodingKey { + case planName + case billingPeriodEnd + case email + case organization + } +} + +// MARK: - Augment Status Snapshot + +public struct AugmentStatusSnapshot: Sendable { + public let creditsRemaining: Double? + public let creditsUsed: Double? + public let creditsLimit: Double? + public let billingCycleEnd: Date? + public let accountEmail: String? + public let accountPlan: String? + public let rawJSON: String? + + public init( + creditsRemaining: Double?, + creditsUsed: Double?, + creditsLimit: Double?, + billingCycleEnd: Date?, + accountEmail: String?, + accountPlan: String?, + rawJSON: String?) + { + self.creditsRemaining = creditsRemaining + self.creditsUsed = creditsUsed + self.creditsLimit = creditsLimit + self.billingCycleEnd = billingCycleEnd + self.accountEmail = accountEmail + self.accountPlan = accountPlan + self.rawJSON = rawJSON + } + + public func toUsageSnapshot() -> UsageSnapshot { + let percentUsed: Double = if let used = self.creditsUsed, let limit = self.creditsLimit, limit > 0 { + (used / limit) * 100.0 + } else if let remaining = self.creditsRemaining, let limit = self.creditsLimit, limit > 0 { + ((limit - remaining) / limit) * 100.0 + } else { + 0 + } + + let primary = RateWindow( + usedPercent: percentUsed, + windowMinutes: nil, + resetsAt: self.billingCycleEnd, + resetDescription: self.billingCycleEnd.map { "Resets \(Self.formatResetDate($0))" }) + + let identity = ProviderIdentitySnapshot( + providerID: .augment, + accountEmail: self.accountEmail, + accountOrganization: nil, + loginMethod: self.accountPlan) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: Date(), + identity: identity) + } + + private static func formatResetDate(_ date: Date) -> String { + let formatter = RelativeDateTimeFormatter() + formatter.locale = Locale(identifier: "en_US") + formatter.unitsStyle = .full + return formatter.localizedString(for: date, relativeTo: Date()) + } +} + +// MARK: - Augment Status Probe Error + +public enum AugmentStatusProbeError: LocalizedError, Sendable { + case notLoggedIn + case networkError(String) + case parseFailed(String) + case noSessionCookie + case sessionExpired + + public var errorDescription: String? { + switch self { + case .notLoggedIn: + "Not logged in to Augment. Please log in via the CodexBar menu." + case let .networkError(msg): + "Augment API error: \(msg)" + case let .parseFailed(msg): + "Could not parse Augment usage: \(msg)" + case .noSessionCookie: + "No Augment session found. Please log in to app.augmentcode.com in \(augmentCookieImportOrder.loginHint)." + case .sessionExpired: + "Augment session expired. Please log in again." + } + } +} + +// MARK: - Augment Session Store + +public actor AugmentSessionStore { + public static let shared = AugmentSessionStore() + + private var sessionCookies: [HTTPCookie] = [] + private var hasLoadedFromDisk = false + private let fileURL: URL + + private init() { + let fm = FileManager.default + let appSupport = fm.urls(for: .applicationSupportDirectory, in: .userDomainMask).first + ?? fm.temporaryDirectory + let dir = appSupport.appendingPathComponent("CodexBar", isDirectory: true) + try? fm.createDirectory(at: dir, withIntermediateDirectories: true) + self.fileURL = dir.appendingPathComponent("augment-session.json") + + // Load saved cookies on init + Task { await self.loadFromDiskIfNeeded() } + } + + public func setCookies(_ cookies: [HTTPCookie]) { + self.hasLoadedFromDisk = true + self.sessionCookies = cookies + self.saveToDisk() + } + + public func getCookies() -> [HTTPCookie] { + self.loadFromDiskIfNeeded() + return self.sessionCookies + } + + public func clearCookies() { + self.hasLoadedFromDisk = true + self.sessionCookies = [] + try? FileManager.default.removeItem(at: self.fileURL) + } + + public func hasValidSession() -> Bool { + self.loadFromDiskIfNeeded() + return !self.sessionCookies.isEmpty + } + + #if DEBUG + func resetForTesting(clearDisk: Bool = true) { + self.hasLoadedFromDisk = false + self.sessionCookies = [] + if clearDisk { + try? FileManager.default.removeItem(at: self.fileURL) + } + } + #endif + + private func loadFromDiskIfNeeded() { + guard !self.hasLoadedFromDisk else { return } + self.hasLoadedFromDisk = true + self.loadFromDisk() + } + + private func saveToDisk() { + // Convert cookie properties to JSON-serializable format + // Date values must be converted to TimeInterval (Double) + let cookieData = self.sessionCookies.compactMap { cookie -> [String: Any]? in + guard let props = cookie.properties else { return nil } + var serializable: [String: Any] = [:] + for (key, value) in props { + let keyString = key.rawValue + if let date = value as? Date { + // Convert Date to TimeInterval for JSON compatibility + serializable[keyString] = date.timeIntervalSince1970 + serializable[keyString + "_isDate"] = true + } else if let url = value as? URL { + serializable[keyString] = url.absoluteString + serializable[keyString + "_isURL"] = true + } else if JSONSerialization.isValidJSONObject([value]) || + value is String || + value is Bool || + value is NSNumber + { + serializable[keyString] = value + } + } + return serializable + } + guard !cookieData.isEmpty, + let data = try? JSONSerialization.data(withJSONObject: cookieData, options: [.prettyPrinted]) + else { + return + } + try? data.write(to: self.fileURL) + } + + private func loadFromDisk() { + guard let data = try? Data(contentsOf: self.fileURL), + let cookieArray = try? JSONSerialization.jsonObject(with: data) as? [[String: Any]] + else { return } + + self.sessionCookies = cookieArray.compactMap { props in + // Convert back to HTTPCookiePropertyKey dictionary + var cookieProps: [HTTPCookiePropertyKey: Any] = [:] + for (key, value) in props { + // Skip marker keys + if key.hasSuffix("_isDate") || key.hasSuffix("_isURL") { continue } + + let propKey = HTTPCookiePropertyKey(key) + + // Check if this was a Date + if props[key + "_isDate"] as? Bool == true, let interval = value as? TimeInterval { + cookieProps[propKey] = Date(timeIntervalSince1970: interval) + } + // Check if this was a URL + else if props[key + "_isURL"] as? Bool == true, let urlString = value as? String { + cookieProps[propKey] = URL(string: urlString) + } else { + cookieProps[propKey] = value + } + } + return HTTPCookie(properties: cookieProps) + } + } +} + +// MARK: - Augment Status Probe + +public struct AugmentStatusProbe: Sendable { + public let baseURL: URL + public var timeout: TimeInterval = 15.0 + + public init(baseURL: URL = URL(string: "https://app.augmentcode.com")!, timeout: TimeInterval = 15.0) { + self.baseURL = baseURL + self.timeout = timeout + } + + /// Fetch Augment usage with manual cookie header (for debugging). + public func fetchWithManualCookies(_ cookieHeader: String) async throws -> AugmentStatusSnapshot { + try await self.fetchWithCookieHeader(cookieHeader) + } + + /// Fetch Augment usage using browser cookies with fallback to stored session. + public func fetch(cookieHeaderOverride: String? = nil, logger: ((String) -> Void)? = nil) + async throws -> AugmentStatusSnapshot + { + let log: (String) -> Void = { msg in logger?("[augment] \(msg)") } + + if let override = CookieHeaderNormalizer.normalize(cookieHeaderOverride) { + log("Using manual cookie header") + return try await self.fetchWithCookieHeader(override) + } + + if let cached = CookieHeaderCache.load(provider: .augment), + !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + log("Using cached cookie header from \(cached.sourceLabel)") + do { + return try await self.fetchWithCookieHeader(cached.cookieHeader) + } catch let error as AugmentStatusProbeError { + switch error { + case .notLoggedIn, .sessionExpired: + CookieHeaderCache.clear(provider: .augment) + default: + throw error + } + } catch { + throw error + } + } + + // Try importing cookies from the configured browser order first. + do { + let session = try AugmentCookieImporter.importSession(logger: log) + log("Using cookies from \(session.sourceLabel)") + let snapshot = try await self.fetchWithCookieHeader(session.cookieHeader) + + // SUCCESS: Save cookies to fallback store for future use + await AugmentSessionStore.shared.setCookies(session.cookies) + log("Saved session cookies to fallback store") + CookieHeaderCache.store( + provider: .augment, + cookieHeader: session.cookieHeader, + sourceLabel: session.sourceLabel) + + return snapshot + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + log("Browser cookie import failed: \(error.localizedDescription)") + } + + // Fall back to stored session cookies (from previous successful fetch or "Add Account" login flow) + let storedCookies = await AugmentSessionStore.shared.getCookies() + if !storedCookies.isEmpty { + log("Using stored session cookies") + let cookieHeader = storedCookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + do { + return try await self.fetchWithCookieHeader(cookieHeader) + } catch { + if case AugmentStatusProbeError.notLoggedIn = error { + // Clear only when auth is invalid; keep for transient failures. + await AugmentSessionStore.shared.clearCookies() + log("Stored session invalid, cleared") + } else if case AugmentStatusProbeError.sessionExpired = error { + await AugmentSessionStore.shared.clearCookies() + log("Stored session expired, cleared") + } else { + log("Stored session failed: \(error.localizedDescription)") + } + } + } + + throw AugmentStatusProbeError.noSessionCookie + } + + private func fetchWithCookieHeader(_ cookieHeader: String) async throws -> AugmentStatusSnapshot { + // Fetch credits (required) + let (creditsResponse, creditsJSON) = try await self.fetchCredits(cookieHeader: cookieHeader) + + // Fetch subscription (optional - provides plan name and billing cycle) + let subscriptionResult: (AugmentSubscriptionResponse?, String?) = await { + do { + let (response, json) = try await self.fetchSubscription(cookieHeader: cookieHeader) + return (response, json) + } catch { + // Subscription API is optional - don't fail the whole fetch if it's unavailable + return (nil, nil) + } + }() + + return self.parseResponse( + credits: creditsResponse, + subscription: subscriptionResult.0, + creditsJSON: creditsJSON, + subscriptionJSON: subscriptionResult.1) + } + + private func fetchCredits(cookieHeader: String) async throws -> (AugmentCreditsResponse, String) { + let url = self.baseURL.appendingPathComponent("/api/credits") + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + + let (data, response) = try await ProviderHTTPClient.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw AugmentStatusProbeError.networkError("Invalid response") + } + + if httpResponse.statusCode == 401 { + // Session expired - trigger automatic recovery + throw AugmentStatusProbeError.sessionExpired + } + + if httpResponse.statusCode == 403 { + let responseBody = String(data: data, encoding: .utf8) ?? "" + throw AugmentStatusProbeError.networkError("HTTP \(httpResponse.statusCode): \(responseBody)") + } + + guard httpResponse.statusCode == 200 else { + let responseBody = String(data: data, encoding: .utf8) ?? "" + throw AugmentStatusProbeError.networkError("HTTP \(httpResponse.statusCode): \(responseBody)") + } + + let rawJSON = String(data: data, encoding: .utf8) ?? "" + let decoder = JSONDecoder() + do { + let response = try decoder.decode(AugmentCreditsResponse.self, from: data) + return (response, rawJSON) + } catch { + throw AugmentStatusProbeError.parseFailed("Credits response: \(error.localizedDescription)") + } + } + + private func fetchSubscription(cookieHeader: String) async throws -> (AugmentSubscriptionResponse, String) { + let url = self.baseURL.appendingPathComponent("/api/subscription") + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + + let (data, response) = try await ProviderHTTPClient.shared.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw AugmentStatusProbeError.networkError("Invalid response") + } + + if httpResponse.statusCode == 401 { + // Session expired - trigger automatic recovery + throw AugmentStatusProbeError.sessionExpired + } + + if httpResponse.statusCode == 403 { + throw AugmentStatusProbeError.notLoggedIn + } + + guard httpResponse.statusCode == 200 else { + throw AugmentStatusProbeError.networkError("HTTP \(httpResponse.statusCode)") + } + + let rawJSON = String(data: data, encoding: .utf8) ?? "" + let decoder = JSONDecoder() + do { + let response = try decoder.decode(AugmentSubscriptionResponse.self, from: data) + return (response, rawJSON) + } catch { + throw AugmentStatusProbeError.parseFailed("Subscription response: \(error.localizedDescription)") + } + } + + private func parseResponse( + credits: AugmentCreditsResponse, + subscription: AugmentSubscriptionResponse?, + creditsJSON: String, + subscriptionJSON: String?) -> AugmentStatusSnapshot + { + // Combine both API responses for debugging + var combinedJSON = "Credits API:\n\(creditsJSON)" + if let subJSON = subscriptionJSON { + combinedJSON += "\n\nSubscription API:\n\(subJSON)" + } + + // Parse billing period end date from ISO8601 string + let billingCycleEnd: Date? = { + guard let dateString = subscription?.billingPeriodEnd else { return nil } + let formatter = ISO8601DateFormatter() + return formatter.date(from: dateString) + }() + + return AugmentStatusSnapshot( + creditsRemaining: credits.credits, + creditsUsed: credits.creditsUsed, + creditsLimit: credits.creditsLimit, + billingCycleEnd: billingCycleEnd, + accountEmail: subscription?.email, + accountPlan: subscription?.planName, + rawJSON: combinedJSON) + } + + /// Debug probe that returns raw API responses + public func debugRawProbe(cookieHeaderOverride: String? = nil) async -> String { + let stamp = ISO8601DateFormatter().string(from: Date()) + var lines: [String] = [] + lines.append("=== Augment Debug Probe @ \(stamp) ===") + lines.append("") + + do { + let snapshot = try await self.fetch( + cookieHeaderOverride: cookieHeaderOverride, + logger: { msg in lines.append("[log] \(msg)") }) + lines.append("") + lines.append("Probe Success") + lines.append("") + lines.append("Credits Balance:") + lines.append(" Remaining: \(snapshot.creditsRemaining?.description ?? "nil")") + lines.append(" Used: \(snapshot.creditsUsed?.description ?? "nil")") + lines.append(" Limit: \(snapshot.creditsLimit?.description ?? "nil")") + lines.append("") + lines.append("Billing Cycle End: \(snapshot.billingCycleEnd?.description ?? "nil")") + lines.append("Account Email: \(snapshot.accountEmail ?? "nil")") + lines.append("Account Plan: \(snapshot.accountPlan ?? "nil")") + + if let rawJSON = snapshot.rawJSON { + lines.append("") + lines.append("Raw API Response:") + lines.append(rawJSON) + } + + let output = lines.joined(separator: "\n") + await MainActor.run { Self.recordDump(output) } + return output + } catch { + lines.append("") + lines.append("Probe Failed: \(error.localizedDescription)") + let output = lines.joined(separator: "\n") + await MainActor.run { Self.recordDump(output) } + return output + } + } + + // MARK: - Dump storage (in-memory ring buffer) + + @MainActor private static var recentDumps: [String] = [] + + @MainActor private static func recordDump(_ text: String) { + if self.recentDumps.count >= 5 { self.recentDumps.removeFirst() } + self.recentDumps.append(text) + } + + public static func latestDumps() async -> String { + await MainActor.run { + let result = Self.recentDumps.joined(separator: "\n\n---\n\n") + return result.isEmpty ? "No Augment probe dumps captured yet." : result + } + } +} + +#else + +// MARK: - Augment (Unsupported) + +public enum AugmentStatusProbeError: LocalizedError, Sendable { + case notSupported + + public var errorDescription: String? { + "Augment is only supported on macOS." + } +} + +public struct AugmentStatusSnapshot: Sendable { + public init() {} + + public func toUsageSnapshot() -> UsageSnapshot { + UsageSnapshot( + primary: RateWindow(usedPercent: 0, windowMinutes: nil, resetsAt: nil, resetDescription: nil), + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: Date(), + identity: nil) + } +} + +public struct AugmentStatusProbe: Sendable { + public init(baseURL: URL = URL(string: "https://app.augmentcode.com")!, timeout: TimeInterval = 15.0) { + _ = baseURL + _ = timeout + } + + public func fetch(cookieHeaderOverride: String? = nil, logger: ((String) -> Void)? = nil) + async throws -> AugmentStatusSnapshot + { + _ = cookieHeaderOverride + _ = logger + throw AugmentStatusProbeError.notSupported + } +} + +#endif diff --git a/Sources/CodexBarCore/Providers/AzureOpenAI/AzureOpenAIProviderDescriptor.swift b/Sources/CodexBarCore/Providers/AzureOpenAI/AzureOpenAIProviderDescriptor.swift new file mode 100644 index 0000000..52859fa --- /dev/null +++ b/Sources/CodexBarCore/Providers/AzureOpenAI/AzureOpenAIProviderDescriptor.swift @@ -0,0 +1,92 @@ +import Foundation + +public enum AzureOpenAIProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .azureopenai, + metadata: ProviderMetadata( + id: .azureopenai, + displayName: "Azure OpenAI", + sessionLabel: "Status", + weeklyLabel: "Deployment", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Azure OpenAI status", + cliName: "azure-openai", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://ai.azure.com", + statusPageURL: nil, + statusLinkURL: "https://azure.status.microsoft/en-us/status"), + branding: ProviderBranding( + iconStyle: .openai, + iconResourceName: "ProviderIcon-codex", + color: ProviderColor(red: 0, green: 120 / 255, blue: 212 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Azure OpenAI usage history is not exposed by the deployment validation probe." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [AzureOpenAIAPIFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "azure-openai", + aliases: ["azureopenai", "aoai"], + versionDetector: nil)) + } +} + +struct AzureOpenAIAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "azureopenai.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + _ = context + // Keep the strategy available so missing partial configuration surfaces + // as a precise settings error instead of a generic no-strategy failure. + return true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveAPIKey(environment: context.env) else { + throw AzureOpenAIUsageError.missingAPIKey + } + try AzureOpenAISettingsReader.validateEndpointOverrides(environment: context.env) + guard let endpoint = Self.resolveEndpoint(environment: context.env) else { + throw AzureOpenAIUsageError.missingEndpoint + } + guard let deploymentName = Self.resolveDeploymentName(environment: context.env) else { + throw AzureOpenAIUsageError.missingDeploymentName + } + + let usage = try await AzureOpenAIUsageFetcher.fetchUsage( + apiKey: apiKey, + endpoint: endpoint, + deploymentName: deploymentName, + apiVersion: AzureOpenAISettingsReader.apiVersion(environment: context.env)) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "deployment") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func resolveAPIKey(environment: [String: String]) -> String? { + ProviderTokenResolver.azureOpenAIToken(environment: environment) + } + + private static func resolveEndpoint(environment: [String: String]) -> URL? { + AzureOpenAISettingsReader.endpoint(environment: environment) + } + + private static func resolveDeploymentName(environment: [String: String]) -> String? { + AzureOpenAISettingsReader.deploymentName(environment: environment) + } +} diff --git a/Sources/CodexBarCore/Providers/AzureOpenAI/AzureOpenAISettingsReader.swift b/Sources/CodexBarCore/Providers/AzureOpenAI/AzureOpenAISettingsReader.swift new file mode 100644 index 0000000..53706dd --- /dev/null +++ b/Sources/CodexBarCore/Providers/AzureOpenAI/AzureOpenAISettingsReader.swift @@ -0,0 +1,82 @@ +import Foundation + +public enum AzureOpenAISettingsReader { + public static let apiKeyEnvironmentKey = "AZURE_OPENAI_API_KEY" + public static let endpointEnvironmentKey = "AZURE_OPENAI_ENDPOINT" + public static let deploymentNameEnvironmentKey = "AZURE_OPENAI_DEPLOYMENT_NAME" + public static let apiVersionEnvironmentKey = "AZURE_OPENAI_API_VERSION" + public static let defaultAPIVersion = "2024-10-21" + + public static func apiKey(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.cleaned(environment[self.apiKeyEnvironmentKey]) + } + + public static func endpoint(environment: [String: String] = ProcessInfo.processInfo.environment) -> URL? { + guard let rawEndpoint = self.rawEndpoint(environment: environment) else { return nil } + return self.endpointURL(from: rawEndpoint) + } + + public static func rawEndpoint(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.cleaned(environment[self.endpointEnvironmentKey]) + } + + public static func deploymentName(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.cleaned(environment[self.deploymentNameEnvironmentKey]) + } + + public static func apiVersion(environment: [String: String] = ProcessInfo.processInfo.environment) -> String { + self.cleaned(environment[self.apiVersionEnvironmentKey]) ?? self.defaultAPIVersion + } + + public static func endpointURL(from rawEndpoint: String) -> URL? { + let trimmed = rawEndpoint.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: trimmed) + } + + public static func validateEndpointOverrides( + environment: [String: String] = ProcessInfo.processInfo.environment) throws + { + guard let rawEndpoint = self.rawEndpoint(environment: environment) else { return } + guard self.endpointURL(from: rawEndpoint) != nil else { + throw AzureOpenAISettingsError.invalidEndpointOverride(self.endpointEnvironmentKey) + } + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} + +public enum AzureOpenAISettingsError: LocalizedError, Sendable, Equatable { + case missingAPIKey + case missingEndpoint + case missingDeploymentName + case invalidEndpointOverride(String) + + public var errorDescription: String? { + switch self { + case .missingAPIKey: + "Azure OpenAI API key not configured. Set AZURE_OPENAI_API_KEY or configure an API key in Settings." + case .missingEndpoint: + "Azure OpenAI endpoint not configured. Set AZURE_OPENAI_ENDPOINT or configure an endpoint in Settings." + case .missingDeploymentName: + "Azure OpenAI deployment not configured. Set AZURE_OPENAI_DEPLOYMENT_NAME or configure a deployment " + + "in Settings." + case let .invalidEndpointOverride(key): + "Azure OpenAI endpoint override \(key) is not allowed. " + + "Use an HTTPS endpoint without user info or encoded host tricks." + } + } +} diff --git a/Sources/CodexBarCore/Providers/AzureOpenAI/AzureOpenAIUsageFetcher.swift b/Sources/CodexBarCore/Providers/AzureOpenAI/AzureOpenAIUsageFetcher.swift new file mode 100644 index 0000000..716d75c --- /dev/null +++ b/Sources/CodexBarCore/Providers/AzureOpenAI/AzureOpenAIUsageFetcher.swift @@ -0,0 +1,244 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct AzureOpenAIUsageSnapshot: Codable, Sendable, Equatable { + public let endpointHost: String + public let deploymentName: String + public let model: String? + public let apiVersion: String + public let updatedAt: Date + + public init( + endpointHost: String, + deploymentName: String, + model: String?, + apiVersion: String, + updatedAt: Date) + { + self.endpointHost = endpointHost + self.deploymentName = deploymentName + self.model = model + self.apiVersion = apiVersion + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let detail = Self.detailText(deploymentName: self.deploymentName, model: self.model) + return UsageSnapshot( + primary: RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: nil, + resetDescription: detail), + secondary: nil, + tertiary: nil, + updatedAt: self.updatedAt, + identity: ProviderIdentitySnapshot( + providerID: .azureopenai, + accountEmail: nil, + accountOrganization: self.endpointHost, + loginMethod: "Deployment: \(self.deploymentName)")) + } + + private static func detailText(deploymentName: String, model: String?) -> String { + let cleanedModel = model?.trimmingCharacters(in: .whitespacesAndNewlines) + guard let cleanedModel, !cleanedModel.isEmpty else { + return "Deployment: \(deploymentName)" + } + return "Deployment: \(deploymentName) · Model: \(cleanedModel)" + } +} + +public enum AzureOpenAIUsageError: LocalizedError, Sendable, Equatable { + case missingAPIKey + case missingEndpoint + case missingDeploymentName + case invalidEndpointOverride(String) + case invalidURL + case networkError(String) + case apiError(statusCode: Int, message: String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingAPIKey: + AzureOpenAISettingsError.missingAPIKey.errorDescription + case .missingEndpoint: + AzureOpenAISettingsError.missingEndpoint.errorDescription + case .missingDeploymentName: + AzureOpenAISettingsError.missingDeploymentName.errorDescription + case let .invalidEndpointOverride(key): + AzureOpenAISettingsError.invalidEndpointOverride(key).errorDescription + case .invalidURL: + "Azure OpenAI validation URL is invalid." + case let .networkError(message): + "Azure OpenAI network error: \(message)" + case let .apiError(statusCode, message): + if message.isEmpty { + "Azure OpenAI API error: HTTP \(statusCode)" + } else { + "Azure OpenAI API error: HTTP \(statusCode): \(message)" + } + case let .parseFailed(message): + "Azure OpenAI response parse error: \(message)" + } + } +} + +private struct AzureOpenAIChatCompletionResponse: Decodable { + let model: String? +} + +public enum AzureOpenAIUsageFetcher { + private static let timeoutSeconds: TimeInterval = 20 + private static let maxErrorBodyLength = 240 + + public static func fetchUsage( + apiKey: String, + endpoint: URL, + deploymentName: String, + apiVersion: String = AzureOpenAISettingsReader.defaultAPIVersion, + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + updatedAt: Date = Date()) async throws -> AzureOpenAIUsageSnapshot + { + let apiKey = apiKey.trimmingCharacters(in: .whitespacesAndNewlines) + let deploymentName = deploymentName.trimmingCharacters(in: .whitespacesAndNewlines) + let apiVersion = apiVersion.trimmingCharacters(in: .whitespacesAndNewlines) + guard !apiKey.isEmpty else { throw AzureOpenAIUsageError.missingAPIKey } + guard !deploymentName.isEmpty else { throw AzureOpenAIUsageError.missingDeploymentName } + guard let endpoint = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: endpoint.absoluteString) else { + throw AzureOpenAIUsageError.invalidEndpointOverride(AzureOpenAISettingsReader.endpointEnvironmentKey) + } + let effectiveAPIVersion = apiVersion.isEmpty ? AzureOpenAISettingsReader.defaultAPIVersion : apiVersion + + var request = try URLRequest(url: self.chatCompletionsURL( + endpoint: endpoint, + deploymentName: deploymentName, + apiVersion: effectiveAPIVersion)) + request.httpMethod = "POST" + request.timeoutInterval = Self.timeoutSeconds + request.setValue(apiKey, forHTTPHeaderField: "api-key") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.httpBody = try self.validationRequestBody( + deploymentName: deploymentName, + apiVersion: effectiveAPIVersion) + + let response: ProviderHTTPResponse + do { + response = try await transport.response(for: request) + } catch { + throw AzureOpenAIUsageError.networkError(error.localizedDescription) + } + + guard (200..<300).contains(response.statusCode) else { + throw AzureOpenAIUsageError.apiError( + statusCode: response.statusCode, + message: self.responseSummary(response.data)) + } + + let model: String? + do { + model = try JSONDecoder().decode(AzureOpenAIChatCompletionResponse.self, from: response.data).model + } catch { + throw AzureOpenAIUsageError.parseFailed(error.localizedDescription) + } + + return AzureOpenAIUsageSnapshot( + endpointHost: endpoint.host ?? endpoint.absoluteString, + deploymentName: deploymentName, + model: model, + apiVersion: effectiveAPIVersion, + updatedAt: updatedAt) + } + + public static func _chatCompletionsURLForTesting( + endpoint: URL, + deploymentName: String, + apiVersion: String) throws -> URL + { + try self.chatCompletionsURL(endpoint: endpoint, deploymentName: deploymentName, apiVersion: apiVersion) + } + + private static func chatCompletionsURL( + endpoint: URL, + deploymentName: String, + apiVersion: String) throws -> URL + { + if self.usesV1API(apiVersion) { + let base = self.apiRoot(endpoint: endpoint, pathComponents: ["openai", "v1"]) + .appendingPathComponent("chat") + .appendingPathComponent("completions") + guard let url = URLComponents(url: base, resolvingAgainstBaseURL: false)?.url else { + throw AzureOpenAIUsageError.invalidURL + } + return url + } + + let base = self.apiRoot(endpoint: endpoint, pathComponents: ["openai"]) + .appendingPathComponent("deployments") + .appendingPathComponent(deploymentName) + .appendingPathComponent("chat") + .appendingPathComponent("completions") + guard var components = URLComponents(url: base, resolvingAgainstBaseURL: false) else { + throw AzureOpenAIUsageError.invalidURL + } + components.queryItems = [URLQueryItem(name: "api-version", value: apiVersion)] + guard let url = components.url else { throw AzureOpenAIUsageError.invalidURL } + return url + } + + private static func apiRoot(endpoint: URL, pathComponents expectedComponents: [String]) -> URL { + let existingComponents = endpoint.pathComponents + .filter { $0 != "/" } + .map { $0.lowercased() } + let expectedComponents = expectedComponents.map { $0.lowercased() } + let sharedCount = stride( + from: min(existingComponents.count, expectedComponents.count), + through: 0, + by: -1) + .first { count in + count == 0 || Array(existingComponents.suffix(count)) == Array(expectedComponents.prefix(count)) + } ?? 0 + return expectedComponents.dropFirst(sharedCount).reduce(endpoint) { url, component in + url.appendingPathComponent(component) + } + } + + private static func validationRequestBody( + deploymentName: String, + apiVersion: String) throws -> Data + { + var payload: [String: Any] = [ + "messages": [ + ["role": "user", "content": "ping"], + ], + ] + if self.usesV1API(apiVersion) { + payload["model"] = deploymentName + payload["max_completion_tokens"] = 1 + } else { + payload["max_tokens"] = 1 + } + return try JSONSerialization.data(withJSONObject: payload, options: []) + } + + private static func usesV1API(_ apiVersion: String) -> Bool { + apiVersion.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() == "v1" + } + + private static func responseSummary(_ data: Data) -> String { + guard let body = String(data: data, encoding: .utf8)? + .replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespacesAndNewlines), + !body.isEmpty + else { + return "" + } + guard body.count > Self.maxErrorBodyLength else { return body } + let index = body.index(body.startIndex, offsetBy: Self.maxErrorBodyLength) + return "\(body[.. SignedHeadersInfo { + var headers: [(String, String)] = [] + if let allHeaders = request.allHTTPHeaderFields { + for (key, value) in allHeaders { + headers.append((key.lowercased(), value.trimmingCharacters(in: .whitespaces))) + } + } + headers.sort { $0.0 < $1.0 } + + let keys = headers.map(\.0).joined(separator: ";") + let canonical = headers.map { "\($0.0):\($0.1)" }.joined(separator: "\n") + return SignedHeadersInfo(keys: keys, canonical: canonical) + } + + private static func canonicalRequest( + request: URLRequest, + signedHeaders: SignedHeadersInfo, + bodyHash: String) -> String + { + let method = request.httpMethod ?? "GET" + let url = request.url! + let path = url.path.isEmpty ? "/" : url.path + let query = Self.canonicalQueryString(url: url) + + return [ + method, + Self.uriEncodePath(path), + query, + signedHeaders.canonical + "\n", + signedHeaders.keys, + bodyHash, + ].joined(separator: "\n") + } + + private static func canonicalQueryString(url: URL) -> String { + guard let components = URLComponents(url: url, resolvingAgainstBaseURL: false), + let queryItems = components.queryItems, + !queryItems.isEmpty + else { + return "" + } + + return queryItems + .map { item in + let key = Self.uriEncode(item.name) + let value = Self.uriEncode(item.value ?? "") + return "\(key)=\(value)" + } + .sorted() + .joined(separator: "&") + } + + private static func calculateSignature( + secretKey: String, + dateStamp: String, + region: String, + service: String, + stringToSign: String) -> String + { + let kDate = Self.hmacSHA256(key: Data("AWS4\(secretKey)".utf8), data: Data(dateStamp.utf8)) + let kRegion = Self.hmacSHA256(key: kDate, data: Data(region.utf8)) + let kService = Self.hmacSHA256(key: kRegion, data: Data(service.utf8)) + let kSigning = Self.hmacSHA256(key: kService, data: Data("aws4_request".utf8)) + let signature = Self.hmacSHA256(key: kSigning, data: Data(stringToSign.utf8)) + return signature.map { String(format: "%02x", $0) }.joined() + } + + private static func hmacSHA256(key: Data, data: Data) -> Data { + let symmetricKey = SymmetricKey(data: key) + let mac = HMAC.authenticationCode(for: data, using: symmetricKey) + return Data(mac) + } + + private static func sha256Hex(_ data: Data) -> String { + let digest = SHA256.hash(data: data) + return digest.map { String(format: "%02x", $0) }.joined() + } + + private static func dateFormatter() -> DateFormatter { + let formatter = DateFormatter() + formatter.dateFormat = "yyyyMMdd'T'HHmmss'Z'" + formatter.timeZone = TimeZone(identifier: "UTC") + formatter.locale = Locale(identifier: "en_US_POSIX") + return formatter + } + + private static func dateStamp(date: Date) -> String { + let formatter = DateFormatter() + formatter.dateFormat = "yyyyMMdd" + formatter.timeZone = TimeZone(identifier: "UTC") + formatter.locale = Locale(identifier: "en_US_POSIX") + return formatter.string(from: date) + } + + private static func uriEncode(_ string: String) -> String { + var allowed = CharacterSet.alphanumerics + allowed.insert(charactersIn: "-._~") + return string.addingPercentEncoding(withAllowedCharacters: allowed) ?? string + } + + private static func uriEncodePath(_ path: String) -> String { + path.split(separator: "/", omittingEmptySubsequences: false) + .map { self.uriEncode(String($0)) } + .joined(separator: "/") + } +} diff --git a/Sources/CodexBarCore/Providers/Bedrock/BedrockCloudWatchUsage.swift b/Sources/CodexBarCore/Providers/Bedrock/BedrockCloudWatchUsage.swift new file mode 100644 index 0000000..98c8bc3 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Bedrock/BedrockCloudWatchUsage.swift @@ -0,0 +1,231 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +struct BedrockClaudeActivity: Equatable { + let inputTokens: Int + let outputTokens: Int + let requestCount: Int +} + +enum BedrockCloudWatchUsageFetcher { + static let lookbackDays = 14 + + private static let maxResponseBytes = 4 * 1024 * 1024 + private static let maxPages = 20 + private static let requestTimeoutSeconds: TimeInterval = 15 + + private enum Metric: String, CaseIterable { + case inputTokens + case outputTokens + case requests + + var cloudWatchName: String { + switch self { + case .inputTokens: "InputTokenCount" + case .outputTokens: "OutputTokenCount" + case .requests: "Invocations" + } + } + } + + private struct RequestContext { + let credentials: BedrockAWSSigner.Credentials + let region: String + let now: Date + let endpoint: URL + let transport: any ProviderHTTPTransport + } + + static func fetch( + credentials: BedrockAWSSigner.Credentials, + region: String, + now: Date, + endpointOverride: String?, + transport: any ProviderHTTPTransport) async throws -> BedrockClaudeActivity + { + let totals = try await self.fetchTotals( + credentials: credentials, + region: region, + now: now, + endpointOverride: endpointOverride, + transport: transport) + + return BedrockClaudeActivity( + inputTokens: totals[.inputTokens] ?? 0, + outputTokens: totals[.outputTokens] ?? 0, + requestCount: totals[.requests] ?? 0) + } + + private static func fetchTotals( + credentials: BedrockAWSSigner.Credentials, + region: String, + now: Date, + endpointOverride: String?, + transport: any ProviderHTTPTransport) async throws -> [Metric: Int] + { + let context = try RequestContext( + credentials: credentials, + region: region, + now: now, + endpoint: self.endpoint(region: region, override: endpointOverride), + transport: transport) + var totals: [Metric: Double] = [:] + var nextToken: String? + var seenTokens: Set = [] + var pageCount = 0 + + repeat { + pageCount += 1 + guard pageCount <= self.maxPages else { + throw BedrockUsageError.cloudWatchParseFailed("too many response pages") + } + + let response = try await self.callPage( + context: context, + nextToken: nextToken) + for (metric, value) in response.totals { + totals[metric, default: 0] += value + } + nextToken = response.nextToken + if let nextToken, !seenTokens.insert(nextToken).inserted { + throw BedrockUsageError.cloudWatchParseFailed("repeated NextToken") + } + } while nextToken != nil + + var converted: [Metric: Int] = [:] + for metric in Metric.allCases { + let total = totals[metric] ?? 0 + guard total.isFinite, total >= 0, total <= Double(Int.max) else { + throw BedrockUsageError.cloudWatchParseFailed("invalid metric total") + } + converted[metric] = Int(total.rounded()) + } + return converted + } + + private static func callPage( + context: RequestContext, + nextToken: String?) async throws + -> (totals: [Metric: Double], nextToken: String?) + { + let start = context.now.addingTimeInterval(-Double(self.lookbackDays) * 24 * 60 * 60) + let queries: [[String: Any]] = Metric.allCases.map { metric in + let search = "SEARCH('{AWS/Bedrock,ModelId} " + + "MetricName=\"\(metric.cloudWatchName)\" claude', 'Sum', 86400)" + return [ + "Id": metric.rawValue, + "Expression": "SUM(\(search))", + "ReturnData": true, + ] + } + var payload: [String: Any] = [ + "StartTime": start.timeIntervalSince1970, + "EndTime": context.now.timeIntervalSince1970, + "ScanBy": "TimestampAscending", + "MetricDataQueries": queries, + ] + if let nextToken { + payload["NextToken"] = nextToken + } + + var request = URLRequest(url: context.endpoint) + request.httpMethod = "POST" + request.httpBody = try JSONSerialization.data(withJSONObject: payload) + request.timeoutInterval = self.requestTimeoutSeconds + request.setValue("application/x-amz-json-1.0", forHTTPHeaderField: "Content-Type") + request.setValue( + "GraniteServiceVersion20100801.GetMetricData", + forHTTPHeaderField: "X-Amz-Target") + BedrockAWSSigner.sign( + request: &request, + credentials: context.credentials, + region: context.region, + service: "monitoring") + + let response = try await context.transport.response(for: request) + guard response.statusCode == 200 else { + throw BedrockUsageError.cloudWatchAPIError("HTTP \(response.statusCode)") + } + guard response.data.count <= self.maxResponseBytes else { + throw BedrockUsageError.cloudWatchParseFailed("response exceeds 4 MiB") + } + return try self.parsePage(response.data) + } + + private static func endpoint(region: String, override: String?) throws -> URL { + if override != nil { + guard let override = BedrockSettingsReader.cleaned(override), + let url = ProviderEndpointOverrideValidator() + .validatedURLAllowingLoopbackHTTP(override) + else { + throw BedrockUsageError.cloudWatchParseFailed("invalid endpoint override") + } + return url + } + guard region.range( + of: #"^[a-z0-9]+(?:-[a-z0-9]+)+-[0-9]+$"#, + options: .regularExpression) != nil + else { + throw BedrockUsageError.cloudWatchParseFailed("invalid region endpoint") + } + // Match the CloudWatch partition suffixes published by the AWS SDK endpoint resolver. + let suffix = switch region { + case let region where region.hasPrefix("cn-"): + "amazonaws.com.cn" + case let region where region.hasPrefix("eusc-"): + "amazonaws.eu" + case let region where region.hasPrefix("us-iso-"): + "c2s.ic.gov" + case let region where region.hasPrefix("us-isob-"): + "sc2s.sgov.gov" + case let region where region.hasPrefix("eu-isoe-"): + "cloud.adc-e.uk" + case let region where region.hasPrefix("us-isof-"): + "csp.hci.ic.gov" + default: + "amazonaws.com" + } + guard let url = URL(string: "https://monitoring.\(region).\(suffix)") else { + throw BedrockUsageError.cloudWatchParseFailed("invalid region endpoint") + } + return url + } + + private static func parsePage(_ data: Data) throws + -> (totals: [Metric: Double], nextToken: String?) + { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw BedrockUsageError.cloudWatchParseFailed("invalid JSON response") + } + + if let messages = json["Messages"] as? [[String: Any]], !messages.isEmpty { + throw BedrockUsageError.cloudWatchParseFailed("CloudWatch reported incomplete results") + } + + let results = json["MetricDataResults"] as? [[String: Any]] ?? [] + var totals: [Metric: Double] = [:] + for result in results { + guard let id = result["Id"] as? String, let metric = Metric(rawValue: id) else { + throw BedrockUsageError.cloudWatchParseFailed("metric result had an unknown ID") + } + guard result["StatusCode"] as? String == "Complete" else { + throw BedrockUsageError.cloudWatchParseFailed("metric result was incomplete") + } + guard let values = result["Values"] as? [Any] else { continue } + for value in values { + guard let number = value as? NSNumber else { + throw BedrockUsageError.cloudWatchParseFailed("metric value was not numeric") + } + let double = number.doubleValue + guard double.isFinite, double >= 0 else { + throw BedrockUsageError.cloudWatchParseFailed("metric value was invalid") + } + totals[metric, default: 0] += double + } + } + + return (totals, BedrockSettingsReader.cleaned(json["NextToken"] as? String)) + } +} diff --git a/Sources/CodexBarCore/Providers/Bedrock/BedrockCredentialResolver.swift b/Sources/CodexBarCore/Providers/Bedrock/BedrockCredentialResolver.swift new file mode 100644 index 0000000..c5f3cac --- /dev/null +++ b/Sources/CodexBarCore/Providers/Bedrock/BedrockCredentialResolver.swift @@ -0,0 +1,78 @@ +import Foundation + +/// Resolves Bedrock signing credentials + region for the configured auth mode. +/// +/// Shared by the main usage fetch (`BedrockAPIFetchStrategy`) and the daily +/// cost-history refresh (`CostUsageFetcher`) so both honor AWS-profile auth +/// identically instead of the history path silently requiring static keys. +enum BedrockCredentialResolver { + struct Resolved { + let credentials: BedrockAWSSigner.Credentials + let region: String + } + + static func resolve( + environment: [String: String], + resolveAWSBinary: ([String: String]) -> String? = { BinaryLocator.resolveAWSBinary(env: $0) }, + makeProvider: (String) -> BedrockProfileCredentialProvider = { + BedrockProfileCredentialProvider.live(awsBinaryPath: $0) + }) async throws -> Resolved + { + switch BedrockSettingsReader.authMode(environment: environment) { + case .keys: + guard let accessKeyID = BedrockSettingsReader.accessKeyID(environment: environment), + let secretAccessKey = BedrockSettingsReader.secretAccessKey(environment: environment) + else { + throw BedrockUsageError.missingCredentials + } + let credentials = BedrockAWSSigner.Credentials( + accessKeyID: accessKeyID, + secretAccessKey: secretAccessKey, + sessionToken: BedrockSettingsReader.sessionToken(environment: environment)) + return Resolved( + credentials: credentials, + region: BedrockSettingsReader.region(environment: environment)) + + case .profile: + guard let profile = BedrockSettingsReader.profile(environment: environment) else { + throw BedrockUsageError.missingCredentials + } + guard let awsBinary = resolveAWSBinary(environment) else { + throw BedrockUsageError.awsCLINotFound + } + let cliEnvironment = Self.profileCLIEnvironment(environment) + let provider = makeProvider(awsBinary) + let credentials = try await provider.exportCredentials(profile: profile, environment: cliEnvironment) + let region = try await Self.resolveRegion( + provider: provider, + profile: profile, + environment: cliEnvironment) + return Resolved(credentials: credentials, region: region) + } + } + + private static func profileCLIEnvironment(_ environment: [String: String]) -> [String: String] { + var cliEnvironment = environment + // `--profile` selects the requested profile. Leaving AWS_PROFILE set can + // break assume-role profiles that use `credential_source = Environment`; + // keep the source credential variables themselves available. + cliEnvironment[BedrockSettingsReader.profileKey] = nil + return cliEnvironment + } + + private static func resolveRegion( + provider: BedrockProfileCredentialProvider, + profile: String, + environment: [String: String]) async throws -> String + { + if let explicit = BedrockSettingsReader.cleaned(environment[BedrockSettingsReader.regionKeys[0]]) + ?? BedrockSettingsReader.cleaned(environment[BedrockSettingsReader.regionKeys[1]]) + { + return explicit + } + if let derived = try await provider.resolveRegion(profile: profile, environment: environment) { + return derived + } + return BedrockSettingsReader.defaultRegion + } +} diff --git a/Sources/CodexBarCore/Providers/Bedrock/BedrockProfileCredentialProvider.swift b/Sources/CodexBarCore/Providers/Bedrock/BedrockProfileCredentialProvider.swift new file mode 100644 index 0000000..1f692f9 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Bedrock/BedrockProfileCredentialProvider.swift @@ -0,0 +1,87 @@ +import Foundation + +/// Resolves AWS credentials for a named profile by shelling out to the AWS CLI. +/// +/// Uses `aws configure export-credentials`, which transparently resolves static +/// credentials, SSO sessions, assume-role chains, and `credential_process` profiles. +/// The runner is injected so tests never invoke a real `aws` binary. +struct BedrockProfileCredentialProvider { + typealias Runner = @Sendable ( + _ arguments: [String], + _ environment: [String: String]) async throws -> SubprocessResult + + let awsBinaryPath: String + let run: Runner + + /// Production provider that drives the real `aws` binary via `SubprocessRunner`. + static func live(awsBinaryPath: String) -> BedrockProfileCredentialProvider { + BedrockProfileCredentialProvider(awsBinaryPath: awsBinaryPath) { arguments, environment in + try await SubprocessRunner.run( + binary: awsBinaryPath, + arguments: arguments, + environment: environment, + timeout: 20, + label: "aws-bedrock-credentials") + } + } + + func exportCredentials( + profile: String, + environment: [String: String] = [:]) async throws -> BedrockAWSSigner.Credentials + { + let result: SubprocessResult + do { + result = try await self.run( + ["configure", "export-credentials", "--profile", profile, "--format", "process"], + environment) + } catch let SubprocessRunnerError.nonZeroExit(_, stderr) { + throw Self.mapExportError(stderr: stderr, profile: profile) + } + + guard let data = result.stdout.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let accessKeyID = Self.nonEmpty(json["AccessKeyId"] as? String), + let secretAccessKey = Self.nonEmpty(json["SecretAccessKey"] as? String) + else { + throw BedrockUsageError.parseFailed("Could not parse AWS CLI export-credentials output") + } + + return BedrockAWSSigner.Credentials( + accessKeyID: accessKeyID, + secretAccessKey: secretAccessKey, + sessionToken: Self.nonEmpty(json["SessionToken"] as? String)) + } + + /// Returns the profile's configured region, or `nil` when unset. + /// `aws configure get region` exits non-zero when the value is not configured, + /// which is a normal case rather than an error. + func resolveRegion( + profile: String, + environment: [String: String] = [:]) async throws -> String? + { + do { + let result = try await self.run( + ["configure", "get", "region", "--profile", profile], + environment) + return Self.nonEmpty(result.stdout) + } catch SubprocessRunnerError.nonZeroExit { + return nil + } + } + + static func mapExportError(stderr: String, profile: String) -> BedrockUsageError { + let lower = stderr.lowercased() + if lower.contains("sso login") || lower.contains("expired") || lower.contains("token has expired") { + return .profileSessionExpired(profile) + } + let trimmed = stderr.trimmingCharacters(in: .whitespacesAndNewlines) + return .apiError(trimmed.isEmpty ? "AWS CLI failed to export credentials" : trimmed) + } + + private static func nonEmpty(_ raw: String?) -> String? { + guard let value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + return value + } +} diff --git a/Sources/CodexBarCore/Providers/Bedrock/BedrockProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Bedrock/BedrockProviderDescriptor.swift new file mode 100644 index 0000000..d0e3310 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Bedrock/BedrockProviderDescriptor.swift @@ -0,0 +1,74 @@ +import Foundation + +public enum BedrockProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .bedrock, + metadata: ProviderMetadata( + id: .bedrock, + displayName: "AWS Bedrock", + sessionLabel: "Budget", + weeklyLabel: "Cost", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show AWS Bedrock usage", + cliName: "bedrock", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + dashboardURL: "https://console.aws.amazon.com/bedrock", + statusPageURL: nil, + statusLinkURL: "https://health.aws.amazon.com/health/status"), + branding: ProviderBranding( + iconStyle: .bedrock, + iconResourceName: "ProviderIcon-bedrock", + color: ProviderColor(red: 1, green: 0.6, blue: 0)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: true, + noDataMessage: { "No AWS Bedrock cost data available. Check your AWS access keys " + + "or profile, and that the AWS CLI is installed for profile auth." + }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [BedrockAPIFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "bedrock", + aliases: ["aws-bedrock"], + versionDetector: nil)) + } +} + +struct BedrockAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "bedrock.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + switch BedrockSettingsReader.authMode(environment: context.env) { + case .keys: + BedrockSettingsReader.hasCredentials(environment: context.env) + case .profile: + BedrockSettingsReader.profile(environment: context.env) != nil + } + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let resolved = try await BedrockCredentialResolver.resolve(environment: context.env) + let budget = BedrockSettingsReader.budget(environment: context.env) + let usage = try await BedrockUsageFetcher.fetchUsage( + credentials: resolved.credentials, + region: resolved.region, + budget: budget, + environment: context.env) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: any Error, context _: ProviderFetchContext) -> Bool { + false + } +} diff --git a/Sources/CodexBarCore/Providers/Bedrock/BedrockSettingsReader.swift b/Sources/CodexBarCore/Providers/Bedrock/BedrockSettingsReader.swift new file mode 100644 index 0000000..406a8b5 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Bedrock/BedrockSettingsReader.swift @@ -0,0 +1,107 @@ +import Foundation + +public enum BedrockAuthMode: String, Codable, Sendable, CaseIterable { + case keys + case profile +} + +public enum BedrockSettingsReader { + public static let accessKeyIDKey = "AWS_ACCESS_KEY_ID" + public static let secretAccessKeyKey = "AWS_SECRET_ACCESS_KEY" + public static let sessionTokenKey = "AWS_SESSION_TOKEN" + public static let regionKeys = ["AWS_REGION", "AWS_DEFAULT_REGION"] + public static let budgetKey = "CODEXBAR_BEDROCK_BUDGET" + public static let apiURLKey = "CODEXBAR_BEDROCK_API_URL" + public static let cloudWatchAPIURLKey = "CODEXBAR_BEDROCK_CLOUDWATCH_API_URL" + public static let profileKey = "AWS_PROFILE" + public static let authModeKey = "CODEXBAR_BEDROCK_AUTH_MODE" + public static let defaultRegion = "us-east-1" + + public static func accessKeyID(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.cleaned(environment[self.accessKeyIDKey]) + } + + public static func secretAccessKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.secretAccessKeyKey]) + } + + public static func sessionToken( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.sessionTokenKey]) + } + + public static func region(environment: [String: String] = ProcessInfo.processInfo.environment) -> String { + for key in self.regionKeys { + if let value = self.cleaned(environment[key]) { + return value + } + } + return self.defaultRegion + } + + public static func budget(environment: [String: String] = ProcessInfo.processInfo.environment) -> Double? { + guard let raw = self.cleaned(environment[self.budgetKey]), + let value = Double(raw), + value > 0 + else { + return nil + } + return value + } + + public static func profile( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.profileKey]) + } + + public static func authMode( + environment: [String: String] = ProcessInfo.processInfo.environment) -> BedrockAuthMode + { + if let raw = self.cleaned(environment[self.authModeKey])?.lowercased(), + let mode = BedrockAuthMode(rawValue: raw) + { + return mode + } + if self.profile(environment: environment) != nil, + !self.hasStaticKeys(environment: environment) + { + return .profile + } + return .keys + } + + static func hasStaticKeys(environment: [String: String]) -> Bool { + self.accessKeyID(environment: environment) != nil && + self.secretAccessKey(environment: environment) != nil + } + + public static func hasCredentials( + environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool + { + switch self.authMode(environment: environment) { + case .keys: + self.hasStaticKeys(environment: environment) + case .profile: + self.profile(environment: environment) != nil + } + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} diff --git a/Sources/CodexBarCore/Providers/Bedrock/BedrockUsageStats.swift b/Sources/CodexBarCore/Providers/Bedrock/BedrockUsageStats.swift new file mode 100644 index 0000000..5852169 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Bedrock/BedrockUsageStats.swift @@ -0,0 +1,503 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct BedrockUsageSnapshot: Codable, Sendable { + public let monthlySpend: Double + public let monthlyBudget: Double? + public let inputTokens: Int? + public let outputTokens: Int? + public let requestCount: Int? + public let region: String + public let updatedAt: Date + + public init( + monthlySpend: Double, + monthlyBudget: Double?, + inputTokens: Int? = nil, + outputTokens: Int? = nil, + requestCount: Int? = nil, + region: String, + updatedAt: Date) + { + self.monthlySpend = monthlySpend + self.monthlyBudget = monthlyBudget + self.inputTokens = inputTokens + self.outputTokens = outputTokens + self.requestCount = requestCount + self.region = region + self.updatedAt = updatedAt + } + + public var budgetUsedPercent: Double? { + guard let budget = self.monthlyBudget, budget > 0 else { return nil } + return min(100, max(0, (self.monthlySpend / budget) * 100)) + } + + public var totalTokens: Int? { + guard let input = self.inputTokens, let output = self.outputTokens else { return nil } + return input + output + } +} + +extension BedrockUsageSnapshot { + public func toUsageSnapshot() -> UsageSnapshot { + let primary: RateWindow? = if let usedPercent = self.budgetUsedPercent { + RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: Self.endOfCurrentMonth(), + resetDescription: "Monthly budget") + } else { + nil + } + + let cost = ProviderCostSnapshot( + used: self.monthlySpend, + limit: self.monthlyBudget ?? 0, + currencyCode: "USD", + period: "Monthly", + resetsAt: Self.endOfCurrentMonth(), + updatedAt: self.updatedAt) + + var loginParts: [String] = [] + loginParts.append(String(format: "Spend: $%.2f", self.monthlySpend)) + if let budget = self.monthlyBudget { + loginParts.append(String(format: "Budget: $%.2f", budget)) + } + if let total = self.totalTokens { + loginParts.append("Claude 14d: \(Self.formattedTokenCount(total)) tokens") + } + if let requestCount = self.requestCount { + loginParts.append("Requests: \(Self.formattedTokenCount(requestCount))") + } + + let identity = ProviderIdentitySnapshot( + providerID: .bedrock, + accountEmail: nil, + accountOrganization: nil, + loginMethod: loginParts.joined(separator: " - ")) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: cost, + updatedAt: self.updatedAt, + identity: identity) + } + + private static func endOfCurrentMonth() -> Date? { + let calendar = Calendar.current + guard let range = calendar.range(of: .day, in: .month, for: Date()) else { return nil } + let components = calendar.dateComponents([.year, .month], from: Date()) + guard let startOfMonth = calendar.date(from: components) else { return nil } + return calendar.date(byAdding: .day, value: range.count, to: startOfMonth) + } + + static func formattedTokenCount(_ count: Int) -> String { + if count >= 1_000_000 { + return String(format: "%.1fM", Double(count) / 1_000_000) + } else if count >= 1000 { + return String(format: "%.1fK", Double(count) / 1000) + } + return "\(count)" + } +} + +enum BedrockUsageFetcher { + private static let log = CodexBarLog.logger(LogCategories.bedrockUsage) + private static let requestTimeoutSeconds: TimeInterval = 15 + + private struct CostExplorerQuery { + let startDate: String + let endDate: String + let granularity: String + let nextPageToken: String? + } + + static func fetchUsage( + credentials: BedrockAWSSigner.Credentials, + region: String, + budget: Double?, + environment: [String: String] = ProcessInfo.processInfo.environment, + now: Date = Date(), + cloudWatchTransport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws + -> BedrockUsageSnapshot + { + let spend = try await Self.fetchMonthlyCost( + credentials: credentials, + environment: environment) + + var activity: BedrockClaudeActivity? + let cloudWatchOverride = environment[BedrockSettingsReader.cloudWatchAPIURLKey] + let shouldFetchCloudWatch = environment[BedrockSettingsReader.apiURLKey] == nil || cloudWatchOverride != nil + if shouldFetchCloudWatch { + do { + activity = try await BedrockCloudWatchUsageFetcher.fetch( + credentials: credentials, + region: region, + now: now, + endpointOverride: cloudWatchOverride, + transport: cloudWatchTransport) + } catch { + if Task.isCancelled { throw error } + Self.log.debug("AWS CloudWatch Claude activity unavailable; keeping Cost Explorer usage.") + } + } + + return BedrockUsageSnapshot( + monthlySpend: spend, + monthlyBudget: budget, + inputTokens: activity?.inputTokens, + outputTokens: activity?.outputTokens, + requestCount: activity?.requestCount, + region: region, + updatedAt: now) + } + + static func fetchDailyReport( + credentials: BedrockAWSSigner.Credentials, + since: Date, + until: Date, + environment: [String: String] = ProcessInfo.processInfo.environment) async throws + -> CostUsageDailyReport + { + let formatter = Self.dateFormatter() + let startDate = formatter.string(from: since) + let inclusiveEnd = Self.utcCalendar().date(byAdding: .day, value: 1, to: until) ?? until + let endDate = formatter.string(from: inclusiveEnd) + + let pages = try await Self.callCostExplorerPages( + startDate: startDate, + endDate: endDate, + granularity: "DAILY", + credentials: credentials, + environment: environment) + + let entries = try Self.parseDailyResponses(pages) + return CostUsageDailyReport(data: entries, summary: nil) + } + + private static func fetchMonthlyCost( + credentials: BedrockAWSSigner.Credentials, + environment: [String: String]) async throws -> Double + { + let (startDate, endDate) = Self.currentMonthRange() + + let pages = try await Self.callCostExplorerPages( + startDate: startDate, + endDate: endDate, + granularity: "MONTHLY", + credentials: credentials, + environment: environment) + + return try Self.parseTotalCost(pages) + } + + private static func callCostExplorerPages( + startDate: String, + endDate: String, + granularity: String, + credentials: BedrockAWSSigner.Credentials, + environment: [String: String]) async throws -> [Data] + { + var pages: [Data] = [] + var nextPageToken: String? + var seenPageTokens: Set = [] + + repeat { + let page = try await Self.callCostExplorerPage( + query: CostExplorerQuery( + startDate: startDate, + endDate: endDate, + granularity: granularity, + nextPageToken: nextPageToken), + credentials: credentials, + environment: environment) + pages.append(page) + nextPageToken = try Self.nextPageToken(from: page) + if let nextPageToken, !seenPageTokens.insert(nextPageToken).inserted { + throw BedrockUsageError.parseFailed("Cost Explorer returned repeated NextPageToken") + } + } while nextPageToken != nil + + return pages + } + + private static func callCostExplorerPage( + query: CostExplorerQuery, + credentials: BedrockAWSSigner.Credentials, + environment: [String: String]) async throws -> Data + { + let ceRegion = "us-east-1" + let baseURL: URL = if environment[BedrockSettingsReader.apiURLKey] != nil { + if let override = BedrockSettingsReader.cleaned(environment[BedrockSettingsReader.apiURLKey]), + let url = ProviderEndpointOverrideValidator().validatedURLAllowingLoopbackHTTP(override) + { + url + } else { + throw BedrockUsageError.parseFailed("invalid endpoint override") + } + } else { + URL(string: "https://ce.\(ceRegion).amazonaws.com")! + } + + var requestBody: [String: Any] = [ + "TimePeriod": [ + "Start": query.startDate, + "End": query.endDate, + ], + "Granularity": query.granularity, + "Metrics": ["UnblendedCost"], + "GroupBy": [ + ["Type": "DIMENSION", "Key": "SERVICE"], + ], + ] + if let nextPageToken = query.nextPageToken { + requestBody["NextPageToken"] = nextPageToken + } + + let bodyData = try JSONSerialization.data(withJSONObject: requestBody) + + var request = URLRequest(url: baseURL) + request.httpMethod = "POST" + request.httpBody = bodyData + request.setValue("application/x-amz-json-1.1", forHTTPHeaderField: "Content-Type") + request.setValue( + "AWSInsightsIndexService.GetCostAndUsage", + forHTTPHeaderField: "X-Amz-Target") + request.timeoutInterval = Self.requestTimeoutSeconds + + BedrockAWSSigner.sign( + request: &request, + credentials: credentials, + region: ceRegion, + service: "ce") + + let response = try await ProviderHTTPClient.shared.response(for: request) + guard response.statusCode == 200 else { + if Self.isDataUnavailableResponse(statusCode: response.statusCode, data: response.data) { + Self.log.info("AWS Cost Explorer data unavailable, assuming zero usage.") + return Data(#"{"ResultsByTime":[]}"#.utf8) + } + let summary = Self.sanitizedResponseBody(response.data) + Self.log.error("AWS Cost Explorer returned \(response.statusCode): \(summary)") + throw BedrockUsageError.apiError("HTTP \(response.statusCode)") + } + + return response.data + } + + private static func nextPageToken(from data: Data) throws -> String? { + guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw BedrockUsageError.parseFailed("Invalid Cost Explorer response") + } + return BedrockSettingsReader.cleaned(json["NextPageToken"] as? String) + } + + private static func parseTotalCost(_ pages: [Data]) throws -> Double { + var total = 0.0 + for page in pages { + total += try Self.parseTotalCost(page) + } + return total + } + + private static func parseTotalCost(_ data: Data) throws -> Double { + var total = 0.0 + for (_, cost, _) in try Self.parseGroupedResults(data) { + total += cost + } + return total + } + + private static func parseDailyResponses(_ pages: [Data]) throws -> [CostUsageDailyReport.Entry] { + let reports = try pages.map { page in + try CostUsageDailyReport(data: Self.parseDailyResponse(page), summary: nil) + } + return CostUsageDailyReport.merged(reports).data + } + + private static func parseDailyResponse(_ data: Data) throws -> [CostUsageDailyReport.Entry] { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let results = json["ResultsByTime"] as? [[String: Any]] + else { + throw BedrockUsageError.parseFailed("Missing ResultsByTime in Cost Explorer response") + } + + var entries: [CostUsageDailyReport.Entry] = [] + for result in results { + guard let timePeriod = result["TimePeriod"] as? [String: String], + let dateStr = timePeriod["Start"] + else { continue } + + var dayCost = 0.0 + var breakdowns: [CostUsageDailyReport.ModelBreakdown] = [] + + if let groups = result["Groups"] as? [[String: Any]] { + for group in groups { + guard let keys = group["Keys"] as? [String], + let serviceName = keys.first, + serviceName.localizedCaseInsensitiveContains("Bedrock") + else { continue } + + if let metrics = group["Metrics"] as? [String: Any], + let unblended = metrics["UnblendedCost"] as? [String: Any], + let amountStr = unblended["Amount"] as? String, + let amount = Double(amountStr), + amount > 0 + { + dayCost += amount + breakdowns.append(CostUsageDailyReport.ModelBreakdown( + modelName: serviceName, + costUSD: amount)) + } + } + } + + guard dayCost > 0 else { continue } + + entries.append(CostUsageDailyReport.Entry( + date: dateStr, + inputTokens: nil, + outputTokens: nil, + totalTokens: nil, + costUSD: dayCost, + modelsUsed: breakdowns.map(\.modelName), + modelBreakdowns: breakdowns.isEmpty ? nil : breakdowns)) + } + + return entries + } + + private static func parseGroupedResults(_ data: Data) throws + -> [(service: String, cost: Double, date: String)] + { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let results = json["ResultsByTime"] as? [[String: Any]] + else { + throw BedrockUsageError.parseFailed("Missing ResultsByTime in Cost Explorer response") + } + + var items: [(service: String, cost: Double, date: String)] = [] + for result in results { + let dateStr = (result["TimePeriod"] as? [String: String])?["Start"] ?? "" + guard let groups = result["Groups"] as? [[String: Any]] else { continue } + for group in groups { + guard let keys = group["Keys"] as? [String], + let serviceName = keys.first, + serviceName.localizedCaseInsensitiveContains("Bedrock") + else { continue } + + if let metrics = group["Metrics"] as? [String: Any], + let unblended = metrics["UnblendedCost"] as? [String: Any], + let amountStr = unblended["Amount"] as? String, + let amount = Double(amountStr) + { + items.append((serviceName, amount, dateStr)) + } + } + } + return items + } + + private static func dateFormatter() -> DateFormatter { + let formatter = DateFormatter() + formatter.dateFormat = "yyyy-MM-dd" + formatter.timeZone = TimeZone(identifier: "UTC") + formatter.locale = Locale(identifier: "en_US_POSIX") + return formatter + } + + private static func utcCalendar() -> Calendar { + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = TimeZone(secondsFromGMT: 0)! + return calendar + } + + static func currentMonthRange(now: Date = Date()) -> (start: String, end: String) { + let calendar = Self.utcCalendar() + let components = calendar.dateComponents([.year, .month], from: now) + let startOfMonth = calendar.date(from: components)! + + let formatter = Self.dateFormatter() + let startOfToday = calendar.startOfDay(for: now) + let tomorrow = calendar.date(byAdding: .day, value: 1, to: startOfToday)! + return (formatter.string(from: startOfMonth), formatter.string(from: tomorrow)) + } + + private static func isDataUnavailableResponse(statusCode: Int, data: Data) -> Bool { + guard statusCode == 400, + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return false + } + + let nestedError = json["Error"] as? [String: Any] + let candidates = [ + json["__type"], + json["code"], + json["Code"], + nestedError?["Code"], + ] + return candidates.compactMap { $0 as? String }.contains { rawCode in + rawCode.split(separator: "#").last == "DataUnavailableException" + } + } + + private static func sanitizedResponseBody(_ data: Data) -> String { + guard !data.isEmpty, + let body = String(bytes: data, encoding: .utf8) + else { + return "empty body" + } + + let trimmed = body.replacingOccurrences( + of: #"\s+"#, + with: " ", + options: .regularExpression) + .trimmingCharacters(in: .whitespacesAndNewlines) + + if trimmed.count > 240 { + let index = trimmed.index(trimmed.startIndex, offsetBy: 240) + return "\(trimmed[.. ProviderDescriptor { + ProviderDescriptor( + id: .chutes, + metadata: ProviderMetadata( + id: .chutes, + displayName: "Chutes", + sessionLabel: "4-hour quota", + weeklyLabel: "Monthly quota", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "Subscription usage from the Chutes API.", + toggleTitle: "Show Chutes usage", + cliName: "chutes", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://chutes.ai", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .chutes, + iconResourceName: "ProviderIcon-chutes", + color: ProviderColor(red: 49 / 255, green: 132 / 255, blue: 255 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Chutes cost history is not available from CodexBar." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [ChutesAPIFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "chutes", + aliases: ["chutes.ai"], + versionDetector: nil)) + } +} + +struct ChutesAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "chutes.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + ChutesSettingsReader.apiKey(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = ChutesSettingsReader.apiKey(environment: context.env) else { + throw ChutesSettingsError.missingToken + } + + let usage = try await ChutesUsageFetcher.fetchUsage( + apiKey: apiKey, + environment: context.env) + + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} diff --git a/Sources/CodexBarCore/Providers/Chutes/ChutesSettingsReader.swift b/Sources/CodexBarCore/Providers/Chutes/ChutesSettingsReader.swift new file mode 100644 index 0000000..f1bb5af --- /dev/null +++ b/Sources/CodexBarCore/Providers/Chutes/ChutesSettingsReader.swift @@ -0,0 +1,63 @@ +import Foundation + +public struct ChutesSettingsReader: Sendable { + public static let apiKeyEnvironmentKey = "CHUTES_API_KEY" + public static let apiURLEnvironmentKey = "CHUTES_API_URL" + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.apiKeyEnvironmentKey]) + } + + public static func apiURL( + environment: [String: String] = ProcessInfo.processInfo.environment) -> URL + { + if let override = self.validAPIURL(environment: environment) { + return override + } + return URL(string: "https://api.chutes.ai")! + } + + public static func validateEndpointOverrides( + environment: [String: String] = ProcessInfo.processInfo.environment) throws + { + guard let raw = self.cleaned(environment[self.apiURLEnvironmentKey]) else { return } + guard ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) == nil else { return } + throw ChutesSettingsError.invalidEndpointOverride(self.apiURLEnvironmentKey) + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } + + private static func validAPIURL(environment: [String: String]) -> URL? { + guard let raw = self.cleaned(environment[self.apiURLEnvironmentKey]) else { return nil } + return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) + } +} + +public enum ChutesSettingsError: LocalizedError, Sendable, Equatable { + case missingToken + case invalidEndpointOverride(String) + + public var errorDescription: String? { + switch self { + case .missingToken: + "Chutes API key not found. Set apiKey in ~/.codexbar/config.json or CHUTES_API_KEY." + case let .invalidEndpointOverride(key): + "Chutes endpoint override \(key) must use HTTPS or a bare host." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Chutes/ChutesUsageStats.swift b/Sources/CodexBarCore/Providers/Chutes/ChutesUsageStats.swift new file mode 100644 index 0000000..3dd9a88 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Chutes/ChutesUsageStats.swift @@ -0,0 +1,1045 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum ChutesUsageError: LocalizedError, Sendable { + case missingCredentials + case invalidCredentials + case invalidURL + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Chutes API key. Set apiKey in ~/.codexbar/config.json or CHUTES_API_KEY." + case .invalidCredentials: + "Chutes API key was rejected. Check the API key in Settings." + case .invalidURL: + "Chutes usage URL is invalid." + case let .apiError(message): + "Chutes usage API error: \(message)" + case let .parseFailed(message): + "Chutes usage parse error: \(message)" + } + } +} + +public enum ChutesSubscriptionState: String, Sendable, Codable { + case active + case inactive + case unknown +} + +public struct ChutesQuotaWindow: Sendable, Equatable { + public let label: String? + public let used: Double? + public let limit: Double? + public let remaining: Double? + public let usedPercent: Double? + public let windowMinutes: Int? + public let resetsAt: Date? + public let unit: String? + + public init( + label: String?, + used: Double?, + limit: Double?, + remaining: Double?, + usedPercent: Double?, + windowMinutes: Int?, + resetsAt: Date?, + unit: String? = nil) + { + self.label = label + self.used = used + self.limit = limit + self.remaining = remaining + self.usedPercent = usedPercent + self.windowMinutes = windowMinutes + self.resetsAt = resetsAt + self.unit = unit + } + + public func rateWindow(defaultWindowMinutes: Int?) -> RateWindow? { + guard let percent = self.usagePercent else { return nil } + return RateWindow( + usedPercent: max(0, min(percent, 100)), + windowMinutes: self.windowMinutes ?? defaultWindowMinutes, + resetsAt: self.resetsAt, + resetDescription: self.usageDescription) + } + + public var usagePercent: Double? { + if let usedPercent { return max(0, min(usedPercent, 100)) } + + var used = self.used + var limit = self.limit + var remaining = self.remaining + + if limit == nil, let used, let remaining { + limit = used + remaining + } + if used == nil, let limit, let remaining { + used = limit - remaining + } + if remaining == nil, let limit, let used { + remaining = max(0, limit - used) + } + + guard let used, let limit, limit > 0 else { return nil } + return (used / limit) * 100 + } + + private var usageDescription: String? { + guard let limit, limit > 0 else { return nil } + + let used: Double + if let explicitUsed = self.used { + used = explicitUsed + } else if let remaining { + used = max(0, limit - remaining) + } else { + return nil + } + + let unit = self.unit?.trimmingCharacters(in: .whitespacesAndNewlines) + let suffix = if let unit, !unit.isEmpty { + " \(unit)" + } else { + "" + } + return "\(Self.formatAmount(used))/\(Self.formatAmount(limit))\(suffix)" + } + + private static func formatAmount(_ value: Double) -> String { + let rounded = value.rounded() + if abs(value - rounded) < 0.0001 { + return String(Int(rounded)) + } + + var text = String(format: "%.2f", value) + while text.contains("."), text.hasSuffix("0") { + text.removeLast() + } + if text.hasSuffix(".") { + text.removeLast() + } + return text + } +} + +public struct ChutesUsageSnapshot: Sendable, Equatable { + public let rollingWindow: ChutesQuotaWindow? + public let monthlyWindow: ChutesQuotaWindow? + public let fallbackWindows: [ChutesQuotaWindow] + public let subscriptionState: ChutesSubscriptionState + public let planName: String? + public let subscriptionRenewsAt: Date? + public let updatedAt: Date + + public init( + rollingWindow: ChutesQuotaWindow?, + monthlyWindow: ChutesQuotaWindow?, + fallbackWindows: [ChutesQuotaWindow] = [], + subscriptionState: ChutesSubscriptionState, + planName: String?, + subscriptionRenewsAt: Date?, + updatedAt: Date) + { + self.rollingWindow = rollingWindow + self.monthlyWindow = monthlyWindow + self.fallbackWindows = fallbackWindows + self.subscriptionState = subscriptionState + self.planName = planName + self.subscriptionRenewsAt = subscriptionRenewsAt + self.updatedAt = updatedAt + } + + public var hasUsageData: Bool { + self.toUsageSnapshot().hasRateLimitWindows + } + + public func toUsageSnapshot() -> UsageSnapshot { + let fallbackRateWindows = self.fallbackWindows.compactMap { + $0.rateWindow(defaultWindowMinutes: $0.windowMinutes) + } + let monthly = self.monthlyWindow?.rateWindow(defaultWindowMinutes: Self.monthlyWindowMinutes) + let rolling = self.rollingWindow?.rateWindow(defaultWindowMinutes: Self.rollingWindowMinutes) + + let primary = rolling ?? (monthly == nil ? fallbackRateWindows.first : nil) + let secondary: RateWindow? = if let monthly { + monthly + } else if rolling != nil { + fallbackRateWindows.first + } else if primary != nil { + fallbackRateWindows.dropFirst().first + } else { + nil + } + + let identity = ProviderIdentitySnapshot( + providerID: .chutes, + accountEmail: nil, + accountOrganization: nil, + loginMethod: self.loginMethod(hasWindows: primary != nil || secondary != nil)) + + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: nil, + providerCost: nil, + subscriptionRenewsAt: self.subscriptionRenewsAt ?? monthly?.resetsAt, + updatedAt: self.updatedAt, + identity: identity) + } + + func preservingSubscriptionContext(from fallback: ChutesUsageSnapshot) -> ChutesUsageSnapshot { + ChutesUsageSnapshot( + rollingWindow: fallback.rollingWindow ?? self.rollingWindow, + monthlyWindow: fallback.monthlyWindow ?? self.monthlyWindow, + fallbackWindows: fallback.fallbackWindows + self.fallbackWindows, + subscriptionState: fallback.subscriptionState, + planName: fallback.planName, + subscriptionRenewsAt: fallback.subscriptionRenewsAt, + updatedAt: fallback.updatedAt) + } + + private func loginMethod(hasWindows: Bool) -> String? { + if let planName = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines), !planName.isEmpty { + return planName + } + switch self.subscriptionState { + case .active: + return nil + case .inactive: + return "No active subscription" + case .unknown: + return hasWindows ? nil : "No usage data" + } + } + + static let rollingWindowMinutes = 4 * 60 + static let monthlyWindowMinutes = 30 * 24 * 60 +} + +public struct ChutesUsageFetcher: Sendable { + private static let requestTimeoutSeconds: TimeInterval = 15 + + public static func fetchUsage( + apiKey: String, + environment: [String: String] = ProcessInfo.processInfo.environment, + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + now: Date = Date()) async throws -> ChutesUsageSnapshot + { + guard let token = ChutesSettingsReader.cleaned(apiKey) else { + throw ChutesUsageError.missingCredentials + } + try ChutesSettingsReader.validateEndpointOverrides(environment: environment) + + let baseURL = ChutesSettingsReader.apiURL(environment: environment) + let subscription = try await self.fetchSnapshot( + pathComponents: ["users", "me", "subscription_usage"], + apiKey: token, + baseURL: baseURL, + transport: transport, + now: now) + + guard subscription.rollingWindow == nil || subscription.monthlyWindow == nil else { + return subscription + } + + do { + let quotas = try await self.fetchQuotaSnapshot( + apiKey: token, + baseURL: baseURL, + transport: transport, + now: now) + return quotas.hasUsageData ? quotas.preservingSubscriptionContext(from: subscription) : subscription + } catch ChutesUsageError.invalidCredentials { + throw ChutesUsageError.invalidCredentials + } catch { + return subscription + } + } + + private static func fetchSnapshot( + pathComponents: [String], + apiKey: String, + baseURL: URL, + transport: any ProviderHTTPTransport, + now: Date) async throws -> ChutesUsageSnapshot + { + let data = try await self.fetchData( + pathComponents: pathComponents, + apiKey: apiKey, + baseURL: baseURL, + transport: transport) + do { + return try ChutesUsageParser.parse(data: data, now: now) + } catch let error as ChutesUsageError { + throw error + } catch { + throw ChutesUsageError.parseFailed(error.localizedDescription) + } + } + + private static func fetchQuotaSnapshot( + apiKey: String, + baseURL: URL, + transport: any ProviderHTTPTransport, + now: Date) async throws -> ChutesUsageSnapshot + { + let data = try await self.fetchData( + pathComponents: ["users", "me", "quotas"], + apiKey: apiKey, + baseURL: baseURL, + transport: transport) + let fallback = try ChutesUsageParser.parse(data: data, now: now) + let definitions = try self.quotaDefinitions(from: data) + guard !definitions.isEmpty else { return fallback } + + var enrichedDefinitions: [[String: Any]] = [] + for definition in definitions { + guard let identifier = self.quotaIdentifier(from: definition) else { + enrichedDefinitions.append(definition) + continue + } + + do { + let usageData = try await self.fetchData( + pathComponents: ["users", "me", "quota_usage", identifier], + apiKey: apiKey, + baseURL: baseURL, + transport: transport) + guard let usage = try self.responseDictionary(from: usageData) else { + enrichedDefinitions.append(definition) + continue + } + enrichedDefinitions.append(definition.merging(usage) { _, usageValue in usageValue }) + } catch ChutesUsageError.invalidCredentials { + throw ChutesUsageError.invalidCredentials + } catch { + enrichedDefinitions.append(definition) + } + } + + let enrichedData = try JSONSerialization.data( + withJSONObject: ["quotas": enrichedDefinitions], + options: []) + let enriched = try ChutesUsageParser.parse(data: enrichedData, now: now) + return enriched.hasUsageData ? enriched : fallback + } + + private static func fetchData( + pathComponents: [String], + apiKey: String, + baseURL: URL, + transport: any ProviderHTTPTransport) async throws -> Data + { + let url = pathComponents.reduce(baseURL) { partial, component in + partial.appendingPathComponent(component) + } + guard url.scheme?.lowercased() == "https" else { + throw ChutesUsageError.invalidURL + } + + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.timeoutInterval = Self.requestTimeoutSeconds + + let response = try await transport.response(for: request) + guard (200..<300).contains(response.statusCode) else { + if response.statusCode == 401 || response.statusCode == 403 { + throw ChutesUsageError.invalidCredentials + } + throw ChutesUsageError.apiError("HTTP \(response.statusCode): \(Self.responseSummary(response.data))") + } + return response.data + } + + private static func quotaDefinitions(from data: Data) throws -> [[String: Any]] { + let object = try JSONSerialization.jsonObject(with: data, options: []) + if let array = object as? [Any] { + return array.compactMap { $0 as? [String: Any] } + } + guard let dictionary = object as? [String: Any] else { return [] } + if let quotas = dictionary["quotas"] as? [Any] { + return quotas.compactMap { $0 as? [String: Any] } + } + if let data = dictionary["data"] as? [Any] { + return data.compactMap { $0 as? [String: Any] } + } + if let data = dictionary["data"] as? [String: Any], + let quotas = data["quotas"] as? [Any] + { + return quotas.compactMap { $0 as? [String: Any] } + } + return [] + } + + private static func quotaIdentifier(from definition: [String: Any]) -> String? { + for key in ["chute_id", "chuteId", "id"] { + if let value = definition[key] as? String { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return trimmed } + } + if let value = definition[key] as? NSNumber { + return value.stringValue + } + } + return nil + } + + private static func responseDictionary(from data: Data) throws -> [String: Any]? { + guard let dictionary = try JSONSerialization.jsonObject(with: data, options: []) as? [String: Any] else { + return nil + } + return dictionary["data"] as? [String: Any] + ?? dictionary["result"] as? [String: Any] + ?? dictionary + } + + private static func responseSummary(_ data: Data) -> String { + String(bytes: data.prefix(240), encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) + ?? "" + } +} + +private final class ChutesISO8601FormatterBox: @unchecked Sendable { + let lock = NSLock() + let withFractional: ISO8601DateFormatter = { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + return formatter + }() + + let plain: ISO8601DateFormatter = { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime] + return formatter + }() +} + +private enum ChutesTimestampParser { + static let box = ChutesISO8601FormatterBox() + + static func parse(_ value: Any?) -> Date? { + switch value { + case let date as Date: + date + case let number as NSNumber: + self.date(fromEpochValue: number.doubleValue) + case let text as String: + self.parseString(text) + default: + nil + } + } + + private static func parseString(_ raw: String) -> Date? { + let text = raw.trimmingCharacters(in: .whitespacesAndNewlines) + guard !text.isEmpty else { return nil } + if let number = Double(text) { + return self.date(fromEpochValue: number) + } + + self.box.lock.lock() + defer { self.box.lock.unlock() } + return self.box.withFractional.date(from: text) ?? self.box.plain.date(from: text) + } + + private static func date(fromEpochValue value: Double) -> Date? { + guard value.isFinite, value > 0 else { return nil } + let seconds = value > 10_000_000_000 ? value / 1000 : value + return Date(timeIntervalSince1970: seconds) + } +} + +enum ChutesUsageParser { + static func parse(data: Data, now: Date = Date()) throws -> ChutesUsageSnapshot { + let object = try JSONSerialization.jsonObject(with: data, options: []) + return self.parse(object: object, now: now) + } + + private static func parse(object: Any, now: Date) -> ChutesUsageSnapshot { + let root: [String: Any] = { + if let dict = object as? [String: Any] { return dict } + if let array = object as? [Any] { return ["quotas": array] } + return [:] + }() + let dataRoot = self.dictionaryValue(self.value(in: root, keys: ["data", "result"])) ?? root + let subscription = self.subscriptionPayload(root: root, dataRoot: dataRoot) + + let explicitRolling = self.firstDictionary( + in: root, + dataRoot: dataRoot, + keys: self.rollingPayloadKeys) + .flatMap { + self.parseQuota( + $0, + defaultLabel: "4-hour quota", + defaultWindowMinutes: ChutesUsageSnapshot.rollingWindowMinutes) + } + let explicitMonthly = self.firstDictionary( + in: root, + dataRoot: dataRoot, + keys: self.monthlyPayloadKeys) + .flatMap { + self.parseQuota( + $0, + defaultLabel: "Monthly quota", + defaultWindowMinutes: ChutesUsageSnapshot.monthlyWindowMinutes) + } + + let quotaWindows = self.fallbackQuotaObjects(from: root, dataRoot: dataRoot).compactMap { + self.parseQuota($0, defaultLabel: nil, defaultWindowMinutes: nil) + } + let classifiedRolling = explicitRolling ?? quotaWindows.first { self.kind(for: $0) == .rolling } + let classifiedMonthly = explicitMonthly ?? quotaWindows.first { self.kind(for: $0) == .monthly } + + let fallbackWindows = quotaWindows.filter { window in + Optional.some(window) != classifiedRolling && Optional.some(window) != classifiedMonthly + } + + return ChutesUsageSnapshot( + rollingWindow: classifiedRolling, + monthlyWindow: classifiedMonthly, + fallbackWindows: fallbackWindows, + subscriptionState: self.subscriptionState(root: root, dataRoot: dataRoot, subscription: subscription), + planName: self.planName(root: root, dataRoot: dataRoot, subscription: subscription), + subscriptionRenewsAt: self.subscriptionRenewsAt( + root: root, + dataRoot: dataRoot, + subscription: subscription), + updatedAt: now) + } + + private enum WindowKind { + case rolling + case monthly + } + + private static func parseQuota( + _ payload: [String: Any], + defaultLabel: String?, + defaultWindowMinutes: Int?) -> ChutesQuotaWindow? + { + let label = self.firstString(in: payload, keys: self.labelKeys) ?? defaultLabel + let limit = self.firstDouble(in: payload, keys: self.limitKeys) + let used = self.firstDouble(in: payload, keys: self.usedKeys) + let remaining = self.firstDouble(in: payload, keys: self.remainingKeys) + var usedPercent = self.normalizedPercent(self.firstDouble(in: payload, keys: self.percentUsedKeys)) + if usedPercent == nil, + let remainingPercent = self.normalizedPercent(self.firstDouble(in: payload, keys: percentRemainingKeys)) + { + usedPercent = 100 - remainingPercent + } + + let window = self.windowMinutes(from: payload) ?? defaultWindowMinutes + let resetsAt = self.firstDate(in: payload, keys: self.resetKeys) + let unit = self.firstString(in: payload, keys: self.unitKeys) ?? "credits" + + let quota = ChutesQuotaWindow( + label: label, + used: used, + limit: limit, + remaining: remaining, + usedPercent: usedPercent, + windowMinutes: window, + resetsAt: resetsAt, + unit: unit) + return quota.usagePercent == nil ? nil : quota + } + + private static func subscriptionPayload(root: [String: Any], dataRoot: [String: Any]) -> [String: Any]? { + self.firstDictionary(in: root, dataRoot: dataRoot, keys: [ + "subscription", + "subscription_usage", + "subscriptionUsage", + "current_subscription", + "currentSubscription", + "plan", + ]) + } + + private static func subscriptionState( + root: [String: Any], + dataRoot: [String: Any], + subscription: [String: Any]?) -> ChutesSubscriptionState + { + if let active = self.firstBool(in: root, keys: activeKeys) + ?? self.firstBool(in: dataRoot, keys: activeKeys) + ?? subscription.flatMap({ self.firstBool(in: $0, keys: activeKeys) }) + { + return active ? .active : .inactive + } + + let status = self.firstString(in: root, keys: self.statusKeys) + ?? self.firstString(in: dataRoot, keys: self.statusKeys) + ?? subscription.flatMap { self.firstString(in: $0, keys: self.statusKeys) } + let normalizedStatus = status?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + guard let normalizedStatus, !normalizedStatus.isEmpty else { return .unknown } + if normalizedStatus.contains("active") && !normalizedStatus.contains("inactive") { + return .active + } + if normalizedStatus.contains("free") || + normalizedStatus.contains("inactive") || + normalizedStatus.contains("cancel") || + normalizedStatus.contains("none") || + normalizedStatus.contains("expired") + { + return .inactive + } + return .unknown + } + + private static func planName( + root: [String: Any], + dataRoot: [String: Any], + subscription: [String: Any]?) -> String? + { + self.firstString(in: root, keys: self.planKeys) + ?? self.firstString(in: dataRoot, keys: self.planKeys) + ?? subscription.flatMap { self.firstString(in: $0, keys: self.planKeys) } + } + + private static func subscriptionRenewsAt( + root: [String: Any], + dataRoot: [String: Any], + subscription: [String: Any]?) -> Date? + { + self.firstDate(in: root, keys: self.resetKeys) + ?? self.firstDate(in: dataRoot, keys: self.resetKeys) + ?? subscription.flatMap { self.firstDate(in: $0, keys: self.resetKeys) } + } + + private static func firstDictionary( + in root: [String: Any], + dataRoot: [String: Any], + keys: [String]) -> [String: Any]? + { + self.dictionaryValue(self.value(in: root, keys: keys)) + ?? self.dictionaryValue(self.value(in: dataRoot, keys: keys)) + } + + private static func fallbackQuotaObjects(from root: [String: Any], dataRoot: [String: Any]) -> [[String: Any]] { + let candidates: [Any?] = [ + self.value(in: root, keys: self.quotaContainerKeys), + self.value(in: dataRoot, keys: self.quotaContainerKeys), + dataRoot, + root, + ] + + var results: [[String: Any]] = [] + for candidate in candidates { + results.append(contentsOf: self.extractQuotaObjects(from: candidate)) + } + return self.deduplicated(results) + } + + private static func extractQuotaObjects(from candidate: Any?) -> [[String: Any]] { + if let array = candidate as? [Any] { + return array.flatMap { self.extractQuotaObjects(from: $0) } + } + + guard let dict = self.dictionaryValue(candidate) else { return [] } + var results: [[String: Any]] = self.isQuotaPayload(dict) ? [dict] : [] + for value in dict.values { + results.append(contentsOf: self.extractQuotaObjects(from: value)) + } + return results + } + + private static func deduplicated(_ objects: [[String: Any]]) -> [[String: Any]] { + var seen: Set = [] + var unique: [[String: Any]] = [] + for object in objects { + guard let key = try? JSONSerialization.data(withJSONObject: object, options: [.sortedKeys]) else { + unique.append(object) + continue + } + guard seen.insert(key).inserted else { continue } + unique.append(object) + } + return unique + } + + private static func isQuotaPayload(_ payload: [String: Any]) -> Bool { + self.firstDouble(in: payload, keys: self.limitKeys) != nil || + self.firstDouble(in: payload, keys: self.usedKeys) != nil || + self.firstDouble(in: payload, keys: self.remainingKeys) != nil || + self.firstDouble(in: payload, keys: self.percentUsedKeys) != nil || + self.firstDouble(in: payload, keys: self.percentRemainingKeys) != nil + } + + private static func kind(for window: ChutesQuotaWindow) -> WindowKind? { + let label = [ + window.label, + window.unit, + ] + .compactMap { $0?.lowercased() } + .joined(separator: " ") + + if label.contains("rolling") || + label.contains("4h") || + label.contains("4 h") || + label.contains("4-hour") || + label.contains("four hour") || + label.contains("four-hour") || + window.windowMinutes == ChutesUsageSnapshot.rollingWindowMinutes + { + return .rolling + } + + if label.contains("month") || + label.contains("billing") || + label.contains("subscription") || + (window.windowMinutes ?? 0) >= 28 * 24 * 60 + { + return .monthly + } + + return nil + } + + private static func windowMinutes(from payload: [String: Any]) -> Int? { + if let minutes = self.firstDouble(in: payload, keys: windowMinuteKeys) { + return Int(minutes.rounded()) + } + if let hours = self.firstDouble(in: payload, keys: windowHourKeys) { + return Int((hours * 60).rounded()) + } + if let days = self.firstDouble(in: payload, keys: windowDayKeys) { + return Int((days * 24 * 60).rounded()) + } + if let seconds = self.firstDouble(in: payload, keys: windowSecondKeys) { + return Int((seconds / 60).rounded()) + } + if let text = self.firstString(in: payload, keys: windowStringKeys) { + return self.windowMinutes(fromText: text) + } + return nil + } + + static func windowMinutes(fromText raw: String) -> Int? { + let text = raw.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + guard !text.isEmpty else { return nil } + let compact = text.replacingOccurrences(of: " ", with: "") + let scanner = Scanner(string: compact) + guard let value = scanner.scanDouble(), value > 0 else { return nil } + let suffix = String(compact[scanner.currentIndex...]) + if suffix.hasPrefix("min") || suffix == "m" { + return Int(value.rounded()) + } + if suffix.hasPrefix("hour") || suffix.hasPrefix("hr") || suffix == "h" { + return Int((value * 60).rounded()) + } + if suffix.hasPrefix("day") || suffix == "d" { + return Int((value * 24 * 60).rounded()) + } + if suffix.hasPrefix("month") || suffix == "mo" { + return Int((value * 30 * 24 * 60).rounded()) + } + return nil + } + + private static func firstString(in dict: [String: Any], keys: [String]) -> String? { + guard let value = self.value(in: dict, keys: keys) else { return nil } + if let string = value as? String { + let cleaned = string.trimmingCharacters(in: .whitespacesAndNewlines) + return cleaned.isEmpty ? nil : cleaned + } + if let number = value as? NSNumber { + return number.stringValue + } + return nil + } + + private static func firstBool(in dict: [String: Any], keys: [String]) -> Bool? { + guard let value = self.value(in: dict, keys: keys) else { return nil } + if let bool = value as? Bool { return bool } + if let number = value as? NSNumber { return number.boolValue } + if let string = value as? String { + switch string.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() { + case "true", "1", "yes", "active": + return true + case "false", "0", "no", "inactive", "none": + return false + default: + return nil + } + } + return nil + } + + private static func firstDouble(in dict: [String: Any], keys: [String]) -> Double? { + self.doubleValue(self.value(in: dict, keys: keys)) + } + + private static func firstDate(in dict: [String: Any], keys: [String]) -> Date? { + ChutesTimestampParser.parse(self.value(in: dict, keys: keys)) + } + + private static func normalizedPercent(_ value: Double?) -> Double? { + guard let value, value.isFinite else { return nil } + let percent = abs(value) < 1 ? value * 100 : value + return max(0, min(percent, 100)) + } + + private static func value(in dict: [String: Any], keys: [String]) -> Any? { + for key in keys { + let normalized = self.normalizedKey(key) + for (candidateKey, value) in dict where self.normalizedKey(candidateKey) == normalized { + return value + } + } + return nil + } + + private static func dictionaryValue(_ value: Any?) -> [String: Any]? { + value as? [String: Any] + } + + private static func doubleValue(_ value: Any?) -> Double? { + switch value { + case let double as Double: + double.isFinite ? double : nil + case let int as Int: + Double(int) + case let number as NSNumber: + number.doubleValue.isFinite ? number.doubleValue : nil + case let string as String: + self.doubleValue(from: string) + default: + nil + } + } + + private static func doubleValue(from raw: String) -> Double? { + let text = raw.trimmingCharacters(in: .whitespacesAndNewlines) + .replacingOccurrences(of: ",", with: "") + .replacingOccurrences(of: "$", with: "") + .replacingOccurrences(of: "%", with: "") + guard !text.isEmpty, let value = Double(text), value.isFinite else { return nil } + return value + } + + private static func normalizedKey(_ key: String) -> String { + key.lowercased().filter { $0.isLetter || $0.isNumber } + } + + private static let rollingPayloadKeys = [ + "rolling", + "rolling_window", + "rollingWindow", + "rolling_4h", + "rolling4h", + "four_hour", + "fourHour", + "four_hour_usage", + "fourHourUsage", + "window_4h", + "window4h", + ] + private static let monthlyPayloadKeys = [ + "monthly", + "monthly_usage", + "monthlyUsage", + "subscription", + "subscription_usage", + "subscriptionUsage", + "billing_period", + "billingPeriod", + ] + private static let quotaContainerKeys = [ + "quotas", + "quota", + "quota_usage", + "quotaUsage", + "limits", + "usage", + "entries", + "subscription_usage", + "subscriptionUsage", + ] + private static let labelKeys = [ + "label", + "name", + "title", + "type", + "quota_type", + "quotaType", + "period", + "window", + "window_name", + "windowName", + "chute_id", + "chuteId", + ] + private static let limitKeys = [ + "limit", + "cap", + "max", + "maximum", + "quota", + "quota_limit", + "quotaLimit", + "monthly_cap", + "monthlyCap", + "monthly_limit", + "monthlyLimit", + "request_limit", + "requestLimit", + "token_limit", + "tokenLimit", + "hard_limit", + "hardLimit", + "total", + ] + private static let usedKeys = [ + "used", + "usage", + "used_amount", + "usedAmount", + "consumed", + "consumed_amount", + "consumedAmount", + "current", + "current_usage", + "currentUsage", + "requests", + "request_count", + "requestCount", + "tokens", + "token_usage", + "tokenUsage", + "monthly_usage", + "monthlyUsage", + ] + private static let remainingKeys = [ + "remaining", + "available", + "balance", + "left", + "remaining_amount", + "remainingAmount", + "available_amount", + "availableAmount", + ] + private static let percentUsedKeys = [ + "percent_used", + "percentUsed", + "usage_percent", + "usagePercent", + "used_percent", + "usedPercent", + "utilization", + "utilization_percent", + "utilizationPercent", + ] + private static let percentRemainingKeys = [ + "percent_remaining", + "percentRemaining", + "remaining_percent", + "remainingPercent", + ] + private static let resetKeys = [ + "reset_at", + "resetAt", + "resets_at", + "resetsAt", + "reset_time", + "resetTime", + "next_reset_at", + "nextResetAt", + "renews_at", + "renewsAt", + "renewal_at", + "renewalAt", + "period_end", + "periodEnd", + "current_period_end", + "currentPeriodEnd", + "expires_at", + "expiresAt", + "window_end", + "windowEnd", + "end_time", + "endTime", + ] + private static let unitKeys = [ + "unit", + "units", + "currency", + "quota_unit", + "quotaUnit", + ] + private static let activeKeys = [ + "active", + "is_active", + "isActive", + "subscription_active", + "subscriptionActive", + "has_subscription", + "hasSubscription", + ] + private static let statusKeys = [ + "status", + "state", + "subscription_status", + "subscriptionStatus", + ] + private static let planKeys = [ + "plan_name", + "planName", + "plan", + "tier", + "subscription_plan", + "subscriptionPlan", + "subscription_tier", + "subscriptionTier", + ] + private static let windowMinuteKeys = [ + "window_minutes", + "windowMinutes", + "period_minutes", + "periodMinutes", + "duration_minutes", + "durationMinutes", + ] + private static let windowHourKeys = [ + "window_hours", + "windowHours", + "period_hours", + "periodHours", + "duration_hours", + "durationHours", + ] + private static let windowDayKeys = [ + "window_days", + "windowDays", + "period_days", + "periodDays", + "duration_days", + "durationDays", + ] + private static let windowSecondKeys = [ + "window_seconds", + "windowSeconds", + "period_seconds", + "periodSeconds", + "duration_seconds", + "durationSeconds", + ] + private static let windowStringKeys = [ + "window", + "period", + "interval", + "duration", + ] +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeAdminAPISettingsReader.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeAdminAPISettingsReader.swift new file mode 100644 index 0000000..cbb1791 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeAdminAPISettingsReader.swift @@ -0,0 +1,43 @@ +import Foundation + +public enum ClaudeAdminAPISettingsReader { + public static let adminAPIKeyEnvironmentKey = "ANTHROPIC_ADMIN_KEY" + public static let alternateAdminAPIKeyEnvironmentKey = "ANTHROPIC_ADMIN_API_KEY" + public static let apiKeyEnvironmentKeys = [ + Self.adminAPIKeyEnvironmentKey, + Self.alternateAdminAPIKeyEnvironmentKey, + ] + + public static func apiKey(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + for key in self.apiKeyEnvironmentKeys { + if let token = self.cleaned(environment[key]) { return token } + } + return nil + } + + public static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} + +public enum ClaudeAdminAPISettingsError: LocalizedError, Sendable { + case missingToken + + public var errorDescription: String? { + switch self { + case .missingToken: + "Claude API usage needs an Anthropic Admin API key." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeAdminAPIUsageFetcher.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeAdminAPIUsageFetcher.swift new file mode 100644 index 0000000..37df52a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeAdminAPIUsageFetcher.swift @@ -0,0 +1,432 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum ClaudeAdminAPIUsageError: LocalizedError, Sendable, Equatable { + case missingCredentials + case networkError(String) + case apiError(endpoint: String, statusCode: Int) + case parseFailed(endpoint: String, message: String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Anthropic Admin API key." + case let .networkError(message): + "Claude API usage network error: \(message)" + case let .apiError(endpoint, statusCode): + "Claude API usage \(endpoint) error: HTTP \(statusCode)" + case let .parseFailed(endpoint, message): + "Failed to parse Claude API usage \(endpoint): \(message)" + } + } +} + +public enum ClaudeAdminAPIUsageFetcher { + public static let costReportURL = URL(string: "https://api.anthropic.com/v1/organizations/cost_report")! + public static let messagesUsageURL = + URL(string: "https://api.anthropic.com/v1/organizations/usage_report/messages")! + + private static let anthropicVersion = "2023-06-01" + private static let timeoutSeconds: TimeInterval = 20 + private static let maxDailyBuckets = 31 + + public static func fetchUsage( + apiKey: String, + costURL: URL = Self.costReportURL, + messagesURL: URL = Self.messagesUsageURL, + session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + now: Date = Date()) async throws -> ClaudeAdminAPIUsageSnapshot + { + let trimmed = apiKey.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { + throw ClaudeAdminAPIUsageError.missingCredentials + } + + let calendar = Self.utcCalendar + let range = Self.dailyRange(now: now, calendar: calendar) + let costs = try await Self.fetchCostReport( + apiKey: trimmed, + baseURL: costURL, + range: range, + transport: transport) + let messages = try await Self.fetchMessagesUsage( + apiKey: trimmed, + baseURL: messagesURL, + range: range, + transport: transport) + + return Self.makeSnapshot(costs: costs, messages: messages, now: now, calendar: calendar) + } + + static func _parseSnapshotForTesting( + costs: Data, + messages: Data, + now: Date, + calendar: Calendar = Self.utcCalendar) throws -> ClaudeAdminAPIUsageSnapshot + { + let costs = try Self.decodeCosts(costs) + let messages = try Self.decodeMessages(messages) + return Self.makeSnapshot(costs: costs, messages: messages, now: now, calendar: calendar) + } + + private static func fetchCostReport( + apiKey: String, + baseURL: URL, + range: DateRange, + transport: any ProviderHTTPTransport) async throws -> CostReportResponse + { + let url = Self.url( + baseURL: baseURL, + range: range, + queryItems: [ + URLQueryItem(name: "group_by[]", value: "description"), + ]) + let data = try await Self.fetchData(url: url, apiKey: apiKey, endpoint: "cost_report", transport: transport) + return try Self.decodeCosts(data) + } + + private static func fetchMessagesUsage( + apiKey: String, + baseURL: URL, + range: DateRange, + transport: any ProviderHTTPTransport) async throws -> MessagesUsageResponse + { + let url = Self.url( + baseURL: baseURL, + range: range, + queryItems: [ + URLQueryItem(name: "group_by[]", value: "model"), + ]) + let data = try await Self.fetchData(url: url, apiKey: apiKey, endpoint: "messages", transport: transport) + return try Self.decodeMessages(data) + } + + private static func fetchData( + url: URL, + apiKey: String, + endpoint: String, + transport: any ProviderHTTPTransport) async throws -> Data + { + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.timeoutInterval = Self.timeoutSeconds + request.setValue(Self.anthropicVersion, forHTTPHeaderField: "anthropic-version") + request.setValue(apiKey, forHTTPHeaderField: "x-api-key") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("CodexBar/1.0", forHTTPHeaderField: "User-Agent") + + let response: ProviderHTTPResponse + do { + response = try await transport.response(for: request) + } catch { + throw ClaudeAdminAPIUsageError.networkError(error.localizedDescription) + } + + guard response.statusCode == 200 else { + throw ClaudeAdminAPIUsageError.apiError(endpoint: endpoint, statusCode: response.statusCode) + } + return response.data + } + + private static func decodeCosts(_ data: Data) throws -> CostReportResponse { + do { + return try JSONDecoder().decode(CostReportResponse.self, from: data) + } catch { + throw ClaudeAdminAPIUsageError.parseFailed(endpoint: "cost_report", message: error.localizedDescription) + } + } + + private static func decodeMessages(_ data: Data) throws -> MessagesUsageResponse { + do { + return try JSONDecoder().decode(MessagesUsageResponse.self, from: data) + } catch { + throw ClaudeAdminAPIUsageError.parseFailed(endpoint: "messages", message: error.localizedDescription) + } + } + + private static func makeSnapshot( + costs: CostReportResponse, + messages: MessagesUsageResponse, + now: Date, + calendar: Calendar) -> ClaudeAdminAPIUsageSnapshot + { + var accumulators: [String: DailyAccumulator] = [:] + + for bucket in costs.data { + var accumulator = accumulators[bucket.startingAt] ?? DailyAccumulator( + startingAt: bucket.startingAt, + endingAt: bucket.endingAt) + for result in bucket.results { + // Anthropic Usage & Cost API docs define `amount` as a decimal string in lowest USD units. + let value = Self.usdFromAnthropicLowestUnitAmount(result.amount) + accumulator.costUSD += value + let item = Self.displayName(result.description ?? result.costType, fallback: "Claude API") + accumulator.costItems[item, default: 0] += value + } + accumulators[bucket.startingAt] = accumulator + } + + for bucket in messages.data { + var accumulator = accumulators[bucket.startingAt] ?? DailyAccumulator( + startingAt: bucket.startingAt, + endingAt: bucket.endingAt) + for result in bucket.results { + let input = result.uncachedInputTokens ?? 0 + let cacheCreation = result.cacheCreation?.totalInputTokens ?? 0 + let cacheRead = result.cacheReadInputTokens ?? 0 + let output = result.outputTokens ?? 0 + let total = input + cacheCreation + cacheRead + output + accumulator.inputTokens += input + accumulator.cacheCreationInputTokens += cacheCreation + accumulator.cacheReadInputTokens += cacheRead + accumulator.outputTokens += output + accumulator.totalTokens += total + let modelName = Self.displayName(result.model, fallback: "Claude API") + accumulator.models[modelName, default: ModelAccumulator()].add( + inputTokens: input, + cacheCreationInputTokens: cacheCreation, + cacheReadInputTokens: cacheRead, + outputTokens: output, + totalTokens: total) + } + accumulators[bucket.startingAt] = accumulator + } + + let daily = accumulators.values + .compactMap { $0.makeBucket(calendar: calendar) } + .filter { $0.startTime <= now } + .sorted { $0.startTime < $1.startTime } + return ClaudeAdminAPIUsageSnapshot(daily: daily, updatedAt: now) + } + + private static func displayName(_ raw: String?, fallback: String) -> String { + guard let trimmed = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return fallback + } + return trimmed + } + + private static func usdFromAnthropicLowestUnitAmount(_ raw: String) -> Double { + (Double(raw) ?? 0) / 100 + } + + private static func url(baseURL: URL, range: DateRange, queryItems extraItems: [URLQueryItem]) -> URL { + var components = URLComponents(url: baseURL, resolvingAgainstBaseURL: false)! + components.queryItems = [ + URLQueryItem(name: "starting_at", value: Self.rfc3339String(from: range.start)), + URLQueryItem(name: "ending_at", value: Self.rfc3339String(from: range.end)), + URLQueryItem(name: "bucket_width", value: "1d"), + URLQueryItem(name: "limit", value: String(Self.maxDailyBuckets)), + ] + extraItems + return components.url! + } + + private static func dailyRange(now: Date, calendar: Calendar) -> DateRange { + let today = calendar.startOfDay(for: now) + let start = calendar.date(byAdding: .day, value: -(Self.maxDailyBuckets - 1), to: today) ?? today + let end = calendar.date(byAdding: .day, value: 1, to: today) ?? now + return DateRange(start: start, end: end) + } + + private static var utcCalendar: Calendar { + var calendar = Calendar(identifier: .gregorian) + calendar.locale = Locale(identifier: "en_US_POSIX") + calendar.timeZone = TimeZone(secondsFromGMT: 0)! + return calendar + } + + private static func rfc3339Formatter() -> ISO8601DateFormatter { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime] + formatter.timeZone = TimeZone(secondsFromGMT: 0) + return formatter + } + + private static func rfc3339String(from date: Date) -> String { + self.rfc3339Formatter().string(from: date) + } + + fileprivate static func dayKey(from date: Date, calendar: Calendar) -> String { + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = calendar.timeZone + formatter.dateFormat = "yyyy-MM-dd" + return formatter.string(from: date) + } + + fileprivate static func parseDate(_ raw: String) -> Date? { + self.rfc3339Formatter().date(from: raw) + } +} + +private struct DateRange { + let start: Date + let end: Date +} + +private struct DailyAccumulator { + let startingAt: String + let endingAt: String + var costUSD: Double = 0 + var inputTokens: Int = 0 + var cacheCreationInputTokens: Int = 0 + var cacheReadInputTokens: Int = 0 + var outputTokens: Int = 0 + var totalTokens: Int = 0 + var costItems: [String: Double] = [:] + var models: [String: ModelAccumulator] = [:] + + func makeBucket(calendar: Calendar) -> ClaudeAdminAPIUsageSnapshot.DailyBucket? { + guard let start = ClaudeAdminAPIUsageFetcher.parseDate(self.startingAt), + let end = ClaudeAdminAPIUsageFetcher.parseDate(self.endingAt) + else { return nil } + return ClaudeAdminAPIUsageSnapshot.DailyBucket( + day: ClaudeAdminAPIUsageFetcher.dayKey(from: start, calendar: calendar), + startTime: start, + endTime: end, + costUSD: self.costUSD, + inputTokens: self.inputTokens, + cacheCreationInputTokens: self.cacheCreationInputTokens, + cacheReadInputTokens: self.cacheReadInputTokens, + outputTokens: self.outputTokens, + totalTokens: self.totalTokens, + costItems: self.costItems + .map { ClaudeAdminAPIUsageSnapshot.CostBreakdown(name: $0.key, costUSD: $0.value) } + .sorted { + if $0.costUSD == $1.costUSD { return $0.name < $1.name } + return $0.costUSD > $1.costUSD + }, + models: self.models + .map { name, total in total.makeModel(name: name) } + .sorted { + if $0.totalTokens == $1.totalTokens { return $0.name < $1.name } + return $0.totalTokens > $1.totalTokens + }) + } +} + +private struct ModelAccumulator { + var inputTokens = 0 + var cacheCreationInputTokens = 0 + var cacheReadInputTokens = 0 + var outputTokens = 0 + var totalTokens = 0 + + mutating func add( + inputTokens: Int, + cacheCreationInputTokens: Int, + cacheReadInputTokens: Int, + outputTokens: Int, + totalTokens: Int) + { + self.inputTokens += inputTokens + self.cacheCreationInputTokens += cacheCreationInputTokens + self.cacheReadInputTokens += cacheReadInputTokens + self.outputTokens += outputTokens + self.totalTokens += totalTokens + } + + func makeModel(name: String) -> ClaudeAdminAPIUsageSnapshot.ModelBreakdown { + ClaudeAdminAPIUsageSnapshot.ModelBreakdown( + name: name, + inputTokens: self.inputTokens, + cacheCreationInputTokens: self.cacheCreationInputTokens, + cacheReadInputTokens: self.cacheReadInputTokens, + outputTokens: self.outputTokens, + totalTokens: self.totalTokens) + } +} + +private struct CostReportResponse: Decodable { + let data: [CostBucket] + let hasMore: Bool? + let nextPage: String? + + private enum CodingKeys: String, CodingKey { + case data + case hasMore = "has_more" + case nextPage = "next_page" + } +} + +private struct CostBucket: Decodable { + let startingAt: String + let endingAt: String + let results: [CostResult] + + private enum CodingKeys: String, CodingKey { + case startingAt = "starting_at" + case endingAt = "ending_at" + case results + } +} + +private struct CostResult: Decodable { + let currency: String? + let amount: String + let description: String? + let costType: String? + + private enum CodingKeys: String, CodingKey { + case currency + case amount + case description + case costType = "cost_type" + } +} + +private struct MessagesUsageResponse: Decodable { + let data: [MessagesBucket] + let hasMore: Bool? + let nextPage: String? + + private enum CodingKeys: String, CodingKey { + case data + case hasMore = "has_more" + case nextPage = "next_page" + } +} + +private struct MessagesBucket: Decodable { + let startingAt: String + let endingAt: String + let results: [MessagesResult] + + private enum CodingKeys: String, CodingKey { + case startingAt = "starting_at" + case endingAt = "ending_at" + case results + } +} + +private struct MessagesResult: Decodable { + let uncachedInputTokens: Int? + let cacheCreation: CacheCreation? + let cacheReadInputTokens: Int? + let outputTokens: Int? + let model: String? + + private enum CodingKeys: String, CodingKey { + case uncachedInputTokens = "uncached_input_tokens" + case cacheCreation = "cache_creation" + case cacheReadInputTokens = "cache_read_input_tokens" + case outputTokens = "output_tokens" + case model + } +} + +private struct CacheCreation: Decodable { + let ephemeral1HInputTokens: Int? + let ephemeral5MInputTokens: Int? + + var totalInputTokens: Int { + (self.ephemeral1HInputTokens ?? 0) + (self.ephemeral5MInputTokens ?? 0) + } + + private enum CodingKeys: String, CodingKey { + case ephemeral1HInputTokens = "ephemeral_1h_input_tokens" + case ephemeral5MInputTokens = "ephemeral_5m_input_tokens" + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeAdminAPIUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeAdminAPIUsageSnapshot.swift new file mode 100644 index 0000000..96d30a2 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeAdminAPIUsageSnapshot.swift @@ -0,0 +1,242 @@ +import Foundation + +public struct ClaudeAdminAPIUsageSnapshot: Codable, Equatable, Sendable { + public struct DailyBucket: Codable, Equatable, Sendable, Identifiable { + public let day: String + public let startTime: Date + public let endTime: Date + public let costUSD: Double + public let inputTokens: Int + public let cacheCreationInputTokens: Int + public let cacheReadInputTokens: Int + public let outputTokens: Int + public let totalTokens: Int + public let costItems: [CostBreakdown] + public let models: [ModelBreakdown] + + public var id: String { + self.day + } + + public init( + day: String, + startTime: Date, + endTime: Date, + costUSD: Double, + inputTokens: Int, + cacheCreationInputTokens: Int, + cacheReadInputTokens: Int, + outputTokens: Int, + totalTokens: Int, + costItems: [CostBreakdown], + models: [ModelBreakdown]) + { + self.day = day + self.startTime = startTime + self.endTime = endTime + self.costUSD = costUSD + self.inputTokens = inputTokens + self.cacheCreationInputTokens = cacheCreationInputTokens + self.cacheReadInputTokens = cacheReadInputTokens + self.outputTokens = outputTokens + self.totalTokens = totalTokens + self.costItems = costItems + self.models = models + } + } + + public struct CostBreakdown: Codable, Equatable, Sendable, Identifiable { + public let name: String + public let costUSD: Double + + public var id: String { + self.name + } + + public init(name: String, costUSD: Double) { + self.name = name + self.costUSD = costUSD + } + } + + public struct ModelBreakdown: Codable, Equatable, Sendable, Identifiable { + public let name: String + public let inputTokens: Int + public let cacheCreationInputTokens: Int + public let cacheReadInputTokens: Int + public let outputTokens: Int + public let totalTokens: Int + + public var id: String { + self.name + } + + public init( + name: String, + inputTokens: Int, + cacheCreationInputTokens: Int, + cacheReadInputTokens: Int, + outputTokens: Int, + totalTokens: Int) + { + self.name = name + self.inputTokens = inputTokens + self.cacheCreationInputTokens = cacheCreationInputTokens + self.cacheReadInputTokens = cacheReadInputTokens + self.outputTokens = outputTokens + self.totalTokens = totalTokens + } + } + + public struct Summary: Equatable, Sendable { + public let costUSD: Double + public let inputTokens: Int + public let cacheCreationInputTokens: Int + public let cacheReadInputTokens: Int + public let outputTokens: Int + public let totalTokens: Int + + public init( + costUSD: Double, + inputTokens: Int, + cacheCreationInputTokens: Int, + cacheReadInputTokens: Int, + outputTokens: Int, + totalTokens: Int) + { + self.costUSD = costUSD + self.inputTokens = inputTokens + self.cacheCreationInputTokens = cacheCreationInputTokens + self.cacheReadInputTokens = cacheReadInputTokens + self.outputTokens = outputTokens + self.totalTokens = totalTokens + } + } + + public let daily: [DailyBucket] + public let updatedAt: Date + + public init(daily: [DailyBucket], updatedAt: Date) { + self.daily = daily.sorted { $0.startTime < $1.startTime } + self.updatedAt = updatedAt + } + + public var last30Days: Summary { + self.summary(days: 30) + } + + public var last7Days: Summary { + self.summary(days: 7) + } + + public var currentDay: Summary { + self.summary(forLocalDayContaining: self.updatedAt) + } + + public var latestDay: Summary { + self.summary(days: 1) + } + + public func summary(forLocalDayContaining date: Date, calendar _: Calendar = .current) -> Summary { + let selected = self.daily.filter { bucket in + CostUsageBucketInterval.contains( + date, + startTime: bucket.startTime, + endTime: bucket.endTime) + } + return Summary( + costUSD: selected.reduce(0) { $0 + $1.costUSD }, + inputTokens: selected.reduce(0) { $0 + $1.inputTokens }, + cacheCreationInputTokens: selected.reduce(0) { $0 + $1.cacheCreationInputTokens }, + cacheReadInputTokens: selected.reduce(0) { $0 + $1.cacheReadInputTokens }, + outputTokens: selected.reduce(0) { $0 + $1.outputTokens }, + totalTokens: selected.reduce(0) { $0 + $1.totalTokens }) + } + + public func summary(days: Int) -> Summary { + let selected = self.daily.suffix(max(1, days)) + return Summary( + costUSD: selected.reduce(0) { $0 + $1.costUSD }, + inputTokens: selected.reduce(0) { $0 + $1.inputTokens }, + cacheCreationInputTokens: selected.reduce(0) { $0 + $1.cacheCreationInputTokens }, + cacheReadInputTokens: selected.reduce(0) { $0 + $1.cacheReadInputTokens }, + outputTokens: selected.reduce(0) { $0 + $1.outputTokens }, + totalTokens: selected.reduce(0) { $0 + $1.totalTokens }) + } + + public var topModels: [ModelBreakdown] { + var totals: [String: ModelAccumulator] = [:] + for day in self.daily { + for model in day.models { + totals[model.name, default: ModelAccumulator()].add(model) + } + } + return totals + .map { name, total in total.makeModel(name: name) } + .sorted { + if $0.totalTokens == $1.totalTokens { return $0.name < $1.name } + return $0.totalTokens > $1.totalTokens + } + } + + public var topCostItems: [CostBreakdown] { + var totals: [String: Double] = [:] + for day in self.daily { + for item in day.costItems { + totals[item.name, default: 0] += item.costUSD + } + } + return totals + .map { CostBreakdown(name: $0.key, costUSD: $0.value) } + .sorted { + if $0.costUSD == $1.costUSD { return $0.name < $1.name } + return $0.costUSD > $1.costUSD + } + } + + public func toUsageSnapshot() -> UsageSnapshot { + let total = self.last30Days + return UsageSnapshot( + primary: nil, + secondary: nil, + providerCost: ProviderCostSnapshot( + used: total.costUSD, + limit: 0, + currencyCode: "USD", + period: "Last 30 days", + updatedAt: self.updatedAt), + claudeAdminAPIUsage: self, + updatedAt: self.updatedAt, + identity: ProviderIdentitySnapshot( + providerID: .claude, + accountEmail: nil, + accountOrganization: nil, + loginMethod: "Admin API")) + } + + private struct ModelAccumulator { + var inputTokens = 0 + var cacheCreationInputTokens = 0 + var cacheReadInputTokens = 0 + var outputTokens = 0 + var totalTokens = 0 + + mutating func add(_ model: ModelBreakdown) { + self.inputTokens += model.inputTokens + self.cacheCreationInputTokens += model.cacheCreationInputTokens + self.cacheReadInputTokens += model.cacheReadInputTokens + self.outputTokens += model.outputTokens + self.totalTokens += model.totalTokens + } + + func makeModel(name: String) -> ModelBreakdown { + ModelBreakdown( + name: name, + inputTokens: self.inputTokens, + cacheCreationInputTokens: self.cacheCreationInputTokens, + cacheReadInputTokens: self.cacheReadInputTokens, + outputTokens: self.outputTokens, + totalTokens: self.totalTokens) + } + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeCLIAuthStatusProbe.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeCLIAuthStatusProbe.swift new file mode 100644 index 0000000..4f3b6f5 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeCLIAuthStatusProbe.swift @@ -0,0 +1,35 @@ +import Foundation + +enum ClaudeCLIAuthStatusProbe { + private struct Response: Decodable { + let loggedIn: Bool + } + + static func isLoggedIn( + binary: String, + environment: [String: String], + timeout: TimeInterval = 5) async -> Bool + { + do { + let result = try await SubprocessRunner.run( + binary: binary, + arguments: ["auth", "status", "--json"], + environment: ClaudeCLISession.launchEnvironment(baseEnv: environment), + timeout: timeout, + standardInput: FileHandle.nullDevice, + label: "claude-auth-status") + return self.parseLoggedIn(result.stdout) + } catch { + return false + } + } + + static func parseLoggedIn(_ output: String) -> Bool { + guard let data = output.data(using: .utf8), + let response = try? JSONDecoder().decode(Response.self, from: data) + else { + return false + } + return response.loggedIn + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeCLIRateLimitGate.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeCLIRateLimitGate.swift new file mode 100644 index 0000000..5c06ec0 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeCLIRateLimitGate.swift @@ -0,0 +1,64 @@ +import Foundation + +enum ClaudeCLIRateLimitGate { + private static let blockedUntilKey = "claudeCLIUsageRateLimitBlockedUntilV1" + private static let defaultCooldown: TimeInterval = 60 * 5 + + static let message = "Claude CLI usage endpoint is rate limited right now. Please try again later." + + static func blockedUntil( + interaction: ProviderInteraction = ProviderInteractionContext.current, + now: Date = Date()) -> Date? + { + guard interaction != .userInitiated else { return nil } + return self.currentBlockedUntil(now: now) + } + + static func currentBlockedUntil(now: Date = Date()) -> Date? { + guard let raw = UserDefaults.standard.object(forKey: self.blockedUntilKey) as? Double else { + return nil + } + + let blockedUntil = Date(timeIntervalSince1970: raw) + guard blockedUntil > now else { + UserDefaults.standard.removeObject(forKey: self.blockedUntilKey) + return nil + } + return blockedUntil + } + + static func recordRateLimit(now: Date = Date()) { + UserDefaults.standard.set( + now.addingTimeInterval(self.defaultCooldown).timeIntervalSince1970, + forKey: self.blockedUntilKey) + } + + static func recordSuccess() { + UserDefaults.standard.removeObject(forKey: self.blockedUntilKey) + } + + static func isRateLimitError(_ error: Error) -> Bool { + if case let ClaudeStatusProbeError.parseFailed(message) = error { + return self.isRateLimitMessage(message, allowRawRateLimitToken: true) + } + if case let ClaudeUsageError.parseFailed(message) = error { + return self.isRateLimitMessage(message, allowRawRateLimitToken: true) + } + return self.isRateLimitMessage(error.localizedDescription, allowRawRateLimitToken: false) + } + + private static func isRateLimitMessage(_ message: String, allowRawRateLimitToken: Bool) -> Bool { + let lower = message.lowercased() + return lower.contains(Self.message.lowercased()) || + (allowRawRateLimitToken && lower.contains("rate_limit_error")) || + (lower.contains("claude cli") && + lower.contains("usage") && + lower.contains("rate limited")) + } + + #if DEBUG + static func resetForTesting() { + UserDefaults.standard.removeObject(forKey: self.blockedUntilKey) + } + #endif +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeCLISession.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeCLISession.swift new file mode 100644 index 0000000..0866b35 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeCLISession.swift @@ -0,0 +1,499 @@ +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif +import Foundation + +actor ClaudeCLISession { + static let shared = ClaudeCLISession() + private static let log = CodexBarLog.logger(LogCategories.claudeCLI) + #if DEBUG + @TaskLocal private static var sessionOverrideForTesting: ClaudeCLISession? + + static var current: ClaudeCLISession { + self.sessionOverrideForTesting ?? self.shared + } + + static func withIsolatedSessionForTesting(operation: () async throws -> T) async rethrows -> T { + let session = ClaudeCLISession() + defer { Task { await session.reset() } } + return try await self.$sessionOverrideForTesting.withValue(session) { + try await operation() + } + } + #else + static var current: ClaudeCLISession { + self.shared + } + #endif + + enum SessionError: LocalizedError { + case launchFailed(String) + case ioFailed(String) + case timedOut + case processExited + + var errorDescription: String? { + switch self { + case let .launchFailed(msg): "Failed to launch Claude CLI session: \(msg)" + case let .ioFailed(msg): "Claude CLI PTY I/O failed: \(msg)" + case .timedOut: "Claude CLI session timed out." + case .processExited: "Claude CLI session exited." + } + } + } + + private var process: Process? + private var primaryFD: Int32 = -1 + private var primaryHandle: FileHandle? + private var secondaryHandle: FileHandle? + private var processGroup: pid_t? + private var binaryPath: String? + private var startedAt: Date? + + private let promptSends: [String: String] = [ + "Do you trust the files in this folder?": "y\r", + "Quick safety check:": "\r", + "Yes, I trust this folder": "\r", + "Ready to code here?": "\r", + "Press Enter to continue": "\r", + ] + + private struct RollingBuffer { + private let maxNeedle: Int + private var tail = Data() + + init(maxNeedle: Int) { + self.maxNeedle = max(0, maxNeedle) + } + + mutating func append(_ data: Data) -> Data { + guard !data.isEmpty else { return Data() } + var combined = Data() + combined.reserveCapacity(self.tail.count + data.count) + combined.append(self.tail) + combined.append(data) + if self.maxNeedle > 1 { + if combined.count >= self.maxNeedle - 1 { + self.tail = combined.suffix(self.maxNeedle - 1) + } else { + self.tail = combined + } + } else { + self.tail.removeAll(keepingCapacity: true) + } + return combined + } + } + + private static func normalizedNeedle(_ text: String) -> String { + String(text.lowercased().filter { !$0.isWhitespace }) + } + + private static func commandPaletteSends(for subcommand: String) -> [String: String] { + let normalized = subcommand.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + switch normalized { + case "/usage": + // Claude's command palette can render several "Show ..." actions together; only auto-confirm the + // usage-related actions here so we do not accidentally execute /status. + return [ + "Show plan": "\r", + "Show plan usage limits": "\r", + ] + case "/status": + return [ + "Show Claude Code": "\r", + "Show Claude Code status": "\r", + ] + default: + return [:] + } + } + + func capture( + subcommand: String, + binary: String, + timeout: TimeInterval, + idleTimeout: TimeInterval? = 3.0, + stopOnSubstrings: [String] = [], + stopWhenNormalized: (@Sendable (String) -> Bool)? = nil, + settleAfterStop: TimeInterval = 0.25, + sendEnterEvery: TimeInterval? = nil) async throws -> String + { + try self.ensureStarted(binary: binary) + if let startedAt { + let sinceStart = Date().timeIntervalSince(startedAt) + // Claude's TUI can drop early keystrokes while it's still initializing. Wait a bit longer than the + // original 0.4s to ensure slash commands reliably open their panels. + if sinceStart < 2.0 { + let delay = UInt64((2.0 - sinceStart) * 1_000_000_000) + try await Task.sleep(nanoseconds: delay) + } + } + self.drainOutput() + + let trimmed = subcommand.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { + try self.send(trimmed) + try self.send("\r") + } + + let stopNeedles = stopOnSubstrings.map { Self.normalizedNeedle($0) } + var sendMap = self.promptSends + for (needle, keys) in Self.commandPaletteSends(for: trimmed) { + sendMap[needle] = keys + } + let sendNeedles = sendMap.map { (needle: Self.normalizedNeedle($0.key), keys: $0.value) } + let cursorQuery = Data([0x1B, 0x5B, 0x36, 0x6E]) + let needleLengths = + stopOnSubstrings.map(\.utf8.count) + + sendMap.keys.map(\.utf8.count) + + [cursorQuery.count] + let maxNeedle = needleLengths.max() ?? cursorQuery.count + var scanBuffer = RollingBuffer(maxNeedle: maxNeedle) + var triggeredSends = Set() + + var buffer = Data() + var scanTailText = "" + var normalizedScan = "" + var utf8Carry = Data() + let deadline = Date().addingTimeInterval(timeout) + var lastOutputAt = Date() + var lastEnterAt = Date() + var stoppedEarly = false + // Only send periodic Enter when the caller explicitly asks for it (used for /usage rendering). + // For /status, periodic input can keep producing output and prevent idle-timeout short-circuiting. + let effectiveEnterEvery: TimeInterval? = sendEnterEvery + + while Date() < deadline { + let newData = self.readChunk() + if !newData.isEmpty { + buffer.append(newData) + lastOutputAt = Date() + Self.appendScanText(newData: newData, scanTailText: &scanTailText, utf8Carry: &utf8Carry) + if scanTailText.count > 8192 { + scanTailText = String(scanTailText.suffix(8192)) + } + normalizedScan = Self.normalizedNeedle(TextParsing.stripANSICodes(scanTailText)) + + let scanData = scanBuffer.append(newData) + if scanData.range(of: cursorQuery) != nil { + try? self.send("\u{1b}[1;1R") + } + + for item in sendNeedles where !triggeredSends.contains(item.needle) { + if normalizedScan.contains(item.needle) { + try? self.send(item.keys) + triggeredSends.insert(item.needle) + } + } + + if stopNeedles + .contains(where: normalizedScan.contains) || (stopWhenNormalized?(normalizedScan) == true) + { + stoppedEarly = true + break + } + } + + if self.shouldStopForIdleTimeout( + idleTimeout: idleTimeout, + bufferIsEmpty: buffer.isEmpty, + lastOutputAt: lastOutputAt) + { + stoppedEarly = true + break + } + + self.sendPeriodicEnterIfNeeded(every: effectiveEnterEvery, lastEnterAt: &lastEnterAt) + + if let proc = self.process, !proc.isRunning { + throw SessionError.processExited + } + + try await Task.sleep(nanoseconds: 60_000_000) + } + + if stoppedEarly { + let settle = max(0, min(settleAfterStop, deadline.timeIntervalSinceNow)) + if settle > 0 { + let settleDeadline = Date().addingTimeInterval(settle) + while Date() < settleDeadline { + let newData = self.readChunk() + if !newData.isEmpty { + buffer.append(newData) + } + try await Task.sleep(nanoseconds: 50_000_000) + } + } + } + + guard !buffer.isEmpty, let text = String(data: buffer, encoding: .utf8) else { + throw SessionError.timedOut + } + return text + } + + private static func appendScanText(newData: Data, scanTailText: inout String, utf8Carry: inout Data) { + // PTY reads can split multibyte UTF-8 sequences. Keep a small carry buffer so prompt/stop scanning doesn't + // drop chunks when the decode fails due to an incomplete trailing sequence. + var combined = Data() + combined.reserveCapacity(utf8Carry.count + newData.count) + combined.append(utf8Carry) + combined.append(newData) + + if let chunk = String(data: combined, encoding: .utf8) { + scanTailText.append(chunk) + utf8Carry.removeAll(keepingCapacity: true) + return + } + + for trimCount in 1...3 where combined.count > trimCount { + let prefix = combined.dropLast(trimCount) + if let chunk = String(data: prefix, encoding: .utf8) { + scanTailText.append(chunk) + utf8Carry = Data(combined.suffix(trimCount)) + return + } + } + + // If the data is still not UTF-8 decodable, keep only a small suffix to avoid unbounded growth. + utf8Carry = Data(combined.suffix(12)) + } + + func reset() { + self.cleanup() + } + + private func ensureStarted(binary: String) throws { + if let proc = self.process, proc.isRunning, self.binaryPath == binary { + Self.log.debug("Claude CLI session reused") + return + } + self.cleanup() + + var primaryFD: Int32 = -1 + var secondaryFD: Int32 = -1 + var win = winsize(ws_row: 50, ws_col: 160, ws_xpixel: 0, ws_ypixel: 0) + guard openpty(&primaryFD, &secondaryFD, nil, nil, &win) == 0 else { + Self.log.warning("Claude CLI PTY openpty failed") + throw SessionError.launchFailed("openpty failed") + } + _ = fcntl(primaryFD, F_SETFL, O_NONBLOCK) + + let primaryHandle = FileHandle(fileDescriptor: primaryFD, closeOnDealloc: true) + let secondaryHandle = FileHandle(fileDescriptor: secondaryFD, closeOnDealloc: true) + + let proc = Process() + let resolvedURL = URL(fileURLWithPath: binary) + let disableWatchdog = ProcessInfo.processInfo.environment["CODEXBAR_DISABLE_CLAUDE_WATCHDOG"] == "1" + if !disableWatchdog, + resolvedURL.lastPathComponent == "claude", + let watchdog = TTYCommandRunner.locateBundledHelper("CodexBarClaudeWatchdog") + { + proc.executableURL = URL(fileURLWithPath: watchdog) + proc.arguments = ["--", binary, "--allowed-tools", ""] + } else { + proc.executableURL = resolvedURL + proc.arguments = ["--allowed-tools", ""] + } + proc.standardInput = secondaryHandle + proc.standardOutput = secondaryHandle + proc.standardError = secondaryHandle + + let workingDirectory = ClaudeStatusProbe.preparedProbeWorkingDirectoryURL() + proc.currentDirectoryURL = workingDirectory + var env = Self.launchEnvironment() + env["PWD"] = workingDirectory.path + proc.environment = env + + guard TTYCommandRunner.beginActiveProcessLaunchForAppShutdown() else { + try? primaryHandle.close() + try? secondaryHandle.close() + throw SessionError.launchFailed("App shutdown in progress") + } + defer { TTYCommandRunner.endActiveProcessLaunchForAppShutdown() } + + do { + try proc.run() + Self.log.debug( + "Claude CLI session started", + metadata: ["binary": URL(fileURLWithPath: binary).lastPathComponent]) + } catch { + Self.log.warning("Claude CLI launch failed", metadata: ["error": error.localizedDescription]) + try? primaryHandle.close() + try? secondaryHandle.close() + throw SessionError.launchFailed(error.localizedDescription) + } + + let pid = proc.processIdentifier + guard TTYCommandRunner.registerActiveProcessForAppShutdown( + pid: pid, + binary: URL(fileURLWithPath: binary).lastPathComponent) + else { + proc.terminate() + kill(pid, SIGKILL) + try? primaryHandle.close() + try? secondaryHandle.close() + throw SessionError.launchFailed("App shutdown in progress") + } + + var processGroup: pid_t? + if setpgid(pid, pid) == 0 { + processGroup = pid + TTYCommandRunner.updateActiveProcessGroupForAppShutdown(pid: pid, processGroup: processGroup) + } + + self.process = proc + self.primaryFD = primaryFD + self.primaryHandle = primaryHandle + self.secondaryHandle = secondaryHandle + self.processGroup = processGroup + self.binaryPath = binary + self.startedAt = Date() + } + + static func launchEnvironment(baseEnv: [String: String] = ProcessInfo.processInfo.environment) -> [String: String] { + var env = self.scrubbedClaudeEnvironment(from: TTYCommandRunner.enrichedEnvironment(baseEnv: baseEnv)) + // Passive status and auth probes must not mutate or update the user's Claude CLI installation. + env["DISABLE_AUTOUPDATER"] = "1" + return env + } + + private static func scrubbedClaudeEnvironment(from base: [String: String]) -> [String: String] { + var env = base + let explicitKeys: [String] = [ + ClaudeOAuthCredentialsStore.environmentTokenKey, + ClaudeOAuthCredentialsStore.environmentScopesKey, + ] + for key in explicitKeys { + env.removeValue(forKey: key) + } + for key in env.keys where key.hasPrefix("ANTHROPIC_") { + env.removeValue(forKey: key) + } + return env + } + + private func cleanup() { + if self.process != nil { + Self.log.debug("Claude CLI session stopping") + } + if let proc = self.process, proc.isRunning { + try? self.writeAllToPrimary(Data("/exit\r".utf8)) + } + try? self.primaryHandle?.close() + try? self.secondaryHandle?.close() + + let descendants = self.process.map { TTYProcessTreeTerminator.descendantPIDs(of: $0.processIdentifier) } ?? [] + if let proc = self.process, proc.isRunning { + proc.terminate() + } + if let proc = self.process { + TTYProcessTreeTerminator.terminateProcessTree( + rootPID: proc.processIdentifier, + processGroup: self.processGroup, + signal: SIGTERM, + knownDescendants: descendants) + } + let waitDeadline = Date().addingTimeInterval(1.0) + if let proc = self.process { + while proc.isRunning, Date() < waitDeadline { + usleep(100_000) + } + if proc.isRunning { + TTYProcessTreeTerminator.terminateProcessTree( + rootPID: proc.processIdentifier, + processGroup: self.processGroup, + signal: SIGKILL, + knownDescendants: descendants) + } else { + for pid in descendants where pid > 0 { + kill(pid, SIGKILL) + } + } + TTYCommandRunner.unregisterActiveProcessForAppShutdown(pid: proc.processIdentifier) + } + + self.process = nil + self.primaryHandle = nil + self.secondaryHandle = nil + self.primaryFD = -1 + self.processGroup = nil + self.startedAt = nil + } + + private func readChunk() -> Data { + guard self.primaryFD >= 0 else { return Data() } + var appended = Data() + while true { + var tmp = [UInt8](repeating: 0, count: 8192) + let n = read(self.primaryFD, &tmp, tmp.count) + if n > 0 { + appended.append(contentsOf: tmp.prefix(n)) + continue + } + break + } + return appended + } + + private func drainOutput() { + _ = self.readChunk() + } + + private func shouldStopForIdleTimeout( + idleTimeout: TimeInterval?, + bufferIsEmpty: Bool, + lastOutputAt: Date) -> Bool + { + guard let idleTimeout, !bufferIsEmpty else { return false } + return Date().timeIntervalSince(lastOutputAt) >= idleTimeout + } + + private func sendPeriodicEnterIfNeeded(every: TimeInterval?, lastEnterAt: inout Date) { + guard let every, Date().timeIntervalSince(lastEnterAt) >= every else { return } + try? self.send("\r") + lastEnterAt = Date() + } + + private func send(_ text: String) throws { + guard let data = text.data(using: .utf8) else { return } + guard self.primaryFD >= 0 else { throw SessionError.processExited } + try self.writeAllToPrimary(data) + } + + private func writeAllToPrimary(_ data: Data) throws { + guard self.primaryFD >= 0 else { throw SessionError.processExited } + try data.withUnsafeBytes { rawBytes in + guard let baseAddress = rawBytes.baseAddress else { return } + var offset = 0 + var retries = 0 + while offset < rawBytes.count { + let written = write(self.primaryFD, baseAddress.advanced(by: offset), rawBytes.count - offset) + if written > 0 { + offset += written + retries = 0 + continue + } + if written == 0 { + break + } + + let err = errno + if err == EINTR || err == EAGAIN || err == EWOULDBLOCK { + retries += 1 + if retries > 200 { + throw SessionError.ioFailed("write to PTY would block") + } + usleep(5000) + continue + } + throw SessionError.ioFailed("write to PTY failed: \(String(cString: strerror(err)))") + } + } + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeCredentialRouting.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeCredentialRouting.swift new file mode 100644 index 0000000..2472de2 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeCredentialRouting.swift @@ -0,0 +1,102 @@ +import Foundation + +public enum ClaudeCredentialRouting: Sendable, Equatable { + case none + case oauth(accessToken: String) + case webCookie(header: String) + case adminAPIKey(String) + + public static func resolve(tokenAccountToken: String?, manualCookieHeader: String?) -> Self { + if let tokenAccountToken, + let route = self.resolvePrimaryCredential(tokenAccountToken) + { + return route + } + + guard let manualCookieHeader = self.normalizedWebCookie(manualCookieHeader) else { + return .none + } + return .webCookie(header: manualCookieHeader) + } + + public var oauthAccessToken: String? { + guard case let .oauth(accessToken) = self else { return nil } + return accessToken + } + + public var manualCookieHeader: String? { + guard case let .webCookie(header) = self else { return nil } + return header + } + + public var isOAuth: Bool { + if case .oauth = self { return true } + return false + } + + public var adminAPIKey: String? { + guard case let .adminAPIKey(key) = self else { return nil } + return key + } + + private static func resolvePrimaryCredential(_ raw: String) -> Self? { + if let adminKey = self.normalizedAdminAPIKey(raw) { + return .adminAPIKey(adminKey) + } + if let accessToken = self.normalizedOAuthToken(raw) { + return .oauth(accessToken: accessToken) + } + if let cookieHeader = self.normalizedWebCookie(raw) { + return .webCookie(header: cookieHeader) + } + return nil + } + + private static func normalizedOAuthToken(_ raw: String?) -> String? { + guard let trimmed = self.cleaned(raw) else { return nil } + let lower = trimmed.lowercased() + if lower.contains("cookie:") || trimmed.contains("=") { + return nil + } + if lower.hasPrefix("bearer ") { + let bearerTrimmed = trimmed.dropFirst("bearer ".count) + .trimmingCharacters(in: .whitespacesAndNewlines) + guard !bearerTrimmed.isEmpty else { return nil } + let lowerBearerTrimmed = bearerTrimmed.lowercased() + return lowerBearerTrimmed.hasPrefix("sk-ant-oat") ? bearerTrimmed : nil + } + return lower.hasPrefix("sk-ant-oat") ? trimmed : nil + } + + private static func normalizedAdminAPIKey(_ raw: String?) -> String? { + guard let trimmed = self.cleaned(raw) else { return nil } + let lower = trimmed.lowercased() + if lower.contains("cookie:") || trimmed.contains("=") { + return nil + } + if lower.hasPrefix("bearer ") { + let bearerTrimmed = trimmed.dropFirst("bearer ".count) + .trimmingCharacters(in: .whitespacesAndNewlines) + guard !bearerTrimmed.isEmpty else { return nil } + return bearerTrimmed.lowercased().hasPrefix("sk-ant-admin") ? bearerTrimmed : nil + } + return lower.hasPrefix("sk-ant-admin") ? trimmed : nil + } + + private static func normalizedWebCookie(_ raw: String?) -> String? { + guard let normalized = CookieHeaderNormalizer.normalize(raw) else { return nil } + if normalized.contains("=") { + return normalized + } + return "sessionKey=\(normalized)" + } + + private static func cleaned(_ raw: String?) -> String? { + guard let trimmed = raw?.trimmingCharacters(in: .whitespacesAndNewlines), + !trimmed.isEmpty + else { + return nil + } + return trimmed + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeDesktopProjectsLocator.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeDesktopProjectsLocator.swift new file mode 100644 index 0000000..9b886e2 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeDesktopProjectsLocator.swift @@ -0,0 +1,89 @@ +import Foundation + +public enum ClaudeDesktopProjectsLocator { + private static let sessionDirectoryNames = [ + "local-agent-mode-sessions", + "claude-code-sessions", + ] + + private static let skippedDirectoryNames = Set([ + ".build", + ".git", + "build", + "DerivedData", + "node_modules", + "outputs", + "target", + ]) + + public static func roots( + homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser, + fileManager: FileManager = .default) -> [URL] + { + let claudeApplicationSupport = homeDirectory + .appendingPathComponent("Library", isDirectory: true) + .appendingPathComponent("Application Support", isDirectory: true) + .appendingPathComponent("Claude", isDirectory: true) + let sessionRoots = self.sessionDirectoryNames.map { + claudeApplicationSupport.appendingPathComponent($0, isDirectory: true) + } + + var roots: [URL] = [] + var queue = sessionRoots.map { (url: $0, depth: 0) } + var visited = Set(sessionRoots.map(\.standardizedFileURL.path)) + var nextIndex = 0 + // Current Desktop entries under claude-code-sessions are metadata whose cliSessionId maps to the + // shared ~/.claude/projects logs. Seed that root anyway so embedded project stores are found when present. + // Covers observed Desktop layouts through account/workspace, session, agent, and local_agent + // without descending into arbitrary checked-out workspaces. + let maxDepth = 4 + + while nextIndex < queue.count { + let current = queue[nextIndex] + nextIndex += 1 + if let projects = self.projectsRoot(under: current.url, fileManager: fileManager) { + roots.append(projects) + } + + guard current.depth < maxDepth else { continue } + for child in self.childDirectories(at: current.url, fileManager: fileManager) { + let standardized = child.standardizedFileURL + guard visited.insert(standardized.path).inserted else { continue } + queue.append((standardized, current.depth + 1)) + } + } + return roots + } + + private static func projectsRoot(under base: URL, fileManager: FileManager) -> URL? { + let projects = base + .appendingPathComponent(".claude", isDirectory: true) + .appendingPathComponent("projects", isDirectory: true) + var isDirectory: ObjCBool = false + guard fileManager.fileExists(atPath: projects.path, isDirectory: &isDirectory), + isDirectory.boolValue + else { return nil } + return projects.standardizedFileURL + } + + private static func childDirectories(at url: URL, fileManager: FileManager) -> [URL] { + guard let children = try? fileManager.contentsOfDirectory( + at: url, + includingPropertiesForKeys: [.isDirectoryKey, .isSymbolicLinkKey], + options: [.skipsHiddenFiles, .skipsPackageDescendants]) + else { + return [] + } + + return children.compactMap { child in + guard !self.skippedDirectoryNames.contains(child.lastPathComponent) else { return nil } + guard let values = try? child.resourceValues(forKeys: [.isDirectoryKey, .isSymbolicLinkKey]), + values.isSymbolicLink != true, + values.isDirectory == true + else { + return nil + } + return child + } + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentialModels.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentialModels.swift new file mode 100644 index 0000000..1b3401c --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentialModels.swift @@ -0,0 +1,301 @@ +import Foundation + +#if canImport(CryptoKit) +import CryptoKit +#else +import Crypto +#endif + +#if os(macOS) +import Security +#endif + +public struct ClaudeOAuthCredentials: Sendable { + public let accessToken: String + public let refreshToken: String? + public let expiresAt: Date? + public let scopes: [String] + public let rateLimitTier: String? + public let subscriptionType: String? + + public init( + accessToken: String, + refreshToken: String?, + expiresAt: Date?, + scopes: [String], + rateLimitTier: String?, + subscriptionType: String? = nil) + { + self.accessToken = accessToken + self.refreshToken = refreshToken + self.expiresAt = expiresAt + self.scopes = scopes + self.rateLimitTier = rateLimitTier + self.subscriptionType = subscriptionType + } + + public var isExpired: Bool { + guard let expiresAt else { return true } + return Date() >= expiresAt + } + + public var expiresIn: TimeInterval? { + guard let expiresAt else { return nil } + return expiresAt.timeIntervalSinceNow + } + + /// A one-way discriminator for history owned by this credential. + /// + /// Prefer the refresh token because access tokens routinely rotate for the same principal. If a provider + /// supplies only an access token, rotating that token intentionally starts a new history bucket rather than + /// risking that two identityless accounts share one. The source secret never leaves this computation. + var historyOwnerIdentifier: String? { + let normalizedRefreshToken = self.refreshToken?.trimmingCharacters(in: .whitespacesAndNewlines) + if let normalizedRefreshToken, !normalizedRefreshToken.isEmpty { + return Self.makeHistoryOwnerIdentifier(secretKind: "refresh", secret: normalizedRefreshToken) + } + + let normalizedAccessToken = self.accessToken.trimmingCharacters(in: .whitespacesAndNewlines) + guard !normalizedAccessToken.isEmpty else { return nil } + return Self.makeHistoryOwnerIdentifier(secretKind: "access", secret: normalizedAccessToken) + } + + static func historyOwnerIdentifier(forRefreshToken refreshToken: String) -> String? { + let normalized = refreshToken.trimmingCharacters(in: .whitespacesAndNewlines) + guard !normalized.isEmpty else { return nil } + return Self.makeHistoryOwnerIdentifier(secretKind: "refresh", secret: normalized) + } + + private static func makeHistoryOwnerIdentifier(secretKind: String, secret: String) -> String? { + let material = Data("codexbar:claude-oauth-history-owner:v1\0\(secretKind)\0\(secret)".utf8) + return SHA256.hash(data: material).map { String(format: "%02x", $0) }.joined() + } + + static func normalizedHistoryOwnerIdentifier(_ identifier: String?) -> String? { + guard let identifier else { return nil } + let normalized = identifier.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + guard normalized.count == 64, + normalized.unicodeScalars.allSatisfy({ scalar in + switch scalar.value { + case 48...57, 97...102: + true + default: + false + } + }) + else { return nil } + return normalized + } + + public static func isMcpOAuthOnlyPayload(data: Data) -> Bool { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + return false + } + return json["claudeAiOauth"] == nil && json["mcpOAuth"] != nil + } + + public static func parse(data: Data) throws -> ClaudeOAuthCredentials { + if ClaudeOAuthCredentials.isMcpOAuthOnlyPayload(data: data) { + throw ClaudeOAuthCredentialsError.mcpOAuthOnlyKeychain + } + + let decoder = JSONDecoder() + guard let root = try? decoder.decode(Root.self, from: data) else { + throw ClaudeOAuthCredentialsError.decodeFailed + } + guard let oauth = root.claudeAiOauth else { + throw ClaudeOAuthCredentialsError.missingOAuth + } + let accessToken = oauth.accessToken?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + guard !accessToken.isEmpty else { + throw ClaudeOAuthCredentialsError.missingAccessToken + } + let expiresAt = oauth.expiresAt.map { millis in + Date(timeIntervalSince1970: millis / 1000.0) + } + return ClaudeOAuthCredentials( + accessToken: accessToken, + refreshToken: oauth.refreshToken, + expiresAt: expiresAt, + scopes: oauth.scopes ?? [], + rateLimitTier: oauth.rateLimitTier, + subscriptionType: oauth.subscriptionType) + } + + private struct Root: Decodable { + let claudeAiOauth: OAuth? + } + + private struct OAuth: Decodable { + let accessToken: String? + let refreshToken: String? + let expiresAt: Double? + let scopes: [String]? + let rateLimitTier: String? + let subscriptionType: String? + + enum CodingKeys: String, CodingKey { + case accessToken + case refreshToken + case expiresAt + case scopes + case rateLimitTier + case subscriptionType + } + } +} + +extension ClaudeOAuthCredentials { + func diagnosticsMetadata(now: Date = Date()) -> [String: String] { + let hasRefreshToken = !(self.refreshToken?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty ?? true) + let hasUserProfileScope = self.scopes.contains("user:profile") + + var metadata: [String: String] = [ + "hasRefreshToken": "\(hasRefreshToken)", + "scopesCount": "\(self.scopes.count)", + "hasUserProfileScope": "\(hasUserProfileScope)", + ] + + if let expiresAt = self.expiresAt { + let expiresAtMs = Int(expiresAt.timeIntervalSince1970 * 1000.0) + let expiresInSec = Int(expiresAt.timeIntervalSince(now).rounded()) + metadata["expiresAtMs"] = "\(expiresAtMs)" + metadata["expiresInSec"] = "\(expiresInSec)" + metadata["isExpired"] = "\(now >= expiresAt)" + } else { + metadata["expiresAtMs"] = "nil" + metadata["expiresInSec"] = "nil" + metadata["isExpired"] = "true" + } + + return metadata + } +} + +public enum ClaudeOAuthCredentialOwner: String, Codable, Sendable { + case claudeCLI + case codexbar + case environment +} + +public enum ClaudeOAuthCredentialSource: String, Sendable { + case environment + case memoryCache + case cacheKeychain + case credentialsFile + case claudeKeychain +} + +enum ClaudeKeychainCredentialMatch: Equatable, Sendable { + case notApplicable + case absent + case unavailable + case mismatch + case matched(persistentRefHash: String) + + var persistentRefHash: String? { + guard case let .matched(persistentRefHash) = self else { return nil } + return persistentRefHash + } + + var isMismatch: Bool { + self == .mismatch + } + + var isAbsent: Bool { + self == .absent + } + + var isUnavailable: Bool { + self == .unavailable + } +} + +public struct ClaudeOAuthCredentialRecord: Sendable { + public let credentials: ClaudeOAuthCredentials + public let owner: ClaudeOAuthCredentialOwner + public let source: ClaudeOAuthCredentialSource + private let inheritedHistoryOwnerIdentifier: String? + + /// An opaque, one-way owner identifier that survives a refresh-token rotation proven by a successful refresh. + /// Records from unrelated credential sources do not inherit this value and derive a fresh identifier instead. + var historyOwnerIdentifier: String? { + self.inheritedHistoryOwnerIdentifier ?? self.credentials.historyOwnerIdentifier + } + + public init( + credentials: ClaudeOAuthCredentials, + owner: ClaudeOAuthCredentialOwner, + source: ClaudeOAuthCredentialSource) + { + self.credentials = credentials + self.owner = owner + self.source = source + self.inheritedHistoryOwnerIdentifier = nil + } + + init( + credentials: ClaudeOAuthCredentials, + owner: ClaudeOAuthCredentialOwner, + source: ClaudeOAuthCredentialSource, + historyOwnerIdentifier: String?) + { + self.credentials = credentials + self.owner = owner + self.source = source + self.inheritedHistoryOwnerIdentifier = ClaudeOAuthCredentials.normalizedHistoryOwnerIdentifier( + historyOwnerIdentifier) + } +} + +public enum ClaudeOAuthCredentialsError: LocalizedError, Sendable { + case decodeFailed + case missingOAuth + case mcpOAuthOnlyKeychain + case missingAccessToken + case notFound + case keychainError(Int) + case readFailed(String) + case refreshFailed(String) + case noRefreshToken + case refreshDelegatedToClaudeCLI + + public var errorDescription: String? { + switch self { + case .decodeFailed: + return "Claude OAuth credentials are invalid." + case .missingOAuth: + return "Claude OAuth credentials missing. Run `claude` to authenticate." + case .mcpOAuthOnlyKeychain: + return "Claude keychain contains MCP OAuth state only (no claudeAiOauth). " + + "Claude Code may store subscription OAuth elsewhere now. " + + "Open the CodexBar menu and click Refresh to re-authenticate, " + + "or switch Claude Usage source to Web/CLI." + case .missingAccessToken: + return "Claude OAuth access token missing. Run `claude` to authenticate." + case .notFound: + return "Claude OAuth credentials not found. Run `claude` to authenticate." + case let .keychainError(status): + #if os(macOS) + if status == Int(errSecUserCanceled) + || status == Int(errSecAuthFailed) + || status == Int(errSecInteractionNotAllowed) + || status == Int(errSecNoAccessForItem) + { + return "Claude Keychain access was denied. CodexBar will back off in the background until you retry " + + "via a user action (menu open / manual refresh). " + + "Switch Claude Usage source to Web/CLI, or allow access in Keychain Access." + } + #endif + return "Claude OAuth keychain error: \(status)" + case let .readFailed(message): + return "Claude OAuth credentials read failed: \(message)" + case let .refreshFailed(message): + return "Claude OAuth token refresh failed: \(message)" + case .noRefreshToken: + return "Claude OAuth refresh token missing. Run `claude` to authenticate." + case .refreshDelegatedToClaudeCLI: + return "Claude OAuth refresh is delegated to Claude CLI." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials+Hashing.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials+Hashing.swift new file mode 100644 index 0000000..5e83619 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials+Hashing.swift @@ -0,0 +1,18 @@ +import Foundation + +#if canImport(CryptoKit) +import CryptoKit +#endif + +extension ClaudeOAuthCredentialsStore { + static func sha256Prefix(_ data: Data) -> String? { + #if canImport(CryptoKit) + let digest = SHA256.hash(data: data) + let hex = digest.compactMap { String(format: "%02x", $0) }.joined() + return String(hex.prefix(12)) + #else + _ = data + return nil + #endif + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials+SecurityCLIReader.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials+SecurityCLIReader.swift new file mode 100644 index 0000000..76df9cb --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials+SecurityCLIReader.swift @@ -0,0 +1,375 @@ +import Dispatch +import Foundation + +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif + +extension ClaudeOAuthCredentialsStore { + private static let securityBinaryPath = "/usr/bin/security" + private static let securityCLIReadTimeout: TimeInterval = 1.5 + static let isolatedSecurityCLIKeychainEnvironmentKey = "CODEXBAR_CLAUDE_SECURITY_CLI_KEYCHAIN" + + struct SecurityCLIReadRequest { + let account: String? + } + + static func shouldPreferSecurityCLIKeychainRead( + readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current()) + -> Bool + { + readStrategy == .securityCLIExperimental + } + + #if os(macOS) + private enum SecurityCLIReadError: Error { + case binaryUnavailable + case isolatedKeychainUnavailable + case launchFailed + case timedOut + case nonZeroExit(status: Int32, stderrLength: Int) + } + + private struct SecurityCLIReadCommandResult { + let status: Int32 + let stdout: Data + let stderrLength: Int + let durationMs: Double + } + + /// Attempts a Claude keychain read via `/usr/bin/security` when the experimental reader is enabled. + /// - Important: `interaction` is diagnostics context only and does not gate CLI execution. + static func loadFromClaudeKeychainViaSecurityCLIIfEnabled( + interaction: ProviderInteraction, + readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current()) + -> Data? + { + guard let sanitized = self.readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled( + interaction: interaction, + readStrategy: readStrategy) + else { + return nil + } + + let interactionMetadata = interaction == .userInitiated ? "user" : "background" + let parsedCredentials: ClaudeOAuthCredentials + do { + parsedCredentials = try ClaudeOAuthCredentials.parse(data: sanitized) + } catch { + self.log.warning( + "Claude keychain security CLI output invalid; falling back", + metadata: [ + "reader": "securityCLI", + "callerInteraction": interactionMetadata, + "payload_bytes": "\(sanitized.count)", + "parse_error_type": String(describing: type(of: error)), + ]) + return nil + } + + var metadata: [String: String] = [ + "reader": "securityCLI", + "callerInteraction": interactionMetadata, + "payload_bytes": "\(sanitized.count)", + ] + for (key, value) in parsedCredentials.diagnosticsMetadata(now: Date()) { + metadata[key] = value + } + self.log.debug( + "Claude keychain security CLI read succeeded", + metadata: metadata) + return sanitized + } + + static func readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled( + interaction: ProviderInteraction, + readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current(), + environment: [String: String] = ProcessInfo.processInfo.environment) + -> Data? + { + guard self.shouldPreferSecurityCLIKeychainRead(readStrategy: readStrategy) else { return nil } + let interactionMetadata = interaction == .userInitiated ? "user" : "background" + + do { + let preferredAccount = self.preferredClaudeKeychainAccountForSecurityCLIRead( + interaction: interaction) + let output: Data + let status: Int32 + let stderrLength: Int + let durationMs: Double + #if DEBUG + if let override = self.taskSecurityCLIReadOverride ?? self.securityCLIReadOverride { + switch override { + case let .data(data): + output = data ?? Data() + status = 0 + stderrLength = 0 + durationMs = 0 + case .timedOut: + throw SecurityCLIReadError.timedOut + case .nonZeroExit: + throw SecurityCLIReadError.nonZeroExit(status: 1, stderrLength: 0) + case let .dynamic(read): + output = read(SecurityCLIReadRequest(account: preferredAccount)) ?? Data() + status = 0 + stderrLength = 0 + durationMs = 0 + } + } else { + let result = try self.runClaudeSecurityCLIRead( + timeout: self.securityCLIReadTimeout, + account: preferredAccount, + environment: environment) + output = result.stdout + status = result.status + stderrLength = result.stderrLength + durationMs = result.durationMs + } + #else + let result = try self.runClaudeSecurityCLIRead( + timeout: self.securityCLIReadTimeout, + account: preferredAccount, + environment: environment) + output = result.stdout + status = result.status + stderrLength = result.stderrLength + durationMs = result.durationMs + #endif + + let sanitized = self.sanitizeSecurityCLIOutput(output) + guard !sanitized.isEmpty else { return nil } + if ClaudeOAuthCredentials.isMcpOAuthOnlyPayload(data: sanitized) { + self.log.warning( + "Claude keychain security CLI output is MCP OAuth only; falling back", + metadata: [ + "reader": "securityCLI", + "callerInteraction": interactionMetadata, + "status": "\(status)", + "duration_ms": String(format: "%.2f", durationMs), + "stderr_length": "\(stderrLength)", + "payload_bytes": "\(sanitized.count)", + ]) + } else { + self.log.debug( + "Claude keychain security CLI raw read succeeded", + metadata: [ + "reader": "securityCLI", + "callerInteraction": interactionMetadata, + "status": "\(status)", + "duration_ms": String(format: "%.2f", durationMs), + "stderr_length": "\(stderrLength)", + "payload_bytes": "\(sanitized.count)", + ]) + } + return sanitized + } catch let error as SecurityCLIReadError { + var metadata: [String: String] = [ + "reader": "securityCLI", + "callerInteraction": interactionMetadata, + "error_type": String(describing: type(of: error)), + ] + switch error { + case .binaryUnavailable: + metadata["reason"] = "binaryUnavailable" + case .isolatedKeychainUnavailable: + metadata["reason"] = "isolatedKeychainUnavailable" + case .launchFailed: + metadata["reason"] = "launchFailed" + case .timedOut: + metadata["reason"] = "timedOut" + case let .nonZeroExit(status, stderrLength): + metadata["reason"] = "nonZeroExit" + metadata["status"] = "\(status)" + metadata["stderr_length"] = "\(stderrLength)" + } + self.log.warning("Claude keychain security CLI read failed; falling back", metadata: metadata) + return nil + } catch { + self.log.warning( + "Claude keychain security CLI read failed; falling back", + metadata: [ + "reader": "securityCLI", + "callerInteraction": interactionMetadata, + "error_type": String(describing: type(of: error)), + ]) + return nil + } + } + + private static func sanitizeSecurityCLIOutput(_ data: Data) -> Data { + var sanitized = data + while let last = sanitized.last, last == 0x0A || last == 0x0D { + sanitized.removeLast() + } + return sanitized + } + + private static func runClaudeSecurityCLIRead( + timeout: TimeInterval, + account: String?, + environment: [String: String]) throws -> SecurityCLIReadCommandResult + { + guard FileManager.default.isExecutableFile(atPath: self.securityBinaryPath) else { + throw SecurityCLIReadError.binaryUnavailable + } + guard let arguments = self.securityCLIReadArguments(account: account, environment: environment) else { + throw SecurityCLIReadError.isolatedKeychainUnavailable + } + + let process = Process() + process.executableURL = URL(fileURLWithPath: self.securityBinaryPath) + process.arguments = arguments + + let stdoutPipe = Pipe() + let stderrPipe = Pipe() + process.standardOutput = stdoutPipe + process.standardError = stderrPipe + process.standardInput = nil + + let startedAt = DispatchTime.now().uptimeNanoseconds + do { + try process.run() + } catch { + throw SecurityCLIReadError.launchFailed + } + + var processGroup: pid_t? + let pid = process.processIdentifier + if setpgid(pid, pid) == 0 { + processGroup = pid + } + + let deadline = Date().addingTimeInterval(timeout) + while process.isRunning, Date() < deadline { + Thread.sleep(forTimeInterval: 0.02) + } + + if process.isRunning { + self.terminate(process: process, processGroup: processGroup) + throw SecurityCLIReadError.timedOut + } + + let stdout = stdoutPipe.fileHandleForReading.readDataToEndOfFile() + let stderr = stderrPipe.fileHandleForReading.readDataToEndOfFile() + let status = process.terminationStatus + let durationMs = Double(DispatchTime.now().uptimeNanoseconds - startedAt) / 1_000_000.0 + guard status == 0 else { + throw SecurityCLIReadError.nonZeroExit(status: status, stderrLength: stderr.count) + } + + return SecurityCLIReadCommandResult( + status: status, + stdout: stdout, + stderrLength: stderr.count, + durationMs: durationMs) + } + + static func securityCLIReadArguments( + account: String?, + environment: [String: String]) -> [String]? + { + let isolatedKeychainPath = self.isolatedSecurityCLIKeychainPath(environment: environment) + if KeychainTestSafety.shouldBlockRealKeychainAccess(environment: environment), + isolatedKeychainPath == nil + { + return nil + } + + var arguments = [ + "find-generic-password", + "-s", + self.claudeKeychainService, + ] + if let account, !account.isEmpty { + arguments.append(contentsOf: ["-a", account]) + } + arguments.append("-w") + + let rawIsolatedPath = environment[self.isolatedSecurityCLIKeychainEnvironmentKey]? + .trimmingCharacters(in: .whitespacesAndNewlines) + if rawIsolatedPath != nil || KeychainAccessGate.isDisabledByEnvironment(environment) { + guard let isolatedKeychainPath else { + return nil + } + arguments.append(isolatedKeychainPath) + } + return arguments + } + + private static func terminate(process: Process, processGroup: pid_t?) { + guard process.isRunning else { return } + process.terminate() + if let processGroup { + kill(-processGroup, SIGTERM) + } + let deadline = Date().addingTimeInterval(0.4) + while process.isRunning, Date() < deadline { + usleep(50000) + } + if process.isRunning { + if let processGroup { + kill(-processGroup, SIGKILL) + } + kill(process.processIdentifier, SIGKILL) + } + } + #else + static func loadFromClaudeKeychainViaSecurityCLIIfEnabled( + interaction _: ProviderInteraction, + readStrategy _: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current()) + -> Data? + { + nil + } + + static func readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled( + interaction _: ProviderInteraction, + readStrategy _: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current(), + environment _: [String: String] = ProcessInfo.processInfo.environment) + -> Data? + { + nil + } + #endif + + private static func isolatedSecurityCLIKeychainPath(environment: [String: String]) -> String? { + guard KeychainAccessGate.isDisabledByEnvironment(environment) else { return nil } + guard let rawPath = environment[self.isolatedSecurityCLIKeychainEnvironmentKey]? + .trimmingCharacters(in: .whitespacesAndNewlines), + !rawPath.isEmpty, + (rawPath as NSString).isAbsolutePath + else { + return nil + } + return URL(fileURLWithPath: rawPath).standardizedFileURL.path + } + + static func isMcpOAuthOnlyClaudeKeychainPayloadPresent( + interaction: ProviderInteraction, + readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current(), + keychainAccessDisabled: Bool = KeychainAccessGate.isDisabled, + environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool + { + guard !keychainAccessDisabled || self.isolatedSecurityCLIKeychainPath(environment: environment) != nil else { + return false + } + guard ClaudeOAuthKeychainPromptPreference.effectiveMode(readStrategy: readStrategy) != .never else { + return false + } + let payload: Data? = switch readStrategy { + case .securityFramework: + self.readRawClaudeKeychainPayloadViaSecurityFrameworkWithoutPrompt() + case .securityCLIExperimental: + self.readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled( + interaction: interaction, + readStrategy: readStrategy, + environment: environment) + } + guard let payload else { return false } + return ClaudeOAuthCredentials.isMcpOAuthOnlyPayload(data: payload) + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials+TestingOverrides.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials+TestingOverrides.swift new file mode 100644 index 0000000..4510bab --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials+TestingOverrides.swift @@ -0,0 +1,430 @@ +import Foundation + +#if DEBUG +extension ClaudeOAuthCredentialsStore { + nonisolated(unsafe) static var claudeKeychainDataOverride: Data? + nonisolated(unsafe) static var claudeKeychainFingerprintOverride: ClaudeKeychainFingerprint? + @TaskLocal static var taskClaudeKeychainDataOverride: Data? + @TaskLocal static var taskClaudeKeychainFingerprintOverride: ClaudeKeychainFingerprint? + @TaskLocal static var taskMemoryCacheStoreOverride: MemoryCacheStore? + @TaskLocal static var taskClaudeKeychainFingerprintStoreOverride: ClaudeKeychainFingerprintStore? + @TaskLocal static var taskPendingCacheClearStoreOverride: ClaudeOAuthPendingCacheClearStore? + + typealias OAuthCacheOperation = KeychainCacheStore.Operation + typealias OAuthCacheOperationRecorder = KeychainCacheStore.OperationRecorder + + final class PendingCacheClearMemoryStore: ClaudeOAuthPendingCacheClearStore, @unchecked Sendable { + private let lock = NSLock() + private var pending: Bool + + init(isPending: Bool = false) { + self.pending = isPending + } + + var isPending: Bool { + self.lock.lock() + defer { self.lock.unlock() } + return self.pending + } + + func markPending() { + self.lock.lock() + self.pending = true + self.lock.unlock() + } + + func withCacheTransaction(_ operation: (inout Bool) -> Void) { + self.lock.lock() + defer { self.lock.unlock() } + operation(&self.pending) + } + } + + final class ClaudeKeychainFingerprintStore: @unchecked Sendable { + var fingerprint: ClaudeKeychainFingerprint? + + init(fingerprint: ClaudeKeychainFingerprint? = nil) { + self.fingerprint = fingerprint + } + } + + final class MemoryCacheStore: @unchecked Sendable { + var record: ClaudeOAuthCredentialRecord? + var timestamp: Date? + } + + static func setClaudeKeychainDataOverrideForTesting(_ data: Data?) { + self.claudeKeychainDataOverride = data + } + + static func setClaudeKeychainFingerprintOverrideForTesting(_ fingerprint: ClaudeKeychainFingerprint?) { + self.claudeKeychainFingerprintOverride = fingerprint + } + + static func withClaudeKeychainOverridesForTesting( + data: Data?, + fingerprint: ClaudeKeychainFingerprint?, + operation: () throws -> T) rethrows -> T + { + try self.$taskClaudeKeychainDataOverride.withValue(data) { + try self.$taskClaudeKeychainFingerprintOverride.withValue(fingerprint) { + try operation() + } + } + } + + static func withClaudeKeychainOverridesForTesting( + data: Data?, + fingerprint: ClaudeKeychainFingerprint?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskClaudeKeychainDataOverride.withValue(data) { + try await self.$taskClaudeKeychainFingerprintOverride.withValue(fingerprint) { + try await operation() + } + } + } + + static func withClaudeKeychainFingerprintStoreOverrideForTesting( + _ store: ClaudeKeychainFingerprintStore?, + operation: () throws -> T) rethrows -> T + { + try self.$taskClaudeKeychainFingerprintStoreOverride.withValue(store) { + try operation() + } + } + + static func withClaudeKeychainFingerprintStoreOverrideForTesting( + _ store: ClaudeKeychainFingerprintStore?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskClaudeKeychainFingerprintStoreOverride.withValue(store) { + try await operation() + } + } + + static func withIsolatedMemoryCacheForTesting(operation: () throws -> T) rethrows -> T { + let store = MemoryCacheStore() + let preAlertStore = ClaudeOAuthKeychainPreAlertGate.StateStore() + return try ClaudeOAuthKeychainPreAlertGate.withStateStoreOverrideForTesting(preAlertStore) { + try self.$taskMemoryCacheStoreOverride.withValue(store) { + try operation() + } + } + } + + static func withIsolatedMemoryCacheForTesting(operation: () async throws -> T) async rethrows -> T { + let store = MemoryCacheStore() + let preAlertStore = ClaudeOAuthKeychainPreAlertGate.StateStore() + return try await ClaudeOAuthKeychainPreAlertGate.withStateStoreOverrideForTesting(preAlertStore) { + try await self.$taskMemoryCacheStoreOverride.withValue(store) { + try await operation() + } + } + } + + final class CredentialsFileFingerprintStore: @unchecked Sendable { + var fingerprint: CredentialsFileFingerprint? + + init(fingerprint: CredentialsFileFingerprint? = nil) { + self.fingerprint = fingerprint + } + + func load() -> CredentialsFileFingerprint? { + self.fingerprint + } + + func save(_ fingerprint: CredentialsFileFingerprint?) { + self.fingerprint = fingerprint + } + } + + enum SecurityCLIReadOverride { + case data(Data?) + case timedOut + case nonZeroExit + case dynamic(@Sendable (SecurityCLIReadRequest) -> Data?) + } + + @TaskLocal static var taskKeychainAccessOverride: Bool? + @TaskLocal static var taskCredentialsFileFingerprintStoreOverride: CredentialsFileFingerprintStore? + @TaskLocal static var taskSecurityCLIReadOverride: SecurityCLIReadOverride? + @TaskLocal static var taskSecurityCLIReadAccountOverride: String? + nonisolated(unsafe) static var securityCLIReadOverride: SecurityCLIReadOverride? + + public struct TestingOverridesSnapshot: Sendable { + let keychainOverrideStore: ClaudeKeychainOverrideStore? + let keychainData: Data? + let keychainFingerprint: ClaudeKeychainFingerprint? + let memoryCacheStore: MemoryCacheStore? + let fingerprintStore: ClaudeKeychainFingerprintStore? + let keychainAccessOverride: Bool? + let credentialsFileFingerprintStore: CredentialsFileFingerprintStore? + let securityCLIReadOverride: SecurityCLIReadOverride? + let securityCLIReadAccountOverride: String? + let pendingCacheClearStore: ClaudeOAuthPendingCacheClearStore? + let oauthCacheOperationRecorder: OAuthCacheOperationRecorder? + + init( + keychainOverrideStore: ClaudeKeychainOverrideStore?, + keychainData: Data?, + keychainFingerprint: ClaudeKeychainFingerprint?, + memoryCacheStore: MemoryCacheStore?, + fingerprintStore: ClaudeKeychainFingerprintStore?, + keychainAccessOverride: Bool?, + credentialsFileFingerprintStore: CredentialsFileFingerprintStore?, + securityCLIReadOverride: SecurityCLIReadOverride?, + securityCLIReadAccountOverride: String?, + pendingCacheClearStore: ClaudeOAuthPendingCacheClearStore?, + oauthCacheOperationRecorder: OAuthCacheOperationRecorder?) + { + self.keychainOverrideStore = keychainOverrideStore + self.keychainData = keychainData + self.keychainFingerprint = keychainFingerprint + self.memoryCacheStore = memoryCacheStore + self.fingerprintStore = fingerprintStore + self.keychainAccessOverride = keychainAccessOverride + self.credentialsFileFingerprintStore = credentialsFileFingerprintStore + self.securityCLIReadOverride = securityCLIReadOverride + self.securityCLIReadAccountOverride = securityCLIReadAccountOverride + self.pendingCacheClearStore = pendingCacheClearStore + self.oauthCacheOperationRecorder = oauthCacheOperationRecorder + } + } + + static func withKeychainAccessOverrideForTesting( + _ disabled: Bool?, + operation: () throws -> T) rethrows -> T + { + try self.$taskKeychainAccessOverride.withValue(disabled) { + try operation() + } + } + + static func withKeychainAccessOverrideForTesting( + _ disabled: Bool?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskKeychainAccessOverride.withValue(disabled) { + try await operation() + } + } + + fileprivate static func withCredentialsFileFingerprintStoreOverrideForTesting( + _ store: CredentialsFileFingerprintStore?, + operation: () throws -> T) rethrows -> T + { + try self.$taskCredentialsFileFingerprintStoreOverride.withValue(store) { + try operation() + } + } + + fileprivate static func withCredentialsFileFingerprintStoreOverrideForTesting( + _ store: CredentialsFileFingerprintStore?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskCredentialsFileFingerprintStoreOverride.withValue(store) { + try await operation() + } + } + + static func withIsolatedCredentialsFileTrackingForTesting( + operation: () throws -> T) rethrows -> T + { + let store = CredentialsFileFingerprintStore() + return try self.$taskCredentialsFileFingerprintStoreOverride.withValue(store) { + try operation() + } + } + + static func withPendingCacheClearStoreOverrideForTesting( + _ store: ClaudeOAuthPendingCacheClearStore?, + operation: () throws -> T) rethrows -> T + { + try self.$taskPendingCacheClearStoreOverride.withValue(store) { + try operation() + } + } + + static func withPendingCacheClearStoreOverrideForTesting( + _ store: ClaudeOAuthPendingCacheClearStore?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskPendingCacheClearStoreOverride.withValue(store) { + try await operation() + } + } + + static func withOAuthCacheOperationRecorderForTesting( + _ recorder: OAuthCacheOperationRecorder?, + operation: () throws -> T) rethrows -> T + { + try KeychainCacheStore.withOperationRecorderForTesting(recorder, operation: operation) + } + + static func withOAuthCacheOperationRecorderForTesting( + _ recorder: OAuthCacheOperationRecorder?, + operation: () async throws -> T) async rethrows -> T + { + try await KeychainCacheStore.withOperationRecorderForTesting(recorder, operation: operation) + } + + static func withIsolatedCredentialsFileTrackingForTesting( + operation: () async throws -> T) async rethrows -> T + { + let store = CredentialsFileFingerprintStore() + return try await self.$taskCredentialsFileFingerprintStoreOverride.withValue(store) { + try await operation() + } + } + + static func withSecurityCLIReadOverrideForTesting( + _ readOverride: SecurityCLIReadOverride?, + operation: () throws -> T) rethrows -> T + { + try self.$taskSecurityCLIReadOverride.withValue(readOverride) { + try operation() + } + } + + static func withSecurityCLIReadOverrideForTesting( + _ readOverride: SecurityCLIReadOverride?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskSecurityCLIReadOverride.withValue(readOverride) { + try await operation() + } + } + + static func currentSecurityCLIReadOverrideForTesting() -> SecurityCLIReadOverride? { + self.taskSecurityCLIReadOverride ?? self.securityCLIReadOverride + } + + static func withSecurityCLIReadAccountOverrideForTesting( + _ account: String?, + operation: () throws -> T) rethrows -> T + { + try self.$taskSecurityCLIReadAccountOverride.withValue(account) { + try operation() + } + } + + static func withSecurityCLIReadAccountOverrideForTesting( + _ account: String?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskSecurityCLIReadAccountOverride.withValue(account) { + try await operation() + } + } + + public static func withCurrentTestingOverridesForTask( + operation: () async throws -> T) async rethrows -> T + { + try await self.withTestingOverridesSnapshotForTask( + self.currentTestingOverridesSnapshotForTask, + operation: operation) + } + + public static func withCurrentTestingOverridesForTask( + operation: () throws -> T) rethrows -> T + { + try self.withTestingOverridesSnapshotForTask( + self.currentTestingOverridesSnapshotForTask, + operation: operation) + } + + public static var currentTestingOverridesSnapshotForTask: TestingOverridesSnapshot { + TestingOverridesSnapshot( + keychainOverrideStore: self.taskClaudeKeychainOverrideStore, + keychainData: self.taskClaudeKeychainDataOverride, + keychainFingerprint: self.taskClaudeKeychainFingerprintOverride, + memoryCacheStore: self.taskMemoryCacheStoreOverride, + fingerprintStore: self.taskClaudeKeychainFingerprintStoreOverride, + keychainAccessOverride: self.taskKeychainAccessOverride, + credentialsFileFingerprintStore: self.taskCredentialsFileFingerprintStoreOverride, + securityCLIReadOverride: self.taskSecurityCLIReadOverride, + securityCLIReadAccountOverride: self.taskSecurityCLIReadAccountOverride, + pendingCacheClearStore: self.taskPendingCacheClearStoreOverride, + oauthCacheOperationRecorder: KeychainCacheStore.currentOperationRecorderForTesting) + } + + public static func withTestingOverridesSnapshotForTask( + _ snapshot: TestingOverridesSnapshot, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskPendingCacheClearStoreOverride.withValue(snapshot.pendingCacheClearStore) { + try await KeychainCacheStore.withOperationRecorderForTesting(snapshot.oauthCacheOperationRecorder) { + try await self.$taskClaudeKeychainOverrideStore.withValue(snapshot.keychainOverrideStore) { + try await self.$taskClaudeKeychainDataOverride.withValue(snapshot.keychainData) { + try await self.$taskClaudeKeychainFingerprintOverride.withValue(snapshot.keychainFingerprint) { + try await self.$taskMemoryCacheStoreOverride.withValue(snapshot.memoryCacheStore) { + try await self.$taskClaudeKeychainFingerprintStoreOverride + .withValue(snapshot.fingerprintStore) { + try await self.$taskKeychainAccessOverride + .withValue(snapshot.keychainAccessOverride) { + try await self.$taskCredentialsFileFingerprintStoreOverride.withValue( + snapshot.credentialsFileFingerprintStore) + { + try await self.$taskSecurityCLIReadOverride.withValue( + snapshot.securityCLIReadOverride) + { + try await self.$taskSecurityCLIReadAccountOverride.withValue( + snapshot.securityCLIReadAccountOverride) + { + try await operation() + } + } + } + } + } + } + } + } + } + } + } + } + + public static func withTestingOverridesSnapshotForTask( + _ snapshot: TestingOverridesSnapshot, + operation: () throws -> T) rethrows -> T + { + try self.$taskPendingCacheClearStoreOverride.withValue(snapshot.pendingCacheClearStore) { + try KeychainCacheStore.withOperationRecorderForTesting(snapshot.oauthCacheOperationRecorder) { + try self.$taskClaudeKeychainOverrideStore.withValue(snapshot.keychainOverrideStore) { + try self.$taskClaudeKeychainDataOverride.withValue(snapshot.keychainData) { + try self.$taskClaudeKeychainFingerprintOverride.withValue(snapshot.keychainFingerprint) { + try self.$taskMemoryCacheStoreOverride.withValue(snapshot.memoryCacheStore) { + try self.$taskClaudeKeychainFingerprintStoreOverride + .withValue(snapshot.fingerprintStore) { + try self.$taskKeychainAccessOverride + .withValue(snapshot.keychainAccessOverride) { + try self.$taskCredentialsFileFingerprintStoreOverride.withValue( + snapshot.credentialsFileFingerprintStore) + { + try self.$taskSecurityCLIReadOverride.withValue( + snapshot.securityCLIReadOverride) + { + try self.$taskSecurityCLIReadAccountOverride.withValue( + snapshot.securityCLIReadAccountOverride) + { + try operation() + } + } + } + } + } + } + } + } + } + } + } + } + + static func setSecurityCLIReadOverrideForTesting(_ readOverride: SecurityCLIReadOverride?) { + self.securityCLIReadOverride = readOverride + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials.swift new file mode 100644 index 0000000..bef352f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthCredentials.swift @@ -0,0 +1,2563 @@ +import Dispatch +import Foundation + +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +#if os(macOS) +import LocalAuthentication +import Security +#endif + +// swiftlint:disable type_body_length file_length +public enum ClaudeOAuthCredentialsStore { + private static let credentialsPath = ".claude/.credentials.json" + static let claudeKeychainService = "Claude Code-credentials" + private static let cacheKey = KeychainCacheStore.Key.oauth(provider: .claude) + public static let environmentTokenKey = "CODEXBAR_CLAUDE_OAUTH_TOKEN" + public static let environmentScopesKey = "CODEXBAR_CLAUDE_OAUTH_SCOPES" + + // Claude CLI's OAuth client ID - this is a public identifier (not a secret). + // It's the same client ID used by Claude Code CLI for OAuth PKCE flow. + // Can be overridden via environment variable if Anthropic ever changes it. + public static let defaultOAuthClientID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e" + public static let environmentClientIDKey = "CODEXBAR_CLAUDE_OAUTH_CLIENT_ID" + private static let tokenRefreshEndpoint = "https://platform.claude.com/v1/oauth/token" + + private static var oauthClientID: String { + ProcessInfo.processInfo.environment[self.environmentClientIDKey]? + .trimmingCharacters(in: .whitespacesAndNewlines) + ?? self.defaultOAuthClientID + } + + static let log = CodexBarLog.logger(LogCategories.claudeUsage) + private static let fileFingerprintKey = "ClaudeOAuthCredentialsFileFingerprintV2" + private static let claudeKeychainPromptLock = NSLock() + private static let claudeKeychainFingerprintKey = "ClaudeOAuthClaudeKeychainFingerprintV2" + private static let claudeKeychainFingerprintLegacyKey = "ClaudeOAuthClaudeKeychainFingerprintV1" + private static let pendingCodexBarOAuthKeychainCacheClearKey = + "ClaudeOAuthPendingCodexBarOAuthKeychainCacheClearV1" + private static let pendingCodexBarOAuthKeychainCacheClearStore: ClaudeOAuthPendingCacheClearStore = + ClaudeOAuthPendingCacheClearUserDefaultsStore( + // The cache service is shared by release/debug apps and their CLIs, so its tombstone is shared too. + domain: "com.steipete.codexbar", + key: ClaudeOAuthCredentialsStore.pendingCodexBarOAuthKeychainCacheClearKey) + private static let claudeKeychainChangeCheckLock = NSLock() + private nonisolated(unsafe) static var lastClaudeKeychainChangeCheckAt: Date? + private static let claudeKeychainChangeCheckMinimumInterval: TimeInterval = 60 + private static let reauthenticateHint = "Run `claude` to re-authenticate." + + struct ClaudeKeychainFingerprint: Codable, Equatable { + let modifiedAt: Int? + let createdAt: Int? + let persistentRefHash: String? + } + + private struct ClaudeKeychainCredentialEvidence { + let credentials: ClaudeOAuthCredentials + let persistentRefHash: String + } + + struct CredentialsFileFingerprint: Codable, Equatable { + let modifiedAtMs: Int? + let size: Int + } + + struct CacheEntry: Codable { + let data: Data + let storedAt: Date + let owner: ClaudeOAuthCredentialOwner? + let historyOwnerIdentifier: String? + + init( + data: Data, + storedAt: Date, + owner: ClaudeOAuthCredentialOwner? = nil, + historyOwnerIdentifier: String? = nil) + { + self.data = data + self.storedAt = storedAt + self.owner = owner + self.historyOwnerIdentifier = ClaudeOAuthCredentials.normalizedHistoryOwnerIdentifier( + historyOwnerIdentifier) + } + } + + private nonisolated(unsafe) static var credentialsURLOverride: URL? + #if DEBUG + @TaskLocal private static var taskCredentialsURLOverride: URL? + #endif + @TaskLocal static var allowBackgroundPromptBootstrap: Bool = false + // In-memory cache (nonisolated for synchronous access) + private static let memoryCacheLock = NSLock() + private nonisolated(unsafe) static var cachedCredentialRecord: ClaudeOAuthCredentialRecord? + private nonisolated(unsafe) static var cacheTimestamp: Date? + private static let memoryCacheValidityDuration: TimeInterval = 1800 + + private static func readMemoryCache() -> (record: ClaudeOAuthCredentialRecord?, timestamp: Date?) { + #if DEBUG + if let store = self.taskMemoryCacheStoreOverride { + return (store.record, store.timestamp) + } + #endif + self.memoryCacheLock.lock() + defer { self.memoryCacheLock.unlock() } + return (self.cachedCredentialRecord, self.cacheTimestamp) + } + + private static func writeMemoryCache(record: ClaudeOAuthCredentialRecord?, timestamp: Date?) { + #if DEBUG + if let store = self.taskMemoryCacheStoreOverride { + store.record = record + store.timestamp = timestamp + return + } + #endif + self.memoryCacheLock.lock() + self.cachedCredentialRecord = record + self.cacheTimestamp = timestamp + self.memoryCacheLock.unlock() + } + + private struct CollaboratorContext { + let allowBackgroundPromptBootstrap: Bool + #if DEBUG + let credentialsURLOverride: URL? + let testingOverrides: TestingOverridesSnapshot + #endif + + func run(_ operation: () throws -> T) rethrows -> T { + try ClaudeOAuthCredentialsStore.$allowBackgroundPromptBootstrap + .withValue(self.allowBackgroundPromptBootstrap) { + #if DEBUG + try ClaudeOAuthCredentialsStore.withTestingOverridesSnapshotForTask(self.testingOverrides) { + try ClaudeOAuthCredentialsStore + .withCredentialsURLOverrideForTesting(self.credentialsURLOverride) { + try operation() + } + } + #else + try operation() + #endif + } + } + + func run(_ operation: () async throws -> T) async rethrows -> T { + try await ClaudeOAuthCredentialsStore.$allowBackgroundPromptBootstrap + .withValue(self.allowBackgroundPromptBootstrap) { + #if DEBUG + try await ClaudeOAuthCredentialsStore.withTestingOverridesSnapshotForTask(self.testingOverrides) { + try await ClaudeOAuthCredentialsStore.withCredentialsURLOverrideForTesting( + self.credentialsURLOverride) + { + try await operation() + } + } + #else + try await operation() + #endif + } + } + } + + private static func currentCollaboratorContext() -> CollaboratorContext { + #if DEBUG + CollaboratorContext( + allowBackgroundPromptBootstrap: self.allowBackgroundPromptBootstrap, + credentialsURLOverride: self.taskCredentialsURLOverride, + testingOverrides: self.currentTestingOverridesSnapshotForTask) + #else + CollaboratorContext( + allowBackgroundPromptBootstrap: self.allowBackgroundPromptBootstrap) + #endif + } + + private struct Repository { + let context: CollaboratorContext + + func load(environment: [String: String], allowKeychainPrompt: Bool, respectKeychainPromptCooldown: Bool) throws + -> ClaudeOAuthCredentials + { + try self.loadRecord( + environment: environment, + allowKeychainPrompt: allowKeychainPrompt, + respectKeychainPromptCooldown: respectKeychainPromptCooldown, + allowClaudeKeychainRepairWithoutPrompt: true).credentials + } + + func loadRecord( + environment: [String: String], + allowKeychainPrompt: Bool, + respectKeychainPromptCooldown: Bool, + allowClaudeKeychainRepairWithoutPrompt: Bool) throws -> ClaudeOAuthCredentialRecord + { + try self.context.run { + let shouldRespectKeychainPromptCooldownForSilentProbes = + respectKeychainPromptCooldown || !allowKeychainPrompt + + if let credentials = ClaudeOAuthCredentialsStore.loadFromEnvironment(environment) { + return ClaudeOAuthCredentialRecord( + credentials: credentials, + owner: .environment, + source: .environment) + } + + _ = self.invalidateCacheIfCredentialsFileChanged() + + let recovery = Recovery(context: self.context) + let memory = ClaudeOAuthCredentialsStore.readMemoryCache() + if ClaudeOAuthCredentialsStore.shouldUseCodexBarOAuthKeychainCache, + !ClaudeOAuthCredentialsStore.hasPendingCodexBarOAuthKeychainCacheClear, + let cachedRecord = memory.record, + let timestamp = memory.timestamp, + Date().timeIntervalSince(timestamp) < ClaudeOAuthCredentialsStore.memoryCacheValidityDuration, + !cachedRecord.credentials.isExpired + { + let owner = self.resolvedCacheOwner(cachedRecord.owner) + let record = ClaudeOAuthCredentialRecord( + credentials: cachedRecord.credentials, + owner: owner, + source: .memoryCache, + historyOwnerIdentifier: cachedRecord.historyOwnerIdentifier) + if recovery.shouldAttemptFreshnessSyncFromClaudeKeychain(cached: record), + let synced = recovery.syncWithClaudeKeychainIfChanged( + cached: record, + respectKeychainPromptCooldown: shouldRespectKeychainPromptCooldownForSilentProbes) + { + return synced + } + return record + } + + var lastError: Error? + var expiredRecord: ClaudeOAuthCredentialRecord? + var cacheTemporarilyUnavailable = false + + switch ClaudeOAuthCredentialsStore.loadCodexBarOAuthKeychainCache() { + case let .found(entry): + if let creds = try? ClaudeOAuthCredentials.parse(data: entry.data) { + let owner = self.resolvedCacheOwner(entry.owner ?? .claudeCLI) + let record = ClaudeOAuthCredentialRecord( + credentials: creds, + owner: owner, + source: .cacheKeychain, + historyOwnerIdentifier: entry.historyOwnerIdentifier) + if creds.isExpired { + expiredRecord = record + } else { + if recovery.shouldAttemptFreshnessSyncFromClaudeKeychain(cached: record), + let synced = recovery.syncWithClaudeKeychainIfChanged( + cached: record, + respectKeychainPromptCooldown: shouldRespectKeychainPromptCooldownForSilentProbes) + { + return synced + } + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: owner, + source: .memoryCache, + historyOwnerIdentifier: record.historyOwnerIdentifier), + timestamp: Date()) + return record + } + } else { + ClaudeOAuthCredentialsStore.clearCacheKeychain() + } + case .invalid: + ClaudeOAuthCredentialsStore.clearCacheKeychain() + case .temporarilyUnavailable: + cacheTemporarilyUnavailable = true + case .missing: + break + } + + do { + let fileData = try ClaudeOAuthCredentialsStore.loadFromFile() + let creds = try ClaudeOAuthCredentials.parse(data: fileData) + let record = ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .credentialsFile) + if creds.isExpired { + expiredRecord = record + } else { + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: Date()) + if !cacheTemporarilyUnavailable { + ClaudeOAuthCredentialsStore.saveToCacheKeychain(fileData, owner: .claudeCLI) + } + return record + } + } catch let error as ClaudeOAuthCredentialsError { + if case .notFound = error { + } else { + lastError = error + } + } catch { + lastError = error + } + + if allowClaudeKeychainRepairWithoutPrompt, !allowKeychainPrompt { + if let repaired = recovery.repairFromClaudeKeychainWithoutPromptIfAllowed( + now: Date(), + respectKeychainPromptCooldown: shouldRespectKeychainPromptCooldownForSilentProbes, + allowCacheKeychainWrite: !cacheTemporarilyUnavailable) + { + return repaired + } + } + + if let prompted = self.loadFromClaudeKeychainWithPromptIfAllowed( + allowKeychainPrompt: allowKeychainPrompt, + respectKeychainPromptCooldown: respectKeychainPromptCooldown, + allowCacheKeychainWrite: !cacheTemporarilyUnavailable, + lastError: &lastError) + { + return prompted + } + + if let expiredRecord { + return expiredRecord + } + if let lastError { + throw lastError + } + throw ClaudeOAuthCredentialsError.notFound + } + } + + private func loadFromClaudeKeychainWithPromptIfAllowed( + allowKeychainPrompt: Bool, + respectKeychainPromptCooldown: Bool, + allowCacheKeychainWrite: Bool, + lastError: inout Error?) -> ClaudeOAuthCredentialRecord? + { + let shouldApplyPromptCooldown = + ClaudeOAuthCredentialsStore.isPromptPolicyApplicable && respectKeychainPromptCooldown + let promptAllowed = + allowKeychainPrompt + && (!shouldApplyPromptCooldown || ClaudeOAuthKeychainAccessGate.shouldAllowPrompt()) + guard promptAllowed else { return nil } + + do { + ClaudeOAuthCredentialsStore.claudeKeychainPromptLock.lock() + defer { ClaudeOAuthCredentialsStore.claudeKeychainPromptLock.unlock() } + + let memory = ClaudeOAuthCredentialsStore.readMemoryCache() + if ClaudeOAuthCredentialsStore.shouldUseCodexBarOAuthKeychainCache, + !ClaudeOAuthCredentialsStore.hasPendingCodexBarOAuthKeychainCacheClear, + let cachedRecord = memory.record, + let timestamp = memory.timestamp, + Date().timeIntervalSince(timestamp) < ClaudeOAuthCredentialsStore.memoryCacheValidityDuration, + !cachedRecord.credentials.isExpired + { + let owner = self.resolvedCacheOwner(cachedRecord.owner) + return ClaudeOAuthCredentialRecord( + credentials: cachedRecord.credentials, + owner: owner, + source: .memoryCache, + historyOwnerIdentifier: cachedRecord.historyOwnerIdentifier) + } + if case let .found(entry) = ClaudeOAuthCredentialsStore.loadCodexBarOAuthKeychainCache(), + let creds = try? ClaudeOAuthCredentials.parse(data: entry.data), + !creds.isExpired + { + let owner = self.resolvedCacheOwner(entry.owner ?? .claudeCLI) + return ClaudeOAuthCredentialRecord( + credentials: creds, + owner: owner, + source: .cacheKeychain, + historyOwnerIdentifier: entry.historyOwnerIdentifier) + } + + let promptMode = ClaudeOAuthKeychainPromptPreference.current() + guard ClaudeOAuthCredentialsStore.shouldAllowClaudeCodeKeychainAccess(mode: promptMode) else { + return nil + } + + if ClaudeOAuthCredentialsStore.shouldPreferSecurityCLIKeychainRead(), + let keychainData = ClaudeOAuthCredentialsStore.loadFromClaudeKeychainViaSecurityCLIIfEnabled( + interaction: ProviderInteractionContext.current) + { + let creds = try ClaudeOAuthCredentials.parse(data: keychainData) + let record = ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .claudeKeychain) + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: Date()) + if allowCacheKeychainWrite { + ClaudeOAuthCredentialsStore.saveToCacheKeychain(keychainData, owner: .claudeCLI) + } + return record + } + + let shouldPreferSecurityCLIKeychainRead = + ClaudeOAuthCredentialsStore.shouldPreferSecurityCLIKeychainRead() + var fallbackPromptMode = promptMode + if shouldPreferSecurityCLIKeychainRead { + fallbackPromptMode = ClaudeOAuthKeychainPromptPreference.securityFrameworkFallbackMode() + let fallbackDecision = ClaudeOAuthCredentialsStore.securityFrameworkFallbackPromptDecision( + promptMode: fallbackPromptMode, + allowKeychainPrompt: allowKeychainPrompt, + respectKeychainPromptCooldown: respectKeychainPromptCooldown) + ClaudeOAuthCredentialsStore.log.debug( + "Claude keychain Security.framework fallback prompt policy evaluated", + metadata: [ + "reader": "securityFrameworkFallback", + "fallbackPromptMode": fallbackPromptMode.rawValue, + "fallbackPromptAllowed": "\(fallbackDecision.allowed)", + "fallbackBlockedReason": fallbackDecision.blockedReason ?? "none", + ]) + guard fallbackDecision.allowed else { return nil } + } + + if ClaudeOAuthCredentialsStore.shouldNotifyClaudeKeychainPreAlert() { + ClaudeOAuthKeychainPreAlertGate.presentIfNeeded { + KeychainPromptHandler.notifyIfHandled( + KeychainPromptContext( + kind: .claudeOAuth, + service: ClaudeOAuthCredentialsStore.claudeKeychainService, + account: nil)) + } + } + let keychainData: Data = if shouldPreferSecurityCLIKeychainRead { + try ClaudeOAuthCredentialsStore.loadFromClaudeKeychainUsingSecurityFramework( + promptMode: fallbackPromptMode, + allowKeychainPrompt: true) + } else { + try ClaudeOAuthCredentialsStore.loadFromClaudeKeychain() + } + let creds = try ClaudeOAuthCredentials.parse(data: keychainData) + let record = ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .claudeKeychain) + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: Date()) + if allowCacheKeychainWrite { + ClaudeOAuthCredentialsStore.saveToCacheKeychain(keychainData, owner: .claudeCLI) + } + return record + } catch let error as ClaudeOAuthCredentialsError { + if case .notFound = error { + } else { + lastError = error + } + } catch { + lastError = error + } + return nil + } + + private func resolvedCacheOwner(_ owner: ClaudeOAuthCredentialOwner) -> ClaudeOAuthCredentialOwner { + guard owner == .codexbar else { return owner } + guard self.hasClaudeCLIStorageWithoutPrompt() else { return owner } + // Claude Code rotates refresh tokens; when its storage exists, it owns the refresh lifecycle. + return .claudeCLI + } + + private func hasClaudeCLIStorageWithoutPrompt() -> Bool { + if ClaudeOAuthCredentialsStore.currentFileFingerprint() != nil { + return true + } + guard ClaudeOAuthKeychainPromptPreference.storedMode() != .never else { return false } + return ClaudeOAuthCredentialsStore.hasClaudeKeychainItemWithoutPrompt() + } + + @discardableResult + func invalidateCacheIfCredentialsFileChanged() -> Bool { + self.context.run { + let current = ClaudeOAuthCredentialsStore.currentFileFingerprint() + let stored = ClaudeOAuthCredentialsStore.loadFileFingerprint() + guard current != stored else { return false } + ClaudeOAuthCredentialsStore.log.info("Claude OAuth credentials file changed; invalidating cache") + + ClaudeOAuthCredentialsStore.writeMemoryCache(record: nil, timestamp: nil) + + var shouldClearKeychainCache = false + var shouldSaveFileFingerprint = true + if ClaudeOAuthCredentialsStore.shouldUseCodexBarOAuthKeychainCache { + if ClaudeOAuthCredentialsStore.hasPendingCodexBarOAuthKeychainCacheClear { + // The next cache access owns the deferred clear. Avoid repeated delete attempts inside one + // load and leave the fingerprint pending until that clear succeeds. + shouldSaveFileFingerprint = false + } else if let current { + if let modifiedAtMs = current.modifiedAtMs { + let modifiedAt = Date( + timeIntervalSince1970: TimeInterval(Double(modifiedAtMs) / 1000.0)) + switch ClaudeOAuthCredentialsStore.loadCodexBarOAuthKeychainCache() { + case let .found(entry): + if entry.storedAt < modifiedAt { + shouldClearKeychainCache = true + } + case .missing, .invalid: + shouldClearKeychainCache = true + case .temporarilyUnavailable: + shouldClearKeychainCache = false + shouldSaveFileFingerprint = false + } + } else { + shouldClearKeychainCache = true + } + } else { + shouldClearKeychainCache = true + } + } else { + ClaudeOAuthCredentialsStore.markPendingCodexBarOAuthKeychainCacheClear() + } + + if shouldClearKeychainCache { + ClaudeOAuthCredentialsStore.clearCacheKeychain() + } + if shouldSaveFileFingerprint { + ClaudeOAuthCredentialsStore.saveFileFingerprint(current) + } + return true + } + } + + func invalidateCache() { + self.context.run { + ClaudeOAuthCredentialsStore.writeMemoryCache(record: nil, timestamp: nil) + ClaudeOAuthCredentialsStore.clearCacheKeychain() + } + } + + func hasCachedCredentials(environment: [String: String]) -> Bool { + self.context.run { + func isRefreshableOrValid(_ record: ClaudeOAuthCredentialRecord) -> Bool { + let creds = record.credentials + if !creds.isExpired { + return true + } + switch record.owner { + case .claudeCLI: + return true + case .codexbar: + let refreshToken = creds.refreshToken?.trimmingCharacters( + in: .whitespacesAndNewlines) ?? "" + return !refreshToken.isEmpty + case .environment: + return false + } + } + + if let creds = ClaudeOAuthCredentialsStore.loadFromEnvironment(environment), + isRefreshableOrValid( + ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .environment, + source: .environment)) + { + return true + } + + let memory = ClaudeOAuthCredentialsStore.readMemoryCache() + if ClaudeOAuthCredentialsStore.shouldUseCodexBarOAuthKeychainCache, + !ClaudeOAuthCredentialsStore.hasPendingCodexBarOAuthKeychainCacheClear, + let timestamp = memory.timestamp, + let cached = memory.record, + Date().timeIntervalSince(timestamp) < ClaudeOAuthCredentialsStore.memoryCacheValidityDuration, + isRefreshableOrValid(cached) + { + return true + } + + switch ClaudeOAuthCredentialsStore.loadCodexBarOAuthKeychainCache() { + case let .found(entry): + guard let creds = try? ClaudeOAuthCredentials.parse(data: entry.data) else { return false } + let record = ClaudeOAuthCredentialRecord( + credentials: creds, + owner: entry.owner ?? .claudeCLI, + source: .cacheKeychain) + return isRefreshableOrValid(record) + case .temporarilyUnavailable: + if ClaudeOAuthCredentialsStore.hasPendingCodexBarOAuthKeychainCacheClear { + break + } + return true + case .missing, .invalid: + break + } + + if let fileData = try? ClaudeOAuthCredentialsStore.loadFromFile(), + let creds = try? ClaudeOAuthCredentials.parse(data: fileData), + isRefreshableOrValid( + ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .credentialsFile)) + { + return true + } + return false + } + } + + func hasClaudeKeychainCredentialsWithoutPrompt() -> Bool { + self.context.run { + #if os(macOS) + let mode = ClaudeOAuthKeychainPromptPreference.current() + guard ClaudeOAuthCredentialsStore.shouldAllowClaudeCodeKeychainAccess(mode: mode) else { return false } + if ClaudeOAuthCredentialsStore.loadFromClaudeKeychainViaSecurityCLIIfEnabled( + interaction: ProviderInteractionContext.current) != nil + { + return true + } + + let fallbackPromptMode = ClaudeOAuthKeychainPromptPreference.securityFrameworkFallbackMode() + guard ClaudeOAuthCredentialsStore.shouldAllowClaudeCodeKeychainAccess(mode: fallbackPromptMode) else { + return false + } + if ProviderInteractionContext.current == .background, + !ClaudeOAuthKeychainAccessGate.shouldAllowPrompt() + { + return false + } + #if DEBUG + if let store = ClaudeOAuthCredentialsStore.taskClaudeKeychainOverrideStore, + let data = store.data + { + return (try? ClaudeOAuthCredentials.parse(data: data)) != nil + } + if let data = ClaudeOAuthCredentialsStore.taskClaudeKeychainDataOverride + ?? ClaudeOAuthCredentialsStore.claudeKeychainDataOverride + { + return (try? ClaudeOAuthCredentials.parse(data: data)) != nil + } + #endif + + var query: [String: Any] = [ + kSecClass as String: kSecClassGenericPassword, + kSecAttrService as String: ClaudeOAuthCredentialsStore.claudeKeychainService, + kSecMatchLimit as String: kSecMatchLimitOne, + kSecReturnAttributes as String: true, + ] + KeychainNoUIQuery.apply(to: &query) + + let (status, _, durationMs) = ClaudeOAuthKeychainQueryTiming.copyMatching(query) + if ClaudeOAuthKeychainQueryTiming + .backoffIfSlowNoUIQuery( + durationMs, + ClaudeOAuthCredentialsStore.claudeKeychainService, + ClaudeOAuthCredentialsStore.log) + { + return false + } + switch status { + case errSecSuccess, errSecInteractionNotAllowed: + return true + case errSecUserCanceled, errSecAuthFailed, errSecNoAccessForItem: + ClaudeOAuthKeychainAccessGate.recordDenied() + return false + default: + return false + } + #else + return false + #endif + } + } + } + + private struct Recovery { + let context: CollaboratorContext + + func shouldAttemptFreshnessSyncFromClaudeKeychain(cached: ClaudeOAuthCredentialRecord) -> Bool { + guard !cached.credentials.isExpired else { return false } + guard cached.owner == .claudeCLI else { return false } + guard ClaudeOAuthCredentialsStore.keychainAccessAllowed else { return false } + + let mode = ClaudeOAuthKeychainPromptPreference.storedMode() + switch mode { + case .never: + return false + case .onlyOnUserAction: + if ProviderInteractionContext.current != .userInitiated { + if ProcessInfo.processInfo.environment["CODEXBAR_DEBUG_CLAUDE_OAUTH_FLOW"] == "1" { + ClaudeOAuthCredentialsStore.log.debug( + "Claude OAuth keychain freshness sync skipped (background)", + metadata: ["promptMode": mode.rawValue, "owner": String(describing: cached.owner)]) + } + return false + } + return true + case .always: + return true + } + } + + func syncWithClaudeKeychainIfChanged( + cached: ClaudeOAuthCredentialRecord, + respectKeychainPromptCooldown: Bool, + now: Date = Date()) -> ClaudeOAuthCredentialRecord? + { + #if os(macOS) + let mode = ClaudeOAuthKeychainPromptPreference.current() + guard ClaudeOAuthCredentialsStore.shouldAllowClaudeCodeKeychainAccess(mode: mode) else { return nil } + if ClaudeOAuthCredentialsStore.isPromptPolicyApplicable, + respectKeychainPromptCooldown, + !ClaudeOAuthKeychainAccessGate.shouldAllowPrompt(now: now) + { + return nil + } + + if ClaudeOAuthCredentialsStore.shouldShowClaudeKeychainPreAlert() { + return nil + } + + if !ClaudeOAuthCredentialsStore.shouldCheckClaudeKeychainChange(now: now) { + return nil + } + + guard let currentFingerprint = ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPrompt() + else { + return nil + } + let storedFingerprint = ClaudeOAuthCredentialsStore.loadClaudeKeychainFingerprint() + guard currentFingerprint != storedFingerprint else { return nil } + + do { + guard let data = try ClaudeOAuthCredentialsStore.loadFromClaudeKeychainNonInteractive() else { + return nil + } + guard let keychainCreds = try? ClaudeOAuthCredentials.parse(data: data) else { + ClaudeOAuthCredentialsStore.saveClaudeKeychainFingerprint(currentFingerprint) + return nil + } + ClaudeOAuthCredentialsStore.saveClaudeKeychainFingerprint(currentFingerprint) + + guard keychainCreds.accessToken != cached.credentials.accessToken else { return nil } + if keychainCreds.isExpired, !cached.credentials.isExpired { + return nil + } + + ClaudeOAuthCredentialsStore.log.info("Claude keychain credentials changed; syncing OAuth cache") + let synced = ClaudeOAuthCredentialRecord( + credentials: keychainCreds, + owner: .claudeCLI, + source: .claudeKeychain) + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: keychainCreds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: now) + ClaudeOAuthCredentialsStore.saveToCacheKeychain(data, owner: .claudeCLI) + return synced + } catch let error as ClaudeOAuthCredentialsError { + if case let .keychainError(status) = error, + status == Int(errSecUserCanceled) + || status == Int(errSecAuthFailed) + || status == Int(errSecInteractionNotAllowed) + || status == Int(errSecNoAccessForItem) + { + ClaudeOAuthKeychainAccessGate.recordDenied(now: now) + } + return nil + } catch { + return nil + } + #else + _ = cached + _ = respectKeychainPromptCooldown + _ = now + return nil + #endif + } + + func repairFromClaudeKeychainWithoutPromptIfAllowed( + now: Date, + respectKeychainPromptCooldown: Bool, + allowCacheKeychainWrite: Bool = true) -> ClaudeOAuthCredentialRecord? + { + #if os(macOS) + let mode = ClaudeOAuthKeychainPromptPreference.current() + guard ClaudeOAuthCredentialsStore.shouldAllowClaudeCodeKeychainAccess(mode: mode) else { return nil } + + if ClaudeOAuthCredentialsStore.shouldShowClaudeKeychainPreAlert() { + return nil + } + + if ClaudeOAuthCredentialsStore.isPromptPolicyApplicable, + respectKeychainPromptCooldown, + ProviderInteractionContext.current != .userInitiated, + !ClaudeOAuthKeychainAccessGate.shouldAllowPrompt(now: now) + { + return nil + } + + do { + if ClaudeOAuthCredentialsStore.shouldPreferSecurityCLIKeychainRead(), + let securityData = ClaudeOAuthCredentialsStore.loadFromClaudeKeychainViaSecurityCLIIfEnabled( + interaction: ProviderInteractionContext.current), + !securityData.isEmpty + { + guard let creds = try? ClaudeOAuthCredentials.parse(data: securityData) else { return nil } + if creds.isExpired { + return ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .claudeKeychain) + } + + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: now) + if allowCacheKeychainWrite { + ClaudeOAuthCredentialsStore.saveToCacheKeychain(securityData, owner: .claudeCLI) + } + + ClaudeOAuthCredentialsStore.log.info( + "Claude keychain credentials loaded without prompt; syncing OAuth cache", + metadata: ["interaction": ProviderInteractionContext.current == .userInitiated + ? "user" : "background"]) + return ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .claudeKeychain) + } + + guard let data = try ClaudeOAuthCredentialsStore.loadFromClaudeKeychainNonInteractive(), + !data.isEmpty + else { + return nil + } + let fingerprint = ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPrompt() + guard let creds = try? ClaudeOAuthCredentials.parse(data: data) else { + ClaudeOAuthCredentialsStore.saveClaudeKeychainFingerprint(fingerprint) + return nil + } + + if creds.isExpired { + ClaudeOAuthCredentialsStore.saveClaudeKeychainFingerprint(fingerprint) + return ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .claudeKeychain) + } + + ClaudeOAuthCredentialsStore.saveClaudeKeychainFingerprint(fingerprint) + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: now) + if allowCacheKeychainWrite { + ClaudeOAuthCredentialsStore.saveToCacheKeychain(data, owner: .claudeCLI) + } + + ClaudeOAuthCredentialsStore.log.info( + "Claude keychain credentials loaded without prompt; syncing OAuth cache", + metadata: ["interaction": ProviderInteractionContext.current == .userInitiated + ? "user" : "background"]) + return ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .claudeKeychain) + } catch let error as ClaudeOAuthCredentialsError { + if case let .keychainError(status) = error, + status == Int(errSecUserCanceled) + || status == Int(errSecAuthFailed) + || status == Int(errSecInteractionNotAllowed) + || status == Int(errSecNoAccessForItem) + { + ClaudeOAuthKeychainAccessGate.recordDenied(now: now) + } + return nil + } catch { + return nil + } + #else + _ = now + _ = respectKeychainPromptCooldown + return nil + #endif + } + + @discardableResult + func syncFromClaudeKeychainWithoutPrompt(now: Date = Date()) -> Bool { + self.context.run { + #if os(macOS) + let mode = ClaudeOAuthKeychainPromptPreference.current() + guard ClaudeOAuthCredentialsStore.shouldAllowClaudeCodeKeychainAccess(mode: mode) else { return false } + + if let data = ClaudeOAuthCredentialsStore.loadFromClaudeKeychainViaSecurityCLIIfEnabled( + interaction: ProviderInteractionContext.current), + !data.isEmpty + { + if let creds = try? ClaudeOAuthCredentials.parse(data: data), !creds.isExpired { + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: now) + ClaudeOAuthCredentialsStore.saveToCacheKeychain(data, owner: .claudeCLI) + return true + } + } + + let fallbackPromptMode = ClaudeOAuthKeychainPromptPreference.securityFrameworkFallbackMode() + guard ClaudeOAuthCredentialsStore.shouldAllowClaudeCodeKeychainAccess(mode: fallbackPromptMode) else { + return false + } + + if ProviderInteractionContext.current == .background, + !ClaudeOAuthKeychainAccessGate.shouldAllowPrompt(now: now) + { + return false + } + + #if DEBUG + let override = ClaudeOAuthCredentialsStore.taskClaudeKeychainOverrideStore?.data + ?? ClaudeOAuthCredentialsStore.taskClaudeKeychainDataOverride + ?? ClaudeOAuthCredentialsStore.claudeKeychainDataOverride + if let override, + !override.isEmpty, + let creds = try? ClaudeOAuthCredentials.parse(data: override), + !creds.isExpired + { + ClaudeOAuthCredentialsStore.saveClaudeKeychainFingerprint( + ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPrompt()) + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: now) + ClaudeOAuthCredentialsStore.saveToCacheKeychain(override, owner: .claudeCLI) + return true + } + #endif + + if ClaudeOAuthCredentialsStore.shouldShowClaudeKeychainPreAlert() { + return false + } + + if let candidate = ClaudeOAuthCredentialsStore.claudeKeychainCandidatesWithoutPrompt( + promptMode: fallbackPromptMode).first, + let data = try? ClaudeOAuthCredentialsStore.loadClaudeKeychainData( + candidate: candidate, + allowKeychainPrompt: false), + !data.isEmpty + { + let fingerprint = ClaudeKeychainFingerprint( + modifiedAt: candidate.modifiedAt.map { Int($0.timeIntervalSince1970) }, + createdAt: candidate.createdAt.map { Int($0.timeIntervalSince1970) }, + persistentRefHash: ClaudeOAuthCredentialsStore.sha256Prefix(candidate.persistentRef)) + + if let creds = try? ClaudeOAuthCredentials.parse(data: data), !creds.isExpired { + ClaudeOAuthCredentialsStore.saveClaudeKeychainFingerprint(fingerprint) + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: now) + ClaudeOAuthCredentialsStore.saveToCacheKeychain(data, owner: .claudeCLI) + return true + } + + ClaudeOAuthCredentialsStore.saveClaudeKeychainFingerprint(fingerprint) + } + + let legacyData = try? ClaudeOAuthCredentialsStore.loadClaudeKeychainLegacyData( + allowKeychainPrompt: false, + promptMode: fallbackPromptMode) + if let legacyData, + !legacyData.isEmpty, + let creds = try? ClaudeOAuthCredentials.parse(data: legacyData), + !creds.isExpired + { + ClaudeOAuthCredentialsStore.saveClaudeKeychainFingerprint( + ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPrompt()) + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: creds, + owner: .claudeCLI, + source: .memoryCache), + timestamp: now) + ClaudeOAuthCredentialsStore.saveToCacheKeychain(legacyData, owner: .claudeCLI) + return true + } + + return false + #else + _ = now + return false + #endif + } + } + } + + private struct Refresher { + let context: CollaboratorContext + + func refreshAccessToken( + refreshToken: String, + existingScopes: [String], + existingRateLimitTier: String?, + existingSubscriptionType: String? = nil, + historyOwnerIdentifier: String?) async throws -> ClaudeOAuthCredentials + { + try await self.context.run { + let newCredentials = try await self.refreshAccessTokenCore( + refreshToken: refreshToken, + existingScopes: existingScopes, + existingRateLimitTier: existingRateLimitTier, + existingSubscriptionType: existingSubscriptionType) + + ClaudeOAuthCredentialsStore.saveRefreshedCredentialsToCache( + newCredentials, + historyOwnerIdentifier: historyOwnerIdentifier) + ClaudeOAuthCredentialsStore.writeMemoryCache( + record: ClaudeOAuthCredentialRecord( + credentials: newCredentials, + owner: .codexbar, + source: .memoryCache, + historyOwnerIdentifier: historyOwnerIdentifier), + timestamp: Date()) + ClaudeOAuthRefreshFailureGate.recordSuccess() + + return newCredentials + } + } + + private func refreshAccessTokenCore( + refreshToken: String, + existingScopes: [String], + existingRateLimitTier: String?, + existingSubscriptionType: String?) async throws -> ClaudeOAuthCredentials + { + guard ClaudeOAuthRefreshFailureGate.shouldAttempt() else { + let status = ClaudeOAuthRefreshFailureGate.currentBlockStatus() + let message = switch status { + case .terminal: + "Claude OAuth refresh blocked until auth changes. \(ClaudeOAuthCredentialsStore.reauthenticateHint)" + case .transient: + "Claude OAuth refresh temporarily backed off due to prior failures; will retry automatically." + case nil: + "Claude OAuth refresh temporarily suppressed due to prior failures; will retry automatically." + } + throw ClaudeOAuthCredentialsError.refreshFailed(message) + } + + guard let url = URL(string: ClaudeOAuthCredentialsStore.tokenRefreshEndpoint) else { + throw ClaudeOAuthCredentialsError.refreshFailed("Invalid token endpoint URL") + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.timeoutInterval = 30 + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + var components = URLComponents() + components.queryItems = [ + URLQueryItem(name: "grant_type", value: "refresh_token"), + URLQueryItem(name: "refresh_token", value: refreshToken), + URLQueryItem(name: "client_id", value: ClaudeOAuthCredentialsStore.oauthClientID), + ] + request.httpBody = (components.percentEncodedQuery ?? "").data(using: .utf8) + + let response = try await ProviderHTTPClient.shared.response(for: request) + let data = response.data + guard response.statusCode == 200 else { + if let disposition = ClaudeOAuthCredentialsStore.refreshFailureDisposition( + statusCode: response.statusCode, + data: data) + { + let oauthError = ClaudeOAuthCredentialsStore.extractOAuthErrorCode(from: data) + ClaudeOAuthCredentialsStore.log.info( + "Claude OAuth refresh rejected", + metadata: [ + "httpStatus": "\(response.statusCode)", + "oauthError": oauthError ?? "nil", + "disposition": disposition.rawValue, + ]) + + switch disposition { + case .terminalInvalidGrant: + ClaudeOAuthRefreshFailureGate.recordTerminalAuthFailure() + Repository(context: self.context).invalidateCache() + let message = "HTTP \(response.statusCode) invalid_grant. " + + ClaudeOAuthCredentialsStore.reauthenticateHint + throw ClaudeOAuthCredentialsError.refreshFailed( + message) + case .transientBackoff: + ClaudeOAuthRefreshFailureGate.recordTransientFailure() + let suffix = oauthError.map { " (\($0))" } ?? "" + throw ClaudeOAuthCredentialsError.refreshFailed("HTTP \(response.statusCode)\(suffix)") + } + } + throw ClaudeOAuthCredentialsError.refreshFailed("HTTP \(response.statusCode)") + } + + let tokenResponse = try JSONDecoder().decode(TokenRefreshResponse.self, from: data) + let expiresAt = Date(timeIntervalSinceNow: TimeInterval(tokenResponse.expiresIn)) + + return ClaudeOAuthCredentials( + accessToken: tokenResponse.accessToken, + refreshToken: tokenResponse.refreshToken ?? refreshToken, + expiresAt: expiresAt, + scopes: existingScopes, + rateLimitTier: existingRateLimitTier, + subscriptionType: existingSubscriptionType) + } + } + + public static func load( + environment: [String: String] = ProcessInfo.processInfo.environment, + allowKeychainPrompt: Bool = true, + respectKeychainPromptCooldown: Bool = false) throws -> ClaudeOAuthCredentials + { + let context = self.currentCollaboratorContext() + return try Repository(context: context).load( + environment: environment, + allowKeychainPrompt: allowKeychainPrompt, + respectKeychainPromptCooldown: respectKeychainPromptCooldown) + } + + public static func loadRecord( + environment: [String: String] = ProcessInfo.processInfo.environment, + allowKeychainPrompt: Bool = true, + respectKeychainPromptCooldown: Bool = false, + allowClaudeKeychainRepairWithoutPrompt: Bool = true) throws -> ClaudeOAuthCredentialRecord + { + let context = self.currentCollaboratorContext() + return try Repository(context: context).loadRecord( + environment: environment, + allowKeychainPrompt: allowKeychainPrompt, + respectKeychainPromptCooldown: respectKeychainPromptCooldown, + allowClaudeKeychainRepairWithoutPrompt: allowClaudeKeychainRepairWithoutPrompt) + } + + /// Async version of load that handles expired tokens based on credential ownership. + /// - Claude CLI-owned credentials delegate refresh to Claude CLI. + /// - CodexBar-owned credentials refresh directly via token endpoint. + public static func loadWithAutoRefresh( + environment: [String: String] = ProcessInfo.processInfo.environment, + allowKeychainPrompt: Bool = true, + respectKeychainPromptCooldown: Bool = false) async throws -> ClaudeOAuthCredentials + { + try await self.loadRecordWithAutoRefresh( + environment: environment, + allowKeychainPrompt: allowKeychainPrompt, + respectKeychainPromptCooldown: respectKeychainPromptCooldown).credentials + } + + /// Record-preserving variant used when callers must distinguish the credential that actually won routing. + public static func loadRecordWithAutoRefresh( + environment: [String: String] = ProcessInfo.processInfo.environment, + allowKeychainPrompt: Bool = true, + respectKeychainPromptCooldown: Bool = false) async throws -> ClaudeOAuthCredentialRecord + { + let context = self.currentCollaboratorContext() + let repository = Repository(context: context) + let refresher = Refresher(context: context) + let record = try repository.loadRecord( + environment: environment, + allowKeychainPrompt: allowKeychainPrompt, + respectKeychainPromptCooldown: respectKeychainPromptCooldown, + allowClaudeKeychainRepairWithoutPrompt: true) + let credentials = record.credentials + let now = Date() + var expiryMetadata = credentials.diagnosticsMetadata(now: now) + expiryMetadata["source"] = record.source.rawValue + expiryMetadata["owner"] = record.owner.rawValue + expiryMetadata["allowKeychainPrompt"] = "\(allowKeychainPrompt)" + expiryMetadata["respectPromptCooldown"] = "\(respectKeychainPromptCooldown)" + expiryMetadata["readStrategy"] = ClaudeOAuthKeychainReadStrategyPreference.current().rawValue + + let isExpired: Bool = if let expiresAt = credentials.expiresAt { + now >= expiresAt + } else { + true + } + + // If not expired, return as-is + guard isExpired else { + self.log.debug("Claude OAuth credentials loaded for usage", metadata: expiryMetadata) + return record + } + + self.log.info("Claude OAuth credentials considered expired", metadata: expiryMetadata) + + switch record.owner { + case .claudeCLI: + if ProviderInteractionContext.current != .userInitiated, + ClaudeOAuthCredentialsStore.isMcpOAuthOnlyClaudeKeychainPayloadPresent( + interaction: ProviderInteractionContext.current, + environment: environment) + { + self.log.warning( + "Claude OAuth credentials expired; Claude keychain has MCP OAuth state only", + metadata: expiryMetadata) + throw ClaudeOAuthCredentialsError.mcpOAuthOnlyKeychain + } + self.log.info( + "Claude OAuth credentials expired; delegating refresh to Claude CLI", + metadata: expiryMetadata) + throw ClaudeOAuthCredentialsError.refreshDelegatedToClaudeCLI + case .environment: + self.log.warning("Environment OAuth token expired and cannot be auto-refreshed") + throw ClaudeOAuthCredentialsError.noRefreshToken + case .codexbar: + break + } + + // Try to refresh if we have a refresh token. + guard let refreshToken = credentials.refreshToken, !refreshToken.isEmpty else { + self.log.warning("Token expired but no refresh token available") + throw ClaudeOAuthCredentialsError.noRefreshToken + } + self.log.info("Access token expired, attempting auto-refresh") + + do { + let refreshed = try await refresher.refreshAccessToken( + refreshToken: refreshToken, + existingScopes: credentials.scopes, + existingRateLimitTier: credentials.rateLimitTier, + existingSubscriptionType: credentials.subscriptionType, + historyOwnerIdentifier: record.historyOwnerIdentifier) + self.log.info("Token refresh successful, expires in \(refreshed.expiresIn ?? 0) seconds") + return ClaudeOAuthCredentialRecord( + credentials: refreshed, + owner: .codexbar, + source: .memoryCache, + historyOwnerIdentifier: record.historyOwnerIdentifier) + } catch { + self.log.error("Token refresh failed: \(error.localizedDescription)") + throw error + } + } + + /// Save refreshed credentials to CodexBar's keychain cache + private static func saveRefreshedCredentialsToCache( + _ credentials: ClaudeOAuthCredentials, + historyOwnerIdentifier: String?) + { + var oauth: [String: Any] = [ + "accessToken": credentials.accessToken, + "expiresAt": (credentials.expiresAt?.timeIntervalSince1970 ?? 0) * 1000, + "scopes": credentials.scopes, + ] + + if let refreshToken = credentials.refreshToken { + oauth["refreshToken"] = refreshToken + } + if let rateLimitTier = credentials.rateLimitTier { + oauth["rateLimitTier"] = rateLimitTier + } + if let subscriptionType = credentials.subscriptionType { + oauth["subscriptionType"] = subscriptionType + } + + let oauthData: [String: Any] = ["claudeAiOauth": oauth] + + guard let jsonData = try? JSONSerialization.data(withJSONObject: oauthData) else { + self.log.error("Failed to serialize refreshed credentials for cache") + return + } + + self.saveToCacheKeychain( + jsonData, + owner: .codexbar, + historyOwnerIdentifier: historyOwnerIdentifier) + self.log.debug("Saved refreshed credentials to CodexBar keychain cache") + } + + /// Response from the OAuth token refresh endpoint + private struct TokenRefreshResponse: Decodable { + let accessToken: String + let refreshToken: String? + let expiresIn: Int + let tokenType: String? + + enum CodingKeys: String, CodingKey { + case accessToken = "access_token" + case refreshToken = "refresh_token" + case expiresIn = "expires_in" + case tokenType = "token_type" + } + } + + public static func loadFromFile() throws -> Data { + let url = self.credentialsFileURL() + do { + return try Data(contentsOf: url) + } catch { + if (error as NSError).code == NSFileReadNoSuchFileError { + throw ClaudeOAuthCredentialsError.notFound + } + throw ClaudeOAuthCredentialsError.readFailed(error.localizedDescription) + } + } + + public static func credentialsFileFingerprintToken() -> String? { + guard let fingerprint = self.currentFileFingerprint() else { return nil } + let modifiedAt = fingerprint.modifiedAtMs.map(String.init) ?? "nil" + return "\(modifiedAt):\(fingerprint.size)" + } + + public static func authFingerprintToken() -> String { + let file = self.credentialsFileFingerprintToken() ?? "nil" + let keychain = self.claudeKeychainFingerprintToken() ?? "nil" + return "file=\(file)|keychain=\(keychain)" + } + + public static func consumeClaudeKeychainFingerprintChangeWithoutPrompt() -> Bool { + let current: ClaudeKeychainFingerprint? + switch self.probeClaudeKeychainFingerprintWithoutPrompt() { + case .unavailable: + return false + case let .value(fingerprint): + current = fingerprint + } + let stored = self.loadClaudeKeychainFingerprint() + guard current != stored else { return false } + self.saveClaudeKeychainFingerprint(current) + return true + } + + public static func claudeKeychainFingerprintChangedWithoutConsuming() -> Bool { + let current: ClaudeKeychainFingerprint? + switch self.probeClaudeKeychainFingerprintWithoutPrompt() { + case .unavailable: + return false + case let .value(fingerprint): + current = fingerprint + } + return current != self.loadClaudeKeychainFingerprint() + } + + public static func claudeKeychainFingerprintToken() -> String? { + let fingerprint: ClaudeKeychainFingerprint? = switch self.probeClaudeKeychainFingerprintWithoutPrompt() { + case .unavailable: + self.loadClaudeKeychainFingerprint() + case let .value(probed): + probed + } + guard let fingerprint else { return nil } + let modifiedAt = fingerprint.modifiedAt.map(String.init) ?? "nil" + let createdAt = fingerprint.createdAt.map(String.init) ?? "nil" + let persistentRefHash = fingerprint.persistentRefHash ?? "nil" + return "\(modifiedAt):\(createdAt):\(persistentRefHash)" + } + + /// Returns the current Claude Code Keychain item's opaque persistent-reference hash without + /// falling back to a stored fingerprint. History ownership must prefer no identity over a + /// potentially stale identity when a non-interactive Keychain probe is unavailable. + public static func claudeKeychainPersistentRefHashWithoutPrompt() -> String? { + switch self.probeClaudeKeychainFingerprintWithoutPrompt() { + case .unavailable: + nil + case let .value(fingerprint): + fingerprint?.persistentRefHash + } + } + + /// Returns the current Keychain item's persistent-reference hash only when it owns the credential + /// that actually won OAuth routing. Token material is compared in memory and is never hashed or persisted. + public static func matchingClaudeKeychainPersistentRefHashWithoutPrompt( + for record: ClaudeOAuthCredentialRecord) -> String? + { + self.claudeKeychainCredentialMatchWithoutPrompt(for: record).persistentRefHash + } + + static func claudeKeychainCredentialMatchWithoutPrompt( + for record: ClaudeOAuthCredentialRecord) -> ClaudeKeychainCredentialMatch + { + guard record.owner == .claudeCLI else { return .notApplicable } + let evidence: ClaudeKeychainCredentialEvidence + switch self.newestClaudeKeychainCredentialEvidenceWithoutPrompt() { + case .unavailable: + return .unavailable + case .value(nil): + return .absent + case let .value(value?): + evidence = value + } + guard evidence.credentials.accessToken == record.credentials.accessToken else { + return .mismatch + } + return .matched(persistentRefHash: evidence.persistentRefHash) + } + + private static func matchingClaudeKeychainPersistentRefHash( + for record: ClaudeOAuthCredentialRecord, + evidence: ClaudeKeychainCredentialEvidence?) -> String? + { + guard let evidence, + evidence.credentials.accessToken == record.credentials.accessToken + else { + return nil + } + return evidence.persistentRefHash + } + + private static func newestClaudeKeychainCredentialEvidenceWithoutPrompt() + -> ClaudeKeychainProbe + { + #if DEBUG + if let store = self.taskClaudeKeychainOverrideStore { + guard store.data != nil || store.fingerprint != nil else { return .value(nil) } + return self.makeClaudeKeychainCredentialEvidence( + data: store.data, + persistentRefHash: store.fingerprint?.persistentRefHash).map { .value($0) } ?? .unavailable + } + let overrideData = self.taskClaudeKeychainDataOverride ?? self.claudeKeychainDataOverride + let overrideFingerprint = self.taskClaudeKeychainFingerprintOverride ?? self.claudeKeychainFingerprintOverride + if overrideData != nil || overrideFingerprint != nil { + return self.makeClaudeKeychainCredentialEvidence( + data: overrideData, + persistentRefHash: overrideFingerprint?.persistentRefHash).map { .value($0) } ?? .unavailable + } + if self.taskSecurityCLIReadOverride != nil || self.securityCLIReadOverride != nil { + // A security(1) result cannot be bound to a persistent reference without an exact candidate read. + return .unavailable + } + #endif + + #if os(macOS) + let promptMode = ClaudeOAuthKeychainPromptPreference.current() + let newest: ClaudeKeychainCandidate? + switch self.claudeKeychainCandidatesProbeWithoutPrompt(promptMode: promptMode) { + case .unavailable: + return .unavailable + case let .value(candidates): + if let first = candidates.first { + newest = first + } else { + switch self.claudeKeychainLegacyCandidateProbeWithoutPrompt(promptMode: promptMode) { + case .unavailable: + return .unavailable + case let .value(candidate): + newest = candidate + } + } + } + guard let newest else { return .value(nil) } + guard let persistentRefHash = self.sha256Prefix(newest.persistentRef), + let data = try? self.loadClaudeKeychainData( + candidate: newest, + allowKeychainPrompt: false, + promptMode: promptMode) + else { + return .unavailable + } + guard let evidence = self.makeClaudeKeychainCredentialEvidence( + data: data, + persistentRefHash: persistentRefHash) + else { return .unavailable } + return .value(evidence) + #else + return .unavailable + #endif + } + + private static func makeClaudeKeychainCredentialEvidence( + data: Data?, + persistentRefHash: String?) -> ClaudeKeychainCredentialEvidence? + { + guard let data, + let persistentRefHash, + let credentials = try? ClaudeOAuthCredentials.parse(data: data) + else { + return nil + } + return ClaudeKeychainCredentialEvidence( + credentials: credentials, + persistentRefHash: persistentRefHash) + } + + #if DEBUG + static func _matchingClaudeKeychainPersistentRefHashForTesting( + record: ClaudeOAuthCredentialRecord, + candidateCredentials: ClaudeOAuthCredentials, + persistentRefHash: String) -> String? + { + self.matchingClaudeKeychainPersistentRefHash( + for: record, + evidence: ClaudeKeychainCredentialEvidence( + credentials: candidateCredentials, + persistentRefHash: persistentRefHash)) + } + #endif + + private enum ClaudeKeychainProbe { + case unavailable + case value(Value) + } + + @discardableResult + public static func invalidateCacheIfCredentialsFileChanged() -> Bool { + Repository(context: self.currentCollaboratorContext()).invalidateCacheIfCredentialsFileChanged() + } + + /// Invalidate the credentials cache (call after login/logout) + public static func invalidateCache() { + Repository(context: self.currentCollaboratorContext()).invalidateCache() + } + + /// Check if CodexBar has cached credentials (in memory or keychain cache) + public static func hasCachedCredentials(environment: [String: String] = ProcessInfo.processInfo + .environment) -> Bool + { + Repository(context: self.currentCollaboratorContext()).hasCachedCredentials(environment: environment) + } + + public static func hasClaudeKeychainCredentialsWithoutPrompt() -> Bool { + Repository(context: self.currentCollaboratorContext()).hasClaudeKeychainCredentialsWithoutPrompt() + } + + private static func hasClaudeKeychainItemWithoutPrompt() -> Bool { + #if DEBUG + if let store = self.taskClaudeKeychainOverrideStore { + if let data = store.data, !data.isEmpty { + return true + } + if store.fingerprint != nil { + return true + } + } + if let data = self.taskClaudeKeychainDataOverride ?? self.claudeKeychainDataOverride, + !data.isEmpty + { + return true + } + if self.taskClaudeKeychainFingerprintOverride ?? self.claudeKeychainFingerprintOverride != nil { + return true + } + #endif + + #if os(macOS) + switch self.claudeKeychainCandidatesProbeWithoutPrompt(enforcePromptPolicy: false) { + case let .value(candidates) where !candidates.isEmpty: + return true + case .value, .unavailable: + break + } + switch self.claudeKeychainLegacyCandidateProbeWithoutPrompt(enforcePromptPolicy: false) { + case let .value(candidate): + return candidate != nil + case .unavailable: + return false + } + #else + return false + #endif + } + + private static func shouldCheckClaudeKeychainChange(now: Date = Date()) -> Bool { + #if DEBUG + // Unit tests can supply TaskLocal overrides for the Claude keychain data/fingerprint. Those tests often run + // concurrently with other suites, so the global throttle becomes nondeterministic. When an override is + // present, bypass the throttle so test expectations don't depend on unrelated activity. + if self.taskClaudeKeychainOverrideStore != nil || self.taskClaudeKeychainFingerprintOverride != nil + || self.claudeKeychainFingerprintOverride != nil + { + return true + } + #endif + + self.claudeKeychainChangeCheckLock.lock() + defer { self.claudeKeychainChangeCheckLock.unlock() } + if let last = self.lastClaudeKeychainChangeCheckAt, + now.timeIntervalSince(last) < self.claudeKeychainChangeCheckMinimumInterval + { + return false + } + self.lastClaudeKeychainChangeCheckAt = now + return true + } + + private static func loadClaudeKeychainFingerprint() -> ClaudeKeychainFingerprint? { + #if DEBUG + if let store = taskClaudeKeychainFingerprintStoreOverride { + return store.fingerprint + } + #endif + // Proactively remove the legacy V1 key (it included the keychain account string, which can be identifying). + UserDefaults.standard.removeObject(forKey: self.claudeKeychainFingerprintLegacyKey) + + guard let data = UserDefaults.standard.data(forKey: self.claudeKeychainFingerprintKey) else { + return nil + } + return try? JSONDecoder().decode(ClaudeKeychainFingerprint.self, from: data) + } + + private static func saveClaudeKeychainFingerprint(_ fingerprint: ClaudeKeychainFingerprint?) { + #if DEBUG + if let store = taskClaudeKeychainFingerprintStoreOverride { + store.fingerprint = fingerprint + return + } + #endif + // Proactively remove the legacy V1 key (it included the keychain account string, which can be identifying). + UserDefaults.standard.removeObject(forKey: self.claudeKeychainFingerprintLegacyKey) + + guard let fingerprint else { + UserDefaults.standard.removeObject(forKey: self.claudeKeychainFingerprintKey) + return + } + if let data = try? JSONEncoder().encode(fingerprint) { + UserDefaults.standard.set(data, forKey: self.claudeKeychainFingerprintKey) + } + } + + private static func currentClaudeKeychainFingerprintWithoutPrompt() -> ClaudeKeychainFingerprint? { + switch self.probeClaudeKeychainFingerprintWithoutPrompt() { + case .unavailable: + nil + case let .value(fingerprint): + fingerprint + } + } + + private static func probeClaudeKeychainFingerprintWithoutPrompt() + -> ClaudeKeychainProbe { + let mode = ClaudeOAuthKeychainPromptPreference.current() + #if DEBUG + if let store = taskClaudeKeychainOverrideStore { + return .value(store.fingerprint) + } + if let override = taskClaudeKeychainFingerprintOverride ?? self + .claudeKeychainFingerprintOverride + { + return .value(override) + } + #endif + guard self.shouldAllowClaudeCodeKeychainAccess(mode: mode) else { return .unavailable } + if self.isPromptPolicyApplicable, + ProviderInteractionContext.current == .background, + !ClaudeOAuthKeychainAccessGate.shouldAllowPrompt() + { + return .unavailable + } + #if os(macOS) + let candidatesProbe = self.claudeKeychainCandidatesProbeWithoutPrompt(promptMode: mode) + let newest: ClaudeKeychainCandidate? + switch candidatesProbe { + case .unavailable: + return .unavailable + case let .value(candidates): + if let first = candidates.first { + newest = first + } else { + switch self.claudeKeychainLegacyCandidateProbeWithoutPrompt(promptMode: mode) { + case .unavailable: + return .unavailable + case let .value(candidate): + newest = candidate + } + } + } + guard let newest else { return .value(nil) } + + let modifiedAt = newest.modifiedAt.map { Int($0.timeIntervalSince1970) } + let createdAt = newest.createdAt.map { Int($0.timeIntervalSince1970) } + let persistentRefHash = Self.sha256Prefix(newest.persistentRef) + return .value(ClaudeKeychainFingerprint( + modifiedAt: modifiedAt, + createdAt: createdAt, + persistentRefHash: persistentRefHash)) + #else + return .unavailable + #endif + } + + static func currentClaudeKeychainFingerprintWithoutPromptForAuthGate() -> ClaudeKeychainFingerprint? { + self.currentClaudeKeychainFingerprintWithoutPrompt() + } + + static func currentCredentialsFileFingerprintWithoutPromptForAuthGate() -> String? { + guard let fingerprint = self.currentFileFingerprint() else { return nil } + let modifiedAt = fingerprint.modifiedAtMs ?? 0 + return "\(modifiedAt):\(fingerprint.size)" + } + + private static func loadFromClaudeKeychainNonInteractive( + readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current()) + throws -> Data? + { + #if os(macOS) + let fallbackPromptMode = ClaudeOAuthKeychainPromptPreference.securityFrameworkFallbackMode( + readStrategy: readStrategy) + if let data = self.loadFromClaudeKeychainViaSecurityCLIIfEnabled( + interaction: ProviderInteractionContext.current, + readStrategy: readStrategy) + { + return data + } + + // For experimental strategy, enforce stored prompt policy before any Security.framework fallback probes. + guard self.shouldAllowClaudeCodeKeychainAccess(mode: fallbackPromptMode) else { return nil } + + #if DEBUG + if let store = taskClaudeKeychainOverrideStore { + return store.data + } + if let override = taskClaudeKeychainDataOverride ?? self.claudeKeychainDataOverride { + return override + } + #endif + + // Keep semantics aligned with fingerprinting: if there are multiple entries, we only ever consult the newest + // candidate (same as currentClaudeKeychainFingerprintWithoutPrompt()) to avoid syncing from a different item. + let candidates = self.claudeKeychainCandidatesWithoutPrompt(promptMode: fallbackPromptMode) + if let newest = candidates.first { + if let data = try self.loadClaudeKeychainData(candidate: newest, allowKeychainPrompt: false), + !data.isEmpty + { + return data + } + return nil + } + + let legacyData = try self.loadClaudeKeychainLegacyData( + allowKeychainPrompt: false, + promptMode: fallbackPromptMode) + if let legacyData, !legacyData.isEmpty { + return legacyData + } + return nil + #else + return nil + #endif + } + + static func readRawClaudeKeychainPayloadViaSecurityFrameworkWithoutPrompt() -> Data? { + #if os(macOS) + guard self.keychainAccessAllowed else { return nil } + #if DEBUG + if let store = self.taskClaudeKeychainOverrideStore { + return store.data + } + if let override = self.taskClaudeKeychainDataOverride ?? self.claudeKeychainDataOverride { + return override + } + #endif + + // This probe must work under the default `onlyOnUserAction` policy, but must never show Keychain UI. + // The candidate and data queries both use KeychainNoUIQuery; `.always` only bypasses the prompt-policy gate. + switch self.claudeKeychainCandidatesProbeWithoutPrompt( + promptMode: .always, + enforcePromptPolicy: false) + { + case .unavailable: + return nil + case let .value(candidates): + if let newest = candidates.first { + return try? self.loadClaudeKeychainData( + candidate: newest, + allowKeychainPrompt: false, + promptMode: .always) + } + } + return try? self.loadClaudeKeychainLegacyData( + allowKeychainPrompt: false, + promptMode: .always) + #else + return nil + #endif + } + + public static func loadFromClaudeKeychain() throws -> Data { + guard self.shouldAllowClaudeCodeKeychainAccess(mode: ClaudeOAuthKeychainPromptPreference.current()) else { + throw ClaudeOAuthCredentialsError.notFound + } + #if DEBUG + if let store = taskClaudeKeychainOverrideStore, let override = store.data { + return override + } + if let override = taskClaudeKeychainDataOverride ?? self.claudeKeychainDataOverride { + return override + } + #endif + if let data = self.loadFromClaudeKeychainViaSecurityCLIIfEnabled( + interaction: ProviderInteractionContext.current) + { + return data + } + if self.shouldPreferSecurityCLIKeychainRead() { + let fallbackPromptMode = ClaudeOAuthKeychainPromptPreference.securityFrameworkFallbackMode() + let fallbackDecision = self.securityFrameworkFallbackPromptDecision( + promptMode: fallbackPromptMode, + allowKeychainPrompt: true, + respectKeychainPromptCooldown: false) + self.log.debug( + "Claude keychain Security.framework fallback prompt policy evaluated", + metadata: [ + "reader": "securityFrameworkFallback", + "fallbackPromptMode": fallbackPromptMode.rawValue, + "fallbackPromptAllowed": "\(fallbackDecision.allowed)", + "fallbackBlockedReason": fallbackDecision.blockedReason ?? "none", + ]) + guard fallbackDecision.allowed else { + throw ClaudeOAuthCredentialsError.notFound + } + return try self.loadFromClaudeKeychainUsingSecurityFramework( + promptMode: fallbackPromptMode, + allowKeychainPrompt: true) + } + return try self.loadFromClaudeKeychainUsingSecurityFramework() + } + + /// Legacy alias for backward compatibility + public static func loadFromKeychain() throws -> Data { + try self.loadFromClaudeKeychain() + } + + private static func loadFromClaudeKeychainUsingSecurityFramework( + promptMode: ClaudeOAuthKeychainPromptMode = ClaudeOAuthKeychainPromptPreference.current(), + allowKeychainPrompt: Bool = true) throws -> Data + { + #if DEBUG + if let store = taskClaudeKeychainOverrideStore, let override = store.data { + return override + } + if let override = taskClaudeKeychainDataOverride ?? self.claudeKeychainDataOverride { + return override + } + #endif + #if os(macOS) + let candidates = self.claudeKeychainCandidatesWithoutPrompt(promptMode: promptMode) + if let newest = candidates.first { + do { + if let data = try self.loadClaudeKeychainData( + candidate: newest, + allowKeychainPrompt: allowKeychainPrompt, + promptMode: promptMode), + !data.isEmpty + { + // Store fingerprint after a successful interactive read so we don't immediately try to + // "sync" in the background (which can still show UI on some systems). + let modifiedAt = newest.modifiedAt.map { Int($0.timeIntervalSince1970) } + let createdAt = newest.createdAt.map { Int($0.timeIntervalSince1970) } + let persistentRefHash = Self.sha256Prefix(newest.persistentRef) + self.saveClaudeKeychainFingerprint( + ClaudeKeychainFingerprint( + modifiedAt: modifiedAt, + createdAt: createdAt, + persistentRefHash: persistentRefHash)) + return data + } + } catch let error as ClaudeOAuthCredentialsError { + if case .keychainError = error { + ClaudeOAuthKeychainAccessGate.recordDenied() + } + throw error + } + } + + // Fallback: legacy query (may pick an arbitrary duplicate). + do { + if let data = try self.loadClaudeKeychainLegacyData( + allowKeychainPrompt: allowKeychainPrompt, + promptMode: promptMode), + !data.isEmpty + { + // Same as above: store fingerprint after interactive read to avoid background "sync" reads. + self.saveClaudeKeychainFingerprint(self.currentClaudeKeychainFingerprintWithoutPrompt()) + return data + } + } catch let error as ClaudeOAuthCredentialsError { + if case .keychainError = error { + ClaudeOAuthKeychainAccessGate.recordDenied() + } + throw error + } + throw ClaudeOAuthCredentialsError.notFound + #else + throw ClaudeOAuthCredentialsError.notFound + #endif + } + + #if os(macOS) + private struct ClaudeKeychainCandidate { + let persistentRef: Data + let account: String? + let modifiedAt: Date? + let createdAt: Date? + } + + private static func claudeKeychainCandidatesProbeWithoutPrompt( + promptMode: ClaudeOAuthKeychainPromptMode = ClaudeOAuthKeychainPromptPreference + .current(), + enforcePromptPolicy: Bool = true) -> ClaudeKeychainProbe<[ClaudeKeychainCandidate]> + { + if enforcePromptPolicy { + guard self.shouldAllowClaudeCodeKeychainAccess(mode: promptMode) else { return .unavailable } + if self.isPromptPolicyApplicable, + ProviderInteractionContext.current == .background, + !ClaudeOAuthKeychainAccessGate.shouldAllowPrompt() + { + return .unavailable + } + } else { + guard self.keychainAccessAllowed else { return .unavailable } + } + var query: [String: Any] = [ + kSecClass as String: kSecClassGenericPassword, + kSecAttrService as String: self.claudeKeychainService, + kSecMatchLimit as String: kSecMatchLimitAll, + kSecReturnAttributes as String: true, + kSecReturnPersistentRef as String: true, + ] + KeychainNoUIQuery.apply(to: &query) + + let (status, result, durationMs) = ClaudeOAuthKeychainQueryTiming.copyMatching(query) + if ClaudeOAuthKeychainQueryTiming + .backoffIfSlowNoUIQuery(durationMs, self.claudeKeychainService, self.log) + { + return .unavailable + } + if status == errSecUserCanceled || status == errSecAuthFailed || status == errSecNoAccessForItem { + ClaudeOAuthKeychainAccessGate.recordDenied() + } + if status == errSecItemNotFound { + return .value([]) + } + guard status == errSecSuccess else { return .unavailable } + guard let rows = result as? [[String: Any]], !rows.isEmpty else { return .value([]) } + + let candidates: [ClaudeKeychainCandidate] = rows.compactMap { row in + guard let persistentRef = row[kSecValuePersistentRef as String] as? Data else { return nil } + return ClaudeKeychainCandidate( + persistentRef: persistentRef, + account: row[kSecAttrAccount as String] as? String, + modifiedAt: row[kSecAttrModificationDate as String] as? Date, + createdAt: row[kSecAttrCreationDate as String] as? Date) + } + + let sorted = candidates.sorted { lhs, rhs in + let lhsDate = lhs.modifiedAt ?? lhs.createdAt ?? Date.distantPast + let rhsDate = rhs.modifiedAt ?? rhs.createdAt ?? Date.distantPast + return lhsDate > rhsDate + } + return .value(sorted) + } + + private static func claudeKeychainCandidatesWithoutPrompt( + promptMode: ClaudeOAuthKeychainPromptMode = ClaudeOAuthKeychainPromptPreference + .current()) -> [ClaudeKeychainCandidate] + { + switch self.claudeKeychainCandidatesProbeWithoutPrompt(promptMode: promptMode) { + case .unavailable: + [] + case let .value(candidates): + candidates + } + } + + private static func claudeKeychainLegacyCandidateProbeWithoutPrompt( + promptMode: ClaudeOAuthKeychainPromptMode = ClaudeOAuthKeychainPromptPreference + .current(), + enforcePromptPolicy: Bool = true) -> ClaudeKeychainProbe + { + if enforcePromptPolicy { + guard self.shouldAllowClaudeCodeKeychainAccess(mode: promptMode) else { return .unavailable } + if self.isPromptPolicyApplicable, + ProviderInteractionContext.current == .background, + !ClaudeOAuthKeychainAccessGate.shouldAllowPrompt() + { + return .unavailable + } + } else { + guard self.keychainAccessAllowed else { return .unavailable } + } + var query: [String: Any] = [ + kSecClass as String: kSecClassGenericPassword, + kSecAttrService as String: self.claudeKeychainService, + kSecMatchLimit as String: kSecMatchLimitOne, + kSecReturnAttributes as String: true, + kSecReturnPersistentRef as String: true, + ] + KeychainNoUIQuery.apply(to: &query) + + let (status, result, durationMs) = ClaudeOAuthKeychainQueryTiming.copyMatching(query) + if ClaudeOAuthKeychainQueryTiming + .backoffIfSlowNoUIQuery(durationMs, self.claudeKeychainService, self.log) + { + return .unavailable + } + if status == errSecUserCanceled || status == errSecAuthFailed || status == errSecNoAccessForItem { + ClaudeOAuthKeychainAccessGate.recordDenied() + } + if status == errSecItemNotFound { + return .value(nil) + } + guard status == errSecSuccess else { return .unavailable } + guard let row = result as? [String: Any] else { return .value(nil) } + guard let persistentRef = row[kSecValuePersistentRef as String] as? Data else { return .value(nil) } + return .value(ClaudeKeychainCandidate( + persistentRef: persistentRef, + account: row[kSecAttrAccount as String] as? String, + modifiedAt: row[kSecAttrModificationDate as String] as? Date, + createdAt: row[kSecAttrCreationDate as String] as? Date)) + } + + private static func claudeKeychainLegacyCandidateWithoutPrompt( + promptMode: ClaudeOAuthKeychainPromptMode = ClaudeOAuthKeychainPromptPreference + .current()) -> ClaudeKeychainCandidate? + { + switch self.claudeKeychainLegacyCandidateProbeWithoutPrompt(promptMode: promptMode) { + case .unavailable: + nil + case let .value(candidate): + candidate + } + } + + private static func loadClaudeKeychainData( + candidate: ClaudeKeychainCandidate, + allowKeychainPrompt: Bool, + promptMode: ClaudeOAuthKeychainPromptMode = ClaudeOAuthKeychainPromptPreference.current()) throws -> Data? + { + guard self.shouldAllowClaudeCodeKeychainAccess(mode: promptMode) else { return nil } + self.log.debug( + "Claude keychain data read start", + metadata: [ + "service": self.claudeKeychainService, + "interactive": "\(allowKeychainPrompt)", + "process": ProcessInfo.processInfo.processName, + ]) + + var query: [String: Any] = [ + kSecClass as String: kSecClassGenericPassword, + kSecValuePersistentRef as String: candidate.persistentRef, + kSecMatchLimit as String: kSecMatchLimitOne, + kSecReturnData as String: true, + ] + + if !allowKeychainPrompt { + KeychainNoUIQuery.apply(to: &query) + } + + var result: AnyObject? + let startedAtNs = DispatchTime.now().uptimeNanoseconds + let status = KeychainSecurity.copyMatching(query as CFDictionary, &result) + let durationMs = Double(DispatchTime.now().uptimeNanoseconds - startedAtNs) / 1_000_000.0 + self.log.debug( + "Claude keychain data read result", + metadata: [ + "service": self.claudeKeychainService, + "interactive": "\(allowKeychainPrompt)", + "status": "\(status)", + "duration_ms": String(format: "%.2f", durationMs), + "process": ProcessInfo.processInfo.processName, + ]) + switch status { + case errSecSuccess: + if let data = result as? Data { + return data + } + return nil + case errSecItemNotFound: + return nil + case errSecInteractionNotAllowed: + if allowKeychainPrompt { + ClaudeOAuthKeychainAccessGate.recordDenied() + throw ClaudeOAuthCredentialsError.keychainError(Int(status)) + } + return nil + case errSecUserCanceled, errSecAuthFailed: + ClaudeOAuthKeychainAccessGate.recordDenied() + throw ClaudeOAuthCredentialsError.keychainError(Int(status)) + case errSecNoAccessForItem: + ClaudeOAuthKeychainAccessGate.recordDenied() + throw ClaudeOAuthCredentialsError.keychainError(Int(status)) + default: + throw ClaudeOAuthCredentialsError.keychainError(Int(status)) + } + } + + private static func loadClaudeKeychainLegacyData( + allowKeychainPrompt: Bool, + promptMode: ClaudeOAuthKeychainPromptMode = ClaudeOAuthKeychainPromptPreference.current()) throws -> Data? + { + guard self.shouldAllowClaudeCodeKeychainAccess(mode: promptMode) else { return nil } + self.log.debug( + "Claude keychain legacy data read start", + metadata: [ + "service": self.claudeKeychainService, + "interactive": "\(allowKeychainPrompt)", + "process": ProcessInfo.processInfo.processName, + ]) + + var query: [String: Any] = [ + kSecClass as String: kSecClassGenericPassword, + kSecAttrService as String: self.claudeKeychainService, + kSecMatchLimit as String: kSecMatchLimitOne, + kSecReturnData as String: true, + ] + + if !allowKeychainPrompt { + KeychainNoUIQuery.apply(to: &query) + } + + var result: AnyObject? + let startedAtNs = DispatchTime.now().uptimeNanoseconds + let status = KeychainSecurity.copyMatching(query as CFDictionary, &result) + let durationMs = Double(DispatchTime.now().uptimeNanoseconds - startedAtNs) / 1_000_000.0 + self.log.debug( + "Claude keychain legacy data read result", + metadata: [ + "service": self.claudeKeychainService, + "interactive": "\(allowKeychainPrompt)", + "status": "\(status)", + "duration_ms": String(format: "%.2f", durationMs), + "process": ProcessInfo.processInfo.processName, + ]) + switch status { + case errSecSuccess: + return result as? Data + case errSecItemNotFound: + return nil + case errSecInteractionNotAllowed: + if allowKeychainPrompt { + ClaudeOAuthKeychainAccessGate.recordDenied() + throw ClaudeOAuthCredentialsError.keychainError(Int(status)) + } + return nil + case errSecUserCanceled, errSecAuthFailed: + ClaudeOAuthKeychainAccessGate.recordDenied() + throw ClaudeOAuthCredentialsError.keychainError(Int(status)) + case errSecNoAccessForItem: + ClaudeOAuthKeychainAccessGate.recordDenied() + throw ClaudeOAuthCredentialsError.keychainError(Int(status)) + default: + throw ClaudeOAuthCredentialsError.keychainError(Int(status)) + } + } + #endif + + private static func loadFromEnvironment(_ environment: [String: String]) + -> ClaudeOAuthCredentials? + { + guard + let token = environment[self.environmentTokenKey]?.trimmingCharacters( + in: .whitespacesAndNewlines), + !token.isEmpty + else { + return nil + } + + let scopes: [String] = { + guard let raw = environment[self.environmentScopesKey] else { return ["user:profile"] } + let parsed = + raw + .split(separator: ",") + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + .filter { !$0.isEmpty } + return parsed.isEmpty ? ["user:profile"] : parsed + }() + + return ClaudeOAuthCredentials( + accessToken: token, + refreshToken: nil, + expiresAt: Date.distantFuture, + scopes: scopes, + rateLimitTier: nil) + } + + static func setCredentialsURLOverrideForTesting(_ url: URL?) { + self.credentialsURLOverride = url + } + + #if DEBUG + public static func withCredentialsURLOverrideForTesting(_ url: URL?, operation: () throws -> T) rethrows -> T { + try self.$taskCredentialsURLOverride.withValue(url) { + try operation() + } + } + + public static func withCredentialsURLOverrideForTesting(_ url: URL?, operation: () async throws -> T) + async rethrows -> T { + try await self.$taskCredentialsURLOverride.withValue(url) { + try await operation() + } + } + + public static var currentCredentialsURLOverrideForTesting: URL? { + self.taskCredentialsURLOverride + } + #endif + + private static func saveToCacheKeychain( + _ data: Data, + owner: ClaudeOAuthCredentialOwner? = nil, + historyOwnerIdentifier: String? = nil) + { + guard self.shouldUseCodexBarOAuthKeychainCache else { + self.markPendingCodexBarOAuthKeychainCacheClear() + return + } + let entry = CacheEntry( + data: data, + storedAt: Date(), + owner: owner, + historyOwnerIdentifier: historyOwnerIdentifier) + self.currentPendingCodexBarOAuthKeychainCacheClearStore.withCacheTransaction { pending in + if pending { + switch KeychainCacheStore.clearResult(key: self.cacheKey) { + case .removed, .missing: + pending = false + case .failed: + break + } + } + pending = !KeychainCacheStore.storeResult(key: self.cacheKey, entry: entry) + } + } + + private static func clearCacheKeychain() { + if self.shouldUseCodexBarOAuthKeychainCache { + self.currentPendingCodexBarOAuthKeychainCacheClearStore.withCacheTransaction { pending in + switch KeychainCacheStore.clearResult(key: self.cacheKey) { + case .removed, .missing: + pending = false + case .failed: + pending = true + } + } + } else { + self.markPendingCodexBarOAuthKeychainCacheClear() + } + } + + private static func loadCodexBarOAuthKeychainCache() -> KeychainCacheStore.LoadResult { + guard self.shouldUseCodexBarOAuthKeychainCache else { return .missing } + var result: KeychainCacheStore.LoadResult = .temporarilyUnavailable + self.currentPendingCodexBarOAuthKeychainCacheClearStore.withCacheTransaction { pending in + if pending { + switch KeychainCacheStore.clearResult(key: self.cacheKey) { + case .removed, .missing: + pending = false + case .failed: + return + } + } + result = KeychainCacheStore.load(key: self.cacheKey, as: CacheEntry.self) + } + return result + } + + private static var shouldUseCodexBarOAuthKeychainCache: Bool { + ClaudeOAuthKeychainPromptPreference.storedMode() != .never + } + + private static func markPendingCodexBarOAuthKeychainCacheClear() { + self.currentPendingCodexBarOAuthKeychainCacheClearStore.markPending() + } + + private static var hasPendingCodexBarOAuthKeychainCacheClear: Bool { + self.currentPendingCodexBarOAuthKeychainCacheClearStore.isPending + } + + private static var currentPendingCodexBarOAuthKeychainCacheClearStore: ClaudeOAuthPendingCacheClearStore { + #if DEBUG + if let store = self.taskPendingCacheClearStoreOverride { + return store + } + #endif + return self.pendingCodexBarOAuthKeychainCacheClearStore + } + + private static var keychainAccessAllowed: Bool { + #if DEBUG + if let override = self.taskKeychainAccessOverride { + return !override + } + if KeychainAccessGate.currentOverrideForTesting == true { + return false + } + if self.hasTaskKeychainTestingOverride { + return true + } + #endif + return !KeychainAccessGate.isDisabled + } + + #if DEBUG + private static var hasTaskKeychainTestingOverride: Bool { + self.taskClaudeKeychainOverrideStore != nil + || self.taskClaudeKeychainDataOverride != nil + || self.taskClaudeKeychainFingerprintOverride != nil + || self.taskSecurityCLIReadOverride != nil + || self.taskSecurityCLIReadAccountOverride != nil + } + #endif + + private static var isPromptPolicyApplicable: Bool { + ClaudeOAuthKeychainPromptPreference.isApplicable() + } + + private static func securityFrameworkFallbackPromptDecision( + promptMode: ClaudeOAuthKeychainPromptMode, + allowKeychainPrompt: Bool, + respectKeychainPromptCooldown: Bool) -> (allowed: Bool, blockedReason: String?) + { + guard allowKeychainPrompt else { + return (allowed: false, blockedReason: "allowKeychainPromptFalse") + } + guard self.shouldAllowClaudeCodeKeychainAccess(mode: promptMode) else { + return (allowed: false, blockedReason: self.fallbackBlockedReason(promptMode: promptMode)) + } + if respectKeychainPromptCooldown, + !ClaudeOAuthKeychainAccessGate.shouldAllowPrompt() + { + return (allowed: false, blockedReason: "cooldown") + } + return (allowed: true, blockedReason: nil) + } + + private static func fallbackBlockedReason(promptMode: ClaudeOAuthKeychainPromptMode) -> String { + if !self.keychainAccessAllowed { + return "keychainDisabled" + } + switch promptMode { + case .never: + return "never" + case .onlyOnUserAction: + return "onlyOnUserAction-background" + case .always: + return "disallowed" + } + } + + private static func shouldAllowClaudeCodeKeychainAccess( + mode: ClaudeOAuthKeychainPromptMode = ClaudeOAuthKeychainPromptPreference.current()) -> Bool + { + guard self.keychainAccessAllowed else { return false } + switch mode { + case .never: return false + case .onlyOnUserAction: + return ProviderInteractionContext.current == .userInitiated || self.allowBackgroundPromptBootstrap + case .always: return true + } + } + + static func preferredClaudeKeychainAccountForSecurityCLIRead( + interaction: ProviderInteraction = ProviderInteractionContext.current) -> String? + { + // Keep the experimental background path fully on /usr/bin/security by default. + // Account pinning requires Security.framework candidate probing, so only allow it on explicit user actions. + guard interaction == .userInitiated else { return nil } + #if DEBUG + if let override = self.taskSecurityCLIReadAccountOverride { + return override + } + #endif + #if os(macOS) + let mode = ClaudeOAuthKeychainPromptPreference.current() + guard self.shouldAllowClaudeCodeKeychainAccess(mode: mode) else { return nil } + // Keep experimental mode prompt-safe: avoid Security.framework candidate probes when preflight says + // interaction is likely. + if self.shouldShowClaudeKeychainPreAlert() { + return nil + } + guard let account = self.claudeKeychainCandidatesWithoutPrompt(promptMode: mode).first?.account, + !account.isEmpty + else { + return nil + } + return account + #else + return nil + #endif + } + + private static func credentialsFileURL() -> URL { + #if DEBUG + if let override = self.taskCredentialsURLOverride { + return override + } + #endif + return self.credentialsURLOverride ?? self.defaultCredentialsURL() + } + + private static func loadFileFingerprint() -> CredentialsFileFingerprint? { + #if DEBUG + if let store = self.taskCredentialsFileFingerprintStoreOverride { + return store.load() + } + #endif + guard let data = UserDefaults.standard.data(forKey: self.fileFingerprintKey) else { + return nil + } + return try? JSONDecoder().decode(CredentialsFileFingerprint.self, from: data) + } + + private static func saveFileFingerprint(_ fingerprint: CredentialsFileFingerprint?) { + #if DEBUG + if let store = self.taskCredentialsFileFingerprintStoreOverride { + store.save(fingerprint); return + } + #endif + guard let fingerprint else { + UserDefaults.standard.removeObject(forKey: self.fileFingerprintKey) + return + } + if let data = try? JSONEncoder().encode(fingerprint) { + UserDefaults.standard.set(data, forKey: self.fileFingerprintKey) + } + } + + private static func currentFileFingerprint() -> CredentialsFileFingerprint? { + let url = self.credentialsFileURL() + guard let attrs = try? FileManager.default.attributesOfItem(atPath: url.path) else { + return nil + } + let size = (attrs[.size] as? NSNumber)?.intValue ?? 0 + let modifiedAtMs = (attrs[.modificationDate] as? Date).map { Int($0.timeIntervalSince1970 * 1000) } + return CredentialsFileFingerprint(modifiedAtMs: modifiedAtMs, size: size) + } + + #if DEBUG + static func _resetCredentialsFileTrackingForTesting() { + if let store = self.taskCredentialsFileFingerprintStoreOverride { + store.save(nil) + } else { + UserDefaults.standard.removeObject(forKey: self.fileFingerprintKey) + } + if self.taskPendingCacheClearStoreOverride != nil { + self.currentPendingCodexBarOAuthKeychainCacheClearStore.withCacheTransaction { pending in + pending = false + } + } + } + + static func _resetClaudeKeychainChangeTrackingForTesting() { + UserDefaults.standard.removeObject(forKey: self.claudeKeychainFingerprintKey) + UserDefaults.standard.removeObject(forKey: self.claudeKeychainFingerprintLegacyKey) + self.setClaudeKeychainDataOverrideForTesting(nil) + self.setClaudeKeychainFingerprintOverrideForTesting(nil) + self.claudeKeychainChangeCheckLock.lock() + self.lastClaudeKeychainChangeCheckAt = nil + self.claudeKeychainChangeCheckLock.unlock() + } + + static func _resetClaudeKeychainChangeThrottleForTesting() { + self.claudeKeychainChangeCheckLock.lock() + self.lastClaudeKeychainChangeCheckAt = nil + self.claudeKeychainChangeCheckLock.unlock() + } + #endif + + private static func defaultCredentialsURL() -> URL { + let home = FileManager.default.homeDirectoryForCurrentUser + return home.appendingPathComponent(self.credentialsPath) + } +} + +// swiftlint:enable type_body_length + +extension ClaudeOAuthCredentialsStore { + /// After delegated Claude CLI refresh, re-load the Claude keychain entry without prompting and sync it into + /// CodexBar's caches. This is used to avoid triggering a second OS keychain dialog during the OAuth retry. + @discardableResult + static func syncFromClaudeKeychainWithoutPrompt(now: Date = Date()) -> Bool { + Recovery(context: self.currentCollaboratorContext()).syncFromClaudeKeychainWithoutPrompt(now: now) + } + + private static func shouldShowClaudeKeychainPreAlert() -> Bool { + #if DEBUG + // Synthetic Claude Keychain fixtures must not fall through to the real preflight. Tests that explicitly + // override the preflight still exercise its prompt-policy branches. + if self.hasTaskKeychainTestingOverride, + !KeychainAccessPreflight.hasCheckGenericPasswordOverrideForTesting + { + return false + } + #endif + let mode = ClaudeOAuthKeychainPromptPreference.current() + guard self.shouldAllowClaudeCodeKeychainAccess(mode: mode) else { return false } + return switch KeychainAccessPreflight.checkGenericPassword(service: self.claudeKeychainService, account: nil) { + case .interactionRequired: + true + case .failure: + // If preflight fails, we can't be sure whether interaction is required (or if the preflight itself + // is impacted by a misbehaving Keychain configuration). Be conservative and show the pre-alert. + true + case .allowed, .notFound: + false + } + } + + private static func shouldNotifyClaudeKeychainPreAlert() -> Bool { + let mode = ClaudeOAuthKeychainPromptPreference.current() + guard self.shouldAllowClaudeCodeKeychainAccess(mode: mode) else { return false } + // Attribute-only preflight can report success even when reading the secret will prompt. Explicit user + // actions are rare and intentional, so always explain the read before Security.framework can show UI. + return ProviderInteractionContext.current == .userInitiated || self.shouldShowClaudeKeychainPreAlert() + } + + /// Refresh the access token using a refresh token. + /// Updates CodexBar's keychain cache with the new credentials. + public static func refreshAccessToken( + refreshToken: String, + existingScopes: [String], + existingRateLimitTier: String?, + existingSubscriptionType: String? = nil) async throws -> ClaudeOAuthCredentials + { + let historyOwnerIdentifier = ClaudeOAuthCredentials.historyOwnerIdentifier(forRefreshToken: refreshToken) + return try await Refresher(context: self.currentCollaboratorContext()).refreshAccessToken( + refreshToken: refreshToken, + existingScopes: existingScopes, + existingRateLimitTier: existingRateLimitTier, + existingSubscriptionType: existingSubscriptionType, + historyOwnerIdentifier: historyOwnerIdentifier) + } + + private enum RefreshFailureDisposition: String { + case terminalInvalidGrant + case transientBackoff + } + + private static func extractOAuthErrorCode(from data: Data) -> String? { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + return nil + } + return json["error"] as? String + } + + private static func refreshFailureDisposition(statusCode: Int, data: Data) -> RefreshFailureDisposition? { + guard statusCode == 400 || statusCode == 401 else { return nil } + if let error = self.extractOAuthErrorCode(from: data)?.lowercased(), error == "invalid_grant" { + return .terminalInvalidGrant + } + return .transientBackoff + } + + #if DEBUG + static func extractOAuthErrorCodeForTesting(from data: Data) -> String? { + self.extractOAuthErrorCode(from: data) + } + + static func refreshFailureDispositionForTesting(statusCode: Int, data: Data) -> String? { + self.refreshFailureDisposition(statusCode: statusCode, data: data)?.rawValue + } + #endif +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthDelegatedRefreshCoordinator.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthDelegatedRefreshCoordinator.swift new file mode 100644 index 0000000..15a5c42 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthDelegatedRefreshCoordinator.swift @@ -0,0 +1,580 @@ +import Foundation + +public enum ClaudeOAuthDelegatedRefreshCoordinator { + private final class AttemptStateStorage: @unchecked Sendable { + let lock = NSLock() + let persistsCooldown: Bool + var hasLoadedState = false + var lastAttemptAt: Date? + var lastCooldownInterval: TimeInterval? + var inFlightAttemptID: UInt64? + var inFlightInteraction: ProviderInteraction? + var inFlightTask: Task? + var nextAttemptID: UInt64 = 0 + + init(persistsCooldown: Bool) { + self.persistsCooldown = persistsCooldown + } + } + + public enum Outcome: Sendable, Equatable { + case skippedByCooldown + case cliUnavailable + case attemptedSucceeded + case attemptedFailed(String) + } + + private static let log = CodexBarLog.logger(LogCategories.claudeUsage) + private static let cooldownDefaultsKey = "claudeOAuthDelegatedRefreshLastAttemptAtV1" + private static let cooldownIntervalDefaultsKey = "claudeOAuthDelegatedRefreshCooldownIntervalSecondsV1" + private static let defaultCooldownInterval: TimeInterval = 60 * 5 + private static let shortCooldownInterval: TimeInterval = 20 + + private static let sharedState = AttemptStateStorage(persistsCooldown: true) + + public static func attempt( + now: Date = Date(), + timeout: TimeInterval = 8, + environment: [String: String] = ProcessInfo.processInfo.environment) async -> Outcome + { + if Task.isCancelled { + return .attemptedFailed("Cancelled.") + } + + let decision = self.inFlightDecision( + now: now, + timeout: timeout, + environment: environment, + interaction: ProviderInteractionContext.current) + #if DEBUG + if case .joinThenRetry = decision { + self.userInitiatedBackgroundJoinObserverForTesting?() + } + #endif + + switch decision { + case let .join(task): + return await task.value + case let .joinThenRetry(id, task, state): + let outcome = await task.value + self.clearInFlightTaskIfStillCurrent(id: id, state: state) + switch outcome { + case .attemptedFailed, .skippedByCooldown, .cliUnavailable: + return await self.attempt(now: now, timeout: timeout, environment: environment) + case .attemptedSucceeded: + return outcome + } + case let .start(id, task, state): + let outcome = await task.value + self.clearInFlightTaskIfStillCurrent(id: id, state: state) + return outcome + } + } + + private enum InFlightDecision { + case join(Task) + case joinThenRetry(UInt64, Task, AttemptStateStorage) + case start(UInt64, Task, AttemptStateStorage) + } + + private struct AttemptConfiguration { + let environment: [String: String] + let interaction: ProviderInteraction + let readStrategy: ClaudeOAuthKeychainReadStrategy + let keychainAccessDisabled: Bool + #if DEBUG + let cliAvailableOverride: Bool? + let touchAuthPathOverride: (@Sendable (TimeInterval, [String: String]) async throws -> Void)? + let keychainFingerprintOverride: (@Sendable () -> ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint?)? + #endif + } + + private static func inFlightDecision( + now: Date, + timeout: TimeInterval, + environment: [String: String], + interaction: ProviderInteraction) -> InFlightDecision + { + let state = self.currentStateStorage + state.lock.lock() + defer { state.lock.unlock() } + + if let existing = state.inFlightTask { + if interaction == .userInitiated, + state.inFlightInteraction != .userInitiated, + let existingID = state.inFlightAttemptID + { + return .joinThenRetry(existingID, existing, state) + } + return .join(existing) + } + + state.nextAttemptID += 1 + let attemptID = state.nextAttemptID + // Detached to avoid inheriting the caller's executor context (e.g. MainActor) and cancellation state. + #if DEBUG + let configuration = AttemptConfiguration( + environment: environment, + interaction: interaction, + readStrategy: ClaudeOAuthKeychainReadStrategyPreference.current(), + keychainAccessDisabled: KeychainAccessGate.isDisabled, + cliAvailableOverride: self.cliAvailableOverrideForTesting, + touchAuthPathOverride: self.touchAuthPathOverrideForTesting, + keychainFingerprintOverride: self.keychainFingerprintOverrideForTesting) + let securityCLIReadOverride = ClaudeOAuthCredentialsStore.currentSecurityCLIReadOverrideForTesting() + #else + let configuration = AttemptConfiguration( + environment: environment, + interaction: interaction, + readStrategy: ClaudeOAuthKeychainReadStrategyPreference.current(), + keychainAccessDisabled: KeychainAccessGate.isDisabled) + #endif + let task = Task.detached(priority: .utility) { + #if DEBUG + return await ClaudeOAuthCredentialsStore.withSecurityCLIReadOverrideForTesting(securityCLIReadOverride) { + await self.performAttempt( + now: now, + timeout: timeout, + configuration: configuration, + state: state) + } + #else + await self.performAttempt( + now: now, + timeout: timeout, + configuration: configuration, + state: state) + #endif + } + state.inFlightAttemptID = attemptID + state.inFlightInteraction = interaction + state.inFlightTask = task + return .start(attemptID, task, state) + } + + private static func performAttempt( + now: Date, + timeout: TimeInterval, + configuration: AttemptConfiguration, + state: AttemptStateStorage) async -> Outcome + { + guard self.isClaudeCLIAvailable(environment: configuration.environment, configuration: configuration) else { + self.log.info("Claude OAuth delegated refresh skipped: claude CLI unavailable") + return .cliUnavailable + } + + // Atomically reserve an attempt under the lock so concurrent callers don't race past isInCooldown() and start + // multiple touches/poll loops. + guard self.reserveAttemptIfNotInCooldown( + now: now, + bypassCooldown: configuration.interaction == .userInitiated, + state: state) + else { + self.log.debug("Claude OAuth delegated refresh skipped by cooldown") + return .skippedByCooldown + } + + if let mcpOAuthOnlyFailure = self.mcpOAuthOnlyKeychainFailureIfPresent( + interaction: configuration.interaction, + readStrategy: configuration.readStrategy, + keychainAccessDisabled: configuration.keychainAccessDisabled, + environment: configuration.environment) + { + self.recordAttempt(now: now, cooldown: self.defaultCooldownInterval, state: state) + self.log.warning( + "Claude OAuth delegated refresh skipped: Claude keychain has MCP OAuth state only", + metadata: ["readStrategy": configuration.readStrategy.rawValue]) + return .attemptedFailed(mcpOAuthOnlyFailure) + } + + let baseline = self.currentKeychainChangeObservationBaseline( + readStrategy: configuration.readStrategy, + keychainAccessDisabled: configuration.keychainAccessDisabled, + configuration: configuration) + var touchError: Error? + + do { + try await self.touchOAuthAuthPath( + timeout: timeout, + environment: configuration.environment, + configuration: configuration) + } catch { + touchError = error + } + + // "Touch succeeded" must mean we actually observed the Claude keychain entry change. + // Otherwise we end up in a long cooldown with still-expired credentials. + let changed = await self.waitForClaudeKeychainChange( + from: baseline, + readStrategy: configuration.readStrategy, + keychainAccessDisabled: configuration.keychainAccessDisabled, + configuration: configuration, + timeout: min(max(timeout, 1), 2)) + if changed { + self.recordAttempt(now: now, cooldown: self.defaultCooldownInterval, state: state) + self.log.info("Claude OAuth delegated refresh touch succeeded") + return .attemptedSucceeded + } + + self.recordAttempt(now: now, cooldown: self.shortCooldownInterval, state: state) + if let touchError { + let errorType = String(describing: type(of: touchError)) + self.log.warning( + "Claude OAuth delegated refresh touch failed", + metadata: ["errorType": errorType]) + self.log.debug("Claude OAuth delegated refresh touch error: \(touchError.localizedDescription)") + return .attemptedFailed(touchError.localizedDescription) + } + + self.log.warning("Claude OAuth delegated refresh touch did not update Claude keychain") + return .attemptedFailed("Claude keychain did not update after Claude CLI touch.") + } + + public static func isInCooldown(now: Date = Date()) -> Bool { + let state = self.currentStateStorage + state.lock.lock() + defer { state.lock.unlock() } + self.loadStateIfNeededLocked(state: state) + guard let lastAttemptAt = state.lastAttemptAt else { return false } + let cooldown = state.lastCooldownInterval ?? self.defaultCooldownInterval + return now.timeIntervalSince(lastAttemptAt) < cooldown + } + + public static func cooldownRemainingSeconds(now: Date = Date()) -> Int? { + let state = self.currentStateStorage + state.lock.lock() + defer { state.lock.unlock() } + self.loadStateIfNeededLocked(state: state) + guard let lastAttemptAt = state.lastAttemptAt else { return nil } + let cooldown = state.lastCooldownInterval ?? self.defaultCooldownInterval + let remaining = cooldown - now.timeIntervalSince(lastAttemptAt) + guard remaining > 0 else { return nil } + return Int(remaining.rounded(.up)) + } + + public static func isClaudeCLIAvailable( + environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool + { + self.isClaudeCLIAvailable( + environment: environment, + configuration: nil) + } + + private static func isClaudeCLIAvailable( + environment: [String: String], + configuration: AttemptConfiguration?) -> Bool + { + #if DEBUG + if let override = configuration?.cliAvailableOverride ?? self.cliAvailableOverrideForTesting { + return override + } + #endif + return ClaudeCLIResolver.isAvailable(environment: environment) + } + + private static func touchOAuthAuthPath( + timeout: TimeInterval, + environment: [String: String], + configuration: AttemptConfiguration?) async throws + { + #if DEBUG + if let override = configuration?.touchAuthPathOverride ?? self.touchAuthPathOverrideForTesting { + try await override(timeout, environment) + return + } + #endif + try await ClaudeStatusProbe.touchOAuthAuthPath(timeout: timeout, environment: environment) + } + + private enum KeychainChangeObservationBaseline { + case securityFramework(fingerprint: ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint?) + case securityCLI(data: Data?) + } + + private static func currentKeychainChangeObservationBaseline( + readStrategy: ClaudeOAuthKeychainReadStrategy, + keychainAccessDisabled: Bool, + configuration: AttemptConfiguration?) -> KeychainChangeObservationBaseline + { + if readStrategy == .securityCLIExperimental { + return .securityCLI(data: self.currentClaudeKeychainDataViaSecurityCLIForObservation( + readStrategy: readStrategy, + keychainAccessDisabled: keychainAccessDisabled, + interaction: .background)) + } + return .securityFramework(fingerprint: self.currentClaudeKeychainFingerprint(configuration: configuration)) + } + + private static func waitForClaudeKeychainChange( + from baseline: KeychainChangeObservationBaseline, + readStrategy: ClaudeOAuthKeychainReadStrategy, + keychainAccessDisabled: Bool, + configuration: AttemptConfiguration?, + timeout: TimeInterval) async -> Bool + { + // Prefer correctness but bound the delay. Keychain writes can be slightly delayed after the CLI touch. + // Keep this short to avoid "prompt storms" on configurations where "no UI" queries can still surface UI. + let clampedTimeout = max(0, min(timeout, 2)) + if clampedTimeout == 0 { return false } + + let delays: [TimeInterval] = [0.2, 0.5, 0.8].filter { $0 <= clampedTimeout } + let deadline = Date().addingTimeInterval(clampedTimeout) + + func isObservedChange() -> Bool { + switch baseline { + case let .securityFramework(fingerprintBefore): + // Treat "no fingerprint" as "not observed"; we only succeed if we can read a fingerprint and it + // differs. + guard let current = self.currentClaudeKeychainFingerprintForObservation(configuration: configuration) + else { + return false + } + return current != fingerprintBefore + case let .securityCLI(dataBefore): + // In experimental mode, avoid Security.framework observation entirely and detect change from + // /usr/bin/security output only. + // If baseline capture failed (nil), treat observation as inconclusive and do not infer a change from + // a later successful read. + guard let dataBefore else { return false } + guard let current = self.currentClaudeKeychainDataViaSecurityCLIForObservation( + readStrategy: readStrategy, + keychainAccessDisabled: keychainAccessDisabled, + interaction: .background) + else { return false } + return current != dataBefore + } + } + + if isObservedChange() { + return true + } + + for delay in delays { + if Date() >= deadline { break } + do { + try Task.checkCancellation() + try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000)) + } catch { + return false + } + + if isObservedChange() { + return true + } + } + + return false + } + + private static func currentClaudeKeychainFingerprint( + configuration: AttemptConfiguration?) -> ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint? + { + #if DEBUG + if let override = configuration?.keychainFingerprintOverride ?? self.keychainFingerprintOverrideForTesting { + return override() + } + #endif + return ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPromptForAuthGate() + } + + private static func currentClaudeKeychainFingerprintForObservation() -> ClaudeOAuthCredentialsStore + .ClaudeKeychainFingerprint? + { + self.currentClaudeKeychainFingerprintForObservation(configuration: nil) + } + + private static func currentClaudeKeychainFingerprintForObservation( + configuration: AttemptConfiguration?) -> ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint? + { + #if DEBUG + if let override = configuration?.keychainFingerprintOverride ?? self.keychainFingerprintOverrideForTesting { + return override() + } + #endif + + // Observation should not be blocked by the background cooldown gate; otherwise we can "false fail" even when + // the CLI refreshed successfully but we couldn't observe it due to a previous denied prompt/cooldown. + // + // This temporarily classifies the observation query as "user initiated" so it bypasses the gate that only + // applies to background probes. The query remains "no UI" and does not clear cooldown state itself. + return ProviderInteractionContext.$current.withValue(.userInitiated) { + ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPromptForAuthGate() + } + } + + private static func currentClaudeKeychainDataViaSecurityCLIForObservation( + readStrategy: ClaudeOAuthKeychainReadStrategy, + keychainAccessDisabled: Bool, + interaction: ProviderInteraction) -> Data? + { + guard !keychainAccessDisabled else { return nil } + return ClaudeOAuthCredentialsStore.readRawClaudeKeychainPayloadViaSecurityCLIIfEnabled( + interaction: interaction, + readStrategy: readStrategy) + } + + private static func mcpOAuthOnlyKeychainFailureIfPresent( + interaction: ProviderInteraction, + readStrategy: ClaudeOAuthKeychainReadStrategy, + keychainAccessDisabled: Bool, + environment: [String: String]) -> String? + { + guard interaction != .userInitiated else { return nil } + guard ClaudeOAuthCredentialsStore.isMcpOAuthOnlyClaudeKeychainPayloadPresent( + interaction: interaction, + readStrategy: readStrategy, + keychainAccessDisabled: keychainAccessDisabled, + environment: environment) + else { + return nil + } + return ClaudeOAuthCredentialsError.mcpOAuthOnlyKeychain.errorDescription + ?? "Claude keychain contains MCP OAuth state only." + } + + private static func clearInFlightTaskIfStillCurrent(id: UInt64, state: AttemptStateStorage) { + state.lock.lock() + if state.inFlightAttemptID == id { + state.inFlightAttemptID = nil + state.inFlightInteraction = nil + state.inFlightTask = nil + } + state.lock.unlock() + } + + private static func recordAttempt(now: Date, cooldown: TimeInterval, state: AttemptStateStorage) { + state.lock.lock() + defer { state.lock.unlock() } + self.loadStateIfNeededLocked(state: state) + state.lastAttemptAt = now + state.lastCooldownInterval = cooldown + guard state.persistsCooldown else { return } + UserDefaults.standard.set(now.timeIntervalSince1970, forKey: self.cooldownDefaultsKey) + UserDefaults.standard.set(cooldown, forKey: self.cooldownIntervalDefaultsKey) + } + + private static func reserveAttemptIfNotInCooldown( + now: Date, + bypassCooldown: Bool, + state: AttemptStateStorage) -> Bool + { + state.lock.lock() + defer { state.lock.unlock() } + self.loadStateIfNeededLocked(state: state) + + let cooldown = state.lastCooldownInterval ?? self.defaultCooldownInterval + if !bypassCooldown, + let lastAttemptAt = state.lastAttemptAt, + now.timeIntervalSince(lastAttemptAt) < cooldown + { + return false + } + + // Reserve with a short cooldown; the final outcome will extend or keep it short. + state.lastAttemptAt = now + state.lastCooldownInterval = self.shortCooldownInterval + guard state.persistsCooldown else { return true } + UserDefaults.standard.set(now.timeIntervalSince1970, forKey: self.cooldownDefaultsKey) + UserDefaults.standard.set(self.shortCooldownInterval, forKey: self.cooldownIntervalDefaultsKey) + return true + } + + private static func loadStateIfNeededLocked(state: AttemptStateStorage) { + guard !state.hasLoadedState else { return } + state.hasLoadedState = true + guard state.persistsCooldown else { + state.lastAttemptAt = nil + state.lastCooldownInterval = nil + return + } + guard let raw = UserDefaults.standard.object(forKey: self.cooldownDefaultsKey) as? Double else { + state.lastAttemptAt = nil + state.lastCooldownInterval = nil + return + } + state.lastAttemptAt = Date(timeIntervalSince1970: raw) + if let interval = UserDefaults.standard.object(forKey: self.cooldownIntervalDefaultsKey) as? Double { + state.lastCooldownInterval = interval + } else { + state.lastCooldownInterval = nil + } + } + + #if DEBUG + @TaskLocal private static var stateStorageForTesting: AttemptStateStorage? + @TaskLocal static var cliAvailableOverrideForTesting: Bool? + @TaskLocal static var touchAuthPathOverrideForTesting: (@Sendable ( + TimeInterval, + [String: String]) async throws -> Void)? + @TaskLocal static var keychainFingerprintOverrideForTesting: (@Sendable () -> ClaudeOAuthCredentialsStore + .ClaudeKeychainFingerprint?)? + @TaskLocal static var userInitiatedBackgroundJoinObserverForTesting: (@Sendable () -> Void)? + + static func withCLIAvailableOverrideForTesting( + _ override: Bool?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$cliAvailableOverrideForTesting.withValue(override) { + try await operation() + } + } + + static func withTouchAuthPathOverrideForTesting( + _ override: (@Sendable (TimeInterval, [String: String]) async throws -> Void)?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$touchAuthPathOverrideForTesting.withValue(override) { + try await operation() + } + } + + static func withKeychainFingerprintOverrideForTesting( + _ override: (@Sendable () -> ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint?)?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$keychainFingerprintOverrideForTesting.withValue(override) { + try await operation() + } + } + + static func withUserInitiatedBackgroundJoinObserverForTesting( + _ observer: (@Sendable () -> Void)?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$userInitiatedBackgroundJoinObserverForTesting.withValue(observer) { + try await operation() + } + } + + static func withIsolatedStateForTesting(operation: () async throws -> T) async rethrows -> T { + let state = AttemptStateStorage(persistsCooldown: false) + return try await self.$stateStorageForTesting.withValue(state) { + try await operation() + } + } + + static func resetForTesting() { + let state = self.currentStateStorage + state.lock.lock() + state.hasLoadedState = true + state.lastAttemptAt = nil + state.lastCooldownInterval = nil + state.inFlightAttemptID = nil + state.inFlightInteraction = nil + state.inFlightTask = nil + state.nextAttemptID = 0 + state.lock.unlock() + guard state.persistsCooldown else { return } + UserDefaults.standard.removeObject(forKey: self.cooldownDefaultsKey) + UserDefaults.standard.removeObject(forKey: self.cooldownIntervalDefaultsKey) + } + #endif + + private static var currentStateStorage: AttemptStateStorage { + #if DEBUG + self.stateStorageForTesting ?? self.sharedState + #else + self.sharedState + #endif + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainAccessGate.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainAccessGate.swift new file mode 100644 index 0000000..6d43c45 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainAccessGate.swift @@ -0,0 +1,184 @@ +import Foundation + +#if os(macOS) +import os.lock + +public enum ClaudeOAuthKeychainAccessGate { + private struct State { + var loaded = false + var deniedUntil: Date? + } + + private static let lock = OSAllocatedUnfairLock(initialState: State()) + private static let defaultsKey = "claudeOAuthKeychainDeniedUntil" + private static let cooldownInterval: TimeInterval = 60 * 60 * 6 + @TaskLocal private static var taskOverrideShouldAllowPromptForTesting: Bool? + #if DEBUG + public final class DeniedUntilStore: @unchecked Sendable { + public var deniedUntil: Date? + + public init() {} + } + + @TaskLocal private static var taskDeniedUntilStoreOverrideForTesting: DeniedUntilStore? + #endif + + public static func shouldAllowPrompt(now: Date = Date()) -> Bool { + guard !KeychainAccessGate.isDisabled else { return false } + if let override = self.taskOverrideShouldAllowPromptForTesting { return override } + #if DEBUG + if let store = self.taskDeniedUntilStoreOverrideForTesting { + if let deniedUntil = store.deniedUntil, deniedUntil > now { + return false + } + store.deniedUntil = nil + return true + } + #endif + return self.lock.withLock { state in + self.loadIfNeeded(&state) + if let deniedUntil = state.deniedUntil { + if deniedUntil > now { + return false + } + state.deniedUntil = nil + self.persist(state) + } + return true + } + } + + public static func recordDenied(now: Date = Date()) { + let deniedUntil = now.addingTimeInterval(self.cooldownInterval) + #if DEBUG + if let store = self.taskDeniedUntilStoreOverrideForTesting { + store.deniedUntil = deniedUntil + return + } + #endif + self.lock.withLock { state in + self.loadIfNeeded(&state) + state.deniedUntil = deniedUntil + self.persist(state) + } + } + + /// Clears the cooldown so the next attempt can proceed. Intended for user-initiated repairs. + /// - Returns: true if a cooldown was present and cleared. + public static func clearDenied(now: Date = Date()) -> Bool { + #if DEBUG + if let store = self.taskDeniedUntilStoreOverrideForTesting { + guard let deniedUntil = store.deniedUntil, deniedUntil > now else { + store.deniedUntil = nil + return false + } + store.deniedUntil = nil + return true + } + #endif + return self.lock.withLock { state in + self.loadIfNeeded(&state) + guard let deniedUntil = state.deniedUntil, deniedUntil > now else { + state.deniedUntil = nil + self.persist(state) + return false + } + state.deniedUntil = nil + self.persist(state) + return true + } + } + + #if DEBUG + static func withShouldAllowPromptOverrideForTesting( + _ value: Bool?, + operation: () throws -> T) rethrows -> T + { + try self.$taskOverrideShouldAllowPromptForTesting.withValue(value) { + try operation() + } + } + + static func withShouldAllowPromptOverrideForTesting( + _ value: Bool?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskOverrideShouldAllowPromptForTesting.withValue(value) { + try await operation() + } + } + + public static func withDeniedUntilStoreOverrideForTesting( + _ store: DeniedUntilStore?, + operation: () throws -> T) rethrows -> T + { + try self.$taskDeniedUntilStoreOverrideForTesting.withValue(store) { + try operation() + } + } + + public static func withDeniedUntilStoreOverrideForTesting( + _ store: DeniedUntilStore?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskDeniedUntilStoreOverrideForTesting.withValue(store) { + try await operation() + } + } + + public static var currentDeniedUntilStoreOverrideForTesting: DeniedUntilStore? { + self.taskDeniedUntilStoreOverrideForTesting + } + + public static func resetForTesting() { + self.lock.withLock { state in + // Keep deterministic during tests: avoid re-loading UserDefaults written by unrelated code paths. + state.loaded = true + state.deniedUntil = nil + UserDefaults.standard.removeObject(forKey: self.defaultsKey) + } + } + + public static func resetInMemoryForTesting() { + self.lock.withLock { state in + state.loaded = false + state.deniedUntil = nil + } + } + #endif + + private static func loadIfNeeded(_ state: inout State) { + guard !state.loaded else { return } + state.loaded = true + if let raw = UserDefaults.standard.object(forKey: self.defaultsKey) as? Double { + state.deniedUntil = Date(timeIntervalSince1970: raw) + } + } + + private static func persist(_ state: State) { + if let deniedUntil = state.deniedUntil { + UserDefaults.standard.set(deniedUntil.timeIntervalSince1970, forKey: self.defaultsKey) + } else { + UserDefaults.standard.removeObject(forKey: self.defaultsKey) + } + } +} +#else +public enum ClaudeOAuthKeychainAccessGate { + public static func shouldAllowPrompt(now _: Date = Date()) -> Bool { + true + } + + public static func recordDenied(now _: Date = Date()) {} + + public static func clearDenied(now _: Date = Date()) -> Bool { + false + } + + #if DEBUG + public static func resetForTesting() {} + + public static func resetInMemoryForTesting() {} + #endif +} +#endif diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainPreAlertGate.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainPreAlertGate.swift new file mode 100644 index 0000000..285947a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainPreAlertGate.swift @@ -0,0 +1,169 @@ +import Foundation + +#if os(macOS) +import os.lock + +enum ClaudeOAuthKeychainPreAlertGate { + fileprivate struct State { + var loaded = false + var acknowledgedUntil: Date? + var presentationInFlight = false + } + + private static let lock = OSAllocatedUnfairLock(initialState: State()) + private static let defaultsKey = "claudeOAuthKeychainPreAlertAcknowledgedUntilV1" + static let cooldownInterval: TimeInterval = 60 * 60 * 6 + + #if DEBUG + final class StateStore: @unchecked Sendable { + fileprivate let lock = OSAllocatedUnfairLock(initialState: State(loaded: true)) + } + + @TaskLocal private static var taskStateStoreOverrideForTesting: StateStore? + #endif + + /// Presents at most one explanatory alert and starts the cooldown only when it reaches a handler. + @discardableResult + static func presentIfNeeded( + now: Date = Date(), + completedAt: Date? = nil, + present: () -> Bool) -> Bool + { + guard self.beginPresentation(now: now) else { return false } + let wasPresented = present() + self.finishPresentation(wasPresented: wasPresented, now: completedAt ?? Date()) + return wasPresented + } + + private static func beginPresentation(now: Date) -> Bool { + #if DEBUG + if let store = self.taskStateStoreOverrideForTesting { + return store.lock.withLock { state in + self.reservePresentation(state: &state, now: now) + } + } + #endif + return self.lock.withLock { state in + self.loadIfNeeded(&state) + guard self.reservePresentation(state: &state, now: now) else { return false } + self.persist(state) + return true + } + } + + private static func finishPresentation(wasPresented: Bool, now: Date) { + #if DEBUG + if let store = self.taskStateStoreOverrideForTesting { + store.lock.withLock { state in + self.completePresentation(state: &state, wasPresented: wasPresented, now: now) + } + return + } + #endif + self.lock.withLock { state in + self.loadIfNeeded(&state) + self.completePresentation(state: &state, wasPresented: wasPresented, now: now) + self.persist(state) + } + } + + #if DEBUG + static func withStateStoreOverrideForTesting( + _ store: StateStore?, + operation: () throws -> T) rethrows -> T + { + try self.$taskStateStoreOverrideForTesting.withValue(store) { + try operation() + } + } + + static func withStateStoreOverrideForTesting( + _ store: StateStore?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskStateStoreOverrideForTesting.withValue(store) { + try await operation() + } + } + + static func resetForTesting() { + self.lock.withLock { state in + state = State(loaded: true) + UserDefaults.standard.removeObject(forKey: self.defaultsKey) + } + } + + static func resetInMemoryForTesting() { + self.lock.withLock { state in + state = State() + } + } + #endif + + private static func loadIfNeeded(_ state: inout State) { + guard !state.loaded else { return } + state.loaded = true + if let raw = UserDefaults.standard.object(forKey: self.defaultsKey) as? Double { + state.acknowledgedUntil = Date(timeIntervalSince1970: raw) + } + } + + private static func reservePresentation(state: inout State, now: Date) -> Bool { + guard !state.presentationInFlight else { return false } + if let acknowledgedUntil = state.acknowledgedUntil, acknowledgedUntil > now { + return false + } + state.acknowledgedUntil = nil + state.presentationInFlight = true + return true + } + + private static func completePresentation(state: inout State, wasPresented: Bool, now: Date) { + state.presentationInFlight = false + if wasPresented { + state.acknowledgedUntil = now.addingTimeInterval(self.cooldownInterval) + } + } + + private static func persist(_ state: State) { + if let acknowledgedUntil = state.acknowledgedUntil { + UserDefaults.standard.set(acknowledgedUntil.timeIntervalSince1970, forKey: self.defaultsKey) + } else { + UserDefaults.standard.removeObject(forKey: self.defaultsKey) + } + } +} +#else +enum ClaudeOAuthKeychainPreAlertGate { + static let cooldownInterval: TimeInterval = 60 * 60 * 6 + + #if DEBUG + final class StateStore: @unchecked Sendable {} + #endif + + @discardableResult + static func presentIfNeeded( + now _: Date = Date(), + completedAt _: Date? = nil, + present _: () -> Bool) -> Bool + { + false + } + + #if DEBUG + static func withStateStoreOverrideForTesting( + _: StateStore?, + operation: () throws -> T) rethrows -> T + { + try operation() + } + + static func withStateStoreOverrideForTesting( + _: StateStore?, + operation: () async throws -> T) async rethrows -> T + { + try await operation() + } + #endif +} +#endif diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainPromptMode.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainPromptMode.swift new file mode 100644 index 0000000..308684a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainPromptMode.swift @@ -0,0 +1,177 @@ +import Foundation + +public enum ClaudeOAuthKeychainPromptMode: String, Sendable, Codable, CaseIterable { + case never + case onlyOnUserAction + case always +} + +public enum ClaudeOAuthKeychainPromptPreference { + static let releaseApplicationDefaultsDomain = "com.steipete.codexbar" + static let debugApplicationDefaultsDomain = "com.steipete.codexbar.debug" + private static let userDefaultsKey = "claudeOAuthKeychainPromptMode" + + static var applicationDefaultsDomain: String { + self.resolveApplicationDefaultsDomain( + bundleIdentifier: Bundle.main.bundleIdentifier, + bundleURL: Bundle.main.bundleURL, + executableURL: Bundle.main.executableURL, + invocationURL: CommandLine.arguments.first.map(URL.init(fileURLWithPath:))) + } + + #if DEBUG + private final class UserDefaultsBox: @unchecked Sendable { + let value: UserDefaults + + init(_ value: UserDefaults) { + self.value = value + } + } + + @TaskLocal private static var taskOverride: ClaudeOAuthKeychainPromptMode? + @TaskLocal private static var taskApplicationUserDefaultsOverride: UserDefaultsBox? + #endif + + public static func current(userDefaults: UserDefaults? = nil) -> ClaudeOAuthKeychainPromptMode { + self.effectiveMode(userDefaults: userDefaults) + } + + public static func storedMode(userDefaults: UserDefaults? = nil) -> ClaudeOAuthKeychainPromptMode { + #if DEBUG + if let taskOverride { + return taskOverride + } + #endif + let userDefaults = userDefaults ?? self.applicationUserDefaults + if let raw = userDefaults.string(forKey: self.userDefaultsKey), + let mode = ClaudeOAuthKeychainPromptMode(rawValue: raw) + { + return mode + } + return .onlyOnUserAction + } + + public static func isApplicable( + readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current()) -> Bool + { + readStrategy == .securityFramework + } + + public static func effectiveMode( + userDefaults: UserDefaults? = nil, + readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current()) + -> ClaudeOAuthKeychainPromptMode + { + guard self.isApplicable(readStrategy: readStrategy) else { + return .always + } + return self.storedMode(userDefaults: userDefaults) + } + + public static func securityFrameworkFallbackMode( + userDefaults: UserDefaults? = nil, + readStrategy: ClaudeOAuthKeychainReadStrategy = ClaudeOAuthKeychainReadStrategyPreference.current()) + -> ClaudeOAuthKeychainPromptMode + { + if readStrategy == .securityCLIExperimental { + return self.storedMode(userDefaults: userDefaults) + } + return self.effectiveMode(userDefaults: userDefaults, readStrategy: readStrategy) + } + + static var applicationUserDefaults: UserDefaults { + #if DEBUG + if let taskApplicationUserDefaultsOverride { + return taskApplicationUserDefaultsOverride.value + } + #endif + return UserDefaults(suiteName: self.applicationDefaultsDomain) ?? .standard + } + + static func resolveApplicationDefaultsDomain( + bundleIdentifier: String?, + bundleURL: URL?, + executableURL: URL?, + invocationURL: URL?, + bundleIdentifierForApp: (URL) -> String? = { Bundle(url: $0)?.bundleIdentifier }) -> String + { + if let domain = self.defaultsDomain(forBundleIdentifier: bundleIdentifier) { + return domain + } + + let candidates = [bundleURL, executableURL, invocationURL].compactMap(\.self) + var visitedPaths = Set() + for candidate in candidates { + var current = candidate.standardizedFileURL.resolvingSymlinksInPath() + let ancestorCount = current.pathComponents.count + for _ in 0.. String? { + guard let bundleIdentifier else { return nil } + // Check debug first because its identifier is a child of the release identifier. + if bundleIdentifier == self.debugApplicationDefaultsDomain + || bundleIdentifier.hasPrefix("\(self.debugApplicationDefaultsDomain).") + { + return self.debugApplicationDefaultsDomain + } + if bundleIdentifier == self.releaseApplicationDefaultsDomain + || bundleIdentifier.hasPrefix("\(self.releaseApplicationDefaultsDomain).") + { + return self.releaseApplicationDefaultsDomain + } + return nil + } + + #if DEBUG + public static func withTaskOverrideForTesting( + _ mode: ClaudeOAuthKeychainPromptMode?, + operation: () throws -> T) rethrows -> T + { + try self.$taskOverride.withValue(mode) { + try operation() + } + } + + public static func withTaskOverrideForTesting( + _ mode: ClaudeOAuthKeychainPromptMode?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskOverride.withValue(mode) { + try await operation() + } + } + + public static var currentTaskOverrideForTesting: ClaudeOAuthKeychainPromptMode? { + self.taskOverride + } + + static func withApplicationUserDefaultsOverrideForTesting( + _ userDefaults: UserDefaults?, + operation: () throws -> T) rethrows -> T + { + try self.$taskApplicationUserDefaultsOverride.withValue(userDefaults.map(UserDefaultsBox.init)) { + try operation() + } + } + + static func withApplicationUserDefaultsOverrideForTesting( + _ userDefaults: UserDefaults?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskApplicationUserDefaultsOverride.withValue(userDefaults.map(UserDefaultsBox.init)) { + try await operation() + } + } + #endif +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainQueryTiming.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainQueryTiming.swift new file mode 100644 index 0000000..eb49938 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainQueryTiming.swift @@ -0,0 +1,31 @@ +import Dispatch +import Foundation + +#if os(macOS) +import Security + +enum ClaudeOAuthKeychainQueryTiming { + static func copyMatching(_ query: [String: Any]) -> (status: OSStatus, result: AnyObject?, durationMs: Double) { + var result: AnyObject? + let startedAtNs = DispatchTime.now().uptimeNanoseconds + let status = KeychainSecurity.copyMatching(query as CFDictionary, &result) + let durationMs = Double(DispatchTime.now().uptimeNanoseconds - startedAtNs) / 1_000_000.0 + return (status, result, durationMs) + } + + static func backoffIfSlowNoUIQuery(_ durationMs: Double, _ service: String, _ log: CodexBarLogger) -> Bool { + // Intentionally no longer treats "slow" no-UI Keychain queries as a denial. Some systems can have + // non-deterministic timing characteristics that would make this backoff too aggressive and surprising. + // + // Keep this hook so call sites can cheaply log slow queries during debugging without changing behavior. + guard ProviderInteractionContext.current == .background, durationMs > 1000 else { return false } + log.debug( + "Claude keychain no-UI query was slow", + metadata: [ + "service": service, + "duration_ms": String(format: "%.2f", durationMs), + ]) + return false + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainReadStrategy.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainReadStrategy.swift new file mode 100644 index 0000000..057f3ab --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthKeychainReadStrategy.swift @@ -0,0 +1,49 @@ +import Foundation + +public enum ClaudeOAuthKeychainReadStrategy: String, Sendable, Codable, CaseIterable { + case securityFramework + case securityCLIExperimental +} + +public enum ClaudeOAuthKeychainReadStrategyPreference { + private static let userDefaultsKey = "claudeOAuthKeychainReadStrategy" + + #if DEBUG + @TaskLocal private static var taskOverride: ClaudeOAuthKeychainReadStrategy? + #endif + + public static func current(userDefaults: UserDefaults = .standard) -> ClaudeOAuthKeychainReadStrategy { + #if DEBUG + if let taskOverride { return taskOverride } + #endif + if let raw = userDefaults.string(forKey: self.userDefaultsKey) { + let strategy = ClaudeOAuthKeychainReadStrategy(rawValue: raw) ?? .securityFramework + return strategy == .securityCLIExperimental ? .securityFramework : strategy + } + return .securityFramework + } + + #if DEBUG + public static func withTaskOverrideForTesting( + _ strategy: ClaudeOAuthKeychainReadStrategy?, + operation: () throws -> T) rethrows -> T + { + try self.$taskOverride.withValue(strategy) { + try operation() + } + } + + public static func withTaskOverrideForTesting( + _ strategy: ClaudeOAuthKeychainReadStrategy?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskOverride.withValue(strategy) { + try await operation() + } + } + + public static var currentTaskOverrideForTesting: ClaudeOAuthKeychainReadStrategy? { + self.taskOverride + } + #endif +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthMutableKeychainOverrides.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthMutableKeychainOverrides.swift new file mode 100644 index 0000000..d1b53f9 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthMutableKeychainOverrides.swift @@ -0,0 +1,35 @@ +import Foundation + +#if DEBUG +extension ClaudeOAuthCredentialsStore { + final class ClaudeKeychainOverrideStore: @unchecked Sendable { + var data: Data? + var fingerprint: ClaudeKeychainFingerprint? + + init(data: Data? = nil, fingerprint: ClaudeKeychainFingerprint? = nil) { + self.data = data + self.fingerprint = fingerprint + } + } + + @TaskLocal static var taskClaudeKeychainOverrideStore: ClaudeKeychainOverrideStore? + + static func withMutableClaudeKeychainOverrideStoreForTesting( + _ store: ClaudeKeychainOverrideStore?, + operation: () throws -> T) rethrows -> T + { + try self.$taskClaudeKeychainOverrideStore.withValue(store) { + try operation() + } + } + + static func withMutableClaudeKeychainOverrideStoreForTesting( + _ store: ClaudeKeychainOverrideStore?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskClaudeKeychainOverrideStore.withValue(store) { + try await operation() + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthPendingCacheClearStore.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthPendingCacheClearStore.swift new file mode 100644 index 0000000..d9c84f8 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthPendingCacheClearStore.swift @@ -0,0 +1,142 @@ +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif +import Foundation + +protocol ClaudeOAuthPendingCacheClearStore: Sendable { + var isPending: Bool { get } + + func markPending() + func withCacheTransaction(_ operation: (inout Bool) -> Void) +} + +final class ClaudeOAuthPendingCacheClearUserDefaultsStore: ClaudeOAuthPendingCacheClearStore, @unchecked Sendable { + private static let processLock = NSLock() + private static let log = CodexBarLog.logger(LogCategories.claudeUsage) + + private let domain: String + private let key: String + private let lockURL: URL + + init( + domain: String, + key: String, + lockURL: URL = URL(fileURLWithPath: NSHomeDirectory(), isDirectory: true) + .appendingPathComponent(".codexbar", isDirectory: true) + .appendingPathComponent("claude-oauth-cache.lock")) + { + self.domain = domain + self.key = key + self.lockURL = lockURL + } + + var isPending: Bool { + do { + return try self.withInterprocessLock { + self.currentGeneration() != nil + } + } catch { + Self.log.error("Claude OAuth cache tombstone lock failed: \(error.localizedDescription)") + return true + } + } + + func markPending() { + do { + try self.withInterprocessLock { + self.writeGeneration(UUID().uuidString) + } + } catch { + // A surviving tombstone is safer than allowing a stale cache read. If the lock itself is unavailable, + // write a fresh generation but never clear one through the unlocked fallback path. + Self.log.error("Claude OAuth cache tombstone lock failed: \(error.localizedDescription)") + self.writeGeneration(UUID().uuidString) + } + } + + func withCacheTransaction(_ operation: (inout Bool) -> Void) { + do { + try self.withInterprocessLock { + let initialGeneration = self.currentGeneration() + var pending = initialGeneration != nil + operation(&pending) + self.persist( + pending: pending, + initialGeneration: initialGeneration) + } + } catch { + Self.log.error("Claude OAuth cache transaction lock failed: \(error.localizedDescription)") + // Fail closed: without the shared lock, do not touch the cache and leave a fresh invalidation marker. + self.writeGeneration(UUID().uuidString) + } + } + + private func persist( + pending: Bool, + initialGeneration: String?) + { + let currentGeneration = self.currentGeneration() + if pending { + if currentGeneration == nil { + self.writeGeneration(UUID().uuidString) + } + return + } + + // Compare the observed generation before removal. This is defensive against writers that do not yet honor + // the lock, while the lock serializes all current app and bundled-CLI cache mutations. + guard currentGeneration == initialGeneration else { return } + self.writeGeneration(nil) + } + + private func currentGeneration() -> String? { + let userDefaults = UserDefaults(suiteName: self.domain) ?? .standard + userDefaults.synchronize() + if let generation = userDefaults.string(forKey: self.key), !generation.isEmpty { + return generation + } + // V1 stored a boolean. Preserve an outstanding invalidation across the generation-based upgrade. + if userDefaults.object(forKey: self.key) as? Bool == true { + return "legacy-boolean" + } + return nil + } + + private func writeGeneration(_ generation: String?) { + let userDefaults = UserDefaults(suiteName: self.domain) ?? .standard + if let generation { + userDefaults.set(generation, forKey: self.key) + } else { + userDefaults.removeObject(forKey: self.key) + } + userDefaults.synchronize() + } + + private func withInterprocessLock(_ operation: () throws -> T) throws -> T { + Self.processLock.lock() + defer { Self.processLock.unlock() } + + try FileManager.default.createDirectory( + at: self.lockURL.deletingLastPathComponent(), + withIntermediateDirectories: true) + let fd = open(self.lockURL.path, O_CREAT | O_RDWR | O_CLOEXEC, S_IRUSR | S_IWUSR) + guard fd >= 0 else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + defer { + _ = flock(fd, LOCK_UN) + close(fd) + } + + while flock(fd, LOCK_EX) != 0 { + guard errno == EINTR else { + throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO) + } + } + return try operation() + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthRefreshFailureGate.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthRefreshFailureGate.swift new file mode 100644 index 0000000..e9caa25 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthRefreshFailureGate.swift @@ -0,0 +1,371 @@ +import Foundation + +#if os(macOS) +import os.lock + +public enum ClaudeOAuthRefreshFailureGate { + public enum BlockStatus: Equatable, Sendable { + case terminal(reason: String?, failures: Int) + case transient(until: Date, failures: Int) + } + + struct AuthFingerprint: Codable, Equatable { + let keychain: ClaudeOAuthCredentialsStore.ClaudeKeychainFingerprint? + let credentialsFile: String? + } + + private struct State { + var loaded = false + var terminalFailureCount = 0 + var transientFailureCount = 0 + var isTerminalBlocked = false + var transientBlockedUntil: Date? + var fingerprintAtFailure: AuthFingerprint? + var lastCredentialsRecheckAt: Date? + var terminalReason: String? + } + + private static let lock = OSAllocatedUnfairLock(initialState: State()) + private static let blockedUntilKey = "claudeOAuthRefreshBackoffBlockedUntilV1" // legacy (migration) + private static let failureCountKey = "claudeOAuthRefreshBackoffFailureCountV1" // legacy + terminal count + private static let fingerprintKey = "claudeOAuthRefreshBackoffFingerprintV2" + private static let terminalBlockedKey = "claudeOAuthRefreshTerminalBlockedV1" + private static let terminalReasonKey = "claudeOAuthRefreshTerminalReasonV1" + private static let transientBlockedUntilKey = "claudeOAuthRefreshTransientBlockedUntilV1" + private static let transientFailureCountKey = "claudeOAuthRefreshTransientFailureCountV1" + + private static let log = CodexBarLog.logger(LogCategories.claudeUsage) + private static let minimumCredentialsRecheckInterval: TimeInterval = 15 + private static let unknownFingerprint = AuthFingerprint(keychain: nil, credentialsFile: nil) + private static let transientBaseInterval: TimeInterval = 60 * 5 + private static let transientMaxInterval: TimeInterval = 60 * 60 * 6 + + #if DEBUG + @TaskLocal static var shouldAttemptOverride: Bool? + private nonisolated(unsafe) static var fingerprintProviderOverride: (() -> AuthFingerprint?)? + + static func setFingerprintProviderOverrideForTesting(_ provider: (() -> AuthFingerprint?)?) { + self.fingerprintProviderOverride = provider + } + + public static func resetInMemoryStateForTesting() { + self.lock.withLock { state in + state.loaded = false + state.terminalFailureCount = 0 + state.transientFailureCount = 0 + state.isTerminalBlocked = false + state.transientBlockedUntil = nil + state.fingerprintAtFailure = nil + state.lastCredentialsRecheckAt = nil + state.terminalReason = nil + } + } + + public static func resetForTesting() { + self.lock.withLock { state in + state.loaded = false + state.terminalFailureCount = 0 + state.transientFailureCount = 0 + state.isTerminalBlocked = false + state.transientBlockedUntil = nil + state.fingerprintAtFailure = nil + state.lastCredentialsRecheckAt = nil + state.terminalReason = nil + UserDefaults.standard.removeObject(forKey: self.blockedUntilKey) + UserDefaults.standard.removeObject(forKey: self.failureCountKey) + UserDefaults.standard.removeObject(forKey: self.fingerprintKey) + UserDefaults.standard.removeObject(forKey: self.terminalBlockedKey) + UserDefaults.standard.removeObject(forKey: self.terminalReasonKey) + UserDefaults.standard.removeObject(forKey: self.transientBlockedUntilKey) + UserDefaults.standard.removeObject(forKey: self.transientFailureCountKey) + } + } + #endif + + public static func shouldAttempt(now: Date = Date()) -> Bool { + #if DEBUG + if let override = self.shouldAttemptOverride { return override } + #endif + + return self.lock.withLock { state in + let didMigrate = self.loadIfNeeded(&state, now: now) + if didMigrate { + self.persist(state) + } + + if state.isTerminalBlocked { + guard self.shouldRecheckCredentials(now: now, state: state) else { return false } + + state.lastCredentialsRecheckAt = now + if self.hasCredentialsChangedSinceFailure(state) { + self.resetState(&state) + self.persist(state) + return true + } + + self.log.debug( + "Claude OAuth refresh blocked until auth changes", + metadata: [ + "terminalFailures": "\(state.terminalFailureCount)", + "reason": state.terminalReason ?? "nil", + ]) + return false + } + + if let blockedUntil = state.transientBlockedUntil { + if blockedUntil <= now { + self.clearTransientState(&state) + // Once transient backoff expires, forget its auth baseline so future failures capture fresh + // fingerprints and so we don't ratchet backoff across unrelated intermittent failures. + state.fingerprintAtFailure = nil + state.lastCredentialsRecheckAt = nil + self.persist(state) + return true + } + + if self.shouldRecheckCredentials(now: now, state: state) { + state.lastCredentialsRecheckAt = now + if self.hasCredentialsChangedSinceFailure(state) { + self.resetState(&state) + self.persist(state) + return true + } + } + + self.log.debug( + "Claude OAuth refresh transient backoff active", + metadata: [ + "until": "\(blockedUntil.timeIntervalSince1970)", + "transientFailures": "\(state.transientFailureCount)", + ]) + return false + } + + return true + } + } + + public static func currentBlockStatus(now: Date = Date()) -> BlockStatus? { + self.lock.withLock { state in + _ = self.loadIfNeeded(&state, now: now) + if state.isTerminalBlocked { + return .terminal(reason: state.terminalReason, failures: state.terminalFailureCount) + } + if let blockedUntil = state.transientBlockedUntil, blockedUntil > now { + return .transient(until: blockedUntil, failures: state.transientFailureCount) + } + return nil + } + } + + public static func recordTerminalAuthFailure(now: Date = Date()) { + self.lock.withLock { state in + _ = self.loadIfNeeded(&state, now: now) + state.terminalFailureCount += 1 + state.isTerminalBlocked = true + state.terminalReason = "invalid_grant" + state.fingerprintAtFailure = self.currentFingerprint() ?? self.unknownFingerprint + state.lastCredentialsRecheckAt = now + self.clearTransientState(&state) + self.persist(state) + } + } + + public static func recordTransientFailure(now: Date = Date()) { + self.lock.withLock { state in + _ = self.loadIfNeeded(&state, now: now) + + // Keep terminal blocking monotonic: once we know auth is rejected (e.g. invalid_grant), + // do not downgrade it to time-based backoff unless auth changes (fingerprint) or we record success. + guard !state.isTerminalBlocked else { return } + + self.clearTerminalState(&state) + + state.transientFailureCount += 1 + let interval = self.transientCooldownInterval(failures: state.transientFailureCount) + state.transientBlockedUntil = now.addingTimeInterval(interval) + state.fingerprintAtFailure = self.currentFingerprint() ?? self.unknownFingerprint + state.lastCredentialsRecheckAt = now + self.persist(state) + } + } + + public static func recordAuthFailure(now: Date = Date()) { + // Legacy shim: treat as terminal auth failure. + self.recordTerminalAuthFailure(now: now) + } + + public static func recordSuccess() { + self.lock.withLock { state in + _ = self.loadIfNeeded(&state, now: Date()) + self.resetState(&state) + self.persist(state) + } + } + + private static func shouldRecheckCredentials(now: Date, state: State) -> Bool { + guard let last = state.lastCredentialsRecheckAt else { return true } + return now.timeIntervalSince(last) >= self.minimumCredentialsRecheckInterval + } + + private static func hasCredentialsChangedSinceFailure(_ state: State) -> Bool { + guard let current = self.currentFingerprint() else { return false } + guard let prior = state.fingerprintAtFailure else { return false } + return current != prior + } + + private static func currentFingerprint() -> AuthFingerprint? { + #if DEBUG + if let override = self.fingerprintProviderOverride { return override() } + #endif + return AuthFingerprint( + keychain: ClaudeOAuthCredentialsStore.currentClaudeKeychainFingerprintWithoutPromptForAuthGate(), + credentialsFile: ClaudeOAuthCredentialsStore.currentCredentialsFileFingerprintWithoutPromptForAuthGate()) + } + + private static func loadIfNeeded(_ state: inout State, now: Date) -> Bool { + state.loaded = true + var didMutate = false + + // Always refresh persisted fields from UserDefaults, even after first load. + // + // This avoids stale state when UserDefaults are modified while the app is running (or during tests), + // while still keeping ephemeral throttling state (like lastCredentialsRecheckAt) in memory. + state.terminalFailureCount = UserDefaults.standard.integer(forKey: self.failureCountKey) + state.transientFailureCount = UserDefaults.standard.integer(forKey: self.transientFailureCountKey) + + if let raw = UserDefaults.standard.object(forKey: self.transientBlockedUntilKey) as? Double { + state.transientBlockedUntil = Date(timeIntervalSince1970: raw) + } + + let legacyBlockedUntil = (UserDefaults.standard.object(forKey: self.blockedUntilKey) as? Double) + .map { Date(timeIntervalSince1970: $0) } + let legacyFailureCount = UserDefaults.standard.integer(forKey: self.failureCountKey) + + if let data = UserDefaults.standard.data(forKey: self.fingerprintKey) { + state.fingerprintAtFailure = (try? JSONDecoder().decode(AuthFingerprint.self, from: data)) + } else { + state.fingerprintAtFailure = nil + } + + if UserDefaults.standard.object(forKey: self.terminalBlockedKey) != nil { + state.isTerminalBlocked = UserDefaults.standard.bool(forKey: self.terminalBlockedKey) + state.terminalReason = UserDefaults.standard.string(forKey: self.terminalReasonKey) + if legacyBlockedUntil != nil { + didMutate = true + } + } else { + // Migration: legacy keys represented a time-based backoff. Migrate to transient backoff (never terminal) + // unless we already have new transient keys persisted. + if UserDefaults.standard.object(forKey: self.transientFailureCountKey) == nil, + UserDefaults.standard.object(forKey: self.transientBlockedUntilKey) == nil, + legacyBlockedUntil != nil || legacyFailureCount > 0 + { + state.isTerminalBlocked = false + state.terminalReason = nil + state.terminalFailureCount = 0 + + if let legacyBlockedUntil, legacyBlockedUntil > now { + state.transientFailureCount = max(legacyFailureCount, 0) + state.transientBlockedUntil = legacyBlockedUntil + } else { + state.transientFailureCount = 0 + state.transientBlockedUntil = nil + } + didMutate = true + } + } + + if state.isTerminalBlocked || state.transientBlockedUntil != nil, state.fingerprintAtFailure == nil { + state.fingerprintAtFailure = self.unknownFingerprint + didMutate = true + } + + if legacyBlockedUntil != nil { + didMutate = true + } + + return didMutate + } + + private static func persist(_ state: State) { + UserDefaults.standard.set(state.terminalFailureCount, forKey: self.failureCountKey) + UserDefaults.standard.set(state.isTerminalBlocked, forKey: self.terminalBlockedKey) + if let reason = state.terminalReason { + UserDefaults.standard.set(reason, forKey: self.terminalReasonKey) + } else { + UserDefaults.standard.removeObject(forKey: self.terminalReasonKey) + } + + UserDefaults.standard.set(state.transientFailureCount, forKey: self.transientFailureCountKey) + if let blockedUntil = state.transientBlockedUntil { + UserDefaults.standard.set(blockedUntil.timeIntervalSince1970, forKey: self.transientBlockedUntilKey) + } else { + UserDefaults.standard.removeObject(forKey: self.transientBlockedUntilKey) + } + + UserDefaults.standard.removeObject(forKey: self.blockedUntilKey) + + if let fingerprint = state.fingerprintAtFailure, + let data = try? JSONEncoder().encode(fingerprint) + { + UserDefaults.standard.set(data, forKey: self.fingerprintKey) + } else { + UserDefaults.standard.removeObject(forKey: self.fingerprintKey) + } + } + + private static func transientCooldownInterval(failures: Int) -> TimeInterval { + guard failures > 0 else { return 0 } + let factor = pow(2.0, Double(failures - 1)) + return min(self.transientBaseInterval * factor, self.transientMaxInterval) + } + + private static func clearTerminalState(_ state: inout State) { + state.terminalFailureCount = 0 + state.isTerminalBlocked = false + state.terminalReason = nil + } + + private static func clearTransientState(_ state: inout State) { + state.transientFailureCount = 0 + state.transientBlockedUntil = nil + } + + private static func resetState(_ state: inout State) { + self.clearTerminalState(&state) + self.clearTransientState(&state) + state.fingerprintAtFailure = nil + state.lastCredentialsRecheckAt = nil + } +} +#else +public enum ClaudeOAuthRefreshFailureGate { + public enum BlockStatus: Equatable, Sendable { + case terminal(reason: String?, failures: Int) + case transient(until: Date, failures: Int) + } + + public static func shouldAttempt(now _: Date = Date()) -> Bool { + true + } + + public static func currentBlockStatus(now _: Date = Date()) -> BlockStatus? { + nil + } + + public static func recordTerminalAuthFailure(now _: Date = Date()) {} + + public static func recordTransientFailure(now _: Date = Date()) {} + + public static func recordAuthFailure(now _: Date = Date()) {} + + public static func recordSuccess() {} + + #if DEBUG + static func setFingerprintProviderOverrideForTesting(_: (() -> Any?)?) {} + public static func resetInMemoryStateForTesting() {} + public static func resetForTesting() {} + #endif +} +#endif diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthUsageFetcher.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthUsageFetcher.swift new file mode 100644 index 0000000..ca7e64d --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthUsageFetcher.swift @@ -0,0 +1,324 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum ClaudeOAuthFetchError: LocalizedError, Sendable { + case unauthorized + case rateLimited(retryAfter: Date?) + case invalidResponse + case serverError(Int, String?) + case networkError(Error) + + public var errorDescription: String? { + switch self { + case .unauthorized: + return "Claude OAuth request unauthorized. Run `claude` to re-authenticate." + case .rateLimited: + return "Claude OAuth usage endpoint is rate limited by Anthropic right now. Wait a few minutes, " + + "then click Refresh. If it keeps happening, run `claude logout && claude login`, then try again." + case .invalidResponse: + return "Claude OAuth response was invalid." + case let .serverError(code, body): + if let body, !body.isEmpty { + let cleaned = body + .replacingOccurrences(of: "\n", with: " ") + .trimmingCharacters(in: .whitespacesAndNewlines) + let shortened = cleaned.count > 400 ? String(cleaned.prefix(400)) + "…" : cleaned + return "Claude OAuth error: HTTP \(code) – \(shortened)" + } + return "Claude OAuth error: HTTP \(code)" + case let .networkError(error): + return "Claude OAuth network error: \(error.localizedDescription)" + } + } +} + +enum ClaudeOAuthUsageFetcher { + private static let baseURL = "https://api.anthropic.com" + private static let usagePath = "/api/oauth/usage" + private static let betaHeader = "oauth-2025-04-20" + private static let fallbackClaudeCodeVersion = "2.1.0" + + static func fetchUsage( + accessToken: String, + detectClaudeVersion: Bool = true) async throws -> OAuthUsageResponse + { + if let blockedUntil = ClaudeOAuthUsageRateLimitGate.blockedUntil() { + throw ClaudeOAuthFetchError.rateLimited(retryAfter: blockedUntil) + } + + guard let url = URL(string: baseURL + usagePath) else { + throw ClaudeOAuthFetchError.invalidResponse + } + + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.timeoutInterval = 30 + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + // OAuth usage endpoint currently requires the beta header. + request.setValue(Self.betaHeader, forHTTPHeaderField: "anthropic-beta") + request.setValue( + Self.claudeCodeUserAgent(detectClaudeVersion: detectClaudeVersion), + forHTTPHeaderField: "User-Agent") + + do { + let response = try await ProviderHTTPClient.shared.response(for: request) + let data = response.data + switch response.statusCode { + case 200: + let usage = try Self.decodeUsageResponse(data) + ClaudeOAuthUsageRateLimitGate.recordSuccess() + return usage + case 401: + throw ClaudeOAuthFetchError.unauthorized + case 429: + let retryAfter = Self.retryAfterDate(from: response.response) + ClaudeOAuthUsageRateLimitGate.recordRateLimit(retryAfter: retryAfter) + throw ClaudeOAuthFetchError.rateLimited( + retryAfter: ClaudeOAuthUsageRateLimitGate.currentBlockedUntil() ?? retryAfter) + case 403: + let body = String(data: data, encoding: .utf8) + throw ClaudeOAuthFetchError.serverError(response.statusCode, body) + default: + let body = String(data: data, encoding: .utf8) + throw ClaudeOAuthFetchError.serverError(response.statusCode, body) + } + } catch let error as ClaudeOAuthFetchError { + throw error + } catch { + throw ClaudeOAuthFetchError.networkError(error) + } + } + + static func decodeUsageResponse(_ data: Data) throws -> OAuthUsageResponse { + let decoder = JSONDecoder() + return try decoder.decode(OAuthUsageResponse.self, from: data) + } + + static func parseISO8601Date(_ string: String?) -> Date? { + guard let string, !string.isEmpty else { return nil } + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = formatter.date(from: string) { return date } + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: string) + } + + private static func retryAfterDate(from response: HTTPURLResponse, now: Date = Date()) -> Date? { + guard let raw = response.value(forHTTPHeaderField: "Retry-After")? + .trimmingCharacters(in: .whitespacesAndNewlines), + !raw.isEmpty + else { return nil } + + if let seconds = TimeInterval(raw), seconds >= 0 { + return now.addingTimeInterval(seconds) + } + + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = TimeZone(secondsFromGMT: 0) + formatter.dateFormat = "EEE',' dd MMM yyyy HH':'mm':'ss zzz" + return formatter.date(from: raw) + } + + private static func claudeCodeUserAgent( + detectClaudeVersion: Bool, + versionDetector: () -> String? = { ProviderVersionDetector.claudeVersion() }) -> String + { + self.claudeCodeUserAgent(versionString: detectClaudeVersion ? versionDetector() : nil) + } + + private static func claudeCodeUserAgent(versionString: String?) -> String { + let version = self.normalizedClaudeCodeVersion(versionString) ?? self.fallbackClaudeCodeVersion + return "claude-code/\(version)" + } + + private static func normalizedClaudeCodeVersion(_ versionString: String?) -> String? { + guard let raw = versionString?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + let token = raw.split(whereSeparator: \.isWhitespace).first.map(String.init) ?? raw + let trimmed = token.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + } +} + +struct OAuthUsageResponse: Decodable { + let fiveHour: OAuthUsageWindow? + let sevenDay: OAuthUsageWindow? + let sevenDayOAuthApps: OAuthUsageWindow? + let sevenDayOpus: OAuthUsageWindow? + let sevenDaySonnet: OAuthUsageWindow? + let sevenDayRoutines: OAuthUsageWindow? + let sevenDayRoutinesSourceKey: String? + let iguanaNecktie: OAuthUsageWindow? + let extraUsage: OAuthExtraUsage? + /// Newer shape (superseding the flat `seven_day_*` fields above for scoped weekly + /// windows): a flat list of limit entries, each optionally naming the model it scopes + /// to via `scope.model.display_name` (e.g. "Fable" during a promotional access window). + let limits: [OAuthLimitEntry]? + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: DynamicCodingKey.self) + self.fiveHour = Self.decodeWindow(in: container, keys: ["five_hour"]) + self.sevenDay = Self.decodeWindow(in: container, keys: ["seven_day"]) + self.sevenDayOAuthApps = Self.decodeWindow(in: container, keys: ["seven_day_oauth_apps"]) + self.sevenDayOpus = Self.decodeWindow(in: container, keys: ["seven_day_opus"]) + self.sevenDaySonnet = Self.decodeWindow(in: container, keys: ["seven_day_sonnet"]) + let routines = Self.decodeWindowWithSource(in: container, keys: [ + "seven_day_routines", + "seven_day_claude_routines", + "claude_routines", + "routines", + "routine", + "seven_day_cowork", + "cowork", + ]) + self.sevenDayRoutines = routines.window + self.sevenDayRoutinesSourceKey = routines.sourceKey + self.iguanaNecktie = Self.decodeWindow(in: container, keys: ["iguana_necktie"]) + self.extraUsage = Self.decodeValue(in: container, keys: ["extra_usage"]) + self.limits = Self.decodeValue(in: container, keys: ["limits"]) + } + + private static func decodeWindow( + in container: KeyedDecodingContainer, + keys: [String]) -> OAuthUsageWindow? + { + self.decodeValue(in: container, keys: keys) + } + + private static func decodeWindowWithSource( + in container: KeyedDecodingContainer, + keys: [String]) -> (window: OAuthUsageWindow?, sourceKey: String?) + { + var firstNullKey: String? + for keyName in keys { + guard let key = DynamicCodingKey(stringValue: keyName) else { continue } + guard container.contains(key) else { continue } + if let value = try? container.decodeIfPresent(OAuthUsageWindow.self, forKey: key) { + return (value, keyName) + } + if firstNullKey == nil { + firstNullKey = keyName + } + } + return (nil, firstNullKey) + } + + private static func decodeValue( + in container: KeyedDecodingContainer, + keys: [String]) -> T? + { + for keyName in keys { + guard let key = DynamicCodingKey(stringValue: keyName) else { continue } + if let value = try? container.decodeIfPresent(T.self, forKey: key) { + return value + } + } + return nil + } +} + +private struct DynamicCodingKey: CodingKey { + let stringValue: String + let intValue: Int? + + init?(stringValue: String) { + self.stringValue = stringValue + self.intValue = nil + } + + init?(intValue: Int) { + nil + } +} + +struct OAuthUsageWindow: Decodable { + let utilization: Double? + let resetsAt: String? + + enum CodingKeys: String, CodingKey { + case utilization + case resetsAt = "resets_at" + } +} + +/// A single entry from the `limits` array. `kind`/`group` classify the limit +/// (e.g. `kind: "weekly_scoped"`, `group: "weekly"`); `scope.model.display_name` names the +/// model it applies to when the limit is scoped to one, rather than the whole account. +struct OAuthLimitEntry: Decodable { + let kind: String? + let group: String? + let percent: Double? + let resetsAt: String? + let scope: OAuthLimitScope? + let isActive: Bool? + + enum CodingKeys: String, CodingKey { + case kind + case group + case percent + case resetsAt = "resets_at" + case scope + case isActive = "is_active" + } +} + +struct OAuthLimitScope: Decodable { + let model: OAuthLimitScopeModel? +} + +struct OAuthLimitScopeModel: Decodable { + let id: String? + let displayName: String? + + enum CodingKeys: String, CodingKey { + case id + case displayName = "display_name" + } +} + +struct OAuthExtraUsage: Decodable { + let isEnabled: Bool? + let monthlyLimit: Double? + let usedCredits: Double? + let utilization: Double? + let currency: String? + + enum CodingKeys: String, CodingKey { + case isEnabled = "is_enabled" + case monthlyLimit = "monthly_limit" + case usedCredits = "used_credits" + case utilization + case currency + } +} + +#if DEBUG +extension ClaudeOAuthUsageFetcher { + static func _decodeUsageResponseForTesting(_ data: Data) throws -> OAuthUsageResponse { + try self.decodeUsageResponse(data) + } + + static func _userAgentForTesting(versionString: String?) -> String { + self.claudeCodeUserAgent(versionString: versionString) + } + + static func _userAgentForTesting( + detectClaudeVersion: Bool, + versionDetector: () -> String?) -> String + { + self.claudeCodeUserAgent( + detectClaudeVersion: detectClaudeVersion, + versionDetector: versionDetector) + } + + static func _retryAfterDateForTesting(from response: HTTPURLResponse, now: Date) -> Date? { + self.retryAfterDate(from: response, now: now) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthUsageRateLimitGate.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthUsageRateLimitGate.swift new file mode 100644 index 0000000..01ee9ec --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeOAuth/ClaudeOAuthUsageRateLimitGate.swift @@ -0,0 +1,46 @@ +import Foundation + +enum ClaudeOAuthUsageRateLimitGate { + private static let blockedUntilKey = "claudeOAuthUsageRateLimitBlockedUntilV1" + private static let defaultCooldown: TimeInterval = 60 * 5 + + static func blockedUntil( + interaction: ProviderInteraction = ProviderInteractionContext.current, + now: Date = Date()) -> Date? + { + guard interaction != .userInitiated else { return nil } + return self.currentBlockedUntil(now: now) + } + + static func currentBlockedUntil(now: Date = Date()) -> Date? { + guard let raw = UserDefaults.standard.object(forKey: self.blockedUntilKey) as? Double else { + return nil + } + + let blockedUntil = Date(timeIntervalSince1970: raw) + guard blockedUntil > now else { + UserDefaults.standard.removeObject(forKey: self.blockedUntilKey) + return nil + } + return blockedUntil + } + + static func recordRateLimit(retryAfter: Date?, now: Date = Date()) { + let blockedUntil = if let retryAfter, retryAfter > now { + retryAfter + } else { + now.addingTimeInterval(self.defaultCooldown) + } + UserDefaults.standard.set(blockedUntil.timeIntervalSince1970, forKey: self.blockedUntilKey) + } + + static func recordSuccess() { + UserDefaults.standard.removeObject(forKey: self.blockedUntilKey) + } + + #if DEBUG + static func resetForTesting() { + UserDefaults.standard.removeObject(forKey: self.blockedUntilKey) + } + #endif +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudePlan.swift b/Sources/CodexBarCore/Providers/Claude/ClaudePlan.swift new file mode 100644 index 0000000..2f58530 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudePlan.swift @@ -0,0 +1,179 @@ +import Foundation + +public enum ClaudePlan: String, CaseIterable, Sendable { + case max + case pro + case team + case enterprise + case ultra + + public var brandedLoginMethod: String { + switch self { + case .max: + "Claude Max" + case .pro: + "Claude Pro" + case .team: + "Claude Team" + case .enterprise: + "Claude Enterprise" + case .ultra: + "Claude Ultra" + } + } + + /// Branded label including the Max usage multiplier when the rate-limit tier carries one. + public func brandedLoginMethod(rateLimitTier: String?) -> String { + guard self == .max, let multiplier = Self.maxUsageMultiplier(rateLimitTier: rateLimitTier) else { + return self.brandedLoginMethod + } + return "\(self.brandedLoginMethod) \(multiplier)" + } + + public var compactLoginMethod: String { + switch self { + case .max: + "Max" + case .pro: + "Pro" + case .team: + "Team" + case .enterprise: + "Enterprise" + case .ultra: + "Ultra" + } + } + + public var countsAsSubscription: Bool { + switch self { + case .max, .pro, .team, .ultra: + true + case .enterprise: + false + } + } + + public static func fromOAuthRateLimitTier(_ rateLimitTier: String?) -> Self? { + self.fromRateLimitTier(rateLimitTier) + } + + public static func fromOAuthCredentials(subscriptionType: String?, rateLimitTier: String?) -> Self? { + self.fromCompatibilityLoginMethod(subscriptionType) + ?? self.fromOAuthRateLimitTier(rateLimitTier) + } + + public static func fromWebAccount(rateLimitTier: String?, billingType: String?) -> Self? { + if let plan = self.fromRateLimitTier(rateLimitTier) { + return plan + } + + let tier = Self.normalized(rateLimitTier) + let billing = Self.normalized(billingType) + if billing.contains("stripe"), tier.contains("claude") { + return .pro + } + return nil + } + + public static func fromCompatibilityLoginMethod(_ loginMethod: String?) -> Self? { + let words = Self.normalizedWords(loginMethod) + if words.isEmpty { + return nil + } + if words.contains("max") { + return .max + } + if words.contains("pro") { + return .pro + } + if words.contains("team") { + return .team + } + if words.contains("enterprise") { + return .enterprise + } + if words.contains("ultra") { + return .ultra + } + return nil + } + + public static func oauthLoginMethod(rateLimitTier: String?) -> String? { + self.fromOAuthRateLimitTier(rateLimitTier)?.brandedLoginMethod(rateLimitTier: rateLimitTier) + } + + public static func oauthLoginMethod(subscriptionType: String?, rateLimitTier: String?) -> String? { + self.fromOAuthCredentials( + subscriptionType: subscriptionType, + rateLimitTier: rateLimitTier)?.brandedLoginMethod(rateLimitTier: rateLimitTier) + } + + public static func webLoginMethod(rateLimitTier: String?, billingType: String?) -> String? { + self.fromWebAccount( + rateLimitTier: rateLimitTier, + billingType: billingType)?.brandedLoginMethod(rateLimitTier: rateLimitTier) + } + + public static func cliCompatibilityLoginMethod(_ loginMethod: String?) -> String? { + guard let loginMethod = loginMethod?.trimmingCharacters(in: .whitespacesAndNewlines), + !loginMethod.isEmpty + else { + return nil + } + + if let plan = self.fromCompatibilityLoginMethod(loginMethod) { + return plan.compactLoginMethod + } + + return loginMethod + } + + public static func isSubscriptionLoginMethod(_ loginMethod: String?) -> Bool { + self.fromCompatibilityLoginMethod(loginMethod)?.countsAsSubscription ?? false + } + + private static func fromRateLimitTier(_ rateLimitTier: String?) -> Self? { + let tier = Self.normalized(rateLimitTier) + if tier.contains("max") { + return .max + } + if tier.contains("pro") { + return .pro + } + if tier.contains("team") { + return .team + } + if tier.contains("enterprise") { + return .enterprise + } + return nil + } + + private static func maxUsageMultiplier(rateLimitTier: String?) -> String? { + let words = Self.normalizedWords(rateLimitTier) + guard let maxIndex = words.firstIndex(of: Self.max.rawValue), + words.indices.contains(maxIndex + 1) + else { + return nil + } + + let multiplier = words[maxIndex + 1] + guard multiplier.last == "x", Int(multiplier.dropLast()) != nil else { + return nil + } + return multiplier + } + + private static func normalized(_ text: String?) -> String { + text? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() ?? "" + } + + private static func normalizedWords(_ text: String?) -> [String] { + self.normalized(text) + .split(whereSeparator: { !$0.isLetter && !$0.isNumber }) + .map(String.init) + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeProbeSessionArtifactCleaner.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeProbeSessionArtifactCleaner.swift new file mode 100644 index 0000000..fa7cce3 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeProbeSessionArtifactCleaner.swift @@ -0,0 +1,105 @@ +import Foundation + +enum ClaudeProbeSessionArtifactCleaner { + private static let log = CodexBarLog.logger(LogCategories.claudeProbe) + private static let maxProjectDirectoryNameLength = 200 + + @discardableResult + static func cleanupProbeSessionArtifacts( + probeDirectory: URL = ClaudeStatusProbe.probeWorkingDirectoryURL(), + environment: [String: String] = ProcessInfo.processInfo.environment, + fileManager fm: FileManager = .default) -> [URL] + { + let projectDirectoryName = self.claudeProjectDirectoryName(for: probeDirectory) + var visitedDirectories = Set() + var removedFiles: [URL] = [] + + for root in self.claudeConfigRoots(environment: environment, fileManager: fm) { + let projectsRoot = root.appendingPathComponent("projects", isDirectory: true) + let directories = [projectsRoot.appendingPathComponent(projectDirectoryName, isDirectory: true)] + + for directory in directories where visitedDirectories.insert(directory.path).inserted { + guard let entries = try? fm.contentsOfDirectory( + at: directory, + includingPropertiesForKeys: [.isRegularFileKey], + options: [.skipsHiddenFiles]) + else { continue } + + for entry in entries where entry.pathExtension == "jsonl" { + let values = try? entry.resourceValues(forKeys: [.isRegularFileKey]) + guard values?.isRegularFile == true else { continue } + do { + try fm.removeItem(at: entry) + removedFiles.append(entry) + } catch { + Self.log.debug( + "Claude probe session artifact cleanup skipped file", + metadata: ["error": error.localizedDescription]) + } + } + + if (try? fm.contentsOfDirectory(atPath: directory.path).isEmpty) == true { + try? fm.removeItem(at: directory) + } + } + } + + return removedFiles + } + + static func claudeProjectDirectoryName(for directory: URL) -> String { + let path = directory.path.precomposedStringWithCanonicalMapping + let sanitized = String(path.utf16.map { codeUnit in + switch codeUnit { + case 48...57, 65...90, 97...122: + Character(UnicodeScalar(codeUnit)!) + default: + "-" + } + }) + + guard sanitized.count > self.maxProjectDirectoryNameLength else { return sanitized } + return "\(sanitized.prefix(self.maxProjectDirectoryNameLength))-\(self.javaScriptHashBase36(path))" + } + + private static func javaScriptHashBase36(_ string: String) -> String { + var hash: Int32 = 0 + for codeUnit in string.utf16 { + hash = hash &* 31 &+ Int32(truncatingIfNeeded: codeUnit) + } + + let magnitude = hash < 0 ? -Int64(hash) : Int64(hash) + return String(magnitude, radix: 36) + } + + private static func claudeConfigRoots( + environment: [String: String], + fileManager fm: FileManager) -> [URL] + { + var roots: [URL] = [] + var seen = Set() + + func append(_ url: URL) { + let standardized = url.standardizedFileURL + guard seen.insert(standardized.path).inserted else { return } + roots.append(standardized) + } + + if let raw = environment["CLAUDE_CONFIG_DIR"] { + for part in raw.split(separator: ",") { + let path = part.trimmingCharacters(in: .whitespacesAndNewlines) + guard !path.isEmpty else { continue } + append(URL(fileURLWithPath: path)) + } + } + + let home = environment["HOME"].flatMap { $0.isEmpty ? nil : $0 } ?? NSHomeDirectory() + append(URL(fileURLWithPath: home).appendingPathComponent(".claude", isDirectory: true)) + append(URL(fileURLWithPath: home).appendingPathComponent(".config/claude", isDirectory: true)) + + if roots.isEmpty { + append(fm.homeDirectoryForCurrentUser.appendingPathComponent(".claude", isDirectory: true)) + } + return roots + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeProviderDescriptor.swift new file mode 100644 index 0000000..1868f5c --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeProviderDescriptor.swift @@ -0,0 +1,679 @@ +import Foundation + +public enum ClaudeProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .claude, + metadata: ProviderMetadata( + id: .claude, + displayName: "Claude", + sessionLabel: "Session", + weeklyLabel: "Weekly", + opusLabel: "Sonnet", + supportsOpus: true, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Claude Code usage", + cliName: "claude", + defaultEnabled: false, + isPrimaryProvider: true, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder, + dashboardURL: "https://console.anthropic.com/settings/billing", + subscriptionDashboardURL: "https://claude.ai/settings/usage", + changelogURL: "https://github.com/anthropics/claude-code/releases", + statusPageURL: "https://status.claude.com/"), + branding: ProviderBranding( + iconStyle: .claude, + iconResourceName: "ProviderIcon-claude", + color: ProviderColor(red: 204 / 255, green: 124 / 255, blue: 94 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: true, + noDataMessage: self.noDataMessage), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api, .web, .cli, .oauth], + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), + cli: ProviderCLIConfig( + name: "claude", + versionDetector: { browserDetection in + ClaudeUsageFetcher(browserDetection: browserDetection).detectVersion() + })) + } + + private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] { + if context.sourceMode == .api || self.hasAutoAdminAPIKey(context: context) { + return [ClaudeAdminAPIFetchStrategy()] + } + if ClaudeAdminAPIFetchStrategy.isSelectedAdminAPIAccount(context: context) { + return [ClaudeAdminAPIFetchStrategy()] + } + + let planningInput = Self.makePlanningInput(context: context) + let plan = ClaudeSourcePlanner.resolve(input: planningInput) + let manualCookieHeader = Self.manualCookieHeader(from: context) + + return plan.orderedSteps.map { step in + let strategy: any ProviderFetchStrategy = switch step.dataSource { + case .api: + ClaudeAdminAPIFetchStrategy() + case .oauth: + ClaudeOAuthFetchStrategy() + case .web: + ClaudeWebFetchStrategy(browserDetection: context.browserDetection) + case .cli: + ClaudeCLIFetchStrategy( + useWebExtras: context.runtime == .app + && planningInput.webExtrasEnabled, + manualCookieHeader: manualCookieHeader, + browserDetection: context.browserDetection, + hasWebFallback: planningInput.hasWebSession) + case .auto: + fatalError("Planner must not emit .auto as an executable step.") + } + return ClaudePlannedFetchStrategy(base: strategy, plannedStep: step) + } + } + + private static func hasAutoAdminAPIKey(context: ProviderFetchContext) -> Bool { + context.sourceMode == .auto && ClaudeAdminAPISettingsReader.apiKey(environment: context.env) != nil + } + + private static func makePlanningInput(context: ProviderFetchContext) -> ClaudeSourcePlanningInput { + let webExtrasEnabled = context.settings?.claude?.webExtrasEnabled ?? false + let needsOAuthAvailability = context.runtime == .app && context.sourceMode == .auto + let hasWebSession = Self.hasPlausibleWebSession(context: context) + + return ClaudeSourcePlanningInput( + runtime: context.runtime, + selectedDataSource: Self.sourceDataSource(from: context.sourceMode), + webExtrasEnabled: webExtrasEnabled, + hasWebSession: hasWebSession, + hasCLI: ClaudeCLIResolver.isAvailable(environment: context.env), + hasOAuthCredentials: needsOAuthAvailability && ClaudeOAuthPlanningAvailability.isAvailable( + runtime: context.runtime, + sourceMode: context.sourceMode, + environment: context.env)) + } + + private static func hasPlausibleWebSession(context: ProviderFetchContext) -> Bool { + switch context.sourceMode { + case .api, .oauth, .cli: + return false + case .web: + // Explicit web performs its cookie/session work inside the bounded fetch. + return context.settings?.claude?.cookieSource != .off + case .auto: + break + } + + guard ClaudeWebFetchStrategy.isSupportedOnCurrentPlatform else { return false } + + switch context.settings?.claude?.cookieSource { + case .off?: + return false + case .manual?: + return ClaudeWebFetchStrategy.hasManualSessionKey(context: context) + case .auto?, nil: + // Browser/Keychain inspection can block. Keep planning synchronous and let the web step + // perform the bounded availability check if app-auto actually reaches its last fallback. + // CLI auto continues to perform its real session import inside the bounded web fetch. + return true + } + } + + private static func manualCookieHeader(from context: ProviderFetchContext) -> String? { + guard context.settings?.claude?.cookieSource == .manual else { return nil } + return CookieHeaderNormalizer.normalize(context.settings?.claude?.manualCookieHeader) + } + + private static func noDataMessage() -> String { + "No Claude usage logs found in ~/.config/claude/projects, ~/.claude/projects, " + + "or Claude Desktop sessions." + } + + public static func resolveUsageStrategy( + selectedDataSource: ClaudeUsageDataSource, + webExtrasEnabled: Bool, + hasWebSession: Bool, + hasCLI: Bool, + hasOAuthCredentials: Bool) -> ClaudeUsageStrategy + { + let plan = ClaudeSourcePlanner.resolve(input: ClaudeSourcePlanningInput( + runtime: .app, + selectedDataSource: selectedDataSource, + webExtrasEnabled: webExtrasEnabled, + hasWebSession: hasWebSession, + hasCLI: hasCLI, + hasOAuthCredentials: hasOAuthCredentials)) + return plan.compatibilityStrategy ?? ClaudeUsageStrategy(dataSource: selectedDataSource, useWebExtras: false) + } + + private static func sourceDataSource(from mode: ProviderSourceMode) -> ClaudeUsageDataSource { + switch mode { + case .auto: + .auto + case .api: + .api + case .web: + .web + case .cli: + .cli + case .oauth: + .oauth + } + } +} + +public struct ClaudeUsageStrategy: Equatable, Sendable { + public let dataSource: ClaudeUsageDataSource + public let useWebExtras: Bool +} + +public enum ClaudeOAuthPlanningAvailability { + public static func isAvailable( + runtime: ProviderRuntime, + sourceMode: ProviderSourceMode, + environment: [String: String]) -> Bool + { + ClaudeOAuthFetchStrategy.isPlausiblyAvailable( + runtime: runtime, + sourceMode: sourceMode, + environment: environment) + } +} + +private struct ClaudePlannedFetchStrategy: ProviderFetchStrategy { + let base: any ProviderFetchStrategy + let plannedStep: ClaudeFetchPlanStep + + var id: String { + self.base.id + } + + var kind: ProviderFetchKind { + self.base.kind + } + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + guard context.sourceMode == .auto else { + return await self.base.isAvailable(context) + } + guard self.plannedStep.isPlausiblyAvailable else { return false } + if self.plannedStep.dataSource == .cli || + (context.runtime == .app && self.plannedStep.dataSource == .web) + { + return await self.base.isAvailable(context) + } + return true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + try await self.base.fetch(context) + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + self.base.shouldFallback(on: error, context: context) + } +} + +struct ClaudeAdminAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "claude.admin-api" + let kind: ProviderFetchKind = .apiToken + let usageFetcher: @Sendable (String) async throws -> ClaudeAdminAPIUsageSnapshot + + init( + usageFetcher: @escaping @Sendable (String) async throws -> ClaudeAdminAPIUsageSnapshot = { apiKey in + try await ClaudeAdminAPIUsageFetcher.fetchUsage(apiKey: apiKey) + }) + { + self.usageFetcher = usageFetcher + } + + static func isSelectedAdminAPIAccount(context: ProviderFetchContext) -> Bool { + guard context.selectedTokenAccountID != nil else { return false } + return self.resolveToken(environment: context.env) != nil + } + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolveToken(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveToken(environment: context.env) else { + throw ClaudeAdminAPISettingsError.missingToken + } + let usage = try await self.usageFetcher(apiKey) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "admin-api") + } + + func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool { + context.runtime == .app && + context.sourceMode == .auto && + !Self.isSelectedAdminAPIAccount(context: context) + } + + private static func resolveToken(environment: [String: String]) -> String? { + ProviderTokenResolver.claudeAdminAPIToken(environment: environment) + } +} + +struct ClaudeOAuthFetchStrategy: ProviderFetchStrategy { + let id: String = "claude.oauth" + let kind: ProviderFetchKind = .oauth + + #if DEBUG + @TaskLocal static var nonInteractiveCredentialRecordOverride: ClaudeOAuthCredentialRecord? + @TaskLocal static var claudeCLIAvailableOverride: Bool? + #endif + + private func loadNonInteractiveCredentialRecord(environment: [String: String]) -> ClaudeOAuthCredentialRecord? { + #if DEBUG + if let override = Self.nonInteractiveCredentialRecordOverride { + return override + } + #endif + + return try? ClaudeOAuthCredentialsStore.loadRecord( + environment: environment, + allowKeychainPrompt: false, + respectKeychainPromptCooldown: true, + allowClaudeKeychainRepairWithoutPrompt: false) + } + + private func isClaudeCLIAvailable(environment: [String: String]) -> Bool { + #if DEBUG + if let override = Self.claudeCLIAvailableOverride { + return override + } + #endif + return ClaudeCLIResolver.isAvailable(environment: environment) + } + + static func isPlausiblyAvailable( + runtime: ProviderRuntime, + sourceMode: ProviderSourceMode, + environment: [String: String]) -> Bool + { + let hasEnvironmentOAuthToken = !(environment[ClaudeOAuthCredentialsStore.environmentTokenKey]? + .trimmingCharacters(in: .whitespacesAndNewlines) + .isEmpty ?? true) + if hasEnvironmentOAuthToken { + return true + } + + let strategy = ClaudeOAuthFetchStrategy() + let nonInteractiveRecord = strategy.loadNonInteractiveCredentialRecord(environment: environment) + let nonInteractiveCredentials = nonInteractiveRecord?.credentials + let hasRequiredScopeWithoutPrompt = nonInteractiveCredentials?.scopes.contains("user:profile") == true + if hasRequiredScopeWithoutPrompt, nonInteractiveCredentials?.isExpired == false { + return true + } + + let claudeCLIAvailable = strategy.isClaudeCLIAvailable(environment: environment) + + if let nonInteractiveRecord, hasRequiredScopeWithoutPrompt, nonInteractiveRecord.credentials.isExpired { + switch nonInteractiveRecord.owner { + case .codexbar: + let refreshToken = nonInteractiveRecord.credentials.refreshToken? + .trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + if sourceMode == .auto { + return !refreshToken.isEmpty + } + return true + case .claudeCLI: + guard sourceMode == .auto else { return true } + guard claudeCLIAvailable else { return false } + guard ProviderInteractionContext.current == .background else { return true } + return !Self.hasMcpOAuthOnlyClaudeKeychainPayload(environment: environment) + case .environment: + return sourceMode != .auto + } + } + + guard sourceMode == .auto else { return true } + + let fallbackPromptMode = ClaudeOAuthKeychainPromptPreference.securityFrameworkFallbackMode() + let promptPolicyApplicable = ClaudeOAuthKeychainPromptPreference.isApplicable() + if ProviderInteractionContext.current == .userInitiated { + _ = ClaudeOAuthKeychainAccessGate.clearDenied() + } + + let shouldAllowStartupBootstrap = runtime == .app && + ProviderRefreshContext.current == .startup && + ProviderInteractionContext.current == .background && + fallbackPromptMode == .onlyOnUserAction && + !ClaudeOAuthCredentialsStore.hasCachedCredentials(environment: environment) + if shouldAllowStartupBootstrap { + return ClaudeOAuthKeychainAccessGate.shouldAllowPrompt() + } + + if promptPolicyApplicable, + !ClaudeOAuthKeychainAccessGate.shouldAllowPrompt() + { + return false + } + return ClaudeOAuthCredentialsStore.hasClaudeKeychainCredentialsWithoutPrompt() + } + + private static func hasMcpOAuthOnlyClaudeKeychainPayload(environment: [String: String]) -> Bool { + ClaudeOAuthCredentialsStore.isMcpOAuthOnlyClaudeKeychainPayloadPresent( + interaction: ProviderInteractionContext.current, + environment: environment) + } + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.isPlausiblyAvailable( + runtime: context.runtime, + sourceMode: context.sourceMode, + environment: context.env) + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let fetcher = ClaudeUsageFetcher( + browserDetection: context.browserDetection, + environment: context.env, + runtime: context.runtime, + dataSource: .oauth, + oauthKeychainPromptCooldownEnabled: context.sourceMode == .auto, + allowBackgroundDelegatedRefresh: false, + allowStartupBootstrapPrompt: context.runtime == .app && + (context.sourceMode == .auto || context.sourceMode == .oauth), + useWebExtras: false) + let usage = try await fetcher.loadLatestUsage(model: "sonnet") + return ProviderFetchResult( + usage: Self.snapshot(from: usage), + credits: nil, + dashboard: nil, + sourceLabel: "oauth", + strategyID: self.id, + strategyKind: self.kind, + claudeOAuthKeychainPersistentRefHash: usage.oauthKeychainPersistentRefHash, + claudeOAuthHistoryOwnerIdentifier: usage.oauthHistoryOwnerIdentifier, + claudeOAuthKeychainCredentialMismatch: usage.oauthKeychainCredentialMismatch, + claudeOAuthKeychainCredentialAbsent: usage.oauthKeychainCredentialAbsent, + claudeOAuthKeychainCredentialUnavailable: usage.oauthKeychainCredentialUnavailable) + } + + func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool { + // In Auto mode, fall back to the next strategy (cli/web) if OAuth fails (e.g. user cancels keychain prompt + // or auth breaks). + context.runtime == .app && context.sourceMode == .auto + } + + fileprivate static func snapshot(from usage: ClaudeUsageSnapshot) -> UsageSnapshot { + let identity = ProviderIdentitySnapshot( + providerID: .claude, + accountEmail: usage.accountEmail, + accountOrganization: usage.accountOrganization, + loginMethod: usage.loginMethod) + let primary = usage.primaryWindowKind == .spendLimit ? nil : usage.primary + return UsageSnapshot( + primary: primary, + secondary: usage.secondary, + tertiary: usage.opus, + extraRateWindows: usage.extraRateWindows.isEmpty ? nil : usage.extraRateWindows, + providerCost: usage.providerCost, + updatedAt: usage.updatedAt, + identity: identity) + } + + static func _snapshotForTesting(from usage: ClaudeUsageSnapshot) -> UsageSnapshot { + self.snapshot(from: usage) + } +} + +private final class ClaudeWebFetchDeadlineState: @unchecked Sendable { + private let lock = NSLock() + private var deadline: ContinuousClock.Instant? + + func remainingBudget( + fullBudget: Duration, + now: ContinuousClock.Instant) -> Duration + { + self.lock.lock() + defer { self.lock.unlock() } + + let deadline: ContinuousClock.Instant + if let existingDeadline = self.deadline { + deadline = existingDeadline + } else { + deadline = now.advanced(by: fullBudget) + self.deadline = deadline + } + return max(.zero, now.duration(to: deadline)) + } +} + +struct ClaudeWebFetchStrategy: ProviderFetchStrategy { + typealias UsageLoader = @Sendable (ProviderFetchContext) async throws -> ClaudeUsageSnapshot + typealias DeadlineNow = @Sendable () -> ContinuousClock.Instant + + #if DEBUG + @TaskLocal static var availabilityProbeOverrideForTesting: + (@Sendable (ProviderFetchContext, BrowserDetection) -> Bool)? + @TaskLocal static var usageLoaderOverrideForTesting: UsageLoader? + #endif + + let id: String = "claude.web" + let kind: ProviderFetchKind = .web + let browserDetection: BrowserDetection + private let usageLoader: UsageLoader? + private let deadlineState: ClaudeWebFetchDeadlineState + private let deadlineNow: DeadlineNow + + init( + browserDetection: BrowserDetection, + usageLoader: UsageLoader? = nil, + deadlineNow: @escaping DeadlineNow = { ContinuousClock.now }) + { + self.browserDetection = browserDetection + self.usageLoader = usageLoader + self.deadlineState = ClaudeWebFetchDeadlineState() + self.deadlineNow = deadlineNow + } + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + let appAutoRemainingBudget: Duration? + if context.runtime == .app, context.sourceMode == .auto { + guard let fullBudget = Self.timeoutDuration(context.webTimeout) else { return false } + let remainingBudget = self.deadlineState.remainingBudget( + fullBudget: fullBudget, + now: self.deadlineNow()) + guard remainingBudget > .zero else { return false } + appAutoRemainingBudget = remainingBudget + } else { + appAutoRemainingBudget = nil + } + + return switch context.settings?.claude?.cookieSource { + case .off?: + false + case .manual?: + Self.hasManualSessionKey(context: context) + case .auto?, nil: + if let appAutoRemainingBudget { + await Self.hasBrowserSessionKey(context: context, before: appAutoRemainingBudget) + } else { + // Explicit web and CLI auto perform the real browser import inside the bounded fetch. + true + } + } + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let usage = try await self.loadUsage(before: context.webTimeout, context: context) + return self.makeResult( + usage: ClaudeOAuthFetchStrategy.snapshot(from: usage), + sourceLabel: "web") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + guard context.sourceMode == .auto else { return false } + guard !Task.isCancelled, + !(error is CancellationError), + (error as? URLError)?.code != .cancelled + else { + return false + } + // In CLI runtime auto mode, web comes before CLI so fallback is required. + // In app runtime auto mode, web is terminal and should surface its concrete error. + return context.runtime == .cli + } + + fileprivate static func hasManualSessionKey(context: ProviderFetchContext) -> Bool { + ClaudeWebAPIFetcher.hasSessionKey(cookieHeader: self.manualCookieHeader(from: context)) + } + + fileprivate static func hasBrowserSessionKey( + context: ProviderFetchContext, + before timeout: Duration) async -> Bool + { + let browserDetection = context.browserDetection + let sourceTask = Task { + #if DEBUG + if let override = Self.availabilityProbeOverrideForTesting { + return override(context, browserDetection) + } + #endif + return ClaudeWebAPIFetcher.hasSessionKey(browserDetection: browserDetection) + } + let race = BoundedTaskJoin(sourceTask: sourceTask) + switch await race.value(joinGrace: timeout) { + case let .value(isAvailable): + return isAvailable + case .failure, .timedOut: + return false + } + } + + private static func manualCookieHeader(from context: ProviderFetchContext) -> String? { + guard context.settings?.claude?.cookieSource == .manual else { return nil } + return CookieHeaderNormalizer.normalize(context.settings?.claude?.manualCookieHeader) + } + + private func loadUsage( + before timeout: TimeInterval, + context: ProviderFetchContext) async throws -> ClaudeUsageSnapshot + { + guard let timeoutDuration = Self.timeoutDuration(timeout) else { + throw ClaudeWebFetchStrategyError.invalidTimeout + } + let remainingBudget = self.deadlineState.remainingBudget( + fullBudget: timeoutDuration, + now: self.deadlineNow()) + guard remainingBudget > .zero else { + try Task.checkCancellation() + throw ClaudeWebFetchStrategyError.timedOut(seconds: timeout) + } + let sourceTask = Task { + if let usageLoader = self.usageLoader { + return try await usageLoader(context) + } + #if DEBUG + if let usageLoader = Self.usageLoaderOverrideForTesting { + return try await usageLoader(context) + } + #endif + let fetcher = ClaudeUsageFetcher( + browserDetection: self.browserDetection, + dataSource: .web, + useWebExtras: false, + manualCookieHeader: Self.manualCookieHeader(from: context), + webOrganizationID: context.settings?.claude?.organizationID) + return try await fetcher.loadLatestUsage(model: "sonnet") + } + let race = BoundedTaskJoin(sourceTask: sourceTask) + switch await race.value(joinGrace: remainingBudget) { + case let .value(usage): + try Task.checkCancellation() + return usage + case let .failure(error): + throw error + case .timedOut: + try Task.checkCancellation() + throw ClaudeWebFetchStrategyError.timedOut(seconds: timeout) + } + } + + private static func timeoutDuration(_ timeout: TimeInterval) -> Duration? { + guard timeout.isFinite, + timeout >= 0, + timeout <= TimeInterval(Int64.max) + else { + return nil + } + return .seconds(timeout) + } + + fileprivate static var isSupportedOnCurrentPlatform: Bool { + #if os(macOS) + true + #else + false + #endif + } +} + +public enum ClaudeWebFetchStrategyError: LocalizedError, Equatable, Sendable { + case invalidTimeout + case timedOut(seconds: TimeInterval) + + public var errorDescription: String? { + switch self { + case .invalidTimeout: + "Claude web usage fetch timeout must be a finite, nonnegative value within the supported range." + case let .timedOut(seconds): + "Claude web usage fetch timed out after \(seconds.formatted()) seconds." + } + } +} + +struct ClaudeCLIFetchStrategy: ProviderFetchStrategy { + let id: String = "claude.cli" + let kind: ProviderFetchKind = .cli + let useWebExtras: Bool + let manualCookieHeader: String? + let browserDetection: BrowserDetection + let hasWebFallback: Bool + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + // The interactive Claude REPL can open browser OAuth when it starts logged out. CLI runtime and background + // app Auto refreshes must establish authentication through the non-interactive status command first. + let requiresAuthPreflight = context.runtime == .cli || ( + context.runtime == .app && + context.sourceMode == .auto && + ProviderInteractionContext.current == .background) + guard requiresAuthPreflight else { return true } + guard let binary = ClaudeCLIResolver.resolvedBinaryPath(environment: context.env) else { return false } + return await ClaudeCLIAuthStatusProbe.isLoggedIn(binary: binary, environment: context.env) + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let keepAlive = context.settings?.debugKeepCLISessionsAlive ?? false + let fetcher = ClaudeUsageFetcher( + browserDetection: browserDetection, + environment: context.env, + dataSource: .cli, + useWebExtras: self.useWebExtras, + manualCookieHeader: self.manualCookieHeader, + webOrganizationID: context.settings?.claude?.organizationID, + keepCLISessionsAlive: keepAlive) + let usage = try await fetcher.loadLatestUsage(model: "sonnet") + return self.makeResult( + usage: ClaudeOAuthFetchStrategy.snapshot(from: usage), + sourceLabel: "claude") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + guard context.runtime == .app, context.sourceMode == .auto else { return false } + guard !ClaudeStatusProbe.isSubscriptionQuotaUnavailableDescription(error.localizedDescription) else { + return false + } + // Reuse the bounded planning result instead of repeating browser/Keychain work after CLI failure. + return self.hasWebFallback + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeScopedWeeklyLimitMapper.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeScopedWeeklyLimitMapper.swift new file mode 100644 index 0000000..0abfbaa --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeScopedWeeklyLimitMapper.swift @@ -0,0 +1,62 @@ +import Foundation + +/// Shared mapping for the model-scoped weekly limits returned by both Claude usage APIs. +enum ClaudeScopedWeeklyLimitMapper { + struct Limit { + let kind: String? + let group: String? + let percent: Double? + let resetsAt: Date? + let modelID: String? + let modelName: String? + } + + static func extraRateWindows( + from limits: [Limit]?, + resetDescription: ((Date) -> String)? = nil) -> [NamedRateWindow] + { + guard let limits else { return [] } + var seenIDs: Set = [] + + return limits.compactMap { limit in + guard limit.group == "weekly", limit.kind == "weekly_scoped" else { return nil } + guard let percent = limit.percent, percent.isFinite else { return nil } + guard let modelName = self.nonEmpty(limit.modelName) else { return nil } + let identity = self.nonEmpty(limit.modelID) ?? modelName + let slug = self.slug(identity) + guard !slug.isEmpty else { return nil } + + let id = "claude-weekly-scoped-\(slug)" + guard seenIDs.insert(id).inserted else { return nil } + + return NamedRateWindow( + id: id, + title: "\(modelName) only", + window: RateWindow( + usedPercent: percent, + windowMinutes: 7 * 24 * 60, + resetsAt: limit.resetsAt, + resetDescription: limit.resetsAt.flatMap { resetDescription?($0) })) + } + } + + private static func nonEmpty(_ value: String?) -> String? { + let trimmed = value?.trimmingCharacters(in: .whitespacesAndNewlines) + return if let trimmed, !trimmed.isEmpty { trimmed } else { nil } + } + + private static func slug(_ value: String) -> String { + var result = "" + var lastWasDash = false + for scalar in value.lowercased().unicodeScalars { + if CharacterSet.alphanumerics.contains(scalar) { + result.unicodeScalars.append(scalar) + lastWasDash = false + } else if !lastWasDash { + result.append("-") + lastWasDash = true + } + } + return result.trimmingCharacters(in: CharacterSet(charactersIn: "-")) + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeSourcePlanner.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeSourcePlanner.swift new file mode 100644 index 0000000..24ae2c0 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeSourcePlanner.swift @@ -0,0 +1,224 @@ +import Foundation + +public struct ClaudeSourcePlanningInput: Equatable, Sendable { + public let runtime: ProviderRuntime + public let selectedDataSource: ClaudeUsageDataSource + public let webExtrasEnabled: Bool + public let hasWebSession: Bool + public let hasCLI: Bool + public let hasOAuthCredentials: Bool + + public init( + runtime: ProviderRuntime, + selectedDataSource: ClaudeUsageDataSource, + webExtrasEnabled: Bool, + hasWebSession: Bool, + hasCLI: Bool, + hasOAuthCredentials: Bool) + { + self.runtime = runtime + self.selectedDataSource = selectedDataSource + self.webExtrasEnabled = webExtrasEnabled + self.hasWebSession = hasWebSession + self.hasCLI = hasCLI + self.hasOAuthCredentials = hasOAuthCredentials + } +} + +public enum ClaudeSourcePlanReason: String, Equatable, Sendable { + case explicitSourceSelection = "explicit-source-selection" + case appAutoPreferredOAuth = "app-auto-preferred-oauth" + case appAutoFallbackCLI = "app-auto-fallback-cli" + case appAutoFallbackWeb = "app-auto-fallback-web" + case cliAutoPreferredWeb = "cli-auto-preferred-web" + case cliAutoFallbackCLI = "cli-auto-fallback-cli" +} + +public struct ClaudeFetchPlanStep: Equatable, Sendable { + public let dataSource: ClaudeUsageDataSource + public let inclusionReason: ClaudeSourcePlanReason + public let isPlausiblyAvailable: Bool + + public init( + dataSource: ClaudeUsageDataSource, + inclusionReason: ClaudeSourcePlanReason, + isPlausiblyAvailable: Bool) + { + self.dataSource = dataSource + self.inclusionReason = inclusionReason + self.isPlausiblyAvailable = isPlausiblyAvailable + } +} + +public struct ClaudeFetchPlan: Equatable, Sendable { + public let input: ClaudeSourcePlanningInput + public let orderedSteps: [ClaudeFetchPlanStep] + + public init(input: ClaudeSourcePlanningInput, orderedSteps: [ClaudeFetchPlanStep]) { + self.input = input + self.orderedSteps = orderedSteps + } + + public var availableSteps: [ClaudeFetchPlanStep] { + self.orderedSteps.filter(\.isPlausiblyAvailable) + } + + public var isNoSourceAvailable: Bool { + self.availableSteps.isEmpty + } + + public var preferredStep: ClaudeFetchPlanStep? { + switch self.input.selectedDataSource { + case .auto: + self.availableSteps.first + case .api, .oauth, .web, .cli: + self.orderedSteps.first + } + } + + public var executionSteps: [ClaudeFetchPlanStep] { + switch self.input.selectedDataSource { + case .auto: + self.availableSteps + case .api, .oauth, .web, .cli: + self.orderedSteps + } + } + + public var compatibilityStrategy: ClaudeUsageStrategy? { + guard let preferredStep else { return nil } + let useWebExtras = self.input.runtime == .app + && preferredStep.dataSource == .cli + && self.input.webExtrasEnabled + return ClaudeUsageStrategy( + dataSource: preferredStep.dataSource, + useWebExtras: useWebExtras) + } + + public var orderLabel: String { + self.orderedSteps.map(\.dataSource.sourceLabel).joined(separator: "→") + } + + public func debugLines() -> [String] { + var lines = ["planner_order=\(self.orderLabel)"] + lines.append("planner_selected=\(self.preferredStep?.dataSource.rawValue ?? "none")") + lines.append("planner_no_source=\(self.isNoSourceAvailable)") + for step in self.orderedSteps { + let availability = step.isPlausiblyAvailable ? "available" : "unavailable" + lines.append( + "planner_step.\(step.dataSource.rawValue)=\(availability) reason=\(step.inclusionReason.rawValue)") + } + return lines + } +} + +public enum ClaudeCLIResolver { + #if DEBUG + @TaskLocal static var resolvedBinaryPathOverrideForTesting: String? + + public static var currentResolvedBinaryPathOverrideForTesting: String? { + self.resolvedBinaryPathOverrideForTesting + } + + public static func withResolvedBinaryPathOverrideForTesting( + _ path: String?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$resolvedBinaryPathOverrideForTesting.withValue(path) { + try await operation() + } + } + #endif + + public static func resolvedBinaryPath( + environment: [String: String] = ProcessInfo.processInfo.environment, + loginPATH: [String]? = LoginShellPathCache.shared.current) + -> String? + { + #if DEBUG + if let override = self.resolvedBinaryPathOverrideForTesting { + return FileManager.default.isExecutableFile(atPath: override) ? override : nil + } + #endif + + var normalizedEnvironment = environment + if let override = environment["CLAUDE_CLI_PATH"]?.trimmingCharacters(in: .whitespacesAndNewlines) { + if override.isEmpty { + normalizedEnvironment.removeValue(forKey: "CLAUDE_CLI_PATH") + } else { + normalizedEnvironment["CLAUDE_CLI_PATH"] = override + if FileManager.default.isExecutableFile(atPath: override) { + return override + } + } + } + + return BinaryLocator.resolveClaudeBinary( + env: normalizedEnvironment, + loginPATH: loginPATH) + } + + public static func isAvailable(environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool { + self.resolvedBinaryPath(environment: environment) != nil + } +} + +public enum ClaudeSourcePlanner { + public static func resolve(input: ClaudeSourcePlanningInput) -> ClaudeFetchPlan { + ClaudeFetchPlan(input: input, orderedSteps: self.makeSteps(input: input)) + } + + private static func makeSteps(input: ClaudeSourcePlanningInput) -> [ClaudeFetchPlanStep] { + switch input.selectedDataSource { + case .auto: + switch input.runtime { + case .app: + [ + self.step(.oauth, reason: .appAutoPreferredOAuth, input: input), + self.step(.cli, reason: .appAutoFallbackCLI, input: input), + self.step(.web, reason: .appAutoFallbackWeb, input: input), + ] + case .cli: + [ + self.step(.web, reason: .cliAutoPreferredWeb, input: input), + self.step(.cli, reason: .cliAutoFallbackCLI, input: input), + ] + } + case .api: + [self.step(.api, reason: .explicitSourceSelection, input: input)] + case .oauth: + [self.step(.oauth, reason: .explicitSourceSelection, input: input)] + case .web: + [self.step(.web, reason: .explicitSourceSelection, input: input)] + case .cli: + [self.step(.cli, reason: .explicitSourceSelection, input: input)] + } + } + + private static func step( + _ dataSource: ClaudeUsageDataSource, + reason: ClaudeSourcePlanReason, + input: ClaudeSourcePlanningInput) -> ClaudeFetchPlanStep + { + ClaudeFetchPlanStep( + dataSource: dataSource, + inclusionReason: reason, + isPlausiblyAvailable: self.isPlausiblyAvailable(dataSource, input: input)) + } + + private static func isPlausiblyAvailable( + _ dataSource: ClaudeUsageDataSource, + input: ClaudeSourcePlanningInput) -> Bool + { + switch dataSource { + case .auto, .api: + false + case .oauth: + input.hasOAuthCredentials + case .web: + input.hasWebSession + case .cli: + input.hasCLI + } + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeStatusProbe.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeStatusProbe.swift new file mode 100644 index 0000000..9c255ee --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeStatusProbe.swift @@ -0,0 +1,1353 @@ +import Foundation + +public struct ClaudeStatusSnapshot: Sendable { + public let sessionPercentLeft: Int? + public let weeklyPercentLeft: Int? + public let opusPercentLeft: Int? + public let extraRateWindows: [NamedRateWindow] + public let accountEmail: String? + public let accountOrganization: String? + public let loginMethod: String? + public let primaryResetDescription: String? + public let secondaryResetDescription: String? + public let opusResetDescription: String? + public let rawText: String + + public init( + sessionPercentLeft: Int?, + weeklyPercentLeft: Int?, + opusPercentLeft: Int?, + accountEmail: String?, + accountOrganization: String?, + loginMethod: String?, + primaryResetDescription: String?, + secondaryResetDescription: String?, + opusResetDescription: String?, + rawText: String, + extraRateWindows: [NamedRateWindow] = []) + { + self.sessionPercentLeft = sessionPercentLeft + self.weeklyPercentLeft = weeklyPercentLeft + self.opusPercentLeft = opusPercentLeft + self.extraRateWindows = extraRateWindows + self.accountEmail = accountEmail + self.accountOrganization = accountOrganization + self.loginMethod = loginMethod + self.primaryResetDescription = primaryResetDescription + self.secondaryResetDescription = secondaryResetDescription + self.opusResetDescription = opusResetDescription + self.rawText = rawText + } +} + +public struct ClaudeAccountIdentity: Sendable { + public let accountEmail: String? + public let accountOrganization: String? + public let loginMethod: String? + + public init(accountEmail: String?, accountOrganization: String?, loginMethod: String?) { + self.accountEmail = accountEmail + self.accountOrganization = accountOrganization + self.loginMethod = loginMethod + } +} + +public enum ClaudeStatusProbeError: LocalizedError, Sendable { + case claudeNotInstalled + case parseFailed(String) + case timedOut + + public var errorDescription: String? { + switch self { + case .claudeNotInstalled: + "Claude CLI is not installed or not on PATH." + case let .parseFailed(msg): + "Could not parse Claude usage: \(msg)" + case .timedOut: + "Claude usage probe timed out." + } + } +} + +/// Runs `claude` inside a PTY, sends `/usage`, and parses the rendered text panel. +public struct ClaudeStatusProbe: Sendable { + public static let subscriptionQuotaUnavailableDescription = + "Claude CLI /usage returned a subscription notice without session quota data. " + + "Local cost and token history remain available." + + public var claudeBinary: String = "claude" + public var timeout: TimeInterval = 20.0 + public var keepCLISessionsAlive: Bool = false + private static let log = CodexBarLog.logger(LogCategories.claudeProbe) + #if DEBUG + public typealias FetchOverride = @Sendable (String, TimeInterval, Bool) async throws -> ClaudeStatusSnapshot + @TaskLocal static var fetchOverride: FetchOverride? + #endif + + public init(claudeBinary: String = "claude", timeout: TimeInterval = 20.0, keepCLISessionsAlive: Bool = false) { + self.claudeBinary = claudeBinary + self.timeout = timeout + self.keepCLISessionsAlive = keepCLISessionsAlive + } + + #if DEBUG + public static var currentFetchOverrideForTesting: FetchOverride? { + self.fetchOverride + } + + public static func withFetchOverrideForTesting( + _ override: FetchOverride?, + operation: () async throws -> T) async rethrows -> T + { + try await self.$fetchOverride.withValue(override) { + try await operation() + } + } + + public static func withFetchOverrideForTesting( + _ override: FetchOverride?, + operation: () async -> T) async -> T + { + await self.$fetchOverride.withValue(override) { + await operation() + } + } + #endif + + public func fetch() async throws -> ClaudeStatusSnapshot { + let resolved = Self.resolvedBinaryPath(binaryName: self.claudeBinary) + guard let resolved, Self.isBinaryAvailable(resolved) else { + throw ClaudeStatusProbeError.claudeNotInstalled + } + + // Run commands sequentially through a shared Claude session to avoid warm-up churn. + let timeout = self.timeout + let keepAlive = self.keepCLISessionsAlive + #if DEBUG + if let override = Self.fetchOverride { + return try await override(resolved, timeout, keepAlive) + } + #endif + do { + var usage = try await Self.capture(subcommand: "/usage", binary: resolved, timeout: timeout) + if !Self.usageOutputLooksRelevant(usage) { + Self.log.debug("Claude CLI /usage looked like startup output; retrying once") + usage = try await Self.capture(subcommand: "/usage", binary: resolved, timeout: max(timeout, 14)) + } + // `/status` only enriches a valid usage snapshot with identity. Terminal usage errors and loading stalls + // cannot be repaired by it, so fail now instead of paying for another interactive CLI round trip. + try Self.validateUsageBeforeStatusProbe(usage) + let status = try? await Self.capture(subcommand: "/status", binary: resolved, timeout: min(timeout, 12)) + let snap = try Self.parse(text: usage, statusText: status) + + Self.log.info("Claude CLI scrape ok", metadata: [ + "sessionPercentLeft": "\(snap.sessionPercentLeft ?? -1)", + "weeklyPercentLeft": "\(snap.weeklyPercentLeft ?? -1)", + "opusPercentLeft": "\(snap.opusPercentLeft ?? -1)", + ]) + if !keepAlive { + await Self.resetTransientCLISessionAndCleanupProbeArtifacts() + } + return snap + } catch { + if !keepAlive { + await Self.resetTransientCLISessionAndCleanupProbeArtifacts() + } + throw error + } + } + + private static func resetTransientCLISessionAndCleanupProbeArtifacts() async { + await ClaudeCLISession.current.reset() + let removed = ClaudeProbeSessionArtifactCleaner.cleanupProbeSessionArtifacts() + guard !removed.isEmpty else { return } + Self.log.debug("Claude probe session artifacts removed", metadata: ["count": "\(removed.count)"]) + } +} + +extension ClaudeStatusProbe { + // MARK: - Parsing helpers + + private struct LabelSearchContext { + let lines: [String] + let normalizedLines: [String] + let normalizedData: Data + + init(text: String) { + self.lines = text.components(separatedBy: .newlines) + self.normalizedLines = self.lines.map { ClaudeStatusProbe.normalizedForLabelSearch($0) } + let normalized = ClaudeStatusProbe.normalizedForLabelSearch(text) + self.normalizedData = Data(normalized.utf8) + } + + func contains(_ needle: String) -> Bool { + self.normalizedData.range(of: Data(needle.utf8)) != nil + } + } + + public static func parse(text: String, statusText: String? = nil) throws -> ClaudeStatusSnapshot { + let clean = TextParsing.stripANSICodes(text) + let statusClean = statusText.map(TextParsing.stripANSICodes) + guard !clean.isEmpty else { throw ClaudeStatusProbeError.timedOut } + + let shouldDump = ProcessInfo.processInfo.environment["DEBUG_CLAUDE_DUMP"] == "1" + + if let usageError = self.extractUsageError(text: clean) { + Self.dumpIfNeeded( + enabled: shouldDump, + reason: "usageError: \(usageError)", + usage: clean, + status: statusText) + throw ClaudeStatusProbeError.parseFailed(usageError) + } + + let latestUsagePanel = self.trimToLatestUsagePanel(clean) + if self.isUsageStillLoading(text: latestUsagePanel ?? clean) { + Self.dumpIfNeeded( + enabled: shouldDump, + reason: "usage still loading", + usage: clean, + status: statusText) + throw ClaudeStatusProbeError.parseFailed("Claude CLI /usage is still loading usage data.") + } + + // Claude CLI renders /usage as a TUI. Our PTY capture includes earlier screen fragments (including a status + // line + // with a "0%" context meter) before the usage panel is drawn. To keep parsing stable, trim to the last + // Settings/Usage panel when present. + let usagePanelText = latestUsagePanel ?? clean + let labelContext = LabelSearchContext(text: usagePanelText) + + var sessionPct = self.extractPercent(labelSubstring: "Current session", context: labelContext) + let weeklyPct = self.extractPercent(labelSubstring: "Current week (all models)", context: labelContext) + let opusLabels = [ + "Current week (Opus)", + "Current week (Sonnet only)", + "Current week (Sonnet)", + ] + let opusPct = self.extractPercent(labelSubstrings: opusLabels, context: labelContext) + + // Fallback: order-based percent scraping when labels are present but the surrounding layout moved. + // Only apply the fallback when the corresponding label exists in the rendered panel; enterprise accounts + // may omit the weekly panel entirely, and we should treat that as "unavailable" rather than guessing. + let weeklyModels = Set(labelContext.lines.compactMap(self.weeklyModelName).map(self.normalizedForLabelSearch)) + let hasAllModelsWeeklyLabel = weeklyModels.contains("allmodels") + let opusModels = Set(opusLabels.compactMap(self.weeklyModelName).map(self.normalizedForLabelSearch)) + let hasOpusLabel = !weeklyModels.isDisjoint(with: opusModels) + + if sessionPct == nil { + let ordered = self.allPercents(usagePanelText) + if sessionPct == nil, ordered.indices.contains(0) { + sessionPct = ordered[0] + } + } + + let identity = Self.parseIdentity(usageText: clean, statusText: statusClean) + + guard let sessionPct else { + Self.dumpIfNeeded( + enabled: shouldDump, + reason: "missing session label", + usage: clean, + status: statusText) + if shouldDump { + let tail = usagePanelText.suffix(1800) + let snippet = tail.isEmpty ? "(empty)" : String(tail) + throw ClaudeStatusProbeError.parseFailed( + "Missing Current session.\n\n--- Clean usage tail ---\n\(snippet)") + } + throw ClaudeStatusProbeError.parseFailed("Missing Current session.") + } + + let sessionReset = self.extractReset(labelSubstring: "Current session", context: labelContext) + let weeklyReset = hasAllModelsWeeklyLabel + ? self.extractReset(labelSubstring: "Current week (all models)", context: labelContext) + : nil + let opusReset = hasOpusLabel + ? self.extractReset(labelSubstrings: opusLabels, context: labelContext) + : nil + let scopedWeeklyUsages = self.extractScopedWeeklyUsages(context: labelContext) + let extraRateWindows = self.extraRateWindows( + fromScopedWeeklyUsages: scopedWeeklyUsages, + fallbackResetDescription: weeklyReset) + + return ClaudeStatusSnapshot( + sessionPercentLeft: sessionPct, + weeklyPercentLeft: weeklyPct, + opusPercentLeft: opusPct, + accountEmail: identity.accountEmail, + accountOrganization: identity.accountOrganization, + loginMethod: identity.loginMethod, + primaryResetDescription: sessionReset, + secondaryResetDescription: weeklyReset, + opusResetDescription: opusReset, + rawText: text + (statusText ?? ""), + extraRateWindows: extraRateWindows) + } + + public static func parseIdentity(usageText: String?, statusText: String?) -> ClaudeAccountIdentity { + let usageClean = usageText.map(TextParsing.stripANSICodes) ?? "" + let statusClean = statusText.map(TextParsing.stripANSICodes) + return self.extractIdentity(usageText: usageClean, statusText: statusClean) + } + + public static func fetchIdentity( + timeout: TimeInterval = 12.0, + environment: [String: String] = ProcessInfo.processInfo.environment) async throws -> ClaudeAccountIdentity + { + let resolved = self.resolvedBinaryPath(binaryName: "claude", environment: environment) + guard let resolved, self.isBinaryAvailable(resolved) else { + throw ClaudeStatusProbeError.claudeNotInstalled + } + let statusText = try await Self.capture(subcommand: "/status", binary: resolved, timeout: timeout) + return Self.parseIdentity(usageText: nil, statusText: statusText) + } + + public static func touchOAuthAuthPath( + timeout: TimeInterval = 8, + environment: [String: String] = ProcessInfo.processInfo.environment) async throws + { + let resolved = self.resolvedBinaryPath(binaryName: "claude", environment: environment) + guard let resolved, self.isBinaryAvailable(resolved) else { + throw ClaudeStatusProbeError.claudeNotInstalled + } + do { + // Use a more robust capture configuration than the standard `/status` scrape: + // - Avoid the short idle-timeout which can terminate the session while CLI auth checks are still running. + // - We intentionally do not parse output here; success is "the command ran without timing out". + _ = try await ClaudeCLISession.shared.capture( + subcommand: "/status", + binary: resolved, + timeout: timeout, + idleTimeout: nil, + stopOnSubstrings: [], + settleAfterStop: 0.8, + sendEnterEvery: 0.8) + await ClaudeCLISession.shared.reset() + } catch { + await ClaudeCLISession.shared.reset() + throw error + } + } + + public static func isClaudeBinaryAvailable( + environment: [String: String] = ProcessInfo.processInfo.environment) -> Bool + { + let resolved = self.resolvedBinaryPath(binaryName: "claude", environment: environment) + return self.isBinaryAvailable(resolved) + } + + private static func extractPercent(labelSubstring: String, context: LabelSearchContext) -> Int? { + let lines = context.lines + let label = self.normalizedForLabelSearch(labelSubstring) + for (idx, line) in lines.enumerated() { + let normalizedLine = context.normalizedLines[idx] + guard self.matchesLabel( + line: line, + normalizedLine: normalizedLine, + labelSubstring: labelSubstring, + normalizedLabel: label) + else { continue } + + // Claude's usage panel can take a moment to render percentages (especially on enterprise accounts), + // so scan a larger window than the original 3–4 lines. + let window = lines.dropFirst(idx).prefix(12) + for candidate in window { + if self.crossesLabelBoundary( + line: candidate, + labelSubstring: labelSubstring, + normalizedLabel: label) + { + break + } + if let pct = self.percentFromLine(candidate) { + return pct + } + } + } + return nil + } + + private static func usageOutputLooksRelevant(_ text: String) -> Bool { + let normalized = TextParsing.stripANSICodes(text).lowercased().filter { !$0.isWhitespace } + return normalized.contains("currentsession") + || normalized.contains("currentweek") + || normalized.contains("loadingusage") + || normalized.contains("failedtoloadusagedata") + || self.usageCaptureHasSubscriptionNotice(normalized) + } + + private static func validateUsageBeforeStatusProbe(_ text: String) throws { + let clean = TextParsing.stripANSICodes(text) + if let usageError = self.extractUsageError(text: clean) { + throw ClaudeStatusProbeError.parseFailed(usageError) + } + + let latestUsagePanel = self.trimToLatestUsagePanel(clean) + if self.isUsageStillLoading(text: latestUsagePanel ?? clean) { + throw ClaudeStatusProbeError.parseFailed("Claude CLI /usage is still loading usage data.") + } + } + + private static func extractPercent(labelSubstrings: [String], context: LabelSearchContext) -> Int? { + for label in labelSubstrings { + if let value = self.extractPercent(labelSubstring: label, context: context) { + return value + } + } + return nil + } + + private struct ScopedWeeklyUsage { + let modelName: String + let percentLeft: Int + let resetDescription: String? + } + + private static func extractScopedWeeklyUsages(context: LabelSearchContext) -> [ScopedWeeklyUsage] { + var usageIndexByModel: [String: Int] = [:] + var usages: [ScopedWeeklyUsage] = [] + for (index, line) in context.lines.enumerated() { + guard let modelName = self.weeklyModelName(from: line) else { continue } + let normalizedModel = self.normalizedForLabelSearch(modelName) + guard normalizedModel != "allmodels", !normalizedModel.isEmpty else { continue } + + let window = context.lines.dropFirst(index).prefix(14) + var percentLeft: Int? + var resetDescription: String? + for candidate in window { + if self.crossesLabelBoundary( + line: candidate, + labelSubstring: line, + normalizedLabel: self.normalizedForLabelSearch(line)) + { + break + } + if percentLeft == nil { + percentLeft = self.percentFromLine(candidate) + } + if resetDescription == nil { + resetDescription = self.resetFromLine(candidate) + } + } + guard let percentLeft else { continue } + let usage = ScopedWeeklyUsage( + modelName: modelName, + percentLeft: percentLeft, + resetDescription: resetDescription) + if let existingIndex = usageIndexByModel[normalizedModel] { + usages[existingIndex] = usage + } else { + usageIndexByModel[normalizedModel] = usages.endIndex + usages.append(usage) + } + } + return usages + } + + private static func extraRateWindows( + fromScopedWeeklyUsages usages: [ScopedWeeklyUsage], + fallbackResetDescription: String?) -> [NamedRateWindow] + { + usages.compactMap { usage in + let normalizedModel = self.normalizedForLabelSearch(usage.modelName) + guard normalizedModel != "opus", + normalizedModel != "sonnet", + normalizedModel != "sonnetonly" + else { return nil } + + let usedPercent = max(0, min(100, 100 - Double(usage.percentLeft))) + let modelTitle = normalizedModel.hasSuffix("only") + ? usage.modelName + : "\(usage.modelName) only" + let resetDescription = usage.resetDescription ?? fallbackResetDescription + return NamedRateWindow( + id: "claude-weekly-scoped-\(self.slug(usage.modelName))", + title: modelTitle, + window: RateWindow( + usedPercent: usedPercent, + windowMinutes: 7 * 24 * 60, + resetsAt: self.parseResetDate( + from: resetDescription, + expectedWindow: 7 * 24 * 60 * 60), + resetDescription: resetDescription)) + } + } + + private static func slug(_ value: String) -> String { + var result = "" + var lastWasDash = false + for scalar in value.lowercased().unicodeScalars { + if CharacterSet.alphanumerics.contains(scalar) { + result.unicodeScalars.append(scalar) + lastWasDash = false + } else if !lastWasDash { + result.append("-") + lastWasDash = true + } + } + return result.trimmingCharacters(in: CharacterSet(charactersIn: "-")) + } + + private static func percentFromLine(_ line: String, assumeRemainingWhenUnclear: Bool = false) -> Int? { + if self.isLikelyStatusContextLine(line) { + return nil + } + + // Allow optional Unicode whitespace before % to handle CLI formatting changes. + let pattern = #"([0-9]{1,3}(?:\.[0-9]+)?)\p{Zs}*%"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) else { return nil } + let range = NSRange(line.startIndex..= 2, + let valRange = Range(match.range(at: 1), in: line) + else { return nil } + let rawVal = Double(line[valRange]) ?? 0 + let clamped = max(0, min(100, rawVal)) + let lower = line.lowercased() + let usedKeywords = ["used", "spent", "consumed"] + let remainingKeywords = ["left", "remaining", "available"] + if usedKeywords.contains(where: lower.contains) { + return Int(max(0, min(100, 100 - clamped)).rounded()) + } + if remainingKeywords.contains(where: lower.contains) { + return Int(clamped.rounded()) + } + return assumeRemainingWhenUnclear ? Int(clamped.rounded()) : nil + } + + private static func isLikelyStatusContextLine(_ line: String) -> Bool { + guard line.contains("|") else { return false } + let lower = line.lowercased() + let modelTokens = ["opus", "sonnet", "haiku", "default"] + return modelTokens.contains(where: lower.contains) + } + + private static func extractFirst(pattern: String, text: String) -> String? { + guard let regex = try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) else { return nil } + let range = NSRange(text.startIndex..= 2, + let r = Range(match.range(at: 1), in: text) else { return nil } + return String(text[r]).trimmingCharacters(in: .whitespacesAndNewlines) + } + + private static func extractIdentity(usageText: String, statusText: String?) -> ClaudeAccountIdentity { + let emailPatterns = [ + #"(?i)Account:\s+([^\s@]+@[^\s@]+)"#, + #"(?i)Email:\s+([^\s@]+@[^\s@]+)"#, + ] + let looseEmailPatterns = [ + #"(?i)Account:\s+(\S+)"#, + #"(?i)Email:\s+(\S+)"#, + ] + let email = emailPatterns + .compactMap { self.extractFirst(pattern: $0, text: usageText) } + .first + ?? emailPatterns + .compactMap { self.extractFirst(pattern: $0, text: statusText ?? "") } + .first + ?? looseEmailPatterns + .compactMap { self.extractFirst(pattern: $0, text: usageText) } + .first + ?? looseEmailPatterns + .compactMap { self.extractFirst(pattern: $0, text: statusText ?? "") } + .first + ?? self.extractFirst( + pattern: #"(?i)[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}"#, + text: usageText) + ?? self.extractFirst( + pattern: #"(?i)[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}"#, + text: statusText ?? "") + let orgPatterns = [ + #"(?i)Org:\s*(.+)"#, + #"(?i)Organization:\s*(.+)"#, + ] + let orgRaw = orgPatterns + .compactMap { self.extractFirst(pattern: $0, text: usageText) } + .first + ?? orgPatterns + .compactMap { self.extractFirst(pattern: $0, text: statusText ?? "") } + .first + let org: String? = { + guard let orgText = orgRaw?.trimmingCharacters(in: .whitespacesAndNewlines), !orgText.isEmpty else { + return nil + } + // Suppress org if it’s just the email prefix (common in CLI panels). + if let email, orgText.lowercased().hasPrefix(email.lowercased()) { + return nil + } + return orgText + }() + // Prefer explicit login method from /status, then fall back to /usage header heuristics. + let login = self.extractLoginMethod(text: statusText ?? "") ?? self.extractLoginMethod(text: usageText) + return ClaudeAccountIdentity(accountEmail: email, accountOrganization: org, loginMethod: login) + } + + private static func extractUsageError(text: String) -> String? { + if let jsonHint = self.extractUsageErrorJSON(text: text) { + return jsonHint + } + + let lower = text.lowercased() + let compact = lower.filter { !$0.isWhitespace } + if lower.contains("do you trust the files in this folder?"), !lower.contains("current session") { + let folder = self.extractFirst( + pattern: #"Do you trust the files in this folder\?\s*(?:\r?\n)+\s*([^\r\n]+)"#, + text: text) + let folderHint = folder.flatMap { value in + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + } + if let folderHint { + return """ + Claude CLI is waiting for a folder trust prompt (\(folderHint)). CodexBar tries to auto-accept this, \ + but if it keeps appearing run: `cd "\(folderHint)" && claude` and choose “Yes, proceed”, then retry. + """ + } + return """ + Claude CLI is waiting for a folder trust prompt. CodexBar tries to auto-accept this, but if it keeps \ + appearing open `claude` once, choose “Yes, proceed”, then retry. + """ + } + if lower.contains("token_expired") || lower.contains("token has expired") { + return "Claude CLI token expired. Run `claude login` to refresh." + } + if lower.contains("authentication_error") { + return "Claude CLI authentication error. Run `claude login`." + } + if lower.contains("rate_limit_error") + || lower.contains("rate limited") + || compact.contains("ratelimited") + { + return "Claude CLI usage endpoint is rate limited right now. Please try again later." + } + if self.isSubscriptionNoticeOnly(text: text) { + return self.subscriptionQuotaUnavailableDescription + } + if lower.contains("failed to load usage data") { + return "Claude CLI could not load usage data. Open the CLI and retry `/usage`." + } + if compact.contains("failedtoloadusagedata") { + return "Claude CLI could not load usage data. Open the CLI and retry `/usage`." + } + return nil + } + + private static func isUsageStillLoading(text: String) -> Bool { + let normalized = TextParsing.stripANSICodes(text).lowercased().filter { !$0.isWhitespace } + guard normalized.contains("loadingusage") else { return false } + return !self.usageCaptureHasSessionValue(normalized) && self.allPercents(text).isEmpty + } + + /// Returns true when the text contains only a subscription notice with no session/weekly quota data. + /// CLI 2.1+ can return "You are currently using your subscription to power your Claude Code usage" + /// which lacks the "Current session" / "Current week" labels and percentage values needed for quota display. + /// A PTY capture may contain both an intermediate "Loading usage data…" panel and the final subscription + /// notice; `loadingusage` is not treated as quota data in this check so mixed captures surface correctly. + private static func isSubscriptionNoticeOnly(text: String) -> Bool { + let normalized = text.lowercased().filter { !$0.isWhitespace } + guard normalized.contains("currentlyusingyoursubscription") else { return false } + guard normalized.contains("claudecodeusage") else { return false } + // Only real session/week labels and actual percentage values count as quota data. + // `loadingusage` is not quota data — a mixed loading+subscription PTY capture should + // surface the subscription error, not a still-loading stall. + let hasQuotaData = normalized.contains("currentsession") || normalized.contains("currentweek") + || normalized.contains("%used") || normalized.contains("%left") || normalized.contains("%remaining") + || normalized.contains("%available") + return !hasQuotaData + } + + public static func isSubscriptionQuotaUnavailableDescription(_ text: String?) -> Bool { + text?.localizedCaseInsensitiveContains("subscription notice without session quota data") == true + } + + /// Collect remaining percentages in the order they appear; used as a backup when labels move/rename. + private static func allPercents(_ text: String) -> [Int] { + let lines = text.components(separatedBy: .newlines) + let normalized = text.lowercased().filter { !$0.isWhitespace } + let hasUsageWindows = normalized.contains("currentsession") || normalized.contains("currentweek") + let hasLoading = normalized.contains("loadingusage") + let hasUsagePercentKeywords = normalized.contains("used") || normalized.contains("left") + || normalized.contains("remaining") || normalized.contains("available") + let loadingOnly = hasLoading && !hasUsageWindows + guard hasUsageWindows || hasLoading else { return [] } + if loadingOnly { + return [] + } + guard hasUsagePercentKeywords else { return [] } + + // Keep this strict to avoid matching Claude's status-line context meter (e.g. "0%") as session usage when the + // /usage panel is still rendering. + return lines.compactMap { self.percentFromLine($0, assumeRemainingWhenUnclear: false) } + } + + /// Attempts to isolate the most recent /usage panel output from a PTY capture. + /// The Claude TUI draws a "Settings: … Usage …" header; we slice from its last occurrence to avoid earlier screen + /// fragments (like the status bar) contaminating percent scraping. + private static func trimToLatestUsagePanel(_ text: String) -> String? { + guard let settingsRange = text.range(of: "Settings:", options: [.caseInsensitive, .backwards]) else { + return nil + } + let tail = text[settingsRange.lowerBound...] + guard tail.range(of: "Usage", options: .caseInsensitive) != nil else { return nil } + let lower = tail.lowercased() + let hasPercent = lower.contains("%") + let hasUsageWords = lower.contains("used") || lower.contains("left") || lower.contains("remaining") + || lower.contains("available") + let hasLoading = lower.contains("loading usage") + guard (hasPercent && hasUsageWords) || hasLoading else { return nil } + return String(tail) + } + + private static func extractReset(labelSubstring: String, context: LabelSearchContext) -> String? { + let lines = context.lines + let label = self.normalizedForLabelSearch(labelSubstring) + for (idx, line) in lines.enumerated() { + let normalizedLine = context.normalizedLines[idx] + guard self.matchesLabel( + line: line, + normalizedLine: normalizedLine, + labelSubstring: labelSubstring, + normalizedLabel: label) + else { continue } + + let window = lines.dropFirst(idx).prefix(14) + for candidate in window { + if self.crossesLabelBoundary( + line: candidate, + labelSubstring: labelSubstring, + normalizedLabel: label) + { + break + } + if let reset = self.resetFromLine(candidate) { + return reset + } + } + } + return nil + } + + private static func matchesLabel( + line: String, + normalizedLine: String, + labelSubstring: String, + normalizedLabel: String) -> Bool + { + guard let expectedModel = self.weeklyModelName(from: labelSubstring) else { + return normalizedLine.contains(normalizedLabel) + } + guard let actualModel = self.weeklyModelName(from: line) else { return false } + return self.normalizedForLabelSearch(actualModel) == self.normalizedForLabelSearch(expectedModel) + } + + private static func crossesLabelBoundary( + line: String, + labelSubstring: String, + normalizedLabel: String) -> Bool + { + let normalizedLine = self.normalizedForLabelSearch(line) + guard normalizedLine.hasPrefix("current") else { return false } + guard let expectedModel = self.weeklyModelName(from: labelSubstring) else { + return !normalizedLine.contains(normalizedLabel) + } + guard let actualModel = self.weeklyModelName(from: line) else { return true } + return self.normalizedForLabelSearch(actualModel) != self.normalizedForLabelSearch(expectedModel) + } + + private static func weeklyModelName(from line: String) -> String? { + guard let regex = try? NSRegularExpression( + pattern: #"current\s*week\s*\(([^)]+)\)"#, + options: [.caseInsensitive]) + else { return nil } + let range = NSRange(line.startIndex..= 2, + let modelRange = Range(match.range(at: 1), in: line) + else { return nil } + return String(line[modelRange]).trimmingCharacters(in: .whitespacesAndNewlines) + } + + private static func extractReset(labelSubstrings: [String], context: LabelSearchContext) -> String? { + for label in labelSubstrings { + if let value = self.extractReset(labelSubstring: label, context: context) { + return value + } + } + return nil + } + + private static func resetFromLine(_ line: String) -> String? { + let range = line.range(of: #"(?i)\bresets?\b"#, options: .regularExpression) + ?? line.range(of: #"\b(?:Reset|Resets)(?=[A-Z0-9])"#, options: .regularExpression) + guard let range else { return nil } + let raw = String(line[range.lowerBound...]).trimmingCharacters(in: .whitespacesAndNewlines) + return self.cleanResetLine(raw) + } + + private static func normalizedForLabelSearch(_ text: String) -> String { + String(text.lowercased().unicodeScalars.filter(CharacterSet.alphanumerics.contains)) + } + + /// Capture all "Reset"/"Resets" strings to surface in the menu. + private static func allResets(_ text: String) -> [String] { + let pat = #"\bResets?\b[^\r\n]*"# + guard let regex = try? NSRegularExpression(pattern: pat, options: [.caseInsensitive]) else { return [] } + let nsrange = NSRange(text.startIndex.. String { + // TTY capture sometimes appends a stray ")" at line ends; trim it to keep snapshots stable. + var cleaned = raw.trimmingCharacters(in: .whitespacesAndNewlines) + cleaned = cleaned.trimmingCharacters(in: CharacterSet(charactersIn: " )")) + cleaned = cleaned.replacingOccurrences( + of: #"(?i)\b([A-Za-z]{3}\s+\d{1,2})\s+t\s+(\d)"#, + with: "$1 at $2", + options: .regularExpression) + let openCount = cleaned.count(where: { $0 == "(" }) + let closeCount = cleaned.count(where: { $0 == ")" }) + if openCount > closeCount { + cleaned.append(")") + } + return cleaned + } + + /// Parses explicit-year resets at their stated occurrence; recurring forms resolve forward in any explicit + /// timezone. + public static func parseResetDate(from text: String?, now: Date = .init()) -> Date? { + self.parseResetDate(from: text, now: now, expectedWindow: nil) + } + + /// Quota surfaces may accept a recently stale occurrence when the next occurrence cannot belong to that window. + package static func parseResetDate( + from text: String?, + now: Date = .init(), + expectedWindow: TimeInterval?) -> Date? + { + guard let normalized = self.normalizeResetInput(text) else { return nil } + let (raw, timeZone) = normalized + + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = timeZone ?? TimeZone.current + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = formatter.timeZone + // Parse yearless Feb 29 independently of the current year, then resolve validated calendar occurrences. + formatter.defaultDate = calendar.date(from: DateComponents(year: 2000, month: 1, day: 1)) + + if let date = self.parseDate(raw, formats: Self.resetDateTimeWithExplicitYear, formatter: formatter) { + return self.resolveOccurrence( + self.explicitOccurrences(for: date, calendar: calendar), + now: now, + expectedWindow: expectedWindow) + } + if let date = self.parseDate(raw, formats: Self.resetDateTimeWithMinutes, formatter: formatter) { + let comps = calendar.dateComponents([.month, .day, .hour, .minute], from: date) + return self.resolveOccurrence( + self.yearlyOccurrences(for: comps, now: now, calendar: calendar), + now: now, + expectedWindow: expectedWindow) + } + if let date = self.parseDate(raw, formats: Self.resetDateTimeHourOnly, formatter: formatter) { + let comps = calendar.dateComponents([.month, .day, .hour], from: date) + return self.resolveOccurrence( + self.yearlyOccurrences(for: comps, now: now, calendar: calendar), + now: now, + expectedWindow: expectedWindow) + } + + if let time = self.parseDate(raw, formats: Self.resetTimeWithMinutes, formatter: formatter) { + let comps = calendar.dateComponents([.hour, .minute], from: time) + return self.resolveOccurrence( + self.dailyOccurrences(for: comps, now: now, calendar: calendar), + now: now, + expectedWindow: expectedWindow) + } + + guard let time = self.parseDate(raw, formats: Self.resetTimeHourOnly, formatter: formatter) else { return nil } + let comps = calendar.dateComponents([.hour], from: time) + return self.resolveOccurrence( + self.dailyOccurrences(for: comps, now: now, calendar: calendar), + now: now, + expectedWindow: expectedWindow) + } + + private static func explicitOccurrences(for date: Date, calendar: Calendar) -> [Date] { + let components = calendar.dateComponents([.year, .month, .day, .hour, .minute], from: date) + guard let year = components.year, + let month = components.month, + let day = components.day, + let hour = components.hour, + let minute = components.minute + else { return [] } + return self.exactLocalOccurrences( + DateComponents( + timeZone: calendar.timeZone, + year: year, + month: month, + day: day, + hour: hour, + minute: minute, + second: 0), + calendar: calendar) + } + + private static func resolveOccurrence( + _ candidates: [Date], + now: Date, + expectedWindow: TimeInterval?) -> Date? + { + let sortedCandidates = candidates.sorted() + guard let future = sortedCandidates.first(where: { $0 >= now }) else { return sortedCandidates.last } + guard let expectedWindow else { return future } + + let horizon = max(0, expectedWindow) + let past = sortedCandidates.last(where: { $0 < now }) + let pastIsPlausible = past.map { now.timeIntervalSince($0) <= horizon } ?? false + let futureIsPlausible = future.timeIntervalSince(now) <= horizon + return pastIsPlausible && !futureIsPlausible ? past : future + } + + private static func yearlyOccurrences( + for components: DateComponents, + now: Date, + calendar: Calendar) -> [Date] + { + guard let month = components.month, + let day = components.day, + let hour = components.hour + else { return [] } + let currentYear = calendar.component(.year, from: now) + // Eight years covers the Gregorian century gap between leap days (for example, 2096 to 2104). + return ((currentYear - 8)...(currentYear + 8)).flatMap { year in + self.exactLocalOccurrences( + DateComponents( + timeZone: calendar.timeZone, + year: year, + month: month, + day: day, + hour: hour, + minute: components.minute ?? 0, + second: 0), + calendar: calendar) + } + } + + private static func dailyOccurrences( + for components: DateComponents, + now: Date, + calendar: Calendar) -> [Date] + { + guard let hour = components.hour else { return [] } + let startOfToday = calendar.startOfDay(for: now) + return (-1...1).flatMap { offset -> [Date] in + guard let day = calendar.date(byAdding: .day, value: offset, to: startOfToday) else { return [] } + let dayComponents = calendar.dateComponents([.year, .month, .day], from: day) + guard let year = dayComponents.year, + let month = dayComponents.month, + let dayOfMonth = dayComponents.day + else { return [] } + return self.exactLocalOccurrences( + DateComponents( + timeZone: calendar.timeZone, + year: year, + month: month, + day: dayOfMonth, + hour: hour, + minute: components.minute ?? 0, + second: 0), + calendar: calendar) + } + } + + /// Returns both instants when a local wall time repeats during a daylight-saving transition. + private static func exactLocalOccurrences( + _ target: DateComponents, + calendar: Calendar) -> [Date] + { + guard let year = target.year, + let month = target.month, + let day = target.day, + let hour = target.hour, + let minute = target.minute + else { return [] } + guard let firstApproximation = calendar.date(from: target) else { return [] } + let approximationComponents = calendar.dateComponents( + [.year, .month, .day, .hour, .minute, .second], + from: firstApproximation) + guard approximationComponents.year == year, + approximationComponents.month == month, + approximationComponents.day == day, + approximationComponents.hour == hour, + approximationComponents.minute == minute, + approximationComponents.second == 0 + else { return [] } + let startOfDay = calendar.startOfDay(for: firstApproximation) + guard let searchStart = calendar.date(byAdding: .second, value: -1, to: startOfDay) else { return [] } + + var results: [Date] = [] + for policy in [Calendar.RepeatedTimePolicy.first, .last] { + guard let date = calendar.nextDate( + after: searchStart, + matching: target, + matchingPolicy: .strict, + repeatedTimePolicy: policy, + direction: .forward) + else { continue } + let resolved = calendar.dateComponents([.year, .month, .day, .hour, .minute, .second], from: date) + guard resolved.year == year, + resolved.month == month, + resolved.day == day, + resolved.hour == hour, + resolved.minute == minute, + resolved.second == 0, + !results.contains(date) + else { continue } + results.append(date) + } + return results + } + + private static let resetTimeWithMinutes = ["h:mma", "h:mm a", "HH:mm", "H:mm"] + private static let resetTimeHourOnly = ["ha", "h a"] + + private static let resetDateTimeWithExplicitYear = [ + "MMM d, yyyy, h:mma", + "MMM d, yyyy, h:mm a", + "MMM d, yyyy, HH:mm", + "MMM d, yyyy h:mma", + "MMM d, yyyy h:mm a", + "MMM d, yyyy HH:mm", + "MMM d yyyy h:mma", + "MMM d yyyy h:mm a", + "MMM d yyyy HH:mm", + "MMM d, yyyy, ha", + "MMM d, yyyy, h a", + "MMM d, yyyy, HH", + "MMM d, yyyy ha", + "MMM d, yyyy h a", + "MMM d, yyyy HH", + "MMM d yyyy ha", + "MMM d yyyy h a", + "MMM d yyyy HH", + ] + + private static let resetDateTimeWithMinutes = [ + "MMM d, h:mma", + "MMM d, h:mm a", + "MMM d h:mma", + "MMM d h:mm a", + "MMM d, HH:mm", + "MMM d HH:mm", + ] + + private static let resetDateTimeHourOnly = [ + "MMM d, ha", + "MMM d, h a", + "MMM d ha", + "MMM d h a", + ] + + private static func normalizeResetInput(_ text: String?) -> (String, TimeZone?)? { + guard var raw = text?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { return nil } + raw = raw.replacingOccurrences(of: #"(?i)^resets?:?\s*"#, with: "", options: .regularExpression) + raw = raw.replacingOccurrences(of: " at ", with: " ", options: .caseInsensitive) + raw = raw.replacingOccurrences(of: #"(?i)\b([A-Za-z]{3})(\d)"#, with: "$1 $2", options: .regularExpression) + raw = raw.replacingOccurrences(of: #",(\d)"#, with: ", $1", options: .regularExpression) + raw = raw.replacingOccurrences(of: #"(?i)(\d)at(?=\d)"#, with: "$1 ", options: .regularExpression) + raw = raw.replacingOccurrences( + of: #"(?<=\d)\.(\d{2})\b"#, + with: ":$1", + options: .regularExpression) + + let timeZone = self.extractTimeZone(from: &raw) + raw = raw.replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespacesAndNewlines) + return raw.isEmpty ? nil : (raw, timeZone) + } + + private static func extractTimeZone(from text: inout String) -> TimeZone? { + guard let tzRange = text.range(of: #"\(([^)]+)\)"#, options: .regularExpression) else { return nil } + let tzID = String(text[tzRange]).trimmingCharacters(in: CharacterSet(charactersIn: "() ")) + text.removeSubrange(tzRange) + text = text.trimmingCharacters(in: .whitespacesAndNewlines) + return TimeZone(identifier: tzID) + } + + private static func parseDate(_ text: String, formats: [String], formatter: DateFormatter) -> Date? { + for pattern in formats { + formatter.dateFormat = pattern + if let date = formatter.date(from: text) { + return date + } + } + return nil + } + + /// Extract login/plan string from CLI output. + private static func extractLoginMethod(text: String) -> String? { + guard !text.isEmpty else { return nil } + if let explicit = self.extractFirst(pattern: #"(?i)login\s+method:\s*(.+)"#, text: text) { + return ClaudePlan.cliCompatibilityLoginMethod(self.cleanPlan(explicit)) + } + // Capture any "Claude <...>" phrase (e.g., Max/Pro/Ultra/Team) to avoid future plan-name churn. + // Strip any leading ANSI that may have survived (rare) before matching. + // Use horizontal whitespace ([ \t]) rather than \s so the match stays on a single rendered line: + // \s spans newlines, which let "…use Claude" bridge into the /usage panel's "d → today" hint and + // produced a bogus "Dtoday" plan label after ANSI stripping glued the lines together. + let planPattern = #"(?i)(claude[ \t]+[a-z0-9][a-z0-9 \t._-]{0,24})"# + var candidates: [String] = [] + if let regex = try? NSRegularExpression(pattern: planPattern, options: []) { + let nsrange = NSRange(text.startIndex..= 2, + let r = Range(match.range(at: 1), in: text) else { return } + let raw = String(text[r]) + let val = ClaudePlan.cliCompatibilityLoginMethod(Self.cleanPlan(raw)) ?? Self.cleanPlan(raw) + candidates.append(val) + } + } + if let plan = candidates.first(where: { cand in + let lower = cand.lowercased() + return !lower.contains("code v") && !lower.contains("code version") && !lower.contains("code") + }) { + return plan + } + return nil + } + + /// Strips ANSI and stray bracketed codes like "[22m" that can survive CLI output. + private static func cleanPlan(_ text: String) -> String { + UsageFormatter.cleanPlanName(text) + } + + private static func dumpIfNeeded(enabled: Bool, reason: String, usage: String, status: String?) { + guard enabled else { return } + let stamp = ISO8601DateFormatter().string(from: Date()) + var parts = [ + "=== Claude parse dump @ \(stamp) ===", + "Reason: \(reason)", + "", + "--- usage (clean) ---", + usage, + "", + ] + if let status { + parts.append(contentsOf: [ + "--- status (raw/optional) ---", + status, + "", + ]) + } + let body = parts.joined(separator: "\n") + Task { @MainActor in self.recordDump(body) } + } + + // MARK: - Dump storage (in-memory ring buffer) + + @MainActor private static var recentDumps: [String] = [] + + @MainActor private static func recordDump(_ text: String) { + if self.recentDumps.count >= 5 { + self.recentDumps.removeFirst() + } + self.recentDumps.append(text) + } + + public static func latestDumps() async -> String { + await MainActor.run { + let result = Self.recentDumps.joined(separator: "\n\n---\n\n") + return result.isEmpty ? "No Claude parse dumps captured yet." : result + } + } + + #if DEBUG + public static func _replaceDumpsForTesting(_ dumps: [String]) async { + await MainActor.run { + self.recentDumps = dumps + } + } + #endif + + private static func extractUsageErrorJSON(text: String) -> String? { + let pattern = #"Failed\s*to\s*load\s*usage\s*data:\s*(\{.*\})"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: [.dotMatchesLineSeparators]) else { + return nil + } + let range = NSRange(text.startIndex..= 2, + let jsonRange = Range(match.range(at: 1), in: text) + else { + return nil + } + + let jsonString = String(text[jsonRange]) + let compactJSON = jsonString.replacingOccurrences(of: "\r", with: "").replacingOccurrences(of: "\n", with: "") + let data = (compactJSON.isEmpty ? jsonString : compactJSON).data(using: .utf8) + guard let data, + let payload = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let error = payload["error"] as? [String: Any] + else { + return nil + } + + let message = (error["message"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + let details = error["details"] as? [String: Any] + let code = (details?["error_code"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + let type = (error["type"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + + if type == "rate_limit_error" { + return "Claude CLI usage endpoint is rate limited right now. Please try again later." + } + + var parts: [String] = [] + if let message, !message.isEmpty { + parts.append(message) + } + if let code, !code.isEmpty { + parts.append("(\(code))") + } + + guard !parts.isEmpty else { return nil } + let hint = parts.joined(separator: " ") + + if let code, code.lowercased().contains("token") { + return "\(hint). Run `claude login` to refresh." + } + return "Claude CLI error: \(hint)" + } + + // MARK: - Process helpers + + private static func resolvedBinaryPath( + binaryName: String, + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + if binaryName.contains("/") { + return binaryName + } + return ClaudeCLIResolver.resolvedBinaryPath(environment: environment) + } + + private static func isBinaryAvailable(_ binaryPathOrName: String?) -> Bool { + guard let binaryPathOrName else { return false } + return FileManager.default.isExecutableFile(atPath: binaryPathOrName) + || TTYCommandRunner.which(binaryPathOrName) != nil + } + + static func probeWorkingDirectoryURL() -> URL { + let fm = FileManager.default + let base = fm.urls(for: .applicationSupportDirectory, in: .userDomainMask).first ?? fm.temporaryDirectory + let dir = base + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("ClaudeProbe", isDirectory: true) + do { + try fm.createDirectory(at: dir, withIntermediateDirectories: true) + return dir + } catch { + return fm.temporaryDirectory + } + } + + static func preparedProbeWorkingDirectoryURL() -> URL { + let directory = self.probeWorkingDirectoryURL() + do { + try self.prepareProbeWorkingDirectory(at: directory) + } catch { + Self.log.warning( + "Claude probe local settings unavailable", + metadata: ["error": error.localizedDescription]) + } + return directory + } + + static func prepareProbeWorkingDirectory(at directory: URL, fileManager fm: FileManager = .default) throws { + try fm.createDirectory(at: directory, withIntermediateDirectories: true) + let claudeDirectory = directory.appendingPathComponent(".claude", isDirectory: true) + try fm.createDirectory(at: claudeDirectory, withIntermediateDirectories: true) + + let settingsURL = claudeDirectory.appendingPathComponent("settings.local.json") + var settings = (try? self.readSettingsObject(from: settingsURL, fileManager: fm)) ?? [:] + settings["disableDeepLinkRegistration"] = "disable" + let data = try JSONSerialization.data( + withJSONObject: settings, + options: [.prettyPrinted, .sortedKeys]) + try data.write(to: settingsURL, options: .atomic) + } + + private static func readSettingsObject(from url: URL, fileManager fm: FileManager) throws -> [String: Any] { + guard fm.fileExists(atPath: url.path) else { + return [:] + } + let data = try Data(contentsOf: url) + guard !data.isEmpty else { + return [:] + } + let object = try JSONSerialization.jsonObject(with: data) + return object as? [String: Any] ?? [:] + } + + /// Run claude CLI inside a PTY so we can respond to interactive permission prompts. + private static func capture(subcommand: String, binary: String, timeout: TimeInterval) async throws -> String { + let stopOnSubstrings = subcommand == "/usage" + ? [ + "Failed to load usage data", + "failed to load usage data", + "Failedto loadusagedata", + "failedtoloadusagedata", + ] + : [] + let idleTimeout: TimeInterval? = subcommand == "/usage" ? nil : 3.0 + let sendEnterEvery: TimeInterval? = subcommand == "/usage" ? 0.8 : nil + let stopWhenNormalized: (@Sendable (String) -> Bool)? = subcommand == "/usage" + ? { @Sendable normalizedScan in + Self.usageCaptureHasSessionValue(normalizedScan) + || Self.usageCaptureHasSubscriptionNotice(normalizedScan) + } + : nil + do { + return try await ClaudeCLISession.current.capture( + subcommand: subcommand, + binary: binary, + timeout: timeout, + idleTimeout: idleTimeout, + stopOnSubstrings: stopOnSubstrings, + stopWhenNormalized: stopWhenNormalized, + settleAfterStop: subcommand == "/usage" ? 2.0 : 0.25, + sendEnterEvery: sendEnterEvery) + } catch ClaudeCLISession.SessionError.processExited { + await ClaudeCLISession.current.reset() + throw ClaudeStatusProbeError.timedOut + } catch ClaudeCLISession.SessionError.timedOut { + await ClaudeCLISession.current.reset() + throw ClaudeStatusProbeError.timedOut + } catch ClaudeCLISession.SessionError.launchFailed(_) { + throw ClaudeStatusProbeError.claudeNotInstalled + } catch { + await ClaudeCLISession.current.reset() + throw error + } + } + + private static func usageCaptureHasSessionValue(_ normalizedText: String) -> Bool { + guard let labelRange = normalizedText.range(of: "currentsession") else { return false } + let tail = normalizedText[labelRange.upperBound...] + return tail.range(of: #"[0-9]{1,3}(?:\.[0-9]+)?%"#, options: .regularExpression) != nil + } + + private static func usageCaptureHasSubscriptionNotice(_ normalizedText: String) -> Bool { + normalizedText.contains("currentlyusingyoursubscription") + && normalizedText.contains("claudecodeusage") + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountList.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountList.swift new file mode 100644 index 0000000..097b118 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountList.swift @@ -0,0 +1,216 @@ +import CoreFoundation +import Foundation + +/// Strictly parsed result of `cswap --list --json` (schema v1). +/// +/// Only the fields allow-listed in `docs/claude-multi-account-and-status-items.md` +/// are decoded: slot number, display email, active state, usage status, and the +/// 5-hour/7-day windows (percent + reset timestamp). Everything else in the +/// payload is ignored; unknown schema versions and partial top-level shapes are +/// rejected. +public struct ClaudeSwapAccountList: Equatable, Sendable { + public let activeAccountNumber: Int? + public let accounts: [ClaudeSwapAccountRow] + + public init(activeAccountNumber: Int?, accounts: [ClaudeSwapAccountRow]) { + self.activeAccountNumber = activeAccountNumber + self.accounts = accounts + } +} + +public struct ClaudeSwapAccountRow: Equatable, Sendable { + public let number: Int + /// Display-only sensitive value; never logged or persisted. + public let email: String + public let isActive: Bool + public let usageStatus: ClaudeSwapUsageStatus + public let fiveHour: ClaudeSwapUsageWindow? + public let sevenDay: ClaudeSwapUsageWindow? + + public init( + number: Int, + email: String, + isActive: Bool, + usageStatus: ClaudeSwapUsageStatus, + fiveHour: ClaudeSwapUsageWindow?, + sevenDay: ClaudeSwapUsageWindow?) + { + self.number = number + self.email = email + self.isActive = isActive + self.usageStatus = usageStatus + self.fiveHour = fiveHour + self.sevenDay = sevenDay + } +} + +public struct ClaudeSwapUsageWindow: Equatable, Sendable { + public let usedPercent: Double + public let resetsAt: Date? + + public init(usedPercent: Double, resetsAt: Date?) { + self.usedPercent = usedPercent + self.resetsAt = resetsAt + } +} + +/// The sentinel statuses `cswap` emits per account. Unknown values from newer +/// claude-swap releases are preserved rather than failing the whole payload. +public enum ClaudeSwapUsageStatus: Equatable, Sendable { + case ok + case tokenExpired + case apiKey + case keychainUnavailable + case noCredentials + case unavailable + case unknown(String) + + init(rawValue: String) { + switch rawValue { + case "ok": self = .ok + case "token_expired": self = .tokenExpired + case "api_key": self = .apiKey + case "keychain_unavailable": self = .keychainUnavailable + case "no_credentials": self = .noCredentials + case "unavailable": self = .unavailable + default: self = .unknown(rawValue) + } + } +} + +public enum ClaudeSwapListParserError: LocalizedError, Equatable, Sendable { + case notJSONObject + case missingSchemaVersion + case unsupportedSchemaVersion(Int) + case reportedError(type: String, message: String) + case malformedShape(String) + + public var errorDescription: String? { + switch self { + case .notJSONObject: + "claude-swap returned output that is not a JSON object." + case .missingSchemaVersion: + "claude-swap output has no schemaVersion field." + case let .unsupportedSchemaVersion(version): + "claude-swap output uses unsupported schema version \(version); CodexBar supports version 1." + case let .reportedError(type, message): + "claude-swap reported \(type): \(message)" + case let .malformedShape(details): + "claude-swap output is malformed: \(details)" + } + } +} + +public enum ClaudeSwapListParser { + public static let supportedSchemaVersion = 1 + + public static func parse(_ data: Data) throws -> ClaudeSwapAccountList { + let raw: Any + do { + raw = try JSONSerialization.jsonObject(with: data) + } catch { + throw ClaudeSwapListParserError.notJSONObject + } + guard let object = raw as? [String: Any] else { + throw ClaudeSwapListParserError.notJSONObject + } + guard let schemaVersion = object["schemaVersion"] as? Int else { + throw ClaudeSwapListParserError.missingSchemaVersion + } + guard schemaVersion == self.supportedSchemaVersion else { + throw ClaudeSwapListParserError.unsupportedSchemaVersion(schemaVersion) + } + if let errorObject = object["error"] as? [String: Any] { + throw ClaudeSwapListParserError.reportedError( + type: errorObject["type"] as? String ?? "Error", + message: errorObject["message"] as? String ?? "unknown error") + } + guard let rawAccounts = object["accounts"] as? [Any] else { + throw ClaudeSwapListParserError.malformedShape("missing accounts array") + } + guard let rawActiveAccountNumber = object["activeAccountNumber"] else { + throw ClaudeSwapListParserError.malformedShape("missing activeAccountNumber") + } + let activeAccountNumber: Int? = switch rawActiveAccountNumber { + case is NSNull: nil + case let number as Int where number > 0: number + default: + throw ClaudeSwapListParserError.malformedShape("activeAccountNumber is not a numeric slot or null") + } + var seenSlots: Set = [] + let accounts = try rawAccounts.map { rawRow -> ClaudeSwapAccountRow in + guard let row = rawRow as? [String: Any] else { + throw ClaudeSwapListParserError.malformedShape("account row is not an object") + } + let account = try self.parseRow(row) + guard seenSlots.insert(account.number).inserted else { + throw ClaudeSwapListParserError.malformedShape("duplicate account slot \(account.number)") + } + return account + } + let activeSlots = accounts.filter(\.isActive).map(\.number) + guard activeSlots == (activeAccountNumber.map { [$0] } ?? []) else { + throw ClaudeSwapListParserError.malformedShape("active account fields disagree") + } + return ClaudeSwapAccountList(activeAccountNumber: activeAccountNumber, accounts: accounts) + } + + private static func parseRow(_ row: [String: Any]) throws -> ClaudeSwapAccountRow { + guard let number = row["number"] as? Int else { + throw ClaudeSwapListParserError.malformedShape("account row has no numeric slot") + } + guard number > 0 else { + throw ClaudeSwapListParserError.malformedShape("account slot must be positive") + } + guard let isActive = row["active"] as? Bool else { + throw ClaudeSwapListParserError.malformedShape("slot \(number) has no active flag") + } + guard let rawStatus = row["usageStatus"] as? String else { + throw ClaudeSwapListParserError.malformedShape("slot \(number) has no usageStatus") + } + let email = (row["email"] as? String ?? "").trimmingCharacters(in: .whitespacesAndNewlines) + let usage = row["usage"] as? [String: Any] + return try ClaudeSwapAccountRow( + number: number, + email: email, + isActive: isActive, + usageStatus: ClaudeSwapUsageStatus(rawValue: rawStatus), + fiveHour: self.parseWindow(usage?["fiveHour"], slot: number, name: "fiveHour"), + sevenDay: self.parseWindow(usage?["sevenDay"], slot: number, name: "sevenDay")) + } + + private static func parseWindow(_ raw: Any?, slot: Int, name: String) throws -> ClaudeSwapUsageWindow? { + guard let raw else { return nil } + guard let window = raw as? [String: Any] else { + throw ClaudeSwapListParserError.malformedShape("slot \(slot) \(name) window is not an object") + } + guard let pct = Self.finiteDouble(window["pct"]) else { + throw ClaudeSwapListParserError.malformedShape("slot \(slot) \(name) percent is not a finite number") + } + var resetsAt: Date? + if let rawResetsAt = window["resetsAt"] { + guard let text = rawResetsAt as? String, let date = Self.parseTimestamp(text) else { + throw ClaudeSwapListParserError.malformedShape("slot \(slot) \(name) resetsAt is not a timestamp") + } + resetsAt = date + } + return ClaudeSwapUsageWindow(usedPercent: min(max(pct, 0), 100), resetsAt: resetsAt) + } + + private static func finiteDouble(_ raw: Any?) -> Double? { + guard let number = raw as? NSNumber else { return nil } + // JSON booleans bridge to NSNumber too; only accept genuine numbers. + guard CFGetTypeID(number) != CFBooleanGetTypeID() else { return nil } + let value = number.doubleValue + return value.isFinite ? value : nil + } + + private static func parseTimestamp(_ text: String) -> Date? { + let withFraction = ISO8601DateFormatter() + withFraction.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = withFraction.date(from: text) { return date } + let plain = ISO8601DateFormatter() + plain.formatOptions = [.withInternetDateTime] + return plain.date(from: text) + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountProjection.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountProjection.swift new file mode 100644 index 0000000..6fa1366 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountProjection.swift @@ -0,0 +1,102 @@ +import Foundation + +/// Projects parsed `cswap --list --json` rows into the provider-neutral +/// account snapshot consumed by menus. Identity uses the source-issued numeric +/// slot (`claude-swap:`), never email or credential-derived values. +public enum ClaudeSwapAccountProjection { + public static let sourceName = "claude-swap" + public static let sourceLabel = "claude-swap" + static let fiveHourWindowMinutes = 5 * 60 + static let sevenDayWindowMinutes = 7 * 24 * 60 + + public static func accountSnapshots( + from list: ClaudeSwapAccountList, + now: Date = Date()) -> [ProviderAccountUsageSnapshot] + { + let ordered = list.accounts.sorted { lhs, rhs in + if lhs.isActive != rhs.isActive { return lhs.isActive } + return lhs.number < rhs.number + } + return ordered.map { row in + ProviderAccountUsageSnapshot( + id: ProviderAccountIdentity(source: self.sourceName, opaqueID: String(row.number)), + provider: .claude, + displayLabel: self.displayLabel(for: row), + isActive: row.isActive, + canActivate: !row.isActive && self.canActivate(row), + snapshot: self.usageSnapshot(for: row, now: now), + error: self.errorText(for: row), + sourceLabel: self.sourceLabel) + } + } + + public static func displayError( + accountError: String?, + adapterError: String?, + switchError: String? = nil) -> String? + { + switchError.map { "Account switch failed: \($0)" } + ?? accountError + ?? adapterError.map { "Showing the last successful update: \($0)" } + } + + static func displayLabel(for row: ClaudeSwapAccountRow) -> String { + row.email.isEmpty ? "Account \(row.number)" : row.email + } + + private static func usageSnapshot(for row: ClaudeSwapAccountRow, now: Date) -> UsageSnapshot? { + guard row.usageStatus == .ok else { return nil } + let primary = row.fiveHour.map { window in + RateWindow( + usedPercent: window.usedPercent, + windowMinutes: self.fiveHourWindowMinutes, + resetsAt: window.resetsAt, + resetDescription: nil) + } + let secondary = row.sevenDay.map { window in + RateWindow( + usedPercent: window.usedPercent, + windowMinutes: self.sevenDayWindowMinutes, + resetsAt: window.resetsAt, + resetDescription: nil) + } + guard primary != nil || secondary != nil else { return nil } + return UsageSnapshot( + primary: primary, + secondary: secondary, + updatedAt: now, + identity: ProviderIdentitySnapshot( + providerID: .claude, + accountEmail: self.displayLabel(for: row), + accountOrganization: nil, + loginMethod: self.sourceLabel)) + } + + private static func errorText(for row: ClaudeSwapAccountRow) -> String? { + switch row.usageStatus { + case .ok: + row.fiveHour == nil && row.sevenDay == nil ? "No usage windows reported." : nil + case .tokenExpired: + "Token expired. Switch to this account in claude-swap to refresh it." + case .apiKey: + "API-key account; subscription usage is unavailable." + case .keychainUnavailable: + "claude-swap could not read the active account's Keychain entry." + case .noCredentials: + "No stored credentials for this account slot." + case .unavailable: + "Usage fetch failed." + case let .unknown(raw): + "Unrecognized claude-swap status: \(raw)" + } + } + + private static func canActivate(_ row: ClaudeSwapAccountRow) -> Bool { + switch row.usageStatus { + case .ok, .apiKey, .unavailable: + true + case .tokenExpired, .keychainUnavailable, .noCredentials, .unknown: + false + } + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountReader.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountReader.swift new file mode 100644 index 0000000..5896400 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountReader.swift @@ -0,0 +1,141 @@ +import Foundation + +public enum ClaudeSwapAccountReaderError: LocalizedError, Sendable { + case executablePathNotConfigured + case outputTooLarge(byteCount: Int) + + public var errorDescription: String? { + switch self { + case .executablePathNotConfigured: + "No claude-swap executable path is configured." + case let .outputTooLarge(byteCount): + "claude-swap produced \(byteCount) bytes of output; refusing to parse more than " + + "\(ClaudeSwapAccountReader.maxOutputBytes)." + } + } +} + +/// Bounded adapter over the external `claude-swap` executable. +/// +/// Executes only the fixed list and explicit switch argument arrays (never a +/// shell or config-defined passthrough arguments), per the contracts in +/// `docs/claude-multi-account-and-status-items.md`. CodexBar never reads +/// claude-swap or Claude Code credential storage; the subprocess is solely +/// responsible for its own credential access. +public enum ClaudeSwapAccountReader { + public static let maxOutputBytes = 262_144 + public static let defaultTimeout: TimeInterval = 30 + + public static func readAccountList( + executablePath: String, + timeout: TimeInterval = ClaudeSwapAccountReader.defaultTimeout) async throws -> ClaudeSwapAccountList + { + // Handled claude-swap failures print a schema-v1 error envelope to stdout + // and exit non-zero, so non-zero exits still parse; the parser surfaces + // the envelope as `ClaudeSwapListParserError.reportedError`. + let result = try await self.run( + executablePath: executablePath, + arguments: ["--list", "--json"], + timeout: timeout, + acceptsNonZeroExit: true, + label: "claude-swap list") + return try ClaudeSwapListParser.parse(Data(result.utf8)) + } + + /// Activates one source-issued numeric slot through claude-swap. + /// + /// The external tool owns the credential transaction. CodexBar supplies no + /// credentials and accepts no config-defined arguments. + public static func switchAccount( + executablePath: String, + accountNumber: Int) async throws + -> ClaudeSwapAccountSwitchResult + { + guard accountNumber > 0 else { + throw ClaudeSwapSwitchParserError.malformedShape("requested account slot must be positive") + } + let result = try await self.runCredentialTransaction( + executablePath: executablePath, + arguments: ["--switch-to", String(accountNumber), "--json"], + acceptsNonZeroExit: true, + label: "claude-swap switch") + let parsed = try ClaudeSwapSwitchParser.parse(Data(result.utf8)) + guard parsed.toAccountNumber == accountNumber else { + throw ClaudeSwapSwitchParserError.mismatchedTarget( + expected: accountNumber, + actual: parsed.toAccountNumber) + } + return parsed + } + + /// Best-effort version probe (`cswap --version` prints `cswap `). + public static func readVersion( + executablePath: String, + timeout: TimeInterval = 10) async -> String? + { + guard let output = try? await self.run( + executablePath: executablePath, + arguments: ["--version"], + timeout: timeout, + label: "claude-swap version") + else { + return nil + } + let trimmed = output.trimmingCharacters(in: .whitespacesAndNewlines) + guard let version = trimmed.split(whereSeparator: \.isWhitespace).last, !version.isEmpty else { + return nil + } + return String(version) + } + + public static func resolvedExecutablePath(_ configuredPath: String) throws -> String { + let trimmed = configuredPath.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { + throw ClaudeSwapAccountReaderError.executablePathNotConfigured + } + return (trimmed as NSString).expandingTildeInPath + } + + private static func run( + executablePath: String, + arguments: [String], + timeout: TimeInterval, + acceptsNonZeroExit: Bool = false, + label: String) async throws -> String + { + try Task.checkCancellation() + let binary = try self.resolvedExecutablePath(executablePath) + let result = try await SubprocessRunner.run( + binary: binary, + arguments: arguments, + environment: ProcessInfo.processInfo.environment, + timeout: timeout, + acceptsNonZeroExit: acceptsNonZeroExit, + label: label) + guard result.stdout.utf8.count <= self.maxOutputBytes else { + throw ClaudeSwapAccountReaderError.outputTooLarge(byteCount: result.stdout.utf8.count) + } + return result.stdout + } + + /// Once launched, a credential mutation must reach the external tool's + /// natural exit; caller cancellation and read-probe timeouts must not kill it. + private static func runCredentialTransaction( + executablePath: String, + arguments: [String], + acceptsNonZeroExit: Bool, + label: String) async throws -> String + { + let binary = try self.resolvedExecutablePath(executablePath) + let result = try await SubprocessRunner.runToCompletion( + binary: binary, + arguments: arguments, + environment: ProcessInfo.processInfo.environment, + acceptsNonZeroExit: acceptsNonZeroExit, + label: label) + guard result.stdout.utf8.count <= self.maxOutputBytes else { + throw ClaudeSwapAccountReaderError.outputTooLarge(byteCount: result.stdout.utf8.count) + } + return result.stdout + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountSwitch.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountSwitch.swift new file mode 100644 index 0000000..de539ed --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeSwap/ClaudeSwapAccountSwitch.swift @@ -0,0 +1,117 @@ +import CoreFoundation +import Foundation + +/// Allow-listed result of `cswap --switch-to --json` (schema v1). +public struct ClaudeSwapAccountSwitchResult: Equatable, Sendable { + public let switched: Bool + public let fromAccountNumber: Int? + public let toAccountNumber: Int + public let reason: String + + public init(switched: Bool, fromAccountNumber: Int?, toAccountNumber: Int, reason: String) { + self.switched = switched + self.fromAccountNumber = fromAccountNumber + self.toAccountNumber = toAccountNumber + self.reason = reason + } +} + +public enum ClaudeSwapSwitchParserError: LocalizedError, Equatable, Sendable { + case notJSONObject + case missingSchemaVersion + case unsupportedSchemaVersion(Int) + case reportedError(type: String, message: String) + case malformedShape(String) + case mismatchedTarget(expected: Int, actual: Int) + + public var errorDescription: String? { + switch self { + case .notJSONObject: + "claude-swap returned switch output that is not a JSON object." + case .missingSchemaVersion: + "claude-swap switch output has no schemaVersion field." + case let .unsupportedSchemaVersion(version): + "claude-swap switch output uses unsupported schema version \(version); CodexBar supports version 1." + case let .reportedError(type, message): + "claude-swap reported \(type): \(message)" + case let .malformedShape(details): + "claude-swap switch output is malformed: \(details)" + case let .mismatchedTarget(expected, actual): + "claude-swap reported account slot \(actual) after CodexBar requested slot \(expected)." + } + } +} + +public enum ClaudeSwapSwitchParser { + public static func parse(_ data: Data) throws -> ClaudeSwapAccountSwitchResult { + let raw: Any + do { + raw = try JSONSerialization.jsonObject(with: data) + } catch { + throw ClaudeSwapSwitchParserError.notJSONObject + } + guard let object = raw as? [String: Any] else { + throw ClaudeSwapSwitchParserError.notJSONObject + } + guard let schemaVersion = self.integer(object["schemaVersion"]) else { + throw ClaudeSwapSwitchParserError.missingSchemaVersion + } + guard schemaVersion == ClaudeSwapListParser.supportedSchemaVersion else { + throw ClaudeSwapSwitchParserError.unsupportedSchemaVersion(schemaVersion) + } + if let errorObject = object["error"] as? [String: Any] { + throw ClaudeSwapSwitchParserError.reportedError( + type: errorObject["type"] as? String ?? "Error", + message: errorObject["message"] as? String ?? "unknown error") + } + guard let switched = object["switched"] as? Bool else { + throw ClaudeSwapSwitchParserError.malformedShape("missing switched flag") + } + guard let reason = (object["reason"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines), + !reason.isEmpty + else { + throw ClaudeSwapSwitchParserError.malformedShape("missing reason") + } + let fromAccountNumber = try self.accountNumber( + in: object["from"], + field: "from", + allowsNull: true) + guard let toAccountNumber = try self.accountNumber( + in: object["to"], + field: "to", + allowsNull: false) + else { + throw ClaudeSwapSwitchParserError.malformedShape("to account has no numeric slot") + } + return ClaudeSwapAccountSwitchResult( + switched: switched, + fromAccountNumber: fromAccountNumber, + toAccountNumber: toAccountNumber, + reason: reason) + } + + private static func accountNumber(in raw: Any?, field: String, allowsNull: Bool) throws -> Int? { + if raw is NSNull, allowsNull { return nil } + guard let account = raw as? [String: Any] else { + throw ClaudeSwapSwitchParserError.malformedShape("missing \(field) account") + } + guard let rawNumber = account["number"] else { + throw ClaudeSwapSwitchParserError.malformedShape("\(field) account has no number") + } + if rawNumber is NSNull, allowsNull { return nil } + guard let number = self.integer(rawNumber), number > 0 + else { + throw ClaudeSwapSwitchParserError.malformedShape("\(field) account number is not a positive slot") + } + return number + } + + private static func integer(_ raw: Any?) -> Int? { + guard let number = raw as? NSNumber, + CFGetTypeID(number) != CFBooleanGetTypeID() + else { + return nil + } + return raw as? Int + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeUsageDataSource.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeUsageDataSource.swift new file mode 100644 index 0000000..063de7f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeUsageDataSource.swift @@ -0,0 +1,38 @@ +import Foundation + +public enum ClaudeUsageDataSource: String, CaseIterable, Identifiable, Sendable { + case auto + case api + case oauth + case web + case cli + + public var id: String { + self.rawValue + } + + public var displayName: String { + switch self { + case .auto: "Auto" + case .api: "API (Admin key)" + case .oauth: "OAuth API" + case .web: "Web API (cookies)" + case .cli: "CLI (PTY)" + } + } + + public var sourceLabel: String { + switch self { + case .auto: + "auto" + case .api: + "api" + case .oauth: + "oauth" + case .web: + "web" + case .cli: + "cli" + } + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeUsageFetcher.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeUsageFetcher.swift new file mode 100644 index 0000000..1bfd3aa --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeUsageFetcher.swift @@ -0,0 +1,1533 @@ +import Foundation + +public protocol ClaudeUsageFetching: Sendable { + func loadLatestUsage(model: String) async throws -> ClaudeUsageSnapshot + func debugRawProbe(model: String) async -> String + func detectVersion() -> String? +} + +public struct ClaudeUsageSnapshot: Sendable { + public enum PrimaryWindowKind: Equatable, Sendable { + case usage + case spendLimit + } + + public let primary: RateWindow + public let primaryWindowKind: PrimaryWindowKind + public let secondary: RateWindow? + public let opus: RateWindow? + public let extraRateWindows: [NamedRateWindow] + public let providerCost: ProviderCostSnapshot? + public let updatedAt: Date + public let accountEmail: String? + public let accountOrganization: String? + public let loginMethod: String? + public let rawText: String? + /// Present only when the credential used for this OAuth fetch matches the current Claude Keychain item. + public let oauthKeychainPersistentRefHash: String? + /// One-way, high-entropy ownership evidence derived from the exact credential used for this OAuth fetch. + public let oauthHistoryOwnerIdentifier: String? + /// True when a prompt-free comparison proved this credential differs from Claude Code's Keychain entry. + public let oauthKeychainCredentialMismatch: Bool + /// True when a prompt-free probe proved Claude Code has no Keychain credential. + public let oauthKeychainCredentialAbsent: Bool + /// True when a Claude CLI credential won routing but the prompt-free Keychain comparison was unavailable. + public let oauthKeychainCredentialUnavailable: Bool + + public init( + primary: RateWindow, + primaryWindowKind: PrimaryWindowKind = .usage, + secondary: RateWindow?, + opus: RateWindow?, + extraRateWindows: [NamedRateWindow] = [], + providerCost: ProviderCostSnapshot? = nil, + updatedAt: Date, + accountEmail: String?, + accountOrganization: String?, + loginMethod: String?, + rawText: String?, + oauthKeychainPersistentRefHash: String? = nil, + oauthHistoryOwnerIdentifier: String? = nil, + oauthKeychainCredentialMismatch: Bool = false, + oauthKeychainCredentialAbsent: Bool = false, + oauthKeychainCredentialUnavailable: Bool = false) + { + self.primary = primary + self.primaryWindowKind = primaryWindowKind + self.secondary = secondary + self.opus = opus + self.extraRateWindows = extraRateWindows + self.providerCost = providerCost + self.updatedAt = updatedAt + self.accountEmail = accountEmail + self.accountOrganization = accountOrganization + self.loginMethod = loginMethod + self.rawText = rawText + self.oauthKeychainPersistentRefHash = oauthKeychainPersistentRefHash + self.oauthHistoryOwnerIdentifier = oauthHistoryOwnerIdentifier + self.oauthKeychainCredentialMismatch = oauthKeychainCredentialMismatch + self.oauthKeychainCredentialAbsent = oauthKeychainCredentialAbsent + self.oauthKeychainCredentialUnavailable = oauthKeychainCredentialUnavailable + } +} + +public enum ClaudeUsageError: LocalizedError, Sendable { + case claudeNotInstalled + case parseFailed(String) + case oauthFailed(String) + + public var errorDescription: String? { + switch self { + case .claudeNotInstalled: + "Claude CLI is not installed. Install it from https://code.claude.com/docs/en/overview." + case let .parseFailed(details): + "Could not parse Claude usage: \(details)" + case let .oauthFailed(details): + details + } + } +} + +public struct ClaudeUsageFetcher: ClaudeUsageFetching, Sendable { + private static let sessionWindowMinutes = 5 * 60 + private static let weeklyWindowMinutes = 7 * 24 * 60 + private static let cliAutoProbeTimeout: TimeInterval = 12 + private static let cliProbeTimeout: TimeInterval = 24 + private static let cliRetryProbeTimeout: TimeInterval = 60 + private struct Configuration { + let environment: [String: String] + let runtime: ProviderRuntime + let dataSource: ClaudeUsageDataSource + let oauthKeychainPromptCooldownEnabled: Bool + let allowBackgroundDelegatedRefresh: Bool + let allowStartupBootstrapPrompt: Bool + let useWebExtras: Bool + let manualCookieHeader: String? + let webOrganizationID: String? + let keepCLISessionsAlive: Bool + let browserDetection: BrowserDetection + } + + private let configuration: Configuration + private static let log = CodexBarLog.logger(LogCategories.claudeUsage) + private static var isClaudeOAuthFlowDebugEnabled: Bool { + ProcessInfo.processInfo.environment["CODEXBAR_DEBUG_CLAUDE_OAUTH_FLOW"] == "1" + } + + private var environment: [String: String] { + self.configuration.environment + } + + private var runtime: ProviderRuntime { + self.configuration.runtime + } + + private var dataSource: ClaudeUsageDataSource { + self.configuration.dataSource + } + + private var oauthKeychainPromptCooldownEnabled: Bool { + self.configuration.oauthKeychainPromptCooldownEnabled + } + + private var allowsDelegatedOAuthRefresh: Bool { + self.runtime == .app + } + + private var allowsOAuthClaudeVersionDetection: Bool { + self.runtime == .app + } + + private var allowBackgroundDelegatedRefresh: Bool { + self.configuration.allowBackgroundDelegatedRefresh + } + + private var allowStartupBootstrapPrompt: Bool { + self.configuration.allowStartupBootstrapPrompt + } + + private var useWebExtras: Bool { + self.configuration.useWebExtras + } + + private var manualCookieHeader: String? { + self.configuration.manualCookieHeader + } + + private var webOrganizationID: String? { + self.configuration.webOrganizationID + } + + private var keepCLISessionsAlive: Bool { + self.configuration.keepCLISessionsAlive + } + + private var browserDetection: BrowserDetection { + self.configuration.browserDetection + } + + private struct ClaudeOAuthKeychainPromptPolicy { + let mode: ClaudeOAuthKeychainPromptMode + let isApplicable: Bool + let interaction: ProviderInteraction + + var canPromptNow: Bool { + switch self.mode { + case .never: + false + case .onlyOnUserAction: + self.interaction == .userInitiated + case .always: + true + } + } + + /// Respect the Keychain prompt cooldown for background operations to avoid spamming system dialogs. + /// User actions (menu open / refresh / settings) are allowed to bypass the cooldown. + var shouldRespectKeychainPromptCooldown: Bool { + self.interaction != .userInitiated + } + + var interactionLabel: String { + self.interaction == .userInitiated ? "user" : "background" + } + } + + private static func currentClaudeOAuthInteractivePromptPolicy() -> ClaudeOAuthKeychainPromptPolicy { + let policy = ClaudeOAuthKeychainPromptPolicy( + mode: ClaudeOAuthKeychainPromptPreference.securityFrameworkFallbackMode(), + isApplicable: true, + interaction: ProviderInteractionContext.current) + + // User actions should be able to immediately retry a Security.framework fallback repair after a background + // cooldown was recorded, even when /usr/bin/security is the primary reader. + if policy.interaction == .userInitiated { + if ClaudeOAuthKeychainAccessGate.clearDenied() { + Self.log.info("Claude OAuth keychain cooldown cleared by user action") + } + } + return policy + } + + private static func currentClaudeOAuthDelegatedRefreshPolicy() -> ClaudeOAuthKeychainPromptPolicy { + // Delegated refresh must honor the stored prompt mode for every keychain read strategy. + // securityCLIExperimental is not "prompt policy N/A" for CLI touch repair. + ClaudeOAuthKeychainPromptPolicy( + mode: ClaudeOAuthKeychainPromptPreference.storedMode(), + isApplicable: true, + interaction: ProviderInteractionContext.current) + } + + private static func assertDelegatedRefreshAllowedInCurrentInteraction( + policy: ClaudeOAuthKeychainPromptPolicy, + allowBackgroundDelegatedRefresh: Bool) throws + { + if policy.mode == .never { + throw ClaudeUsageError.oauthFailed("Delegated refresh is disabled by 'never' keychain policy.") + } + if policy.mode == .onlyOnUserAction, + policy.interaction != .userInitiated, + !allowBackgroundDelegatedRefresh + { + throw ClaudeUsageError.oauthFailed( + "Claude OAuth token expired, but background repair is suppressed when Keychain prompt policy " + + "is set to only prompt on user action. Open the CodexBar menu or click Refresh to retry.") + } + } + + #if DEBUG + @TaskLocal static var loadOAuthCredentialsOverride: (@Sendable ( + [String: String], + Bool, + Bool) async throws -> ClaudeOAuthCredentials)? + @TaskLocal static var fetchOAuthUsageOverride: (@Sendable ( + String, + Bool) async throws -> OAuthUsageResponse)? + @TaskLocal static var delegatedRefreshAttemptOverride: (@Sendable ( + Date, + TimeInterval, + [String: String]) async -> ClaudeOAuthDelegatedRefreshCoordinator.Outcome)? + @TaskLocal static var hasCachedCredentialsOverride: Bool? + #endif + + /// Creates a new ClaudeUsageFetcher. + /// - Parameters: + /// - environment: Process environment (default: current process environment) + /// - dataSource: Usage data source (default: OAuth API). + /// - useWebExtras: If true, attempts to enrich usage with Claude web data (cookies). + public init( + browserDetection: BrowserDetection, + environment: [String: String] = ProcessInfo.processInfo.environment, + runtime: ProviderRuntime = .app, + dataSource: ClaudeUsageDataSource = .oauth, + oauthKeychainPromptCooldownEnabled: Bool = false, + allowBackgroundDelegatedRefresh: Bool = false, + allowStartupBootstrapPrompt: Bool = false, + useWebExtras: Bool = false, + manualCookieHeader: String? = nil, + webOrganizationID: String? = nil, + keepCLISessionsAlive: Bool = false) + { + self.configuration = Configuration( + environment: environment, + runtime: runtime, + dataSource: dataSource, + oauthKeychainPromptCooldownEnabled: oauthKeychainPromptCooldownEnabled, + allowBackgroundDelegatedRefresh: allowBackgroundDelegatedRefresh, + allowStartupBootstrapPrompt: allowStartupBootstrapPrompt, + useWebExtras: useWebExtras, + manualCookieHeader: manualCookieHeader, + webOrganizationID: webOrganizationID, + keepCLISessionsAlive: keepCLISessionsAlive, + browserDetection: browserDetection) + } + + private struct OAuthExecutor { + let fetcher: ClaudeUsageFetcher + + func load(allowDelegatedRetry: Bool) async throws -> ClaudeUsageSnapshot { + do { + let promptPolicy = ClaudeUsageFetcher.currentClaudeOAuthInteractivePromptPolicy() + + #if DEBUG + let hasCache = if let hasCachedCredentialsOverride = ClaudeUsageFetcher.hasCachedCredentialsOverride { + hasCachedCredentialsOverride + } else if ClaudeUsageFetcher.loadOAuthCredentialsOverride != nil { + false + } else { + ClaudeOAuthCredentialsStore.hasCachedCredentials(environment: self.fetcher.environment) + } + #else + let hasCache = ClaudeOAuthCredentialsStore.hasCachedCredentials(environment: self.fetcher.environment) + #endif + + let startupBootstrapOverride = self.shouldAllowStartupBootstrapPrompt( + policy: promptPolicy, + hasCache: hasCache) + let allowKeychainPrompt = (promptPolicy.canPromptNow || startupBootstrapOverride) && !hasCache + ClaudeUsageFetcher.logOAuthBootstrapPromptDecision( + allowKeychainPrompt: allowKeychainPrompt, + policy: promptPolicy, + hasCache: hasCache, + startupBootstrapOverride: startupBootstrapOverride) + + let credentialRecord = try await ClaudeOAuthCredentialsStore.$allowBackgroundPromptBootstrap + .withValue(startupBootstrapOverride) { + try await ClaudeUsageFetcher.loadOAuthCredentialRecord( + environment: self.fetcher.environment, + allowKeychainPrompt: allowKeychainPrompt, + respectKeychainPromptCooldown: promptPolicy.shouldRespectKeychainPromptCooldown) + } + let credentials = credentialRecord.credentials + + try self.validateRequiredOAuthScope(credentials) + let usage = try await ClaudeUsageFetcher.fetchOAuthUsage( + accessToken: credentials.accessToken, + detectClaudeVersion: self.fetcher.allowsOAuthClaudeVersionDetection) + let keychainMatch = ClaudeOAuthCredentialsStore + .claudeKeychainCredentialMatchWithoutPrompt(for: credentialRecord) + return try ClaudeUsageFetcher.mapOAuthUsage( + usage, + credentials: credentials, + oauthKeychainPersistentRefHash: keychainMatch.persistentRefHash, + oauthHistoryOwnerIdentifier: credentialRecord.historyOwnerIdentifier, + oauthKeychainCredentialMismatch: keychainMatch.isMismatch, + oauthKeychainCredentialAbsent: keychainMatch.isAbsent, + oauthKeychainCredentialUnavailable: keychainMatch.isUnavailable) + } catch let error as CancellationError { + throw error + } catch let error as ClaudeUsageError { + throw error + } catch let error as ClaudeOAuthCredentialsError { + if case .refreshDelegatedToClaudeCLI = error { + return try await self.loadAfterDelegatedRefresh(allowDelegatedRetry: allowDelegatedRetry) + } + throw ClaudeUsageError.oauthFailed(error.localizedDescription) + } catch let error as ClaudeOAuthFetchError { + if case .rateLimited = error { + throw ClaudeUsageError.oauthFailed(error.localizedDescription) + } + ClaudeOAuthCredentialsStore.invalidateCache() + if case let .serverError(statusCode, body) = error, + statusCode == 403, + body?.contains("user:profile") ?? false + { + throw ClaudeUsageError.oauthFailed( + "Claude OAuth token does not meet scope requirement 'user:profile'. " + + "Run `claude setup-token` to re-generate credentials, or switch Claude Source to " + + "Web/CLI.") + } + throw ClaudeUsageError.oauthFailed(error.localizedDescription) + } catch { + throw ClaudeUsageError.oauthFailed(error.localizedDescription) + } + } + + private func shouldAllowStartupBootstrapPrompt( + policy: ClaudeOAuthKeychainPromptPolicy, + hasCache: Bool) -> Bool + { + guard self.fetcher.allowStartupBootstrapPrompt else { return false } + guard !hasCache else { return false } + guard ClaudeOAuthKeychainPromptPreference.securityFrameworkFallbackMode() == .onlyOnUserAction else { + return false + } + guard policy.interaction == .background else { return false } + return ProviderRefreshContext.current == .startup + } + + private func loadAfterDelegatedRefresh(allowDelegatedRetry: Bool) async throws -> ClaudeUsageSnapshot { + guard allowDelegatedRetry else { + throw ClaudeUsageError.oauthFailed( + "Claude OAuth token expired. CodexBar CLI does not launch Claude to refresh credentials. " + + "Run `claude login`, then retry.") + } + + try Task.checkCancellation() + + let delegatedPromptPolicy = ClaudeUsageFetcher.currentClaudeOAuthDelegatedRefreshPolicy() + try ClaudeUsageFetcher.assertDelegatedRefreshAllowedInCurrentInteraction( + policy: delegatedPromptPolicy, + allowBackgroundDelegatedRefresh: self.fetcher.allowBackgroundDelegatedRefresh) + + let delegatedOutcome = await ClaudeUsageFetcher.attemptDelegatedRefresh( + environment: self.fetcher.environment) + ClaudeUsageFetcher.log.info( + "Claude OAuth delegated refresh attempted", + metadata: [ + "outcome": ClaudeUsageFetcher.delegatedRefreshOutcomeLabel(delegatedOutcome), + ]) + + do { + if self.fetcher.oauthKeychainPromptCooldownEnabled { + switch delegatedOutcome { + case .skippedByCooldown, .cliUnavailable: + throw ClaudeUsageError.oauthFailed( + "Claude OAuth token expired; delegated refresh is unavailable (outcome=" + + "\(ClaudeUsageFetcher.delegatedRefreshOutcomeLabel(delegatedOutcome))).") + case .attemptedSucceeded, .attemptedFailed: + break + } + } + + try Task.checkCancellation() + + _ = ClaudeOAuthCredentialsStore.invalidateCacheIfCredentialsFileChanged() + + let didSyncSilently = delegatedOutcome == .attemptedSucceeded + && ClaudeOAuthCredentialsStore.syncFromClaudeKeychainWithoutPrompt(now: Date()) + + let promptPolicy = ClaudeUsageFetcher.currentClaudeOAuthInteractivePromptPolicy() + ClaudeUsageFetcher.logDeferredBackgroundDelegatedRecoveryIfNeeded( + delegatedOutcome: delegatedOutcome, + didSyncSilently: didSyncSilently, + policy: promptPolicy) + let retryAllowKeychainPrompt = promptPolicy.canPromptNow && !didSyncSilently + if retryAllowKeychainPrompt { + ClaudeUsageFetcher.log.info( + "Claude OAuth keychain prompt allowed (post-delegation retry)", + metadata: [ + "interaction": promptPolicy.interactionLabel, + "promptMode": promptPolicy.mode.rawValue, + "promptPolicyApplicable": "\(promptPolicy.isApplicable)", + "delegatedOutcome": ClaudeUsageFetcher.delegatedRefreshOutcomeLabel(delegatedOutcome), + "didSyncSilently": "\(didSyncSilently)", + ]) + } + if ClaudeUsageFetcher.isClaudeOAuthFlowDebugEnabled { + ClaudeUsageFetcher.log.debug( + "Claude OAuth credential load (post-delegation retry start)", + metadata: [ + "cooldownEnabled": "\(self.fetcher.oauthKeychainPromptCooldownEnabled)", + "didSyncSilently": "\(didSyncSilently)", + "allowKeychainPrompt": "\(retryAllowKeychainPrompt)", + "delegatedOutcome": ClaudeUsageFetcher.delegatedRefreshOutcomeLabel(delegatedOutcome), + "interaction": promptPolicy.interactionLabel, + "promptMode": promptPolicy.mode.rawValue, + "promptPolicyApplicable": "\(promptPolicy.isApplicable)", + ]) + } + + let refreshedRecord = try await ClaudeUsageFetcher.loadOAuthCredentialRecord( + environment: self.fetcher.environment, + allowKeychainPrompt: retryAllowKeychainPrompt, + respectKeychainPromptCooldown: promptPolicy.shouldRespectKeychainPromptCooldown) + let refreshedCredentials = refreshedRecord.credentials + if ClaudeUsageFetcher.isClaudeOAuthFlowDebugEnabled { + ClaudeUsageFetcher.log.debug( + "Claude OAuth credential load (post-delegation retry)", + metadata: [ + "cooldownEnabled": "\(self.fetcher.oauthKeychainPromptCooldownEnabled)", + "didSyncSilently": "\(didSyncSilently)", + "allowKeychainPrompt": "\(retryAllowKeychainPrompt)", + "delegatedOutcome": ClaudeUsageFetcher.delegatedRefreshOutcomeLabel(delegatedOutcome), + "interaction": promptPolicy.interactionLabel, + "promptMode": promptPolicy.mode.rawValue, + "promptPolicyApplicable": "\(promptPolicy.isApplicable)", + ]) + } + + try self.validateRequiredOAuthScope(refreshedCredentials) + let usage = try await ClaudeUsageFetcher.fetchOAuthUsage( + accessToken: refreshedCredentials.accessToken, + detectClaudeVersion: self.fetcher.allowsOAuthClaudeVersionDetection) + let keychainMatch = ClaudeOAuthCredentialsStore + .claudeKeychainCredentialMatchWithoutPrompt(for: refreshedRecord) + return try ClaudeUsageFetcher.mapOAuthUsage( + usage, + credentials: refreshedCredentials, + oauthKeychainPersistentRefHash: keychainMatch.persistentRefHash, + oauthHistoryOwnerIdentifier: refreshedRecord.historyOwnerIdentifier, + oauthKeychainCredentialMismatch: keychainMatch.isMismatch, + oauthKeychainCredentialAbsent: keychainMatch.isAbsent, + oauthKeychainCredentialUnavailable: keychainMatch.isUnavailable) + } catch { + ClaudeUsageFetcher.log.debug( + "Claude OAuth post-delegation retry failed", + metadata: ClaudeUsageFetcher.delegatedRetryFailureMetadata( + error: error, + oauthKeychainPromptCooldownEnabled: self.fetcher.oauthKeychainPromptCooldownEnabled, + delegatedOutcome: delegatedOutcome)) + throw ClaudeUsageError.oauthFailed( + ClaudeUsageFetcher.delegatedRefreshFailureMessage( + for: delegatedOutcome, + retryError: error)) + } + } + + private func validateRequiredOAuthScope(_ credentials: ClaudeOAuthCredentials) throws { + guard credentials.scopes.contains("user:profile") else { + let scopes = credentials.scopes.joined(separator: ", ") + let detail = scopes.isEmpty + ? "Claude OAuth token missing 'user:profile' scope." + : "Claude OAuth token missing 'user:profile' scope (has: \(scopes))." + throw ClaudeUsageError.oauthFailed( + detail + " Run `claude setup-token` to re-generate credentials, or switch Claude Source to " + + "Web/CLI.") + } + } + } + + private struct StepExecutor { + let fetcher: ClaudeUsageFetcher + + func loadLatestUsage(model: String) async throws -> ClaudeUsageSnapshot { + switch self.fetcher.dataSource { + case .auto: + return try await self.executeAuto(model: model) + case .api: + throw ClaudeUsageError.parseFailed("Claude Admin API usage is handled by the provider descriptor.") + case .oauth: + var snapshot = try await self.fetcher.loadViaOAuth( + allowDelegatedRetry: self.fetcher.allowsDelegatedOAuthRefresh) + snapshot = await self.fetcher.applyWebExtrasIfNeeded(to: snapshot) + return snapshot + case .web: + return try await self.fetcher.loadViaWebAPI() + case .cli: + return try await self.loadViaCLIWithRetry(model: model) + } + } + + private func executeAuto(model: String) async throws -> ClaudeUsageSnapshot { + let plan = await self.makeAutoFetchPlan() + self.logAutoPlan(plan) + + let executionSteps = plan.executionSteps + for (index, step) in executionSteps.enumerated() { + do { + return try await self.execute(step: step, model: model) + } catch { + if index < executionSteps.count - 1 { + ClaudeUsageFetcher.log.debug( + "Claude planner step failed; falling back to next step", + metadata: [ + "step": step.dataSource.rawValue, + "reason": step.inclusionReason.rawValue, + "errorType": String(describing: type(of: error)), + ]) + continue + } + throw error + } + } + throw ClaudeUsageError.parseFailed("Claude planner produced no executable steps.") + } + + private func makeAutoFetchPlan() async -> ClaudeFetchPlan { + let hasWebSession = + if let header = self.fetcher.manualCookieHeader { + ClaudeWebAPIFetcher.hasSessionKey(cookieHeader: header) + } else { + ClaudeWebAPIFetcher.hasSessionKey(browserDetection: self.fetcher.browserDetection) + } + let hasCLI = ClaudeCLIResolver.isAvailable(environment: self.fetcher.environment) + return ClaudeSourcePlanner.resolve(input: ClaudeSourcePlanningInput( + runtime: self.fetcher.runtime, + selectedDataSource: .auto, + webExtrasEnabled: self.fetcher.useWebExtras, + hasWebSession: hasWebSession, + hasCLI: hasCLI, + hasOAuthCredentials: ClaudeOAuthPlanningAvailability.isAvailable( + runtime: self.fetcher.runtime, + sourceMode: .auto, + environment: self.fetcher.environment))) + } + + private func logAutoPlan(_ plan: ClaudeFetchPlan) { + var metadata: [String: String] = [ + "plannerOrder": plan.orderLabel, + "selected": plan.preferredStep?.dataSource.rawValue ?? "none", + "noSourceAvailable": "\(plan.isNoSourceAvailable)", + "webExtrasEnabled": "\(self.fetcher.useWebExtras)", + "oauthReadStrategy": ClaudeOAuthKeychainReadStrategyPreference.current().rawValue, + ] + for (index, step) in plan.orderedSteps.enumerated() { + metadata["step\(index)"] = + "\(step.dataSource.rawValue):\(step.inclusionReason.rawValue):\(step.isPlausiblyAvailable)" + } + ClaudeUsageFetcher.log.debug("Claude auto source planner", metadata: metadata) + } + + private func execute(step: ClaudeFetchPlanStep, model: String) async throws -> ClaudeUsageSnapshot { + switch step.dataSource { + case .api: + throw ClaudeUsageError.parseFailed("Planner emitted invalid api execution step.") + case .oauth: + var snapshot = try await self.fetcher.loadViaOAuth( + allowDelegatedRetry: self.fetcher.allowsDelegatedOAuthRefresh) + snapshot = await self.fetcher.applyWebExtrasIfNeeded(to: snapshot) + return snapshot + case .web: + return try await self.fetcher.loadViaWebAPI() + case .cli: + return try await self.loadViaAutoCLI(model: model) + case .auto: + throw ClaudeUsageError.parseFailed("Planner emitted invalid auto execution step.") + } + } + + private func loadViaAutoCLI(model: String) async throws -> ClaudeUsageSnapshot { + do { + return try await self.loadViaCLI(model: model, timeout: ClaudeUsageFetcher.cliAutoProbeTimeout) + } catch { + if error is CancellationError { + throw error + } + guard Self.shouldRetryCLIProbe(after: error) else { throw error } + return try await self.loadViaCLI(model: model, timeout: ClaudeUsageFetcher.cliRetryProbeTimeout) + } + } + + private func loadViaCLIWithRetry(model: String) async throws -> ClaudeUsageSnapshot { + do { + return try await self.loadViaCLI(model: model, timeout: ClaudeUsageFetcher.cliProbeTimeout) + } catch { + if error is CancellationError { + throw error + } + guard Self.shouldRetryCLIProbe(after: error) else { throw error } + return try await self.loadViaCLI(model: model, timeout: ClaudeUsageFetcher.cliRetryProbeTimeout) + } + } + + private func loadViaCLI(model: String, timeout: TimeInterval) async throws -> ClaudeUsageSnapshot { + if ClaudeCLIRateLimitGate.blockedUntil() != nil { + throw ClaudeUsageError.parseFailed(ClaudeCLIRateLimitGate.message) + } + + var snapshot: ClaudeUsageSnapshot + do { + snapshot = try await self.fetcher.loadViaPTY(model: model, timeout: timeout) + } catch { + if error is CancellationError { + throw error + } + if ClaudeUsageFetcher.isCLIRateLimitError(error) { + ClaudeCLIRateLimitGate.recordRateLimit() + throw error + } + guard Self.shouldTryDirectCLIUsage(after: error) else { throw error } + let ptyError = error + do { + snapshot = try await self.fetcher.loadViaDirectCLI( + timeout: Self.directCLIUsageTimeout(for: timeout)) + } catch let directError { + if directError is CancellationError { + throw directError + } + if ClaudeUsageFetcher.isCLIRateLimitError(directError) { + ClaudeCLIRateLimitGate.recordRateLimit() + throw directError + } + guard Self.directCLIErrorShouldReplacePTYError(directError) else { throw ptyError } + throw directError + } + } + ClaudeCLIRateLimitGate.recordSuccess() + snapshot = await self.fetcher.applyWebExtrasIfNeeded(to: snapshot) + return snapshot + } + + private static func directCLIUsageTimeout(for ptyTimeout: TimeInterval) -> TimeInterval { + min(max(ptyTimeout / 3, 6), 8) + } + + private static func directCLIErrorShouldReplacePTYError(_ error: Error) -> Bool { + if case let ClaudeStatusProbeError.parseFailed(message) = error { + return message.lowercased().contains("subscription") + } + if case let ClaudeUsageError.parseFailed(message) = error { + return message.lowercased().contains("subscription") + } + return false + } + + private static func shouldTryDirectCLIUsage(after error: Error) -> Bool { + if case ClaudeStatusProbeError.timedOut = error { + return true + } + if case let ClaudeStatusProbeError.parseFailed(message) = error { + let lower = message.lowercased() + return lower.contains("still loading usage") || lower.contains("could not load usage data") + } + let message = error.localizedDescription.lowercased() + return message.contains("timed out") || message.contains("timeout") + } + + private static func shouldRetryCLIProbe(after error: Error) -> Bool { + if case ClaudeStatusProbeError.timedOut = error { + return true + } + if case let ClaudeStatusProbeError.parseFailed(message) = error { + return message.lowercased().contains("still loading usage") + } + let message = error.localizedDescription.lowercased() + return message.contains("timed out") || message.contains("timeout") + } + } +} + +extension ClaudeUsageFetcher { + // MARK: - Parsing helpers + + public static func parse(json: Data) -> ClaudeUsageSnapshot? { + self.parse(json: json, now: .init()) + } + + package static func parse(json: Data, now: Date) -> ClaudeUsageSnapshot? { + guard let output = String(data: json, encoding: .utf8) else { return nil } + return try? Self.parse(output: output, now: now) + } + + private static func parse(output: String, now: Date) throws -> ClaudeUsageSnapshot { + guard + let data = output.data(using: .utf8), + let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + throw ClaudeUsageError.parseFailed(output.prefix(500).description) + } + + if let ok = obj["ok"] as? Bool, !ok { + let hint = obj["hint"] as? String ?? (obj["pane_preview"] as? String ?? "") + throw ClaudeUsageError.parseFailed(hint) + } + + func firstWindowDict(_ keys: [String]) -> [String: Any]? { + for key in keys { + if let dict = obj[key] as? [String: Any] { + return dict + } + } + return nil + } + + func makeWindow(_ dict: [String: Any]?, windowMinutes: Int) -> RateWindow? { + guard let dict else { return nil } + let pct = (dict["pct_used"] as? NSNumber)?.doubleValue ?? 0 + let resetText = dict["resets"] as? String + return RateWindow( + usedPercent: pct, + windowMinutes: windowMinutes, + resetsAt: ClaudeStatusProbe.parseResetDate( + from: resetText, + now: now, + expectedWindow: TimeInterval(windowMinutes * 60)), + resetDescription: resetText) + } + + guard let session = makeWindow( + firstWindowDict(["session_5h"]), + windowMinutes: Self.sessionWindowMinutes) + else { + throw ClaudeUsageError.parseFailed("missing session data") + } + let weekAll = makeWindow( + firstWindowDict(["week_all_models", "week_all"]), + windowMinutes: Self.weeklyWindowMinutes) + + let rawEmail = (obj["account_email"] as? String)?.trimmingCharacters( + in: .whitespacesAndNewlines) + let email = (rawEmail?.isEmpty ?? true) ? nil : rawEmail + let rawOrg = (obj["account_org"] as? String)?.trimmingCharacters( + in: .whitespacesAndNewlines) + let org = (rawOrg?.isEmpty ?? true) ? nil : rawOrg + let loginMethod = (obj["login_method"] as? String)?.trimmingCharacters( + in: .whitespacesAndNewlines) + let opusWindow = makeWindow( + firstWindowDict([ + "week_sonnet", + "week_sonnet_only", + "week_opus", + ]), + windowMinutes: Self.weeklyWindowMinutes) + return ClaudeUsageSnapshot( + primary: session, + secondary: weekAll, + opus: opusWindow, + providerCost: nil, + updatedAt: now, + accountEmail: email, + accountOrganization: org, + loginMethod: loginMethod, + rawText: output) + } + + // MARK: - Public API + + public func detectVersion() -> String? { + ProviderVersionDetector.claudeVersion() + } + + public func debugRawProbe(model: String = "sonnet") async -> String { + do { + let snap = try await self.loadViaPTY(model: model, timeout: Self.cliProbeTimeout) + let opus = snap.opus?.remainingPercent ?? -1 + let email = snap.accountEmail ?? "nil" + let org = snap.accountOrganization ?? "nil" + let weekly = snap.secondary?.remainingPercent ?? -1 + let primary = snap.primary.remainingPercent + return """ + session_left=\(primary) weekly_left=\(weekly) + opus_left=\(opus) email \(email) org \(org) + \(snap) + """ + } catch { + return "Probe failed: \(error)" + } + } + + public func loadLatestUsage(model: String = "sonnet") async throws -> ClaudeUsageSnapshot { + try await StepExecutor(fetcher: self).loadLatestUsage(model: model) + } + + public static func isCLIRateLimitError(_ error: Error) -> Bool { + ClaudeCLIRateLimitGate.isRateLimitError(error) + } +} + +extension ClaudeUsageFetcher { + // MARK: - OAuth API path + + private static func logOAuthBootstrapPromptDecision( + allowKeychainPrompt: Bool, + policy: ClaudeOAuthKeychainPromptPolicy, + hasCache: Bool, + startupBootstrapOverride: Bool) + { + guard allowKeychainPrompt else { return } + self.log.info( + "Claude OAuth keychain prompt allowed (bootstrap)", + metadata: [ + "interaction": policy.interactionLabel, + "promptMode": policy.mode.rawValue, + "promptPolicyApplicable": "\(policy.isApplicable)", + "hasCache": "\(hasCache)", + "startupBootstrapOverride": "\(startupBootstrapOverride)", + ]) + } + + private static func logDeferredBackgroundDelegatedRecoveryIfNeeded( + delegatedOutcome: ClaudeOAuthDelegatedRefreshCoordinator.Outcome, + didSyncSilently: Bool, + policy: ClaudeOAuthKeychainPromptPolicy) + { + guard delegatedOutcome == .attemptedSucceeded else { return } + guard !didSyncSilently else { return } + guard policy.mode == .onlyOnUserAction else { return } + guard policy.interaction == .background else { return } + self.log.info( + "Claude OAuth delegated refresh completed; background recovery deferred until user action", + metadata: [ + "interaction": policy.interactionLabel, + "promptMode": policy.mode.rawValue, + "delegatedOutcome": self.delegatedRefreshOutcomeLabel(delegatedOutcome), + ]) + } + + private func loadViaOAuth(allowDelegatedRetry: Bool) async throws -> ClaudeUsageSnapshot { + try await OAuthExecutor(fetcher: self).load(allowDelegatedRetry: allowDelegatedRetry) + } + + private static func loadOAuthCredentialRecord( + environment: [String: String], + allowKeychainPrompt: Bool, + respectKeychainPromptCooldown: Bool) async throws -> ClaudeOAuthCredentialRecord + { + #if DEBUG + if let override = loadOAuthCredentialsOverride { + return try await ClaudeOAuthCredentialRecord( + credentials: override(environment, allowKeychainPrompt, respectKeychainPromptCooldown), + owner: .environment, + source: .environment) + } + #endif + return try await ClaudeOAuthCredentialsStore.loadRecordWithAutoRefresh( + environment: environment, + allowKeychainPrompt: allowKeychainPrompt, + respectKeychainPromptCooldown: respectKeychainPromptCooldown) + } + + private static func fetchOAuthUsage( + accessToken: String, + detectClaudeVersion: Bool) async throws -> OAuthUsageResponse + { + #if DEBUG + if let override = fetchOAuthUsageOverride { + return try await override(accessToken, detectClaudeVersion) + } + #endif + return try await ClaudeOAuthUsageFetcher.fetchUsage( + accessToken: accessToken, + detectClaudeVersion: detectClaudeVersion) + } + + private static func attemptDelegatedRefresh( + now: Date = Date(), + timeout: TimeInterval = 15, + environment: [String: String] = ProcessInfo.processInfo.environment) + async -> ClaudeOAuthDelegatedRefreshCoordinator.Outcome + { + #if DEBUG + if let override = delegatedRefreshAttemptOverride { + return await override(now, timeout, environment) + } + #endif + return await ClaudeOAuthDelegatedRefreshCoordinator.attempt( + now: now, + timeout: timeout, + environment: environment) + } + + private static func delegatedRefreshOutcomeLabel( + _ outcome: ClaudeOAuthDelegatedRefreshCoordinator.Outcome) -> String + { + switch outcome { + case .skippedByCooldown: + "skippedByCooldown" + case .cliUnavailable: + "cliUnavailable" + case .attemptedSucceeded: + "attemptedSucceeded" + case .attemptedFailed: + "attemptedFailed" + } + } + + private static func delegatedRefreshFailureMessage( + for outcome: ClaudeOAuthDelegatedRefreshCoordinator.Outcome, + retryError: Error) -> String + { + if let oauthError = retryError as? ClaudeOAuthFetchError, + case .rateLimited = oauthError + { + return oauthError.localizedDescription + } + + switch outcome { + case .skippedByCooldown: + return "Claude OAuth token expired and delegated refresh is cooling down. " + + "Please retry shortly, or run `claude login`." + case .cliUnavailable: + return "Claude OAuth token expired and Claude CLI is not available for delegated refresh. " + + "Install/configure `claude`, or run `claude login`." + case .attemptedSucceeded: + return "Claude OAuth token is still unavailable after delegated Claude CLI refresh. " + + "Run `claude login`, then retry." + case let .attemptedFailed(message): + return "Claude OAuth token expired and delegated Claude CLI refresh failed: \(message). " + + "Run `claude login`, then retry." + } + } + + private static func delegatedRetryFailureMetadata( + error: Error, + oauthKeychainPromptCooldownEnabled: Bool, + delegatedOutcome: ClaudeOAuthDelegatedRefreshCoordinator.Outcome) -> [String: String] + { + var metadata: [String: String] = [ + "errorType": String(describing: type(of: error)), + "cooldownEnabled": "\(oauthKeychainPromptCooldownEnabled)", + "delegatedOutcome": Self.delegatedRefreshOutcomeLabel(delegatedOutcome), + ] + + // Avoid `localizedDescription` here: some error types include server response bodies in their + // `errorDescription`, which can leak potentially identifying information into logs. + if let oauthError = error as? ClaudeOAuthFetchError { + switch oauthError { + case .unauthorized: + metadata["oauthError"] = "unauthorized" + case let .rateLimited(retryAfter): + metadata["oauthError"] = "rateLimited" + metadata["retryAfter"] = retryAfter.map { "\($0.timeIntervalSince1970)" } ?? "nil" + case .invalidResponse: + metadata["oauthError"] = "invalidResponse" + case let .serverError(statusCode, body): + metadata["oauthError"] = "serverError" + metadata["httpStatus"] = "\(statusCode)" + metadata["bodyLength"] = "\(body?.utf8.count ?? 0)" + case let .networkError(underlying): + metadata["oauthError"] = "networkError" + metadata["underlyingType"] = String(describing: type(of: underlying)) + } + } + + return metadata + } + + private static func mapOAuthUsage( + _ usage: OAuthUsageResponse, + credentials: ClaudeOAuthCredentials, + oauthKeychainPersistentRefHash: String? = nil, + oauthHistoryOwnerIdentifier: String? = nil, + oauthKeychainCredentialMismatch: Bool = false, + oauthKeychainCredentialAbsent: Bool = false, + oauthKeychainCredentialUnavailable: Bool = false) throws -> ClaudeUsageSnapshot + { + let oauthHistoryOwnerIdentifier = ClaudeOAuthCredentials.normalizedHistoryOwnerIdentifier( + oauthHistoryOwnerIdentifier) ?? credentials.historyOwnerIdentifier + + func makeWindow(_ window: OAuthUsageWindow?, windowMinutes: Int?) -> RateWindow? { + guard let window, + let utilization = window.utilization + else { return nil } + let resetDate = ClaudeOAuthUsageFetcher.parseISO8601Date(window.resetsAt) + let resetDescription = resetDate.map(Self.formatResetDate) + return RateWindow( + usedPercent: utilization, + windowMinutes: windowMinutes, + resetsAt: resetDate, + resetDescription: resetDescription) + } + + let loginMethod = ClaudePlan.oauthLoginMethod( + subscriptionType: credentials.subscriptionType, + rateLimitTier: credentials.rateLimitTier) + let primary = makeWindow(usage.fiveHour, windowMinutes: 5 * 60) + ?? makeWindow(usage.sevenDay, windowMinutes: 7 * 24 * 60) + ?? makeWindow(usage.sevenDayOAuthApps, windowMinutes: 7 * 24 * 60) + ?? makeWindow(usage.sevenDaySonnet, windowMinutes: 7 * 24 * 60) + ?? makeWindow(usage.sevenDayOpus, windowMinutes: 7 * 24 * 60) + let treatAsSpendLimit = primary == nil && usage.extraUsage?.isEnabled == true + let providerCost = Self.oauthExtraUsageCost( + usage.extraUsage, + loginMethod: loginMethod, + treatAsSpendLimit: treatAsSpendLimit) + + guard let primary else { + if let spendLimit = Self.oauthSpendLimitWindow(from: providerCost, extraUsage: usage.extraUsage) { + return ClaudeUsageSnapshot( + primary: spendLimit, + primaryWindowKind: .spendLimit, + secondary: nil, + opus: nil, + extraRateWindows: Self.oauthExtraRateWindows(from: usage), + providerCost: providerCost, + updatedAt: Date(), + accountEmail: nil, + accountOrganization: nil, + loginMethod: loginMethod, + rawText: nil, + oauthKeychainPersistentRefHash: oauthKeychainPersistentRefHash, + oauthHistoryOwnerIdentifier: oauthHistoryOwnerIdentifier, + oauthKeychainCredentialMismatch: oauthKeychainCredentialMismatch, + oauthKeychainCredentialAbsent: oauthKeychainCredentialAbsent, + oauthKeychainCredentialUnavailable: oauthKeychainCredentialUnavailable) + } + throw ClaudeUsageError.parseFailed("missing session data") + } + + let weekly = makeWindow(usage.sevenDay, windowMinutes: 7 * 24 * 60) + let modelSpecific = makeWindow( + usage.sevenDaySonnet ?? usage.sevenDayOpus, + windowMinutes: 7 * 24 * 60) + let extraRateWindows = Self.oauthExtraRateWindows(from: usage) + + return ClaudeUsageSnapshot( + primary: primary, + secondary: weekly, + opus: modelSpecific, + extraRateWindows: extraRateWindows, + providerCost: providerCost, + updatedAt: Date(), + accountEmail: nil, + accountOrganization: nil, + loginMethod: loginMethod, + rawText: nil, + oauthKeychainPersistentRefHash: oauthKeychainPersistentRefHash, + oauthHistoryOwnerIdentifier: oauthHistoryOwnerIdentifier, + oauthKeychainCredentialMismatch: oauthKeychainCredentialMismatch, + oauthKeychainCredentialAbsent: oauthKeychainCredentialAbsent, + oauthKeychainCredentialUnavailable: oauthKeychainCredentialUnavailable) + } + + private static func oauthExtraUsageCost( + _ extra: OAuthExtraUsage?, + loginMethod: String?, + treatAsSpendLimit: Bool = false) -> ProviderCostSnapshot? + { + guard let extra, extra.isEnabled == true else { return nil } + guard let used = extra.usedCredits, + let limit = extra.monthlyLimit + else { return nil } + let currency = extra.currency?.trimmingCharacters(in: .whitespacesAndNewlines) + let code = (currency?.isEmpty ?? true) ? "USD" : currency! + let isSpendLimit = treatAsSpendLimit || ClaudePlan.fromCompatibilityLoginMethod(loginMethod) == .enterprise + let normalized = Self.normalizeClaudeExtraUsageAmounts( + used: used, + limit: limit, + treatAsMajorUnits: false) + return ProviderCostSnapshot( + used: normalized.used, + limit: normalized.limit, + currencyCode: code, + period: isSpendLimit ? "Spend limit" : "Monthly cap", + resetsAt: nil, + updatedAt: Date()) + } + + private static func oauthSpendLimitWindow( + from providerCost: ProviderCostSnapshot?, + extraUsage: OAuthExtraUsage?) -> RateWindow? + { + guard let providerCost, + providerCost.limit > 0 + else { return nil } + let usedPercent = extraUsage?.utilization ?? (providerCost.used / providerCost.limit) * 100 + let used = UsageFormatter.currencyString(providerCost.used, currencyCode: providerCost.currencyCode) + let limit = UsageFormatter.currencyString(providerCost.limit, currencyCode: providerCost.currencyCode) + return RateWindow( + usedPercent: min(100, max(0, usedPercent)), + windowMinutes: nil, + resetsAt: providerCost.resetsAt, + resetDescription: "\(providerCost.period ?? "Spend limit"): \(used) / \(limit)") + } + + private static func normalizeClaudeExtraUsageAmounts( + used: Double, + limit: Double, + treatAsMajorUnits: Bool) -> (used: Double, limit: Double) + { + if treatAsMajorUnits { + return (used: used, limit: limit) + } + + // Claude's OAuth API returns values in cents (minor units), same as the Web API. + // Always convert to dollars (major units) for display consistency. + // See: ClaudeWebAPIFetcher.swift which always divides by 100. + return (used: used / 100.0, limit: limit / 100.0) + } + + private static func oauthExtraRateWindows(from usage: OAuthUsageResponse) -> [NamedRateWindow] { + let definitions: [(id: String, title: String, window: OAuthUsageWindow?, sourceKey: String?)] = [ + ( + id: "claude-routines", + title: "Daily Routines", + window: usage.sevenDayRoutines, + sourceKey: usage.sevenDayRoutinesSourceKey), + ] + if let routinesKey = usage.sevenDayRoutinesSourceKey { + Self.log.debug("Claude OAuth extra usage key matched: routines=\(routinesKey)") + } + let routineWindows: [NamedRateWindow] = definitions.compactMap { definition in + let utilization: Double + let resetDate: Date? + if let window = definition.window, let parsedUtilization = window.utilization { + utilization = parsedUtilization + resetDate = ClaudeOAuthUsageFetcher.parseISO8601Date(window.resetsAt) + } else if definition.sourceKey != nil { + // Keep product bars visible when the API returns a known key with null payload. + utilization = 0 + resetDate = nil + } else { + return nil + } + let resetDescription = resetDate.map(Self.formatResetDate) + return NamedRateWindow( + id: definition.id, + title: definition.title, + window: RateWindow( + usedPercent: utilization, + windowMinutes: Self.weeklyWindowMinutes, + resetsAt: resetDate, + resetDescription: resetDescription)) + } + return routineWindows + Self.oauthScopedWeeklyLimitWindows(from: usage) + } + + private static func oauthScopedWeeklyLimitWindows(from usage: OAuthUsageResponse) -> [NamedRateWindow] { + let limits = usage.limits?.map { entry in + ClaudeScopedWeeklyLimitMapper.Limit( + kind: entry.kind, + group: entry.group, + percent: entry.percent, + resetsAt: ClaudeOAuthUsageFetcher.parseISO8601Date(entry.resetsAt), + modelID: entry.scope?.model?.id, + modelName: entry.scope?.model?.displayName) + } + // `is_active` is intentionally not a filter: observed enforceable scoped limits report false. + return ClaudeScopedWeeklyLimitMapper.extraRateWindows( + from: limits, + resetDescription: Self.formatResetDate) + } + + // MARK: - Web API path (uses browser cookies) + + /// The 5-hour session (`primary`) window for a Claude web usage payload. + /// + /// When `five_hour` is null the web fetcher reports a synthesized 0% session (an account with no live + /// session window). Flag that placeholder so lane classifiers — e.g. the combined Session + Weekly + /// menu-bar metric — surface the weekly lane instead of a phantom 5h session. The flag keys off the + /// reported presence of the session object, so a real session that is merely idle (0% used, possibly + /// with no reset) stays unflagged and keeps rendering. + static func webPrimaryWindow(from webData: ClaudeWebAPIFetcher.WebUsageData) -> RateWindow { + RateWindow( + usedPercent: webData.sessionPercentUsed, + windowMinutes: 5 * 60, + resetsAt: webData.sessionResetsAt, + resetDescription: webData.sessionResetsAt.map { Self.formatResetDate($0) }, + isSyntheticPlaceholder: !webData.hasLiveSessionWindow) + } + + private func loadViaWebAPI() async throws -> ClaudeUsageSnapshot { + let webData: ClaudeWebAPIFetcher.WebUsageData = + if let header = self.manualCookieHeader { + try await ClaudeWebAPIFetcher.fetchUsage( + cookieHeader: header, + targetOrganizationID: self.webOrganizationID) + { msg in + Self.log.debug(msg) + } + } else { + try await ClaudeWebAPIFetcher.fetchUsage( + browserDetection: self.browserDetection, + targetOrganizationID: self.webOrganizationID) + { msg in + Self.log.debug(msg) + } + } + // Convert web API data to ClaudeUsageSnapshot format + let primary = Self.webPrimaryWindow(from: webData) + + let secondary: RateWindow? = webData.weeklyPercentUsed.map { pct in + RateWindow( + usedPercent: pct, + windowMinutes: 7 * 24 * 60, + resetsAt: webData.weeklyResetsAt, + resetDescription: webData.weeklyResetsAt.map { Self.formatResetDate($0) }) + } + + let opus: RateWindow? = webData.opusPercentUsed.map { opusPct in + RateWindow( + usedPercent: opusPct, + windowMinutes: 7 * 24 * 60, + resetsAt: webData.weeklyResetsAt, + resetDescription: webData.weeklyResetsAt.map { Self.formatResetDate($0) }) + } + + return ClaudeUsageSnapshot( + primary: primary, + secondary: secondary, + opus: opus, + extraRateWindows: webData.extraRateWindows, + providerCost: webData.extraUsageCost, + updatedAt: Date(), + accountEmail: webData.accountEmail, + accountOrganization: webData.accountOrganization, + loginMethod: webData.loginMethod, + rawText: nil) + } + + private static func formatResetDate(_ date: Date) -> String { + let formatter = DateFormatter() + formatter.dateFormat = "MMM d 'at' h:mma" + formatter.locale = Locale(identifier: "en_US_POSIX") + return formatter.string(from: date) + } + + // MARK: - PTY-based probe (no tmux) + + private func loadViaPTY(model: String, timeout: TimeInterval = 10) async throws -> ClaudeUsageSnapshot { + guard let claudeBinary = ClaudeCLIResolver.resolvedBinaryPath(environment: self.environment) else { + throw ClaudeUsageError.claudeNotInstalled + } + let probe = ClaudeStatusProbe( + claudeBinary: claudeBinary, + timeout: timeout, + keepCLISessionsAlive: self.keepCLISessionsAlive) + let snap = try await probe.fetch() + + return try Self.makeSnapshot(from: snap) + } + + private func loadViaDirectCLI(timeout: TimeInterval) async throws -> ClaudeUsageSnapshot { + guard let claudeBinary = ClaudeCLIResolver.resolvedBinaryPath(environment: self.environment) else { + throw ClaudeUsageError.claudeNotInstalled + } + + let workingDirectory = ClaudeStatusProbe.preparedProbeWorkingDirectoryURL() + var environment = ClaudeCLISession.launchEnvironment(baseEnv: self.environment) + environment["PWD"] = workingDirectory.path + defer { + ClaudeProbeSessionArtifactCleaner.cleanupProbeSessionArtifacts( + probeDirectory: workingDirectory, + environment: environment) + } + + let result = try await SubprocessRunner.run( + binary: claudeBinary, + arguments: ["/usage"], + environment: environment, + timeout: timeout, + standardInput: FileHandle.nullDevice, + currentDirectoryURL: workingDirectory, + label: "claude-direct-usage") + let snap = try ClaudeStatusProbe.parse(text: result.stdout) + return try Self.makeSnapshot(from: snap) + } + + private static func makeSnapshot( + from snap: ClaudeStatusSnapshot, + now: Date = .init()) throws -> ClaudeUsageSnapshot + { + guard let sessionPctLeft = snap.sessionPercentLeft else { + throw ClaudeUsageError.parseFailed("missing session data") + } + + func makeWindow(pctLeft: Int?, reset: String?, windowMinutes: Int) -> RateWindow? { + guard let left = pctLeft else { return nil } + let used = max(0, min(100, 100 - Double(left))) + let resetClean = reset?.trimmingCharacters(in: .whitespacesAndNewlines) + return RateWindow( + usedPercent: used, + windowMinutes: windowMinutes, + resetsAt: ClaudeStatusProbe.parseResetDate( + from: resetClean, + now: now, + expectedWindow: TimeInterval(windowMinutes * 60)), + resetDescription: resetClean) + } + + let primary = makeWindow( + pctLeft: sessionPctLeft, + reset: snap.primaryResetDescription, + windowMinutes: Self.sessionWindowMinutes)! + let weekly = makeWindow( + pctLeft: snap.weeklyPercentLeft, + reset: snap.secondaryResetDescription, + windowMinutes: Self.weeklyWindowMinutes) + let opus = makeWindow( + pctLeft: snap.opusPercentLeft, + reset: snap.opusResetDescription, + windowMinutes: Self.weeklyWindowMinutes) + + return ClaudeUsageSnapshot( + primary: primary, + secondary: weekly, + opus: opus, + extraRateWindows: snap.extraRateWindows, + providerCost: nil, + updatedAt: now, + accountEmail: snap.accountEmail, + accountOrganization: snap.accountOrganization, + loginMethod: snap.loginMethod, + rawText: snap.rawText) + } + + private func applyWebExtrasIfNeeded(to snapshot: ClaudeUsageSnapshot) async -> ClaudeUsageSnapshot { + guard self.useWebExtras, self.dataSource != .web else { return snapshot } + do { + let webData: ClaudeWebAPIFetcher.WebUsageData = + if let header = self.manualCookieHeader { + try await ClaudeWebAPIFetcher.fetchUsage( + cookieHeader: header, + targetOrganizationID: self.webOrganizationID) + { msg in + Self.log.debug(msg) + } + } else { + try await ClaudeWebAPIFetcher.fetchUsage( + browserDetection: self.browserDetection, + targetOrganizationID: self.webOrganizationID) + { msg in + Self.log.debug(msg) + } + } + // Only merge usage/cost extras; keep identity fields from the primary data source. + let mergedExtraRateWindows = Self.mergeExtraRateWindows( + primary: snapshot.extraRateWindows, + web: webData.extraRateWindows) + let mergedProviderCost = snapshot.providerCost ?? webData.extraUsageCost + if mergedProviderCost != snapshot.providerCost || mergedExtraRateWindows != snapshot.extraRateWindows { + return ClaudeUsageSnapshot( + primary: snapshot.primary, + primaryWindowKind: snapshot.primaryWindowKind, + secondary: snapshot.secondary, + opus: snapshot.opus, + extraRateWindows: mergedExtraRateWindows, + providerCost: mergedProviderCost, + updatedAt: snapshot.updatedAt, + accountEmail: snapshot.accountEmail, + accountOrganization: snapshot.accountOrganization, + loginMethod: snapshot.loginMethod, + rawText: snapshot.rawText, + oauthKeychainPersistentRefHash: snapshot.oauthKeychainPersistentRefHash, + oauthHistoryOwnerIdentifier: snapshot.oauthHistoryOwnerIdentifier, + oauthKeychainCredentialMismatch: snapshot.oauthKeychainCredentialMismatch, + oauthKeychainCredentialAbsent: snapshot.oauthKeychainCredentialAbsent, + oauthKeychainCredentialUnavailable: snapshot.oauthKeychainCredentialUnavailable) + } + } catch { + Self.log.debug("Claude web extras fetch failed: \(error.localizedDescription)") + } + return snapshot + } + + private static func mergeExtraRateWindows( + primary: [NamedRateWindow], + web: [NamedRateWindow]) -> [NamedRateWindow] + { + guard !primary.isEmpty else { return web } + guard !web.isEmpty else { return primary } + + let primaryScopedKeys = Set(primary.compactMap(self.scopedExtraRateWindowMergeKey)) + var webScopedIDsByKey: [String: Set] = [:] + for window in web { + guard let scopedKey = self.scopedExtraRateWindowMergeKey(window) else { continue } + webScopedIDsByKey[scopedKey, default: []].insert(window.id) + } + var seenIDs = Set(primary.map(\.id)) + var merged = primary + for window in web where seenIDs.insert(window.id).inserted { + if let scopedKey = self.scopedExtraRateWindowMergeKey(window), + primaryScopedKeys.contains(scopedKey), + webScopedIDsByKey[scopedKey]?.count == 1 + { + continue + } + merged.append(window) + } + return merged + } + + private static func scopedExtraRateWindowMergeKey(_ window: NamedRateWindow) -> String? { + guard window.id.hasPrefix("claude-weekly-scoped-") else { return nil } + + let modelName = window.title.replacingOccurrences( + of: #"\s+only\s*$"#, + with: "", + options: [.regularExpression, .caseInsensitive]) + let normalizedName = String(modelName.lowercased().unicodeScalars.filter(CharacterSet.alphanumerics.contains)) + guard !normalizedName.isEmpty else { return nil } + return normalizedName + } + + // MARK: - Process helpers + + private static func which(_ tool: String) -> String? { + let process = Process() + process.executableURL = URL(fileURLWithPath: "/usr/bin/which") + process.arguments = [tool] + let pipe = Pipe() + process.standardOutput = pipe + try? process.run() + process.waitUntilExit() + guard process.terminationStatus == 0 else { return nil } + let data = pipe.fileHandleForReading.readDataToEndOfFile() + guard + let path = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines), + !path.isEmpty + else { return nil } + return path + } + + private static func readString(cmd: String, args: [String]) -> String? { + let task = Process() + task.executableURL = URL(fileURLWithPath: cmd) + task.arguments = args + let pipe = Pipe() + task.standardOutput = pipe + try? task.run() + task.waitUntilExit() + guard task.terminationStatus == 0 else { return nil } + let data = pipe.fileHandleForReading.readDataToEndOfFile() + return String(data: data, encoding: .utf8) + } + + private static func oauthCredentialProbeErrorLabel(_ error: Error) -> String { + guard let oauthError = error as? ClaudeOAuthCredentialsError else { + return String(describing: type(of: error)) + } + + return switch oauthError { + case .decodeFailed: + "decodeFailed" + case .missingOAuth: + "missingOAuth" + case .mcpOAuthOnlyKeychain: + "mcpOAuthOnlyKeychain" + case .missingAccessToken: + "missingAccessToken" + case .notFound: + "notFound" + case let .keychainError(status): + "keychainError:\(status)" + case .readFailed: + "readFailed" + case .refreshFailed: + "refreshFailed" + case .noRefreshToken: + "noRefreshToken" + case .refreshDelegatedToClaudeCLI: + "refreshDelegatedToClaudeCLI" + } + } +} + +#if DEBUG +extension ClaudeUsageFetcher { + public static func _mergeExtraRateWindowsForTesting( + primary: [NamedRateWindow], + web: [NamedRateWindow]) -> [NamedRateWindow] + { + self.mergeExtraRateWindows(primary: primary, web: web) + } + + public static func _mapOAuthUsageForTesting( + _ data: Data, + rateLimitTier: String? = nil, + subscriptionType: String? = nil) throws -> ClaudeUsageSnapshot + { + let usage = try ClaudeOAuthUsageFetcher.decodeUsageResponse(data) + let creds = ClaudeOAuthCredentials( + accessToken: "test", + refreshToken: nil, + expiresAt: Date().addingTimeInterval(3600), + scopes: [], + rateLimitTier: rateLimitTier, + subscriptionType: subscriptionType) + return try Self.mapOAuthUsage(usage, credentials: creds) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeWeb/ClaudeWebAPIFetcher.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeWeb/ClaudeWebAPIFetcher.swift new file mode 100644 index 0000000..59fac13 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeWeb/ClaudeWebAPIFetcher.swift @@ -0,0 +1,1270 @@ +import Foundation +import SweetCookieKit +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +enum ClaudeWebHTTPTransport { + #if DEBUG + @TaskLocal static var overrideForTesting: (any ProviderHTTPTransport)? + #endif + + static var current: any ProviderHTTPTransport { + #if DEBUG + if let override = self.overrideForTesting { + return override + } + #endif + return ProviderHTTPClient.shared + } +} + +enum ClaudeWebSessionKeyImport { + #if DEBUG + @TaskLocal static var overrideForTesting: ClaudeWebAPIFetcher.SessionKeyInfo? + #endif + + static var currentOverride: ClaudeWebAPIFetcher.SessionKeyInfo? { + #if DEBUG + self.overrideForTesting + #else + nil + #endif + } +} + +private actor ClaudeWebBrowserFetchGate { + private struct Waiter { + let id: UUID + let continuation: CheckedContinuation + } + + private var ownerID: UUID? + private var waiters: [Waiter] = [] + + func acquire(id: UUID) async -> Bool { + if Task.isCancelled { return false } + guard self.ownerID != nil else { + self.ownerID = id + return true + } + return await withCheckedContinuation { continuation in + self.waiters.append(Waiter(id: id, continuation: continuation)) + } + } + + func cancel(id: UUID) { + if self.ownerID == id { return } + guard let index = self.waiters.firstIndex(where: { $0.id == id }) else { return } + let waiter = self.waiters.remove(at: index) + waiter.continuation.resume(returning: false) + } + + func release(id: UUID) { + guard self.ownerID == id else { return } + guard !self.waiters.isEmpty else { + self.ownerID = nil + return + } + let waiter = self.waiters.removeFirst() + self.ownerID = waiter.id + waiter.continuation.resume(returning: true) + } +} + +private enum ClaudeWebBrowserFetchSerialization { + private static let gate = ClaudeWebBrowserFetchGate() + + static func run(_ operation: () async throws -> T) async throws -> T { + let id = UUID() + let acquired = await withTaskCancellationHandler { + await self.gate.acquire(id: id) + } onCancel: { + Task { await self.gate.cancel(id: id) } + } + guard acquired else { throw CancellationError() } + do { + try Task.checkCancellation() + let value = try await operation() + await self.gate.release(id: id) + return value + } catch { + await self.gate.release(id: id) + throw error + } + } +} + +/// Fetches Claude usage data directly from the claude.ai API using browser session cookies. +/// +/// This approach mirrors what Claude Usage Tracker does, but automatically extracts the session key +/// from browser cookies instead of requiring manual setup. +/// +/// API endpoints used: +/// - `GET https://claude.ai/api/organizations` → get org UUID +/// - `GET https://claude.ai/api/organizations/{org_id}/usage` → usage percentages + reset times +public enum ClaudeWebAPIFetcher { + private static let baseURL = "https://claude.ai/api" + private static let maxProbeBytes = 200_000 + #if os(macOS) + private static let cookieClient = BrowserCookieClient() + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.claude]?.browserCookieOrder ?? Browser.defaultImportOrder + #else + private static let cookieImportOrder: BrowserCookieImportOrder = [] + #endif + + public struct OrganizationInfo: Sendable { + public let id: String + public let name: String? + + public init(id: String, name: String?) { + self.id = id + self.name = name + } + } + + public struct SessionKeyInfo: Sendable { + public let key: String + public let sourceLabel: String + public let cookieCount: Int + + public init(key: String, sourceLabel: String, cookieCount: Int) { + self.key = key + self.sourceLabel = sourceLabel + self.cookieCount = cookieCount + } + } + + public enum FetchError: LocalizedError, Sendable { + case noSessionKeyFound + case invalidSessionKey + case notSupportedOnThisPlatform + case networkError(Error) + case invalidResponse + case unauthorized + case serverError(statusCode: Int) + case noOrganization + case organizationNotFound(String) + + public var errorDescription: String? { + switch self { + case .noSessionKeyFound: + "No Claude session key found in browser cookies." + case .invalidSessionKey: + "Invalid Claude session key format." + case .notSupportedOnThisPlatform: + "Claude web fetching is only supported on macOS." + case let .networkError(error): + "Network error: \(error.localizedDescription)" + case .invalidResponse: + "Invalid response from Claude API." + case .unauthorized: + "Sign in to claude.ai (or refresh Claude cookies) to load usage data." + case let .serverError(code): + "Claude API error: HTTP \(code)" + case .noOrganization: + "No Claude organization found for this account." + case let .organizationNotFound(id): + "Claude organization '\(id)' was not found for this session." + } + } + } + + /// Claude usage data from the API + public struct WebUsageData: Sendable { + public let sessionPercentUsed: Double + public let sessionResetsAt: Date? + public let weeklyPercentUsed: Double? + public let weeklyResetsAt: Date? + public let opusPercentUsed: Double? + public let extraRateWindows: [NamedRateWindow] + public let extraUsageCost: ProviderCostSnapshot? + public let accountOrganization: String? + public let accountEmail: String? + public let loginMethod: String? + /// Whether the API reported a `five_hour` session object. When `false` (the API sent + /// `five_hour: null`, as enterprise/credit accounts with no live session do), `sessionPercentUsed` + /// is the synthesized `0` placeholder rather than a real reading. Distinguishing this from a real + /// session that happens to be at `0%` (with or without a reset) lets lane classifiers drop the + /// placeholder without hiding a genuine empty session. + public let hasLiveSessionWindow: Bool + + public init( + sessionPercentUsed: Double, + sessionResetsAt: Date?, + weeklyPercentUsed: Double?, + weeklyResetsAt: Date?, + opusPercentUsed: Double?, + extraRateWindows: [NamedRateWindow], + extraUsageCost: ProviderCostSnapshot?, + accountOrganization: String?, + accountEmail: String?, + loginMethod: String?, + hasLiveSessionWindow: Bool = true) + { + self.sessionPercentUsed = sessionPercentUsed + self.sessionResetsAt = sessionResetsAt + self.weeklyPercentUsed = weeklyPercentUsed + self.weeklyResetsAt = weeklyResetsAt + self.opusPercentUsed = opusPercentUsed + self.extraRateWindows = extraRateWindows + self.extraUsageCost = extraUsageCost + self.accountOrganization = accountOrganization + self.accountEmail = accountEmail + self.loginMethod = loginMethod + self.hasLiveSessionWindow = hasLiveSessionWindow + } + } + + public struct ProbeResult: Sendable { + public let url: String + public let statusCode: Int? + public let contentType: String? + public let topLevelKeys: [String] + public let emails: [String] + public let planHints: [String] + public let notableFields: [String] + public let bodyPreview: String? + } + + // MARK: - Public API + + #if os(macOS) + + /// Attempts to fetch Claude usage data using cookies extracted from browsers. + /// Tries browser cookies using the standard import order. + public static func fetchUsage( + browserDetection: BrowserDetection, + targetOrganizationID: String? = nil, + logger: ((String) -> Void)? = nil) async throws -> WebUsageData + { + try await ClaudeWebBrowserFetchSerialization.run { + try await self.fetchUsageSerialized( + browserDetection: browserDetection, + targetOrganizationID: targetOrganizationID, + logger: logger) + } + } + + public static func fetchUsage( + cookieHeader: String, + targetOrganizationID: String? = nil, + logger: ((String) -> Void)? = nil) async throws -> WebUsageData + { + let log: (String) -> Void = { msg in logger?("[claude-web] \(msg)") } + let sessionInfo = try self.sessionKeyInfo(cookieHeader: cookieHeader) + log("Using manual session key (\(sessionInfo.cookieCount) cookies)") + return try await self.fetchUsage( + using: sessionInfo, + targetOrganizationID: targetOrganizationID, + logger: log) + } + + public static func fetchUsage( + using sessionKeyInfo: SessionKeyInfo, + targetOrganizationID: String? = nil, + logger: ((String) -> Void)? = nil) async throws -> WebUsageData + { + try await self.fetchUsage( + using: sessionKeyInfo, + targetOrganizationID: targetOrganizationID, + logger: logger, + cacheSourceLabel: nil, + expectedCacheObservation: .authoritative(nil)) + } + + private static func fetchUsageAndRenewCache( + cachedEntry: CookieHeaderCache.Entry, + targetOrganizationID: String?, + logger: ((String) -> Void)? = nil) async throws -> WebUsageData + { + let sessionInfo = try self.sessionKeyInfo(cookieHeader: cachedEntry.cookieHeader) + return try await self.fetchUsage( + using: sessionInfo, + targetOrganizationID: targetOrganizationID, + logger: logger, + cacheSourceLabel: cachedEntry.sourceLabel, + expectedCacheObservation: .authoritative(cachedEntry)) + } + + private static func fetchUsage( + using sessionKeyInfo: SessionKeyInfo, + targetOrganizationID: String?, + logger: ((String) -> Void)?, + cacheSourceLabel: String?, + expectedCacheObservation: CookieHeaderCache.ConditionalMutationObservation, + persistInitialSessionKey: Bool = false) async throws -> WebUsageData + { + let log: (String) -> Void = { msg in logger?(msg) } + let sessionKey = sessionKeyInfo.key + let renewalTracker = ClaudeWebSessionKeyRenewalTracker(initialSessionKey: sessionKey) + + // Fetch organization info + let organization = try await fetchOrganizationInfo( + sessionKey: sessionKey, + targetOrganizationID: targetOrganizationID, + logger: log, + renewalTracker: renewalTracker) + log("Organization resolved") + + var usage = try await fetchUsageData( + orgId: organization.id, + sessionKey: renewalTracker.sessionKey, + logger: log, + renewalTracker: renewalTracker) + if usage.extraUsageCost == nil, + let extra = await ClaudeWebExtraUsageCost.fetch( + baseURL: Self.baseURL, + orgId: organization.id, + sessionKey: renewalTracker.sessionKey, + logger: log, + renewalTracker: renewalTracker) + { + usage = WebUsageData( + sessionPercentUsed: usage.sessionPercentUsed, + sessionResetsAt: usage.sessionResetsAt, + weeklyPercentUsed: usage.weeklyPercentUsed, + weeklyResetsAt: usage.weeklyResetsAt, + opusPercentUsed: usage.opusPercentUsed, + extraRateWindows: usage.extraRateWindows, + extraUsageCost: extra, + accountOrganization: usage.accountOrganization, + accountEmail: usage.accountEmail, + loginMethod: usage.loginMethod, + hasLiveSessionWindow: usage.hasLiveSessionWindow) + } + if let account = await fetchAccountInfo( + sessionKey: renewalTracker.sessionKey, + orgId: organization.id, + logger: log, + renewalTracker: renewalTracker) + { + usage = WebUsageData( + sessionPercentUsed: usage.sessionPercentUsed, + sessionResetsAt: usage.sessionResetsAt, + weeklyPercentUsed: usage.weeklyPercentUsed, + weeklyResetsAt: usage.weeklyResetsAt, + opusPercentUsed: usage.opusPercentUsed, + extraRateWindows: usage.extraRateWindows, + extraUsageCost: usage.extraUsageCost, + accountOrganization: usage.accountOrganization, + accountEmail: account.email, + loginMethod: account.loginMethod, + hasLiveSessionWindow: usage.hasLiveSessionWindow) + } + if usage.accountOrganization == nil, let name = organization.name { + usage = WebUsageData( + sessionPercentUsed: usage.sessionPercentUsed, + sessionResetsAt: usage.sessionResetsAt, + weeklyPercentUsed: usage.weeklyPercentUsed, + weeklyResetsAt: usage.weeklyResetsAt, + opusPercentUsed: usage.opusPercentUsed, + extraRateWindows: usage.extraRateWindows, + extraUsageCost: usage.extraUsageCost, + accountOrganization: name, + accountEmail: usage.accountEmail, + loginMethod: usage.loginMethod, + hasLiveSessionWindow: usage.hasLiveSessionWindow) + } + if let cacheSourceLabel { + self.persistSessionKeyIfNeeded( + source: (sessionKey, cacheSourceLabel), + renewedCookieHeader: renewalTracker.renewedCookieHeader, + expectedCacheObservation: expectedCacheObservation, + persistInitialSessionKey: persistInitialSessionKey, + logger: log) + } + return usage + } + + /// Probes a list of endpoints using the current claude.ai session cookies. + /// - Parameters: + /// - endpoints: Absolute URLs or "/api/..." paths. Supports "{orgId}" placeholder. + /// - includePreview: When true, includes a truncated response preview in results. + public static func probeEndpoints( + _ endpoints: [String], + browserDetection: BrowserDetection, + includePreview: Bool = false, + logger: ((String) -> Void)? = nil) async throws -> [ProbeResult] + { + let log: (String) -> Void = { msg in logger?("[claude-probe] \(msg)") } + let sessionInfo = try extractSessionKeyInfo(browserDetection: browserDetection, logger: log) + let sessionKey = sessionInfo.key + let organization = try? await fetchOrganizationInfo(sessionKey: sessionKey, logger: log) + let expanded = endpoints.map { endpoint -> String in + var url = endpoint + if let orgId = organization?.id { + url = url.replacingOccurrences(of: "{orgId}", with: orgId) + } + if url.hasPrefix("/") { + url = "https://claude.ai\(url)" + } + return url + } + + var results: [ProbeResult] = [] + results.reserveCapacity(expanded.count) + + for endpoint in expanded { + guard let url = URL(string: endpoint) else { continue } + var request = URLRequest(url: url) + request.setValue("sessionKey=\(sessionKey)", forHTTPHeaderField: "Cookie") + request.setValue("application/json, text/html;q=0.9, */*;q=0.8", forHTTPHeaderField: "Accept") + request.httpMethod = "GET" + request.timeoutInterval = 20 + + do { + let (data, response) = try await ClaudeWebHTTPTransport.current.data(for: request) + let http = response as? HTTPURLResponse + let contentType = http?.allHeaderFields["Content-Type"] as? String + let truncated = data.prefix(Self.maxProbeBytes) + let body = String(data: truncated, encoding: .utf8) ?? "" + + let parsed = Self.parseProbeBody(data: data, fallbackText: body, contentType: contentType) + let preview = includePreview ? parsed.preview : nil + + results.append(ProbeResult( + url: endpoint, + statusCode: http?.statusCode, + contentType: contentType, + topLevelKeys: parsed.keys, + emails: parsed.emails, + planHints: parsed.planHints, + notableFields: parsed.notableFields, + bodyPreview: preview)) + } catch { + results.append(ProbeResult( + url: endpoint, + statusCode: nil, + contentType: nil, + topLevelKeys: [], + emails: [], + planHints: [], + notableFields: [], + bodyPreview: "Error: \(error.localizedDescription)")) + } + } + + return results + } + + /// Checks if we can find a Claude session key in browser cookies without making API calls. + public static func hasSessionKey(browserDetection: BrowserDetection, logger: ((String) -> Void)? = nil) -> Bool { + if let cached = CookieHeaderCache.load(provider: .claude), + self.hasSessionKey(cookieHeader: cached.cookieHeader) + { + return true + } + do { + _ = try self.sessionKeyInfo(browserDetection: browserDetection, logger: logger) + return true + } catch { + return false + } + } + + public static func hasSessionKey(cookieHeader: String?) -> Bool { + guard let cookieHeader else { return false } + return (try? self.sessionKeyInfo(cookieHeader: cookieHeader)) != nil + } + + public static func sessionKeyInfo( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) throws -> SessionKeyInfo + { + try self.extractSessionKeyInfo(browserDetection: browserDetection, logger: logger) + } + + public static func sessionKeyInfo(cookieHeader: String) throws -> SessionKeyInfo { + let pairs = CookieHeaderNormalizer.pairs(from: cookieHeader) + if let sessionKey = self.findSessionKey(in: pairs) { + return SessionKeyInfo( + key: sessionKey, + sourceLabel: "Manual", + cookieCount: pairs.count) + } + throw FetchError.noSessionKeyFound + } + + // MARK: - Session Key Extraction + + private static func extractSessionKeyInfo( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) throws -> SessionKeyInfo + { + if let override = ClaudeWebSessionKeyImport.currentOverride { return override } + let log: (String) -> Void = { msg in logger?(msg) } + + let cookieDomains = ["claude.ai"] + + // Filter to cookie-eligible browsers to avoid unnecessary keychain prompts + let installedBrowsers = Self.cookieImportOrder.cookieImportCandidates(using: browserDetection) + for browserSource in installedBrowsers { + do { + let query = BrowserCookieQuery(domains: cookieDomains) + let sources = try Self.cookieClient.codexBarRecords( + matching: query, + in: browserSource, + logger: log) + for source in sources { + if let sessionKey = findSessionKey(in: source.records.map { record in + (name: record.name, value: record.value) + }) { + log("Found sessionKey in \(source.label)") + return SessionKeyInfo( + key: sessionKey, + sourceLabel: source.label, + cookieCount: source.records.count) + } + } + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + log("\(browserSource.displayName) cookie load failed: \(error.localizedDescription)") + } + } + + throw FetchError.noSessionKeyFound + } + + private static func findSessionKey(in cookies: [(name: String, value: String)]) -> String? { + for cookie in cookies where cookie.name == "sessionKey" { + let value = cookie.value.trimmingCharacters(in: .whitespacesAndNewlines) + // Validate it looks like a Claude session key + if value.hasPrefix("sk-ant-") { + return value + } + } + return nil + } + + // MARK: - API Calls + + private static func fetchOrganizationInfo( + sessionKey: String, + targetOrganizationID: String? = nil, + logger: ((String) -> Void)? = nil, + renewalTracker: ClaudeWebSessionKeyRenewalTracker? = nil) async throws -> OrganizationInfo + { + let url = URL(string: "\(baseURL)/organizations")! + var request = URLRequest(url: url) + request.setValue("sessionKey=\(sessionKey)", forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.httpMethod = "GET" + request.timeoutInterval = 15 + + let (data, response) = try await ClaudeWebHTTPTransport.current.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw FetchError.invalidResponse + } + renewalTracker?.observe(response: httpResponse) + + logger?("Organizations API status: \(httpResponse.statusCode)") + + switch httpResponse.statusCode { + case 200: + return try self.parseOrganizationResponse(data, targetOrganizationID: targetOrganizationID) + case 401, 403: + throw FetchError.unauthorized + default: + throw FetchError.serverError(statusCode: httpResponse.statusCode) + } + } + + private static func fetchUsageData( + orgId: String, + sessionKey: String, + logger: ((String) -> Void)? = nil, + renewalTracker: ClaudeWebSessionKeyRenewalTracker? = nil) async throws -> WebUsageData + { + let url = URL(string: "\(baseURL)/organizations/\(orgId)/usage")! + var request = URLRequest(url: url) + request.setValue("sessionKey=\(sessionKey)", forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.httpMethod = "GET" + request.timeoutInterval = 15 + + let (data, response) = try await ClaudeWebHTTPTransport.current.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw FetchError.invalidResponse + } + renewalTracker?.observe(response: httpResponse) + + logger?("Usage API status: \(httpResponse.statusCode)") + + switch httpResponse.statusCode { + case 200: + return try self.parseUsageResponse(data, logger: logger) + case 401, 403: + throw FetchError.unauthorized + default: + throw FetchError.serverError(statusCode: httpResponse.statusCode) + } + } + + private static func parseUsageResponse(_ data: Data, logger: ((String) -> Void)? = nil) throws -> WebUsageData { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw FetchError.invalidResponse + } + + // Parse five_hour (session) usage + var sessionPercent: Double? + var sessionResets: Date? + let fiveHour = json["five_hour"] as? [String: Any] + if let fiveHour { + sessionPercent = Self.percentValue(from: fiveHour["utilization"]) + if let resetsAt = fiveHour["resets_at"] as? String { + sessionResets = self.parseISO8601Date(resetsAt) + } + } + // Enterprise/credit-based accounts return null for five_hour; treat as 0% rather than an error. + // Track the object's presence so a real 0% session (with or without a reset) is not mistaken for + // the synthesized null-session placeholder downstream. + let resolvedSessionPercent = sessionPercent ?? 0.0 + let hasLiveSessionWindow = fiveHour != nil + + // Parse seven_day (weekly) usage + var weeklyPercent: Double? + var weeklyResets: Date? + if let sevenDay = json["seven_day"] as? [String: Any] { + weeklyPercent = Self.percentValue(from: sevenDay["utilization"]) + if let resetsAt = sevenDay["resets_at"] as? String { + weeklyResets = self.parseISO8601Date(resetsAt) + } + } + + // Parse seven_day_sonnet (preferred) / seven_day_opus usage + var opusPercent: Double? + if let sevenDaySonnet = json["seven_day_sonnet"] as? [String: Any] { + opusPercent = Self.percentValue(from: sevenDaySonnet["utilization"]) + } else if let sevenDayOpus = json["seven_day_opus"] as? [String: Any] { + opusPercent = Self.percentValue(from: sevenDayOpus["utilization"]) + } + let extraRateParse = ClaudeWebExtraRateWindowParser.parse(from: json) + if let sourceKey = extraRateParse.sourceKeys["claude-routines"] { + logger?("Usage API extra window key matched: routines=\(sourceKey)") + } + let extraUsageCost = ClaudeWebExtraUsageCost.parse(from: json["extra_usage"]) + + return WebUsageData( + sessionPercentUsed: resolvedSessionPercent, + sessionResetsAt: sessionResets, + weeklyPercentUsed: weeklyPercent, + weeklyResetsAt: weeklyResets, + opusPercentUsed: opusPercent, + extraRateWindows: extraRateParse.windows, + extraUsageCost: extraUsageCost, + accountOrganization: nil, + accountEmail: nil, + loginMethod: nil, + hasLiveSessionWindow: hasLiveSessionWindow) + } + + private static func percentValue(from value: Any?) -> Double? { + if let intValue = value as? Int { + return Double(intValue) + } + if let doubleValue = value as? Double { + return doubleValue + } + return nil + } + + #if DEBUG + + // MARK: - Test hooks (DEBUG-only) + + public static func _parseUsageResponseForTesting(_ data: Data) throws -> WebUsageData { + try self.parseUsageResponse(data) + } + + public static func _parseOrganizationsResponseForTesting( + _ data: Data, + targetOrganizationID: String? = nil) throws -> OrganizationInfo + { + try self.parseOrganizationResponse(data, targetOrganizationID: targetOrganizationID) + } + + public static func _parseOverageSpendLimitForTesting(_ data: Data) -> ProviderCostSnapshot? { + ClaudeWebExtraUsageCost.parseOverageSpendLimit(data) + } + + public static func _parseAccountInfoForTesting(_ data: Data, orgId: String?) -> WebAccountInfo? { + self.parseAccountInfo(data, orgId: orgId) + } + #endif + + private static func parseISO8601Date(_ string: String) -> Date? { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = formatter.date(from: string) { + return date + } + // Try without fractional seconds + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: string) + } + + private static func parseOrganizationResponse( + _ data: Data, + targetOrganizationID: String? = nil) throws -> OrganizationInfo + { + guard let organizations = try? JSONDecoder().decode([ClaudeWebOrganizationResponse].self, from: data) else { + throw FetchError.invalidResponse + } + if let targetOrganizationID = targetOrganizationID?.trimmingCharacters(in: .whitespacesAndNewlines), + !targetOrganizationID.isEmpty + { + guard let selected = organizations.first(where: { $0.uuid == targetOrganizationID }) else { + throw FetchError.organizationNotFound(targetOrganizationID) + } + return self.organizationInfo(from: selected) + } + guard let selected = organizations.first(where: { $0.hasChatCapability }) + ?? organizations.first(where: { !$0.isApiOnly }) + ?? organizations.first + else { + throw FetchError.noOrganization + } + return self.organizationInfo(from: selected) + } + + private static func organizationInfo(from selected: ClaudeWebOrganizationResponse) -> OrganizationInfo { + let name = selected.name?.trimmingCharacters(in: .whitespacesAndNewlines) + let sanitized = (name?.isEmpty ?? true) ? nil : name + return OrganizationInfo(id: selected.uuid, name: sanitized) + } + + public struct WebAccountInfo: Sendable { + public let email: String? + public let loginMethod: String? + + public init(email: String?, loginMethod: String?) { + self.email = email + self.loginMethod = loginMethod + } + } + + private struct AccountResponse: Decodable { + let emailAddress: String? + let memberships: [Membership]? + + enum CodingKeys: String, CodingKey { + case emailAddress = "email_address" + case memberships + } + + struct Membership: Decodable { + let organization: Organization + + struct Organization: Decodable { + let uuid: String? + let name: String? + let rateLimitTier: String? + let billingType: String? + + enum CodingKeys: String, CodingKey { + case uuid + case name + case rateLimitTier = "rate_limit_tier" + case billingType = "billing_type" + } + } + } + } + + private static func fetchAccountInfo( + sessionKey: String, + orgId: String?, + logger: ((String) -> Void)? = nil, + renewalTracker: ClaudeWebSessionKeyRenewalTracker? = nil) async -> WebAccountInfo? + { + let url = URL(string: "\(baseURL)/account")! + var request = URLRequest(url: url) + request.setValue("sessionKey=\(sessionKey)", forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.httpMethod = "GET" + request.timeoutInterval = 15 + + do { + let (data, response) = try await ClaudeWebHTTPTransport.current.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { return nil } + renewalTracker?.observe(response: httpResponse) + logger?("Account API status: \(httpResponse.statusCode)") + guard httpResponse.statusCode == 200 else { return nil } + return Self.parseAccountInfo(data, orgId: orgId) + } catch { + return nil + } + } + + private static func parseAccountInfo(_ data: Data, orgId: String?) -> WebAccountInfo? { + guard let response = try? JSONDecoder().decode(AccountResponse.self, from: data) else { return nil } + let email = response.emailAddress?.trimmingCharacters(in: .whitespacesAndNewlines) + let membership = Self.selectMembership(response.memberships, orgId: orgId) + let plan = ClaudePlan.webLoginMethod( + rateLimitTier: membership?.organization.rateLimitTier, + billingType: membership?.organization.billingType) + return WebAccountInfo(email: email, loginMethod: plan) + } + + private static func selectMembership( + _ memberships: [AccountResponse.Membership]?, + orgId: String?) -> AccountResponse.Membership? + { + guard let memberships, !memberships.isEmpty else { return nil } + if let orgId { + if let match = memberships.first(where: { $0.organization.uuid == orgId }) { return match } + } + return memberships.first + } + + private struct ProbeParseResult { + let keys: [String] + let emails: [String] + let planHints: [String] + let notableFields: [String] + let preview: String? + } + + private static func parseProbeBody( + data: Data, + fallbackText: String, + contentType: String?) -> ProbeParseResult + { + let trimmed = fallbackText.trimmingCharacters(in: .whitespacesAndNewlines) + let looksJSON = (contentType?.lowercased().contains("application/json") ?? false) || + trimmed.hasPrefix("{") || trimmed.hasPrefix("[") + + var keys: [String] = [] + var notableFields: [String] = [] + if looksJSON, let json = try? JSONSerialization.jsonObject(with: data) { + if let dict = json as? [String: Any] { + keys = dict.keys.sorted() + } else if let array = json as? [[String: Any]], let first = array.first { + keys = first.keys.sorted() + } + notableFields = Self.extractNotableFields(from: json) + } + + let emails = Self.extractEmails(from: trimmed) + let planHints = Self.extractPlanHints(from: trimmed) + let preview = trimmed.isEmpty ? nil : String(trimmed.prefix(500)) + return ProbeParseResult( + keys: keys, + emails: emails, + planHints: planHints, + notableFields: notableFields, + preview: preview) + } + + private static func extractEmails(from text: String) -> [String] { + let pattern = #"(?i)[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { return [] } + let range = NSRange(text.startIndex.. [String] { + let pattern = #"(?i)\b(max|pro|team|ultra|enterprise)\b"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { return [] } + let range = NSRange(text.startIndex.. [String] { + let pattern = #"(?i)(plan|tier|subscription|seat|billing|product)"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { return [] } + var results: [String] = [] + + func keyMatches(_ key: String) -> Bool { + let range = NSRange(key.startIndex..= 40 { return } + let rendered: String + switch value { + case let str as String: + rendered = str + case let num as NSNumber: + rendered = num.stringValue + case let bool as Bool: + rendered = bool ? "true" : "false" + default: + return + } + let trimmed = rendered.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return } + results.append("\(keyPath)=\(trimmed)") + } + + func walk(_ value: Any, path: String) { + if let dict = value as? [String: Any] { + for (key, nested) in dict { + let nextPath = path.isEmpty ? key : "\(path).\(key)" + if keyMatches(key) { + appendValue(nextPath, value: nested) + } + walk(nested, path: nextPath) + } + } else if let array = value as? [Any] { + for (idx, nested) in array.enumerated() { + let nextPath = "\(path)[\(idx)]" + walk(nested, path: nextPath) + } + } + } + + walk(json, path: "") + return results + } + + #else + + public static func fetchUsage(logger: ((String) -> Void)? = nil) async throws -> WebUsageData { + throw FetchError.notSupportedOnThisPlatform + } + + public static func fetchUsage( + browserDetection: BrowserDetection, + targetOrganizationID: String? = nil, + logger: ((String) -> Void)? = nil) async throws -> WebUsageData + { + _ = browserDetection + _ = targetOrganizationID + _ = logger + throw FetchError.notSupportedOnThisPlatform + } + + public static func fetchUsage( + cookieHeader: String, + targetOrganizationID: String? = nil, + logger: ((String) -> Void)? = nil) async throws -> WebUsageData + { + _ = cookieHeader + _ = targetOrganizationID + _ = logger + throw FetchError.notSupportedOnThisPlatform + } + + public static func fetchUsage( + using sessionKeyInfo: SessionKeyInfo, + targetOrganizationID: String? = nil, + logger: ((String) -> Void)? = nil) async throws -> WebUsageData + { + _ = targetOrganizationID + throw FetchError.notSupportedOnThisPlatform + } + + public static func probeEndpoints( + _ endpoints: [String], + includePreview: Bool = false, + logger: ((String) -> Void)? = nil) async throws -> [ProbeResult] + { + throw FetchError.notSupportedOnThisPlatform + } + + public static func hasSessionKey(browserDetection: BrowserDetection, logger: ((String) -> Void)? = nil) -> Bool { + _ = browserDetection + _ = logger + return false + } + + public static func hasSessionKey(cookieHeader: String?) -> Bool { + guard let cookieHeader else { return false } + for pair in CookieHeaderNormalizer.pairs(from: cookieHeader) where pair.name == "sessionKey" { + let value = pair.value.trimmingCharacters(in: .whitespacesAndNewlines) + if value.hasPrefix("sk-ant-") { + return true + } + } + return false + } + + public static func sessionKeyInfo(logger: ((String) -> Void)? = nil) throws -> SessionKeyInfo { + throw FetchError.notSupportedOnThisPlatform + } + + #endif +} + +extension ClaudeWebAPIFetcher { + private static func persistSessionKeyIfNeeded( + source: (initialSessionKey: String, label: String), + renewedCookieHeader: String?, + expectedCacheObservation: CookieHeaderCache.ConditionalMutationObservation, + persistInitialSessionKey: Bool, + logger: (String) -> Void) + { + let importedCookieHeader = persistInitialSessionKey ? "sessionKey=\(source.initialSessionKey)" : nil + guard let cookieHeader = renewedCookieHeader ?? importedCookieHeader else { return } + let stored = CookieHeaderCache.storeIfObservationCurrent( + provider: .claude, + expected: expectedCacheObservation, + cookieHeader: cookieHeader, + sourceLabel: source.label) + if renewedCookieHeader != nil { + logger(stored ? "Stored renewed Claude session key" : "Skipped renewal because the cache changed") + } + } +} + +private final class ClaudeWebSessionKeyRenewalTracker: @unchecked Sendable { + private let initialSessionKey: String + private let lock = NSLock() + private var currentSessionKey: String + + init(initialSessionKey: String) { + self.initialSessionKey = initialSessionKey + self.currentSessionKey = initialSessionKey + } + + var sessionKey: String { + self.lock.lock() + defer { self.lock.unlock() } + return self.currentSessionKey + } + + var renewedCookieHeader: String? { + self.lock.lock() + defer { self.lock.unlock() } + guard self.currentSessionKey != self.initialSessionKey else { return nil } + return "sessionKey=\(self.currentSessionKey)" + } + + func observe(response: HTTPURLResponse) { + guard response.statusCode == 200, + let sessionKey = Self.sessionKey(fromSetCookieHeaders: response.allHeaderFields) + else { + return + } + self.lock.lock() + self.currentSessionKey = sessionKey + self.lock.unlock() + } + + private static func sessionKey(fromSetCookieHeaders fields: [AnyHashable: Any]) -> String? { + guard let value = fields.first(where: { + String(describing: $0.key).caseInsensitiveCompare("Set-Cookie") == .orderedSame + })?.value else { + return nil + } + var latestSessionKey: String? + for header in self.setCookieHeaderValues(from: value) { + latestSessionKey = self.sessionKey(fromSetCookieHeader: header) ?? latestSessionKey + } + return latestSessionKey + } + + private static func setCookieHeaderValues(from value: Any) -> [String] { + if let values = value as? [String] { + return values + } + if let values = value as? [Any] { + return values.map { String(describing: $0) } + } + return [String(describing: value)] + } + + private static func sessionKey(fromSetCookieHeader header: String) -> String? { + let pattern = #"(?i)(?:^|[,\r\n])\s*sessionKey=([^;,\r\n]+)"# + guard let regex = try? NSRegularExpression(pattern: pattern) else { return nil } + let range = NSRange(header.startIndex..= 2, + let valueRange = Range(match.range(at: 1), in: header) + else { + continue + } + let value = String(header[valueRange]).trimmingCharacters(in: .whitespacesAndNewlines) + if value.hasPrefix("sk-ant-") { + latestSessionKey = value + } + } + return latestSessionKey + } +} + +private enum ClaudeWebExtraUsageCost { + // MARK: - Extra usage cost (Claude "Extra") + + static func parse(from value: Any?) -> ProviderCostSnapshot? { + guard let extraUsage = value as? [String: Any] else { return nil } + guard let used = Self.doubleValue(extraUsage["used_credits"]), + let limit = Self.doubleValue(extraUsage["monthly_limit"] ?? extraUsage["monthly_credit_limit"]), + limit > 0 else { return nil } + let currency = (extraUsage["currency"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + let currencyCode = currency?.isEmpty == false ? currency ?? "USD" : "USD" + return Self.makeExtraUsageCost( + usedCredits: used, + monthlyCreditLimit: limit, + currencyCode: currencyCode) + } + + struct OverageSpendLimitResponse: Decodable { + let monthlyCreditLimit: Double? + let currency: String? + let usedCredits: Double? + let isEnabled: Bool? + + enum CodingKeys: String, CodingKey { + case monthlyCreditLimit = "monthly_credit_limit" + case currency + case usedCredits = "used_credits" + case isEnabled = "is_enabled" + } + } + + /// Best-effort fetch of Claude Extra spend/limit (does not fail the main usage fetch). + static func fetch( + baseURL: String, + orgId: String, + sessionKey: String, + logger: ((String) -> Void)? = nil, + renewalTracker: ClaudeWebSessionKeyRenewalTracker? = nil) async -> ProviderCostSnapshot? + { + let url = URL(string: "\(baseURL)/organizations/\(orgId)/overage_spend_limit")! + var request = URLRequest(url: url) + request.setValue("sessionKey=\(sessionKey)", forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.httpMethod = "GET" + request.timeoutInterval = 15 + + do { + let (data, response) = try await ClaudeWebHTTPTransport.current.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { return nil } + renewalTracker?.observe(response: httpResponse) + logger?("Overage API status: \(httpResponse.statusCode)") + guard httpResponse.statusCode == 200 else { return nil } + return Self.parseOverageSpendLimit(data) + } catch { + return nil + } + } + + static func parseOverageSpendLimit(_ data: Data) -> ProviderCostSnapshot? { + guard let decoded = try? JSONDecoder().decode(OverageSpendLimitResponse.self, from: data) else { return nil } + guard decoded.isEnabled == true else { return nil } + guard let used = decoded.usedCredits, + let limit = decoded.monthlyCreditLimit, + let currency = decoded.currency, + !currency.isEmpty else { return nil } + + return Self.makeExtraUsageCost( + usedCredits: used, + monthlyCreditLimit: limit, + currencyCode: currency) + } + + static func makeExtraUsageCost( + usedCredits: Double, + monthlyCreditLimit: Double, + currencyCode: String) -> ProviderCostSnapshot + { + let usedAmount = usedCredits / 100.0 + let limitAmount = monthlyCreditLimit / 100.0 + + return ProviderCostSnapshot( + used: usedAmount, + limit: limitAmount, + currencyCode: currencyCode, + period: "Monthly cap", + resetsAt: nil, + updatedAt: Date()) + } + + static func doubleValue(_ value: Any?) -> Double? { + switch value { + case let int as Int: + Double(int) + case let double as Double: + double + case let string as String: + Double(string) + default: + nil + } + } +} + +#if os(macOS) +extension ClaudeWebAPIFetcher { + fileprivate static func fetchUsageSerialized( + browserDetection: BrowserDetection, + targetOrganizationID: String?, + logger: ((String) -> Void)?) async throws -> WebUsageData + { + let log: (String) -> Void = { msg in logger?("[claude-web] \(msg)") } + var cacheObservation = CookieHeaderCache.observeForConditionalMutation(provider: .claude) + + if let cached = cacheObservation.entry, + !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + log("Using cached cookie header from \(cached.sourceLabel)") + do { + return try await self.fetchUsageAndRenewCache( + cachedEntry: cached, + targetOrganizationID: targetOrganizationID, + logger: log) + } catch let error as FetchError { + switch error { + case .unauthorized, .noSessionKeyFound, .invalidSessionKey: + let cleared = CookieHeaderCache.clearIfCurrent(provider: .claude, expected: cached) + cacheObservation = .authoritative(cleared ? nil : cached) + default: + throw error + } + } catch { + throw error + } + } + + let sessionInfo = try extractSessionKeyInfo(browserDetection: browserDetection, logger: log) + log("Found session key (\(sessionInfo.cookieCount) cookies)") + + return try await self.fetchUsage( + using: sessionInfo, + targetOrganizationID: targetOrganizationID, + logger: log, + cacheSourceLabel: sessionInfo.sourceLabel, + expectedCacheObservation: cacheObservation, + persistInitialSessionKey: true) + } +} +#endif + +private struct ClaudeWebOrganizationResponse: Decodable { + let uuid: String + let name: String? + let capabilities: [String]? + + var normalizedCapabilities: Set { + Set((self.capabilities ?? []).map { $0.lowercased() }) + } + + var hasChatCapability: Bool { + self.normalizedCapabilities.contains("chat") + } + + var isApiOnly: Bool { + let normalized = self.normalizedCapabilities + return !normalized.isEmpty && normalized == ["api"] + } +} diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeWeb/ClaudeWebExtraRateWindowParser.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeWeb/ClaudeWebExtraRateWindowParser.swift new file mode 100644 index 0000000..4090cf3 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeWeb/ClaudeWebExtraRateWindowParser.swift @@ -0,0 +1,123 @@ +import Foundation + +enum ClaudeWebExtraRateWindowParser { + private static let definitions: [(id: String, title: String, keys: [String])] = [ + ( + id: "claude-routines", + title: "Daily Routines", + keys: [ + "seven_day_routines", + "seven_day_claude_routines", + "claude_routines", + "routines", + "routine", + "seven_day_cowork", + "cowork", + ]), + ] + + static func parse(from json: [String: Any]) -> (windows: [NamedRateWindow], sourceKeys: [String: String]) { + var windows: [NamedRateWindow] = [] + var sourceKeys: [String: String] = [:] + windows.reserveCapacity(Self.definitions.count) + + for definition in Self.definitions { + if let foundWindow = Self.firstUsageWindow(in: json, keys: definition.keys) { + let rawWindow = foundWindow.window + guard let utilization = Self.percentValue(from: rawWindow["utilization"]) else { continue } + let resetsAt = (rawWindow["resets_at"] as? String).flatMap(Self.parseISO8601Date) + windows.append(Self.namedWindow( + id: definition.id, + title: definition.title, + usedPercent: utilization, + resetsAt: resetsAt)) + sourceKeys[definition.id] = foundWindow.sourceKey + continue + } + + // Some accounts expose the key with null payloads (for example `seven_day_cowork: null`). + // Preserve the bar in that case with a 0% window so the product section remains visible. + if let key = Self.firstUsageKey(in: json, keys: definition.keys) { + windows.append(Self.namedWindow( + id: definition.id, + title: definition.title, + usedPercent: 0, + resetsAt: nil)) + sourceKeys[definition.id] = key + } + } + windows.append(contentsOf: Self.scopedWeeklyLimitWindows(from: json)) + return (windows, sourceKeys) + } + + private static func scopedWeeklyLimitWindows(from json: [String: Any]) -> [NamedRateWindow] { + guard let limits = json["limits"] as? [[String: Any]] else { return [] } + let mappedLimits = limits.map { entry in + let scope = entry["scope"] as? [String: Any] + let model = scope?["model"] as? [String: Any] + return ClaudeScopedWeeklyLimitMapper.Limit( + kind: entry["kind"] as? String, + group: entry["group"] as? String, + percent: Self.percentValue(from: entry["percent"]), + resetsAt: (entry["resets_at"] as? String).flatMap(Self.parseISO8601Date), + modelID: model?["id"] as? String, + modelName: model?["display_name"] as? String) + } + return ClaudeScopedWeeklyLimitMapper.extraRateWindows(from: mappedLimits) + } + + private static func namedWindow( + id: String, + title: String, + usedPercent: Double, + resetsAt: Date?) -> NamedRateWindow + { + NamedRateWindow( + id: id, + title: title, + window: RateWindow( + usedPercent: usedPercent, + windowMinutes: 7 * 24 * 60, + resetsAt: resetsAt, + resetDescription: nil)) + } + + private static func firstUsageWindow( + in json: [String: Any], + keys: [String]) -> (window: [String: Any], sourceKey: String)? + { + for key in keys { + if let window = json[key] as? [String: Any] { + return (window, key) + } + } + return nil + } + + private static func firstUsageKey(in json: [String: Any], keys: [String]) -> String? { + for key in keys where json.keys.contains(key) { + return key + } + return nil + } + + private static func percentValue(from value: Any?) -> Double? { + if let intValue = value as? Int { + return Double(intValue) + } + if let doubleValue = value as? Double { + return doubleValue + } + return nil + } + + private static func parseISO8601Date(_ string: String) -> Date? { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = formatter.date(from: string) { + return date + } + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: string) + } +} diff --git a/Sources/CodexBarCore/Providers/ClawRouter/ClawRouterProviderDescriptor.swift b/Sources/CodexBarCore/Providers/ClawRouter/ClawRouterProviderDescriptor.swift new file mode 100644 index 0000000..2b52905 --- /dev/null +++ b/Sources/CodexBarCore/Providers/ClawRouter/ClawRouterProviderDescriptor.swift @@ -0,0 +1,62 @@ +import Foundation + +public enum ClawRouterProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .clawrouter, + metadata: ProviderMetadata( + id: .clawrouter, + displayName: "ClawRouter", + sessionLabel: "Monthly budget", + weeklyLabel: "Requests", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show ClawRouter usage", + cliName: "clawrouter", + defaultEnabled: false, + dashboardURL: "https://clawrouter.openclaw.ai/dashboard/access", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .clawrouter, + iconResourceName: "ProviderIcon-clawrouter", + color: ProviderColor(red: 89 / 255, green: 110 / 255, blue: 246 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "ClawRouter spend is reported by its usage API." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [ClawRouterAPIFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "clawrouter", + aliases: ["claw-router"], + versionDetector: nil)) + } +} + +struct ClawRouterAPIFetchStrategy: ProviderFetchStrategy { + let id = "clawrouter.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + ProviderTokenResolver.clawRouterToken(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = ProviderTokenResolver.clawRouterToken(environment: context.env) else { + throw ClawRouterUsageError.missingCredentials + } + try ClawRouterSettingsReader.validateEndpointOverride(environment: context.env) + let usage = try await ClawRouterUsageFetcher.fetchUsage( + apiKey: apiKey, + baseURL: ClawRouterSettingsReader.baseURL(environment: context.env)) + return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} diff --git a/Sources/CodexBarCore/Providers/ClawRouter/ClawRouterSettingsReader.swift b/Sources/CodexBarCore/Providers/ClawRouter/ClawRouterSettingsReader.swift new file mode 100644 index 0000000..4504b07 --- /dev/null +++ b/Sources/CodexBarCore/Providers/ClawRouter/ClawRouterSettingsReader.swift @@ -0,0 +1,55 @@ +import Foundation + +public enum ClawRouterSettingsError: LocalizedError, Equatable, Sendable { + case invalidEndpointOverride(String) + + public var errorDescription: String? { + switch self { + case let .invalidEndpointOverride(key): + "ClawRouter endpoint override \(key) is invalid. Use an HTTPS URL without embedded credentials." + } + } +} + +public enum ClawRouterSettingsReader { + public static let apiKeyEnvironmentKey = "CLAWROUTER_API_KEY" + public static let baseURLEnvironmentKey = "CLAWROUTER_BASE_URL" + public static let defaultBaseURL = URL(string: "https://clawrouter.openclaw.ai")! + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.apiKeyEnvironmentKey]) + } + + public static func baseURL( + environment: [String: String] = ProcessInfo.processInfo.environment) -> URL + { + guard let raw = self.cleaned(environment[self.baseURLEnvironmentKey]) else { + return self.defaultBaseURL + } + return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) ?? self.defaultBaseURL + } + + public static func validateEndpointOverride( + environment: [String: String] = ProcessInfo.processInfo.environment) throws + { + guard let raw = self.cleaned(environment[self.baseURLEnvironmentKey]) else { return } + guard ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) != nil else { + throw ClawRouterSettingsError.invalidEndpointOverride(self.baseURLEnvironmentKey) + } + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} diff --git a/Sources/CodexBarCore/Providers/ClawRouter/ClawRouterUsageFetcher.swift b/Sources/CodexBarCore/Providers/ClawRouter/ClawRouterUsageFetcher.swift new file mode 100644 index 0000000..609a790 --- /dev/null +++ b/Sources/CodexBarCore/Providers/ClawRouter/ClawRouterUsageFetcher.swift @@ -0,0 +1,255 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum ClawRouterUsageError: LocalizedError, Equatable, Sendable { + case missingCredentials + case invalidCredentials + case apiError(Int) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing ClawRouter API key. Add one in Settings or set CLAWROUTER_API_KEY." + case .invalidCredentials: + "ClawRouter rejected the API key. Check the key and its policy status." + case let .apiError(statusCode): + "ClawRouter API returned HTTP \(statusCode)." + case let .parseFailed(message): + "Could not parse ClawRouter usage: \(message)" + } + } +} + +public struct ClawRouterUsageSnapshot: Codable, Sendable, Equatable { + public struct ProviderSummary: Codable, Sendable, Equatable { + public let provider: String + public let requestCount: Int + public let successCount: Int + public let errorCount: Int + public let totalTokens: Int + public let actualCostUSD: Double + } + + public let budgetConfigured: Bool + public let budgetLedger: String + public let budgetLimitUSD: Double? + public let budgetSpentUSD: Double? + public let budgetRemainingUSD: Double? + public let budgetResetsAt: Date? + public let requestCount: Int + public let successCount: Int + public let errorCount: Int + public let inputTokens: Int + public let outputTokens: Int + public let totalTokens: Int + public let actualCostUSD: Double + public let providers: [ProviderSummary] + public let updatedAt: Date + + public func toUsageSnapshot() -> UsageSnapshot { + let usedPercent: Double? = if let spent = self.budgetSpentUSD, + let limit = self.budgetLimitUSD, + limit > 0 + { + min(100, max(0, spent / limit * 100)) + } else { + nil + } + let providerCost: ProviderCostSnapshot? = if let spent = self.budgetSpentUSD, + let limit = self.budgetLimitUSD + { + ProviderCostSnapshot( + used: spent, + limit: limit, + currencyCode: "USD", + period: "This month", + resetsAt: self.budgetResetsAt, + updatedAt: self.updatedAt) + } else if self.actualCostUSD > 0 { + ProviderCostSnapshot( + used: self.actualCostUSD, + limit: 0, + currencyCode: "USD", + period: "This month", + resetsAt: self.budgetResetsAt, + updatedAt: self.updatedAt) + } else { + nil + } + + return UsageSnapshot( + primary: usedPercent.map { + RateWindow( + usedPercent: $0, + windowMinutes: nil, + resetsAt: self.budgetResetsAt, + resetDescription: nil) + }, + secondary: nil, + providerCost: providerCost, + clawRouterUsage: self, + updatedAt: self.updatedAt, + identity: ProviderIdentitySnapshot( + providerID: .clawrouter, + accountEmail: nil, + accountOrganization: "\(self.providers.count) routed providers", + loginMethod: self.budgetConfigured ? "Managed monthly budget" : "Unmetered"), + dataConfidence: .exact) + } +} + +private struct ClawRouterUsageResponse: Decodable { + struct Budget: Decodable { + let configured: Bool + let ledger: String + let windowKey: String? + let limitMicros: Int64? + let spentMicros: Int64? + let remainingMicros: Int64? + } + + struct Usage: Decodable { + struct Summary: Decodable { + let requestCount: Int + let successCount: Int + let errorCount: Int + let inputTokens: Int + let outputTokens: Int + let totalTokens: Int + let actualCostMicros: Int64 + } + + struct Provider: Decodable { + let provider: String + let requestCount: Int + let successCount: Int + let errorCount: Int + let totalTokens: Int + let actualCostMicros: Int64 + } + + let summary: Summary + let providers: [Provider] + } + + let budget: Budget + let usage: Usage +} + +public enum ClawRouterUsageFetcher { + public static func fetchUsage( + apiKey: String, + baseURL: URL, + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + updatedAt: Date = Date()) async throws -> ClawRouterUsageSnapshot + { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw ClawRouterUsageError.missingCredentials + } + var request = URLRequest(url: self.usageURL(baseURL: baseURL)) + request.httpMethod = "GET" + request.timeoutInterval = 15 + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + let response = try await transport.response(for: request) + if response.statusCode == 401 || response.statusCode == 403 { + throw ClawRouterUsageError.invalidCredentials + } + guard (200..<300).contains(response.statusCode) else { + throw ClawRouterUsageError.apiError(response.statusCode) + } + return try self.parseSnapshot(data: response.data, updatedAt: updatedAt) + } + + public static func _parseSnapshotForTesting( + _ data: Data, + updatedAt: Date) throws -> ClawRouterUsageSnapshot + { + try self.parseSnapshot(data: data, updatedAt: updatedAt) + } + + public static func _usageURLForTesting(baseURL: URL) -> URL { + self.usageURL(baseURL: baseURL) + } + + private static func usageURL(baseURL: URL) -> URL { + let path = baseURL.path.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + let versionedBaseURL = path.split(separator: "/").last == "v1" + ? baseURL + : baseURL.appendingPathComponent("v1") + return versionedBaseURL.appendingPathComponent("usage") + } + + private static func parseSnapshot(data: Data, updatedAt: Date) throws -> ClawRouterUsageSnapshot { + do { + let response = try JSONDecoder().decode(ClawRouterUsageResponse.self, from: data) + return ClawRouterUsageSnapshot( + budgetConfigured: response.budget.configured, + budgetLedger: response.budget.ledger, + budgetLimitUSD: self.dollars(response.budget.limitMicros), + budgetSpentUSD: self.dollars(response.budget.spentMicros), + budgetRemainingUSD: self.dollars(response.budget.remainingMicros), + budgetResetsAt: self.nextMonthlyReset(windowKey: response.budget.windowKey), + requestCount: response.usage.summary.requestCount, + successCount: response.usage.summary.successCount, + errorCount: response.usage.summary.errorCount, + inputTokens: response.usage.summary.inputTokens, + outputTokens: response.usage.summary.outputTokens, + totalTokens: response.usage.summary.totalTokens, + actualCostUSD: self.dollars(response.usage.summary.actualCostMicros), + providers: response.usage.providers.map { + ClawRouterUsageSnapshot.ProviderSummary( + provider: $0.provider, + requestCount: $0.requestCount, + successCount: $0.successCount, + errorCount: $0.errorCount, + totalTokens: $0.totalTokens, + actualCostUSD: self.dollars($0.actualCostMicros)) + }.sorted { + if $0.actualCostUSD != $1.actualCostUSD { return $0.actualCostUSD > $1.actualCostUSD } + if $0.requestCount != $1.requestCount { return $0.requestCount > $1.requestCount } + return $0.provider < $1.provider + }, + updatedAt: updatedAt) + } catch let error as ClawRouterUsageError { + throw error + } catch { + throw ClawRouterUsageError.parseFailed(error.localizedDescription) + } + } + + private static func dollars(_ micros: Int64?) -> Double? { + micros.map(self.dollars) + } + + private static func dollars(_ micros: Int64) -> Double { + Double(micros) / 1_000_000 + } + + private static func nextMonthlyReset(windowKey: String?) -> Date? { + guard let suffix = windowKey?.split(separator: "/").last else { return nil } + let components = suffix.split(separator: "-") + guard components.count == 2, + let year = Int(components[0]), + let month = Int(components[1]), + (1...12).contains(month) + else { return nil } + + var nextYear = year + var nextMonth = month + 1 + if nextMonth == 13 { + nextMonth = 1 + nextYear += 1 + } + return DateComponents( + calendar: Calendar(identifier: .gregorian), + timeZone: TimeZone(secondsFromGMT: 0), + year: nextYear, + month: nextMonth, + day: 1).date + } +} diff --git a/Sources/CodexBarCore/Providers/Codebuff/CodebuffProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Codebuff/CodebuffProviderDescriptor.swift new file mode 100644 index 0000000..6ea5634 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codebuff/CodebuffProviderDescriptor.swift @@ -0,0 +1,95 @@ +import Foundation + +public enum CodebuffProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .codebuff, + metadata: ProviderMetadata( + id: .codebuff, + displayName: "Codebuff", + sessionLabel: "Credits", + weeklyLabel: "Weekly", + opusLabel: nil, + supportsOpus: false, + supportsCredits: true, + creditsHint: "Credit balance from the Codebuff API", + toggleTitle: "Show Codebuff usage", + cliName: "codebuff", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://www.codebuff.com/usage", + statusPageURL: nil, + statusLinkURL: nil), + branding: ProviderBranding( + iconStyle: .codebuff, + iconResourceName: "ProviderIcon-codebuff", + color: ProviderColor(red: 68 / 255, green: 255 / 255, blue: 0 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Codebuff cost summary is not yet supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [CodebuffAPIFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "codebuff", + aliases: ["manicode"], + versionDetector: nil)) + } +} + +struct CodebuffAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "codebuff.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + _ = context + // Keep the strategy available so missing-token surfaces as a user-friendly error + // instead of a generic "no strategy" outcome. + return true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let resolution = Self.resolveToken(environment: context.env) else { + throw CodebuffUsageError.missingCredentials + } + let usage = try await CodebuffUsageFetcher.fetchUsage( + apiKey: resolution.token, + environment: context.env, + includeSubscription: Self.shouldFetchSubscription(for: resolution)) + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + static func shouldFetchSubscription(for resolution: ProviderTokenResolution) -> Bool { + resolution.source == .authFile + } + + private static func resolveToken(environment: [String: String]) -> ProviderTokenResolution? { + ProviderTokenResolver.codebuffResolution(environment: environment) + } +} + +/// Errors related to Codebuff settings. +public enum CodebuffSettingsError: LocalizedError, Sendable, Equatable { + case missingToken + case invalidEndpointOverride(String) + + public var errorDescription: String? { + switch self { + case .missingToken: + "Codebuff API token not configured. Set CODEBUFF_API_KEY or run `codebuff login` to " + + "populate ~/.config/manicode/credentials.json." + case let .invalidEndpointOverride(key): + "Codebuff endpoint override \(key) must use HTTPS or a bare host." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Codebuff/CodebuffSettingsReader.swift b/Sources/CodexBarCore/Providers/Codebuff/CodebuffSettingsReader.swift new file mode 100644 index 0000000..633bf00 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codebuff/CodebuffSettingsReader.swift @@ -0,0 +1,86 @@ +import Foundation + +/// Reads Codebuff settings from the environment or the local credentials file +/// that the `codebuff` CLI (formerly `manicode`) writes when the user logs in. +public enum CodebuffSettingsReader { + /// Environment variable key for the Codebuff API token. + public static let apiTokenKey = "CODEBUFF_API_KEY" + + /// Returns the API token from environment if present and non-empty. + public static func apiKey(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.cleaned(environment[self.apiTokenKey]) + } + + /// Returns the API base URL, defaulting to the production endpoint. + public static func apiURL(environment: [String: String] = ProcessInfo.processInfo.environment) -> URL { + if let override = self.validAPIURL(environment: environment) { + return override + } + return URL(string: "https://www.codebuff.com")! + } + + public static func validateEndpointOverrides( + environment: [String: String] = ProcessInfo.processInfo.environment) throws + { + guard let raw = self.cleaned(environment["CODEBUFF_API_URL"]) else { return } + guard ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) == nil else { return } + throw CodebuffSettingsError.invalidEndpointOverride("CODEBUFF_API_URL") + } + + /// Returns the auth token from the local credentials file if present. + public static func authToken( + authFileURL: URL? = nil, + homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) -> String? + { + let fileURL = authFileURL ?? self.defaultAuthFileURL(homeDirectory: homeDirectory) + guard let data = try? Data(contentsOf: fileURL) else { return nil } + return self.parseAuthToken(data: data) + } + + /// Default on-disk credentials path: `~/.config/manicode/credentials.json`. + static func defaultAuthFileURL(homeDirectory: URL) -> URL { + homeDirectory + .appendingPathComponent(".config", isDirectory: true) + .appendingPathComponent("manicode", isDirectory: true) + .appendingPathComponent("credentials.json", isDirectory: false) + } + + static func parseAuthToken(data: Data) -> String? { + guard let payload = try? JSONDecoder().decode(CredentialsFile.self, from: data) else { + return nil + } + return self.cleaned(payload.default?.authToken) ?? self.cleaned(payload.authToken) + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } + + private static func validAPIURL(environment: [String: String]) -> URL? { + guard let raw = self.cleaned(environment["CODEBUFF_API_URL"]) else { return nil } + return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) + } +} + +private struct CredentialsFile: Decodable { + let `default`: CredentialsProfile? + let authToken: String? +} + +private struct CredentialsProfile: Decodable { + let authToken: String? + let fingerprintId: String? + let email: String? + let name: String? +} diff --git a/Sources/CodexBarCore/Providers/Codebuff/CodebuffUsageError.swift b/Sources/CodexBarCore/Providers/Codebuff/CodebuffUsageError.swift new file mode 100644 index 0000000..729d17f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codebuff/CodebuffUsageError.swift @@ -0,0 +1,40 @@ +import Foundation + +public enum CodebuffUsageError: LocalizedError, Sendable, Equatable { + case missingCredentials + case unauthorized + case endpointNotFound + case serviceUnavailable(Int) + case apiError(Int) + case networkError(String) + case parseFailed(String) + + public static let missingToken: CodebuffUsageError = .missingCredentials + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Codebuff API token not configured. Set CODEBUFF_API_KEY or run `codebuff login` to " + + "populate ~/.config/manicode/credentials.json." + case .unauthorized: + "Unauthorized. Please sign in to Codebuff again." + case .endpointNotFound: + "Codebuff usage endpoint not found." + case let .serviceUnavailable(status): + "Codebuff API is temporarily unavailable (status \(status))." + case let .apiError(status): + "Codebuff API returned an unexpected status (\(status))." + case let .networkError(message): + "Codebuff API error: \(message)" + case let .parseFailed(message): + "Could not parse Codebuff usage: \(message)" + } + } + + public var isAuthRelated: Bool { + switch self { + case .unauthorized, .missingCredentials: true + default: false + } + } +} diff --git a/Sources/CodexBarCore/Providers/Codebuff/CodebuffUsageFetcher.swift b/Sources/CodexBarCore/Providers/Codebuff/CodebuffUsageFetcher.swift new file mode 100644 index 0000000..830fa9b --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codebuff/CodebuffUsageFetcher.swift @@ -0,0 +1,365 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +/// Fetches live credit balance and subscription details from the Codebuff API. +/// Uses Bearer token auth against the public www.codebuff.com endpoints used by +/// the dashboard + the `codebuff` CLI. +public enum CodebuffUsageFetcher { + private static let requestTimeoutSeconds: TimeInterval = 15 + /// Extra grace period to wait for the optional subscription endpoint after the + /// primary usage call returns. Keeps the menu responsive when `/api/user/subscription` + /// is slow or hangs while `/api/v1/usage` succeeds quickly. + private static let subscriptionGraceSeconds: TimeInterval = 2 + + public static func fetchUsage( + apiKey: String, + environment: [String: String] = ProcessInfo.processInfo.environment, + includeSubscription: Bool = true, + session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> CodebuffUsageSnapshot + { + try await self.fetchUsage( + apiKey: apiKey, + environment: environment, + includeSubscription: includeSubscription, + transport: transport, + subscriptionGrace: .seconds(self.subscriptionGraceSeconds)) + } + + static func _fetchUsageForTesting( + apiKey: String, + environment: [String: String] = ProcessInfo.processInfo.environment, + includeSubscription: Bool = true, + transport: any ProviderHTTPTransport, + subscriptionGrace: Duration) async throws -> CodebuffUsageSnapshot + { + try await self.fetchUsage( + apiKey: apiKey, + environment: environment, + includeSubscription: includeSubscription, + transport: transport, + subscriptionGrace: subscriptionGrace) + } + + private static func fetchUsage( + apiKey: String, + environment: [String: String], + includeSubscription: Bool, + transport: any ProviderHTTPTransport, + subscriptionGrace: Duration) async throws -> CodebuffUsageSnapshot + { + let trimmed = apiKey.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { + throw CodebuffUsageError.missingCredentials + } + try CodebuffSettingsReader.validateEndpointOverrides(environment: environment) + + let baseURL = CodebuffSettingsReader.apiURL(environment: environment) + + let (usageValues, subscriptionValues) = try await self.fetchPayloads( + apiKey: trimmed, + baseURL: baseURL, + includeSubscription: includeSubscription, + transport: transport, + subscriptionGrace: subscriptionGrace) + + return CodebuffUsageSnapshot( + creditsUsed: usageValues.used, + creditsTotal: usageValues.total, + creditsRemaining: usageValues.remaining, + weeklyUsed: subscriptionValues?.weeklyUsed, + weeklyLimit: subscriptionValues?.weeklyLimit, + weeklyResetsAt: subscriptionValues?.weeklyResetsAt, + billingPeriodEnd: subscriptionValues?.billingPeriodEnd, + nextQuotaReset: usageValues.nextQuotaReset, + tier: subscriptionValues?.tier, + subscriptionStatus: subscriptionValues?.status, + autoTopUpEnabled: usageValues.autoTopupEnabled, + accountEmail: subscriptionValues?.email, + updatedAt: Date()) + } + + private static func fetchPayloads( + apiKey: String, + baseURL: URL, + includeSubscription: Bool, + transport: any ProviderHTTPTransport, + subscriptionGrace: Duration) async throws -> (UsagePayload, SubscriptionPayload?) + { + guard includeSubscription else { + return try await ( + self.fetchUsagePayload(apiKey: apiKey, baseURL: baseURL, transport: transport), + nil) + } + + let subscriptionTask = Task { + try await self.fetchSubscriptionPayload( + apiKey: apiKey, + baseURL: baseURL, + transport: transport) + } + let usageValues: UsagePayload + do { + usageValues = try await withTaskCancellationHandler { + try await self.fetchUsagePayload( + apiKey: apiKey, + baseURL: baseURL, + transport: transport) + } onCancel: { + subscriptionTask.cancel() + } + } catch { + subscriptionTask.cancel() + throw error + } + + do { + try Task.checkCancellation() + } catch { + subscriptionTask.cancel() + throw error + } + let race = BoundedTaskJoin(sourceTask: subscriptionTask) + switch await race.value(joinGrace: subscriptionGrace) { + case let .value(subscriptionValues): + try Task.checkCancellation() + return (usageValues, subscriptionValues) + case .timedOut: + try Task.checkCancellation() + return (usageValues, nil) + case .failure: + subscriptionTask.cancel() + try Task.checkCancellation() + return (usageValues, nil) + } + } + + // MARK: - Endpoint helpers + + struct UsagePayload { + let used: Double? + let total: Double? + let remaining: Double? + let nextQuotaReset: Date? + let autoTopupEnabled: Bool? + } + + struct SubscriptionPayload { + let status: String? + let tier: String? + let billingPeriodEnd: Date? + let weeklyUsed: Double? + let weeklyLimit: Double? + let weeklyResetsAt: Date? + let email: String? + } + + static func usageURL(baseURL: URL) -> URL { + baseURL.appendingPathComponent("/api/v1/usage") + } + + static func subscriptionURL(baseURL: URL) -> URL { + baseURL.appendingPathComponent("/api/user/subscription") + } + + static func statusError(for statusCode: Int) -> CodebuffUsageError? { + switch statusCode { + case 401, 403: .unauthorized + case 404: .endpointNotFound + case 500...599: .serviceUnavailable(statusCode) + default: nil + } + } + + static func parseUsagePayload(_ data: Data) throws -> UsagePayload { + guard let root = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw CodebuffUsageError.parseFailed("Invalid JSON") + } + + let used = self.double(from: root["usage"]) ?? self.double(from: root["used"]) + let total = self.double(from: root["quota"]) ?? self.double(from: root["limit"]) + let remaining = self.double(from: root["remainingBalance"]) ?? self.double(from: root["remaining"]) + let reset = self.date(from: root["next_quota_reset"]) + let autoTopUp = root["autoTopupEnabled"] as? Bool ?? root["auto_topup_enabled"] as? Bool + + return UsagePayload( + used: used, + total: total, + remaining: remaining, + nextQuotaReset: reset, + autoTopupEnabled: autoTopUp) + } + + static func parseSubscriptionPayload(_ data: Data) throws -> SubscriptionPayload { + guard let root = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw CodebuffUsageError.parseFailed("Invalid JSON") + } + + let subscription = root["subscription"] as? [String: Any] + let rateLimit = root["rateLimit"] as? [String: Any] + + let tier = self.string(from: subscription?["displayName"]) + ?? self.string(from: root["displayName"]) + ?? self.string(from: subscription?["tier"]) + ?? self.string(from: root["tier"]) + ?? self.string(from: subscription?["scheduledTier"]) + let status = subscription?["status"] as? String + let email = root["email"] as? String ?? (root["user"] as? [String: Any])?["email"] as? String + let billingPeriodEnd = self.date(from: subscription?["billingPeriodEnd"]) + ?? self.date(from: subscription?["currentPeriodEnd"]) + let weeklyUsed = self.double(from: rateLimit?["weeklyUsed"]) + ?? self.double(from: rateLimit?["used"]) + let weeklyLimit = self.double(from: rateLimit?["weeklyLimit"]) + ?? self.double(from: rateLimit?["limit"]) + let weeklyResetsAt = self.date(from: rateLimit?["weeklyResetsAt"]) + + return SubscriptionPayload( + status: status, + tier: tier, + billingPeriodEnd: billingPeriodEnd, + weeklyUsed: weeklyUsed, + weeklyLimit: weeklyLimit, + weeklyResetsAt: weeklyResetsAt, + email: email) + } + + // MARK: - Test hooks + + static func _parseUsagePayloadForTesting(_ data: Data) throws -> UsagePayload { + try self.parseUsagePayload(data) + } + + static func _parseSubscriptionPayloadForTesting(_ data: Data) throws -> SubscriptionPayload { + try self.parseSubscriptionPayload(data) + } + + static func _statusErrorForTesting(_ statusCode: Int) -> CodebuffUsageError? { + self.statusError(for: statusCode) + } + + // MARK: - Networking + + private static func fetchUsagePayload( + apiKey: String, + baseURL: URL, + transport: any ProviderHTTPTransport) async throws -> UsagePayload + { + var request = URLRequest(url: self.usageURL(baseURL: baseURL)) + request.httpMethod = "POST" + request.timeoutInterval = self.requestTimeoutSeconds + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.httpBody = try? JSONSerialization.data(withJSONObject: ["fingerprintId": "codexbar-usage"]) + + let response = try await self.send(request: request, transport: transport) + if let err = self.statusError(for: response.statusCode) { + throw err + } + guard response.statusCode == 200 else { + throw CodebuffUsageError.apiError(response.statusCode) + } + return try self.parseUsagePayload(response.data) + } + + private static func fetchSubscriptionPayload( + apiKey: String, + baseURL: URL, + transport: any ProviderHTTPTransport) async throws -> SubscriptionPayload + { + var request = URLRequest(url: self.subscriptionURL(baseURL: baseURL)) + request.httpMethod = "GET" + request.timeoutInterval = self.requestTimeoutSeconds + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + let response = try await self.send(request: request, transport: transport) + if let err = self.statusError(for: response.statusCode) { + throw err + } + guard response.statusCode == 200 else { + throw CodebuffUsageError.apiError(response.statusCode) + } + return try self.parseSubscriptionPayload(response.data) + } + + private static func send( + request: URLRequest, + transport: any ProviderHTTPTransport) async throws -> ProviderHTTPResponse + { + do { + return try await transport.response(for: request) + } catch let error as CodebuffUsageError { + throw error + } catch let error as URLError where error.code == .badServerResponse { + throw CodebuffUsageError.networkError("Invalid response") + } catch { + throw CodebuffUsageError.networkError(error.localizedDescription) + } + } + + // MARK: - Value parsing + + private static func double(from value: Any?) -> Double? { + switch value { + case let number as NSNumber: + let raw = number.doubleValue + return raw.isFinite ? raw : nil + case let string as String: + let trimmed = string.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty, let raw = Double(trimmed), raw.isFinite else { return nil } + return raw + default: + return nil + } + } + + private static func string(from value: Any?) -> String? { + switch value { + case let string as String: + let trimmed = string.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + case let number as NSNumber: + let raw = number.doubleValue + guard raw.isFinite else { return nil } + return number.stringValue + default: + return nil + } + } + + private static func date(from value: Any?) -> Date? { + switch value { + case let string as String: + let trimmed = string.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + let fractional = ISO8601DateFormatter() + fractional.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = fractional.date(from: trimmed) { + return date + } + let plain = ISO8601DateFormatter() + plain.formatOptions = [.withInternetDateTime] + if let date = plain.date(from: trimmed) { + return date + } + if let interval = Double(trimmed), interval.isFinite { + return Self.dateFromNumeric(interval) + } + return nil + case let number as NSNumber: + let raw = number.doubleValue + return raw.isFinite ? Self.dateFromNumeric(raw) : nil + default: + return nil + } + } + + private static func dateFromNumeric(_ value: Double) -> Date? { + if value > 10_000_000_000 { + return Date(timeIntervalSince1970: value / 1000) + } + return Date(timeIntervalSince1970: value) + } +} diff --git a/Sources/CodexBarCore/Providers/Codebuff/CodebuffUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Codebuff/CodebuffUsageSnapshot.swift new file mode 100644 index 0000000..478cc8e --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codebuff/CodebuffUsageSnapshot.swift @@ -0,0 +1,147 @@ +import Foundation + +/// Parsed view of a Codebuff usage + subscription response pair. +public struct CodebuffUsageSnapshot: Sendable { + public let creditsUsed: Double? + public let creditsTotal: Double? + public let creditsRemaining: Double? + public let weeklyUsed: Double? + public let weeklyLimit: Double? + public let weeklyResetsAt: Date? + public let billingPeriodEnd: Date? + public let nextQuotaReset: Date? + public let tier: String? + public let subscriptionStatus: String? + public let autoTopUpEnabled: Bool? + public let accountEmail: String? + public let updatedAt: Date + + public init( + creditsUsed: Double? = nil, + creditsTotal: Double? = nil, + creditsRemaining: Double? = nil, + weeklyUsed: Double? = nil, + weeklyLimit: Double? = nil, + weeklyResetsAt: Date? = nil, + billingPeriodEnd: Date? = nil, + nextQuotaReset: Date? = nil, + tier: String? = nil, + subscriptionStatus: String? = nil, + autoTopUpEnabled: Bool? = nil, + accountEmail: String? = nil, + updatedAt: Date = Date()) + { + self.creditsUsed = creditsUsed + self.creditsTotal = creditsTotal + self.creditsRemaining = creditsRemaining + self.weeklyUsed = weeklyUsed + self.weeklyLimit = weeklyLimit + self.weeklyResetsAt = weeklyResetsAt + self.billingPeriodEnd = billingPeriodEnd + self.nextQuotaReset = nextQuotaReset + self.tier = tier + self.subscriptionStatus = subscriptionStatus + self.autoTopUpEnabled = autoTopUpEnabled + self.accountEmail = accountEmail + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let primary = self.makeCreditsWindow() + let secondary = self.makeWeeklyWindow() + + let identity = ProviderIdentitySnapshot( + providerID: .codebuff, + accountEmail: self.accountEmail, + accountOrganization: nil, + loginMethod: self.makeLoginMethod()) + + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } + + private func makeCreditsWindow() -> RateWindow? { + let total = self.resolvedTotal + guard let total, total > 0 else { + if self.creditsRemaining != nil || self.creditsUsed != nil { + // Degenerate case: no usable quota in the payload. Surface the row as fully + // exhausted so missing quota data is visibly surfaced (matches Kilo's behaviour + // for zero/unknown totals) rather than rendering a misleading healthy bar. + return RateWindow( + usedPercent: 100, + windowMinutes: nil, + resetsAt: self.nextQuotaReset, + resetDescription: nil) + } + return nil + } + let used = self.resolvedUsed + let percent = min(100, max(0, (used / total) * 100)) + // Note: do not stuff the credit balance ("X/Y credits") into `resetDescription` — + // generic renderers (UsageFormatter.resetLine) prepend "Resets " when `resetsAt` + // is absent, which would surface misleading text like "Resets 250/1,000 credits". + // The credits detail is shown via the dedicated Codebuff account panel instead. + return RateWindow( + usedPercent: percent, + windowMinutes: nil, + resetsAt: self.nextQuotaReset, + resetDescription: nil) + } + + private func makeWeeklyWindow() -> RateWindow? { + guard let limit = self.weeklyLimit, limit > 0 else { return nil } + let used = max(0, self.weeklyUsed ?? 0) + let percent = min(100, max(0, (used / limit) * 100)) + // Same reasoning as above: avoid encoding non-reset detail in `resetDescription`. + return RateWindow( + usedPercent: percent, + windowMinutes: 7 * 24 * 60, + resetsAt: self.weeklyResetsAt, + resetDescription: nil) + } + + private var resolvedTotal: Double? { + if let creditsTotal { return max(0, creditsTotal) } + if let creditsUsed, let creditsRemaining { + return max(0, creditsUsed + creditsRemaining) + } + return nil + } + + private var resolvedUsed: Double { + if let creditsUsed { + return max(0, creditsUsed) + } + if let total = self.resolvedTotal, let creditsRemaining { + return max(0, total - creditsRemaining) + } + return 0 + } + + private func makeLoginMethod() -> String? { + var parts: [String] = [] + if let tier = self.tier?.trimmingCharacters(in: .whitespacesAndNewlines), !tier.isEmpty { + parts.append(tier.capitalized) + } + if let remaining = self.creditsRemaining { + parts.append("\(Self.compactNumber(remaining)) remaining") + } + if self.autoTopUpEnabled == true { + parts.append("auto top-up") + } + return parts.isEmpty ? nil : parts.joined(separator: " · ") + } + + static func compactNumber(_ value: Double) -> String { + let formatter = NumberFormatter() + formatter.numberStyle = .decimal + formatter.locale = Locale(identifier: "en_US") + formatter.maximumFractionDigits = value >= 1000 ? 0 : 1 + return formatter.string(from: NSNumber(value: value)) ?? String(format: "%.0f", value) + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexAccountReconciliation.swift b/Sources/CodexBarCore/Providers/Codex/CodexAccountReconciliation.swift new file mode 100644 index 0000000..d157767 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexAccountReconciliation.swift @@ -0,0 +1,474 @@ +import Foundation + +public struct CodexResolvedActiveSource: Equatable, Sendable { + public let persistedSource: CodexActiveSource + public let resolvedSource: CodexActiveSource + + public init(persistedSource: CodexActiveSource, resolvedSource: CodexActiveSource) { + self.persistedSource = persistedSource + self.resolvedSource = resolvedSource + } + + public var requiresPersistenceCorrection: Bool { + self.persistedSource != self.resolvedSource + } +} + +public enum CodexActiveSourceResolver { + public static func resolve(from snapshot: CodexAccountReconciliationSnapshot) -> CodexResolvedActiveSource { + let persistedSource = snapshot.activeSource + let resolvedSource: CodexActiveSource = switch persistedSource { + case .liveSystem: + .liveSystem + case let .managedAccount(id): + if let activeStoredAccount = snapshot.activeStoredAccount { + self.resolvedSource(for: activeStoredAccount, snapshot: snapshot) + } else { + snapshot.liveSystemAccount != nil ? .liveSystem : .managedAccount(id: id) + } + case let .profileHome(path): + if let normalizedPath = snapshot.configuredProfileHomePath(path: path) { + self.resolvedProfileSource(path: normalizedPath, snapshot: snapshot) + } else { + .liveSystem + } + } + + return CodexResolvedActiveSource( + persistedSource: persistedSource, + resolvedSource: resolvedSource) + } + + private static func resolvedProfileSource( + path: String, + snapshot: CodexAccountReconciliationSnapshot) -> CodexActiveSource + { + if let livePath = snapshot.liveSystemAccount.flatMap({ + CodexHomeScope.normalizedHomePath($0.codexHomePath) + }), livePath == path { + return .liveSystem + } + + if let storedAccount = snapshot.storedAccounts.first(where: { + CodexHomeScope.normalizedHomePath($0.managedHomePath) == path + }) { + return self.resolvedSource(for: storedAccount, snapshot: snapshot) + } + + return .profileHome(path: path) + } + + private static func resolvedSource( + for storedAccount: ManagedCodexAccount, + snapshot: CodexAccountReconciliationSnapshot) -> CodexActiveSource + { + self.matchesLiveSystemAccount( + storedAccount: storedAccount, + snapshot: snapshot, + liveSystemAccount: snapshot.liveSystemAccount) ? .liveSystem : .managedAccount(id: storedAccount.id) + } + + private static func matchesLiveSystemAccount( + storedAccount: ManagedCodexAccount, + snapshot: CodexAccountReconciliationSnapshot, + liveSystemAccount: ObservedSystemCodexAccount?) -> Bool + { + guard let liveSystemAccount else { return false } + if let storedFingerprint = storedAccount.authFingerprint, + let liveFingerprint = liveSystemAccount.authFingerprint, + storedFingerprint == liveFingerprint + { + return true + } + return CodexIdentityMatcher.matches( + snapshot.runtimeIdentity(for: storedAccount), + lhsEmail: snapshot.runtimeEmail(for: storedAccount), + snapshot.runtimeIdentity(for: liveSystemAccount), + rhsEmail: liveSystemAccount.email) + } +} + +public struct CodexAccountReconciliationSnapshot: Equatable, Sendable { + public let storedAccounts: [ManagedCodexAccount] + public let activeStoredAccount: ManagedCodexAccount? + public let liveSystemAccount: ObservedSystemCodexAccount? + public let profileHomeAccounts: [ObservedSystemCodexAccount] + public let profileHomePaths: [String] + public let matchingStoredAccountForLiveSystemAccount: ManagedCodexAccount? + public let activeSource: CodexActiveSource + public let hasUnreadableAddedAccountStore: Bool + public let storedAccountRuntimeIdentities: [UUID: CodexIdentity] + public let storedAccountRuntimeEmails: [UUID: String] + + public init( + storedAccounts: [ManagedCodexAccount], + activeStoredAccount: ManagedCodexAccount?, + liveSystemAccount: ObservedSystemCodexAccount?, + profileHomeAccounts: [ObservedSystemCodexAccount] = [], + profileHomePaths: [String]? = nil, + matchingStoredAccountForLiveSystemAccount: ManagedCodexAccount?, + activeSource: CodexActiveSource, + hasUnreadableAddedAccountStore: Bool, + storedAccountRuntimeIdentities: [UUID: CodexIdentity] = [:], + storedAccountRuntimeEmails: [UUID: String] = [:]) + { + self.storedAccounts = storedAccounts + self.activeStoredAccount = activeStoredAccount + self.liveSystemAccount = liveSystemAccount + self.profileHomeAccounts = Self.uniqueProfileHomeAccounts(profileHomeAccounts) + self.profileHomePaths = Self.uniqueNormalizedPaths( + profileHomePaths ?? profileHomeAccounts.map(\.codexHomePath)) + self.matchingStoredAccountForLiveSystemAccount = matchingStoredAccountForLiveSystemAccount + self.activeSource = activeSource + self.hasUnreadableAddedAccountStore = hasUnreadableAddedAccountStore + self.storedAccountRuntimeIdentities = storedAccountRuntimeIdentities + self.storedAccountRuntimeEmails = storedAccountRuntimeEmails + } + + public static func == (lhs: CodexAccountReconciliationSnapshot, rhs: CodexAccountReconciliationSnapshot) -> Bool { + lhs.storedAccounts.map(AccountIdentity.init) == rhs.storedAccounts.map(AccountIdentity.init) + && lhs.activeStoredAccount.map(AccountIdentity.init) == rhs.activeStoredAccount.map(AccountIdentity.init) + && lhs.liveSystemAccount == rhs.liveSystemAccount + && lhs.profileHomeAccounts == rhs.profileHomeAccounts + && lhs.profileHomePaths == rhs.profileHomePaths + && lhs.matchingStoredAccountForLiveSystemAccount.map(AccountIdentity.init) + == rhs.matchingStoredAccountForLiveSystemAccount.map(AccountIdentity.init) + && lhs.activeSource == rhs.activeSource + && lhs.hasUnreadableAddedAccountStore == rhs.hasUnreadableAddedAccountStore + && lhs.storedAccountRuntimeIdentities == rhs.storedAccountRuntimeIdentities + && lhs.storedAccountRuntimeEmails == rhs.storedAccountRuntimeEmails + } + + public func runtimeIdentity(for storedAccount: ManagedCodexAccount) -> CodexIdentity { + self.storedAccountRuntimeIdentities[storedAccount.id] + ?? CodexIdentityResolver.resolve(accountId: nil, email: storedAccount.email) + } + + public func runtimeEmail(for storedAccount: ManagedCodexAccount) -> String { + self.storedAccountRuntimeEmails[storedAccount.id] + ?? Self.normalizeEmail(storedAccount.email) + } + + public func runtimeIdentity(for liveSystemAccount: ObservedSystemCodexAccount) -> CodexIdentity { + CodexIdentityMatcher.normalized( + liveSystemAccount.identity, + fallbackEmail: liveSystemAccount.email) + } + + public func profileHomeAccount(path: String) -> ObservedSystemCodexAccount? { + guard let normalizedPath = CodexHomeScope.normalizedHomePath(path) else { return nil } + return self.profileHomeAccounts.first { + CodexHomeScope.normalizedHomePath($0.codexHomePath) == normalizedPath + } + } + + public func configuredProfileHomePath(path: String) -> String? { + guard let normalizedPath = CodexHomeScope.normalizedHomePath(path), + self.profileHomePaths.contains(normalizedPath) + else { + return nil + } + return normalizedPath + } + + private static func normalizeEmail(_ email: String) -> String { + email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } + + private static func uniqueProfileHomeAccounts( + _ accounts: [ObservedSystemCodexAccount]) -> [ObservedSystemCodexAccount] + { + var seenPaths: Set = [] + var result: [ObservedSystemCodexAccount] = [] + for account in accounts { + let path = CodexHomeScope.normalizedHomePath(account.codexHomePath) ?? account.codexHomePath + guard seenPaths.insert(path).inserted else { continue } + result.append(account) + } + return result + } + + private static func uniqueNormalizedPaths(_ paths: [String]) -> [String] { + var seen: Set = [] + return paths.compactMap { path in + guard let normalizedPath = CodexHomeScope.normalizedHomePath(path), + seen.insert(normalizedPath).inserted + else { + return nil + } + return normalizedPath + } + } +} + +public struct DefaultCodexAccountReconciler: Sendable { + public let storeLoader: @Sendable () throws -> ManagedCodexAccountSet + public let systemObserver: any CodexSystemAccountObserving + public let activeSource: CodexActiveSource + public let baseEnvironment: [String: String] + public let profileHomePaths: [String] + public let managedEnvironmentBuilder: @Sendable ([String: String], ManagedCodexAccount) -> [String: String] + + public init( + storeLoader: @escaping @Sendable () throws -> ManagedCodexAccountSet = { + try FileManagedCodexAccountStore().loadAccounts() + }, + systemObserver: any CodexSystemAccountObserving = DefaultCodexSystemAccountObserver(), + activeSource: CodexActiveSource = .liveSystem, + baseEnvironment: [String: String], + profileHomePaths: [String] = [], + managedEnvironmentBuilder: @escaping @Sendable ([String: String], ManagedCodexAccount) + -> [String: String] = { baseEnvironment, account in + CodexHomeScope.scopedEnvironment(base: baseEnvironment, codexHome: account.managedHomePath) + }) + { + self.storeLoader = storeLoader + self.systemObserver = systemObserver + self.activeSource = activeSource + self.baseEnvironment = baseEnvironment + self.profileHomePaths = Self.uniqueNormalizedPaths(profileHomePaths) + self.managedEnvironmentBuilder = managedEnvironmentBuilder + } + + public func loadSnapshot() -> CodexAccountReconciliationSnapshot { + let liveSystemAccount = self.loadLiveSystemAccount() + let profileHomeAccounts = self.loadProfileHomeAccounts(liveSystemAccount: liveSystemAccount) + + do { + let accounts = try self.storeLoader() + let runtimeAccounts = Dictionary(uniqueKeysWithValues: accounts.accounts.map { account in + let runtimeAccount = self.loadRuntimeAccount(for: account) + return (account.id, runtimeAccount) + }) + let activeStoredAccount: ManagedCodexAccount? = switch self.activeSource { + case let .managedAccount(id): + accounts.account(id: id) + case .liveSystem, .profileHome: + nil + } + let matchingStoredAccountForLiveSystemAccount = liveSystemAccount.flatMap { liveAccount in + if let liveFingerprint = liveAccount.authFingerprint, + let exactFingerprintMatch = accounts.accounts.first(where: { + $0.authFingerprint == liveFingerprint + }) + { + return exactFingerprintMatch + } + return accounts.accounts.first { account in + guard let runtimeAccount = runtimeAccounts[account.id] else { return false } + return CodexIdentityMatcher.matches( + runtimeAccount.identity, + lhsEmail: runtimeAccount.email, + self.runtimeIdentity(for: liveAccount), + rhsEmail: liveAccount.email) + } + } + + return CodexAccountReconciliationSnapshot( + storedAccounts: accounts.accounts, + activeStoredAccount: activeStoredAccount, + liveSystemAccount: liveSystemAccount, + profileHomeAccounts: self.profileHomeAccounts( + profileHomeAccounts, + excludingManagedAccounts: accounts.accounts), + profileHomePaths: self.profileHomePaths, + matchingStoredAccountForLiveSystemAccount: matchingStoredAccountForLiveSystemAccount, + activeSource: self.activeSource, + hasUnreadableAddedAccountStore: false, + storedAccountRuntimeIdentities: runtimeAccounts.mapValues(\.identity), + storedAccountRuntimeEmails: runtimeAccounts.mapValues(\.email)) + } catch { + return CodexAccountReconciliationSnapshot( + storedAccounts: [], + activeStoredAccount: nil, + liveSystemAccount: liveSystemAccount, + profileHomeAccounts: profileHomeAccounts, + profileHomePaths: self.profileHomePaths, + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: self.activeSource, + hasUnreadableAddedAccountStore: true) + } + } + + private func loadProfileHomeAccounts( + liveSystemAccount: ObservedSystemCodexAccount?) -> [ObservedSystemCodexAccount] + { + let livePath = liveSystemAccount.flatMap { CodexHomeScope.normalizedHomePath($0.codexHomePath) } + return self.profileHomePaths.compactMap { path in + guard path != livePath else { return nil } + return self.loadProfileHomeAccount(homePath: path) + } + } + + private func loadProfileHomeAccount(homePath: String) -> ObservedSystemCodexAccount? { + let environment = CodexHomeScope.scopedEnvironment(base: self.baseEnvironment, codexHome: homePath) + let account = UsageFetcher(environment: environment).loadAuthBackedCodexAccount() + + guard let rawEmail = account.email?.trimmingCharacters(in: .whitespacesAndNewlines), + !rawEmail.isEmpty + else { + return nil + } + + let providerAccountID: String? = switch account.identity { + case let .providerAccount(id): + ManagedCodexAccount.normalizeProviderAccountID(id) + case .emailOnly, .unresolved: + nil + } + + return ObservedSystemCodexAccount( + email: rawEmail.lowercased(), + workspaceAccountID: providerAccountID, + authFingerprint: CodexAuthFingerprint.fingerprint(homePath: homePath), + codexHomePath: homePath, + observedAt: Date(), + identity: account.identity) + } + + private func loadLiveSystemAccount() -> ObservedSystemCodexAccount? { + do { + guard let account = try self.systemObserver.loadSystemAccount(environment: self.baseEnvironment) else { + return nil + } + let normalizedEmail = Self.normalizeEmail(account.email) + guard !normalizedEmail.isEmpty else { + return nil + } + return ObservedSystemCodexAccount( + email: normalizedEmail, + workspaceLabel: account.workspaceLabel, + workspaceAccountID: account.workspaceAccountID, + authFingerprint: account.authFingerprint, + codexHomePath: account.codexHomePath, + observedAt: account.observedAt, + identity: self.runtimeIdentity(for: account)) + } catch { + return nil + } + } + + private static func normalizeEmail(_ email: String) -> String { + email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } + + private func loadRuntimeAccount(for account: ManagedCodexAccount) -> RuntimeManagedCodexAccount { + let scopedEnvironment = self.managedEnvironmentBuilder(self.baseEnvironment, account) + let authBackedAccount = UsageFetcher(environment: scopedEnvironment).loadAuthBackedCodexAccount() + let email = Self.normalizeEmail(authBackedAccount.email ?? account.email) + let identity = CodexIdentityMatcher.normalized(authBackedAccount.identity, fallbackEmail: email) + + return RuntimeManagedCodexAccount( + email: email, + identity: identity) + } + + private func profileHomeAccounts( + _ profileHomeAccounts: [ObservedSystemCodexAccount], + excludingManagedAccounts managedAccounts: [ManagedCodexAccount]) -> [ObservedSystemCodexAccount] + { + let managedPaths = Set(managedAccounts.compactMap { CodexHomeScope.normalizedHomePath($0.managedHomePath) }) + return profileHomeAccounts.filter { account in + guard let path = CodexHomeScope.normalizedHomePath(account.codexHomePath) else { return false } + return !managedPaths.contains(path) + } + } + + private func runtimeIdentity(for liveSystemAccount: ObservedSystemCodexAccount) -> CodexIdentity { + CodexIdentityMatcher.normalized( + liveSystemAccount.identity, + fallbackEmail: liveSystemAccount.email) + } + + private static func uniqueNormalizedPaths(_ paths: [String]) -> [String] { + var seen: Set = [] + var result: [String] = [] + for path in paths.compactMap({ CodexHomeScope.normalizedHomePath($0) }) { + guard seen.insert(path).inserted else { continue } + result.append(path) + } + return result + } +} + +public enum CodexIdentityMatcher { + public static func matches(_ lhs: CodexIdentity, _ rhs: CodexIdentity) -> Bool { + switch (lhs, rhs) { + case let (.providerAccount(leftID), .providerAccount(rightID)): + leftID == rightID + case let (.emailOnly(leftEmail), .emailOnly(rightEmail)): + leftEmail == rightEmail + default: + false + } + } + + public static func matches( + _ lhs: CodexIdentity, + lhsEmail: String?, + _ rhs: CodexIdentity, + rhsEmail: String?) -> Bool + { + guard self.matches(lhs, rhs) else { return false } + guard case .providerAccount = lhs, case .providerAccount = rhs else { return true } + guard let normalizedLeftEmail = CodexIdentityResolver.normalizeEmail(lhsEmail), + let normalizedRightEmail = CodexIdentityResolver.normalizeEmail(rhsEmail) + else { + return true + } + return normalizedLeftEmail == normalizedRightEmail + } + + public static func normalized(_ identity: CodexIdentity, fallbackEmail: String) -> CodexIdentity { + switch identity { + case .providerAccount: + identity + case let .emailOnly(normalizedEmail): + CodexIdentityResolver.resolve(accountId: nil, email: normalizedEmail) + case .unresolved: + CodexIdentityResolver.resolve(accountId: nil, email: fallbackEmail) + } + } + + public static func selectionKey(for identity: CodexIdentity, fallbackEmail: String) -> String { + switch self.normalized(identity, fallbackEmail: fallbackEmail) { + case let .providerAccount(id): + "provider:\(id)" + case let .emailOnly(normalizedEmail): + "email:\(normalizedEmail)" + case .unresolved: + "unresolved:\(fallbackEmail)" + } + } +} + +private struct RuntimeManagedCodexAccount { + let email: String + let identity: CodexIdentity +} + +private struct AccountIdentity: Equatable { + let id: UUID + let email: String + let providerAccountID: String? + let workspaceLabel: String? + let workspaceAccountID: String? + let managedHomePath: String + let createdAt: TimeInterval + let updatedAt: TimeInterval + let lastAuthenticatedAt: TimeInterval? + let authFingerprint: String? + + init(_ account: ManagedCodexAccount) { + self.id = account.id + self.email = account.email + self.providerAccountID = account.providerAccountID + self.workspaceLabel = account.workspaceLabel + self.workspaceAccountID = account.workspaceAccountID + self.managedHomePath = account.managedHomePath + self.createdAt = account.createdAt + self.updatedAt = account.updatedAt + self.lastAuthenticatedAt = account.lastAuthenticatedAt + self.authFingerprint = account.authFingerprint + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexAdditionalRateLimitMapper.swift b/Sources/CodexBarCore/Providers/Codex/CodexAdditionalRateLimitMapper.swift new file mode 100644 index 0000000..c37f28c --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexAdditionalRateLimitMapper.swift @@ -0,0 +1,153 @@ +import Foundation + +/// Maps Codex `additional_rate_limits` entries (model-specific limits such as GPT-5.3-Codex-Spark) +/// into named extra rate windows. +/// +/// These limits are reported by the `wham/usage` API alongside, but separately from, the primary and +/// weekly Codex windows, so we surface them through `UsageSnapshot.extraRateWindows` rather than the core +/// primary/secondary lanes. When the field is absent the mapper returns an empty list and the snapshot is +/// unchanged. +package enum CodexAdditionalRateLimitMapper { + /// Stable ids/titles for GPT-5.3-Codex-Spark limits so SwiftUI identity stays constant even if the API + /// label wording shifts. Keep the original 5-hour id for compatibility with the first Spark implementation. + package static let sparkWindowID = "codex-spark" + package static let sparkWeeklyWindowID = "codex-spark-weekly" + static let sparkWindowTitle = "Codex Spark 5-hour" + static let sparkWeeklyWindowTitle = "Codex Spark Weekly" + + static func extraRateWindows( + from additionalRateLimits: [CodexUsageResponse.AdditionalRateLimit]?, + now: Date = Date()) -> [NamedRateWindow] + { + guard let additionalRateLimits, !additionalRateLimits.isEmpty else { return [] } + var usedIDs = Set() + return additionalRateLimits.flatMap { entry in + self.namedWindows(from: entry, usedIDs: &usedIDs, now: now) + } + } + + private static func namedWindows( + from entry: CodexUsageResponse.AdditionalRateLimit, + usedIDs: inout Set, + now: Date) -> [NamedRateWindow] + { + if self.isSpark(entry) { + return self.sparkWindows(from: entry, usedIDs: &usedIDs, now: now) + } + + // Model-specific limits report utilization in the primary window; fall back to the secondary + // window only when a primary one is not present. + guard let snapshot = entry.rateLimit?.primaryWindow ?? entry.rateLimit?.secondaryWindow else { + return [] + } + guard let id = self.windowID(for: entry), usedIDs.insert(id).inserted else { return [] } + return [self.namedWindow( + id: id, + title: self.windowTitle(for: entry), + snapshot: snapshot, + now: now)] + } + + private static func sparkWindows( + from entry: CodexUsageResponse.AdditionalRateLimit, + usedIDs: inout Set, + now: Date) -> [NamedRateWindow] + { + let candidates: [(CodexUsageResponse.WindowSnapshot?, SparkWindowKind)] = [ + (entry.rateLimit?.primaryWindow, .fiveHour), + (entry.rateLimit?.secondaryWindow, .weekly), + ] + + return candidates.compactMap { snapshot, fallbackKind in + guard let snapshot else { return nil } + let kind = self.sparkWindowKind(for: snapshot, fallback: fallbackKind) + guard usedIDs.insert(kind.id).inserted else { return nil } + return self.namedWindow(id: kind.id, title: kind.title, snapshot: snapshot, now: now) + } + } + + private static func namedWindow( + id: String, + title: String, + snapshot: CodexUsageResponse.WindowSnapshot, + now: Date) -> NamedRateWindow + { + let resetDate: Date? = snapshot.resetAt > 0 + ? Date(timeIntervalSince1970: TimeInterval(snapshot.resetAt)) + : nil + let window = RateWindow( + usedPercent: Double(snapshot.usedPercent), + windowMinutes: snapshot.limitWindowSeconds > 0 ? snapshot.limitWindowSeconds / 60 : nil, + resetsAt: resetDate, + resetDescription: resetDate.map { UsageFormatter.resetDescription(from: $0, now: now) }) + return NamedRateWindow(id: id, title: title, window: window) + } + + private enum SparkWindowKind { + case fiveHour + case weekly + + var id: String { + switch self { + case .fiveHour: CodexAdditionalRateLimitMapper.sparkWindowID + case .weekly: CodexAdditionalRateLimitMapper.sparkWeeklyWindowID + } + } + + var title: String { + switch self { + case .fiveHour: CodexAdditionalRateLimitMapper.sparkWindowTitle + case .weekly: CodexAdditionalRateLimitMapper.sparkWeeklyWindowTitle + } + } + } + + private static func sparkWindowKind( + for snapshot: CodexUsageResponse.WindowSnapshot, + fallback: SparkWindowKind) -> SparkWindowKind + { + let minutes = snapshot.limitWindowSeconds > 0 ? snapshot.limitWindowSeconds / 60 : 0 + if minutes > 0, minutes <= 6 * 60 { return .fiveHour } + if minutes >= 6 * 24 * 60 { return .weekly } + return fallback + } + + private static func windowID(for entry: CodexUsageResponse.AdditionalRateLimit) -> String? { + guard let source = self.firstNonEmpty(entry.meteredFeature, entry.limitName) else { return nil } + let slug = self.slug(source) + return slug.isEmpty ? nil : "codex-\(slug)" + } + + private static func windowTitle(for entry: CodexUsageResponse.AdditionalRateLimit) -> String { + self.firstNonEmpty(entry.limitName, entry.meteredFeature) ?? "Codex extra limit" + } + + private static func isSpark(_ entry: CodexUsageResponse.AdditionalRateLimit) -> Bool { + [entry.limitName, entry.meteredFeature] + .compactMap { $0?.lowercased() } + .contains { $0.contains("spark") } + } + + private static func firstNonEmpty(_ values: String?...) -> String? { + for value in values { + let trimmed = value?.trimmingCharacters(in: .whitespacesAndNewlines) + if let trimmed, !trimmed.isEmpty { return trimmed } + } + return nil + } + + private static func slug(_ value: String) -> String { + var result = "" + var lastWasDash = false + for scalar in value.lowercased().unicodeScalars { + if CharacterSet.alphanumerics.contains(scalar) { + result.unicodeScalars.append(scalar) + lastWasDash = false + } else if !lastWasDash { + result.append("-") + lastWasDash = true + } + } + return result.trimmingCharacters(in: CharacterSet(charactersIn: "-")) + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexAuthenticatedHTTPTransport.swift b/Sources/CodexBarCore/Providers/Codex/CodexAuthenticatedHTTPTransport.swift new file mode 100644 index 0000000..12c17a2 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexAuthenticatedHTTPTransport.swift @@ -0,0 +1,33 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +enum CodexAuthenticatedHTTPTransport { + static let shared = Self.makeClient() + + @TaskLocal static var overrideForTesting: (any ProviderHTTPTransport)? + + static var current: any ProviderHTTPTransport { + if let override = self.overrideForTesting { + return override + } + return self.shared + } + + static func makeConfiguration() -> URLSessionConfiguration { + let configuration = URLSessionConfiguration.ephemeral + configuration.urlCache = nil + configuration.requestCachePolicy = .reloadIgnoringLocalCacheData + configuration.httpCookieStorage = nil + configuration.httpShouldSetCookies = false + configuration.urlCredentialStorage = nil + return configuration + } + + static func makeClient(configuration: URLSessionConfiguration? = nil) -> ProviderHTTPClient { + let configuration = configuration ?? self.makeConfiguration() + let session = ProviderHTTPClient.redirectGuardedSession(configuration: configuration) + return ProviderHTTPClient(session: session) + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexCLIDashboardAuthorityContext.swift b/Sources/CodexBarCore/Providers/Codex/CodexCLIDashboardAuthorityContext.swift new file mode 100644 index 0000000..ce34806 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexCLIDashboardAuthorityContext.swift @@ -0,0 +1,62 @@ +import Foundation + +public enum CodexCLIDashboardAuthorityContext { + public static func makeLiveWebInput( + dashboard: OpenAIDashboardSnapshot, + context: ProviderFetchContext, + routingTargetEmail: String?) -> CodexDashboardAuthorityInput + { + let auth = context.fetcher.loadAuthBackedCodexAccount() + return CodexDashboardAuthorityInput( + sourceKind: .liveWeb, + proof: CodexDashboardOwnershipProofContext( + currentIdentity: auth.identity, + expectedScopedEmail: auth.email, + trustedCurrentUsageEmail: nil, + dashboardSignedInEmail: dashboard.signedInEmail, + knownOwners: context.settings?.codex?.dashboardAuthorityKnownOwners ?? []), + routing: CodexDashboardRoutingHints( + targetEmail: CodexIdentityResolver.normalizeEmail(routingTargetEmail), + lastKnownDashboardRoutingEmail: nil)) + } + + public static func makeCachedDashboardInput( + dashboard: OpenAIDashboardSnapshot, + cachedAccountEmail: String, + usage: UsageSnapshot, + sourceLabel: String, + context: ProviderFetchContext) -> CodexDashboardAuthorityInput + { + let auth = context.fetcher.loadAuthBackedCodexAccount() + let trustedCurrentUsageEmail = Self.shouldTrustUsageEmail(sourceLabel: sourceLabel) + ? usage.accountEmail(for: .codex) + : nil + return CodexDashboardAuthorityInput( + sourceKind: .cachedDashboard, + proof: CodexDashboardOwnershipProofContext( + currentIdentity: auth.identity, + expectedScopedEmail: auth.email, + trustedCurrentUsageEmail: trustedCurrentUsageEmail, + dashboardSignedInEmail: dashboard.signedInEmail, + knownOwners: context.settings?.codex?.dashboardAuthorityKnownOwners ?? []), + routing: CodexDashboardRoutingHints( + targetEmail: auth.email, + lastKnownDashboardRoutingEmail: cachedAccountEmail)) + } + + public static func attachmentEmail(from input: CodexDashboardAuthorityInput) -> String? { + CodexIdentityResolver.normalizeEmail( + input.proof.expectedScopedEmail ?? + input.proof.trustedCurrentUsageEmail ?? + input.proof.dashboardSignedInEmail) + } + + public static func shouldTrustUsageEmail(sourceLabel: String) -> Bool { + switch sourceLabel.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() { + case "codex-cli", "oauth": + true + default: + false + } + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexCLILaunchGate.swift b/Sources/CodexBarCore/Providers/Codex/CodexCLILaunchGate.swift new file mode 100644 index 0000000..78ffa8a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexCLILaunchGate.swift @@ -0,0 +1,70 @@ +import Foundation + +final class CodexCLILaunchGate: @unchecked Sendable { + static let shared = CodexCLILaunchGate() + static let cooldown: TimeInterval = 30 * 60 + + private struct Entry { + let message: String + let expiresAt: Date + } + + private let lock = NSLock() + private var entries: [String: Entry] = [:] + + func backgroundSkipMessage( + binary: String, + now: Date = Date(), + interaction: ProviderInteraction = ProviderInteractionContext.current) -> String? + { + guard interaction == .background else { return nil } + + self.lock.lock() + defer { self.lock.unlock() } + + guard let entry = self.entries[binary] else { return nil } + guard entry.expiresAt > now else { + self.entries.removeValue(forKey: binary) + return nil + } + return entry.message + } + + @discardableResult + func recordLaunchFailure(binary: String, message: String, now: Date = Date()) -> String? { + guard Self.shouldThrottleLaunchFailure(message) else { return nil } + let throttled = Self.throttledMessage(binary: binary, originalMessage: message) + let entry = Entry( + message: throttled, + expiresAt: now.addingTimeInterval(Self.cooldown)) + + self.lock.lock() + self.entries[binary] = entry + self.lock.unlock() + + return throttled + } + + func resetForTesting() { + self.lock.lock() + self.entries.removeAll() + self.lock.unlock() + } + + static func shouldThrottleLaunchFailure(_ message: String) -> Bool { + let lower = message.lowercased() + if lower.contains("openpty") || + lower.contains("write to pty") || + lower.contains("app shutdown") + { + return false + } + return true + } + + private static func throttledMessage(binary: String, originalMessage: String) -> String { + "Codex CLI launch failed; background refresh is paused for 30 minutes. " + + "Reinstall or unblock `\(binary)` in macOS security settings, then refresh manually. " + + "Last error: \(originalMessage)" + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift b/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift new file mode 100644 index 0000000..2c88db3 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift @@ -0,0 +1,427 @@ +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif +import Foundation + +actor CodexCLISession { + static let shared = CodexCLISession() + + enum SessionError: LocalizedError { + case launchFailed(String) + case timedOut + case processExited + + var errorDescription: String? { + switch self { + case let .launchFailed(msg): "Failed to launch Codex CLI session: \(msg)" + case .timedOut: "Codex CLI session timed out." + case .processExited: "Codex CLI session exited." + } + } + } + + private var process: Process? + private var primaryFD: Int32 = -1 + private var primaryHandle: FileHandle? + private var secondaryHandle: FileHandle? + private var processGroup: pid_t? + private var binaryPath: String? + private var startedAt: Date? + private var ptyRows: UInt16 = 0 + private var ptyCols: UInt16 = 0 + private var sessionEnvironment: [String: String]? + private var sessionArguments: [String] = [] + private var sessionWorkingDirectory: URL? + + struct CaptureOptions { + let timeout: TimeInterval + let rows: UInt16 + let cols: UInt16 + let environment: [String: String] + let extraArgs: [String] + let workingDirectory: URL? + } + + private struct RollingBuffer { + private let maxNeedle: Int + private var tail = Data() + + init(maxNeedle: Int) { + self.maxNeedle = max(0, maxNeedle) + } + + mutating func append(_ data: Data) -> Data { + guard !data.isEmpty else { return Data() } + var combined = Data() + combined.reserveCapacity(self.tail.count + data.count) + combined.append(self.tail) + combined.append(data) + if self.maxNeedle > 1 { + if combined.count >= self.maxNeedle - 1 { + self.tail = combined.suffix(self.maxNeedle - 1) + } else { + self.tail = combined + } + } else { + self.tail.removeAll(keepingCapacity: true) + } + return combined + } + + mutating func reset() { + self.tail.removeAll(keepingCapacity: true) + } + } + + static func lowercasedASCII(_ data: Data) -> Data { + guard !data.isEmpty else { return data } + var out = Data(count: data.count) + out.withUnsafeMutableBytes { dest in + data.withUnsafeBytes { source in + let src = source.bindMemory(to: UInt8.self) + let dst = dest.bindMemory(to: UInt8.self) + for idx in 0..= 65, byte <= 90 { byte += 32 } + dst[idx] = byte + } + } + } + return out + } + + // swiftlint:disable cyclomatic_complexity + func captureStatus( + binary: String, + options: CaptureOptions) async throws -> String + { + try self.ensureStarted(binary: binary, options: options) + if let startedAt { + let sinceStart = Date().timeIntervalSince(startedAt) + if sinceStart < 0.4 { + let delay = UInt64((0.4 - sinceStart) * 1_000_000_000) + try await Task.sleep(nanoseconds: delay) + } + } + self.drainOutput() + + let script = "/status" + let cursorQuery = Data([0x1B, 0x5B, 0x36, 0x6E]) + let statusMarkers = [ + "Credits:", + "5h limit", + "5-hour limit", + "Weekly limit", + ].map { Data($0.utf8) } + let updateNeedles = ["Update available!", "Run bun install -g @openai/codex", "0.60.1 ->"] + let updateNeedlesLower = updateNeedles.map { Data($0.lowercased().utf8) } + let statusNeedleLengths = statusMarkers.map(\.count) + let updateNeedleLengths = updateNeedlesLower.map(\.count) + let statusMaxNeedle = ([cursorQuery.count] + statusNeedleLengths).max() ?? cursorQuery.count + let updateMaxNeedle = updateNeedleLengths.max() ?? 0 + var statusScanBuffer = RollingBuffer(maxNeedle: statusMaxNeedle) + var updateScanBuffer = RollingBuffer(maxNeedle: updateMaxNeedle) + + var buffer = Data() + let deadline = Date().addingTimeInterval(options.timeout) + var nextCursorCheckAt = Date(timeIntervalSince1970: 0) + + var skippedCodexUpdate = false + var sentScript = false + var updateSkipAttempts = 0 + var lastEnter = Date(timeIntervalSince1970: 0) + var scriptSentAt: Date? + var resendStatusRetries = 0 + var enterRetries = 0 + var sawCodexStatus = false + var sawCodexUpdatePrompt = false + + while Date() < deadline { + let newData = self.readChunk() + if !newData.isEmpty { + buffer.append(newData) + } + let scanData = statusScanBuffer.append(newData) + if Date() >= nextCursorCheckAt, + !scanData.isEmpty, + scanData.range(of: cursorQuery) != nil + { + try? self.send("\u{1b}[1;1R") + nextCursorCheckAt = Date().addingTimeInterval(1.0) + } + if !scanData.isEmpty, !sawCodexStatus { + if statusMarkers.contains(where: { scanData.range(of: $0) != nil }) { + sawCodexStatus = true + } + } + + if !skippedCodexUpdate, !sawCodexUpdatePrompt, !newData.isEmpty { + let lowerData = Self.lowercasedASCII(newData) + let lowerScan = updateScanBuffer.append(lowerData) + if updateNeedlesLower.contains(where: { lowerScan.range(of: $0) != nil }) { + sawCodexUpdatePrompt = true + } + } + + if !skippedCodexUpdate, sawCodexUpdatePrompt { + try? self.send("\u{1b}[B") + try await Task.sleep(nanoseconds: 120_000_000) + try? self.send("\r") + try await Task.sleep(nanoseconds: 150_000_000) + try? self.send("\r") + try? self.send(script) + try? self.send("\r") + updateSkipAttempts += 1 + if updateSkipAttempts >= 1 { + skippedCodexUpdate = true + sentScript = false + scriptSentAt = nil + buffer.removeAll() + statusScanBuffer.reset() + updateScanBuffer.reset() + sawCodexStatus = false + } + try await Task.sleep(nanoseconds: 300_000_000) + } + + if !sentScript, !sawCodexUpdatePrompt || skippedCodexUpdate { + try? self.send(script) + try? self.send("\r") + sentScript = true + scriptSentAt = Date() + lastEnter = Date() + try await Task.sleep(nanoseconds: 200_000_000) + continue + } + if sentScript, !sawCodexStatus { + if Date().timeIntervalSince(lastEnter) >= 1.2, enterRetries < 6 { + try? self.send("\r") + enterRetries += 1 + lastEnter = Date() + try await Task.sleep(nanoseconds: 120_000_000) + continue + } + if let sentAt = scriptSentAt, + Date().timeIntervalSince(sentAt) >= 3.0, + resendStatusRetries < 2 + { + try? self.send(script) + try? self.send("\r") + resendStatusRetries += 1 + buffer.removeAll() + statusScanBuffer.reset() + updateScanBuffer.reset() + sawCodexStatus = false + scriptSentAt = Date() + lastEnter = Date() + try await Task.sleep(nanoseconds: 220_000_000) + continue + } + } + if sawCodexStatus { break } + if let proc = self.process, !proc.isRunning { + throw SessionError.processExited + } + try await Task.sleep(nanoseconds: 120_000_000) + } + + if sawCodexStatus { + let settleDeadline = Date().addingTimeInterval(2.0) + while Date() < settleDeadline { + let newData = self.readChunk() + if !newData.isEmpty { + buffer.append(newData) + } + let scanData = statusScanBuffer.append(newData) + if Date() >= nextCursorCheckAt, + !scanData.isEmpty, + scanData.range(of: cursorQuery) != nil + { + try? self.send("\u{1b}[1;1R") + nextCursorCheckAt = Date().addingTimeInterval(1.0) + } + try await Task.sleep(nanoseconds: 100_000_000) + } + } + + guard !buffer.isEmpty, let text = String(data: buffer, encoding: .utf8) else { + throw SessionError.timedOut + } + return text + } + + // swiftlint:enable cyclomatic_complexity + + func reset() { + self.cleanup() + } + + private func ensureStarted( + binary: String, + options: CaptureOptions) throws + { + if let proc = self.process, + proc.isRunning, + self.binaryPath == binary, + self.ptyRows == options.rows, + self.ptyCols == options.cols, + self.sessionEnvironment == options.environment, + self.sessionArguments == options.extraArgs, + self.sessionWorkingDirectory == options.workingDirectory + { + return + } + self.cleanup() + + var primaryFD: Int32 = -1 + var secondaryFD: Int32 = -1 + var win = winsize(ws_row: options.rows, ws_col: options.cols, ws_xpixel: 0, ws_ypixel: 0) + guard openpty(&primaryFD, &secondaryFD, nil, nil, &win) == 0 else { + throw SessionError.launchFailed("openpty failed") + } + _ = fcntl(primaryFD, F_SETFL, O_NONBLOCK) + + let primaryHandle = FileHandle(fileDescriptor: primaryFD, closeOnDealloc: true) + let secondaryHandle = FileHandle(fileDescriptor: secondaryFD, closeOnDealloc: true) + + let proc = Process() + let resolvedURL = URL(fileURLWithPath: binary) + proc.executableURL = resolvedURL + proc.arguments = options.extraArgs + proc.standardInput = secondaryHandle + proc.standardOutput = secondaryHandle + proc.standardError = secondaryHandle + proc.currentDirectoryURL = options.workingDirectory + + let env = TTYCommandRunner.enrichedEnvironment( + baseEnv: options.environment, + home: options.environment["HOME"] ?? NSHomeDirectory()) + proc.environment = env + + guard TTYCommandRunner.beginActiveProcessLaunchForAppShutdown() else { + try? primaryHandle.close() + try? secondaryHandle.close() + throw SessionError.launchFailed("App shutdown in progress") + } + defer { TTYCommandRunner.endActiveProcessLaunchForAppShutdown() } + + do { + try proc.run() + } catch { + try? primaryHandle.close() + try? secondaryHandle.close() + throw SessionError.launchFailed(error.localizedDescription) + } + + let pid = proc.processIdentifier + guard TTYCommandRunner.registerActiveProcessForAppShutdown( + pid: pid, + binary: resolvedURL.lastPathComponent) + else { + proc.terminate() + kill(pid, SIGKILL) + try? primaryHandle.close() + try? secondaryHandle.close() + throw SessionError.launchFailed("App shutdown in progress") + } + + var processGroup: pid_t? + if setpgid(pid, pid) == 0 { + processGroup = pid + TTYCommandRunner.updateActiveProcessGroupForAppShutdown(pid: pid, processGroup: processGroup) + } + + self.process = proc + self.primaryFD = primaryFD + self.primaryHandle = primaryHandle + self.secondaryHandle = secondaryHandle + self.processGroup = processGroup + self.binaryPath = binary + self.startedAt = Date() + self.ptyRows = options.rows + self.ptyCols = options.cols + self.sessionEnvironment = options.environment + self.sessionArguments = options.extraArgs + self.sessionWorkingDirectory = options.workingDirectory + } + + private func cleanup() { + if let proc = self.process, proc.isRunning, let handle = self.primaryHandle { + try? handle.write(contentsOf: Data("/exit\n".utf8)) + } + try? self.primaryHandle?.close() + try? self.secondaryHandle?.close() + + let descendants = self.process.map { TTYProcessTreeTerminator.descendantPIDs(of: $0.processIdentifier) } ?? [] + if let proc = self.process, proc.isRunning { + proc.terminate() + } + if let proc = self.process { + TTYProcessTreeTerminator.terminateProcessTree( + rootPID: proc.processIdentifier, + processGroup: self.processGroup, + signal: SIGTERM, + knownDescendants: descendants) + } + let waitDeadline = Date().addingTimeInterval(1.0) + if let proc = self.process { + while proc.isRunning, Date() < waitDeadline { + usleep(100_000) + } + if proc.isRunning { + TTYProcessTreeTerminator.terminateProcessTree( + rootPID: proc.processIdentifier, + processGroup: self.processGroup, + signal: SIGKILL, + knownDescendants: descendants) + } else { + for pid in descendants where pid > 0 { + kill(pid, SIGKILL) + } + } + TTYCommandRunner.unregisterActiveProcessForAppShutdown(pid: proc.processIdentifier) + } + + self.process = nil + self.primaryHandle = nil + self.secondaryHandle = nil + self.primaryFD = -1 + self.processGroup = nil + self.binaryPath = nil + self.startedAt = nil + self.ptyRows = 0 + self.ptyCols = 0 + self.sessionEnvironment = nil + self.sessionArguments = [] + self.sessionWorkingDirectory = nil + } + + private func readChunk() -> Data { + guard self.primaryFD >= 0 else { return Data() } + var appended = Data() + while true { + var tmp = [UInt8](repeating: 0, count: 8192) + let n = read(self.primaryFD, &tmp, tmp.count) + if n > 0 { + appended.append(contentsOf: tmp.prefix(n)) + continue + } + break + } + return appended + } + + private func drainOutput() { + _ = self.readChunk() + } + + private func send(_ text: String) throws { + guard let data = text.data(using: .utf8) else { return } + guard let handle = self.primaryHandle else { throw SessionError.processExited } + try handle.write(contentsOf: data) + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexDashboardAuthority.swift b/Sources/CodexBarCore/Providers/Codex/CodexDashboardAuthority.swift new file mode 100644 index 0000000..8d40ccc --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexDashboardAuthority.swift @@ -0,0 +1,332 @@ +import Foundation + +public enum CodexDashboardSourceKind: String, Codable, Sendable { + case liveWeb + case cachedDashboard +} + +public enum CodexDashboardDisposition: String, Codable, Sendable { + case attach + case displayOnly + case failClosed +} + +public enum CodexDashboardAllowedEffect: String, Codable, CaseIterable, Hashable, Sendable { + case usageBackfill + case creditsAttachment + case refreshGuardSeed + case historicalBackfill + case cachedDashboardReuse +} + +public enum CodexDashboardCleanup: String, Codable, CaseIterable, Hashable, Sendable { + case dashboardSnapshot + case dashboardDerivedUsage + case dashboardDerivedCredits + case dashboardRefreshGuardSeed + case dashboardCache +} + +public struct CodexDashboardKnownOwnerCandidate: Equatable, Hashable, Sendable { + public let identity: CodexIdentity + public let normalizedEmail: String? + public let sourceIsolationIdentifier: String? + + public init( + identity: CodexIdentity, + normalizedEmail: String?, + sourceIsolationIdentifier: String? = nil) + { + self.identity = identity + self.normalizedEmail = normalizedEmail + self.sourceIsolationIdentifier = sourceIsolationIdentifier + } + + public func hash(into hasher: inout Hasher) { + switch self.identity { + case let .providerAccount(id): + hasher.combine("providerAccount") + hasher.combine(id) + case let .emailOnly(normalizedEmail): + hasher.combine("emailOnly") + hasher.combine(normalizedEmail) + case .unresolved: + hasher.combine("unresolved") + } + hasher.combine(self.normalizedEmail) + hasher.combine(self.sourceIsolationIdentifier) + } +} + +public struct CodexDashboardOwnershipProofContext: Equatable, Sendable { + public let currentIdentity: CodexIdentity + public let expectedScopedEmail: String? + public let trustedCurrentUsageEmail: String? + public let dashboardSignedInEmail: String? + public let knownOwners: [CodexDashboardKnownOwnerCandidate] + + public init( + currentIdentity: CodexIdentity, + expectedScopedEmail: String?, + trustedCurrentUsageEmail: String?, + dashboardSignedInEmail: String?, + knownOwners: [CodexDashboardKnownOwnerCandidate]) + { + self.currentIdentity = currentIdentity + self.expectedScopedEmail = expectedScopedEmail + self.trustedCurrentUsageEmail = trustedCurrentUsageEmail + self.dashboardSignedInEmail = dashboardSignedInEmail + self.knownOwners = knownOwners + } +} + +public struct CodexDashboardRoutingHints: Equatable, Sendable { + public let targetEmail: String? + public let lastKnownDashboardRoutingEmail: String? + + public init(targetEmail: String?, lastKnownDashboardRoutingEmail: String?) { + self.targetEmail = targetEmail + self.lastKnownDashboardRoutingEmail = lastKnownDashboardRoutingEmail + } +} + +public struct CodexDashboardAuthorityInput: Equatable, Sendable { + public let sourceKind: CodexDashboardSourceKind + public let proof: CodexDashboardOwnershipProofContext + public let routing: CodexDashboardRoutingHints + + public init( + sourceKind: CodexDashboardSourceKind, + proof: CodexDashboardOwnershipProofContext, + routing: CodexDashboardRoutingHints) + { + self.sourceKind = sourceKind + self.proof = proof + self.routing = routing + } +} + +public enum CodexDashboardDecisionReason: Equatable, Sendable { + case exactProviderAccountMatch + case trustedEmailMatchNoCompetingOwner + case trustedContinuityNoCompetingOwner + case wrongEmail(expected: String?, actual: String?) + case sameEmailAmbiguity(email: String) + case unresolvedWithoutTrustedEvidence + case providerAccountMissingScopedEmail + case providerAccountLacksExactOwnershipProof + case missingDashboardSignedInEmail +} + +public struct CodexDashboardAuthorityDecision: Equatable, Sendable { + public let disposition: CodexDashboardDisposition + public let reason: CodexDashboardDecisionReason + public let allowedEffects: Set + public let cleanup: Set + + public init( + disposition: CodexDashboardDisposition, + reason: CodexDashboardDecisionReason, + allowedEffects: Set, + cleanup: Set) + { + self.disposition = disposition + self.reason = reason + self.allowedEffects = allowedEffects + self.cleanup = cleanup + } +} + +public enum CodexDashboardPolicyError: LocalizedError, Equatable, Sendable { + case displayOnly(CodexDashboardAuthorityDecision) + + public var errorDescription: String? { + switch self { + case .displayOnly: + "Codex dashboard may be displayed, but it cannot be attached to the active account." + } + } +} + +public enum CodexDashboardAuthority { + /// Evaluates whether a Codex dashboard snapshot may attach to the active account. + /// + /// App callers may keep `.displayOnly` dashboard data visible, but must not attach usage, + /// credits, refresh guards, or historical backfill. + /// + /// CLI callers should surface `.displayOnly` as `CodexDashboardPolicyError.displayOnly` + /// instead of treating it as a generic failure. + /// + /// Cached dashboard callers may restore data only when the decision includes + /// `.cachedDashboardReuse` in `allowedEffects`. + public static func evaluate(_ input: CodexDashboardAuthorityInput) -> CodexDashboardAuthorityDecision { + let proof = input.proof + let currentIdentity = Self.normalizeIdentity(proof.currentIdentity) + let expectedScopedEmail = CodexIdentityResolver.normalizeEmail(proof.expectedScopedEmail) + let trustedCurrentUsageEmail = CodexIdentityResolver.normalizeEmail(proof.trustedCurrentUsageEmail) + let dashboardSignedInEmail = CodexIdentityResolver.normalizeEmail(proof.dashboardSignedInEmail) + let knownOwners = Self.normalizeKnownOwners(proof.knownOwners) + + // Routing hints are intentionally excluded from ownership proof. They may help fetch or route + // dashboard requests, but they must never influence attach/display/fail-closed policy. + + guard let dashboardSignedInEmail else { + return Self.makeDecision( + disposition: .failClosed, + reason: .missingDashboardSignedInEmail, + sourceKind: input.sourceKind) + } + + if let expectedScopedEmail, dashboardSignedInEmail != expectedScopedEmail { + return Self.makeDecision( + disposition: .failClosed, + reason: .wrongEmail(expected: expectedScopedEmail, actual: dashboardSignedInEmail), + sourceKind: input.sourceKind) + } + + switch currentIdentity { + case let .providerAccount(id): + guard expectedScopedEmail != nil else { + return Self.makeDecision( + disposition: .failClosed, + reason: .providerAccountMissingScopedEmail, + sourceKind: input.sourceKind) + } + if Self.knownOwnerCount(for: dashboardSignedInEmail, in: knownOwners) > 1 { + return Self.makeDecision( + disposition: .displayOnly, + reason: .sameEmailAmbiguity(email: dashboardSignedInEmail), + sourceKind: input.sourceKind) + } + let exactMatch = knownOwners.contains { candidate in + candidate.identity == .providerAccount(id: id) && candidate.normalizedEmail == dashboardSignedInEmail + } + if exactMatch { + return Self.makeDecision( + disposition: .attach, + reason: .exactProviderAccountMatch, + sourceKind: input.sourceKind) + } + return Self.makeDecision( + disposition: .failClosed, + reason: .providerAccountLacksExactOwnershipProof, + sourceKind: input.sourceKind) + + case let .emailOnly(normalizedEmail): + guard dashboardSignedInEmail == normalizedEmail else { + return Self.makeDecision( + disposition: .failClosed, + reason: .wrongEmail(expected: normalizedEmail, actual: dashboardSignedInEmail), + sourceKind: input.sourceKind) + } + if Self.knownOwnerCount(for: normalizedEmail, in: knownOwners) > 1 { + return Self.makeDecision( + disposition: .displayOnly, + reason: .sameEmailAmbiguity(email: normalizedEmail), + sourceKind: input.sourceKind) + } + return Self.makeDecision( + disposition: .attach, + reason: .trustedEmailMatchNoCompetingOwner, + sourceKind: input.sourceKind) + + case .unresolved: + guard let trustedCurrentUsageEmail else { + return Self.makeDecision( + disposition: .failClosed, + reason: .unresolvedWithoutTrustedEvidence, + sourceKind: input.sourceKind) + } + guard dashboardSignedInEmail == trustedCurrentUsageEmail else { + return Self.makeDecision( + disposition: .failClosed, + reason: .wrongEmail(expected: trustedCurrentUsageEmail, actual: dashboardSignedInEmail), + sourceKind: input.sourceKind) + } + if Self.knownOwnerCount(for: trustedCurrentUsageEmail, in: knownOwners) > 1 { + return Self.makeDecision( + disposition: .displayOnly, + reason: .sameEmailAmbiguity(email: trustedCurrentUsageEmail), + sourceKind: input.sourceKind) + } + return Self.makeDecision( + disposition: .attach, + reason: .trustedContinuityNoCompetingOwner, + sourceKind: input.sourceKind) + } + } + + private static func normalizeIdentity(_ identity: CodexIdentity) -> CodexIdentity { + switch identity { + case let .providerAccount(id): + if let normalizedID = CodexIdentityResolver.normalizeAccountID(id) { + return .providerAccount(id: normalizedID) + } + return .unresolved + case let .emailOnly(normalizedEmail): + if let normalizedEmail = CodexIdentityResolver.normalizeEmail(normalizedEmail) { + return .emailOnly(normalizedEmail: normalizedEmail) + } + return .unresolved + case .unresolved: + return .unresolved + } + } + + private static func normalizeKnownOwners( + _ candidates: [CodexDashboardKnownOwnerCandidate]) + -> Set + { + Set(candidates.map { candidate in + let identity = self.normalizeIdentity(candidate.identity) + let sourceIsolationIdentifier: String? = switch identity { + case .providerAccount: + nil + case .emailOnly, .unresolved: + candidate.sourceIsolationIdentifier + } + return CodexDashboardKnownOwnerCandidate( + identity: identity, + normalizedEmail: CodexIdentityResolver.normalizeEmail(candidate.normalizedEmail), + sourceIsolationIdentifier: sourceIsolationIdentifier) + }) + } + + private static func knownOwnerCount( + for email: String, + in candidates: Set) -> Int + { + candidates.count { $0.normalizedEmail == email } + } + + private static func makeDecision( + disposition: CodexDashboardDisposition, + reason: CodexDashboardDecisionReason, + sourceKind: CodexDashboardSourceKind) -> CodexDashboardAuthorityDecision + { + CodexDashboardAuthorityDecision( + disposition: disposition, + reason: reason, + allowedEffects: self.allowedEffects(disposition: disposition, sourceKind: sourceKind), + cleanup: disposition == .attach ? [] : Set(CodexDashboardCleanup.allCases)) + } + + private static func allowedEffects( + disposition: CodexDashboardDisposition, + sourceKind: CodexDashboardSourceKind) -> Set + { + guard disposition == .attach else { return [] } + + switch sourceKind { + case .liveWeb: + return [ + .usageBackfill, + .creditsAttachment, + .refreshGuardSeed, + .historicalBackfill, + ] + case .cachedDashboard: + return [.cachedDashboardReuse] + } + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexIdentity.swift b/Sources/CodexBarCore/Providers/Codex/CodexIdentity.swift new file mode 100644 index 0000000..f83cb19 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexIdentity.swift @@ -0,0 +1,47 @@ +import Foundation + +public enum CodexIdentity: Codable, Equatable, Sendable { + case providerAccount(id: String) + // Normal OAuth auth should resolve to a provider account ID. Email-only identity is kept as a + // migration/hardening fallback for partial auth payloads, not as the primary steady-state path. + case emailOnly(normalizedEmail: String) + case unresolved +} + +public enum CodexIdentityResolver { + public static func resolve(accountId: String?, email: String?) -> CodexIdentity { + if let accountId = normalizeAccountID(accountId) { + return .providerAccount(id: accountId) + } + if let email = Self.normalizeEmail(email) { + return .emailOnly(normalizedEmail: email) + } + return .unresolved + } + + public static func normalizeEmail(_ email: String?) -> String? { + guard let email = email?.trimmingCharacters(in: .whitespacesAndNewlines), !email.isEmpty else { + return nil + } + return email.lowercased() + } + + public static func normalizeAccountID(_ accountId: String?) -> String? { + guard let accountId = accountId?.trimmingCharacters(in: .whitespacesAndNewlines), !accountId.isEmpty else { + return nil + } + return accountId + } +} + +public struct CodexAuthBackedAccount: Equatable, Sendable { + public let identity: CodexIdentity + public let email: String? + public let plan: String? + + public init(identity: CodexIdentity, email: String?, plan: String?) { + self.identity = identity + self.email = email + self.plan = plan + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthCredentials.swift b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthCredentials.swift new file mode 100644 index 0000000..7c180da --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthCredentials.swift @@ -0,0 +1,285 @@ +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif +import Foundation + +public struct CodexOAuthCredentials: Sendable { + public let accessToken: String + public let refreshToken: String + public let idToken: String? + public let accountId: String? + public let lastRefresh: Date? + + public init( + accessToken: String, + refreshToken: String, + idToken: String?, + accountId: String?, + lastRefresh: Date?) + { + self.accessToken = accessToken + self.refreshToken = refreshToken + self.idToken = idToken + self.accountId = accountId + self.lastRefresh = lastRefresh + } + + public var needsRefresh: Bool { + guard let lastRefresh else { return true } + let eightDays: TimeInterval = 8 * 24 * 60 * 60 + return Date().timeIntervalSince(lastRefresh) > eightDays + } +} + +public enum CodexOAuthCredentialsError: LocalizedError, Sendable { + case notFound + case decodeFailed(String) + case missingTokens + + public var errorDescription: String? { + switch self { + case .notFound: + "Codex auth.json not found. Run `codex` to log in." + case let .decodeFailed(message): + "Failed to decode Codex credentials: \(message)" + case .missingTokens: + "Codex auth.json exists but contains no tokens." + } + } +} + +public enum CodexOAuthCredentialsStore { + private static func authFilePath( + env: [String: String] = ProcessInfo.processInfo.environment, + fileManager: FileManager = .default) -> URL + { + CodexHomeScope + .ambientHomeURL(env: env, fileManager: fileManager) + .appendingPathComponent("auth.json") + } + + public static func load(env: [String: String] = ProcessInfo.processInfo + .environment) throws -> CodexOAuthCredentials + { + let url = self.authFilePath(env: env) + guard FileManager.default.fileExists(atPath: url.path) else { + throw CodexOAuthCredentialsError.notFound + } + + let data = try Data(contentsOf: url) + return try self.parse(data: data) + } + + public static func loadOAuthTokens(env: [String: String] = ProcessInfo.processInfo + .environment) throws -> CodexOAuthCredentials + { + let url = self.authFilePath(env: env) + guard FileManager.default.fileExists(atPath: url.path) else { + throw CodexOAuthCredentialsError.notFound + } + + let data = try Data(contentsOf: url) + guard let credentials = try self.tokenCredentials(data: data) else { + throw CodexOAuthCredentialsError.missingTokens + } + return credentials + } + + public static func parse(data: Data) throws -> CodexOAuthCredentials { + guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw CodexOAuthCredentialsError.decodeFailed("Invalid JSON") + } + + if let apiKeyCredentials = Self.apiKeyCredentials(in: json) { + return apiKeyCredentials + } + + if let tokenCredentials = Self.tokenCredentials(in: json) { + return tokenCredentials + } + + throw CodexOAuthCredentialsError.missingTokens + } + + private static func tokenCredentials(data: Data) throws -> CodexOAuthCredentials? { + guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw CodexOAuthCredentialsError.decodeFailed("Invalid JSON") + } + return self.tokenCredentials(in: json) + } + + private static func tokenCredentials(in json: [String: Any]) -> CodexOAuthCredentials? { + guard let tokens = json["tokens"] as? [String: Any], + let accessToken = stringValue( + in: tokens, + snakeCaseKey: "access_token", + camelCaseKey: "accessToken"), + let refreshToken = stringValue( + in: tokens, + snakeCaseKey: "refresh_token", + camelCaseKey: "refreshToken"), + !accessToken.isEmpty + else { + return nil + } + + let idToken = Self.stringValue(in: tokens, snakeCaseKey: "id_token", camelCaseKey: "idToken") + let accountId = Self.stringValue(in: tokens, snakeCaseKey: "account_id", camelCaseKey: "accountId") + let lastRefresh = Self.parseLastRefresh(from: json["last_refresh"]) + + return CodexOAuthCredentials( + accessToken: accessToken, + refreshToken: refreshToken, + idToken: idToken, + accountId: accountId, + lastRefresh: lastRefresh) + } + + private static func apiKeyCredentials(in json: [String: Any]) -> CodexOAuthCredentials? { + guard let apiKey = json["OPENAI_API_KEY"] as? String, + !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + else { + return nil + } + return CodexOAuthCredentials( + accessToken: apiKey, + refreshToken: "", + idToken: nil, + accountId: nil, + lastRefresh: nil) + } + + public static func save( + _ credentials: CodexOAuthCredentials, + env: [String: String] = ProcessInfo.processInfo.environment) throws + { + let url = self.authFilePath(env: env) + + var json: [String: Any] = [:] + if let data = try? Data(contentsOf: url), + let existing = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + { + json = existing + } + + var tokens: [String: Any] = [ + "access_token": credentials.accessToken, + "refresh_token": credentials.refreshToken, + ] + if let idToken = credentials.idToken { + tokens["id_token"] = idToken + } + if let accountId = credentials.accountId { + tokens["account_id"] = accountId + } + + json["tokens"] = tokens + json["last_refresh"] = ISO8601DateFormatter().string(from: Date()) + + let data = try JSONSerialization.data(withJSONObject: json, options: [.prettyPrinted, .sortedKeys]) + let directory = url.deletingLastPathComponent() + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + try self.writePrivateFile(data, to: url) + } + + private static func writePrivateFile( + _ data: Data, + to url: URL, + beforePublish: ((URL) throws -> Void)? = nil) throws + { + let fileManager = FileManager.default + let stagedURL = url.deletingLastPathComponent().appendingPathComponent( + ".\(url.lastPathComponent).codexbar-staged-\(UUID().uuidString)", + isDirectory: false) + let stagedPath = stagedURL.path + let descriptor = stagedPath.withCString { + open($0, O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, mode_t(0o600)) + } + guard descriptor >= 0 else { + throw self.posixError(code: errno, path: stagedPath) + } + + let handle = FileHandle(fileDescriptor: descriptor, closeOnDealloc: true) + var handleIsOpen = true + do { + guard fchmod(descriptor, mode_t(0o600)) == 0 else { + throw self.posixError(code: errno, path: stagedPath) + } + try handle.write(contentsOf: data) + try handle.synchronize() + try handle.close() + handleIsOpen = false + + try beforePublish?(stagedURL) + try self.renameItem(at: stagedURL, to: url) + } catch { + if handleIsOpen { + try? handle.close() + } + try? fileManager.removeItem(at: stagedURL) + throw error + } + } + + private static func renameItem(at sourceURL: URL, to destinationURL: URL) throws { + let result = sourceURL.path.withCString { sourcePath in + destinationURL.path.withCString { destinationPath in + rename(sourcePath, destinationPath) + } + } + guard result == 0 else { + throw self.posixError(code: errno, path: destinationURL.path) + } + } + + private static func posixError(code: Int32, path: String) -> NSError { + NSError( + domain: NSPOSIXErrorDomain, + code: Int(code), + userInfo: [NSFilePathErrorKey: path]) + } + + private static func parseLastRefresh(from raw: Any?) -> Date? { + guard let value = raw as? String, !value.isEmpty else { return nil } + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = formatter.date(from: value) { return date } + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: value) + } + + private static func stringValue( + in dictionary: [String: Any], + snakeCaseKey: String, + camelCaseKey: String) + -> String? + { + if let value = dictionary[snakeCaseKey] as? String, !value.isEmpty { + return value + } + if let value = dictionary[camelCaseKey] as? String, !value.isEmpty { + return value + } + return nil + } +} + +#if DEBUG +extension CodexOAuthCredentialsStore { + static func _authFileURLForTesting(env: [String: String]) -> URL { + self.authFilePath(env: env) + } + + static func _writePrivateFileForTesting( + _ data: Data, + to url: URL, + beforePublish: @escaping (URL) throws -> Void) throws + { + try self.writePrivateFile(data, to: url, beforePublish: beforePublish) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthUsageFetcher.swift b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthUsageFetcher.swift new file mode 100644 index 0000000..2ae5e3e --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthUsageFetcher.swift @@ -0,0 +1,655 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct CodexUsageResponse: Decodable, Sendable { + public let planType: PlanType? + public let rateLimit: RateLimitDetails? + public let credits: CreditDetails? + public let individualLimit: SpendControlLimitSnapshot? + /// Model-specific limits (e.g. GPT-5.3-Codex-Spark) that sit alongside the primary/weekly windows. + public let additionalRateLimits: [AdditionalRateLimit]? + let additionalRateLimitsDecodeFailed: Bool + + enum CodingKeys: String, CodingKey { + case planType = "plan_type" + case rateLimit = "rate_limit" + case credits + case individualLimit = "individual_limit" + case individualLimitCamel = "individualLimit" + case additionalRateLimits = "additional_rate_limits" + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.planType = try? container.decodeIfPresent(PlanType.self, forKey: .planType) + self.rateLimit = try? container.decodeIfPresent(RateLimitDetails.self, forKey: .rateLimit) + self.credits = try? container.decodeIfPresent(CreditDetails.self, forKey: .credits) + self.individualLimit = (try? container.decodeIfPresent( + SpendControlLimitSnapshot.self, + forKey: .individualLimit)) + ?? (try? container.decodeIfPresent(SpendControlLimitSnapshot.self, forKey: .individualLimitCamel)) + // Optional and additive: missing/malformed extra limits must never disturb primary/weekly mapping. + // Decode per element so a single malformed entry cannot discard its valid siblings; a non-array + // value (or absent field) leaves `additionalRateLimits` nil and primary/weekly mapping untouched. + let additionalRateLimitsHadValue = Self.hasNonNilValue(container: container, key: .additionalRateLimits) + do { + let decoded = try container.decodeIfPresent( + [LossyAdditionalRateLimit].self, + forKey: .additionalRateLimits) + self.additionalRateLimits = decoded?.compactMap(\.value) + self.additionalRateLimitsDecodeFailed = decoded?.contains(where: \.decodeFailed) == true + || self.additionalRateLimits?.contains(where: \.hasWindowDecodeFailure) == true + } catch { + self.additionalRateLimits = nil + self.additionalRateLimitsDecodeFailed = additionalRateLimitsHadValue + } + } + + private static func hasNonNilValue( + container: KeyedDecodingContainer, + key: CodingKeys) -> Bool + { + guard container.contains(key) else { return false } + return (try? container.decodeNil(forKey: key)) == false + } + + public enum PlanType: Sendable, Decodable, Equatable { + case guest + case free + case go + case plus + case pro + case freeWorkspace + case team + case business + case education + case quorum + case k12 + case enterprise + case edu + case unknown(String) + + public var rawValue: String { + switch self { + case .guest: "guest" + case .free: "free" + case .go: "go" + case .plus: "plus" + case .pro: "pro" + case .freeWorkspace: "free_workspace" + case .team: "team" + case .business: "business" + case .education: "education" + case .quorum: "quorum" + case .k12: "k12" + case .enterprise: "enterprise" + case .edu: "edu" + case let .unknown(value): value + } + } + + public init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + let value = try container.decode(String.self) + switch value { + case "guest": self = .guest + case "free": self = .free + case "go": self = .go + case "plus": self = .plus + case "pro": self = .pro + case "free_workspace": self = .freeWorkspace + case "team": self = .team + case "business": self = .business + case "education": self = .education + case "quorum": self = .quorum + case "k12": self = .k12 + case "enterprise": self = .enterprise + case "edu": self = .edu + default: + self = .unknown(value) + } + } + } + + public struct RateLimitDetails: Decodable, Sendable { + public let primaryWindow: WindowSnapshot? + public let secondaryWindow: WindowSnapshot? + public let individualLimit: SpendControlLimitSnapshot? + let primaryWindowDecodeFailed: Bool + let secondaryWindowDecodeFailed: Bool + + enum CodingKeys: String, CodingKey { + case primaryWindow = "primary_window" + case secondaryWindow = "secondary_window" + case individualLimit = "individual_limit" + case individualLimitCamel = "individualLimit" + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let primaryHadValue = Self.hasNonNilValue(container: container, key: .primaryWindow) + do { + self.primaryWindow = try container.decodeIfPresent(WindowSnapshot.self, forKey: .primaryWindow) + self.primaryWindowDecodeFailed = false + } catch { + self.primaryWindow = nil + self.primaryWindowDecodeFailed = primaryHadValue + } + + let secondaryHadValue = Self.hasNonNilValue(container: container, key: .secondaryWindow) + do { + self.secondaryWindow = try container.decodeIfPresent(WindowSnapshot.self, forKey: .secondaryWindow) + self.secondaryWindowDecodeFailed = false + } catch { + self.secondaryWindow = nil + self.secondaryWindowDecodeFailed = secondaryHadValue + } + self.individualLimit = (try? container.decodeIfPresent( + SpendControlLimitSnapshot.self, + forKey: .individualLimit)) + ?? (try? container.decodeIfPresent(SpendControlLimitSnapshot.self, forKey: .individualLimitCamel)) + } + + private static func hasNonNilValue( + container: KeyedDecodingContainer, + key: CodingKeys) -> Bool + { + guard container.contains(key) else { return false } + return (try? container.decodeNil(forKey: key)) == false + } + + var hasWindowDecodeFailure: Bool { + self.primaryWindowDecodeFailed || self.secondaryWindowDecodeFailed + } + } + + public struct WindowSnapshot: Decodable, Sendable { + public let usedPercent: Int + public let resetAt: Int + public let limitWindowSeconds: Int + + enum CodingKeys: String, CodingKey { + case usedPercent = "used_percent" + case resetAt = "reset_at" + case limitWindowSeconds = "limit_window_seconds" + } + } + + /// One entry of `additional_rate_limits`: a named, model-specific limit (e.g. GPT-5.3-Codex-Spark) + /// whose windows reuse the same shape as the primary/weekly `RateLimitDetails`. + public struct AdditionalRateLimit: Decodable, Sendable { + public let limitName: String? + public let meteredFeature: String? + public let rateLimit: RateLimitDetails? + let rateLimitDecodeFailed: Bool + + enum CodingKeys: String, CodingKey { + case limitName = "limit_name" + case meteredFeature = "metered_feature" + case rateLimit = "rate_limit" + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.limitName = try? container.decodeIfPresent(String.self, forKey: .limitName) + self.meteredFeature = try? container.decodeIfPresent(String.self, forKey: .meteredFeature) + let rateLimitHadValue = Self.hasNonNilValue(container: container, key: .rateLimit) + do { + self.rateLimit = try container.decodeIfPresent(RateLimitDetails.self, forKey: .rateLimit) + self.rateLimitDecodeFailed = false + } catch { + self.rateLimit = nil + self.rateLimitDecodeFailed = rateLimitHadValue + } + } + + private static func hasNonNilValue( + container: KeyedDecodingContainer, + key: CodingKeys) -> Bool + { + guard container.contains(key) else { return false } + return (try? container.decodeNil(forKey: key)) == false + } + + var hasWindowDecodeFailure: Bool { + self.rateLimitDecodeFailed || self.rateLimit?.hasWindowDecodeFailure == true + } + } + + /// Decodes a single `additional_rate_limits` element without ever throwing, so one malformed + /// entry cannot discard its valid siblings during array decoding. + private struct LossyAdditionalRateLimit: Decodable { + let value: AdditionalRateLimit? + let decodeFailed: Bool + + init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + self.value = try? container.decode(AdditionalRateLimit.self) + self.decodeFailed = self.value == nil + } + } + + public struct SpendControlLimitSnapshot: Decodable, Sendable { + public let limit: Double? + public let used: Double? + public let remainingPercent: Double? + public let resetsAt: Int? + + enum CodingKeys: String, CodingKey { + case limit + case used + case remainingPercent + case remainingPercentSnake = "remaining_percent" + case resetsAt + case resetsAtSnake = "resets_at" + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.limit = Self.decodeFlexibleDouble(container, forKey: .limit) + self.used = Self.decodeFlexibleDouble(container, forKey: .used) + self.remainingPercent = Self.decodeFlexibleDouble(container, forKey: .remainingPercent) + ?? Self.decodeFlexibleDouble(container, forKey: .remainingPercentSnake) + self.resetsAt = Self.decodeFlexibleInt(container, forKey: .resetsAt) + ?? Self.decodeFlexibleInt(container, forKey: .resetsAtSnake) + } + + private static func decodeFlexibleDouble( + _ container: KeyedDecodingContainer, + forKey key: CodingKeys) -> Double? + { + if let value = try? container.decodeIfPresent(Double.self, forKey: key) { + return value + } + if let value = try? container.decodeIfPresent(Int.self, forKey: key) { + return Double(value) + } + if let value = try? container.decodeIfPresent(String.self, forKey: key) { + return Double(value.trimmingCharacters(in: .whitespacesAndNewlines)) + } + return nil + } + + private static func decodeFlexibleInt( + _ container: KeyedDecodingContainer, + forKey key: CodingKeys) -> Int? + { + if let value = try? container.decodeIfPresent(Int.self, forKey: key) { + return value + } + if let value = try? container.decodeIfPresent(Double.self, forKey: key) { + return Int(value) + } + if let value = try? container.decodeIfPresent(String.self, forKey: key) { + return Int(value.trimmingCharacters(in: .whitespacesAndNewlines)) + } + return nil + } + } + + public struct CreditDetails: Decodable, Sendable { + public let hasCredits: Bool + public let unlimited: Bool + public let balance: Double? + + enum CodingKeys: String, CodingKey { + case hasCredits = "has_credits" + case unlimited + case balance + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.hasCredits = (try? container.decode(Bool.self, forKey: .hasCredits)) ?? false + self.unlimited = (try? container.decode(Bool.self, forKey: .unlimited)) ?? false + if let balance = try? container.decode(Double.self, forKey: .balance) { + self.balance = balance + } else if let balance = try? container.decode(String.self, forKey: .balance), + let value = Double(balance) + { + self.balance = value + } else { + self.balance = nil + } + } + } +} + +public enum CodexOAuthFetchError: LocalizedError, Sendable { + case unauthorized + case invalidResponse + case serverError(Int, String?) + case networkError(Error) + + public var errorDescription: String? { + switch self { + case .unauthorized: + return "Codex OAuth token expired or invalid. Run `codex` to re-authenticate." + case .invalidResponse: + return "Invalid response from Codex usage API." + case let .serverError(code, message): + if let message, !message.isEmpty { + return "Codex API error \(code): \(message)" + } + return "Codex API error \(code)." + case let .networkError(error): + return "Network error: \(error.localizedDescription)" + } + } +} + +public enum CodexOAuthUsageFetcher { + private static let defaultChatGPTBaseURL = "https://chatgpt.com/backend-api/" + private static let chatGPTUsagePath = "/wham/usage" + private static let codexUsagePath = "/api/codex/usage" + private static let rateLimitResetCreditsPath = "/wham/rate-limit-reset-credits" + + public static func fetchUsage( + accessToken: String, + accountId: String?, + env: [String: String] = ProcessInfo.processInfo.environment) async throws -> CodexUsageResponse + { + try await self.fetchUsage( + accessToken: accessToken, + accountId: accountId, + env: env, + session: CodexAuthenticatedHTTPTransport.current) + } + + public static func fetchUsage( + accessToken: String, + accountId: String?, + env: [String: String] = ProcessInfo.processInfo.environment, + session transport: any ProviderHTTPTransport) async throws -> CodexUsageResponse + { + var request = URLRequest( + url: Self.resolveUsageURL(env: env), + cachePolicy: .reloadIgnoringLocalCacheData, + timeoutInterval: 30) + request.httpMethod = "GET" + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("CodexBar", forHTTPHeaderField: "User-Agent") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + if let accountId, !accountId.isEmpty { + request.setValue(accountId, forHTTPHeaderField: "ChatGPT-Account-Id") + } + + do { + let response = try await transport.response(for: request) + let data = response.data + + switch response.statusCode { + case 200...299: + do { + return try JSONDecoder().decode(CodexUsageResponse.self, from: data) + } catch { + throw CodexOAuthFetchError.invalidResponse + } + case 401, 403: + throw CodexOAuthFetchError.unauthorized + default: + let body = String(data: data, encoding: .utf8) + throw CodexOAuthFetchError.serverError(response.statusCode, body) + } + } catch let error as CodexOAuthFetchError { + throw error + } catch is CancellationError { + throw CancellationError() + } catch { + if Task.isCancelled || (error as? URLError)?.code == .cancelled { + throw CancellationError() + } + throw CodexOAuthFetchError.networkError(error) + } + } + + public static func fetchRateLimitResetCredits( + accessToken: String, + accountId: String?, + env: [String: String] = ProcessInfo.processInfo.environment, + timeout: TimeInterval = 4) async throws -> CodexRateLimitResetCreditsSnapshot + { + try await self.fetchRateLimitResetCredits( + accessToken: accessToken, + accountId: accountId, + env: env, + timeout: timeout, + session: CodexAuthenticatedHTTPTransport.current) + } + + public static func fetchRateLimitResetCredits( + accessToken: String, + accountId: String?, + env: [String: String] = ProcessInfo.processInfo.environment, + timeout: TimeInterval = 4, + session transport: any ProviderHTTPTransport) async throws + -> CodexRateLimitResetCreditsSnapshot + { + var request = URLRequest( + url: Self.resolveRateLimitResetCreditsURL(env: env), + cachePolicy: .reloadIgnoringLocalCacheData, + timeoutInterval: timeout) + request.httpMethod = "GET" + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("CodexBar", forHTTPHeaderField: "User-Agent") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("codex-1", forHTTPHeaderField: "OpenAI-Beta") + request.setValue("Codex Desktop", forHTTPHeaderField: "originator") + + if let accountId, !accountId.isEmpty { + request.setValue(accountId, forHTTPHeaderField: "ChatGPT-Account-ID") + } + + do { + let response = try await transport.response(for: request) + let data = response.data + + switch response.statusCode { + case 200...299: + do { + let decoder = JSONDecoder() + decoder.dateDecodingStrategy = .custom(Self.decodeISO8601Date) + let payload = try decoder.decode(RateLimitResetCreditsResponse.self, from: data) + guard payload.availableCount >= 0 else { + throw CodexOAuthFetchError.invalidResponse + } + return CodexRateLimitResetCreditsSnapshot( + credits: payload.credits.map(\.model), + availableCount: payload.availableCount, + updatedAt: Date()) + } catch { + throw CodexOAuthFetchError.invalidResponse + } + case 401, 403: + throw CodexOAuthFetchError.unauthorized + default: + let body = String(data: data, encoding: .utf8) + throw CodexOAuthFetchError.serverError(response.statusCode, body) + } + } catch let error as CodexOAuthFetchError { + throw error + } catch is CancellationError { + throw CancellationError() + } catch { + if Task.isCancelled || (error as? URLError)?.code == .cancelled { + throw CancellationError() + } + throw CodexOAuthFetchError.networkError(error) + } + } + + private static func resolveUsageURL(env: [String: String]) -> URL { + self.resolveUsageURL(env: env, configContents: nil) + } + + private static func resolveUsageURL(env: [String: String], configContents: String?) -> URL { + let baseURL = self.resolveChatGPTBaseURL(env: env, configContents: configContents) + let normalized = self.normalizeChatGPTBaseURL(baseURL) + let path = normalized.contains("/backend-api") ? Self.chatGPTUsagePath : Self.codexUsagePath + let full = normalized + path + return URL(string: full) ?? URL(string: Self.defaultChatGPTBaseURL + Self.chatGPTUsagePath)! + } + + private static func resolveRateLimitResetCreditsURL(env: [String: String]) -> URL { + self.resolveRateLimitResetCreditsURL(env: env, configContents: nil) + } + + private static func resolveRateLimitResetCreditsURL(env: [String: String], configContents: String?) -> URL { + let baseURL = self.resolveChatGPTBaseURL(env: env, configContents: configContents) + let normalized = self.normalizeChatGPTBaseURL(baseURL) + let full = normalized + Self.rateLimitResetCreditsPath + return URL(string: full) ?? URL(string: Self.defaultChatGPTBaseURL + Self.rateLimitResetCreditsPath)! + } + + private static func resolveChatGPTBaseURL(env: [String: String], configContents: String?) -> String { + if let configContents, let parsed = self.parseChatGPTBaseURL(from: configContents) { + return parsed + } + if let contents = self.loadConfigContents(env: env), + let parsed = self.parseChatGPTBaseURL(from: contents) + { + return parsed + } + return Self.defaultChatGPTBaseURL + } + + private static func normalizeChatGPTBaseURL(_ value: String) -> String { + var trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.isEmpty { trimmed = Self.defaultChatGPTBaseURL } + while trimmed.hasSuffix("/") { + trimmed.removeLast() + } + if trimmed.hasPrefix("https://chatgpt.com") || trimmed.hasPrefix("https://chat.openai.com"), + !trimmed.contains("/backend-api") + { + trimmed += "/backend-api" + } + return trimmed + } + + private static func parseChatGPTBaseURL(from contents: String) -> String? { + for rawLine in contents.split(whereSeparator: \.isNewline) { + let line = rawLine.split(separator: "#", maxSplits: 1, omittingEmptySubsequences: true).first + let trimmed = line?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + guard !trimmed.isEmpty else { continue } + let parts = trimmed.split(separator: "=", maxSplits: 1, omittingEmptySubsequences: true) + guard parts.count == 2 else { continue } + let key = parts[0].trimmingCharacters(in: .whitespacesAndNewlines) + guard key == "chatgpt_base_url" else { continue } + var value = parts[1].trimmingCharacters(in: .whitespacesAndNewlines) + if value.hasPrefix("\""), value.hasSuffix("\"") { + value = String(value.dropFirst().dropLast()) + } else if value.hasPrefix("'"), value.hasSuffix("'") { + value = String(value.dropFirst().dropLast()) + } + return value.trimmingCharacters(in: .whitespacesAndNewlines) + } + return nil + } + + private static func loadConfigContents(env: [String: String]) -> String? { + let home = FileManager.default.homeDirectoryForCurrentUser + let codexHome = env["CODEX_HOME"]?.trimmingCharacters(in: .whitespacesAndNewlines) + let root = (codexHome?.isEmpty == false) ? URL(fileURLWithPath: codexHome!) : home + .appendingPathComponent(".codex") + let url = root.appendingPathComponent("config.toml") + return try? String(contentsOf: url, encoding: .utf8) + } + + private struct RateLimitResetCreditsResponse: Decodable { + let credits: [RateLimitResetCreditResponse] + let availableCount: Int + + private enum CodingKeys: String, CodingKey { + case credits + case availableCount = "available_count" + } + } + + private struct RateLimitResetCreditResponse: Decodable { + let id: String + let resetType: String + let status: CodexRateLimitResetCreditStatus + let grantedAt: Date + let expiresAt: Date? + let redeemStartedAt: Date? + let redeemedAt: Date? + let title: String? + let description: String? + + private enum CodingKeys: String, CodingKey { + case id + case resetType = "reset_type" + case status + case grantedAt = "granted_at" + case expiresAt = "expires_at" + case redeemStartedAt = "redeem_started_at" + case redeemedAt = "redeemed_at" + case title + case description + } + + var model: CodexRateLimitResetCredit { + CodexRateLimitResetCredit( + id: self.id, + resetType: self.resetType, + status: self.status, + grantedAt: self.grantedAt, + expiresAt: self.expiresAt, + redeemStartedAt: self.redeemStartedAt, + redeemedAt: self.redeemedAt, + title: self.title, + description: self.description) + } + } + + private static func decodeISO8601Date(from decoder: Decoder) throws -> Date { + let container = try decoder.singleValueContainer() + let raw = try container.decode(String.self) + let fractional = ISO8601DateFormatter() + fractional.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + let seconds = ISO8601DateFormatter() + seconds.formatOptions = [.withInternetDateTime] + if let date = fractional.date(from: raw) ?? seconds.date(from: raw) { + return date + } + throw DecodingError.dataCorruptedError( + in: container, + debugDescription: "Invalid ISO-8601 date: \(raw)") + } +} + +#if DEBUG +extension CodexOAuthUsageFetcher { + static func _resolveUsageURLForTesting(env: [String: String] = [:], configContents: String? = nil) -> URL { + self.resolveUsageURL(env: env, configContents: configContents) + } + + static func _decodeUsageResponseForTesting(_ data: Data) throws -> CodexUsageResponse { + try JSONDecoder().decode(CodexUsageResponse.self, from: data) + } + + static func _resolveRateLimitResetCreditsURLForTesting( + env: [String: String] = [:], + configContents: String? = nil) -> URL + { + self.resolveRateLimitResetCreditsURL(env: env, configContents: configContents) + } + + static func _decodeRateLimitResetCreditsForTesting( + _ data: Data, + now: Date = Date()) throws -> CodexRateLimitResetCreditsSnapshot + { + let decoder = JSONDecoder() + decoder.dateDecodingStrategy = .custom(Self.decodeISO8601Date) + let payload = try decoder.decode(RateLimitResetCreditsResponse.self, from: data) + return CodexRateLimitResetCreditsSnapshot( + credits: payload.credits.map(\.model), + availableCount: payload.availableCount, + updatedAt: now) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexTokenRefresher.swift b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexTokenRefresher.swift new file mode 100644 index 0000000..ff79764 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexTokenRefresher.swift @@ -0,0 +1,123 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum CodexTokenRefresher { + private static let refreshEndpoint = URL(string: "https://auth.openai.com/oauth/token")! + private static let clientID = "app_EMoamEEZ73f0CkXaXp7hrann" + + public enum RefreshError: LocalizedError, Sendable { + case expired + case revoked + case reused + case networkError(Error) + case invalidResponse(String) + + public var errorDescription: String? { + switch self { + case .expired: + "Refresh token expired. Please run `codex` to log in again." + case .revoked: + "Refresh token was revoked. Please run `codex` to log in again." + case .reused: + "Refresh token was already used. Please run `codex` to log in again." + case let .networkError(error): + "Network error during token refresh: \(error.localizedDescription)" + case let .invalidResponse(message): + "Invalid refresh response: \(message)" + } + } + } + + public static func refresh(_ credentials: CodexOAuthCredentials) async throws -> CodexOAuthCredentials { + try await self.refresh(credentials, session: CodexAuthenticatedHTTPTransport.current) + } + + static func refresh( + _ credentials: CodexOAuthCredentials, + session transport: any ProviderHTTPTransport) async throws -> CodexOAuthCredentials + { + guard !credentials.refreshToken.isEmpty else { + return credentials + } + + var request = URLRequest( + url: Self.refreshEndpoint, + cachePolicy: .reloadIgnoringLocalCacheData, + timeoutInterval: 30) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + + let body: [String: String] = [ + "client_id": Self.clientID, + "grant_type": "refresh_token", + "refresh_token": credentials.refreshToken, + "scope": "openid profile email", + ] + request.httpBody = try JSONSerialization.data(withJSONObject: body) + + do { + let response = try await transport.response(for: request) + let data = response.data + guard response.statusCode == 200 else { + throw Self.refreshFailureError(statusCode: response.statusCode, data: data) + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw RefreshError.invalidResponse("Invalid JSON") + } + + let newAccessToken = json["access_token"] as? String ?? credentials.accessToken + let newRefreshToken = json["refresh_token"] as? String ?? credentials.refreshToken + let newIdToken = json["id_token"] as? String ?? credentials.idToken + + return CodexOAuthCredentials( + accessToken: newAccessToken, + refreshToken: newRefreshToken, + idToken: newIdToken, + accountId: credentials.accountId, + lastRefresh: Date()) + } catch let error as RefreshError { + throw error + } catch { + throw RefreshError.networkError(error) + } + } + + private static func extractErrorCode(from data: Data) -> String? { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { return nil } + if let error = json["error"] as? [String: Any], let code = error["code"] as? String { return code } + if let error = json["error"] as? String { return error } + return json["code"] as? String + } + + private static func refreshFailureError(statusCode: Int, data: Data) -> RefreshError { + if let errorCode = extractErrorCode(from: data) { + switch errorCode.lowercased() { + case "refresh_token_expired": + return .expired + case "refresh_token_reused": + return .reused + case "invalid_grant", "refresh_token_invalidated": + return .revoked + default: + break + } + } + + if statusCode == 401 { + return .expired + } + + return .invalidResponse("Status \(statusCode)") + } +} + +#if DEBUG +extension CodexTokenRefresher { + static func _refreshFailureErrorForTesting(statusCode: Int, data: Data) -> RefreshError { + self.refreshFailureError(statusCode: statusCode, data: data) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceIdentityCache.swift b/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceIdentityCache.swift new file mode 100644 index 0000000..2809c8c --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceIdentityCache.swift @@ -0,0 +1,115 @@ +import Foundation + +public struct CodexOpenAIWorkspaceIdentityCache: @unchecked Sendable { + public static let currentVersion = 1 + + private struct Payload: Codable { + let version: Int + var labelsByWorkspaceAccountID: [String: String] + + init( + version: Int = CodexOpenAIWorkspaceIdentityCache.currentVersion, + labelsByWorkspaceAccountID: [String: String]) + { + self.version = version + self.labelsByWorkspaceAccountID = labelsByWorkspaceAccountID + } + } + + #if DEBUG + @TaskLocal static var taskFileURLOverride: URL? + #endif + + private let fileURL: URL + private let fileManager: FileManager + + public init(fileURL: URL = Self.defaultURL(), fileManager: FileManager = .default) { + self.fileURL = fileURL + self.fileManager = fileManager + } + + public func workspaceLabel(for workspaceAccountID: String?) -> String? { + guard let normalizedWorkspaceAccountID = CodexOpenAIWorkspaceResolver + .normalizeWorkspaceAccountID(workspaceAccountID) + else { + return nil + } + + return self.loadPayload().labelsByWorkspaceAccountID[normalizedWorkspaceAccountID] + } + + public func store(_ identity: CodexOpenAIWorkspaceIdentity) throws { + guard let workspaceLabel = CodexOpenAIWorkspaceIdentity.normalizeWorkspaceLabel(identity.workspaceLabel) else { + return + } + + var payload = self.loadPayload() + payload.labelsByWorkspaceAccountID[identity.workspaceAccountID] = workspaceLabel + + let encoder = JSONEncoder() + encoder.outputFormatting = [.prettyPrinted, .sortedKeys] + let data = try encoder.encode(payload) + let directory = self.fileURL.deletingLastPathComponent() + if !self.fileManager.fileExists(atPath: directory.path) { + try self.fileManager.createDirectory(at: directory, withIntermediateDirectories: true) + } + try data.write(to: self.fileURL, options: [.atomic]) + try self.applySecurePermissionsIfNeeded() + } + + private func loadPayload() -> Payload { + guard self.fileManager.fileExists(atPath: self.fileURL.path), + let data = try? Data(contentsOf: self.fileURL), + let payload = try? JSONDecoder().decode(Payload.self, from: data), + payload.version == Self.currentVersion + else { + return Payload(labelsByWorkspaceAccountID: [:]) + } + + return payload + } + + private func applySecurePermissionsIfNeeded() throws { + #if os(macOS) + try self.fileManager.setAttributes([ + .posixPermissions: NSNumber(value: Int16(0o600)), + ], ofItemAtPath: self.fileURL.path) + #endif + } + + public static func defaultURL() -> URL { + #if DEBUG + if let override = self.taskFileURLOverride { + return override + } + #endif + + let base = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first + ?? FileManager.default.homeDirectoryForCurrentUser + return base + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("codex-openai-workspaces.json") + } +} + +#if DEBUG +extension CodexOpenAIWorkspaceIdentityCache { + public static func withFileURLOverrideForTesting( + _ fileURL: URL, + operation: () throws -> T) rethrows -> T + { + try self.$taskFileURLOverride.withValue(fileURL) { + try operation() + } + } + + public static func withFileURLOverrideForTesting( + _ fileURL: URL, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskFileURLOverride.withValue(fileURL) { + try await operation() + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceResolver.swift b/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceResolver.swift new file mode 100644 index 0000000..9f50276 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceResolver.swift @@ -0,0 +1,123 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct CodexOpenAIWorkspaceIdentity: Equatable, Sendable { + public let workspaceAccountID: String + public let workspaceLabel: String? + + public init(workspaceAccountID: String, workspaceLabel: String?) { + self.workspaceAccountID = Self.normalizeWorkspaceAccountID(workspaceAccountID) + self.workspaceLabel = Self.normalizeWorkspaceLabel(workspaceLabel) + } + + public static func normalizeWorkspaceAccountID(_ value: String) -> String { + value.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } + + public static func normalizeWorkspaceLabel(_ value: String?) -> String? { + guard let trimmed = value?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed + } +} + +public enum CodexOpenAIWorkspaceResolver { + private struct AccountsResponse: Decodable { + let items: [AccountItem] + } + + private struct AccountItem: Decodable { + let id: String + let name: String? + } + + private static let accountsURL = URL(string: "https://chatgpt.com/backend-api/accounts")! + + public static func resolve(credentials: CodexOAuthCredentials) async throws -> CodexOpenAIWorkspaceIdentity? { + try await self.resolve(credentials: credentials, session: CodexAuthenticatedHTTPTransport.current) + } + + public static func resolve( + credentials: CodexOAuthCredentials, + session transport: any ProviderHTTPTransport) async throws -> CodexOpenAIWorkspaceIdentity? + { + guard let workspaceAccountID = normalizeWorkspaceAccountID(credentials.accountId) else { + return nil + } + + let identities = try await self.listWorkspaces(credentials: credentials, session: transport) + if let identity = identities.first(where: { $0.workspaceAccountID == workspaceAccountID }) { + return identity + } + + return CodexOpenAIWorkspaceIdentity( + workspaceAccountID: workspaceAccountID, + workspaceLabel: nil) + } + + public static func listWorkspaces( + credentials: CodexOAuthCredentials) async throws -> [CodexOpenAIWorkspaceIdentity] + { + try await self.listWorkspaces(credentials: credentials, session: CodexAuthenticatedHTTPTransport.current) + } + + public static func listWorkspaces( + credentials: CodexOAuthCredentials, + session transport: any ProviderHTTPTransport) async throws -> [CodexOpenAIWorkspaceIdentity] + { + var request = URLRequest( + url: self.accountsURL, + cachePolicy: .reloadIgnoringLocalCacheData, + timeoutInterval: 20) + request.httpMethod = "GET" + request.setValue("Bearer \(credentials.accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("codex-cli", forHTTPHeaderField: "User-Agent") + request.setValue("application/json", forHTTPHeaderField: "Accept") + if let workspaceAccountID = normalizeWorkspaceAccountID(credentials.accountId) { + request.setValue(workspaceAccountID, forHTTPHeaderField: "ChatGPT-Account-Id") + } + + let response = try await transport.response(for: request) + guard (200..<300).contains(response.statusCode) + else { + throw CodexOpenAIWorkspaceResolverError.invalidResponse + } + + let decoded = try JSONDecoder().decode(AccountsResponse.self, from: response.data) + return decoded.items.compactMap { account in + guard let workspaceAccountID = self.normalizeWorkspaceAccountID(account.id) else { return nil } + return CodexOpenAIWorkspaceIdentity( + workspaceAccountID: workspaceAccountID, + workspaceLabel: self.resolveWorkspaceLabel(from: account)) + } + } + + public static func normalizeWorkspaceAccountID(_ value: String?) -> String? { + guard let trimmed = value?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed.lowercased() + } + + private static func resolveWorkspaceLabel(from account: AccountItem) -> String? { + let name = account.name?.trimmingCharacters(in: .whitespacesAndNewlines) + if let name, !name.isEmpty { + return name + } + return "Personal" + } +} + +public enum CodexOpenAIWorkspaceResolverError: LocalizedError, Sendable { + case invalidResponse + + public var errorDescription: String? { + switch self { + case .invalidResponse: + "OpenAI account lookup returned an invalid response." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexPlanFormatting.swift b/Sources/CodexBarCore/Providers/Codex/CodexPlanFormatting.swift new file mode 100644 index 0000000..af546a3 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexPlanFormatting.swift @@ -0,0 +1,59 @@ +import Foundation + +public enum CodexPlanFormatting { + private static let exactDisplayNames: [String: String] = [ + "pro": "Pro 20x", + "prolite": "Pro 5x", + "pro_lite": "Pro 5x", + "pro-lite": "Pro 5x", + "pro lite": "Pro 5x", + ] + + private static let uppercaseWords: Set = [ + "cbp", + "k12", + ] + + public static func displayName(_ raw: String?) -> String? { + guard let raw = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + + let lower = raw.lowercased() + if let exact = Self.exactDisplayNames[lower] { + return exact + } + + let cleaned = UsageFormatter.cleanPlanName(raw) + let candidate = cleaned.trimmingCharacters(in: .whitespacesAndNewlines) + guard !candidate.isEmpty else { return raw } + + if let exact = Self.exactDisplayNames[candidate.lowercased()] { + return exact + } + + let components = candidate + .split(whereSeparator: { $0 == "_" || $0 == "-" || $0.isWhitespace }) + .map(String.init) + .filter { !$0.isEmpty } + + guard !components.isEmpty else { return candidate } + + let formatted = components.map(Self.wordDisplayName).joined(separator: " ") + return formatted.isEmpty ? candidate : formatted + } + + private static func wordDisplayName(_ raw: String) -> String { + let lower = raw.lowercased() + if Self.uppercaseWords.contains(lower) { + return lower.uppercased() + } + if raw == raw.uppercased(), raw.contains(where: \.isLetter) { + return raw + } + if let first = raw.first, first.isLowercase { + return raw.prefix(1).uppercased() + raw.dropFirst() + } + return raw + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift new file mode 100644 index 0000000..5a15ee1 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift @@ -0,0 +1,454 @@ +import Foundation + +public enum CodexProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .codex, + metadata: ProviderMetadata( + id: .codex, + displayName: "Codex", + sessionLabel: "Session", + weeklyLabel: "Weekly", + opusLabel: nil, + supportsOpus: false, + supportsCredits: true, + creditsHint: "Credits unavailable; keep Codex running to refresh.", + toggleTitle: "Show Codex usage", + cliName: "codex", + defaultEnabled: true, + isPrimaryProvider: true, + usesAccountFallback: true, + browserCookieOrder: ProviderBrowserCookieDefaults.codexCookieImportOrder + ?? ProviderBrowserCookieDefaults.defaultImportOrder, + dashboardURL: "https://chatgpt.com/codex/settings/usage", + changelogURL: "https://github.com/openai/codex/releases", + statusPageURL: "https://status.openai.com/"), + branding: ProviderBranding( + iconStyle: .codex, + iconResourceName: "ProviderIcon-codex", + color: ProviderColor(red: 73 / 255, green: 163 / 255, blue: 176 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: true, + noDataMessage: self.noDataMessage), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web, .cli, .oauth], + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), + cli: ProviderCLIConfig( + name: "codex", + versionDetector: { _ in ProviderVersionDetector.codexVersion() })) + } + + private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] { + let cli = CodexCLIUsageStrategy() + let oauth = CodexOAuthFetchStrategy() + let web = CodexWebDashboardStrategy() + + switch context.runtime { + case .cli: + switch context.sourceMode { + case .oauth: + return [oauth] + case .web: + return [web] + case .cli: + return [cli] + case .api: + return [] + case .auto: + return [oauth, cli] + } + case .app: + switch context.sourceMode { + case .oauth: + return [oauth] + case .cli: + return [cli] + case .web: + return [web] + case .api: + return [] + case .auto: + return [oauth, cli] + } + } + } + + private static func noDataMessage() -> String { + self.noDataMessage(env: ProcessInfo.processInfo.environment) + } + + private static func noDataMessage(env: [String: String], fileManager: FileManager = .default) -> String { + let base = CodexHomeScope.ambientHomeURL(env: env, fileManager: fileManager).path + let sessions = "\(base)/sessions" + let archived = "\(base)/archived_sessions" + return "No Codex sessions found in \(sessions) or \(archived)." + } + + public static func resolveUsageStrategy( + selectedDataSource: CodexUsageDataSource, + hasOAuthCredentials: Bool) -> CodexUsageStrategy + { + if selectedDataSource == .auto { + if hasOAuthCredentials { + return CodexUsageStrategy(dataSource: .oauth) + } + return CodexUsageStrategy(dataSource: .cli) + } + return CodexUsageStrategy(dataSource: selectedDataSource) + } +} + +public struct CodexUsageStrategy: Equatable, Sendable { + public let dataSource: CodexUsageDataSource +} + +struct CodexCLIUsageStrategy: ProviderFetchStrategy { + let id: String = "codex.cli" + let kind: ProviderFetchKind = .cli + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolvedBinary(env: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let snapshot = try await context.fetcher.loadLatestCLIAccountSnapshot() + guard let usage = snapshot.usage else { + guard context.includeCredits, let credits = snapshot.credits else { + throw UsageError.noRateLimitsFound + } + // Credits refresh can succeed even when RPC omits rate-limit windows. + return self.makeResult( + usage: UsageSnapshot( + primary: nil, + secondary: nil, + updatedAt: credits.updatedAt, + identity: snapshot.identity), + credits: credits, + sourceLabel: "codex-cli") + } + let credits = context.includeCredits ? snapshot.credits : nil + return self.makeResult( + usage: usage, + credits: credits, + sourceLabel: "codex-cli") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + static func resolvedBinary( + env: [String: String], + loginPATH: [String]? = LoginShellPathCache.shared.current, + commandV: (String, String?, TimeInterval, FileManager) -> String? = ShellCommandLocator.commandV, + aliasResolver: (String, String?, TimeInterval, FileManager, String) -> String? = ShellCommandLocator + .resolveAlias, + fileManager: FileManager = .default, + home: String = NSHomeDirectory()) -> String? + { + BinaryLocator.resolveCodexBinary( + env: env, + loginPATH: loginPATH, + commandV: commandV, + aliasResolver: aliasResolver, + fileManager: fileManager, + home: home) + } +} + +struct CodexOAuthFetchStrategy: ProviderFetchStrategy { + let id: String = "codex.oauth" + let kind: ProviderFetchKind = .oauth + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + (try? CodexOAuthCredentialsStore.load(env: context.env)) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + var credentials = try CodexOAuthCredentialsStore.load(env: context.env) + + if credentials.needsRefresh, !credentials.refreshToken.isEmpty { + credentials = try await CodexTokenRefresher.refresh(credentials) + try CodexOAuthCredentialsStore.save(credentials, env: context.env) + } + + let usage = try await CodexOAuthUsageFetcher.fetchUsage( + accessToken: credentials.accessToken, + accountId: credentials.accountId, + env: context.env) + let resetCredits = try await Self.fetchResetCreditsIfRequested( + context: context, + credentials: credentials) + let updatedAt = Date() + let oauthResult = try Self.makeResult( + usageResponse: usage, + resetCredits: resetCredits, + credentials: credentials, + updatedAt: updatedAt, + allowEmptyUsageForResetCreditEnrichment: Self.defersResetCreditFetchToApp(context)) + return try await Self.replacingWithCLIMonthlyLimitIfAvailable(oauthResult, context: context) + } + + private static func shouldFetchResetCredits(_ context: ProviderFetchContext) -> Bool { + guard case .cli = context.runtime else { return false } + return context.includeCredits + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + guard context.sourceMode == .auto else { return false } + // Auto mode may launch the CLI as the next strategy. Keep that fallback + // limited to OAuth states the CLI can actually repair, otherwise + // transient API or decode failures can spawn `codex app-server` + // repeatedly instead of surfacing the original OAuth failure. + if let fetchError = error as? CodexOAuthFetchError { + switch fetchError { + case .unauthorized: + return true + case .invalidResponse, .serverError, .networkError: + return false + } + } + if let credentialsError = error as? CodexOAuthCredentialsError { + switch credentialsError { + case .notFound, .missingTokens: + return true + case .decodeFailed: + return false + } + } + switch error as? CodexTokenRefresher.RefreshError { + case .expired, .revoked, .reused: + return true + case .networkError, .invalidResponse, .none: + return false + } + } + + private static func mapCredits( + response: CodexUsageResponse, + updatedAt: Date) -> CreditsSnapshot? + { + let balance = response.credits?.balance + let creditLimit = (response.individualLimit ?? response.rateLimit?.individualLimit)? + .codexCreditLimitSnapshot(updatedAt: updatedAt) + guard balance != nil || creditLimit != nil else { return nil } + return CreditsSnapshot( + remaining: balance ?? 0, + events: [], + updatedAt: updatedAt, + codexCreditLimit: creditLimit) + } + + private static func makeResult( + usageResponse: CodexUsageResponse, + resetCredits: CodexRateLimitResetCreditsSnapshot? = nil, + credentials: CodexOAuthCredentials, + updatedAt: Date, + allowEmptyUsageForResetCreditEnrichment: Bool = false) throws -> ProviderFetchResult + { + let credits = Self.mapCredits(response: usageResponse, updatedAt: updatedAt) + let reconciled = CodexReconciledState.fromOAuth( + response: usageResponse, + credentials: credentials, + updatedAt: updatedAt) + + if let reconciled { + let dataConfidence: UsageDataConfidence = usageResponse.rateLimit?.hasWindowDecodeFailure == true + || usageResponse.additionalRateLimitsDecodeFailed + ? .unknown + : .exact + return CodexOAuthFetchStrategy().makeResult( + usage: reconciled.toUsageSnapshot() + .withCodexResetCredits(resetCredits) + .withDataConfidence(dataConfidence), + credits: credits, + sourceLabel: "oauth") + } + + guard credits != nil + || (resetCredits?.availableInventory(at: updatedAt).count ?? 0) > 0 + || allowEmptyUsageForResetCreditEnrichment + else { + throw UsageError.noRateLimitsFound + } + + // Credit balances and manual resets remain useful when OAuth omits + // rate-limit windows. Keep the partial result instead of discarding it. + return CodexOAuthFetchStrategy().makeResult( + usage: UsageSnapshot( + primary: nil, + secondary: nil, + tertiary: nil, + codexResetCredits: resetCredits, + updatedAt: updatedAt, + identity: CodexReconciledState.oauthIdentity( + response: usageResponse, + credentials: credentials)), + credits: credits, + sourceLabel: "oauth") + } + + private static func replacingWithCLIMonthlyLimitIfAvailable( + _ oauthResult: ProviderFetchResult, + context: ProviderFetchContext, + cliStrategy: any ProviderFetchStrategy = CodexCLIUsageStrategy()) async throws -> ProviderFetchResult + { + guard context.sourceMode == .auto, + context.includeCredits, + self.shouldTryCLIForMonthlyLimit(oauthResult) + else { return oauthResult } + guard await cliStrategy.isAvailable(context) else { return oauthResult } + let cliResult: ProviderFetchResult + do { + cliResult = try await cliStrategy.fetch(context) + } catch { + if error is CancellationError { throw error } + return oauthResult + } + guard let cliLimit = cliResult.credits?.codexCreditLimit, + self.identitiesAreCompatible(oauth: oauthResult.usage.identity, cli: cliResult.usage.identity), + let oauthCredits = oauthResult.credits + else { return oauthResult } + return ProviderFetchResult( + usage: oauthResult.usage, + credits: CreditsSnapshot( + remaining: oauthCredits.remaining, + events: oauthCredits.events, + updatedAt: oauthCredits.updatedAt, + codexCreditLimit: cliLimit), + dashboard: oauthResult.dashboard, + sourceLabel: oauthResult.sourceLabel, + strategyID: oauthResult.strategyID, + strategyKind: oauthResult.strategyKind) + } + + private static func fetchResetCreditsIfRequested( + context: ProviderFetchContext, + credentials: CodexOAuthCredentials) async throws -> CodexRateLimitResetCreditsSnapshot? + { + try await self.fetchResetCreditsIfRequested( + context: context, + credentials: credentials, + fetcher: { credentials in + try await CodexOAuthUsageFetcher.fetchRateLimitResetCredits( + accessToken: credentials.accessToken, + accountId: credentials.accountId, + env: context.env) + }) + } + + private static func defersResetCreditFetchToApp(_ context: ProviderFetchContext) -> Bool { + if case .app = context.runtime { return true } + return false + } + + private static func fetchResetCreditsIfRequested( + context: ProviderFetchContext, + credentials: CodexOAuthCredentials, + fetcher: @escaping @Sendable (CodexOAuthCredentials) async throws + -> CodexRateLimitResetCreditsSnapshot) async throws -> CodexRateLimitResetCreditsSnapshot? + { + guard self.shouldFetchResetCredits(context) else { return nil } + // The app enriches the winning outcome once in UsageStore. One-shot CLI callers do not + // pass through that app layer, so only their explicit credits flag reaches this request. + do { + return try await fetcher(credentials) + } catch { + if error is CancellationError || Task.isCancelled { throw CancellationError() } + return nil + } + } + + private static func identitiesAreCompatible( + oauth: ProviderIdentitySnapshot?, + cli: ProviderIdentitySnapshot?) -> Bool + { + guard let cliEmail = self.normalizedEmail(cli?.accountEmail), + let oauthEmail = self.normalizedEmail(oauth?.accountEmail) + else { return false } + return cliEmail == oauthEmail + } + + private static func normalizedEmail(_ email: String?) -> String? { + guard let normalized = email?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased(), + !normalized.isEmpty + else { return nil } + return normalized + } + + private static func shouldTryCLIForMonthlyLimit(_ result: ProviderFetchResult) -> Bool { + guard let credits = result.credits else { return false } + return credits.remaining == 0 + && credits.codexCreditLimit == nil + && (result.usage.codexResetCredits?.availableInventory(at: result.usage.updatedAt).count ?? 0) == 0 + } +} + +#if DEBUG +extension CodexOAuthFetchStrategy { + static func _fetchResetCreditsForTesting( + context: ProviderFetchContext, + credentials: CodexOAuthCredentials, + fetcher: @escaping @Sendable (CodexOAuthCredentials) async throws + -> CodexRateLimitResetCreditsSnapshot) async throws -> CodexRateLimitResetCreditsSnapshot? + { + try await self.fetchResetCreditsIfRequested( + context: context, + credentials: credentials, + fetcher: fetcher) + } + + static func _shouldFetchResetCreditsForTesting(_ context: ProviderFetchContext) -> Bool { + self.shouldFetchResetCredits(context) + } +} +#endif + +#if DEBUG +extension CodexOAuthFetchStrategy { + static func _mapUsageForTesting(_ data: Data, credentials: CodexOAuthCredentials) throws -> UsageSnapshot? { + let usage = try JSONDecoder().decode(CodexUsageResponse.self, from: data) + return CodexReconciledState.fromOAuth(response: usage, credentials: credentials)?.toUsageSnapshot() + } + + static func _mapResultForTesting( + _ data: Data, + credentials: CodexOAuthCredentials, + resetCredits: CodexRateLimitResetCreditsSnapshot? = nil, + sourceMode: ProviderSourceMode = .oauth, + allowEmptyUsageForResetCreditEnrichment: Bool = false) throws -> ProviderFetchResult + { + let usageResponse = try JSONDecoder().decode(CodexUsageResponse.self, from: data) + _ = sourceMode + return try Self.makeResult( + usageResponse: usageResponse, + resetCredits: resetCredits, + credentials: credentials, + updatedAt: Date(), + allowEmptyUsageForResetCreditEnrichment: allowEmptyUsageForResetCreditEnrichment) + } + + static func _replaceWithCLIMonthlyLimitForTesting( + oauthResult: ProviderFetchResult, + context: ProviderFetchContext, + cliStrategy: any ProviderFetchStrategy) async throws -> ProviderFetchResult + { + try await self.replacingWithCLIMonthlyLimitIfAvailable( + oauthResult, + context: context, + cliStrategy: cliStrategy) + } + + static func _shouldTryCLIForMonthlyLimitForTesting(_ result: ProviderFetchResult) -> Bool { + self.shouldTryCLIForMonthlyLimit(result) + } +} + +extension CodexProviderDescriptor { + static func _noDataMessageForTesting(env: [String: String]) -> String { + self.noDataMessage(env: env) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Codex/CodexProviderSettingsBuilder.swift b/Sources/CodexBarCore/Providers/Codex/CodexProviderSettingsBuilder.swift new file mode 100644 index 0000000..62de155 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexProviderSettingsBuilder.swift @@ -0,0 +1,93 @@ +import Foundation + +public struct CodexProviderSettingsBuilderInput: Sendable { + public let usageDataSource: CodexUsageDataSource + public let cookieSource: ProviderCookieSource + public let manualCookieHeader: String? + public let reconciliationSnapshot: CodexAccountReconciliationSnapshot + public let resolvedActiveSource: CodexResolvedActiveSource + + public init( + usageDataSource: CodexUsageDataSource, + cookieSource: ProviderCookieSource, + manualCookieHeader: String?, + reconciliationSnapshot: CodexAccountReconciliationSnapshot, + resolvedActiveSource: CodexResolvedActiveSource) + { + self.usageDataSource = usageDataSource + self.cookieSource = cookieSource + self.manualCookieHeader = manualCookieHeader + self.reconciliationSnapshot = reconciliationSnapshot + self.resolvedActiveSource = resolvedActiveSource + } +} + +public enum CodexKnownOwnerCatalog { + public static func candidates( + from snapshot: CodexAccountReconciliationSnapshot) -> [CodexDashboardKnownOwnerCandidate] + { + var candidates = snapshot.storedAccounts.map { account in + CodexDashboardKnownOwnerCandidate( + identity: snapshot.runtimeIdentity(for: account), + normalizedEmail: CodexIdentityResolver.normalizeEmail(snapshot.runtimeEmail(for: account))) + } + + if let liveSystemAccount = snapshot.liveSystemAccount { + candidates.append(CodexDashboardKnownOwnerCandidate( + identity: snapshot.runtimeIdentity(for: liveSystemAccount), + normalizedEmail: CodexIdentityResolver.normalizeEmail(liveSystemAccount.email))) + } + + for profileAccount in snapshot.profileHomeAccounts { + candidates.append(CodexDashboardKnownOwnerCandidate( + identity: snapshot.runtimeIdentity(for: profileAccount), + normalizedEmail: CodexIdentityResolver.normalizeEmail(profileAccount.email), + sourceIsolationIdentifier: CookieHeaderCache.Scope.profileHome(profileAccount.codexHomePath) + .isolationIdentifier)) + } + + return candidates + } +} + +public enum CodexProviderSettingsBuilder { + public static func make(input: CodexProviderSettingsBuilderInput) -> ProviderSettingsSnapshot + .CodexProviderSettings { + let snapshot = input.reconciliationSnapshot + let persistedSource = input.resolvedActiveSource.persistedSource + let managedSourceSelected = switch persistedSource { + case .liveSystem: + false + case .managedAccount: + true + case .profileHome: + false + } + let openAIWebCacheScope: CookieHeaderCache.Scope? = switch input.resolvedActiveSource.resolvedSource { + case .liveSystem: + nil + case let .managedAccount(id): + .managedAccount(id) + case let .profileHome(path): + .profileHome(path) + } + let profileAccountTargetUnavailable = switch input.resolvedActiveSource.resolvedSource { + case .liveSystem, .managedAccount: + false + case let .profileHome(path): + snapshot.profileHomeAccount(path: path) == nil + } + + return ProviderSettingsSnapshot.CodexProviderSettings( + usageDataSource: input.usageDataSource, + cookieSource: input.cookieSource, + manualCookieHeader: input.manualCookieHeader, + managedAccountStoreUnreadable: managedSourceSelected && snapshot.hasUnreadableAddedAccountStore, + managedAccountTargetUnavailable: managedSourceSelected + && snapshot.hasUnreadableAddedAccountStore == false + && snapshot.activeStoredAccount == nil, + profileAccountTargetUnavailable: profileAccountTargetUnavailable, + openAIWebCacheScope: openAIWebCacheScope, + dashboardAuthorityKnownOwners: CodexKnownOwnerCatalog.candidates(from: snapshot)) + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexRateWindowNormalizer.swift b/Sources/CodexBarCore/Providers/Codex/CodexRateWindowNormalizer.swift new file mode 100644 index 0000000..018100e --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexRateWindowNormalizer.swift @@ -0,0 +1,66 @@ +import Foundation + +enum CodexRateWindowNormalizer { + static func normalize( + primary: RateWindow?, + secondary: RateWindow?) + -> (primary: RateWindow?, secondary: RateWindow?) + { + switch (primary, secondary) { + case let (.some(primaryWindow), .some(secondaryWindow)): + switch (self.role(for: primaryWindow), self.role(for: secondaryWindow)) { + case (.session, .weekly), (.session, .unknown), (.unknown, .weekly): + (primaryWindow, secondaryWindow) + case (.weekly, .session), (.weekly, .unknown): + (secondaryWindow, primaryWindow) + default: + (primaryWindow, secondaryWindow) + } + case let (.some(primaryWindow), .none): + switch role(for: primaryWindow) { + case .weekly: + (nil, primaryWindow) + case .session, .unknown: + (primaryWindow, nil) + } + case let (.none, .some(secondaryWindow)): + switch self.role(for: secondaryWindow) { + case .session, .unknown: + (secondaryWindow, nil) + case .weekly: + (nil, secondaryWindow) + } + case (.none, .none): + (nil, nil) + } + } + + private enum WindowRole { + case session + case weekly + case unknown + } + + private static func role(for window: RateWindow) -> WindowRole { + switch window.windowMinutes { + case 300: + .session + case 10080: + .weekly + default: + .unknown + } + } +} + +#if DEBUG +extension CodexRateWindowNormalizer { + static func _normalizeForTesting( + primary: RateWindow?, + secondary: RateWindow?) + -> (primary: RateWindow?, secondary: RateWindow?) + { + self.normalize(primary: primary, secondary: secondary) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Codex/CodexReconciledState.swift b/Sources/CodexBarCore/Providers/Codex/CodexReconciledState.swift new file mode 100644 index 0000000..83f8f28 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexReconciledState.swift @@ -0,0 +1,149 @@ +import Foundation + +public struct CodexReconciledState: Sendable { + public let session: RateWindow? + public let weekly: RateWindow? + /// Named model-specific limits (e.g. Codex Spark) surfaced through `UsageSnapshot.extraRateWindows`. + public let extraRateWindows: [NamedRateWindow] + public let identity: ProviderIdentitySnapshot? + public let updatedAt: Date + + public init( + session: RateWindow?, + weekly: RateWindow?, + extraRateWindows: [NamedRateWindow] = [], + identity: ProviderIdentitySnapshot?, + updatedAt: Date) + { + self.session = session + self.weekly = weekly + self.extraRateWindows = extraRateWindows + self.identity = identity + self.updatedAt = updatedAt + } + + public static func fromCLI( + primary: RateWindow?, + secondary: RateWindow?, + identity: ProviderIdentitySnapshot?, + updatedAt: Date = Date()) -> CodexReconciledState? + { + self.make(primary: primary, secondary: secondary, identity: identity, updatedAt: updatedAt) + } + + public static func fromOAuth( + response: CodexUsageResponse, + credentials: CodexOAuthCredentials, + updatedAt: Date = Date()) -> CodexReconciledState? + { + self.make( + primary: self.makeWindow(response.rateLimit?.primaryWindow), + secondary: self.makeWindow(response.rateLimit?.secondaryWindow), + extraRateWindows: CodexAdditionalRateLimitMapper.extraRateWindows( + from: response.additionalRateLimits, + now: updatedAt), + identity: self.oauthIdentity(response: response, credentials: credentials), + updatedAt: updatedAt) + } + + public static func fromAttachedDashboard( + snapshot: OpenAIDashboardSnapshot, + provider: UsageProvider = .codex, + accountEmail: String? = nil, + accountPlan: String? = nil) -> CodexReconciledState? + { + let resolvedEmail = accountEmail ?? snapshot.signedInEmail + let resolvedPlan = accountPlan ?? snapshot.accountPlan + let identity = ProviderIdentitySnapshot( + providerID: provider, + accountEmail: resolvedEmail, + accountOrganization: nil, + loginMethod: resolvedPlan) + + return self.make( + primary: snapshot.primaryLimit, + secondary: snapshot.secondaryLimit, + extraRateWindows: snapshot.extraRateWindows ?? [], + identity: identity, + updatedAt: snapshot.updatedAt) + } + + public func toUsageSnapshot() -> UsageSnapshot { + UsageSnapshot( + primary: self.session, + secondary: self.weekly, + tertiary: nil, + extraRateWindows: self.extraRateWindows.isEmpty ? nil : self.extraRateWindows, + updatedAt: self.updatedAt, + identity: self.identity) + } + + public static func oauthIdentity( + response: CodexUsageResponse, + credentials: CodexOAuthCredentials) -> ProviderIdentitySnapshot + { + ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: self.resolveAccountEmail(from: credentials), + accountOrganization: nil, + loginMethod: self.resolvePlan(response: response, credentials: credentials)) + } + + private static func make( + primary: RateWindow?, + secondary: RateWindow?, + extraRateWindows: [NamedRateWindow] = [], + identity: ProviderIdentitySnapshot?, + updatedAt: Date) -> CodexReconciledState? + { + let normalized = CodexRateWindowNormalizer.normalize(primary: primary, secondary: secondary) + // Extra windows are supplemental, so they never resurrect a snapshot on their own: keep the + // existing primary/weekly gate to preserve current behavior when only extra limits are present. + guard normalized.primary != nil || normalized.secondary != nil else { + return nil + } + + return CodexReconciledState( + session: normalized.primary, + weekly: normalized.secondary, + extraRateWindows: extraRateWindows, + identity: identity, + updatedAt: updatedAt) + } + + private static func makeWindow(_ window: CodexUsageResponse.WindowSnapshot?) -> RateWindow? { + guard let window else { return nil } + let resetDate = Date(timeIntervalSince1970: TimeInterval(window.resetAt)) + let resetDescription = UsageFormatter.resetDescription(from: resetDate) + return RateWindow( + usedPercent: Double(window.usedPercent), + windowMinutes: window.limitWindowSeconds / 60, + resetsAt: resetDate, + resetDescription: resetDescription) + } + + private static func resolveAccountEmail(from credentials: CodexOAuthCredentials) -> String? { + guard let idToken = credentials.idToken, + let payload = UsageFetcher.parseJWT(idToken) + else { + return nil + } + + let profileDict = payload["https://api.openai.com/profile"] as? [String: Any] + let email = (payload["email"] as? String) ?? (profileDict?["email"] as? String) + return email?.trimmingCharacters(in: .whitespacesAndNewlines) + } + + private static func resolvePlan(response: CodexUsageResponse, credentials: CodexOAuthCredentials) -> String? { + if let plan = response.planType?.rawValue, !plan.isEmpty { return plan } + guard let idToken = credentials.idToken, + let payload = UsageFetcher.parseJWT(idToken) + else { + return nil + } + + let authDict = payload["https://api.openai.com/auth"] as? [String: Any] + let plan = (authDict?["chatgpt_plan_type"] as? String) ?? (payload["chatgpt_plan_type"] as? String) + return plan?.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexSpendControlLimitMapping.swift b/Sources/CodexBarCore/Providers/Codex/CodexSpendControlLimitMapping.swift new file mode 100644 index 0000000..39415da --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexSpendControlLimitMapping.swift @@ -0,0 +1,25 @@ +import Foundation + +extension CodexUsageResponse.SpendControlLimitSnapshot { + func codexCreditLimitSnapshot(updatedAt: Date) -> CodexCreditLimitSnapshot? { + guard let limit, limit > 0 else { return nil } + let used: Double = if let used { + used + } else if let remainingPercent { + limit * max(0, min(100, 100 - remainingPercent)) / 100 + } else { + 0 + } + let remainingPercent = self.remainingPercent ?? max(0, min(100, 100 - (used / limit * 100))) + let resetsAt = self.resetsAt.flatMap { value -> Date? in + guard value > 0 else { return nil } + return Date(timeIntervalSince1970: TimeInterval(value)) + } + return CodexCreditLimitSnapshot( + used: used, + limit: limit, + remainingPercent: remainingPercent, + resetsAt: resetsAt, + updatedAt: updatedAt) + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexStatusProbe.swift b/Sources/CodexBarCore/Providers/Codex/CodexStatusProbe.swift new file mode 100644 index 0000000..a3cb8ee --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexStatusProbe.swift @@ -0,0 +1,294 @@ +import Foundation + +public struct CodexStatusSnapshot: Sendable { + public let credits: Double? + public let fiveHourPercentLeft: Int? + public let weeklyPercentLeft: Int? + public let fiveHourResetDescription: String? + public let weeklyResetDescription: String? + public let fiveHourResetsAt: Date? + public let weeklyResetsAt: Date? + public let codexCreditLimit: CodexCreditLimitSnapshot? + public let rawText: String + + public init( + credits: Double?, + fiveHourPercentLeft: Int?, + weeklyPercentLeft: Int?, + fiveHourResetDescription: String?, + weeklyResetDescription: String?, + fiveHourResetsAt: Date?, + weeklyResetsAt: Date?, + codexCreditLimit: CodexCreditLimitSnapshot? = nil, + rawText: String) + { + self.credits = credits + self.fiveHourPercentLeft = fiveHourPercentLeft + self.weeklyPercentLeft = weeklyPercentLeft + self.fiveHourResetDescription = fiveHourResetDescription + self.weeklyResetDescription = weeklyResetDescription + self.fiveHourResetsAt = fiveHourResetsAt + self.weeklyResetsAt = weeklyResetsAt + self.codexCreditLimit = codexCreditLimit + self.rawText = rawText + } +} + +public enum CodexStatusProbeError: LocalizedError, Sendable { + case codexNotInstalled + case launchBlocked(String) + case parseFailed(String) + case timedOut + case updateRequired(String) + + public var errorDescription: String? { + switch self { + case .codexNotInstalled: + "Codex CLI missing. Install via `npm i -g @openai/codex` (or bun install) and restart." + case let .launchBlocked(message): + message + case .parseFailed: + "Could not parse Codex status; will retry shortly." + case .timedOut: + "Codex status probe timed out." + case let .updateRequired(msg): + "Codex CLI update needed: \(msg)" + } + } +} + +/// Runs `codex` inside a PTY, sends `/status`, captures text, and parses credits/limits. +public struct CodexStatusProbe { + private static let defaultTimeoutSeconds: TimeInterval = 8.0 + private static let parseRetryTimeoutSeconds: TimeInterval = 4.0 + + public var codexBinary: String = "codex" + public var timeout: TimeInterval = Self.defaultTimeoutSeconds + public var keepCLISessionsAlive: Bool = false + public var environment: [String: String] = ProcessInfo.processInfo.environment + + public init() {} + + public init( + codexBinary: String = "codex", + timeout: TimeInterval = 8.0, + keepCLISessionsAlive: Bool = false, + environment: [String: String] = ProcessInfo.processInfo.environment) + { + self.codexBinary = codexBinary + self.timeout = timeout + self.keepCLISessionsAlive = keepCLISessionsAlive + self.environment = environment + } + + public func fetch() async throws -> CodexStatusSnapshot { + let env = self.environment + let resolved = BinaryLocator.resolveCodexBinary(env: env, loginPATH: LoginShellPathCache.shared.current) + ?? self.codexBinary + guard FileManager.default.isExecutableFile(atPath: resolved) || TTYCommandRunner.which(resolved) != nil else { + throw CodexStatusProbeError.codexNotInstalled + } + if let message = CodexCLILaunchGate.shared.backgroundSkipMessage(binary: resolved) { + throw CodexStatusProbeError.launchBlocked(message) + } + do { + return try await self.runAndParse(binary: resolved, rows: 60, cols: 200, timeout: self.timeout) + } catch let error as CodexStatusProbeError { + // Retry only parser-level flakes with a short second attempt. + switch error { + case .parseFailed: + return try await self.runAndParse( + binary: resolved, + rows: 70, + cols: 220, + timeout: Self.parseRetryTimeoutSeconds) + default: + throw error + } + } catch { + throw error + } + } + + // MARK: - Parsing + + public static func parse(text: String, now: Date = .init()) throws -> CodexStatusSnapshot { + let clean = TextParsing.stripANSICodes(text) + guard !clean.isEmpty else { throw CodexStatusProbeError.timedOut } + if clean.localizedCaseInsensitiveContains("data not available yet") { + throw CodexStatusProbeError.parseFailed("data not available yet") + } + if self.containsUpdatePrompt(clean) { + throw CodexStatusProbeError.updateRequired( + "Run `bun install -g @openai/codex` to continue (update prompt blocking /status).") + } + let credits = TextParsing.firstNumber(pattern: #"Credits:\s*([0-9][0-9.,]*)"#, text: clean) + // Pull reset info from the same lines that contain the percentages. + let fiveLine = TextParsing.firstLine(matching: #"5h limit[^\n]*"#, text: clean) + let weekLine = TextParsing.firstLine(matching: #"Weekly limit[^\n]*"#, text: clean) + let fivePct = fiveLine.flatMap(TextParsing.percentLeft(fromLine:)) + let weekPct = weekLine.flatMap(TextParsing.percentLeft(fromLine:)) + let fiveReset = fiveLine.flatMap(TextParsing.resetString(fromLine:)) + let weekReset = weekLine.flatMap(TextParsing.resetString(fromLine:)) + let monthlyLimit = self.parseMonthlyCreditLimit(text: clean, now: now) + if credits == nil, fivePct == nil, weekPct == nil, monthlyLimit == nil { + throw CodexStatusProbeError.parseFailed(clean.prefix(400).description) + } + return CodexStatusSnapshot( + credits: credits, + fiveHourPercentLeft: fivePct, + weeklyPercentLeft: weekPct, + fiveHourResetDescription: fiveReset, + weeklyResetDescription: weekReset, + fiveHourResetsAt: self.parseResetDate(from: fiveReset, now: now), + weeklyResetsAt: self.parseResetDate(from: weekReset, now: now), + codexCreditLimit: monthlyLimit, + rawText: clean) + } + + private static func parseMonthlyCreditLimit(text: String, now: Date) -> CodexCreditLimitSnapshot? { + guard let monthlyLine = TextParsing.firstLine(matching: #"Monthly credit limit[^\n]*"#, text: text) else { + return nil + } + guard let limit = TextParsing.firstNumber( + pattern: #"of\s+([0-9][0-9., ]*)\s+credits used"#, + text: text), + limit > 0 + else { + return nil + } + let used = TextParsing.firstNumber( + pattern: #"([0-9][0-9., ]*)\s+of\s+[0-9][0-9., ]*\s+credits used"#, + text: text) ?? 0 + let remainingPercent = TextParsing.percentLeft(fromLine: monthlyLine) + .map(Double.init) + ?? max(0, min(100, 100 - (used / limit * 100))) + let resetText = TextParsing.resetString(fromLine: monthlyLine) + return CodexCreditLimitSnapshot( + used: used, + limit: limit, + remainingPercent: remainingPercent, + resetsAt: self.parseResetDate(from: resetText, now: now), + updatedAt: now) + } + + private static func parseResetDate(from text: String?, now: Date) -> Date? { + guard var raw = text?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { return nil } + raw = raw.trimmingCharacters(in: CharacterSet(charactersIn: "()")) + raw = raw.replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespacesAndNewlines) + + let calendar = Calendar(identifier: .gregorian) + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = TimeZone.current + formatter.defaultDate = now + + if let match = raw.firstMatch(of: /^([0-9]{1,2}:[0-9]{2}) on ([0-9]{1,2} [A-Za-z]{3})$/) { + raw = "\(match.output.2) \(match.output.1)" + formatter.dateFormat = "d MMM HH:mm" + if let date = formatter.date(from: raw) { + return self.bumpYearIfNeeded(date, now: now, calendar: calendar) + } + } + + if let match = raw.firstMatch(of: /^([0-9]{1,2}:[0-9]{2}) on ([A-Za-z]{3} [0-9]{1,2})$/) { + raw = "\(match.output.2) \(match.output.1)" + formatter.dateFormat = "MMM d HH:mm" + if let date = formatter.date(from: raw) { + return self.bumpYearIfNeeded(date, now: now, calendar: calendar) + } + } + + for format in ["HH:mm", "H:mm"] { + formatter.dateFormat = format + if let time = formatter.date(from: raw) { + let components = calendar.dateComponents([.hour, .minute], from: time) + guard let anchored = calendar.date( + bySettingHour: components.hour ?? 0, + minute: components.minute ?? 0, + second: 0, + of: now) + else { + return nil + } + if anchored >= now { + return anchored + } + return calendar.date(byAdding: .day, value: 1, to: anchored) + } + } + + return nil + } + + private static func bumpYearIfNeeded(_ date: Date, now: Date, calendar: Calendar) -> Date? { + if date >= now { + return date + } + return calendar.date(byAdding: .year, value: 1, to: date) + } + + private func runAndParse( + binary: String, + rows: UInt16, + cols: UInt16, + timeout: TimeInterval) async throws -> CodexStatusSnapshot + { + let stateHome = try CodexStatusProbeIsolation.supportDirectory(environment: self.environment) + let extraArgs = CodexStatusProbeIsolation.codexArguments(stateHome: stateHome) + let workingDirectory = CodexStatusProbeIsolation.workingDirectory(environment: self.environment) + let text: String + if self.keepCLISessionsAlive { + do { + text = try await CodexCLISession.shared.captureStatus( + binary: binary, + options: .init( + timeout: timeout, + rows: rows, + cols: cols, + environment: self.environment, + extraArgs: extraArgs, + workingDirectory: workingDirectory)) + } catch CodexCLISession.SessionError.processExited { + throw CodexStatusProbeError.timedOut + } catch CodexCLISession.SessionError.timedOut { + throw CodexStatusProbeError.timedOut + } catch let CodexCLISession.SessionError.launchFailed(message) { + if let throttled = CodexCLILaunchGate.shared.recordLaunchFailure(binary: binary, message: message) { + throw CodexStatusProbeError.launchBlocked(throttled) + } + throw CodexStatusProbeError.codexNotInstalled + } + } else { + let runner = TTYCommandRunner() + let script = "/status" + let result: TTYCommandRunner.Result + do { + result = try runner.run( + binary: binary, + send: script, + options: .init( + rows: rows, + cols: cols, + timeout: timeout, + workingDirectory: workingDirectory, + extraArgs: extraArgs, + baseEnvironment: self.environment, + forceCodexStatusMode: true)) + } catch let TTYCommandRunner.Error.launchFailed(message) { + if let throttled = CodexCLILaunchGate.shared.recordLaunchFailure(binary: binary, message: message) { + throw CodexStatusProbeError.launchBlocked(throttled) + } + throw CodexStatusProbeError.codexNotInstalled + } + text = result.text + } + return try Self.parse(text: text) + } + + private static func containsUpdatePrompt(_ text: String) -> Bool { + let lower = text.lowercased() + return lower.contains("update available") && lower.contains("codex") + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexStatusProbeIsolation.swift b/Sources/CodexBarCore/Providers/Codex/CodexStatusProbeIsolation.swift new file mode 100644 index 0000000..2401122 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexStatusProbeIsolation.swift @@ -0,0 +1,46 @@ +import Foundation + +enum CodexStatusProbeIsolation { + static func supportDirectory(environment: [String: String]) throws -> URL { + let baseURL: URL = if let tmp = environment["TMPDIR"]?.trimmingCharacters(in: .whitespacesAndNewlines), + !tmp.isEmpty + { + URL(fileURLWithPath: tmp, isDirectory: true) + } else { + FileManager.default.temporaryDirectory + } + + let directory = baseURL + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("CodexStatusProbe", isDirectory: true) + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + return directory + } + + static func workingDirectory(environment: [String: String]) -> URL? { + let home = environment["HOME"]?.trimmingCharacters(in: .whitespacesAndNewlines) + guard let home, !home.isEmpty else { return nil } + return URL(fileURLWithPath: home, isDirectory: true) + } + + static func codexArguments(stateHome: URL) -> [String] { + [ + "-s", + "read-only", + "-a", + "untrusted", + "-c", + "history.persistence=\"none\"", + "-c", + "experimental_thread_store={type=\"in_memory\",id=\"codexbar-status\"}", + "-c", + "sqlite_home=\"\(self.tomlEscaped(stateHome.path))\"", + ] + } + + private static func tomlEscaped(_ value: String) -> String { + value + .replacingOccurrences(of: "\\", with: "\\\\") + .replacingOccurrences(of: "\"", with: "\\\"") + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexSystemAccountObserver.swift b/Sources/CodexBarCore/Providers/Codex/CodexSystemAccountObserver.swift new file mode 100644 index 0000000..5b66209 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexSystemAccountObserver.swift @@ -0,0 +1,69 @@ +import Foundation + +public struct ObservedSystemCodexAccount: Equatable, Sendable { + public let email: String + public let workspaceLabel: String? + public let workspaceAccountID: String? + public let authFingerprint: String? + public let codexHomePath: String + public let observedAt: Date + public let identity: CodexIdentity + + public init( + email: String, + workspaceLabel: String? = nil, + workspaceAccountID: String? = nil, + authFingerprint: String? = nil, + codexHomePath: String, + observedAt: Date, + identity: CodexIdentity = .unresolved) + { + self.email = email + self.workspaceLabel = workspaceLabel + self.workspaceAccountID = workspaceAccountID + self.authFingerprint = CodexAuthFingerprint.normalize(authFingerprint) + self.codexHomePath = codexHomePath + self.observedAt = observedAt + self.identity = identity + } +} + +public protocol CodexSystemAccountObserving: Sendable { + func loadSystemAccount(environment: [String: String]) throws -> ObservedSystemCodexAccount? +} + +public struct DefaultCodexSystemAccountObserver: CodexSystemAccountObserving { + private let workspaceCache: CodexOpenAIWorkspaceIdentityCache + + public init(workspaceCache: CodexOpenAIWorkspaceIdentityCache = CodexOpenAIWorkspaceIdentityCache()) { + self.workspaceCache = workspaceCache + } + + public func loadSystemAccount(environment: [String: String]) throws -> ObservedSystemCodexAccount? { + let homeURL = CodexHomeScope.ambientHomeURL(env: environment) + let fetcher = UsageFetcher(environment: environment) + let account = fetcher.loadAuthBackedCodexAccount() + + guard let rawEmail = account.email?.trimmingCharacters(in: .whitespacesAndNewlines), + !rawEmail.isEmpty + else { + return nil + } + + let providerAccountID: String? = switch account.identity { + case let .providerAccount(id): + ManagedCodexAccount.normalizeProviderAccountID(id) + case .emailOnly, .unresolved: + nil + } + + return ObservedSystemCodexAccount( + email: rawEmail.lowercased(), + workspaceLabel: self.workspaceCache.workspaceLabel(for: providerAccountID), + workspaceAccountID: providerAccountID, + authFingerprint: CodexAuthFingerprint.fingerprint(homePath: homeURL.path), + codexHomePath: homeURL.path, + observedAt: Date(), + identity: account.identity) + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexUsageDataSource.swift b/Sources/CodexBarCore/Providers/Codex/CodexUsageDataSource.swift new file mode 100644 index 0000000..b3bb16c --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexUsageDataSource.swift @@ -0,0 +1,30 @@ +import Foundation + +public enum CodexUsageDataSource: String, CaseIterable, Identifiable, Sendable { + case auto + case oauth + case cli + + public var id: String { + self.rawValue + } + + public var displayName: String { + switch self { + case .auto: "Auto" + case .oauth: "OAuth API" + case .cli: "CLI (RPC/PTY)" + } + } + + public var sourceLabel: String { + switch self { + case .auto: + "auto" + case .oauth: + "oauth" + case .cli: + "cli" + } + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexVisibleAccountProjection.swift b/Sources/CodexBarCore/Providers/Codex/CodexVisibleAccountProjection.swift new file mode 100644 index 0000000..a25f3fe --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexVisibleAccountProjection.swift @@ -0,0 +1,287 @@ +import Foundation + +public struct CodexVisibleAccount: Equatable, Identifiable, Sendable { + public let id: String + public let email: String + public let workspaceLabel: String? + public let workspaceAccountID: String? + public let authFingerprint: String? + public let storedAccountID: UUID? + public let selectionSource: CodexActiveSource + public let isActive: Bool + public let isLive: Bool + public let canReauthenticate: Bool + public let canRemove: Bool + + public init( + id: String, + email: String, + workspaceLabel: String? = nil, + workspaceAccountID: String? = nil, + authFingerprint: String? = nil, + storedAccountID: UUID?, + selectionSource: CodexActiveSource, + isActive: Bool, + isLive: Bool, + canReauthenticate: Bool, + canRemove: Bool) + { + self.id = id + self.email = email + self.workspaceLabel = Self.normalizeWorkspaceLabel(workspaceLabel) + self.workspaceAccountID = workspaceAccountID + self.authFingerprint = CodexAuthFingerprint.normalize(authFingerprint) + self.storedAccountID = storedAccountID + self.selectionSource = selectionSource + self.isActive = isActive + self.isLive = isLive + self.canReauthenticate = canReauthenticate + self.canRemove = canRemove + } + + public var displayName: String { + guard let workspaceLabel else { return self.email } + return "\(self.email) — \(workspaceLabel)" + } + + public var menuDisplayName: String { + guard let menuWorkspaceLabel else { return self.email } + return "\(self.email) — \(menuWorkspaceLabel)" + } + + public var menuWorkspaceLabel: String? { + guard let workspaceLabel, workspaceLabel.compare("Personal", options: [.caseInsensitive]) != .orderedSame else { + return nil + } + return workspaceLabel + } + + public var authenticationHealthLabel: String? { + guard !self.isLive, self.storedAccountID != nil, self.authFingerprint == nil else { return nil } + return "Missing auth" + } + + private static func normalizeWorkspaceLabel(_ workspaceLabel: String?) -> String? { + guard let trimmed = workspaceLabel?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed + } +} + +public struct CodexVisibleAccountProjection: Equatable, Sendable { + public let visibleAccounts: [CodexVisibleAccount] + public let activeVisibleAccountID: String? + public let liveVisibleAccountID: String? + public let hasUnreadableAddedAccountStore: Bool + + public init( + visibleAccounts: [CodexVisibleAccount], + activeVisibleAccountID: String?, + liveVisibleAccountID: String?, + hasUnreadableAddedAccountStore: Bool) + { + self.visibleAccounts = visibleAccounts + self.activeVisibleAccountID = activeVisibleAccountID + self.liveVisibleAccountID = liveVisibleAccountID + self.hasUnreadableAddedAccountStore = hasUnreadableAddedAccountStore + } + + public func source(forVisibleAccountID id: String) -> CodexActiveSource? { + self.visibleAccounts.first { $0.id == id }?.selectionSource + } +} + +extension DefaultCodexAccountReconciler { + public func loadVisibleAccounts() -> CodexVisibleAccountProjection { + CodexVisibleAccountProjection.make(from: self.loadSnapshot()) + } +} + +extension CodexVisibleAccountProjection { + public static func make(from snapshot: CodexAccountReconciliationSnapshot) -> CodexVisibleAccountProjection { + let resolvedActiveSource = CodexActiveSourceResolver.resolve(from: snapshot).resolvedSource + var drafts: [VisibleAccountDraft] = [] + + for storedAccount in snapshot.storedAccounts { + let normalizedEmail = snapshot.runtimeEmail(for: storedAccount) + let runtimeIdentity = snapshot.runtimeIdentity(for: storedAccount) + let runtimeWorkspaceAccountID: String? = switch runtimeIdentity { + case let .providerAccount(id): + ManagedCodexAccount.normalizeWorkspaceAccountID(id) + case .emailOnly, .unresolved: + nil + } + drafts.append(VisibleAccountDraft( + email: normalizedEmail, + workspaceLabel: Self.normalizeWorkspaceLabel(storedAccount.workspaceLabel), + workspaceAccountID: storedAccount.workspaceAccountID ?? runtimeWorkspaceAccountID, + authFingerprint: storedAccount.authFingerprint, + storedAccountID: storedAccount.id, + selectionSource: .managedAccount(id: storedAccount.id), + isLive: false, + canReauthenticate: true, + canRemove: true, + identity: runtimeIdentity)) + } + + if let liveSystemAccount = snapshot.liveSystemAccount { + let normalizedEmail = Self.normalizeVisibleEmail(liveSystemAccount.email) + let liveIdentity = snapshot.runtimeIdentity(for: liveSystemAccount) + if let exactStoredAccountID = snapshot.matchingStoredAccountForLiveSystemAccount?.id, + let exactIndex = drafts.firstIndex(where: { $0.storedAccountID == exactStoredAccountID }) + { + let existingDraft = drafts[exactIndex] + let liveWorkspaceLabel = Self.normalizeWorkspaceLabel(liveSystemAccount.workspaceLabel) + drafts[exactIndex] = VisibleAccountDraft( + email: existingDraft.email, + workspaceLabel: liveWorkspaceLabel ?? existingDraft.workspaceLabel, + workspaceAccountID: liveSystemAccount.workspaceAccountID ?? existingDraft.workspaceAccountID, + authFingerprint: liveSystemAccount.authFingerprint ?? existingDraft.authFingerprint, + storedAccountID: existingDraft.storedAccountID, + selectionSource: .liveSystem, + isLive: true, + canReauthenticate: existingDraft.canReauthenticate, + canRemove: existingDraft.canRemove, + identity: liveIdentity) + } else if let existingIndex = drafts.firstIndex(where: { draft in + CodexIdentityMatcher.matches( + draft.identity, + lhsEmail: draft.email, + liveIdentity, + rhsEmail: normalizedEmail) + }) { + let existingDraft = drafts[existingIndex] + let liveWorkspaceLabel = Self.normalizeWorkspaceLabel(liveSystemAccount.workspaceLabel) + drafts[existingIndex] = VisibleAccountDraft( + email: existingDraft.email, + workspaceLabel: liveWorkspaceLabel ?? existingDraft.workspaceLabel, + workspaceAccountID: liveSystemAccount.workspaceAccountID ?? existingDraft.workspaceAccountID, + authFingerprint: liveSystemAccount.authFingerprint ?? existingDraft.authFingerprint, + storedAccountID: existingDraft.storedAccountID, + selectionSource: .liveSystem, + isLive: true, + canReauthenticate: existingDraft.canReauthenticate, + canRemove: existingDraft.canRemove, + identity: liveIdentity) + } else { + drafts.append(VisibleAccountDraft( + email: normalizedEmail, + workspaceLabel: Self.normalizeWorkspaceLabel(liveSystemAccount.workspaceLabel), + workspaceAccountID: liveSystemAccount.workspaceAccountID, + authFingerprint: liveSystemAccount.authFingerprint, + storedAccountID: nil, + selectionSource: .liveSystem, + isLive: true, + canReauthenticate: true, + canRemove: false, + identity: liveIdentity)) + } + } + + let livePath = snapshot.liveSystemAccount.flatMap { CodexHomeScope.normalizedHomePath($0.codexHomePath) } + let managedPaths = Set(snapshot.storedAccounts.compactMap { + CodexHomeScope.normalizedHomePath($0.managedHomePath) + }) + for profileAccount in snapshot.profileHomeAccounts { + let profilePath = CodexHomeScope.normalizedHomePath(profileAccount.codexHomePath) + guard let profilePath, + profilePath != livePath, + !managedPaths.contains(profilePath) + else { + continue + } + drafts.append(VisibleAccountDraft( + email: Self.normalizeVisibleEmail(profileAccount.email), + workspaceLabel: Self.normalizeWorkspaceLabel(profileAccount.workspaceLabel), + workspaceAccountID: profileAccount.workspaceAccountID, + authFingerprint: profileAccount.authFingerprint, + storedAccountID: nil, + selectionSource: .profileHome(path: profilePath), + isLive: false, + canReauthenticate: false, + canRemove: false, + identity: snapshot.runtimeIdentity(for: profileAccount))) + } + + let groupedByEmail = Dictionary(grouping: drafts.indices, by: { drafts[$0].email }) + let visibleAccounts = drafts.map { draft in + let id = Self.visibleAccountID(for: draft, emailGroupSize: groupedByEmail[draft.email]?.count ?? 0) + let isActive = switch resolvedActiveSource { + case .liveSystem: + draft.selectionSource == .liveSystem + case let .managedAccount(id): + draft.selectionSource == .managedAccount(id: id) + case let .profileHome(path): + draft.selectionSource == .profileHome(path: path) + } + + return CodexVisibleAccount( + id: id, + email: draft.email, + workspaceLabel: draft.workspaceLabel, + workspaceAccountID: draft.workspaceAccountID, + authFingerprint: draft.authFingerprint, + storedAccountID: draft.storedAccountID, + selectionSource: draft.selectionSource, + isActive: isActive, + isLive: draft.isLive, + canReauthenticate: draft.canReauthenticate, + canRemove: draft.canRemove) + }.sorted { lhs, rhs in + if lhs.email != rhs.email { + return lhs.email < rhs.email + } + if lhs.isLive != rhs.isLive { + return lhs.isLive && !rhs.isLive + } + if lhs.displayName != rhs.displayName { + return lhs.displayName < rhs.displayName + } + return lhs.id < rhs.id + } + + return CodexVisibleAccountProjection( + visibleAccounts: visibleAccounts, + activeVisibleAccountID: visibleAccounts.first { $0.isActive }?.id, + liveVisibleAccountID: visibleAccounts.first { $0.isLive }?.id, + hasUnreadableAddedAccountStore: snapshot.hasUnreadableAddedAccountStore) + } + + private static func normalizeVisibleEmail(_ email: String) -> String { + email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } + + private static func normalizeWorkspaceLabel(_ workspaceLabel: String?) -> String? { + guard let trimmed = workspaceLabel?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed + } + + private static func visibleAccountID(for draft: VisibleAccountDraft, emailGroupSize: Int) -> String { + guard emailGroupSize > 1 else { return draft.email } + + switch draft.selectionSource { + case .liveSystem: + return "live:\(CodexIdentityMatcher.selectionKey(for: draft.identity, fallbackEmail: draft.email))" + case let .managedAccount(id): + return "managed:\(id.uuidString.lowercased())" + case let .profileHome(path): + return "profile:\(path)" + } + } +} + +private struct VisibleAccountDraft { + let email: String + let workspaceLabel: String? + let workspaceAccountID: String? + let authFingerprint: String? + let storedAccountID: UUID? + let selectionSource: CodexActiveSource + let isLive: Bool + let canReauthenticate: Bool + let canRemove: Bool + let identity: CodexIdentity +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift b/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift new file mode 100644 index 0000000..12eedd4 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift @@ -0,0 +1,343 @@ +#if os(macOS) +import AppKit +import Foundation + +public struct CodexWebDashboardStrategy: ProviderFetchStrategy { + public let id: String = "codex.web.dashboard" + public let kind: ProviderFetchKind = .webDashboard + + public init() {} + + public func isAvailable(_ context: ProviderFetchContext) async -> Bool { + context.sourceMode.usesWeb && + !Self.managedAccountStoreIsUnreadable(context) && + !Self.managedAccountTargetIsUnavailable(context) && + (!Self.profileAccountTargetIsUnavailable(context) || context.sourceMode == .web) + } + + public func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard !Self.managedAccountStoreIsUnreadable(context) else { + // A fail-closed placeholder CODEX_HOME does not identify a target account. If the managed store + // itself is unreadable, web import must not fall back to "any signed-in browser account". + throw OpenAIDashboardFetcher.FetchError.loginRequired + } + guard !Self.managedAccountTargetIsUnavailable(context) else { + // If the selected managed account no longer exists in a readable store, web import must not + // fall back to "any signed-in browser account" for that stale selection. + throw OpenAIDashboardFetcher.FetchError.loginRequired + } + guard !Self.profileAccountTargetIsUnavailable(context) else { + // A profile without an auth-backed email cannot safely target browser cookie import. + throw OpenAIDashboardFetcher.FetchError.loginRequired + } + + let options = OpenAIWebOptions( + deadline: OpenAIDashboardFetcher.deadline(startingAt: Date(), timeout: context.webTimeout), + debugDumpHTML: context.webDebugDumpHTML, + verbose: context.verbose) + + // Ensure AppKit is initialized before using WebKit in a CLI. + await MainActor.run { + _ = NSApplication.shared + } + + let result: OpenAIWebCodexResult + do { + result = try await Self.fetchOpenAIWebCodex( + context: context, + options: options, + browserDetection: context.browserDetection) + } catch let error as URLError where error.code == .timedOut { + throw OpenAIWebCodexError.timedOut(seconds: context.webTimeout) + } + return self.makeResult( + usage: result.usage, + credits: result.credits, + dashboard: result.dashboard, + sourceLabel: "openai-web") + } + + public func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + _ = error + return context.sourceMode == .auto + } + + private static func managedAccountStoreIsUnreadable(_ context: ProviderFetchContext) -> Bool { + context.settings?.codex?.managedAccountStoreUnreadable == true + } + + private static func managedAccountTargetIsUnavailable(_ context: ProviderFetchContext) -> Bool { + context.settings?.codex?.managedAccountTargetUnavailable == true + } + + private static func profileAccountTargetIsUnavailable(_ context: ProviderFetchContext) -> Bool { + context.settings?.codex?.profileAccountTargetUnavailable == true + } +} + +struct OpenAIWebCodexResult { + let usage: UsageSnapshot + let credits: CreditsSnapshot? + let dashboard: OpenAIDashboardSnapshot +} + +enum OpenAIWebCodexError: LocalizedError, Equatable { + case missingUsage + case policyRejected(CodexDashboardAuthorityDecision) + case timedOut(seconds: TimeInterval) + + var errorDescription: String? { + switch self { + case .missingUsage: + return "OpenAI web dashboard did not include usage limits." + case let .timedOut(seconds): + return "OpenAI web dashboard fetch timed out after \(seconds.formatted()) seconds." + case let .policyRejected(decision): + switch decision.reason { + case let .wrongEmail(expected, actual): + var details: [String] = [] + if let expected { + details.append("expected \(expected)") + } + if let actual { + details.append("got \(actual)") + } + if details.isEmpty { + return "OpenAI web dashboard belonged to the wrong account." + } + return "OpenAI web dashboard belonged to the wrong account (\(details.joined(separator: ", ")))." + case .unresolvedWithoutTrustedEvidence: + return "Active Codex identity is unresolved and no trusted auth-backed continuity exists." + case .providerAccountMissingScopedEmail: + return "Active Codex provider account is missing its scoped email, " + + "so dashboard ownership cannot be proven." + case .providerAccountLacksExactOwnershipProof: + return "OpenAI web dashboard could not be proven to belong to the active provider account." + case .missingDashboardSignedInEmail: + return "OpenAI web dashboard did not expose a signed-in email." + case let .sameEmailAmbiguity(email): + return "OpenAI web dashboard email \(email) is ambiguous across multiple known owners." + default: + return "OpenAI web dashboard was rejected by Codex dashboard authority." + } + } + } +} + +private struct OpenAIWebOptions { + let deadline: Date + let debugDumpHTML: Bool + let verbose: Bool +} + +@MainActor +private final class WebLogBuffer { + private var lines: [String] = [] + private let maxCount: Int + private let verbose: Bool + private let logger = CodexBarLog.logger(LogCategories.openAIWeb) + + init(maxCount: Int = 300, verbose: Bool) { + self.maxCount = maxCount + self.verbose = verbose + } + + func append(_ line: String) { + self.lines.append(line) + if self.lines.count > self.maxCount { + self.lines.removeFirst(self.lines.count - self.maxCount) + } + if self.verbose { + self.logger.verbose(line) + } + } + + func snapshot() -> [String] { + self.lines + } +} + +extension CodexWebDashboardStrategy { + @MainActor + fileprivate static func fetchOpenAIWebCodex( + context: ProviderFetchContext, + options: OpenAIWebOptions, + browserDetection: BrowserDetection) async throws -> OpenAIWebCodexResult + { + let logger = WebLogBuffer(verbose: options.verbose) + let log: @MainActor (String) -> Void = { line in + logger.append(line) + } + do { + let result = try await Self.fetchOpenAIWebDashboard( + context: context, + options: options, + browserDetection: browserDetection, + preferCachedCookieHeader: true, + logger: log) + return try Self.makeAuthorizedDashboardResult( + dashboard: result.dashboard, + context: context, + routingTargetEmail: result.routingTargetEmail) + } catch { + guard Self.shouldRetryWithFreshBrowserImport(after: error) else { + throw error + } + _ = try OpenAIDashboardBrowserCookieImporter.remainingTimeout(until: options.deadline) + log("Retrying OpenAI web dashboard with a fresh browser cookie import.") + let result = try await Self.fetchOpenAIWebDashboard( + context: context, + options: options, + browserDetection: browserDetection, + preferCachedCookieHeader: false, + logger: log) + return try Self.makeAuthorizedDashboardResult( + dashboard: result.dashboard, + context: context, + routingTargetEmail: result.routingTargetEmail) + } + } + + nonisolated static func shouldRetryWithFreshBrowserImport(after error: Error) -> Bool { + if error is OpenAIWebCodexError { + return error as? OpenAIWebCodexError == .missingUsage + } + if case OpenAIDashboardFetcher.FetchError.noDashboardData = error { + return true + } + return false + } + + @MainActor + static func makeAuthorizedDashboardResultForTesting( + dashboard: OpenAIDashboardSnapshot, + context: ProviderFetchContext, + routingTargetEmail: String?) + throws -> OpenAIWebCodexResult + { + try self.makeAuthorizedDashboardResult( + dashboard: dashboard, + context: context, + routingTargetEmail: routingTargetEmail) + } + + @MainActor + private static func makeAuthorizedDashboardResult( + dashboard: OpenAIDashboardSnapshot, + context: ProviderFetchContext, + routingTargetEmail: String?) throws -> OpenAIWebCodexResult + { + let input = CodexCLIDashboardAuthorityContext.makeLiveWebInput( + dashboard: dashboard, + context: context, + routingTargetEmail: routingTargetEmail) + let decision = CodexDashboardAuthority.evaluate(input) + + switch decision.disposition { + case .attach: + let attachedAccountEmail = CodexCLIDashboardAuthorityContext.attachmentEmail(from: input) + let credits = dashboard.toCreditsSnapshot() + let usage = dashboard.toUsageSnapshot(provider: .codex, accountEmail: attachedAccountEmail) + ?? Self.makeCreditsOnlyUsageSnapshot( + dashboard: dashboard, + attachedAccountEmail: attachedAccountEmail, + credits: credits) + guard let usage else { + throw OpenAIWebCodexError.missingUsage + } + if let attachedAccountEmail { + OpenAIDashboardCacheStore.save(OpenAIDashboardCache( + accountEmail: attachedAccountEmail, + snapshot: dashboard)) + } + return OpenAIWebCodexResult(usage: usage, credits: credits, dashboard: dashboard) + case .displayOnly: + if decision.cleanup.contains(.dashboardCache) { + OpenAIDashboardCacheStore.clear() + } + throw CodexDashboardPolicyError.displayOnly(decision) + case .failClosed: + if decision.cleanup.contains(.dashboardCache) { + OpenAIDashboardCacheStore.clear() + } + throw OpenAIWebCodexError.policyRejected(decision) + } + } + + private static func makeCreditsOnlyUsageSnapshot( + dashboard: OpenAIDashboardSnapshot, + attachedAccountEmail: String?, + credits: CreditsSnapshot?) -> UsageSnapshot? + { + guard credits != nil else { return nil } + return UsageSnapshot( + primary: nil, + secondary: nil, + tertiary: nil, + updatedAt: dashboard.updatedAt, + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: attachedAccountEmail ?? dashboard.signedInEmail, + accountOrganization: nil, + loginMethod: dashboard.accountPlan)) + } + + private struct OpenAIWebDashboardFetchResult { + let dashboard: OpenAIDashboardSnapshot + let routingTargetEmail: String? + } + + @MainActor + private static func fetchOpenAIWebDashboard( + context: ProviderFetchContext, + options: OpenAIWebOptions, + browserDetection: BrowserDetection, + preferCachedCookieHeader: Bool, + logger: @MainActor @escaping (String) -> Void) async throws -> OpenAIWebDashboardFetchResult + { + let auth = context.fetcher.loadAuthBackedCodexAccount() + let routingTargetEmail = auth.email?.trimmingCharacters(in: .whitespacesAndNewlines) + let allowAnyAccount = routingTargetEmail == nil + + let importResult = try await OpenAIDashboardBrowserCookieImporter(browserDetection: browserDetection) + .importBestCookies( + intoAccountEmail: routingTargetEmail, + allowAnyAccount: allowAnyAccount, + preferCachedCookieHeader: preferCachedCookieHeader, + cacheScope: context.settings?.codex?.openAIWebCacheScope, + deadline: options.deadline, + logger: logger) + let effectiveEmail = routingTargetEmail ?? importResult.signedInEmail? + .trimmingCharacters(in: .whitespacesAndNewlines) + + let dashboard = try await OpenAIDashboardFetcher().loadLatestDashboard( + accountEmail: effectiveEmail, + cacheScope: context.settings?.codex?.openAIWebCacheScope, + logger: logger, + debugDumpHTML: options.debugDumpHTML, + timeout: OpenAIDashboardBrowserCookieImporter.remainingTimeout(until: options.deadline)) + return OpenAIWebDashboardFetchResult( + dashboard: dashboard, + routingTargetEmail: routingTargetEmail) + } +} +#else +public struct CodexWebDashboardStrategy: ProviderFetchStrategy { + public let id: String = "codex.web.dashboard" + public let kind: ProviderFetchKind = .webDashboard + + public init() {} + + public func isAvailable(_: ProviderFetchContext) async -> Bool { + false + } + + public func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { + throw ProviderFetchError.noAvailableStrategy(.codex) + } + + public func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/CommandCode/CommandCodeCookieHeader.swift b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeCookieHeader.swift new file mode 100644 index 0000000..f5835c8 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeCookieHeader.swift @@ -0,0 +1,85 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +/// A resolved CommandCode session cookie ready to be sent on `Cookie:` headers. +/// +/// CommandCode's API (api.commandcode.ai) authenticates with the better-auth session +/// cookie set by commandcode.ai. better-auth emits either `better-auth.session_token` +/// or `__Secure-better-auth.session_token` depending on whether `useSecureCookies` is +/// enabled (the `__Secure-` variant is required by browsers for HTTPS production). +public struct CommandCodeCookieOverride: Sendable, Equatable { + public let name: String + public let token: String + + public init(name: String, token: String) { + self.name = name + self.token = token + } + + /// `Cookie: name=value` header value. + public var headerValue: String { + "\(self.name)=\(self.token)" + } +} + +public enum CommandCodeCookieHeader { + /// Cookie names used by better-auth in production (HTTPS) and dev (HTTP). + /// The `__Secure-` variant is the standard production deployment. + public static let supportedSessionCookieNames = [ + "__Host-better-auth.session_token", + "__Secure-better-auth.session_token", + "better-auth.session_token", + ] + + /// Extract a session cookie from a list of `HTTPCookie` records. + public static func sessionCookie(from cookies: [HTTPCookie]) -> CommandCodeCookieOverride? { + let pairs = cookies.map { (name: $0.name, value: $0.value) } + return self.extractSessionCookie(from: pairs) + } + + /// Parse a raw `Cookie:` header (or bare token) and extract the session value. + public static func override(from raw: String?) -> CommandCodeCookieOverride? { + guard let raw = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + + // Bare token — assume the production cookie name. + if !raw.contains("="), !raw.contains(";") { + return CommandCodeCookieOverride( + name: "__Secure-better-auth.session_token", + token: raw) + } + + return self.extractSessionCookie(fromHeader: raw) + } + + private static func extractSessionCookie(fromHeader header: String) -> CommandCodeCookieOverride? { + var pairs: [(name: String, value: String)] = [] + for chunk in header.split(separator: ";") { + let trimmed = chunk.trimmingCharacters(in: .whitespacesAndNewlines) + guard let separator = trimmed.firstIndex(of: "=") else { continue } + let key = String(trimmed[.. CommandCodeCookieOverride? { + var byLowerName: [String: (name: String, value: String)] = [:] + for pair in pairs { + byLowerName[pair.name.lowercased()] = pair + } + for expected in self.supportedSessionCookieNames { + if let match = byLowerName[expected.lowercased()] { + return CommandCodeCookieOverride(name: match.name, token: match.value) + } + } + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/CommandCode/CommandCodeCookieImporter.swift b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeCookieImporter.swift new file mode 100644 index 0000000..3db2c52 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeCookieImporter.swift @@ -0,0 +1,230 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit + +/// Imports CommandCode session cookies from installed browsers (Chrome by default). +public enum CommandCodeCookieImporter { + private static let importSessionCacheTTL: TimeInterval = 5 + private static let importSessionCache = ImportSessionCache(ttl: importSessionCacheTTL) + private static let log = CodexBarLog.logger(LogCategories.commandcodeCookie) + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["commandcode.ai", "www.commandcode.ai"] + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.commandcode]?.browserCookieOrder ?? Browser.defaultImportOrder + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var sessionCookie: CommandCodeCookieOverride? { + CommandCodeCookieHeader.sessionCookie(from: self.cookies) + } + + public var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + } + + public static func importSessions( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + if let cached = self.cachedImportSessions() { + return cached + } + + var sessions: [SessionInfo] = [] + let candidates = self.cookieImportOrder.cookieImportCandidates(using: browserDetection) + for browserSource in candidates { + do { + let perSource = try self.importSessions(from: browserSource, logger: logger) + sessions.append(contentsOf: perSource) + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + self.emit( + "\(browserSource.displayName) cookie import failed: \(error.localizedDescription)", + logger: logger) + } + } + + guard !sessions.isEmpty else { + throw CommandCodeCookieImportError.noCookies + } + self.storeImportSessions(sessions) + return sessions + } + + public static func importSessions( + from browserSource: Browser, + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + let query = BrowserCookieQuery(domains: self.cookieDomains) + let log: (String) -> Void = { msg in self.emit(msg, logger: logger) } + let sources = try Self.cookieClient.codexBarRecords( + matching: query, + in: browserSource, + logger: log) + + var sessions: [SessionInfo] = [] + let grouped = Dictionary(grouping: sources, by: { $0.store.profile.id }) + let sortedGroups = grouped.values.sorted { lhs, rhs in + self.mergedLabel(for: lhs) < self.mergedLabel(for: rhs) + } + + for group in sortedGroups where !group.isEmpty { + let label = self.mergedLabel(for: group) + let mergedRecords = self.mergeRecords(group) + guard !mergedRecords.isEmpty else { continue } + let httpCookies = BrowserCookieClient.makeHTTPCookies(mergedRecords, origin: query.origin) + guard !httpCookies.isEmpty else { continue } + + let session = SessionInfo(cookies: httpCookies, sourceLabel: label) + if let sessionCookie = session.sessionCookie { + log("Found \(sessionCookie.name) cookie in \(label)") + } else { + let names = httpCookies.map(\.name).joined(separator: ", ") + log("No known session name in \(label); sending all domain cookies (\(names))") + } + sessions.append(session) + } + return sessions + } + + public static func importSession( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) throws -> SessionInfo + { + let sessions = try self.importSessions(browserDetection: browserDetection, logger: logger) + guard let first = sessions.first else { throw CommandCodeCookieImportError.noCookies } + return first + } + + public static func hasSession( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) -> Bool + { + do { + let session = try self.importSession(browserDetection: browserDetection, logger: logger) + return !session.cookies.isEmpty + } catch { + return false + } + } + + static func invalidateImportSessionCache() { + self.importSessionCache.invalidate() + } + + private static func emit(_ message: String, logger: ((String) -> Void)?) { + logger?("[commandcode-cookie] \(message)") + self.log.debug(message) + } + + private static func cachedImportSessions(now: Date = Date()) -> [SessionInfo]? { + self.importSessionCache.load(now: now) + } + + private static func storeImportSessions(_ sessions: [SessionInfo], now: Date = Date()) { + self.importSessionCache.store(sessions, now: now) + } + + private static func mergedLabel(for sources: [BrowserCookieStoreRecords]) -> String { + guard let base = sources.map(\.label).min() else { return "Unknown" } + if base.hasSuffix(" (Network)") { + return String(base.dropLast(" (Network)".count)) + } + return base + } + + private static func mergeRecords(_ sources: [BrowserCookieStoreRecords]) -> [BrowserCookieRecord] { + let sortedSources = sources.sorted { lhs, rhs in + self.storePriority(lhs.store.kind) < self.storePriority(rhs.store.kind) + } + var mergedByKey: [String: BrowserCookieRecord] = [:] + for source in sortedSources { + for record in source.records { + let key = self.recordKey(record) + if let existing = mergedByKey[key] { + if self.shouldReplace(existing: existing, candidate: record) { + mergedByKey[key] = record + } + } else { + mergedByKey[key] = record + } + } + } + return Array(mergedByKey.values) + } + + private static func storePriority(_ kind: BrowserCookieStoreKind) -> Int { + switch kind { + case .network: 0 + case .primary: 1 + case .safari: 2 + } + } + + private static func recordKey(_ record: BrowserCookieRecord) -> String { + "\(record.name)|\(record.domain)|\(record.path)" + } + + private static func shouldReplace(existing: BrowserCookieRecord, candidate: BrowserCookieRecord) -> Bool { + switch (existing.expires, candidate.expires) { + case let (lhs?, rhs?): rhs > lhs + case (nil, .some): true + case (.some, nil): false + case (nil, nil): false + } + } + + private final class ImportSessionCache: @unchecked Sendable { + private let ttl: TimeInterval + private let lock = NSLock() + private var entry: (sessions: [SessionInfo], expiresAt: Date)? + + init(ttl: TimeInterval) { + self.ttl = ttl + } + + func load(now: Date) -> [SessionInfo]? { + self.lock.lock() + defer { self.lock.unlock() } + guard let entry = self.entry else { return nil } + guard entry.expiresAt > now else { + self.entry = nil + return nil + } + return entry.sessions + } + + func store(_ sessions: [SessionInfo], now: Date) { + self.lock.lock() + defer { self.lock.unlock() } + self.entry = (sessions: sessions, expiresAt: now.addingTimeInterval(self.ttl)) + } + + func invalidate() { + self.lock.lock() + defer { self.lock.unlock() } + self.entry = nil + } + } +} + +public enum CommandCodeCookieImportError: LocalizedError { + case noCookies + + public var errorDescription: String? { + switch self { + case .noCookies: + "No Command Code session cookies found in browsers. Sign in to commandcode.ai." + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/CommandCode/CommandCodePlanCatalog.swift b/Sources/CodexBarCore/Providers/CommandCode/CommandCodePlanCatalog.swift new file mode 100644 index 0000000..16c7bce --- /dev/null +++ b/Sources/CodexBarCore/Providers/CommandCode/CommandCodePlanCatalog.swift @@ -0,0 +1,34 @@ +import Foundation + +/// Static catalog of CommandCode subscription plans → monthly credit allowance (in USD). +/// +/// The `/internal/billing/credits` endpoint exposes the *remaining* `monthlyCredits`, +/// not the plan total. The plan total is published on the public pricing page +/// (https://commandcode.ai/pricing) and is keyed by `planId` returned from +/// `/internal/billing/subscriptions`. +public enum CommandCodePlanCatalog { + public struct Plan: Sendable, Equatable { + public let id: String + public let displayName: String + /// Monthly credit allowance in USD. + public let monthlyCreditsUSD: Double + + public init(id: String, displayName: String, monthlyCreditsUSD: Double) { + self.id = id + self.displayName = displayName + self.monthlyCreditsUSD = monthlyCreditsUSD + } + } + + public static let plans: [Plan] = [ + Plan(id: "individual-go", displayName: "Go", monthlyCreditsUSD: 10), + Plan(id: "individual-pro", displayName: "Pro", monthlyCreditsUSD: 30), + Plan(id: "individual-max", displayName: "Max", monthlyCreditsUSD: 150), + Plan(id: "individual-ultra", displayName: "Ultra", monthlyCreditsUSD: 300), + ] + + public static func plan(forID planID: String) -> Plan? { + let normalized = planID.lowercased() + return self.plans.first(where: { $0.id == normalized }) + } +} diff --git a/Sources/CodexBarCore/Providers/CommandCode/CommandCodeProviderDescriptor.swift b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeProviderDescriptor.swift new file mode 100644 index 0000000..c348d79 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeProviderDescriptor.swift @@ -0,0 +1,98 @@ +import Foundation + +public enum CommandCodeProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .commandcode, + metadata: ProviderMetadata( + id: .commandcode, + displayName: "Command Code", + sessionLabel: "Monthly credits", + weeklyLabel: "Monthly", + opusLabel: nil, + supportsOpus: false, + supportsCredits: true, + creditsHint: "Monthly USD credits from Command Code billing.", + toggleTitle: "Show Command Code usage", + cliName: "commandcode", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder, + dashboardURL: "https://commandcode.ai/studio", + subscriptionDashboardURL: "https://commandcode.ai/sixhobbits/settings/billing", + statusPageURL: nil, + statusLinkURL: nil), + branding: ProviderBranding( + iconStyle: .commandcode, + iconResourceName: "ProviderIcon-commandcode", + color: ProviderColor(red: 0 / 255, green: 0 / 255, blue: 0 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Command Code cost summary is not yet supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [CommandCodeWebFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "commandcode", + aliases: ["command-code"], + versionDetector: nil)) + } +} + +struct CommandCodeWebFetchStrategy: ProviderFetchStrategy { + let id: String = "commandcode.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + guard context.settings?.commandcode?.cookieSource != .off else { return false } + if Self.manualCookieHeader(from: context) != nil { + return true + } + #if os(macOS) + return true + #else + return false + #endif + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let cookieHeader: String + let sourceLabel: String + if let manual = Self.manualCookieHeader(from: context) { + cookieHeader = manual + sourceLabel = "manual" + } else { + #if os(macOS) + let session: CommandCodeCookieImporter.SessionInfo + do { + session = try CommandCodeCookieImporter.importSession() + } catch { + throw CommandCodeUsageError.missingCredentials + } + guard !session.cookies.isEmpty else { + throw CommandCodeUsageError.missingCredentials + } + cookieHeader = session.cookieHeader + sourceLabel = session.sourceLabel + #else + throw CommandCodeUsageError.missingCredentials + #endif + } + let snapshot = try await CommandCodeUsageFetcher.fetchUsage(cookieHeader: cookieHeader) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: sourceLabel) + } + + private static func manualCookieHeader(from context: ProviderFetchContext) -> String? { + guard context.settings?.commandcode?.cookieSource == .manual else { return nil } + return CookieHeaderNormalizer.normalize(context.settings?.commandcode?.manualCookieHeader) + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} diff --git a/Sources/CodexBarCore/Providers/CommandCode/CommandCodeUsageError.swift b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeUsageError.swift new file mode 100644 index 0000000..9e13b89 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeUsageError.swift @@ -0,0 +1,34 @@ +import Foundation + +public enum CommandCodeUsageError: LocalizedError, Sendable, Equatable { + case missingCredentials + case invalidCredentials + case networkError(String) + case apiError(Int) + case parseFailed(String) + case unknownPlan(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Command Code session cookie not found. Sign in to commandcode.ai in Chrome." + case .invalidCredentials: + "Command Code session is invalid or expired. Sign in to commandcode.ai again." + case let .networkError(message): + "Command Code network error: \(message)" + case let .apiError(status): + "Command Code API returned status \(status)." + case let .parseFailed(message): + "Could not parse Command Code response: \(message)" + case let .unknownPlan(planID): + "Unknown Command Code plan: \(planID). Add it to CommandCodePlanCatalog." + } + } + + public var isAuthRelated: Bool { + switch self { + case .missingCredentials, .invalidCredentials: true + default: false + } + } +} diff --git a/Sources/CodexBarCore/Providers/CommandCode/CommandCodeUsageFetcher.swift b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeUsageFetcher.swift new file mode 100644 index 0000000..74b8d92 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeUsageFetcher.swift @@ -0,0 +1,263 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +/// Fetches live billing data from `api.commandcode.ai` using a better-auth session +/// cookie scraped from the user's browser. +public enum CommandCodeUsageFetcher { + private static let log = CodexBarLog.logger(LogCategories.commandcodeUsage) + private static let requestTimeoutSeconds: TimeInterval = 15 + private static let subscriptionGraceSeconds: TimeInterval = 2 + private static let apiBase = URL(string: "https://api.commandcode.ai")! + private static let creditsPath = "/internal/billing/credits" + private static let subscriptionsPath = "/internal/billing/subscriptions" + private static let webOrigin = "https://commandcode.ai" + private static let userAgent = + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36" + + public static func fetchUsage( + cookieHeader: String, + session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + now: Date = Date()) async throws -> CommandCodeUsageSnapshot + { + try await self.fetchUsage( + cookieHeader: cookieHeader, + transport: transport, + now: now, + subscriptionGrace: .seconds(self.subscriptionGraceSeconds)) + } + + static func _fetchUsageForTesting( + cookieHeader: String, + transport: any ProviderHTTPTransport, + now: Date = Date(), + subscriptionGrace: Duration) async throws -> CommandCodeUsageSnapshot + { + try await self.fetchUsage( + cookieHeader: cookieHeader, + transport: transport, + now: now, + subscriptionGrace: subscriptionGrace) + } + + private static func fetchUsage( + cookieHeader: String, + transport: any ProviderHTTPTransport, + now: Date, + subscriptionGrace: Duration) async throws -> CommandCodeUsageSnapshot + { + let (credits, subscription, subscriptionEnrichmentUnavailable) = try await self.fetchPayloads( + cookieHeader: cookieHeader, + transport: transport, + subscriptionGrace: subscriptionGrace) + + let plan: CommandCodePlanCatalog.Plan? = subscription.flatMap { sub in + CommandCodePlanCatalog.plan(forID: sub.planID) + } + + // If we got an active subscription with an unrecognised plan ID, surface that + // explicitly rather than silently dropping the totals row. + if let sub = subscription, sub.status.lowercased() == "active", plan == nil { + Self.log.error("Unknown CommandCode planId: \(sub.planID)") + throw CommandCodeUsageError.unknownPlan(sub.planID) + } + + return CommandCodeUsageSnapshot( + monthlyCreditsRemaining: credits.monthlyCredits, + purchasedCredits: credits.purchasedCredits, + premiumMonthlyCredits: credits.premiumMonthlyCredits, + opensourceMonthlyCredits: credits.opensourceMonthlyCredits, + plan: plan, + billingPeriodEnd: subscription?.currentPeriodEnd, + subscriptionStatus: subscription?.status, + subscriptionEnrichmentUnavailable: subscriptionEnrichmentUnavailable, + updatedAt: now) + } + + private static func fetchPayloads( + cookieHeader: String, + transport: any ProviderHTTPTransport, + subscriptionGrace: Duration) async throws -> (CreditsPayload, SubscriptionPayload?, Bool) + { + let subscriptionTask = Task { + try await self.fetchSubscription(cookieHeader: cookieHeader, transport: transport) + } + let credits: CreditsPayload + do { + credits = try await withTaskCancellationHandler { + try await self.fetchCredits(cookieHeader: cookieHeader, transport: transport) + } onCancel: { + subscriptionTask.cancel() + } + } catch { + subscriptionTask.cancel() + throw error + } + + do { + try Task.checkCancellation() + } catch { + subscriptionTask.cancel() + throw error + } + let race = BoundedTaskJoin(sourceTask: subscriptionTask) + switch await race.value(joinGrace: subscriptionGrace) { + case let .value(subscription): + try Task.checkCancellation() + return (credits, subscription, false) + case .timedOut: + try Task.checkCancellation() + Self.log.warning("Command Code subscription enrichment timed out") + return (credits, nil, true) + case let .failure(error): + subscriptionTask.cancel() + try Task.checkCancellation() + Self.log.warning("Command Code subscription enrichment failed: \(error.localizedDescription)") + return (credits, nil, true) + } + } + + // MARK: - Endpoints + + struct CreditsPayload { + let monthlyCredits: Double + let purchasedCredits: Double + let premiumMonthlyCredits: Double + let opensourceMonthlyCredits: Double + } + + struct SubscriptionPayload { + let planID: String + let status: String + let currentPeriodEnd: Date? + } + + private static func fetchCredits( + cookieHeader: String, + transport: any ProviderHTTPTransport) async throws -> CreditsPayload + { + let url = self.apiBase.appendingPathComponent(self.creditsPath) + let data = try await self.send(url: url, cookieHeader: cookieHeader, transport: transport) + return try self.parseCredits(data: data) + } + + private static func fetchSubscription( + cookieHeader: String, + transport: any ProviderHTTPTransport) async throws -> SubscriptionPayload? + { + let url = self.apiBase.appendingPathComponent(self.subscriptionsPath) + let data = try await self.send(url: url, cookieHeader: cookieHeader, transport: transport) + return try self.parseSubscription(data: data) + } + + private static func send( + url: URL, + cookieHeader: String, + transport: any ProviderHTTPTransport) async throws -> Data + { + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.timeoutInterval = self.requestTimeoutSeconds + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept") + request.setValue("en-US,en;q=0.9", forHTTPHeaderField: "Accept-Language") + request.setValue(self.userAgent, forHTTPHeaderField: "User-Agent") + request.setValue(self.webOrigin, forHTTPHeaderField: "Origin") + request.setValue("\(self.webOrigin)/", forHTTPHeaderField: "Referer") + + let response: ProviderHTTPResponse + do { + response = try await transport.response(for: request) + } catch { + if error is CancellationError || (error as? URLError)?.code == .cancelled || Task.isCancelled { + throw CancellationError() + } + throw CommandCodeUsageError.networkError(error.localizedDescription) + } + if response.statusCode == 401 || response.statusCode == 403 { + throw CommandCodeUsageError.invalidCredentials + } + guard (200..<300).contains(response.statusCode) else { + let body = String(data: response.data, encoding: .utf8) ?? "" + Self.log.error("CommandCode \(url.path) → \(response.statusCode): \(body)") + throw CommandCodeUsageError.apiError(response.statusCode) + } + return response.data + } + + // MARK: - Parsing + + static func parseCredits(data: Data) throws -> CreditsPayload { + guard let root = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw CommandCodeUsageError.parseFailed("Credits: invalid JSON") + } + guard let credits = root["credits"] as? [String: Any] else { + throw CommandCodeUsageError.parseFailed("Credits: missing 'credits' object") + } + guard let monthly = self.double(from: credits["monthlyCredits"]) else { + throw CommandCodeUsageError.parseFailed("Credits: missing monthlyCredits") + } + return CreditsPayload( + monthlyCredits: monthly, + purchasedCredits: self.double(from: credits["purchasedCredits"]) ?? 0, + premiumMonthlyCredits: self.double(from: credits["premiumMonthlyCredits"]) ?? 0, + opensourceMonthlyCredits: self.double(from: credits["opensourceMonthlyCredits"]) ?? 0) + } + + static func parseSubscription(data: Data) throws -> SubscriptionPayload? { + guard let root = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw CommandCodeUsageError.parseFailed("Subscriptions: invalid JSON") + } + // Only an explicit successful null response identifies the free tier. Failure envelopes are transient. + guard let success = root["success"] as? Bool else { + throw CommandCodeUsageError.parseFailed("Subscriptions: missing success flag") + } + guard success else { + throw CommandCodeUsageError.parseFailed("Subscriptions: unsuccessful response") + } + guard let dataValue = root["data"] else { + throw CommandCodeUsageError.parseFailed("Subscriptions: missing data") + } + if dataValue is NSNull { + return nil + } + guard let data = dataValue as? [String: Any] else { + throw CommandCodeUsageError.parseFailed("Subscriptions: invalid data") + } + guard let planID = data["planId"] as? String, !planID.isEmpty else { + throw CommandCodeUsageError.parseFailed("Subscriptions: missing planId") + } + let status = (data["status"] as? String) ?? "unknown" + let periodEnd = self.date(from: data["currentPeriodEnd"]) + return SubscriptionPayload(planID: planID, status: status, currentPeriodEnd: periodEnd) + } + + // MARK: - Value coercion + + private static func double(from value: Any?) -> Double? { + switch value { + case let n as NSNumber: + let d = n.doubleValue + return d.isFinite ? d : nil + case let s as String: + let trimmed = s.trimmingCharacters(in: .whitespacesAndNewlines) + return Double(trimmed) + default: + return nil + } + } + + private static func date(from value: Any?) -> Date? { + guard let s = value as? String else { return nil } + let trimmed = s.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + let fractional = ISO8601DateFormatter() + fractional.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = fractional.date(from: trimmed) { return date } + let plain = ISO8601DateFormatter() + plain.formatOptions = [.withInternetDateTime] + return plain.date(from: trimmed) + } +} diff --git a/Sources/CodexBarCore/Providers/CommandCode/CommandCodeUsageSnapshot.swift b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeUsageSnapshot.swift new file mode 100644 index 0000000..b5747a9 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CommandCode/CommandCodeUsageSnapshot.swift @@ -0,0 +1,124 @@ +import Foundation + +/// Parsed view of CommandCode `/internal/billing/credits` + `/internal/billing/subscriptions`. +public struct CommandCodeUsageSnapshot: Sendable { + /// USD remaining in the current monthly grant (`credits.monthlyCredits`). + public let monthlyCreditsRemaining: Double + /// USD top-up balance carried over (`credits.purchasedCredits`). + public let purchasedCredits: Double + /// USD remaining in the premium monthly grant (`credits.premiumMonthlyCredits`). + public let premiumMonthlyCredits: Double + /// USD remaining in the open-source monthly grant (`credits.opensourceMonthlyCredits`). + public let opensourceMonthlyCredits: Double + /// Subscription plan, or nil when the user is on the free tier. + public let plan: CommandCodePlanCatalog.Plan? + /// `currentPeriodEnd` from the active subscription. + public let billingPeriodEnd: Date? + /// Subscription status (e.g. `active`, `canceled`). + public let subscriptionStatus: String? + /// The optional subscription request timed out or failed for this refresh. + public let subscriptionEnrichmentUnavailable: Bool + public let updatedAt: Date + + public init( + monthlyCreditsRemaining: Double, + purchasedCredits: Double, + premiumMonthlyCredits: Double, + opensourceMonthlyCredits: Double, + plan: CommandCodePlanCatalog.Plan?, + billingPeriodEnd: Date?, + subscriptionStatus: String?, + subscriptionEnrichmentUnavailable: Bool = false, + updatedAt: Date = Date()) + { + self.monthlyCreditsRemaining = monthlyCreditsRemaining + self.purchasedCredits = purchasedCredits + self.premiumMonthlyCredits = premiumMonthlyCredits + self.opensourceMonthlyCredits = opensourceMonthlyCredits + self.plan = plan + self.billingPeriodEnd = billingPeriodEnd + self.subscriptionStatus = subscriptionStatus + self.subscriptionEnrichmentUnavailable = subscriptionEnrichmentUnavailable + self.updatedAt = updatedAt + } + + /// USD allocation for the active monthly grant (from the catalog). + public var monthlyCreditsTotal: Double? { + self.plan?.monthlyCreditsUSD + } + + /// USD spent in the current monthly grant (total – remaining), clamped to [0, total]. + public var monthlyCreditsUsed: Double? { + guard let total = self.monthlyCreditsTotal else { return nil } + return max(0, min(total, total - self.monthlyCreditsRemaining)) + } + + public func toUsageSnapshot() -> UsageSnapshot { + let primary = self.makePrimaryWindow() + + let identity = ProviderIdentitySnapshot( + providerID: .commandcode, + accountEmail: nil, + accountOrganization: nil, + loginMethod: self.makeLoginMethod()) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + commandCodeSubscriptionEnrichmentUnavailable: self.subscriptionEnrichmentUnavailable, + commandCodeHasSubscriptionPlan: self.plan != nil, + commandCodeMonthlyGrantDepleted: self.monthlyCreditsRemaining <= 0, + updatedAt: self.updatedAt, + identity: identity) + } + + private func makePrimaryWindow() -> RateWindow? { + guard let total = self.monthlyCreditsTotal, total > 0 else { + // Free / unknown plan with no allowance — surface 100% so the bar renders empty. + if self.monthlyCreditsRemaining > 0 || self.purchasedCredits > 0 { + return RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: self.billingPeriodEnd, + resetDescription: nil) + } + return nil + } + let used = self.monthlyCreditsUsed ?? 0 + let percent = min(100, max(0, (used / total) * 100)) + return RateWindow( + usedPercent: percent, + windowMinutes: nil, + resetsAt: self.billingPeriodEnd, + resetDescription: nil) + } + + private func makeLoginMethod() -> String? { + var parts: [String] = [] + if let name = self.plan?.displayName, !name.isEmpty { + parts.append(name) + } + if let total = self.monthlyCreditsTotal { + let used = self.monthlyCreditsUsed ?? 0 + parts.append("\(Self.formatUSD(used)) of \(Self.formatUSD(total))") + } else if self.monthlyCreditsRemaining > 0 { + parts.append("\(Self.formatUSD(self.monthlyCreditsRemaining)) remaining") + } + if self.purchasedCredits > 0 { + parts.append("+ \(Self.formatUSD(self.purchasedCredits)) credits") + } + return parts.isEmpty ? nil : parts.joined(separator: " · ") + } + + static func formatUSD(_ value: Double) -> String { + let formatter = NumberFormatter() + formatter.numberStyle = .currency + formatter.currencyCode = "USD" + formatter.locale = Locale(identifier: "en_US") + formatter.maximumFractionDigits = value < 100 ? 2 : 0 + formatter.minimumFractionDigits = value < 100 ? 2 : 0 + return formatter.string(from: NSNumber(value: value)) ?? "$\(value)" + } +} diff --git a/Sources/CodexBarCore/Providers/Copilot/CopilotBudgetWebFetcher.swift b/Sources/CodexBarCore/Providers/Copilot/CopilotBudgetWebFetcher.swift new file mode 100644 index 0000000..4c599fd --- /dev/null +++ b/Sources/CodexBarCore/Providers/Copilot/CopilotBudgetWebFetcher.swift @@ -0,0 +1,801 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif +#if os(macOS) +import SweetCookieKit +#endif + +public struct CopilotBudgetWebFetcher: Sendable { + public enum Error: Swift.Error, LocalizedError, Equatable { + case noSessionCookie + case notLoggedIn + case accountMismatch(expected: String, actual: String?) + case badStatus(Int) + case invalidResponse + + public var errorDescription: String? { + switch self { + case .noSessionCookie: + "No GitHub browser session cookie found." + case .notLoggedIn: + "GitHub browser session is not logged in." + case let .accountMismatch(expected, actual): + "GitHub browser session belongs to \(actual ?? "an unknown account"), expected \(expected)." + case let .badStatus(status): + "GitHub budgets request failed with HTTP \(status)." + case .invalidResponse: + "GitHub budgets response could not be decoded." + } + } + } + + struct BudgetResponse: Decodable { + let budgets: [Budget] + let hasNextPage: Bool? + + private enum CodingKeys: String, CodingKey { + case budgets + case payload + case hasNextPage + case hasNextPageSnake = "has_next_page" + } + + init(budgets: [Budget], hasNextPage: Bool? = nil) { + self.budgets = budgets + self.hasNextPage = hasNextPage + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + if let payload = try container.decodeIfPresent(BudgetResponse.self, forKey: .payload) { + self = payload + return + } + self.budgets = try container.decodeIfPresent([Budget].self, forKey: .budgets) ?? [] + self.hasNextPage = try container.decodeIfPresent(Bool.self, forKey: .hasNextPage) + ?? container.decodeIfPresent(Bool.self, forKey: .hasNextPageSnake) + } + } + + struct Budget: Decodable, Equatable { + let id: String? + let name: String? + let budgetType: String? + let budgetProductSkus: [String] + let budgetScope: String? + let budgetEntityName: String? + let budgetAmount: Double + let currentAmount: Double + + init( + id: String? = nil, + name: String? = nil, + budgetType: String? = nil, + budgetProductSkus: [String] = [], + budgetScope: String? = nil, + budgetEntityName: String? = nil, + budgetAmount: Double, + currentAmount: Double = 0) + { + self.id = id + self.name = name + self.budgetType = budgetType + self.budgetProductSkus = budgetProductSkus + self.budgetScope = budgetScope + self.budgetEntityName = budgetEntityName + self.budgetAmount = budgetAmount + self.currentAmount = currentAmount + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: DynamicCodingKey.self) + self.id = Self.decodeString(container: container, keys: ["id", "uuid", "budget_id", "budgetId"]) + self.name = Self.decodeString(container: container, keys: ["name", "display_name", "displayName", "title"]) + self.budgetType = Self.decodeString( + container: container, + keys: ["budget_type", "budgetType", "type", "pricing_target_type", "pricingTargetType"]) + self.budgetProductSkus = Self.decodeStringArray( + container: container, + keys: [ + "budget_product_skus", + "budgetProductSkus", + "budget_product_sku", + "budgetProductSku", + "product_skus", + "productSkus", + "skus", + "sku", + "product", + "product_name", + "productName", + "pricing_target_id", + "pricingTargetId", + ]) + self.budgetScope = Self.decodeString(container: container, keys: ["budget_scope", "budgetScope", "scope"]) + self.budgetEntityName = Self.decodeString( + container: container, + keys: [ + "budget_entity_name", + "budgetEntityName", + "entity_name", + "entityName", + "target_name", + "targetName", + ]) + self.budgetAmount = Self.decodeDouble( + container: container, + keys: [ + "budget_amount", + "budgetAmount", + "target_amount", + "targetAmount", + "spending_limit", + "spendingLimit", + "limit", + "amount", + "max", + ]) ?? 0 + self.currentAmount = Self.decodeDouble( + container: container, + keys: [ + "current_usage", + "currentUsage", + "current_amount", + "currentAmount", + "usage_amount", + "usageAmount", + "usage", + "spent", + "amount_used", + "amountUsed", + ]) ?? 0 + } + + var normalizedSelectors: Set { + let values = self.budgetProductSkus + [ + self.budgetType, + self.budgetEntityName, + self.name, + ].compactMap(\.self) + return Set(values.compactMap(CopilotBudgetWebFetcher.normalizedBillingIdentifier)) + } + + private static func decodeString( + container: KeyedDecodingContainer, + keys: [String]) -> String? + { + for key in keys { + guard let codingKey = DynamicCodingKey(key) else { continue } + if let value = try? container.decodeIfPresent(String.self, forKey: codingKey), !value.isEmpty { + return value + } + if let value = try? container.decodeIfPresent(Int.self, forKey: codingKey) { + return String(value) + } + } + return nil + } + + private static func decodeStringArray( + container: KeyedDecodingContainer, + keys: [String]) -> [String] + { + for key in keys { + guard let codingKey = DynamicCodingKey(key) else { continue } + if let values = try? container.decodeIfPresent([String].self, forKey: codingKey), !values.isEmpty { + return values + } + if let value = try? container.decodeIfPresent(String.self, forKey: codingKey), !value.isEmpty { + return [value] + } + if let values = try? container.decodeIfPresent([ProductSKU].self, forKey: codingKey), + !values.isEmpty + { + return values.flatMap(\.selectors) + } + } + return [] + } + + private static func decodeDouble( + container: KeyedDecodingContainer, + keys: [String]) -> Double? + { + for key in keys { + guard let codingKey = DynamicCodingKey(key) else { continue } + if let value = try? container.decodeIfPresent(Double.self, forKey: codingKey) { + return value + } + if let value = try? container.decodeIfPresent(Int.self, forKey: codingKey) { + return Double(value) + } + if let value = try? container.decodeIfPresent(String.self, forKey: codingKey), + let parsed = Self.parseAmount(value) + { + return parsed + } + if let value = try? container.decodeIfPresent(AmountValue.self, forKey: codingKey) { + return value.amount + } + } + return nil + } + + fileprivate static func parseAmount(_ value: String) -> Double? { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + let isNegative = trimmed.first == "-" + guard !trimmed.dropFirst(isNegative ? 1 : 0).contains("-") else { return nil } + let unsigned = trimmed.filter { $0.isNumber || $0 == "." } + guard !unsigned.isEmpty else { return nil } + return Double(isNegative ? "-\(unsigned)" : unsigned) + } + } + + struct GitHubWebIdentity: Equatable { + let id: String? + let login: String? + + init(id: String?, login: String?) { + let trimmedID = id?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + let trimmedLogin = login?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + self.id = trimmedID.isEmpty ? nil : trimmedID + self.login = trimmedLogin.isEmpty ? nil : trimmedLogin + } + + var displayName: String? { + self.login ?? self.id.map { "github:user:\($0)" } + } + } + + private struct ProductSKU: Decodable, Equatable { + let selectors: [String] + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: DynamicCodingKey.self) + self.selectors = [ + "sku", + "name", + "display_name", + "displayName", + "product", + "product_name", + "productName", + ].compactMap { key in + guard let codingKey = DynamicCodingKey(key) else { return nil } + return try? container.decodeIfPresent(String.self, forKey: codingKey) + } + } + } + + private struct AmountValue: Decodable { + let amount: Double? + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: DynamicCodingKey.self) + self.amount = [ + "amount", + "value", + "total", + "cents", + "formatted", + ].lazy.compactMap { key -> Double? in + guard let codingKey = DynamicCodingKey(key) else { return nil } + if let value = try? container.decodeIfPresent(Double.self, forKey: codingKey) { + return key == "cents" ? value / 100 : value + } + if let value = try? container.decodeIfPresent(Int.self, forKey: codingKey) { + return key == "cents" ? Double(value) / 100 : Double(value) + } + if let value = try? container.decodeIfPresent(String.self, forKey: codingKey) { + return Budget.parseAmount(value) + } + return nil + }.first + } + } + + private struct DynamicCodingKey: CodingKey { + let stringValue: String + let intValue: Int? + + init?(_ stringValue: String) { + self.stringValue = stringValue + self.intValue = nil + } + + init?(stringValue: String) { + self.init(stringValue) + } + + init?(intValue: Int) { + self.stringValue = String(intValue) + self.intValue = intValue + } + } + + private static let copilotProductID = "copilot" + private static let copilotPremiumRequestSKU = "copilot_premium_request" + private static let copilotAgentPremiumRequestSKU = "copilot_agent_premium_request" + private static let sparkPremiumRequestSKU = "spark_premium_request" + private static let copilotBudgetSelectors: Set = [ + copilotProductID, + copilotPremiumRequestSKU, + copilotAgentPremiumRequestSKU, + sparkPremiumRequestSKU, + ] + + private let cookieHeaderOverride: String? + private let expectedGitHubAccountIdentifier: String? + private let browserDetection: BrowserDetection + private let transport: any ProviderHTTPTransport + private let now: @Sendable () -> Date + + public init( + cookieHeaderOverride: String? = nil, + expectedGitHubAccountIdentifier: String? = nil, + browserDetection: BrowserDetection = BrowserDetection(), + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + now: @escaping @Sendable () -> Date = { Date() }) + { + self.cookieHeaderOverride = CookieHeaderNormalizer.normalize(cookieHeaderOverride ?? "") + self.expectedGitHubAccountIdentifier = Self.normalizedExpectedAccountIdentifier( + expectedGitHubAccountIdentifier) + self.browserDetection = browserDetection + self.transport = transport + self.now = now + } + + public func fetchBudgetWindows() async throws -> [NamedRateWindow] { + if let cookieHeaderOverride, !cookieHeaderOverride.isEmpty { + return try await self.fetchBudgetWindows(cookieHeader: cookieHeaderOverride) + } + + if let cached = CookieHeaderCache.load(provider: .copilot), + !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + do { + return try await self.fetchBudgetWindows(cookieHeader: cached.cookieHeader) + } catch { + if case Error.notLoggedIn = error { + CookieHeaderCache.clear(provider: .copilot) + } else if case Error.accountMismatch = error { + CookieHeaderCache.clear(provider: .copilot) + } else { + throw error + } + } + } + + #if os(macOS) + for session in CopilotGitHubCookieImporter.importSessions(browserDetection: self.browserDetection) { + do { + let windows = try await self.fetchBudgetWindows(cookieHeader: session.cookieHeader) + CookieHeaderCache.store( + provider: .copilot, + cookieHeader: session.cookieHeader, + sourceLabel: session.sourceLabel) + return windows + } catch { + if case Error.notLoggedIn = error { + continue + } + if case Error.accountMismatch = error { + continue + } + throw error + } + } + #endif + + throw Error.noSessionCookie + } + + func fetchBudgetWindows(cookieHeader: String) async throws -> [NamedRateWindow] { + let nonce = try await self.boundFetchNonce(cookieHeader: cookieHeader) + var allBudgets: [Budget] = [] + var page = 1 + let maxPages = 20 + var shouldContinue = true + while shouldContinue, page <= maxPages { + let response = try await self.fetchBudgetPage( + cookieHeader: cookieHeader, + nonce: nonce, + page: page) + allBudgets.append(contentsOf: response.budgets) + shouldContinue = response.hasNextPage == true + page += 1 + } + if shouldContinue { + CodexBarLog.logger(LogCategories.providers).warning( + "Copilot budget pagination reached page cap", + metadata: ["pageCap": "\(maxPages)"]) + } + return Self.extraRateWindows(from: allBudgets, now: self.now()) + } + + private func boundFetchNonce(cookieHeader: String) async throws -> String? { + guard self.expectedGitHubAccountIdentifier != nil else { + return await self.bestEffortFetchNonce(cookieHeader: cookieHeader) + } + let metadata = try await self.fetchBudgetPageMetadata(cookieHeader: cookieHeader) + try self.verifyExpectedGitHubAccount(metadata.identity) + return metadata.nonce + } + + private func bestEffortFetchNonce(cookieHeader: String) async -> String? { + do { + return try await self.fetchBudgetPageMetadata(cookieHeader: cookieHeader).nonce + } catch { + // GitHub accepts some budget requests without a nonce. Keep auth failures and + // page-shape changes non-fatal so a valid Cookie header can still be tried. + return nil + } + } + + private struct BudgetPageMetadata { + let nonce: String? + let identity: GitHubWebIdentity? + } + + private func fetchBudgetPageMetadata(cookieHeader: String) async throws -> BudgetPageMetadata { + guard let url = URL(string: "https://github.com/settings/billing/budgets") else { + throw URLError(.badURL) + } + var request = URLRequest(url: url) + request.timeoutInterval = 15 + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("text/html,application/xhtml+xml", forHTTPHeaderField: "Accept") + request.setValue("CodexBar", forHTTPHeaderField: "User-Agent") + + let response = try await self.transport.response(for: request) + switch response.statusCode { + case 200: + guard let html = String(data: response.data, encoding: .utf8) else { throw Error.invalidResponse } + return BudgetPageMetadata( + nonce: Self.extractFetchNonce(from: html), + identity: Self.extractGitHubWebIdentity(from: html)) + case 401, 403: + throw Error.notLoggedIn + default: + throw Error.badStatus(response.statusCode) + } + } + + private func verifyExpectedGitHubAccount(_ actual: GitHubWebIdentity?) throws { + guard let expected = self.expectedGitHubAccountIdentifier else { return } + guard let actual else { throw Error.accountMismatch(expected: expected, actual: nil) } + guard Self.webIdentity(actual, matches: expected) else { + throw Error.accountMismatch(expected: expected, actual: actual.displayName) + } + } + + private func fetchBudgetPage(cookieHeader: String, nonce: String?, page: Int) async throws -> BudgetResponse { + guard var components = URLComponents(string: "https://github.com/settings/billing/budgets") else { + throw URLError(.badURL) + } + components.queryItems = [ + URLQueryItem(name: "page", value: "\(page)"), + URLQueryItem(name: "page_size", value: "10"), + URLQueryItem(name: "scope", value: "customer"), + ] + guard let url = components.url else { throw URLError(.badURL) } + var request = URLRequest(url: url) + request.timeoutInterval = 15 + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("https://github.com/settings/billing/budgets", forHTTPHeaderField: "Referer") + request.setValue("XMLHttpRequest", forHTTPHeaderField: "X-Requested-With") + request.setValue("true", forHTTPHeaderField: "GitHub-Verified-Fetch") + request.setValue("CodexBar", forHTTPHeaderField: "User-Agent") + if let nonce, !nonce.isEmpty { + request.setValue(nonce, forHTTPHeaderField: "X-Fetch-Nonce") + } + + let response = try await self.transport.response(for: request) + switch response.statusCode { + case 200: + do { + return try JSONDecoder().decode(BudgetResponse.self, from: response.data) + } catch { + throw Error.invalidResponse + } + case 401, 403: + throw Error.notLoggedIn + default: + throw Error.badStatus(response.statusCode) + } + } + + static func extractFetchNonce(from html: String) -> String? { + let patterns = [ + #"x-fetch-nonce"\s+content="([^"]+)""#, + #"X-Fetch-Nonce"\s*:\s*"([^"]+)""#, + #"fetchNonce"\s*:\s*"([^"]+)""#, + #"data-fetch-nonce="([^"]+)""#, + ] + for pattern in patterns { + guard let regex = try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) else { + continue + } + let range = NSRange(html.startIndex.. GitHubWebIdentity? { + let id = self.extractMetaContent( + named: [ + "octolytics-actor-id", + "analytics-user-id", + "user-id", + ], + from: html) + let login = self.extractMetaContent( + named: [ + "user-login", + "octolytics-actor-login", + "analytics-user-login", + ], + from: html) + let identity = GitHubWebIdentity(id: id, login: login) + return identity.id == nil && identity.login == nil ? nil : identity + } + + private static let metaTagRegex = try? NSRegularExpression( + pattern: #"]*>"#, + options: [.caseInsensitive]) + + private static let metaAttributeRegex = try? NSRegularExpression( + pattern: #"([A-Za-z_:][-A-Za-z0-9_:.]*)\s*=\s*(['"])(.*?)\2"#, + options: [.caseInsensitive]) + + private static func extractMetaContent(named names: [String], from html: String) -> String? { + guard let metaTagRegex, let metaAttributeRegex else { return nil } + let expectedNames = Set(names.map { $0.lowercased() }) + var contentByName: [String: String] = [:] + let htmlRange = NSRange(html.startIndex.. Bool { + guard let expected = self.normalizedExpectedAccountIdentifier(expectedIdentifier), + let identity + else { return false } + if let expectedID = self.githubUserID(from: expected) { + return identity.id == expectedID + } + return identity.login?.lowercased() == expected + } + + static func normalizedGitHubAccountIdentifier(for identity: CopilotUsageFetcher.GitHubUserIdentity) -> String { + "github:user:\(identity.id)" + } + + private static func normalizedExpectedAccountIdentifier(_ identifier: String?) -> String? { + let trimmed = identifier?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + return trimmed.isEmpty ? nil : trimmed.lowercased() + } + + private static func githubUserID(from identifier: String) -> String? { + let prefix = "github:user:" + guard identifier.hasPrefix(prefix) else { return nil } + let suffix = String(identifier.dropFirst(prefix.count)) + .trimmingCharacters(in: .whitespacesAndNewlines) + return suffix.isEmpty ? nil : suffix + } + + static func extraRateWindows(from budgets: [Budget], now: Date) -> [NamedRateWindow] { + var usedIDs = Set() + let resetDate = self.approximateNextMonthResetDate(now: now) + return budgets + .compactMap { budget in + let selectors = budget.normalizedSelectors + guard self.isCopilotBudget(budget, selectors: selectors) else { return nil } + let id = self.uniqueWindowID(for: budget, selectors: selectors, usedIDs: &usedIDs) + let usedPercent = budget.budgetAmount > 0 + ? min(999, max(0, budget.currentAmount / budget.budgetAmount * 100)) + : 0 + let window = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: resetDate, + resetDescription: resetDate.map { + UsageFormatter.resetDescription(from: $0, now: now) + }) + return NamedRateWindow( + id: id, + title: self.windowTitle(for: budget, selectors: selectors), + window: window) + } + } + + private static func isCopilotBudget(_ budget: Budget, selectors: Set) -> Bool { + guard budget.budgetAmount > 0 else { return false } + return !selectors.isDisjoint(with: self.copilotBudgetSelectors) + } + + static func normalizedBillingIdentifier(_ value: String?) -> String? { + guard let value else { return nil } + let slug = self.slug(value) + guard !slug.isEmpty else { return nil } + let underscored = slug.replacingOccurrences(of: "-", with: "_") + if underscored == self.copilotProductID { + return self.copilotProductID + } + if underscored == "premium_request" || underscored == "premium_requests" { + return self.copilotPremiumRequestSKU + } + if underscored == "coding_agent_premium_request" || underscored == "coding_agent_premium_requests" { + return self.copilotAgentPremiumRequestSKU + } + if underscored.contains("spark"), underscored.contains("premium"), underscored.contains("request") { + return self.sparkPremiumRequestSKU + } + if underscored.contains("cloud") || underscored.contains("coding"), + underscored.contains("agent"), + underscored.contains("premium"), + underscored.contains("request") + { + return self.copilotAgentPremiumRequestSKU + } + if underscored.contains("bundled"), underscored.contains("premium"), underscored.contains("request") { + return self.copilotPremiumRequestSKU + } + if underscored.contains("copilot"), + underscored.contains("agent"), + underscored.contains("premium"), + underscored.contains("request") + { + return self.copilotAgentPremiumRequestSKU + } + if underscored.contains("copilot"), underscored.contains("premium"), underscored.contains("request") { + return self.copilotPremiumRequestSKU + } + return underscored + } + + private static func windowTitle(for budget: Budget, selectors: Set) -> String { + let budgetType = if selectors == [self.copilotProductID] { + "Copilot" + } else if selectors.contains(self.copilotAgentPremiumRequestSKU) { + "Copilot Agent Premium Requests" + } else if selectors.contains(self.sparkPremiumRequestSKU) { + "Spark Premium Requests" + } else if selectors.contains(self.copilotPremiumRequestSKU) { + "All Premium Request SKUs" + } else if let name = budget.name?.trimmingCharacters(in: .whitespacesAndNewlines), + !name.isEmpty + { + name + } else { + "Copilot Premium Requests" + } + return "Budget - \(budgetType)" + } + + private static func uniqueWindowID( + for budget: Budget, + selectors: Set, + usedIDs: inout Set) -> String + { + let source = budget.id ?? budget.budgetProductSkus.joined(separator: "-") + let slug = self.slug(source.isEmpty ? self.windowTitle(for: budget, selectors: selectors) : source) + let base = slug.isEmpty ? "copilot-budget" : "copilot-budget-\(slug)" + var candidate = base + var suffix = 2 + while !usedIDs.insert(candidate).inserted { + candidate = "\(base)-\(suffix)" + suffix += 1 + } + return candidate + } + + private static func approximateNextMonthResetDate(now: Date) -> Date? { + // GitHub's budget response does not expose a reset instant. Use the local + // start of next month as a display-only approximation. + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = .current + let components = calendar.dateComponents([.year, .month], from: now) + guard let monthStart = calendar.date(from: DateComponents( + year: components.year, + month: components.month, + day: 1)) + else { + return nil + } + return calendar.date(byAdding: .month, value: 1, to: monthStart) + } + + private static func slug(_ value: String) -> String { + var result = "" + var lastWasDash = false + for scalar in value.lowercased().unicodeScalars { + if CharacterSet.alphanumerics.contains(scalar) { + result.unicodeScalars.append(scalar) + lastWasDash = false + } else if !lastWasDash { + result.append("-") + lastWasDash = true + } + } + return result.trimmingCharacters(in: CharacterSet(charactersIn: "-")) + } +} + +#if os(macOS) +private enum CopilotGitHubCookieImporter { + private static let cookieClient = BrowserCookieClient() + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.copilot]?.browserCookieOrder ?? [.chrome] + private static let sessionCookieNames: Set = [ + "user_session", + "__Host-user_session_same_site", + "_gh_sess", + "logged_in", + "dotcom_user", + ] + + struct SessionInfo { + let cookies: [HTTPCookie] + let sourceLabel: String + + var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + } + + static func importSessions(browserDetection: BrowserDetection) -> [SessionInfo] { + let installedBrowsers = self.cookieImportOrder.cookieImportCandidates(using: browserDetection) + return installedBrowsers.flatMap { browser -> [SessionInfo] in + do { + return try self.importSessions(from: browser) + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + return [] + } + } + } + + private static func importSessions(from browser: Browser) throws -> [SessionInfo] { + let query = BrowserCookieQuery(domains: ["github.com", "www.github.com"]) + let sources = try self.cookieClient.codexBarRecords(matching: query, in: browser) + return sources.compactMap { source -> SessionInfo? in + guard !source.records.isEmpty else { return nil } + let cookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin) + guard cookies.contains(where: { self.sessionCookieNames.contains($0.name) }) else { + return nil + } + return SessionInfo(cookies: cookies, sourceLabel: source.label) + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Copilot/CopilotDeviceFlow.swift b/Sources/CodexBarCore/Providers/Copilot/CopilotDeviceFlow.swift new file mode 100644 index 0000000..ed00333 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Copilot/CopilotDeviceFlow.swift @@ -0,0 +1,173 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct CopilotDeviceFlow: Sendable { + public static let defaultHost = "github.com" + + private let clientID = "Iv1.b507a08c87ecfe98" // VS Code Client ID + private let scopes = "read:user" + private let host: String + + public struct DeviceCodeResponse: Decodable, Sendable { + public let deviceCode: String + public let userCode: String + public let verificationUri: String + public let verificationUriComplete: String? + public let expiresIn: Int + public let interval: Int + + public var verificationURLToOpen: String { + self.verificationUriComplete ?? self.verificationUri + } + + enum CodingKeys: String, CodingKey { + case deviceCode = "device_code" + case userCode = "user_code" + case verificationUri = "verification_uri" + case verificationUriComplete = "verification_uri_complete" + case expiresIn = "expires_in" + case interval + } + } + + public struct AccessTokenResponse: Decodable, Sendable { + public let accessToken: String + public let tokenType: String + public let scope: String + + enum CodingKeys: String, CodingKey { + case accessToken = "access_token" + case tokenType = "token_type" + case scope + } + } + + public init(enterpriseHost: String? = nil) { + self.host = Self.normalizedHost(enterpriseHost) + } + + public var deviceCodeURL: URL? { + Self.makeRequestURL(host: self.host, path: "/login/device/code") + } + + public var accessTokenURL: URL? { + Self.makeRequestURL(host: self.host, path: "/login/oauth/access_token") + } + + public static func normalizedHost(_ raw: String?) -> String { + guard var host = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !host.isEmpty else { + return self.defaultHost + } + let componentsValue = host.contains("://") ? host : "https://\(host)" + if let components = URLComponents(string: componentsValue), + let parsedHost = components.host, + !parsedHost.isEmpty + { + host = parsedHost + if let port = components.port { + host += ":\(port)" + } + } else { + if host.hasPrefix("https://") { + host.removeFirst("https://".count) + } else if host.hasPrefix("http://") { + host.removeFirst("http://".count) + } + host = host.split(separator: "/", maxSplits: 1).first.map(String.init) ?? host + } + let normalized = host.trimmingCharacters(in: CharacterSet(charactersIn: ".")).lowercased() + return normalized.isEmpty ? Self.defaultHost : normalized + } + + public func requestDeviceCode() async throws -> DeviceCodeResponse { + guard let deviceCodeURL = self.deviceCodeURL else { + throw URLError(.badURL) + } + let request = URLRequest(url: deviceCodeURL) + + var postRequest = request + postRequest.httpMethod = "POST" + postRequest.setValue("application/json", forHTTPHeaderField: "Accept") + postRequest.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + + let body = [ + "client_id": self.clientID, + "scope": self.scopes, + ] + postRequest.httpBody = Self.formURLEncodedBody(body) + + let response = try await ProviderHTTPClient.shared.response(for: postRequest) + + guard response.statusCode == 200 else { + throw URLError(.badServerResponse) + } + + return try JSONDecoder().decode(DeviceCodeResponse.self, from: response.data) + } + + public func pollForToken(deviceCode: String, interval: Int) async throws -> String { + guard let accessTokenURL = self.accessTokenURL else { + throw URLError(.badURL) + } + var request = URLRequest(url: accessTokenURL) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + + let body = [ + "client_id": self.clientID, + "device_code": deviceCode, + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + ] + request.httpBody = Self.formURLEncodedBody(body) + + while true { + try await Task.sleep(nanoseconds: UInt64(interval) * 1_000_000_000) + try Task.checkCancellation() + + let response = try await ProviderHTTPClient.shared.response(for: request) + + // Check for error in JSON + if let json = try? JSONSerialization.jsonObject(with: response.data) as? [String: Any], + let error = json["error"] as? String + { + if error == "authorization_pending" { + continue + } + if error == "slow_down" { + try await Task.sleep(nanoseconds: 5_000_000_000) // Add 5s + continue + } + if error == "expired_token" { + throw URLError(.timedOut) + } + throw URLError(.userAuthenticationRequired) // Generic failure + } + + if let tokenResponse = try? JSONDecoder().decode(AccessTokenResponse.self, from: response.data) { + return tokenResponse.accessToken + } + } + } + + private static func formURLEncodedBody(_ parameters: [String: String]) -> Data { + let pairs = parameters + .map { key, value in + "\(Self.formEncode(key))=\(Self.formEncode(value))" + } + .joined(separator: "&") + return Data(pairs.utf8) + } + + static func makeRequestURL(host: String, path: String) -> URL? { + URL(string: "https://\(host)\(path)") + } + + private static func formEncode(_ value: String) -> String { + var allowed = CharacterSet.urlQueryAllowed + allowed.remove(charactersIn: "+&=") + return value.addingPercentEncoding(withAllowedCharacters: allowed) ?? value + } +} diff --git a/Sources/CodexBarCore/Providers/Copilot/CopilotProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Copilot/CopilotProviderDescriptor.swift new file mode 100644 index 0000000..5a3b1f3 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Copilot/CopilotProviderDescriptor.swift @@ -0,0 +1,135 @@ +import Foundation + +public enum CopilotProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .copilot, + metadata: ProviderMetadata( + id: .copilot, + displayName: "Copilot", + sessionLabel: "Premium", + weeklyLabel: "Chat", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Copilot usage", + cliName: "copilot", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.copilotCookieImportOrder, + dashboardURL: "https://github.com/settings/copilot", + statusPageURL: "https://www.githubstatus.com/"), + branding: ProviderBranding( + iconStyle: .copilot, + iconResourceName: "ProviderIcon-copilot", + color: ProviderColor(red: 168 / 255, green: 85 / 255, blue: 247 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Copilot cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [CopilotAPIFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "copilot", + versionDetector: nil)) + } +} + +struct CopilotAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "copilot.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolveToken(context: context) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let token = Self.resolveToken(context: context), !token.isEmpty else { + throw URLError(.userAuthenticationRequired) + } + let fetcher = CopilotUsageFetcher( + token: token, + enterpriseHost: context.settings?.copilot?.enterpriseHost) + let usage = try await fetcher.fetch() + let snap = await self.addBudgetWindowsIfNeeded(to: usage, token: token, context: context) + return self.makeResult( + usage: snap, + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func resolveToken(context: ProviderFetchContext) -> String? { + ProviderTokenResolver.copilotToken(environment: context.env) + ?? ProviderTokenResolver.copilotResolution(environment: [ + "COPILOT_API_TOKEN": context.settings?.copilot?.apiToken ?? "", + ])?.token + } + + private func addBudgetWindowsIfNeeded( + to usage: UsageSnapshot, + token: String, + context: ProviderFetchContext) async -> UsageSnapshot + { + guard let settings = context.settings?.copilot, + settings.budgetExtrasEnabled, + settings.budgetCookieSource != .off + else { return usage } + + let manualCookieHeader = Self.budgetCookieHeaderOverride(from: settings) + if settings.budgetCookieSource == .manual, manualCookieHeader == nil { + return usage + } + do { + let expectedAccountIdentifier = try await self.expectedBudgetAccountIdentifier( + token: token, + settings: settings) + let extraRateWindows = try await CopilotBudgetWebFetcher( + cookieHeaderOverride: manualCookieHeader, + expectedGitHubAccountIdentifier: expectedAccountIdentifier, + browserDetection: context.browserDetection) + .fetchBudgetWindows() + guard !extraRateWindows.isEmpty else { return usage } + return usage.with(extraRateWindows: extraRateWindows) + } catch { + CodexBarLog.logger(LogCategories.providers).warning( + "Copilot budget extras unavailable", + metadata: ["error": "\(error.localizedDescription)"]) + return usage + } + } + + static func budgetCookieHeaderOverride( + from settings: ProviderSettingsSnapshot.CopilotProviderSettings) -> String? + { + guard settings.budgetCookieSource == .manual else { return nil } + return CookieHeaderNormalizer.normalize(settings.manualBudgetCookieHeader) + } + + private func expectedBudgetAccountIdentifier( + token: String, + settings: ProviderSettingsSnapshot.CopilotProviderSettings) async throws -> String + { + let identity = try await CopilotUsageFetcher.fetchGitHubIdentity(token: token) + let tokenIdentifier = CopilotBudgetWebFetcher.normalizedGitHubAccountIdentifier(for: identity) + if let selectedIdentifier = Self.normalizedBudgetAccountIdentifier(settings.selectedAccountExternalIdentifier), + selectedIdentifier != tokenIdentifier.lowercased(), + selectedIdentifier != identity.login.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + { + CodexBarLog.logger(LogCategories.providers).warning( + "Ignoring stale Copilot account identifier") + } + return tokenIdentifier + } + + private static func normalizedBudgetAccountIdentifier(_ identifier: String?) -> String? { + let trimmed = identifier?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + return trimmed.isEmpty ? nil : trimmed.lowercased() + } +} diff --git a/Sources/CodexBarCore/Providers/Copilot/CopilotUsageFetcher.swift b/Sources/CodexBarCore/Providers/Copilot/CopilotUsageFetcher.swift new file mode 100644 index 0000000..f9fe0fb --- /dev/null +++ b/Sources/CodexBarCore/Providers/Copilot/CopilotUsageFetcher.swift @@ -0,0 +1,186 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct CopilotUsageFetcher: Sendable { + public struct GitHubUserIdentity: Decodable, Equatable, Sendable { + public let id: Int64 + public let login: String + + public init(id: Int64, login: String) { + self.id = id + self.login = login + } + } + + private let token: String + private let enterpriseHost: String? + private let transport: any ProviderHTTPTransport + + public init( + token: String, + enterpriseHost: String? = nil, + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) + { + self.token = token + self.enterpriseHost = enterpriseHost + self.transport = transport + } + + public static func apiHost(enterpriseHost: String?) -> String { + let host = CopilotDeviceFlow.normalizedHost(enterpriseHost) + if host == CopilotDeviceFlow.defaultHost { + return "api.github.com" + } + if host.hasPrefix("api.") { + return host + } + return "api.\(host)" + } + + public static func usageURL(enterpriseHost: String?) -> URL? { + CopilotDeviceFlow.makeRequestURL( + host: self.apiHost(enterpriseHost: enterpriseHost), + path: "/copilot_internal/user") + } + + public func fetch() async throws -> UsageSnapshot { + guard let url = Self.usageURL(enterpriseHost: self.enterpriseHost) else { + throw URLError(.badURL) + } + + var request = URLRequest(url: url) + // Use the GitHub OAuth token directly, not the Copilot token. + request.setValue("token \(self.token)", forHTTPHeaderField: "Authorization") + self.addCommonHeaders(to: &request) + + let response = try await self.transport.response(for: request) + + if response.statusCode == 401 || response.statusCode == 403 { + throw URLError(.userAuthenticationRequired) + } + + guard response.statusCode == 200 else { + throw URLError(.badServerResponse) + } + + let usage = try JSONDecoder().decode(CopilotUsageResponse.self, from: response.data) + let resetsAt = Self.parseQuotaResetDate(usage.quotaResetDate) + let premium = Self.makeRateWindow(from: usage.quotaSnapshots.premiumInteractions, resetsAt: resetsAt) + let chat = Self.makeRateWindow(from: usage.quotaSnapshots.chat, resetsAt: resetsAt) + + let primary: RateWindow? + let secondary: RateWindow? + if let premium { + primary = premium + secondary = chat + } else if let chatWindow = chat { + // Keep chat in the secondary slot so provider labels remain accurate + // ("Premium" for primary, "Chat" for secondary) on chat-only plans. + primary = nil + secondary = chatWindow + } else if usage.tokenBasedBilling { + // Copilot Business token-based billing currently exposes zero-entitlement + // placeholder quotas on this endpoint, so surface the plan without fake usage. + primary = nil + secondary = nil + } else { + throw URLError(.cannotDecodeRawData) + } + + let identity = ProviderIdentitySnapshot( + providerID: .copilot, + accountEmail: nil, + accountOrganization: nil, + loginMethod: usage.copilotPlan.capitalized) + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: nil, + providerCost: nil, + updatedAt: Date(), + identity: identity) + } + + public static func fetchGitHubUsername(token: String) async throws -> String { + try await self.fetchGitHubIdentity(token: token).login + } + + public static func fetchGitHubIdentity( + token: String, + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) + async throws -> GitHubUserIdentity + { + guard let url = URL(string: "https://api.github.com/user") else { + throw URLError(.badURL) + } + var request = URLRequest(url: url) + request.setValue("token \(token)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + let response = try await transport.response(for: request) + switch response.statusCode { + case 200: + return try JSONDecoder().decode(GitHubUserIdentity.self, from: response.data) + case 401, 403: + throw URLError(.userAuthenticationRequired) + default: + throw URLError(.badServerResponse) + } + } + + private func addCommonHeaders(to request: inout URLRequest) { + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("vscode/1.96.2", forHTTPHeaderField: "Editor-Version") + request.setValue("copilot-chat/0.26.7", forHTTPHeaderField: "Editor-Plugin-Version") + request.setValue("GitHubCopilotChat/0.26.7", forHTTPHeaderField: "User-Agent") + request.setValue("2025-04-01", forHTTPHeaderField: "X-Github-Api-Version") + } + + static func makeRateWindow( + from snapshot: CopilotUsageResponse.QuotaSnapshot?, + resetsAt: Date? = nil) -> RateWindow? + { + guard let snapshot else { return nil } + guard !snapshot.isPlaceholder else { return nil } + guard snapshot.hasPercentRemaining else { return nil } + let usedPercent = snapshot.usedPercent + let overQuotaDescription = snapshot.overQuotaUsedPercent.map { used in + String(format: "%.0f%% used", used) + } + let effectiveResetsAt = snapshot.unlimited ? nil : resetsAt + + return RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: effectiveResetsAt, + resetDescription: overQuotaDescription) + } + + static func parseQuotaResetDate(_ value: String?) -> Date? { + guard let raw = value?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + + let fractionalISO = ISO8601DateFormatter() + fractionalISO.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = fractionalISO.date(from: raw) { + return date + } + + let internetISO = ISO8601DateFormatter() + internetISO.formatOptions = [.withInternetDateTime] + if let date = internetISO.date(from: raw) { + return date + } + + let formatter = DateFormatter() + formatter.calendar = Calendar(identifier: .gregorian) + formatter.isLenient = false + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = TimeZone(secondsFromGMT: 0) + formatter.dateFormat = "yyyy-MM-dd" + return formatter.date(from: raw) + } +} diff --git a/Sources/CodexBarCore/Providers/Crof/CrofProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Crof/CrofProviderDescriptor.swift new file mode 100644 index 0000000..860733d --- /dev/null +++ b/Sources/CodexBarCore/Providers/Crof/CrofProviderDescriptor.swift @@ -0,0 +1,46 @@ +import Foundation + +public enum CrofProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .crof, + metadata: ProviderMetadata( + id: .crof, + displayName: "Crof", + sessionLabel: "Requests", + weeklyLabel: "Credits", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "Credit balance from the Crof usage API", + toggleTitle: "Show Crof usage", + cliName: "crof", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://crof.ai/dashboard", + statusPageURL: nil, + statusLinkURL: nil), + branding: ProviderBranding( + iconStyle: .crof, + iconResourceName: "ProviderIcon-crof", + color: ProviderColor(red: 0.18, green: 0.67, blue: 0.58)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Crof cost summary is not available via API." }), + fetchPlan: .apiToken( + strategyID: "crof.api", + resolveToken: { ProviderTokenResolver.crofToken(environment: $0) }, + missingCredentialsError: { CrofUsageError.missingCredentials }, + loadUsage: { apiKey, _ in + try await CrofUsageFetcher.fetchUsage(apiKey: apiKey).toUsageSnapshot() + }), + cli: ProviderCLIConfig( + name: "crof", + aliases: ["crofai"], + versionDetector: nil)) + } +} diff --git a/Sources/CodexBarCore/Providers/Crof/CrofSettingsReader.swift b/Sources/CodexBarCore/Providers/Crof/CrofSettingsReader.swift new file mode 100644 index 0000000..14680a9 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Crof/CrofSettingsReader.swift @@ -0,0 +1,27 @@ +import Foundation + +public enum CrofSettingsReader { + public static let apiKeyEnvironmentKeys = ["CROF_API_KEY", "CROFAI_API_KEY"] + + public static func apiKey(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + for key in self.apiKeyEnvironmentKeys { + if let value = self.cleaned(environment[key]) { + return value + } + } + return nil + } + + private static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} diff --git a/Sources/CodexBarCore/Providers/Crof/CrofUsageFetcher.swift b/Sources/CodexBarCore/Providers/Crof/CrofUsageFetcher.swift new file mode 100644 index 0000000..c8c3a4c --- /dev/null +++ b/Sources/CodexBarCore/Providers/Crof/CrofUsageFetcher.swift @@ -0,0 +1,89 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct CrofUsageResponse: Decodable, Sendable { + public let credits: Double + public let requestsPlan: Double + public let usableRequests: Double + + enum CodingKeys: String, CodingKey { + case credits + case requestsPlan = "requests_plan" + case usableRequests = "usable_requests" + } +} + +public enum CrofUsageError: LocalizedError, Sendable, Equatable { + case missingCredentials + case networkError(String) + case apiError(Int) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Crof API key." + case let .networkError(message): + "Crof network error: \(message)" + case let .apiError(statusCode): + "Crof API error: HTTP \(statusCode)" + case let .parseFailed(message): + "Failed to parse Crof response: \(message)" + } + } +} + +public enum CrofUsageFetcher { + public static let usageURL = URL(string: "https://crof.ai/usage_api/")! + private static let timeoutSeconds: TimeInterval = 15 + + public static func fetchUsage( + apiKey: String, + session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> CrofUsageSnapshot + { + let trimmed = apiKey.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { + throw CrofUsageError.missingCredentials + } + + var request = URLRequest(url: self.usageURL) + request.httpMethod = "GET" + request.timeoutInterval = Self.timeoutSeconds + request.setValue("Bearer \(trimmed)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + let response: ProviderHTTPResponse + do { + response = try await transport.response(for: request) + } catch { + throw CrofUsageError.networkError(error.localizedDescription) + } + + guard response.statusCode == 200 else { + throw CrofUsageError.apiError(response.statusCode) + } + + return try self.parseSnapshot(response.data) + } + + static func _parseSnapshotForTesting(_ data: Data) throws -> CrofUsageSnapshot { + try self.parseSnapshot(data) + } + + private static func parseSnapshot(_ data: Data) throws -> CrofUsageSnapshot { + let decoded: CrofUsageResponse + do { + decoded = try JSONDecoder().decode(CrofUsageResponse.self, from: data) + } catch { + throw CrofUsageError.parseFailed(error.localizedDescription) + } + + return CrofUsageSnapshot( + credits: decoded.credits, + requestsPlan: decoded.requestsPlan, + usableRequests: decoded.usableRequests, + updatedAt: Date()) + } +} diff --git a/Sources/CodexBarCore/Providers/Crof/CrofUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Crof/CrofUsageSnapshot.swift new file mode 100644 index 0000000..ee76bcb --- /dev/null +++ b/Sources/CodexBarCore/Providers/Crof/CrofUsageSnapshot.swift @@ -0,0 +1,83 @@ +import Foundation + +public struct CrofUsageSnapshot: Sendable { + private static let requestWindowMinutes = 24 * 60 + private static let resetTimeZone = TimeZone(identifier: "America/Chicago") ?? TimeZone(secondsFromGMT: -5)! + + public let credits: Double + public let requestsPlan: Double + public let usableRequests: Double + public let updatedAt: Date + + public init( + credits: Double, + requestsPlan: Double, + usableRequests: Double, + updatedAt: Date = Date()) + { + self.credits = credits + self.requestsPlan = requestsPlan + self.usableRequests = usableRequests + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let usedPercent: Double + if self.requestsPlan > 0 { + let usableRequests = max(0, min(self.requestsPlan, self.usableRequests)) + let remainingPercent = floor(usableRequests / self.requestsPlan * 100).clamped(to: 0...100) + usedPercent = 100 - remainingPercent + } else { + usedPercent = 100 + } + + let primary = RateWindow( + usedPercent: usedPercent, + windowMinutes: Self.requestWindowMinutes, + resetsAt: Self.nextRequestReset(after: self.updatedAt), + resetDescription: Self.formatRequestsLeft(self.usableRequests)) + + let creditsDetail = Self.formatCredits(self.credits) + let secondary = RateWindow( + // Crof returns a balance but no credit cap, so the bar only indicates present vs. exhausted credits. + usedPercent: self.credits > 0 ? 0 : 100, + windowMinutes: nil, + resetsAt: nil, + resetDescription: creditsDetail) + + return UsageSnapshot( + primary: primary, + secondary: secondary, + providerCost: nil, + updatedAt: self.updatedAt, + identity: ProviderIdentitySnapshot( + providerID: .crof, + accountEmail: nil, + accountOrganization: nil, + loginMethod: "API key")) + } + + private static func formatCredits(_ value: Double) -> String { + let clamped = max(0, value) + let cents = floor(clamped * 100) / 100 + return String(format: "$%.2f", cents) + } + + private static func formatRequestsLeft(_ value: Double) -> String { + let clamped = max(0, value) + let formatted = clamped.rounded() == clamped + ? String(format: "%.0f", clamped) + : String(format: "%.2f", clamped) + return "\(formatted) requests left" + } + + private static func nextRequestReset(after date: Date) -> Date? { + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = Self.resetTimeZone + + let startOfDay = calendar.startOfDay(for: date) + return startOfDay <= date + ? calendar.date(byAdding: .day, value: 1, to: startOfDay) + : startOfDay + } +} diff --git a/Sources/CodexBarCore/Providers/CrossModel/CrossModelProviderDescriptor.swift b/Sources/CodexBarCore/Providers/CrossModel/CrossModelProviderDescriptor.swift new file mode 100644 index 0000000..9a9c465 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CrossModel/CrossModelProviderDescriptor.swift @@ -0,0 +1,63 @@ +import Foundation + +public enum CrossModelProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .crossmodel, + metadata: ProviderMetadata( + id: .crossmodel, + displayName: "CrossModel", + sessionLabel: "Credits", + weeklyLabel: "Usage", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show CrossModel usage", + cliName: "crossmodel", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + dashboardURL: "https://crossmodel.ai/console/usage", + statusPageURL: nil, + statusLinkURL: nil), + branding: ProviderBranding( + iconStyle: .crossmodel, + iconResourceName: "ProviderIcon-crossmodel", + color: ProviderColor(red: 124 / 255, green: 58 / 255, blue: 237 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "CrossModel cost summary is not yet supported." }), + fetchPlan: .apiToken( + strategyID: "crossmodel.api", + resolveToken: { ProviderTokenResolver.crossModelToken(environment: $0) }, + missingCredentialsError: { CrossModelSettingsError.missingToken }, + loadUsage: { apiKey, context in + try await CrossModelUsageFetcher.fetchUsage( + apiKey: apiKey, + environment: context.env, + includeOptionalUsage: context.includeOptionalUsage).toUsageSnapshot() + }), + cli: ProviderCLIConfig( + name: "crossmodel", + aliases: ["cm"], + versionDetector: nil)) + } +} + +/// Errors related to CrossModel settings. +public enum CrossModelSettingsError: LocalizedError, Sendable, Equatable { + case missingToken + case invalidEndpointOverride(String) + + public var errorDescription: String? { + switch self { + case .missingToken: + "CrossModel API token not configured. Set CROSSMODEL_API_KEY environment variable or configure in Settings." + case let .invalidEndpointOverride(key): + "CrossModel endpoint override \(key) must use HTTPS (or a loopback HTTP host)." + } + } +} diff --git a/Sources/CodexBarCore/Providers/CrossModel/CrossModelSettingsReader.swift b/Sources/CodexBarCore/Providers/CrossModel/CrossModelSettingsReader.swift new file mode 100644 index 0000000..3e01237 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CrossModel/CrossModelSettingsReader.swift @@ -0,0 +1,52 @@ +import Foundation + +/// Reads CrossModel settings from environment variables. +public enum CrossModelSettingsReader { + /// Environment variable key for the CrossModel API token. + public static let envKey = "CROSSMODEL_API_KEY" + /// Environment variable key for an optional API base URL override. + public static let urlEnvKey = "CROSSMODEL_API_URL" + + /// Returns the API token from environment if present and non-empty. + public static func apiToken(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + self.cleaned(environment[self.envKey]) + } + + /// Returns the API base URL, defaulting to the production endpoint. + public static func apiURL(environment: [String: String] = ProcessInfo.processInfo.environment) -> URL { + if let override = self.validAPIURL(environment: environment) { + return override + } + return URL(string: "https://api.crossmodel.ai/v1")! + } + + public static func validateEndpointOverrides( + environment: [String: String] = ProcessInfo.processInfo.environment) throws + { + guard let raw = self.cleaned(environment[self.urlEnvKey]) else { return } + // Loopback HTTP is allowed so the provider can be exercised end-to-end + // against a locally running gateway during development. + guard ProviderEndpointOverrideValidator().validatedURLAllowingLoopbackHTTP(raw) == nil else { return } + throw CrossModelSettingsError.invalidEndpointOverride(self.urlEnvKey) + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } + + private static func validAPIURL(environment: [String: String]) -> URL? { + guard let raw = self.cleaned(environment[self.urlEnvKey]) else { return nil } + return ProviderEndpointOverrideValidator().validatedURLAllowingLoopbackHTTP(raw) + } +} diff --git a/Sources/CodexBarCore/Providers/CrossModel/CrossModelUsageStats.swift b/Sources/CodexBarCore/Providers/CrossModel/CrossModelUsageStats.swift new file mode 100644 index 0000000..7eb9242 --- /dev/null +++ b/Sources/CodexBarCore/Providers/CrossModel/CrossModelUsageStats.swift @@ -0,0 +1,417 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +// MARK: - Wire DTOs (CrossModel returns integer micro units; 1 major currency unit = 1_000_000 micro) + +/// `GET /v1/credits` response. +struct CrossModelCreditsResponse: Decodable { + let currency: String + let balanceMicro: Int64 + let uncollectedMicro: Int64 + + private enum CodingKeys: String, CodingKey { + case currency + case balanceMicro = "balance_micro" + case uncollectedMicro = "uncollected_micro" + } +} + +/// One usage window from `GET /v1/usage`. +struct CrossModelUsageWindowDTO: Decodable { + let costMicro: Int64 + let promptTokens: Int64 + let completionTokens: Int64 + let totalTokens: Int64 + let requestCount: Int64 + let successCount: Int64 + + private enum CodingKeys: String, CodingKey { + case costMicro = "cost_micro" + case promptTokens = "prompt_tokens" + case completionTokens = "completion_tokens" + case totalTokens = "total_tokens" + case requestCount = "request_count" + case successCount = "success_count" + } +} + +/// `GET /v1/usage` response. +struct CrossModelUsageResponse: Decodable { + let currency: String + let daily: CrossModelUsageWindowDTO + let weekly: CrossModelUsageWindowDTO + let monthly: CrossModelUsageWindowDTO +} + +// MARK: - Snapshot (currency-normalized, persisted in UsageSnapshot) + +/// One usage window with cost converted to the response currency's major unit. +public struct CrossModelUsageWindow: Codable, Sendable, Equatable { + public let cost: Double + public let promptTokens: Int + public let completionTokens: Int + public let totalTokens: Int + public let requestCount: Int + public let successCount: Int + + public init( + cost: Double, + promptTokens: Int, + completionTokens: Int, + totalTokens: Int, + requestCount: Int, + successCount: Int) + { + self.cost = cost + self.promptTokens = promptTokens + self.completionTokens = completionTokens + self.totalTokens = totalTokens + self.requestCount = requestCount + self.successCount = successCount + } + + init(dto: CrossModelUsageWindowDTO) { + self.cost = CrossModelUsageSnapshot.majorUnits(dto.costMicro) + self.promptTokens = Int(dto.promptTokens) + self.completionTokens = Int(dto.completionTokens) + self.totalTokens = Int(dto.totalTokens) + self.requestCount = Int(dto.requestCount) + self.successCount = Int(dto.successCount) + } +} + +/// Complete CrossModel usage snapshot: wallet balance plus UTC day/week/month spend. +public struct CrossModelUsageSnapshot: Codable, Sendable, Equatable { + public let currency: String + public let balance: Double + public let uncollected: Double + public let daily: CrossModelUsageWindow? + public let weekly: CrossModelUsageWindow? + public let monthly: CrossModelUsageWindow? + public let updatedAt: Date + + public init( + currency: String, + balance: Double, + uncollected: Double, + daily: CrossModelUsageWindow?, + weekly: CrossModelUsageWindow?, + monthly: CrossModelUsageWindow?, + updatedAt: Date) + { + self.currency = Self.normalizedCurrency(currency) + self.balance = balance + self.uncollected = uncollected + self.daily = daily + self.weekly = weekly + self.monthly = monthly + self.updatedAt = updatedAt + } + + static func majorUnits(_ micro: Int64) -> Double { + Double(micro) / 1_000_000.0 + } + + static func normalizedCurrency(_ raw: String) -> String { + raw.trimmingCharacters(in: .whitespacesAndNewlines).uppercased() + } + + /// Formatted balance for provider-owned menu/status display. + public var balanceDisplay: String { + self.currencyString(self.balance) + } + + public func currencyString(_ value: Double) -> String { + UsageFormatter.currencyString(value, currencyCode: self.currency) + } +} + +extension CrossModelUsageSnapshot { + public func toUsageSnapshot() -> UsageSnapshot { + let identity = ProviderIdentitySnapshot( + providerID: .crossmodel, + accountEmail: nil, + accountOrganization: nil, + loginMethod: "API key") + + return UsageSnapshot( + primary: nil, + secondary: nil, + tertiary: nil, + crossModelUsage: self, + updatedAt: self.updatedAt, + identity: identity, + dataConfidence: .exact) + } +} + +// MARK: - Fetcher + +/// Fetches balance + usage stats from the CrossModel API. +public struct CrossModelUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.crossModelUsage) + private static let creditsRequestTimeoutSeconds: TimeInterval = 15 + private static let usageRequestTimeoutSeconds: TimeInterval = 3 + private static let maxErrorBodyLength = 240 + + /// Fetches balance (required) and usage windows (best-effort) using the provided API key. + public static func fetchUsage( + apiKey: String, + environment: [String: String] = ProcessInfo.processInfo.environment, + includeOptionalUsage: Bool = true, + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + usageJoinGrace: Duration = .seconds(3)) async throws -> CrossModelUsageSnapshot + { + guard !apiKey.isEmpty else { + throw CrossModelUsageError.invalidCredentials + } + try CrossModelSettingsReader.validateEndpointOverrides(environment: environment) + + let baseURL = CrossModelSettingsReader.apiURL(environment: environment) + let credits = try await self.fetchCredits(apiKey: apiKey, baseURL: baseURL, transport: transport) + let currency = try self.normalizedResponseCurrency(credits.currency, endpoint: "credits") + + guard includeOptionalUsage else { + return CrossModelUsageSnapshot( + currency: currency, + balance: CrossModelUsageSnapshot.majorUnits(credits.balanceMicro), + uncollected: CrossModelUsageSnapshot.majorUnits(credits.uncollectedMicro), + daily: nil, + weekly: nil, + monthly: nil, + updatedAt: Date()) + } + + // Usage windows are best-effort: a slow or failing /usage call should not + // block the balance the user came to see. + let usage = try await self.fetchUsageWindows( + apiKey: apiKey, + baseURL: baseURL, + transport: transport, + joinGrace: usageJoinGrace) + let validatedUsage: CrossModelUsageResponse? + if let usage { + let usageCurrency: String + do { + usageCurrency = try self.normalizedResponseCurrency(usage.currency, endpoint: "usage") + } catch { + Self.log.error("CrossModel /usage omitted: invalid currency") + validatedUsage = nil + return CrossModelUsageSnapshot( + currency: currency, + balance: CrossModelUsageSnapshot.majorUnits(credits.balanceMicro), + uncollected: CrossModelUsageSnapshot.majorUnits(credits.uncollectedMicro), + daily: nil, + weekly: nil, + monthly: nil, + updatedAt: Date()) + } + guard usageCurrency == currency else { + Self.log.error("CrossModel currency mismatch: credits=\(currency) usage=\(usageCurrency)") + return CrossModelUsageSnapshot( + currency: currency, + balance: CrossModelUsageSnapshot.majorUnits(credits.balanceMicro), + uncollected: CrossModelUsageSnapshot.majorUnits(credits.uncollectedMicro), + daily: nil, + weekly: nil, + monthly: nil, + updatedAt: Date()) + } + validatedUsage = usage + } else { + validatedUsage = nil + } + + return CrossModelUsageSnapshot( + currency: currency, + balance: CrossModelUsageSnapshot.majorUnits(credits.balanceMicro), + uncollected: CrossModelUsageSnapshot.majorUnits(credits.uncollectedMicro), + daily: validatedUsage.map { CrossModelUsageWindow(dto: $0.daily) }, + weekly: validatedUsage.map { CrossModelUsageWindow(dto: $0.weekly) }, + monthly: validatedUsage.map { CrossModelUsageWindow(dto: $0.monthly) }, + updatedAt: Date()) + } + + private static func fetchCredits( + apiKey: String, + baseURL: URL, + transport: any ProviderHTTPTransport) async throws -> CrossModelCreditsResponse + { + let url = baseURL.appendingPathComponent("credits") + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.timeoutInterval = Self.creditsRequestTimeoutSeconds + + let response = try await transport.response(for: request) + guard response.statusCode == 200 else { + let summary = LogRedactor.redact(Self.sanitizedResponseBodySummary(response.data)) + Self.log.error("CrossModel /credits returned \(response.statusCode): \(summary)") + if response.statusCode == 401 || response.statusCode == 403 { + throw CrossModelUsageError.invalidCredentials + } + throw CrossModelUsageError.apiError("HTTP \(response.statusCode)") + } + try self.validateSameOrigin(response: response, request: request, endpoint: "credits") + + do { + return try JSONDecoder().decode(CrossModelCreditsResponse.self, from: response.data) + } catch let error as DecodingError { + Self.log.error("CrossModel /credits decode error: \(error.localizedDescription)") + throw CrossModelUsageError.parseFailed(error.localizedDescription) + } + } + + private static func fetchUsageWindows( + apiKey: String, + baseURL: URL, + transport: any ProviderHTTPTransport, + joinGrace: Duration) async throws -> CrossModelUsageResponse? + { + let url = baseURL.appendingPathComponent("usage") + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.timeoutInterval = Self.usageRequestTimeoutSeconds + + let sourceTask = Task { + let response = try await transport.response(for: request) + guard response.statusCode == 200 else { + Self.log.debug("CrossModel /usage enrichment returned \(response.statusCode)") + return nil + } + try self.validateSameOrigin(response: response, request: request, endpoint: "usage") + return try JSONDecoder().decode(CrossModelUsageResponse.self, from: response.data) + } + let join = BoundedTaskJoin(sourceTask: sourceTask) + let outcome = await join.value(joinGrace: joinGrace) + do { + switch outcome { + case let .value(response): + return response + case .timedOut: + Self.log.debug("Timed out fetching CrossModel /usage enrichment") + return nil + case let .failure(error): + throw error + } + } catch is CancellationError { + throw CancellationError() + } catch let error as URLError where error.code == .cancelled { + throw CancellationError() + } catch { + if Task.isCancelled { + throw CancellationError() + } + Self.log.debug("Failed to fetch CrossModel /usage enrichment: \(error.localizedDescription)") + return nil + } + } + + private static func normalizedResponseCurrency(_ raw: String, endpoint: String) throws -> String { + let currency = CrossModelUsageSnapshot.normalizedCurrency(raw) + guard !currency.isEmpty else { + throw CrossModelUsageError.parseFailed("Missing CrossModel /\(endpoint) currency") + } + return currency + } + + private static func validateSameOrigin( + response: ProviderHTTPResponse, + request: URLRequest, + endpoint: String) throws + { + guard let requestURL = request.url, + let responseURL = response.response.url, + let requestHost = requestURL.host?.lowercased(), + let responseHost = responseURL.host?.lowercased(), + requestURL.scheme?.lowercased() == responseURL.scheme?.lowercased(), + requestHost == responseHost, + self.effectivePort(for: requestURL) == self.effectivePort(for: responseURL) + else { + throw CrossModelUsageError.apiError("CrossModel /\(endpoint) redirected to a different origin") + } + } + + private static func effectivePort(for url: URL) -> Int? { + if let port = url.port { + return port + } + switch url.scheme?.lowercased() { + case "https": + return 443 + case "http": + return 80 + default: + return nil + } + } + + private static func sanitizedResponseBodySummary(_ data: Data) -> String { + guard !data.isEmpty else { return "empty body" } + guard let rawBody = String(bytes: data, encoding: .utf8) else { + return "non-text body (\(data.count) bytes)" + } + + let body = Self.redactSensitiveBodyContent(rawBody) + .replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespacesAndNewlines) + + guard !body.isEmpty else { return "non-text body (\(data.count) bytes)" } + guard body.count > Self.maxErrorBodyLength else { return body } + + let index = body.index(body.startIndex, offsetBy: Self.maxErrorBodyLength) + return "\(body[.. String { + let replacements: [(String, String)] = [ + (#"(?i)(bearer\s+)[A-Za-z0-9._\-]+"#, "$1[REDACTED]"), + (#"(?i)(cm-)[A-Za-z0-9._\-]+"#, "$1[REDACTED]"), + ( + #"(?i)(\"(?:api_?key|authorization|token|access_token|refresh_token)\"\s*:\s*\")([^\"]+)(\")"#, + "$1[REDACTED]$3"), + ( + #"(?i)((?:api_?key|authorization|token|access_token|refresh_token)\s*[=:]\s*)([^,\s]+)"#, + "$1[REDACTED]"), + ] + + return replacements.reduce(text) { partial, replacement in + partial.replacingOccurrences( + of: replacement.0, + with: replacement.1, + options: .regularExpression) + } + } + + #if DEBUG + static func _sanitizedResponseBodySummaryForTesting(_ body: String) -> String { + self.sanitizedResponseBodySummary(Data(body.utf8)) + } + #endif +} + +/// Errors that can occur during CrossModel usage fetching. +public enum CrossModelUsageError: LocalizedError, Sendable, Equatable { + case invalidCredentials + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .invalidCredentials: + "Invalid CrossModel API credentials" + case let .networkError(message): + "CrossModel network error: \(message)" + case let .apiError(message): + "CrossModel API error: \(message)" + case let .parseFailed(message): + "Failed to parse CrossModel response: \(message)" + } + } +} diff --git a/Sources/CodexBarCore/Providers/Cursor/CursorProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Cursor/CursorProviderDescriptor.swift new file mode 100644 index 0000000..983cae2 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Cursor/CursorProviderDescriptor.swift @@ -0,0 +1,70 @@ +import Foundation + +public enum CursorProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .cursor, + metadata: ProviderMetadata( + id: .cursor, + displayName: "Cursor", + sessionLabel: "Total", + weeklyLabel: "Auto", + opusLabel: "API", + supportsOpus: true, + supportsCredits: true, + creditsHint: "On-demand usage beyond included plan limits.", + toggleTitle: "Show Cursor usage", + cliName: "cursor", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.cursorCookieImportOrder + ?? ProviderBrowserCookieDefaults.defaultImportOrder, + dashboardURL: "https://cursor.com/dashboard?tab=usage", + statusPageURL: "https://status.cursor.com", + statusLinkURL: nil), + branding: ProviderBranding( + iconStyle: .cursor, + iconResourceName: "ProviderIcon-cursor", + color: ProviderColor(red: 0 / 255, green: 191 / 255, blue: 165 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Cursor cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .cli, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [CursorStatusFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "cursor", + versionDetector: nil)) + } +} + +struct CursorStatusFetchStrategy: ProviderFetchStrategy { + let id: String = "cursor.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + guard context.settings?.cursor?.cookieSource != .off else { return false } + return true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let probe = CursorStatusProbe(browserDetection: context.browserDetection) + let manual = Self.manualCookieHeader(from: context) + let snap = try await probe.fetch(cookieHeaderOverride: manual) + return self.makeResult( + usage: snap.toUsageSnapshot(), + sourceLabel: "web") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func manualCookieHeader(from context: ProviderFetchContext) -> String? { + guard context.settings?.cursor?.cookieSource == .manual else { return nil } + return CookieHeaderNormalizer.normalize(context.settings?.cursor?.manualCookieHeader) + } +} diff --git a/Sources/CodexBarCore/Providers/Cursor/CursorRequestUsage.swift b/Sources/CodexBarCore/Providers/Cursor/CursorRequestUsage.swift new file mode 100644 index 0000000..7bba64f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Cursor/CursorRequestUsage.swift @@ -0,0 +1,23 @@ +import Foundation + +/// Request usage snapshot for legacy Cursor plans (request-based instead of token-based). +public struct CursorRequestUsage: Codable, Sendable { + /// Requests used this billing cycle + public let used: Int + /// Request limit (e.g., 500 for legacy enterprise plans) + public let limit: Int + + public init(used: Int, limit: Int) { + self.used = used + self.limit = limit + } + + public var usedPercent: Double { + guard self.limit > 0 else { return 0 } + return (Double(self.used) / Double(self.limit)) * 100 + } + + public var remainingPercent: Double { + max(0, 100 - self.usedPercent) + } +} diff --git a/Sources/CodexBarCore/Providers/Cursor/CursorStatusProbe.swift b/Sources/CodexBarCore/Providers/Cursor/CursorStatusProbe.swift new file mode 100644 index 0000000..bff2b00 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Cursor/CursorStatusProbe.swift @@ -0,0 +1,1475 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif +import SweetCookieKit +#if canImport(SQLite3) +import SQLite3 +#elseif canImport(CSQLite3) +import CSQLite3 +#endif + +#if os(macOS) || os(Linux) + +#if os(macOS) +private let cursorCookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.cursor]?.browserCookieOrder ?? Browser.defaultImportOrder +#endif + +#if os(macOS) + +// MARK: - Cursor Cookie Importer + +/// Imports Cursor session cookies from browser cookies. +public enum CursorCookieImporter { + private static let cookieClient = BrowserCookieClient() + private static let sessionCookieNames: Set = [ + "WorkosCursorSessionToken", + "__Secure-next-auth.session-token", + "next-auth.session-token", + // WorkOS AuthKit (common default; configurable server-side) + "wos-session", + "__Secure-wos-session", + // Auth.js v5 + "authjs.session-token", + "__Secure-authjs.session-token", + ] + + /// Hosts whose cookies may authenticate Cursor web/API requests. + private static let cookieDomains = [ + "cursor.com", + "www.cursor.com", + "cursor.sh", + "authenticator.cursor.sh", + ] + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + } + + /// Reads Cursor session cookies from one browser if present (no fallback to other browsers). + static func importSessionIfPresent( + browser: Browser, + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) -> SessionInfo? + { + self.importSessionsIfPresent( + browser: browser, + browserDetection: browserDetection, + logger: logger).first + } + + /// Reads all Cursor session-cookie candidates from one browser source order. + static func importSessionsIfPresent( + browser: Browser, + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) -> [SessionInfo] + { + self.importCookiesFromBrowser( + browser: browser, + browserDetection: browserDetection, + requireKnownSessionName: true, + logger: logger) + } + + /// Like ``importSessionIfPresent`` but accepts any non-empty cookie set for Cursor domains so the API can validate + /// (used after the strict name pass fails — e.g. new cookie names or host-only cookies). + static func importDomainCookiesIfPresent( + browser: Browser, + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) -> SessionInfo? + { + self.importDomainCookieSessionsIfPresent( + browser: browser, + browserDetection: browserDetection, + logger: logger).first + } + + /// Reads fallback cookie candidates whose names are not already covered by the strict session-cookie pass. + static func importDomainCookieSessionsIfPresent( + browser: Browser, + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) -> [SessionInfo] + { + self.importCookiesFromBrowser( + browser: browser, + browserDetection: browserDetection, + requireKnownSessionName: false, + logger: logger) + } + + private static func importCookiesFromBrowser( + browser: Browser, + browserDetection: BrowserDetection, + requireKnownSessionName: Bool, + logger: ((String) -> Void)?) -> [SessionInfo] + { + let log: (String) -> Void = { msg in logger?("[cursor-cookie] \(msg)") } + guard browserDetection.isCookieSourceAvailable(browser) else { return [] } + guard BrowserCookieAccessGate.shouldAttempt(browser) else { return [] } + + do { + let query = BrowserCookieQuery(domains: Self.cookieDomains) + let sources = try Self.cookieClient.codexBarRecords( + matching: query, + in: browser, + logger: log) + var sessions: [SessionInfo] = [] + for source in sources where !source.records.isEmpty { + let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin) + let hasNamedSession = httpCookies.contains(where: { Self.sessionCookieNames.contains($0.name) }) + if hasNamedSession { + log("Found \(httpCookies.count) Cursor cookies in \(source.label)") + if requireKnownSessionName { + sessions.append(SessionInfo(cookies: httpCookies, sourceLabel: source.label)) + } + continue + } + if !requireKnownSessionName, !httpCookies.isEmpty { + log( + "Found \(httpCookies.count) Cursor domain cookies in \(source.label) " + + "(no known session name); will validate via API") + sessions.append(SessionInfo( + cookies: httpCookies, + sourceLabel: "\(source.label) (domain cookies)")) + continue + } + log("\(source.label) cookies found, but no Cursor session cookie present") + } + return sessions + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + log("\(browser.displayName) cookie import failed: \(error.localizedDescription)") + } + return [] + } + + /// Attempts to import Cursor cookies using the standard browser import order. + public static func importSession( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) throws -> SessionInfo + { + let installedBrowsers = cursorCookieImportOrder.cookieImportCandidates(using: browserDetection) + for browserSource in installedBrowsers { + if let session = Self.importSessionsIfPresent( + browser: browserSource, + browserDetection: browserDetection, + logger: logger).first + { + return session + } + } + for browserSource in installedBrowsers { + if let session = Self.importDomainCookieSessionsIfPresent( + browser: browserSource, + browserDetection: browserDetection, + logger: logger).first + { + return session + } + } + + throw CursorStatusProbeError.noSessionCookie + } + + /// Check if Cursor session cookies are available + public static func hasSession(browserDetection: BrowserDetection, logger: ((String) -> Void)? = nil) -> Bool { + do { + let session = try self.importSession(browserDetection: browserDetection, logger: logger) + return !session.cookies.isEmpty + } catch { + return false + } + } +} +#endif + +// MARK: - Cursor API Models + +public struct CursorUsageSummary: Codable, Sendable { + public let billingCycleStart: String? + public let billingCycleEnd: String? + public let membershipType: String? + public let limitType: String? + public let isUnlimited: Bool? + public let autoModelSelectedDisplayMessage: String? + public let namedModelSelectedDisplayMessage: String? + public let individualUsage: CursorIndividualUsage? + public let teamUsage: CursorTeamUsage? +} + +public struct CursorIndividualUsage: Codable, Sendable { + public let plan: CursorPlanUsage? + public let onDemand: CursorOnDemandUsage? + /// Enterprise / team-member personal cap. Reported by Cursor when the account is part of a team or + /// enterprise plan with an individual quota. Values follow the same cents-based units as `plan`. + public let overall: CursorOverallUsage? + + public init( + plan: CursorPlanUsage? = nil, + onDemand: CursorOnDemandUsage? = nil, + overall: CursorOverallUsage? = nil) + { + self.plan = plan + self.onDemand = onDemand + self.overall = overall + } +} + +/// Personal cap reported under `individualUsage.overall` for Enterprise/Team members. +/// Mirrors the shape of `CursorOnDemandUsage`; values are in cents. +public struct CursorOverallUsage: Codable, Sendable { + public let enabled: Bool? + /// Usage in cents (e.g., 7384 = $73.84) + public let used: Int? + /// Limit in cents (e.g., 10000 = $100.00). `nil` indicates the API omitted a numeric cap. + public let limit: Int? + /// Remaining in cents. + public let remaining: Int? + + public init(enabled: Bool? = nil, used: Int? = nil, limit: Int? = nil, remaining: Int? = nil) { + self.enabled = enabled + self.used = used + self.limit = limit + self.remaining = remaining + } +} + +public struct CursorPlanUsage: Codable, Sendable { + public let enabled: Bool? + /// Usage in cents (e.g., 2000 = $20.00) + public let used: Int? + /// Limit in cents (e.g., 2000 = $20.00) + public let limit: Int? + /// Remaining in cents + public let remaining: Int? + public let breakdown: CursorPlanBreakdown? + public let autoPercentUsed: Double? + public let apiPercentUsed: Double? + public let totalPercentUsed: Double? +} + +public struct CursorPlanBreakdown: Codable, Sendable { + public let included: Int? + public let bonus: Int? + public let total: Int? +} + +public struct CursorOnDemandUsage: Codable, Sendable { + public let enabled: Bool? + /// Usage in cents + public let used: Int? + /// Limit in cents (nil if unlimited) + public let limit: Int? + /// Remaining in cents (nil if unlimited) + public let remaining: Int? +} + +public struct CursorTeamUsage: Codable, Sendable { + public let onDemand: CursorOnDemandUsage? + /// Shared team/enterprise pool counted across all members. Same cents-based units as the other usage blocks. + public let pooled: CursorPooledUsage? + + public init(onDemand: CursorOnDemandUsage? = nil, pooled: CursorPooledUsage? = nil) { + self.onDemand = onDemand + self.pooled = pooled + } +} + +/// Shared team/enterprise pool reported under `teamUsage.pooled`. Values are in cents. +public struct CursorPooledUsage: Codable, Sendable { + public let enabled: Bool? + /// Pool usage in cents. + public let used: Int? + /// Pool limit in cents. `nil` indicates an unlimited or unreported pool. + public let limit: Int? + /// Pool remaining in cents. + public let remaining: Int? + + public init(enabled: Bool? = nil, used: Int? = nil, limit: Int? = nil, remaining: Int? = nil) { + self.enabled = enabled + self.used = used + self.limit = limit + self.remaining = remaining + } +} + +// MARK: - Cursor Usage API Models (Legacy Request-Based Plans) + +/// Response from `/api/usage?user=ID` endpoint for legacy request-based plans. +public struct CursorUsageResponse: Codable, Sendable { + public let gpt4: CursorModelUsage? + public let startOfMonth: String? + + enum CodingKeys: String, CodingKey { + case gpt4 = "gpt-4" + case startOfMonth + } +} + +public struct CursorModelUsage: Codable, Sendable { + public let numRequests: Int? + public let numRequestsTotal: Int? + public let numTokens: Int? + public let maxRequestUsage: Int? + public let maxTokenUsage: Int? +} + +public struct CursorUserInfo: Codable, Sendable { + public let email: String? + public let emailVerified: Bool? + public let name: String? + public let sub: String? + public let createdAt: String? + public let updatedAt: String? + public let picture: String? + + enum CodingKeys: String, CodingKey { + case email + case emailVerified = "email_verified" + case name + case sub + case createdAt = "created_at" + case updatedAt = "updated_at" + case picture + } +} + +// MARK: - Cursor App Auth + +struct CursorAppAuthSession: Equatable { + let accessToken: String + + var isUsable: Bool { + guard !self.accessToken.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty, + (try? self.userID()) != nil, + let expiresAt = try? self.expiresAt() + else { + return false + } + return expiresAt.timeIntervalSinceNow > 60 + } + + func cookieHeader() throws -> String { + try "WorkosCursorSessionToken=\(self.userID())%3A%3A\(self.accessToken)" + } + + func userID() throws -> String { + let json = try self.payload() + guard let subject = json["sub"] as? String, + let userID = subject.split(separator: "|", omittingEmptySubsequences: true).last.map(String.init), + !userID.isEmpty + else { + throw CursorStatusProbeError.parseFailed("Cursor.app access token is missing a user ID") + } + + let allowed = CharacterSet.alphanumerics.union(CharacterSet(charactersIn: "._-")) + guard userID.unicodeScalars.allSatisfy(allowed.contains) else { + throw CursorStatusProbeError.parseFailed("Cursor.app access token has an invalid user ID") + } + + return userID + } + + private func expiresAt() throws -> Date { + let json = try self.payload() + guard let expiration = json["exp"] as? NSNumber else { + throw CursorStatusProbeError.parseFailed("Cursor.app access token is missing an expiration") + } + return Date(timeIntervalSince1970: expiration.doubleValue) + } + + private func payload() throws -> [String: Any] { + let parts = self.accessToken.split(separator: ".", omittingEmptySubsequences: false) + guard parts.count >= 2 else { + throw CursorStatusProbeError.parseFailed("Cursor.app access token is not a JWT") + } + + var payload = String(parts[1]) + .replacingOccurrences(of: "-", with: "+") + .replacingOccurrences(of: "_", with: "/") + payload += String(repeating: "=", count: (4 - payload.count % 4) % 4) + + guard let data = Data(base64Encoded: payload), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + throw CursorStatusProbeError.parseFailed("Cursor.app access token has an invalid payload") + } + + return json + } +} + +protocol CursorAppAuthSessionProviding: Sendable { + func loadSession() throws -> CursorAppAuthSession? +} + +struct CursorAppAuthStore: CursorAppAuthSessionProviding { + private static let defaultDBPath: String = Self.resolveDefaultDBPath() + + private let dbPath: String + + init(dbPath: String? = nil) { + self.dbPath = dbPath ?? Self.defaultDBPath + } + + static func resolveDefaultDBPath( + home: String = NSHomeDirectory(), + environment: [String: String] = ProcessInfo.processInfo.environment, + fileManager: FileManager = .default) -> String + { + #if os(macOS) + _ = environment + _ = fileManager + return "\(home)/Library/Application Support/Cursor/User/globalStorage/state.vscdb" + #elseif os(Linux) + let configHome = environment[CodexBarConfigStore.xdgConfigHomeEnvironmentKey]? + .trimmingCharacters(in: .whitespacesAndNewlines) + let expandedConfigHome = configHome.map { ($0 as NSString).expandingTildeInPath } + let base: String = if let expandedConfigHome, + !expandedConfigHome.isEmpty, + (expandedConfigHome as NSString).isAbsolutePath + { + expandedConfigHome + } else { + "\(home)/.config" + } + return "\(base)/Cursor/User/globalStorage/state.vscdb" + #else + _ = home + _ = environment + _ = fileManager + return "" + #endif + } + + func loadSession() throws -> CursorAppAuthSession? { + guard FileManager.default.fileExists(atPath: self.dbPath) else { return nil } + + guard let accessToken = try self.value(for: "cursorAuth/accessToken"), + !accessToken.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + else { + return nil + } + + return CursorAppAuthSession(accessToken: accessToken) + } + + private func value(for key: String) throws -> String? { + var db: OpaquePointer? + guard sqlite3_open_v2(self.dbPath, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { + let message = db.flatMap { String(cString: sqlite3_errmsg($0)) } ?? "unknown error" + sqlite3_close(db) + throw CursorStatusProbeError.networkError("SQLite error reading Cursor app auth: \(message)") + } + defer { sqlite3_close(db) } + sqlite3_busy_timeout(db, 250) + + let query = "SELECT value FROM ItemTable WHERE key = ? LIMIT 1;" + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, query, -1, &stmt, nil) == SQLITE_OK else { + let message = db.flatMap { String(cString: sqlite3_errmsg($0)) } ?? "unknown error" + throw CursorStatusProbeError.networkError("SQLite error preparing Cursor app auth read: \(message)") + } + defer { sqlite3_finalize(stmt) } + + sqlite3_bind_text(stmt, 1, key, -1, SQLITE_TRANSIENT) + let stepResult = sqlite3_step(stmt) + guard stepResult == SQLITE_ROW else { + if stepResult == SQLITE_DONE { return nil } + let message = db.flatMap { String(cString: sqlite3_errmsg($0)) } ?? "unknown error" + throw CursorStatusProbeError.networkError("SQLite error reading Cursor app auth: \(message)") + } + + return Self.decodeSQLiteValue(stmt: stmt, index: 0) + } + + private static func decodeSQLiteValue(stmt: OpaquePointer?, index: Int32) -> String? { + switch sqlite3_column_type(stmt, index) { + case SQLITE_TEXT: + guard let c = sqlite3_column_text(stmt, index) else { return nil } + return String(cString: c) + case SQLITE_BLOB: + guard let bytes = sqlite3_column_blob(stmt, index) else { return nil } + let data = Data(bytes: bytes, count: Int(sqlite3_column_bytes(stmt, index))) + return String(data: data, encoding: .utf8) + ?? String(data: data, encoding: .utf16LittleEndian) + default: + return nil + } + } +} + +private let SQLITE_TRANSIENT = unsafeBitCast(-1, to: sqlite3_destructor_type.self) + +// MARK: - Cursor Status Snapshot + +public struct CursorStatusSnapshot: Sendable { + /// Percentage of included plan usage (0-100) — the "Total" headline number from Cursor's UI + public let planPercentUsed: Double + /// Auto + Composer usage percent (0-100), nil when not available + public let autoPercentUsed: Double? + /// API (named model) usage percent (0-100), nil when not available + public let apiPercentUsed: Double? + /// Included plan usage in USD + public let planUsedUSD: Double + /// Included plan limit in USD + public let planLimitUSD: Double + /// On-demand usage in USD + public let onDemandUsedUSD: Double + /// On-demand limit in USD (nil if unlimited) + public let onDemandLimitUSD: Double? + /// Team on-demand usage in USD (for team plans) + public let teamOnDemandUsedUSD: Double? + /// Team on-demand limit in USD + public let teamOnDemandLimitUSD: Double? + /// Billing cycle start date + public let billingCycleStart: Date? + /// Billing cycle reset date + public let billingCycleEnd: Date? + /// Membership type (e.g., "enterprise", "pro", "hobby") + public let membershipType: String? + /// User email + public let accountEmail: String? + /// User name + public let accountName: String? + /// Raw API response for debugging + public let rawJSON: String? + + // MARK: - Legacy Plan (Request-Based) Fields + + /// Requests used this billing cycle (legacy plans only) + public let requestsUsed: Int? + /// Request limit (non-nil indicates legacy request-based plan) + public let requestsLimit: Int? + + /// Whether this is a legacy request-based plan (vs token-based) + public var isLegacyRequestPlan: Bool { + self.requestsLimit != nil + } + + public init( + planPercentUsed: Double, + autoPercentUsed: Double? = nil, + apiPercentUsed: Double? = nil, + planUsedUSD: Double, + planLimitUSD: Double, + onDemandUsedUSD: Double, + onDemandLimitUSD: Double?, + teamOnDemandUsedUSD: Double?, + teamOnDemandLimitUSD: Double?, + billingCycleStart: Date? = nil, + billingCycleEnd: Date?, + membershipType: String?, + accountEmail: String?, + accountName: String?, + rawJSON: String?, + requestsUsed: Int? = nil, + requestsLimit: Int? = nil) + { + self.planPercentUsed = planPercentUsed + self.autoPercentUsed = autoPercentUsed + self.apiPercentUsed = apiPercentUsed + self.planUsedUSD = planUsedUSD + self.planLimitUSD = planLimitUSD + self.onDemandUsedUSD = onDemandUsedUSD + self.onDemandLimitUSD = onDemandLimitUSD + self.teamOnDemandUsedUSD = teamOnDemandUsedUSD + self.teamOnDemandLimitUSD = teamOnDemandLimitUSD + self.billingCycleStart = billingCycleStart + self.billingCycleEnd = billingCycleEnd + self.membershipType = membershipType + self.accountEmail = accountEmail + self.accountName = accountName + self.rawJSON = rawJSON + self.requestsUsed = requestsUsed + self.requestsLimit = requestsLimit + } + + /// Convert to UsageSnapshot for the common provider interface + public func toUsageSnapshot() -> UsageSnapshot { + let cursorRequests: CursorRequestUsage? = if let used = self.requestsUsed, + let limit = self.requestsLimit, + limit > 0 + { + CursorRequestUsage(used: used, limit: limit) + } else { + nil + } + + // Primary: For usable legacy request quotas, use request usage; otherwise preserve plan percentage. + let primaryUsedPercent = cursorRequests?.usedPercent ?? self.planPercentUsed + + let billingCycleWindowMinutes = Self.billingCycleWindowMinutes( + start: self.billingCycleStart, + end: self.billingCycleEnd) + + let primary = RateWindow( + usedPercent: primaryUsedPercent, + windowMinutes: billingCycleWindowMinutes, + resetsAt: self.billingCycleEnd, + resetDescription: self.billingCycleEnd.map { Self.formatResetDate($0) }) + + // Secondary: Auto + Composer usage (shown as its own bar below Total). + // Legacy request-based plans don't have the token-based Auto/API breakdown — those percentages + // come from the new usage-based pricing and are meaningless next to a request quota, so hide them. + let secondary: RateWindow? = cursorRequests != nil ? nil : self.autoPercentUsed.map { pct in + RateWindow( + usedPercent: pct, + windowMinutes: billingCycleWindowMinutes, + resetsAt: self.billingCycleEnd, + resetDescription: self.billingCycleEnd.map { Self.formatResetDate($0) }) + } + + // Tertiary: API (named model) usage — hidden for legacy request-based plans (see above). + let tertiary: RateWindow? = cursorRequests != nil ? nil : self.apiPercentUsed.map { pct in + RateWindow( + usedPercent: pct, + windowMinutes: billingCycleWindowMinutes, + resetsAt: self.billingCycleEnd, + resetDescription: self.billingCycleEnd.map { Self.formatResetDate($0) }) + } + + // Prefer a personal cap. Team accounts with no user cap expose only the shared on-demand budget. + let resolvedOnDemandUsed: Double + let resolvedOnDemandLimit: Double? + if (self.onDemandLimitUSD ?? 0) > 0 { + resolvedOnDemandUsed = self.onDemandUsedUSD + resolvedOnDemandLimit = self.onDemandLimitUSD + } else if (self.teamOnDemandLimitUSD ?? 0) > 0 { + resolvedOnDemandUsed = self.teamOnDemandUsedUSD ?? 0 + resolvedOnDemandLimit = self.teamOnDemandLimitUSD + } else { + resolvedOnDemandUsed = self.onDemandUsedUSD + resolvedOnDemandLimit = self.onDemandLimitUSD + } + + // Your own on-demand spend to surface alongside a shared team pool (nil when the budget is personal). + let personalOnDemandUsed: Double? = if (self.onDemandLimitUSD ?? 0) > 0 { + nil + } else if (self.teamOnDemandLimitUSD ?? 0) > 0 { + self.onDemandUsedUSD > 0 ? self.onDemandUsedUSD : nil + } else { + nil + } + + // Provider cost snapshot for on-demand usage (include budget before first spend) + let providerCost: ProviderCostSnapshot? = if resolvedOnDemandUsed > 0 + || (resolvedOnDemandLimit ?? 0) > 0 + { + ProviderCostSnapshot( + used: resolvedOnDemandUsed, + limit: resolvedOnDemandLimit ?? 0, + currencyCode: "USD", + period: "Monthly", + resetsAt: self.billingCycleEnd, + personalUsed: personalOnDemandUsed, + updatedAt: Date()) + } else { + nil + } + + let identity = ProviderIdentitySnapshot( + providerID: .cursor, + accountEmail: self.accountEmail, + accountOrganization: nil, + loginMethod: self.membershipType.map { Self.formatMembershipType($0) }) + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: tertiary, + providerCost: providerCost, + cursorRequests: cursorRequests, + updatedAt: Date(), + identity: identity) + } + + private static func formatResetDate(_ date: Date) -> String { + let formatter = DateFormatter() + formatter.dateFormat = "MMM d 'at' h:mma" + formatter.locale = Locale(identifier: "en_US_POSIX") + return "Resets " + formatter.string(from: date) + } + + private static func billingCycleWindowMinutes(start: Date?, end: Date?) -> Int? { + guard let start, + let end + else { return nil } + let minutes = Int((end.timeIntervalSince(start) / 60).rounded()) + return minutes > 0 ? minutes : nil + } + + private static func formatMembershipType(_ type: String) -> String { + switch type.lowercased() { + case "enterprise": + "Cursor Enterprise" + case "pro": + "Cursor Pro" + case "hobby": + "Cursor Hobby" + case "team": + "Cursor Team" + default: + "Cursor \(type.capitalized)" + } + } +} + +// MARK: - Cursor Status Probe Error + +public enum CursorStatusProbeError: LocalizedError, Sendable { + case notLoggedIn + case networkError(String) + case parseFailed(String) + case noSessionCookie + + static let safariFullDiskAccessHint = + "If you use Safari, grant CodexBar Full Disk Access in System Settings ▸ Privacy & Security." + + public var errorDescription: String? { + switch self { + case .notLoggedIn: + #if os(macOS) + "Not logged in to Cursor. Please log in via the CodexBar menu." + #else + "Not logged in to Cursor. Sign in to the Cursor app on this machine or paste a Cookie header copied " + + "from cursor.com into ~/.config/codexbar/config.json (legacy: ~/.codexbar/config.json)." + #endif + case let .networkError(msg): + "Cursor API error: \(msg)" + case let .parseFailed(msg): + "Could not parse Cursor usage: \(msg)" + case .noSessionCookie: + #if os(macOS) + "No Cursor session found. \(Self.safariFullDiskAccessHint) " + + "Please log in to cursor.com in \(cursorCookieImportOrder.loginHint). " + + "You can also sign in to Cursor from the CodexBar menu (Add / switch account)." + #else + "No Cursor session found. Sign in to the Cursor app on this machine or paste a Cookie header copied " + + "from cursor.com into ~/.config/codexbar/config.json (legacy: ~/.codexbar/config.json)." + #endif + } + } +} + +// MARK: - Cursor Session Store + +public actor CursorSessionStore { + public static let shared = CursorSessionStore() + + private var sessionCookies: [HTTPCookie] = [] + private var hasLoadedFromDisk = false + private let fileURL: URL + + private init() { + let fm = FileManager.default + let appSupport = fm.urls(for: .applicationSupportDirectory, in: .userDomainMask).first + ?? fm.temporaryDirectory + let dir = appSupport.appendingPathComponent("CodexBar", isDirectory: true) + try? fm.createDirectory(at: dir, withIntermediateDirectories: true) + self.fileURL = dir.appendingPathComponent("cursor-session.json") + + // Load saved cookies on init + Task { await self.loadFromDiskIfNeeded() } + } + + public func setCookies(_ cookies: [HTTPCookie]) { + self.hasLoadedFromDisk = true + self.sessionCookies = cookies + self.saveToDisk() + } + + public func getCookies() -> [HTTPCookie] { + self.loadFromDiskIfNeeded() + return self.sessionCookies + } + + public func clearCookies() { + self.hasLoadedFromDisk = true + self.sessionCookies = [] + try? FileManager.default.removeItem(at: self.fileURL) + } + + public func hasValidSession() -> Bool { + self.loadFromDiskIfNeeded() + return !self.sessionCookies.isEmpty + } + + #if DEBUG + func resetForTesting(clearDisk: Bool = true) { + self.hasLoadedFromDisk = false + self.sessionCookies = [] + if clearDisk { + try? FileManager.default.removeItem(at: self.fileURL) + } + } + #endif + + private func loadFromDiskIfNeeded() { + guard !self.hasLoadedFromDisk else { return } + self.hasLoadedFromDisk = true + self.loadFromDisk() + } + + private func saveToDisk() { + // Convert cookie properties to JSON-serializable format + // Date values must be converted to TimeInterval (Double) + let cookieData = self.sessionCookies.compactMap { cookie -> [String: Any]? in + guard let props = cookie.properties else { return nil } + var serializable: [String: Any] = [:] + for (key, value) in props { + let keyString = key.rawValue + if let date = value as? Date { + // Convert Date to TimeInterval for JSON compatibility + serializable[keyString] = date.timeIntervalSince1970 + serializable[keyString + "_isDate"] = true + } else if let url = value as? URL { + serializable[keyString] = url.absoluteString + serializable[keyString + "_isURL"] = true + } else if JSONSerialization.isValidJSONObject([value]) || + value is String || + value is Bool || + value is NSNumber + { + serializable[keyString] = value + } + } + return serializable + } + guard !cookieData.isEmpty, + let data = try? JSONSerialization.data(withJSONObject: cookieData, options: [.prettyPrinted]) + else { + return + } + try? data.write(to: self.fileURL) + } + + private func loadFromDisk() { + guard let data = try? Data(contentsOf: self.fileURL), + let cookieArray = try? JSONSerialization.jsonObject(with: data) as? [[String: Any]] + else { return } + + self.sessionCookies = cookieArray.compactMap { props in + // Convert back to HTTPCookiePropertyKey dictionary + var cookieProps: [HTTPCookiePropertyKey: Any] = [:] + for (key, value) in props { + // Skip marker keys + if key.hasSuffix("_isDate") || key.hasSuffix("_isURL") { continue } + + let propKey = HTTPCookiePropertyKey(key) + + // Check if this was a Date + if props[key + "_isDate"] as? Bool == true, let interval = value as? TimeInterval { + cookieProps[propKey] = Date(timeIntervalSince1970: interval) + } + // Check if this was a URL + else if props[key + "_isURL"] as? Bool == true, let urlString = value as? String { + cookieProps[propKey] = URL(string: urlString) + } else { + cookieProps[propKey] = value + } + } + return HTTPCookie(properties: cookieProps) + } + } +} + +// MARK: - Cursor Status Probe + +public struct CursorStatusProbe: Sendable { + public let baseURL: URL + public var timeout: TimeInterval = 15.0 + private let browserDetection: BrowserDetection + private let browserCookieImportOrder: BrowserCookieImportOrder + private let urlSession: any ProviderHTTPTransport + private let appAuthStore: any CursorAppAuthSessionProviding + + public init( + baseURL: URL = URL(string: "https://cursor.com")!, + timeout: TimeInterval = 15.0, + browserDetection: BrowserDetection, + urlSession: any ProviderHTTPTransport = ProviderHTTPClient.shared) + { + self.init( + baseURL: baseURL, + timeout: timeout, + browserDetection: browserDetection, + browserCookieImportOrder: Self.defaultBrowserCookieImportOrder, + urlSession: urlSession, + appAuthStore: CursorAppAuthStore()) + } + + init( + baseURL: URL = URL(string: "https://cursor.com")!, + timeout: TimeInterval = 15.0, + browserDetection: BrowserDetection, + browserCookieImportOrder: BrowserCookieImportOrder = Self.defaultBrowserCookieImportOrder, + urlSession: any ProviderHTTPTransport = ProviderHTTPClient.shared, + appAuthStore: any CursorAppAuthSessionProviding) + { + self.baseURL = baseURL + self.timeout = timeout + self.browserDetection = browserDetection + self.browserCookieImportOrder = browserCookieImportOrder + self.urlSession = urlSession + self.appAuthStore = appAuthStore + } + + /// Fetch Cursor usage using a first-party web session derived from Cursor.app's access token. + func fetchWithAppAuthSession(_ session: CursorAppAuthSession) async throws -> CursorStatusSnapshot { + try await self.fetchWithCookieHeader( + session.cookieHeader(), + requestUsageUserIDFallback: session.userID()) + } + + /// Fetch Cursor usage with manual cookie header (for debugging). + public func fetchWithManualCookies(_ cookieHeader: String) async throws -> CursorStatusSnapshot { + try await self.fetchWithCookieHeader(cookieHeader) + } + + /// Fetch Cursor usage using browser cookies with fallback to stored session. + public func fetch( + cookieHeaderOverride: String? = nil, + allowCachedSessions: Bool = true, + allowAppAuthFallback: Bool = true, + logger: ((String) -> Void)? = nil) + async throws -> CursorStatusSnapshot + { + let log: (String) -> Void = { msg in logger?("[cursor] \(msg)") } + var firstRecoverableError: CursorStatusProbeError? + + if let override = CookieHeaderNormalizer.normalize(cookieHeaderOverride) { + log("Using manual cookie header") + return try await self.fetchWithCookieHeader(override) + } + + if allowCachedSessions, + let cached = CookieHeaderCache.load(provider: .cursor), + !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + log("Using cached cookie header from \(cached.sourceLabel)") + do { + return try await self.fetchWithCookieHeader(cached.cookieHeader) + } catch let error as CursorStatusProbeError { + if case .notLoggedIn = error { + CookieHeaderCache.clear(provider: .cursor) + } else { + throw error + } + } catch { + throw error + } + } + + #if os(macOS) + // Try each browser in order. The first browser that *has* session cookie names is not always valid + // (e.g. stale Chrome tokens); keep trying until the API accepts a session or we run out of browsers. + let browserCandidates = self.browserCookieImportOrder.cookieImportCandidates(using: self.browserDetection) + switch await self.scanBrowsers( + browserCandidates, + importSessions: { browser in + CursorCookieImporter.importSessionsIfPresent( + browser: browser, + browserDetection: self.browserDetection, + logger: log) + }, + attemptFetch: { session in + await self.fetchIfSessionAccepted(session, log: log) + }) + { + case let .succeeded(snapshot): + return snapshot + case let .exhausted(error): + firstRecoverableError = error ?? firstRecoverableError + } + + switch await self.scanBrowsers( + browserCandidates, + importSessions: { browser in + CursorCookieImporter.importDomainCookieSessionsIfPresent( + browser: browser, + browserDetection: self.browserDetection, + logger: log) + }, + attemptFetch: { session in + await self.fetchIfSessionAccepted(session, log: log) + }) + { + case let .succeeded(snapshot): + return snapshot + case let .exhausted(error): + firstRecoverableError = error ?? firstRecoverableError + } + #endif + + // Fall back to stored session cookies (from "Add Account" login flow) + if allowCachedSessions { + let storedCookies = await CursorSessionStore.shared.getCookies() + if !storedCookies.isEmpty { + log("Using stored session cookies") + let cookieHeader = storedCookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + do { + return try await self.fetchWithCookieHeader(cookieHeader) + } catch let error as CursorStatusProbeError { + if case .notLoggedIn = error { + // Clear only when auth is invalid; keep for transient failures. + await CursorSessionStore.shared.clearCookies() + log("Stored session invalid, cleared") + } else { + log("Stored session failed: \(error.localizedDescription)") + firstRecoverableError = firstRecoverableError ?? error + } + } catch { + log("Stored session failed: \(error.localizedDescription)") + firstRecoverableError = firstRecoverableError ?? .networkError(error.localizedDescription) + } + } + } + + // A transient failure for an explicitly selected session must not switch to Cursor.app's account. + if let firstRecoverableError { + throw firstRecoverableError + } + + // Last fallback: derive Cursor's first-party web session from the app token in its global state DB. + // Reusing the web flow preserves modern billing, legacy request quotas, and account-scoped identity. + if allowAppAuthFallback, + let appSession = try? self.appAuthStore.loadSession(), + appSession.isUsable + { + log("Using Cursor.app local auth fallback") + do { + return try await self.fetchWithAppAuthSession(appSession) + } catch let error as CursorStatusProbeError { + if case .notLoggedIn = error { + log("Cursor.app local auth was rejected") + } else { + firstRecoverableError = firstRecoverableError ?? error + } + } catch { + firstRecoverableError = firstRecoverableError ?? .networkError(error.localizedDescription) + } + } + + if let firstRecoverableError { + throw firstRecoverableError + } + + throw CursorStatusProbeError.noSessionCookie + } + + #if os(macOS) + enum ImportedSessionFetchOutcome { + case succeeded(CursorStatusSnapshot) + case tryNextBrowser + case failed(CursorStatusProbeError) + } + + enum ImportedSessionScanResult { + case succeeded(CursorStatusSnapshot) + case exhausted(CursorStatusProbeError?) + } + + func scanBrowsers( + _ browsers: [Browser], + importSessions: (Browser) -> [CursorCookieImporter.SessionInfo], + attemptFetch: (CursorCookieImporter.SessionInfo) async -> ImportedSessionFetchOutcome) async + -> ImportedSessionScanResult + { + var firstFailure: CursorStatusProbeError? + + for browser in browsers { + let sessions = importSessions(browser) + guard !sessions.isEmpty else { continue } + for session in sessions { + switch await attemptFetch(session) { + case let .succeeded(snapshot): + return .succeeded(snapshot) + case .tryNextBrowser: + continue + case let .failed(error): + firstFailure = firstFailure ?? error + } + } + } + + return .exhausted(firstFailure) + } + + func scanImportedSessions( + _ sessions: [CursorCookieImporter.SessionInfo], + attemptFetch: (CursorCookieImporter.SessionInfo) async -> ImportedSessionFetchOutcome) async + -> ImportedSessionScanResult + { + var firstFailure: CursorStatusProbeError? + + for session in sessions { + switch await attemptFetch(session) { + case let .succeeded(snapshot): + return .succeeded(snapshot) + case .tryNextBrowser: + continue + case let .failed(error): + firstFailure = firstFailure ?? error + } + } + + return .exhausted(firstFailure) + } + + private func fetchIfSessionAccepted( + _ session: CursorCookieImporter.SessionInfo, + log: @escaping (String) -> Void) async -> ImportedSessionFetchOutcome + { + log("Trying Cursor session from \(session.sourceLabel)") + do { + let snapshot = try await self.fetchWithCookieHeader(session.cookieHeader) + CookieHeaderCache.store( + provider: .cursor, + cookieHeader: session.cookieHeader, + sourceLabel: session.sourceLabel) + return .succeeded(snapshot) + } catch let error as CursorStatusProbeError { + if case .notLoggedIn = error { + log("Cursor API rejected cookies from \(session.sourceLabel); trying next browser if any") + return .tryNextBrowser + } + log("Cursor fetch failed using \(session.sourceLabel): \(error.localizedDescription)") + return .failed(error) + } catch { + log("Cursor fetch failed using \(session.sourceLabel): \(error.localizedDescription)") + return .failed(.networkError(error.localizedDescription)) + } + } + #endif + + private func fetchWithCookieHeader( + _ cookieHeader: String, + requestUsageUserIDFallback: String? = nil) async throws -> CursorStatusSnapshot + { + enum FetchPart: Sendable { + case usageSummary((CursorUsageSummary, String)) + case userInfo(Result) + } + + var usageSummaryResult: (CursorUsageSummary, String)? + var userInfo: CursorUserInfo? + + try await withThrowingTaskGroup(of: FetchPart.self) { group in + group.addTask { + try await .usageSummary(self.fetchUsageSummary(cookieHeader: cookieHeader)) + } + group.addTask { + do { + return try await .userInfo(.success(self.fetchUserInfo(cookieHeader: cookieHeader))) + } catch { + return .userInfo(.failure(error)) + } + } + + while let result = try await group.next() { + switch result { + case let .usageSummary(value): + usageSummaryResult = value + case let .userInfo(value): + userInfo = try? value.get() + } + } + } + + guard let usageSummaryResult else { + throw CursorStatusProbeError.networkError("Cursor usage summary fetch did not complete") + } + + let (usageSummary, rawJSON) = usageSummaryResult + + // Fetch legacy request usage only if user has a sub ID. + // Uses try? to avoid breaking the flow for users where this endpoint fails or returns unexpected data. + var requestUsage: CursorUsageResponse? + var requestUsageRawJSON: String? + if let userId = userInfo?.sub ?? requestUsageUserIDFallback { + do { + let (usage, usageRawJSON) = try await self.fetchRequestUsage(userId: userId, cookieHeader: cookieHeader) + requestUsage = usage + requestUsageRawJSON = usageRawJSON + } catch { + // Silently ignore - not all plans have this endpoint + } + } + + // Combine raw JSON for debugging + var combinedRawJSON: String? = rawJSON + if let usageJSON = requestUsageRawJSON { + combinedRawJSON = (combinedRawJSON ?? "") + "\n\n--- /api/usage response ---\n" + usageJSON + } + + return self.parseUsageSummary( + usageSummary, + userInfo: userInfo, + rawJSON: combinedRawJSON, + requestUsage: requestUsage) + } + + private func fetchUsageSummary(cookieHeader: String) async throws -> (CursorUsageSummary, String) { + let url = self.baseURL.appendingPathComponent("/api/usage-summary") + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + + let (data, response) = try await self.urlSession.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw CursorStatusProbeError.networkError("Invalid response") + } + + if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 { + throw CursorStatusProbeError.notLoggedIn + } + + guard httpResponse.statusCode == 200 else { + throw CursorStatusProbeError.networkError("HTTP \(httpResponse.statusCode)") + } + + let rawJSON = String(data: data, encoding: .utf8) ?? "" + + do { + let decoder = JSONDecoder() + let summary = try decoder.decode(CursorUsageSummary.self, from: data) + return (summary, rawJSON) + } catch { + throw CursorStatusProbeError + .parseFailed("JSON decode failed: \(error.localizedDescription). Raw: \(rawJSON.prefix(200))") + } + } + + private func fetchUserInfo(cookieHeader: String) async throws -> CursorUserInfo { + let url = self.baseURL.appendingPathComponent("/api/auth/me") + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + + let (data, response) = try await self.urlSession.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else { + throw CursorStatusProbeError.networkError("Failed to fetch user info") + } + + let decoder = JSONDecoder() + return try decoder.decode(CursorUserInfo.self, from: data) + } + + private func fetchRequestUsage( + userId: String, + cookieHeader: String) async throws -> (CursorUsageResponse, String) + { + let url = self.baseURL.appendingPathComponent("/api/usage") + .appending(queryItems: [URLQueryItem(name: "user", value: userId)]) + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + + let (data, response) = try await self.urlSession.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else { + throw CursorStatusProbeError.networkError("Failed to fetch request usage") + } + + let rawJSON = String(data: data, encoding: .utf8) ?? "" + let decoder = JSONDecoder() + let usage = try decoder.decode(CursorUsageResponse.self, from: data) + return (usage, rawJSON) + } + + func parseUsageSummary( + _ summary: CursorUsageSummary, + userInfo: CursorUserInfo?, + rawJSON: String?, + requestUsage: CursorUsageResponse? = nil) -> CursorStatusSnapshot + { + func parseBillingCycleDate(_ dateString: String?) -> Date? { + guard let dateString else { return nil } + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + return formatter.date(from: dateString) ?? ISO8601DateFormatter().date(from: dateString) + } + let billingCycleStart = parseBillingCycleDate(summary.billingCycleStart) + let billingCycleEnd = parseBillingCycleDate(summary.billingCycleEnd) + + // Convert cents to USD (plan percent derives from raw values to avoid percent unit mismatches). + // Use plan.limit directly - breakdown.total represents total *used* credits, not the limit. + let planUsedRaw = Double(summary.individualUsage?.plan?.used ?? 0) + let planLimitRaw = Double(summary.individualUsage?.plan?.limit ?? 0) + func normPct(_ value: Double?) -> Double? { + guard let v = value else { return nil } + if v < 0 { return 0 } + if v > 100 { return 100 } + return v + } + + func normalizeTotalPercent(_ v: Double) -> Double { + max(0, min(100, v)) + } + + // Cursor's usage-summary percent fields are already in percentage units, even when they are fractional + // values below 1.0 (for example 0.36 means 0.36%, which the dashboard rounds to 0%). + let autoPercent = normPct(summary.individualUsage?.plan?.autoPercentUsed) + let apiPercent = normPct(summary.individualUsage?.plan?.apiPercentUsed) + + // Enterprise / team-member personal cap (cents). Reported under `individualUsage.overall` for accounts + // that don't get a `plan` block. Falls through to existing logic when absent so non-enterprise paths + // are untouched. + let overallUsedRaw = (summary.individualUsage?.overall?.used).map(Double.init) + let overallLimitRaw = (summary.individualUsage?.overall?.limit).map(Double.init) + + // Shared team/enterprise pool (cents). Last-resort fallback when no individual data is available. + let pooledUsedRaw = (summary.teamUsage?.pooled?.used).map(Double.init) + let pooledLimitRaw = (summary.teamUsage?.pooled?.limit).map(Double.init) + + // Headline "Total" precedence: + // 1. `individualUsage.plan.totalPercentUsed` (existing behavior for Pro/Hobby/etc.) + // 2. averaged `auto` + `api` lane percents (existing behavior) + // 3. either lane alone (existing behavior) + // 4. `individualUsage.plan` ratio (existing behavior) + // 5. NEW: `individualUsage.overall` ratio (Enterprise/Team personal cap) + // 6. NEW: `teamUsage.pooled` ratio (last resort when no individual data is reported) + let planPercentUsed: Double = if let totalPercentUsed = summary.individualUsage?.plan?.totalPercentUsed { + normalizeTotalPercent(totalPercentUsed) + } else if let autoUsed = autoPercent, let apiUsed = apiPercent { + max(0, min(100, (autoUsed + apiUsed) / 2)) + } else if let apiUsed = apiPercent { + max(0, min(100, apiUsed)) + } else if let autoUsed = autoPercent { + max(0, min(100, autoUsed)) + } else if planLimitRaw > 0 { + (planUsedRaw / planLimitRaw) * 100 + } else if let used = overallUsedRaw, let limit = overallLimitRaw, limit > 0 { + normalizeTotalPercent((used / limit) * 100) + } else if let used = pooledUsedRaw, let limit = pooledLimitRaw, limit > 0 { + normalizeTotalPercent((used / limit) * 100) + } else { + 0 + } + + // USD figures: prefer the source the headline ultimately came from. When `plan` is missing but + // `overall` or `pooled` carry the cents, surface those so the on-demand display and downstream + // consumers see real dollar amounts instead of zeros. + let planUsed: Double + let planLimit: Double + if planLimitRaw > 0 || planUsedRaw > 0 { + planUsed = planUsedRaw / 100.0 + planLimit = planLimitRaw / 100.0 + } else if let usedCents = overallUsedRaw, let limitCents = overallLimitRaw { + planUsed = usedCents / 100.0 + planLimit = limitCents / 100.0 + } else if let usedCents = pooledUsedRaw, let limitCents = pooledLimitRaw { + planUsed = usedCents / 100.0 + planLimit = limitCents / 100.0 + } else { + planUsed = 0 + planLimit = 0 + } + + let onDemandUsed = Double(summary.individualUsage?.onDemand?.used ?? 0) / 100.0 + let onDemandLimit: Double? = summary.individualUsage?.onDemand?.limit.map { Double($0) / 100.0 } + + let teamOnDemandUsed: Double? = summary.teamUsage?.onDemand?.used.map { Double($0) / 100.0 } + let teamOnDemandLimit: Double? = summary.teamUsage?.onDemand?.limit.map { Double($0) / 100.0 } + + // Legacy request-based plan: maxRequestUsage being non-nil indicates a request-based plan + let requestsUsed: Int? = requestUsage?.gpt4?.numRequestsTotal ?? requestUsage?.gpt4?.numRequests + let requestsLimit: Int? = requestUsage?.gpt4?.maxRequestUsage + + return CursorStatusSnapshot( + planPercentUsed: planPercentUsed, + autoPercentUsed: autoPercent, + apiPercentUsed: apiPercent, + planUsedUSD: planUsed, + planLimitUSD: planLimit, + onDemandUsedUSD: onDemandUsed, + onDemandLimitUSD: onDemandLimit, + teamOnDemandUsedUSD: teamOnDemandUsed, + teamOnDemandLimitUSD: teamOnDemandLimit, + billingCycleStart: billingCycleStart, + billingCycleEnd: billingCycleEnd, + membershipType: summary.membershipType, + accountEmail: userInfo?.email, + accountName: userInfo?.name, + rawJSON: rawJSON, + requestsUsed: requestsUsed, + requestsLimit: requestsLimit) + } + + #if os(macOS) + private static let defaultBrowserCookieImportOrder: BrowserCookieImportOrder = cursorCookieImportOrder + #else + private static let defaultBrowserCookieImportOrder: BrowserCookieImportOrder = [] + #endif +} + +#else + +// MARK: - Cursor (Unsupported) + +public enum CursorStatusProbeError: LocalizedError, Sendable { + case notSupported + + public var errorDescription: String? { + "Cursor is only supported on macOS." + } +} + +public struct CursorStatusSnapshot: Sendable { + public init() {} + + public func toUsageSnapshot() -> UsageSnapshot { + UsageSnapshot( + primary: RateWindow(usedPercent: 0, windowMinutes: nil, resetsAt: nil, resetDescription: nil), + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: Date(), + identity: nil) + } +} + +public struct CursorStatusProbe: Sendable { + public init( + baseURL: URL = URL(string: "https://cursor.com")!, + timeout: TimeInterval = 15.0, + browserDetection: BrowserDetection, + urlSession: any ProviderHTTPTransport = ProviderHTTPClient.shared) + { + _ = baseURL + _ = timeout + _ = browserDetection + _ = urlSession + } + + public func fetch(logger: ((String) -> Void)? = nil) async throws -> CursorStatusSnapshot { + _ = logger + throw CursorStatusProbeError.notSupported + } + + public func fetch( + cookieHeaderOverride _: String? = nil, + allowCachedSessions _: Bool = true, + logger: ((String) -> Void)? = nil) async throws -> CursorStatusSnapshot + { + try await self.fetch(logger: logger) + } +} + +#endif diff --git a/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekProviderDescriptor.swift b/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekProviderDescriptor.swift new file mode 100644 index 0000000..f2847c7 --- /dev/null +++ b/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekProviderDescriptor.swift @@ -0,0 +1,48 @@ +import Foundation + +public enum DeepSeekProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .deepseek, + metadata: ProviderMetadata( + id: .deepseek, + displayName: "DeepSeek", + sessionLabel: "Balance", + weeklyLabel: "Balance", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show DeepSeek usage", + cliName: "deepseek", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://platform.deepseek.com/usage", + statusPageURL: nil, + statusLinkURL: "https://status.deepseek.com"), + branding: ProviderBranding( + iconStyle: .deepseek, + iconResourceName: "ProviderIcon-deepseek", + color: ProviderColor(red: 0.32, green: 0.49, blue: 0.94)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "DeepSeek per-day cost history is not available via API." }), + fetchPlan: .apiToken( + strategyID: "deepseek.api", + resolveToken: { ProviderTokenResolver.deepseekToken(environment: $0) }, + missingCredentialsError: { DeepSeekUsageError.missingCredentials }, + loadUsage: { apiKey, context in + try await DeepSeekUsageFetcher.fetchUsage( + apiKey: apiKey, + includeOptionalUsage: context.includeOptionalUsage).toUsageSnapshot() + }), + cli: ProviderCLIConfig( + name: "deepseek", + aliases: ["deep-seek", "ds"], + versionDetector: nil)) + } +} diff --git a/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekSettingsReader.swift b/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekSettingsReader.swift new file mode 100644 index 0000000..3c787ca --- /dev/null +++ b/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekSettingsReader.swift @@ -0,0 +1,33 @@ +import Foundation + +public struct DeepSeekSettingsReader: Sendable { + public static let apiKeyEnvironmentKey = "DEEPSEEK_API_KEY" + public static let apiKeyEnvironmentKeys = [Self.apiKeyEnvironmentKey, "DEEPSEEK_KEY"] + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + for key in self.apiKeyEnvironmentKeys { + guard let raw = environment[key]?.trimmingCharacters(in: .whitespacesAndNewlines), + !raw.isEmpty + else { + continue + } + let cleaned = Self.cleaned(raw) + if !cleaned.isEmpty { + return cleaned + } + } + return nil + } + + private static func cleaned(_ raw: String) -> String { + var value = raw + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + return value.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekUsageCostParser.swift b/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekUsageCostParser.swift new file mode 100644 index 0000000..77e6e4f --- /dev/null +++ b/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekUsageCostParser.swift @@ -0,0 +1,709 @@ +import Foundation + +// MARK: - Amount Response Models + +struct DeepSeekAmountPayload: Decodable { + let code: Int? + let msg: String? + let data: DeepSeekAmountData? + + private enum CodingKeys: String, CodingKey { + case code, msg, data + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.code = try container.decodeIfPresent(Int.self, forKey: .code) + self.msg = try container.decodeIfPresent(String.self, forKey: .msg) + if container.contains(.data) { + if let dataValue = try? container.decodeIfPresent(DeepSeekAmountData.self, forKey: .data) { + self.data = dataValue + } else { + self.data = nil + } + } else { + self.data = nil + } + } +} + +struct DeepSeekAmountData: Decodable { + let bizCode: Int? + let bizMsg: String? + let bizData: DeepSeekAmountBizData? + + private enum CodingKeys: String, CodingKey { + case bizCode = "biz_code" + case bizMsg = "biz_msg" + case bizData = "biz_data" + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.bizCode = try container.decodeIfPresent(Int.self, forKey: .bizCode) + self.bizMsg = try container.decodeIfPresent(String.self, forKey: .bizMsg) + if container.contains(.bizData) { + if let dataValue = try? container.decodeIfPresent(DeepSeekAmountBizData.self, forKey: .bizData) { + self.bizData = dataValue + } else { + self.bizData = nil + } + } else { + self.bizData = nil + } + } +} + +struct DeepSeekAmountBizData: Decodable { + let total: [DeepSeekModelUsage]? + let days: [DeepSeekDayUsage]? + + private enum CodingKeys: String, CodingKey { + case total, days + } +} + +// MARK: - Cost Response Models + +struct DeepSeekCostPayload: Decodable { + let code: Int? + let msg: String? + let data: DeepSeekCostData? + + private enum CodingKeys: String, CodingKey { + case code, msg, data + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.code = try container.decodeIfPresent(Int.self, forKey: .code) + self.msg = try container.decodeIfPresent(String.self, forKey: .msg) + if container.contains(.data) { + if let dataValue = try? container.decodeIfPresent(DeepSeekCostData.self, forKey: .data) { + self.data = dataValue + } else { + self.data = nil + } + } else { + self.data = nil + } + } +} + +struct DeepSeekCostData: Decodable { + let bizCode: Int? + let bizMsg: String? + let bizData: [DeepSeekCostBizDataItem]? + + private enum CodingKeys: String, CodingKey { + case bizCode = "biz_code" + case bizMsg = "biz_msg" + case bizData = "biz_data" + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.bizCode = try container.decodeIfPresent(Int.self, forKey: .bizCode) + self.bizMsg = try container.decodeIfPresent(String.self, forKey: .bizMsg) + self.bizData = try container.decodeIfPresent([DeepSeekCostBizDataItem].self, forKey: .bizData) + } +} + +struct DeepSeekCostBizDataItem: Decodable { + let total: [DeepSeekCostModelUsage]? + let days: [DeepSeekCostDayUsage]? + let currency: String? + + private enum CodingKeys: String, CodingKey { + case total, days, currency + } +} + +// MARK: - Shared Models + +struct DeepSeekModelUsage: Decodable { + let model: String? + let usage: [DeepSeekUsageItem]? + + private enum CodingKeys: String, CodingKey { + case model, usage + } +} + +struct DeepSeekDayUsage: Decodable { + let date: String? + let data: [DeepSeekModelUsage]? + + private enum CodingKeys: String, CodingKey { + case date, data + } +} + +struct DeepSeekUsageItem: Decodable { + let type: String? + let amount: String? + + private enum CodingKeys: String, CodingKey { + case type, amount + } +} + +struct DeepSeekCostModelUsage: Decodable { + let model: String? + let usage: [DeepSeekCostItem]? + + private enum CodingKeys: String, CodingKey { + case model, usage + } +} + +struct DeepSeekCostDayUsage: Decodable { + let date: String? + let data: [DeepSeekCostModelUsage]? + + private enum CodingKeys: String, CodingKey { + case date, data + } +} + +struct DeepSeekCostItem: Decodable { + let type: String? + let amount: String? + + private enum CodingKeys: String, CodingKey { + case type, amount + } +} + +// MARK: - Domain Models + +public struct DeepSeekUsageSummary: Sendable, Equatable { + public let todayTokens: Int + public let currentMonthTokens: Int + public let todayCost: Double? + public let currentMonthCost: Double? + public let requestCount: Int + public let currentMonthRequestCount: Int + public let topModel: String? + public let categoryBreakdown: [DeepSeekCategoryBreakdown] + public let daily: [DeepSeekDailyUsage] + public let currency: String + public let updatedAt: Date + + public init( + todayTokens: Int, + currentMonthTokens: Int, + todayCost: Double?, + currentMonthCost: Double?, + requestCount: Int, + currentMonthRequestCount: Int, + topModel: String?, + categoryBreakdown: [DeepSeekCategoryBreakdown], + daily: [DeepSeekDailyUsage], + currency: String, + updatedAt: Date) + { + self.todayTokens = todayTokens + self.currentMonthTokens = currentMonthTokens + self.todayCost = todayCost + self.currentMonthCost = currentMonthCost + self.requestCount = requestCount + self.currentMonthRequestCount = currentMonthRequestCount + self.topModel = topModel + self.categoryBreakdown = categoryBreakdown + self.daily = daily + self.currency = currency + self.updatedAt = updatedAt + } +} + +public struct DeepSeekCategoryBreakdown: Sendable, Equatable { + public let category: DeepSeekUsageCategory + public let tokens: Int + public let cost: Double? + + public init(category: DeepSeekUsageCategory, tokens: Int, cost: Double?) { + self.category = category + self.tokens = tokens + self.cost = cost + } +} + +public enum DeepSeekUsageCategory: String, Sendable, Equatable { + case promptCacheHitToken = "PROMPT_CACHE_HIT_TOKEN" + case promptCacheMissToken = "PROMPT_CACHE_MISS_TOKEN" + case responseToken = "RESPONSE_TOKEN" + case request = "REQUEST" + + public init?(rawValue: String) { + switch rawValue.uppercased() { + case "PROMPT_CACHE_HIT_TOKEN": + self = .promptCacheHitToken + case "PROMPT_CACHE_MISS_TOKEN": + self = .promptCacheMissToken + case "RESPONSE_TOKEN": + self = .responseToken + case "REQUEST": + self = .request + default: + return nil + } + } +} + +public struct DeepSeekDailyUsage: Sendable, Equatable { + public let date: String + public let totalTokens: Int + public let cost: Double? + public let requestCount: Int + + public init(date: String, totalTokens: Int, cost: Double?, requestCount: Int) { + self.date = date + self.totalTokens = totalTokens + self.cost = cost + self.requestCount = requestCount + } +} + +// MARK: - Parsing + +enum DeepSeekUsageCostParser { + static func decodeAmountPayload(data: Data) throws -> DeepSeekAmountPayload { + try JSONDecoder().decode(DeepSeekAmountPayload.self, from: data) + } + + static func decodeCostPayload(data: Data) throws -> DeepSeekCostPayload { + try JSONDecoder().decode(DeepSeekCostPayload.self, from: data) + } + + static func parse( + amountData: Data, + costData: Data, + now: Date = Date(), + calendar: Calendar = .current) throws -> DeepSeekUsageSummary + { + let amountPayload: DeepSeekAmountPayload + let costPayload: DeepSeekCostPayload + do { + amountPayload = try self.decodeAmountPayload(data: amountData) + } catch { + throw DeepSeekUsageError.parseFailed("amount: \(error.localizedDescription)") + } + do { + costPayload = try self.decodeCostPayload(data: costData) + } catch { + throw DeepSeekUsageError.parseFailed("cost: \(error.localizedDescription)") + } + + // Validate responses + if let code = amountPayload.code, code != 0 { + throw DeepSeekUsageError.apiError("amount code \(code)") + } + if let bizCode = amountPayload.data?.bizCode, bizCode != 0 { + throw DeepSeekUsageError.apiError("amount biz_code \(bizCode)") + } + if let code = costPayload.code, code != 0 { + throw DeepSeekUsageError.apiError("cost code \(code)") + } + if let bizCode = costPayload.data?.bizCode, bizCode != 0 { + throw DeepSeekUsageError.apiError("cost biz_code \(bizCode)") + } + + guard let amountBizData = amountPayload.data?.bizData else { + throw DeepSeekUsageError.parseFailed("Missing amount biz_data") + } + + let currency = costPayload.data?.bizData?.first?.currency ?? "CNY" + + // Parse total amounts + let totalAmounts = amountBizData.total ?? [] + let totalCosts = costPayload.data?.bizData?.first?.total ?? [] + + // Parse daily data + let dailyAmounts = amountBizData.days ?? [] + let dailyCosts = costPayload.data?.bizData?.first?.days ?? [] + + return self.aggregate(input: AggregationInput( + totalAmounts: totalAmounts, + totalCosts: totalCosts, + dailyAmounts: dailyAmounts, + dailyCosts: dailyCosts, + currency: currency, + now: now, + calendar: calendar)) + } + + // MARK: - Aggregation + + private struct AggregationContext { + let calendar: Calendar + let todayString: String + let startOfMonth: Date + let now: Date + let dailyAmountMap: [String: [String: [DeepSeekUsageItem]]] + let dailyCostMap: [String: [String: [DeepSeekCostItem]]] + let allDates: Set + + init( + dailyAmounts: [DeepSeekDayUsage], + dailyCosts: [DeepSeekCostDayUsage], + now: Date, + calendar: Calendar) + { + self.calendar = calendar + self.now = now + self.todayString = Self.dayString(now, calendar: calendar) + + var components = calendar.dateComponents([.year, .month], from: now) + components.day = 1 + self.startOfMonth = calendar.date(from: components) ?? now + + self.dailyAmountMap = Self.buildAmountMap(from: dailyAmounts) + self.dailyCostMap = Self.buildCostMap(from: dailyCosts) + + var dates: Set = [] + for date in self.dailyAmountMap.keys { + dates.insert(date) + } + for date in self.dailyCostMap.keys { + dates.insert(date) + } + self.allDates = dates + } + + static func dayString(_ date: Date, calendar: Calendar) -> String { + let components = calendar.dateComponents([.year, .month, .day], from: date) + guard let year = components.year, + let month = components.month, + let day = components.day + else { return "" } + return String(format: "%04d-%02d-%02d", year, month, day) + } + + static func buildAmountMap( + from dailyAmounts: [DeepSeekDayUsage]) -> [String: [String: [DeepSeekUsageItem]]] + { + var result: [String: [String: [DeepSeekUsageItem]]] = [:] + for dayUsage in dailyAmounts { + guard let date = dayUsage.date else { continue } + var modelMap: [String: [DeepSeekUsageItem]] = [:] + for modelUsage in dayUsage.data ?? [] { + guard let model = modelUsage.model else { continue } + let items = modelUsage.usage ?? [] + if !items.isEmpty { + modelMap[model] = items + } + } + if !modelMap.isEmpty { + result[date] = modelMap + } + } + return result + } + + static func buildCostMap( + from dailyCosts: [DeepSeekCostDayUsage]) -> [String: [String: [DeepSeekCostItem]]] + { + var result: [String: [String: [DeepSeekCostItem]]] = [:] + for dayUsage in dailyCosts { + guard let date = dayUsage.date else { continue } + var modelMap: [String: [DeepSeekCostItem]] = [:] + for modelUsage in dayUsage.data ?? [] { + guard let model = modelUsage.model else { continue } + let items = modelUsage.usage ?? [] + if !items.isEmpty { + modelMap[model] = items + } + } + if !modelMap.isEmpty { + result[date] = modelMap + } + } + return result + } + } + + private struct AggregationInput { + let totalAmounts: [DeepSeekModelUsage] + let totalCosts: [DeepSeekCostModelUsage] + let dailyAmounts: [DeepSeekDayUsage] + let dailyCosts: [DeepSeekCostDayUsage] + let currency: String + let now: Date + let calendar: Calendar + } + + private static func aggregate(input: AggregationInput) -> DeepSeekUsageSummary { + let ctx = AggregationContext( + dailyAmounts: input.dailyAmounts, + dailyCosts: input.dailyCosts, + now: input.now, + calendar: input.calendar) + + // Today aggregation + let todayResult = self.aggregateDay( + dateString: ctx.todayString, + amountMap: ctx.dailyAmountMap, + costMap: ctx.dailyCostMap, + calendar: ctx.calendar) + + // Month aggregation + let dailyCtx = DailyAggregationContext( + allDates: ctx.allDates, + startOfMonth: ctx.startOfMonth, + now: ctx.now, + amountMap: ctx.dailyAmountMap, + costMap: ctx.dailyCostMap, + calendar: ctx.calendar) + let monthResult = self.aggregateMonth(ctx: dailyCtx) + + // Model and category breakdown from totals + let (topModel, categoryBreakdown) = self.buildBreakdowns( + totalAmounts: input.totalAmounts, + totalCosts: input.totalCosts) + + // Daily usage array + let dailyUsages = self.buildDailyUsages(ctx: dailyCtx) + + return DeepSeekUsageSummary( + todayTokens: todayResult.tokens, + currentMonthTokens: monthResult.tokens, + todayCost: todayResult.cost, + currentMonthCost: monthResult.cost, + requestCount: todayResult.requests, + currentMonthRequestCount: monthResult.requests, + topModel: topModel, + categoryBreakdown: categoryBreakdown, + daily: dailyUsages, + currency: input.currency, + updatedAt: input.now) + } + + private struct DayAggregationResult { + let tokens: Int + let cost: Double? + let requests: Int + } + + private struct DailyAggregationContext { + let allDates: Set + let startOfMonth: Date + let now: Date + let amountMap: [String: [String: [DeepSeekUsageItem]]] + let costMap: [String: [String: [DeepSeekCostItem]]] + let calendar: Calendar + } + + private static func aggregateDay( + dateString: String, + amountMap: [String: [String: [DeepSeekUsageItem]]], + costMap: [String: [String: [DeepSeekCostItem]]], + calendar: Calendar) -> DayAggregationResult + { + var tokens = 0 + var cost: Double? + var requests = 0 + + if let amounts = amountMap[dateString] { + for items in amounts.values { + for item in items { + guard let category = DeepSeekUsageCategory(rawValue: item.type ?? "") else { continue } + if category == .request { + requests += self.parseTokenAmount(item.amount) + } else { + tokens += self.parseTokenAmount(item.amount) + } + } + } + } + + if let costs = costMap[dateString] { + for items in costs.values { + for item in items { + guard let category = DeepSeekUsageCategory(rawValue: item.type ?? "") else { continue } + if category != .request { + let amount = Self.parseCostAmount(item.amount) + if let existing = cost { + cost = existing + amount + } else { + cost = amount + } + } + } + } + } + + return DayAggregationResult(tokens: tokens, cost: cost, requests: requests) + } + + private static func aggregateMonth(ctx: DailyAggregationContext) -> DayAggregationResult { + var tokens = 0 + var cost: Double? + var requests = 0 + + for date in ctx.allDates { + guard let parsed = self.parseDate(date, calendar: ctx.calendar), + parsed >= ctx.startOfMonth, + parsed <= ctx.now + else { continue } + + if let amounts = ctx.amountMap[date] { + for items in amounts.values { + for item in items { + guard let category = DeepSeekUsageCategory(rawValue: item.type ?? "") else { continue } + if category == .request { + requests += Self.parseTokenAmount(item.amount) + } else { + tokens += Self.parseTokenAmount(item.amount) + } + } + } + } + + if let costs = ctx.costMap[date] { + for items in costs.values { + for item in items { + guard let category = DeepSeekUsageCategory(rawValue: item.type ?? "") else { continue } + if category != .request { + let amount = Self.parseCostAmount(item.amount) + if let existing = cost { + cost = existing + amount + } else { + cost = amount + } + } + } + } + } + } + + return DayAggregationResult(tokens: tokens, cost: cost, requests: requests) + } + + private static func buildBreakdowns( + totalAmounts: [DeepSeekModelUsage], + totalCosts: [DeepSeekCostModelUsage]) -> (String?, [DeepSeekCategoryBreakdown]) + { + var modelTokens: [String: Int] = [:] + var categoryTokens: [DeepSeekUsageCategory: Int] = [:] + var categoryCosts: [DeepSeekUsageCategory: Double] = [:] + + for modelUsage in totalAmounts { + guard let model = modelUsage.model else { continue } + var total = 0 + for item in modelUsage.usage ?? [] { + guard let category = DeepSeekUsageCategory(rawValue: item.type ?? "") else { continue } + if category != .request { + let amount = Self.parseTokenAmount(item.amount) + total += amount + categoryTokens[category, default: 0] += amount + } + } + modelTokens[model] = total + } + + for costUsage in totalCosts { + guard costUsage.model != nil else { continue } + for item in costUsage.usage ?? [] { + guard let category = DeepSeekUsageCategory(rawValue: item.type ?? "") else { continue } + if category != .request { + let amount = Self.parseCostAmount(item.amount) + categoryCosts[category, default: 0] += amount + } + } + } + + let topModel = modelTokens.max { + if $0.value == $1.value { return $0.key > $1.key } + return $0.value < $1.value + }?.key + + var breakdown: [DeepSeekCategoryBreakdown] = [] + for category in [DeepSeekUsageCategory.promptCacheHitToken, .promptCacheMissToken, .responseToken] { + breakdown.append(DeepSeekCategoryBreakdown( + category: category, + tokens: categoryTokens[category] ?? 0, + cost: categoryCosts[category])) + } + + return (topModel, breakdown) + } + + private static func buildDailyUsages(ctx: DailyAggregationContext) -> [DeepSeekDailyUsage] { + var result: [DeepSeekDailyUsage] = [] + + for date in ctx.allDates.sorted() { + guard let parsed = self.parseDate(date, calendar: ctx.calendar), + parsed >= ctx.startOfMonth, + parsed <= ctx.now + else { continue } + + var dayTokens = 0 + var dayCost: Double? + var dayRequests = 0 + + if let amounts = ctx.amountMap[date] { + for items in amounts.values { + for item in items { + guard let category = DeepSeekUsageCategory(rawValue: item.type ?? "") else { continue } + if category == .request { + dayRequests += Self.parseTokenAmount(item.amount) + } else { + dayTokens += Self.parseTokenAmount(item.amount) + } + } + } + } + + if let costs = ctx.costMap[date] { + for items in costs.values { + for item in items { + guard let category = DeepSeekUsageCategory(rawValue: item.type ?? "") else { continue } + if category != .request { + let amount = Self.parseCostAmount(item.amount) + if let existing = dayCost { + dayCost = existing + amount + } else { + dayCost = amount + } + } + } + } + } + + result.append(DeepSeekDailyUsage( + date: date, + totalTokens: dayTokens, + cost: dayCost, + requestCount: dayRequests)) + } + + return result + } + + // MARK: - Helpers + + private static func parseTokenAmount(_ value: String?) -> Int { + guard let value, let intValue = Int64(value.trimmingCharacters(in: .whitespacesAndNewlines)) else { + return 0 + } + return Int(intValue) + } + + private static func parseCostAmount(_ value: String?) -> Double { + guard let value else { return 0 } + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + return Double(trimmed) ?? 0 + } + + private static func parseDate(_ text: String, calendar: Calendar) -> Date? { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.calendar = calendar + formatter.timeZone = calendar.timeZone + formatter.dateFormat = "yyyy-MM-dd" + return formatter.date(from: trimmed) + } +} diff --git a/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekUsageFetcher.swift b/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekUsageFetcher.swift new file mode 100644 index 0000000..d2f956c --- /dev/null +++ b/Sources/CodexBarCore/Providers/DeepSeek/DeepSeekUsageFetcher.swift @@ -0,0 +1,432 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +// MARK: - API response types + +public struct DeepSeekBalanceResponse: Decodable, Sendable { + public let isAvailable: Bool + public let balanceInfos: [DeepSeekBalanceInfo] + + enum CodingKeys: String, CodingKey { + case isAvailable = "is_available" + case balanceInfos = "balance_infos" + } +} + +public struct DeepSeekBalanceInfo: Decodable, Sendable { + public let currency: String + public let totalBalance: String + public let grantedBalance: String + public let toppedUpBalance: String + + enum CodingKeys: String, CodingKey { + case currency + case totalBalance = "total_balance" + case grantedBalance = "granted_balance" + case toppedUpBalance = "topped_up_balance" + } +} + +// MARK: - Domain snapshot + +public struct DeepSeekUsageSnapshot: Sendable { + public let isAvailable: Bool + public let currency: String + public let totalBalance: Double + public let grantedBalance: Double + public let toppedUpBalance: Double + public let usageSummary: DeepSeekUsageSummary? + public let updatedAt: Date + + public init( + isAvailable: Bool, + currency: String, + totalBalance: Double, + grantedBalance: Double, + toppedUpBalance: Double, + usageSummary: DeepSeekUsageSummary? = nil, + updatedAt: Date) + { + self.isAvailable = isAvailable + self.currency = currency + self.totalBalance = totalBalance + self.grantedBalance = grantedBalance + self.toppedUpBalance = toppedUpBalance + self.usageSummary = usageSummary + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + let symbol = self.currency == "CNY" ? "¥" : "$" + + let balanceDetail: String + let usedPercent: Double + if self.totalBalance <= 0 { + balanceDetail = "\(symbol)0.00 — add credits at platform.deepseek.com" + usedPercent = 100 + } else if !self.isAvailable { + balanceDetail = "Balance unavailable for API calls" + usedPercent = 100 + } else { + let total = String(format: "\(symbol)%.2f", self.totalBalance) + let paid = String(format: "\(symbol)%.2f", self.toppedUpBalance) + let granted = String(format: "\(symbol)%.2f", self.grantedBalance) + balanceDetail = "\(total) (Paid: \(paid) / Granted: \(granted))" + usedPercent = 0 + } + + let identity = ProviderIdentitySnapshot( + providerID: .deepseek, + accountEmail: nil, + accountOrganization: nil, + loginMethod: nil) + let balanceWindow = RateWindow( + usedPercent: usedPercent, + windowMinutes: nil, + resetsAt: nil, + resetDescription: balanceDetail) + + return UsageSnapshot( + primary: balanceWindow, + secondary: nil, + tertiary: nil, + providerCost: nil, + deepseekUsage: self.usageSummary, + updatedAt: self.updatedAt, + identity: identity) + } +} + +// MARK: - Errors + +public enum DeepSeekUsageError: LocalizedError, Sendable { + case missingCredentials + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing DeepSeek API key." + case let .networkError(message): + "DeepSeek network error: \(message)" + case let .apiError(message): + "DeepSeek API error: \(message)" + case let .parseFailed(message): + "Failed to parse DeepSeek response: \(message)" + } + } +} + +// MARK: - Fetcher + +public struct DeepSeekUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.deepSeekUsage) + private static let balanceURL = URL(string: "https://api.deepseek.com/user/balance")! + private static let usageAmountURL = URL(string: "https://platform.deepseek.com/api/v0/usage/amount")! + private static let usageCostURL = URL(string: "https://platform.deepseek.com/api/v0/usage/cost")! + private static let timeoutSeconds: TimeInterval = 15 + private static let optionalSummaryJoinGrace: Duration = .seconds(2) + private static var apiCalendar: Calendar { + var calendar = Calendar(identifier: .gregorian) + calendar.locale = Locale(identifier: "en_US_POSIX") + calendar.timeZone = TimeZone(secondsFromGMT: 0) ?? .gmt + return calendar + } + + public static func fetchUsage( + apiKey: String, + includeOptionalUsage: Bool = true) async throws -> DeepSeekUsageSnapshot + { + try await self.fetchUsage( + apiKey: apiKey, + includeOptionalUsage: includeOptionalUsage, + optionalSummaryJoinGrace: self.optionalSummaryJoinGrace, + fetchBalanceData: { key in + try await self.fetchBalanceData(apiKey: key) + }, + fetchSummary: { key in + try await self.fetchUsageSummary(apiKey: key) + }) + } + + static func _fetchUsageForTesting( + apiKey: String, + includeOptionalUsage: Bool, + optionalSummaryJoinGrace: Duration = .zero, + fetchBalanceData: @escaping @Sendable (String) async throws -> Data, + fetchSummary: @escaping @Sendable (String) async throws -> DeepSeekUsageSummary) + async throws -> DeepSeekUsageSnapshot + { + try await self.fetchUsage( + apiKey: apiKey, + includeOptionalUsage: includeOptionalUsage, + optionalSummaryJoinGrace: optionalSummaryJoinGrace, + fetchBalanceData: fetchBalanceData, + fetchSummary: fetchSummary) + } + + private static func fetchUsage( + apiKey: String, + includeOptionalUsage: Bool, + optionalSummaryJoinGrace: Duration, + fetchBalanceData: @escaping @Sendable (String) async throws -> Data, + fetchSummary: @escaping @Sendable (String) async throws -> DeepSeekUsageSummary) + async throws -> DeepSeekUsageSnapshot + { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw DeepSeekUsageError.missingCredentials + } + + let summaryTask: Task? = if includeOptionalUsage { + Task { + try await fetchSummary(apiKey) + } + } else { + nil + } + + let balanceData: Data + do { + balanceData = try await withTaskCancellationHandler { + try await fetchBalanceData(apiKey) + } onCancel: { + summaryTask?.cancel() + } + } catch { + summaryTask?.cancel() + throw error + } + var snapshot: DeepSeekUsageSnapshot + do { + snapshot = try Self.parseSnapshot(data: balanceData) + } catch { + summaryTask?.cancel() + throw error + } + + if let summaryTask { + let summary = try await self.completedOptionalUsageSummary( + from: summaryTask, + joinGrace: optionalSummaryJoinGrace) + if let summary { + snapshot = DeepSeekUsageSnapshot( + isAvailable: snapshot.isAvailable, + currency: snapshot.currency, + totalBalance: snapshot.totalBalance, + grantedBalance: snapshot.grantedBalance, + toppedUpBalance: snapshot.toppedUpBalance, + usageSummary: summary, + updatedAt: snapshot.updatedAt) + } + } + + return snapshot + } + + private static func fetchBalanceData(apiKey: String) async throws -> Data { + var request = URLRequest(url: self.balanceURL) + request.httpMethod = "GET" + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.timeoutInterval = Self.timeoutSeconds + + let response = try await ProviderHTTPClient.shared.response(for: request) + let data = response.data + guard response.statusCode == 200 else { + Self.log.error("DeepSeek balance endpoint returned HTTP \(response.statusCode)") + throw DeepSeekUsageError.apiError("HTTP \(response.statusCode)") + } + + return data + } + + private static func completedOptionalUsageSummary( + from task: Task, + joinGrace: Duration) async throws -> DeepSeekUsageSummary? + { + let race = BoundedTaskJoin(sourceTask: task) + switch await race.value(joinGrace: joinGrace) { + case let .value(summary): + try Task.checkCancellation() + return summary + case .timedOut: + try Task.checkCancellation() + return nil + case let .failure(error): + task.cancel() + if Task.isCancelled { + throw error + } + return nil + } + } + + static func _parseSnapshotForTesting(_ data: Data) throws -> DeepSeekUsageSnapshot { + try self.parseSnapshot(data: data) + } + + public static func fetchUsageSummary( + apiKey: String, + now: Date = Date(), + calendar: Calendar? = nil) async throws -> DeepSeekUsageSummary + { + let calendar = calendar ?? self.apiCalendar + let period = try self.usagePeriod(now: now, calendar: calendar) + + let amountData = try await self.fetchAmount(apiKey: apiKey, month: period.month, year: period.year) + let costData = try await self.fetchCost(apiKey: apiKey, month: period.month, year: period.year) + + return try DeepSeekUsageCostParser.parse( + amountData: amountData, + costData: costData, + now: now, + calendar: calendar) + } + + static func _apiUsagePeriodForTesting(now: Date, calendar: Calendar? = nil) throws -> (month: Int, year: Int) { + try self.usagePeriod(now: now, calendar: calendar ?? self.apiCalendar) + } + + private static func usagePeriod(now: Date, calendar: Calendar) throws -> (month: Int, year: Int) { + let monthComponents = calendar.dateComponents([.month, .year], from: now) + guard let month = monthComponents.month, let year = monthComponents.year else { + throw DeepSeekUsageError.parseFailed("Could not determine current month/year") + } + return (month: month, year: year) + } + + private static func fetchAmount(apiKey: String, month: Int, year: Int) async throws -> Data { + guard var components = URLComponents(url: self.usageAmountURL, resolvingAgainstBaseURL: false) else { + throw DeepSeekUsageError.networkError("Invalid URL") + } + components.queryItems = [ + URLQueryItem(name: "month", value: String(month)), + URLQueryItem(name: "year", value: String(year)), + ] + guard let url = components.url else { + throw DeepSeekUsageError.networkError("Could not construct URL") + } + + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.timeoutInterval = Self.timeoutSeconds + + let response = try await ProviderHTTPClient.shared.response(for: request) + let data = response.data + guard response.statusCode == 200 else { + if response.statusCode == 401 || response.statusCode == 403 { + throw DeepSeekUsageError.missingCredentials + } + throw DeepSeekUsageError.apiError("HTTP \(response.statusCode)") + } + + return data + } + + private static func fetchCost(apiKey: String, month: Int, year: Int) async throws -> Data { + guard var components = URLComponents(url: self.usageCostURL, resolvingAgainstBaseURL: false) else { + throw DeepSeekUsageError.networkError("Invalid URL") + } + components.queryItems = [ + URLQueryItem(name: "month", value: String(month)), + URLQueryItem(name: "year", value: String(year)), + ] + guard let url = components.url else { + throw DeepSeekUsageError.networkError("Could not construct URL") + } + + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.timeoutInterval = Self.timeoutSeconds + + let response = try await ProviderHTTPClient.shared.response(for: request) + let data = response.data + guard response.statusCode == 200 else { + if response.statusCode == 401 || response.statusCode == 403 { + throw DeepSeekUsageError.missingCredentials + } + throw DeepSeekUsageError.apiError("HTTP \(response.statusCode)") + } + + return data + } + + static func _parseUsageSummaryForTesting( + amountData: Data, + costData: Data, + now: Date = Date(), + calendar: Calendar = .current) throws -> DeepSeekUsageSummary + { + try DeepSeekUsageCostParser.parse( + amountData: amountData, + costData: costData, + now: now, + calendar: calendar) + } + + private static func parseSnapshot(data: Data) throws -> DeepSeekUsageSnapshot { + let decoded: DeepSeekBalanceResponse + do { + decoded = try JSONDecoder().decode(DeepSeekBalanceResponse.self, from: data) + } catch { + throw DeepSeekUsageError.parseFailed(error.localizedDescription) + } + + let balances = try decoded.balanceInfos.map(Self.parseBalanceInfo) + guard !balances.isEmpty else { + return DeepSeekUsageSnapshot( + isAvailable: false, + currency: "USD", + totalBalance: 0, + grantedBalance: 0, + toppedUpBalance: 0, + updatedAt: Date()) + } + + // Prefer USD when it is funded, but do not hide a positive CNY balance behind + // an empty USD row returned by the API. + let selected = balances.first { $0.currency == "USD" && $0.totalBalance > 0 } + ?? balances.first { $0.totalBalance > 0 } + ?? balances.first { $0.currency == "USD" } + ?? balances[0] + + return DeepSeekUsageSnapshot( + isAvailable: decoded.isAvailable, + currency: selected.currency, + totalBalance: selected.totalBalance, + grantedBalance: selected.grantedBalance, + toppedUpBalance: selected.toppedUpBalance, + updatedAt: Date()) + } + + private struct ParsedBalanceInfo { + let currency: String + let totalBalance: Double + let grantedBalance: Double + let toppedUpBalance: Double + } + + private static func parseBalanceInfo(_ info: DeepSeekBalanceInfo) throws -> ParsedBalanceInfo { + guard + let total = Double(info.totalBalance), + let granted = Double(info.grantedBalance), + let toppedUp = Double(info.toppedUpBalance) + else { + throw DeepSeekUsageError.parseFailed("Non-numeric balance value in response.") + } + + return ParsedBalanceInfo( + currency: info.currency, + totalBalance: total, + grantedBalance: granted, + toppedUpBalance: toppedUp) + } +} diff --git a/Sources/CodexBarCore/Providers/Deepgram/DeepgramProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Deepgram/DeepgramProviderDescriptor.swift new file mode 100644 index 0000000..5e77cb4 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Deepgram/DeepgramProviderDescriptor.swift @@ -0,0 +1,102 @@ +import Foundation + +public enum DeepgramProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .deepgram, + metadata: ProviderMetadata( + id: .deepgram, + displayName: "Deepgram", + sessionLabel: "Requests", + weeklyLabel: "Usage", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "Usage summary from Deepgram API", + toggleTitle: "Show Deepgram usage", + cliName: "deepgram", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://console.deepgram.com/project/", + statusPageURL: nil, + statusLinkURL: "https://status.deepgram.com"), + branding: ProviderBranding( + iconStyle: .deepgram, + iconResourceName: "ProviderIcon-deepgram", + color: ProviderColor( + red: 100 / 255, + green: 103 / 255, + blue: 242 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { + "Deepgram cost summary is not yet supported." + }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline( + resolveStrategies: { _ in + [DeepgramAPIFetchStrategy()] + })), + cli: ProviderCLIConfig( + name: "deepgram", + aliases: ["dg"], + versionDetector: nil)) + } +} + +struct DeepgramAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "deepgram.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + Self.resolveAPIKey(context) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveAPIKey(context) else { + throw DeepgramSettingsError.missingToken + } + + let usage = try await DeepgramUsageFetcher.fetchUsage( + apiKey: apiKey, + projectID: Self.resolveProjectID(context), + environment: context.env) + + return self.makeResult( + usage: usage.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func resolveAPIKey(_ context: ProviderFetchContext) -> String? { + ProviderTokenResolver.deepgramResolution( + type: .apiKey, + environment: context.env) + } + + private static func resolveProjectID(_ context: ProviderFetchContext) -> String? { + ProviderTokenResolver.deepgramResolution( + type: .projectID, + environment: context.env) + } +} + +/// Errors related to Deepgram settings +public enum DeepgramSettingsError: LocalizedError, Sendable { + case missingToken + + public var errorDescription: String? { + switch self { + case .missingToken: + "Deepgram API token not configured. Set DEEPGRAM_API_KEY environment variable or configure in Settings." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Deepgram/DeepgramSettingsReader.swift b/Sources/CodexBarCore/Providers/Deepgram/DeepgramSettingsReader.swift new file mode 100644 index 0000000..beb8610 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Deepgram/DeepgramSettingsReader.swift @@ -0,0 +1,33 @@ +import Foundation + +public struct DeepgramSettingsReader: Sendable { + public static let apiKeyEnvironmentKey = "DEEPGRAM_API_KEY" + public static let projectIDEnvironmentKey = "DEEPGRAM_PROJECT_ID" + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.apiKeyEnvironmentKey]) + } + + public static func projectID( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.projectIDEnvironmentKey]) + } + + private static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} diff --git a/Sources/CodexBarCore/Providers/Deepgram/DeepgramUsageFetcher.swift b/Sources/CodexBarCore/Providers/Deepgram/DeepgramUsageFetcher.swift new file mode 100644 index 0000000..0cbbc44 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Deepgram/DeepgramUsageFetcher.swift @@ -0,0 +1,561 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum DeepgramUsageError: LocalizedError, Sendable { + case missingAPIKey + case invalidEndpointOverride(String) + case invalidCredentials + case invalidProjectID + case forbidden(String) + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingAPIKey: + "Missing Deepgram API key. Set apiKey in ~/.codexbar/config.json or DEEPGRAM_API_KEY." + case let .invalidEndpointOverride(key): + "Deepgram endpoint override \(key) must use HTTPS or a bare host." + case .invalidCredentials: + "Deepgram API key is invalid or expired." + case .invalidProjectID: + "Deepgram project ID is invalid or no projects were returned for this API key." + case let .forbidden(message): + "Deepgram rejected access: \(message)" + case let .networkError(message): + "Deepgram network error: \(message)" + case let .apiError(message): + "Deepgram API error: \(message)" + case let .parseFailed(message): + "Deepgram parse error: \(message)" + } + } +} + +// MARK: - API Responses + +public struct DeepgramProjectsResponse: Decodable, Sendable { + public let projects: [DeepgramProject] +} + +public struct DeepgramProject: Decodable, Sendable, Equatable { + public let projectID: String + public let name: String? + + private enum CodingKeys: String, CodingKey { + case projectID = "project_id" + case name + } + + public init(projectID: String, name: String? = nil) { + self.projectID = projectID + self.name = name + } +} + +public struct DeepgramUsageResponse: Decodable, Sendable { + public let start: String? + public let end: String? + public let resolution: DeepgramUsageResolution? + public let results: [DeepgramUsageResult] +} + +public struct DeepgramUsageResolution: Decodable, Sendable { + public let units: String? + public let amount: Int? +} + +public struct DeepgramUsageResult: Codable, Sendable { + public let start: String? + public let end: String? + public let hours: Double? + public let totalHours: Double? + public let agentHours: Double? + public let tokensIn: Int? + public let tokensOut: Int? + public let ttsCharacters: Int? + public let requests: Int? + + private enum CodingKeys: String, CodingKey { + case start + case end + case hours + case totalHours = "total_hours" + case agentHours = "agent_hours" + case tokensIn = "tokens_in" + case tokensOut = "tokens_out" + case ttsCharacters = "tts_characters" + case requests + } +} + +// MARK: - Query + +public struct DeepgramUsageQuery: Sendable { + public var start: String? + public var end: String? + + public init( + start: String? = nil, + end: String? = nil) + { + self.start = start + self.end = end + } + + public func queryItems() -> [URLQueryItem] { + var items: [URLQueryItem] = [] + + func add(_ name: String, _ value: String?) { + guard let value = value?.trimmingCharacters(in: .whitespacesAndNewlines), + !value.isEmpty + else { + return + } + items.append(URLQueryItem(name: name, value: value)) + } + + add("start", self.start) + add("end", self.end) + + return items + } +} + +// MARK: - Snapshot + +public struct DeepgramUsageSnapshot: Codable, Sendable { + public let projectID: String + public let projectName: String? + public let projectCount: Int + public let start: String? + public let end: String? + public let hours: Double + public let totalHours: Double + public let agentHours: Double + public let tokensIn: Int + public let tokensOut: Int + public let ttsCharacters: Int + public let requests: Int + public let updatedAt: Date + + public init( + projectID: String, + projectName: String? = nil, + projectCount: Int = 1, + start: String?, + end: String?, + hours: Double, + totalHours: Double, + agentHours: Double = 0, + tokensIn: Int = 0, + tokensOut: Int = 0, + ttsCharacters: Int = 0, + requests: Int, + updatedAt: Date) + { + self.projectID = projectID + self.projectName = projectName + self.projectCount = projectCount + self.start = start + self.end = end + self.hours = hours + self.totalHours = totalHours + self.agentHours = agentHours + self.tokensIn = tokensIn + self.tokensOut = tokensOut + self.ttsCharacters = ttsCharacters + self.requests = requests + self.updatedAt = updatedAt + } +} + +extension DeepgramUsageSnapshot { + public func toUsageSnapshot() -> UsageSnapshot { + let identity = ProviderIdentitySnapshot( + providerID: .deepgram, + accountEmail: nil, + accountOrganization: nil, + loginMethod: self.identityLabel) + + return UsageSnapshot( + primary: nil, + secondary: nil, + tertiary: nil, + providerCost: nil, + deepgramUsage: self, + updatedAt: self.updatedAt, + identity: identity) + } + + public var displayLines: [String] { + var lines: [String] = [] + lines.append("Requests: \(Self.formatInteger(self.requests))") + + var usageParts: [String] = [] + if self.hours > 0 { + usageParts.append("\(Self.formatDecimal(self.hours)) audio hours") + } + if self.totalHours > 0 { + usageParts.append("\(Self.formatDecimal(self.totalHours)) billable hours") + } + if !usageParts.isEmpty { + lines.append(usageParts.joined(separator: " · ")) + } + + var modelParts: [String] = [] + if self.agentHours > 0 { + modelParts.append("\(Self.formatDecimal(self.agentHours)) agent hours") + } + if self.tokensIn > 0 || self.tokensOut > 0 { + modelParts.append("\(Self.formatInteger(self.tokensIn + self.tokensOut)) tokens") + } + if self.ttsCharacters > 0 { + modelParts.append("\(Self.formatInteger(self.ttsCharacters)) TTS chars") + } + if !modelParts.isEmpty { + lines.append(modelParts.joined(separator: " · ")) + } + + if let start, let end { + lines.append("Period: \(start) to \(end)") + } + + return lines + } + + private var identityLabel: String? { + if self.projectCount > 1 { + return "\(self.projectCount) projects" + } + if let projectName = self.projectName?.trimmingCharacters(in: .whitespacesAndNewlines), + !projectName.isEmpty + { + return "Project: \(projectName)" + } + return "Project: \(self.projectID)" + } + + fileprivate static func aggregate( + _ snapshots: [DeepgramUsageSnapshot], + updatedAt: Date) throws -> DeepgramUsageSnapshot + { + guard let first = snapshots.first else { + throw DeepgramUsageError.invalidProjectID + } + if snapshots.count == 1 { return first } + return DeepgramUsageSnapshot( + projectID: "all", + projectName: nil, + projectCount: snapshots.count, + start: snapshots.compactMap(\.start).min(), + end: snapshots.compactMap(\.end).max(), + hours: snapshots.reduce(0) { $0 + $1.hours }, + totalHours: snapshots.reduce(0) { $0 + $1.totalHours }, + agentHours: snapshots.reduce(0) { $0 + $1.agentHours }, + tokensIn: snapshots.reduce(0) { $0 + $1.tokensIn }, + tokensOut: snapshots.reduce(0) { $0 + $1.tokensOut }, + ttsCharacters: snapshots.reduce(0) { $0 + $1.ttsCharacters }, + requests: snapshots.reduce(0) { $0 + $1.requests }, + updatedAt: updatedAt) + } + + private static func formatInteger(_ value: Int) -> String { + let formatter = NumberFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.numberStyle = .decimal + formatter.usesGroupingSeparator = true + formatter.groupingSeparator = "," + formatter.maximumFractionDigits = 0 + return formatter.string(from: NSNumber(value: value)) ?? "\(value)" + } + + private static func formatDecimal(_ value: Double) -> String { + let formatter = NumberFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.numberStyle = .decimal + formatter.usesGroupingSeparator = true + formatter.groupingSeparator = "," + formatter.minimumFractionDigits = value == floor(value) ? 0 : 1 + formatter.maximumFractionDigits = 1 + return formatter.string(from: NSNumber(value: value)) ?? String(format: "%.1f", value) + } +} + +// MARK: - Fetcher + +public struct DeepgramUsageFetcher: Sendable { + public static let apiURLKey = "DEEPGRAM_API_URL" + + private static let log = CodexBarLog.logger(LogCategories.deepgramUsage) + private static let defaultBaseURL = URL(string: "https://api.deepgram.com/v1")! + + private struct FetchContext { + let apiKey: String + let query: DeepgramUsageQuery + let timeout: TimeInterval + let environment: [String: String] + let transport: ProviderHTTPTransport + let updatedAt: Date + } + + public static func fetchUsage( + apiKey: String, + projectID: String? = nil, + query: DeepgramUsageQuery = DeepgramUsageQuery(), + timeout: TimeInterval = 15, + environment: [String: String] = ProcessInfo.processInfo.environment, + transport: ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> DeepgramUsageSnapshot + { + let cleanedAPIKey = apiKey.trimmingCharacters(in: .whitespacesAndNewlines) + guard !cleanedAPIKey.isEmpty else { + throw DeepgramUsageError.missingAPIKey + } + try self.validateEndpointOverrides(environment: environment) + + let updatedAt = Date() + let context = FetchContext( + apiKey: cleanedAPIKey, + query: query, + timeout: timeout, + environment: environment, + transport: transport, + updatedAt: updatedAt) + + if let cleanedProjectID = self.cleaned(projectID) { + return try await self.fetchUsage( + project: DeepgramProject(projectID: cleanedProjectID), + context: context) + } + + let projects = try await self.listProjects( + apiKey: cleanedAPIKey, + timeout: timeout, + environment: environment, + transport: transport) + guard !projects.isEmpty else { + throw DeepgramUsageError.invalidProjectID + } + + var snapshots: [DeepgramUsageSnapshot] = [] + snapshots.reserveCapacity(projects.count) + for project in projects { + let snapshot = try await self.fetchUsage( + project: project, + context: context) + snapshots.append(snapshot) + } + return try DeepgramUsageSnapshot.aggregate(snapshots, updatedAt: updatedAt) + } + + static func _parseSnapshotForTesting( + _ data: Data, + projectID: String = "project-test", + projectName: String? = nil, + updatedAt: Date = Date()) throws -> DeepgramUsageSnapshot + { + try self.parseUsage( + data: data, + project: DeepgramProject(projectID: projectID, name: projectName), + updatedAt: updatedAt) + } + + private static func listProjects( + apiKey: String, + timeout: TimeInterval, + environment: [String: String], + transport: ProviderHTTPTransport) async throws -> [DeepgramProject] + { + let url = self.apiURL(environment: environment).appendingPathComponent("projects") + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.timeoutInterval = timeout + request.setValue("Token \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + let response = try await self.perform(request: request, transport: transport) + do { + return try JSONDecoder().decode(DeepgramProjectsResponse.self, from: response.data).projects + } catch { + Self.log.error("Deepgram projects decode failed: \(error.localizedDescription)") + throw DeepgramUsageError.parseFailed(error.localizedDescription) + } + } + + private static func fetchUsage( + project: DeepgramProject, + context: FetchContext) async throws -> DeepgramUsageSnapshot + { + let url = try self.usageURL( + projectID: project.projectID, + query: context.query, + environment: context.environment) + + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.timeoutInterval = context.timeout + request.setValue("Token \(context.apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + let response = try await self.perform(request: request, transport: context.transport) + return try self.parseUsage(data: response.data, project: project, updatedAt: context.updatedAt) + } + + private static func perform( + request: URLRequest, + transport: ProviderHTTPTransport) async throws -> ProviderHTTPResponse + { + let response: ProviderHTTPResponse + do { + response = try await transport.response(for: request) + } catch { + throw DeepgramUsageError.networkError(error.localizedDescription) + } + + guard response.statusCode == 200 else { + let summary = self.responseSummary(response.data) + Self.log.error("Deepgram returned HTTP \(response.statusCode): \(summary)") + + switch response.statusCode { + case 401: + throw DeepgramUsageError.invalidCredentials + case 403: + throw DeepgramUsageError.forbidden( + "The API key may not have access to the project or the Management API. HTTP 403: \(summary)") + case 400: + throw DeepgramUsageError.apiError("Bad request. HTTP 400: \(summary)") + default: + throw DeepgramUsageError.apiError("HTTP \(response.statusCode): \(summary)") + } + } + + return response + } + + private static func usageURL( + projectID: String, + query: DeepgramUsageQuery, + environment: [String: String]) throws -> URL + { + let usageURL = self.apiURL(environment: environment) + .appendingPathComponent("projects") + .appendingPathComponent(projectID) + .appendingPathComponent("usage") + .appendingPathComponent("breakdown") + + guard var components = URLComponents(url: usageURL, resolvingAgainstBaseURL: false) else { + throw DeepgramUsageError.networkError("Invalid usage URL") + } + + let queryItems = query.queryItems() + if !queryItems.isEmpty { + components.queryItems = queryItems + } + + guard let finalURL = components.url else { + throw DeepgramUsageError.networkError("Invalid usage query") + } + + return finalURL + } + + private static func apiURL(environment: [String: String]) -> URL { + if let raw = self.cleaned(environment[self.apiURLKey]), + let url = ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) + { + return url + } + + return self.defaultBaseURL + } + + public static func validateEndpointOverrides(environment: [String: String] = ProcessInfo.processInfo + .environment) throws + { + guard let raw = self.cleaned(environment[self.apiURLKey]) else { return } + guard ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) == nil else { return } + throw DeepgramUsageError.invalidEndpointOverride(self.apiURLKey) + } + + private static func parseUsage( + data: Data, + project: DeepgramProject, + updatedAt: Date) throws -> DeepgramUsageSnapshot + { + do { + let response = try JSONDecoder().decode(DeepgramUsageResponse.self, from: data) + return DeepgramUsageSnapshot( + projectID: project.projectID, + projectName: project.name, + start: response.start, + end: response.end, + hours: response.results.reduce(0) { $0 + ($1.hours ?? 0) }, + totalHours: response.results.reduce(0) { $0 + ($1.totalHours ?? 0) }, + agentHours: response.results.reduce(0) { $0 + ($1.agentHours ?? 0) }, + tokensIn: response.results.reduce(0) { $0 + ($1.tokensIn ?? 0) }, + tokensOut: response.results.reduce(0) { $0 + ($1.tokensOut ?? 0) }, + ttsCharacters: response.results.reduce(0) { $0 + ($1.ttsCharacters ?? 0) }, + requests: response.results.reduce(0) { $0 + ($1.requests ?? 0) }, + updatedAt: updatedAt) + } catch let error as DecodingError { + Self.log.error("Deepgram decoding error: \(error.localizedDescription)") + Self.log.error("Deepgram raw response: \(self.responseSummary(data))") + throw DeepgramUsageError.parseFailed(error.localizedDescription) + } catch { + Self.log.error("Deepgram parse error: \(error.localizedDescription)") + Self.log.error("Deepgram raw response: \(self.responseSummary(data))") + throw DeepgramUsageError.parseFailed(error.localizedDescription) + } + } + + private static func cleaned(_ raw: String?) -> String? { + guard let value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + return value + } + + private static func responseSummary(_ data: Data) -> String { + guard !data.isEmpty else { return "empty body" } + + guard let text = String(data: data, encoding: .utf8) else { + return "non-text body (\(data.count) bytes)" + } + + let cleaned = text + .replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespacesAndNewlines) + + guard !cleaned.isEmpty else { + return "empty body" + } + + let maxLength = 240 + guard cleaned.count > maxLength else { + return self.redact(cleaned) + } + + let index = cleaned.index(cleaned.startIndex, offsetBy: maxLength) + return self.redact("\(cleaned[.. String { + let replacements: [(String, String)] = [ + (#"(?i)(token\s+)[A-Za-z0-9._\-]+"#, "$1[REDACTED]"), + (#"(?i)(dg_[A-Za-z0-9._\-]+)"#, "[REDACTED]"), + ( + #"(?i)(\"(?:api_?key|authorization|token|access_token|refresh_token)\"\s*:\s*\")([^\"]+)(\")"#, + "$1[REDACTED]$3"), + ] + + return replacements.reduce(text) { partial, replacement in + partial.replacingOccurrences( + of: replacement.0, + with: replacement.1, + options: .regularExpression) + } + } +} diff --git a/Sources/CodexBarCore/Providers/Devin/DevinProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Devin/DevinProviderDescriptor.swift new file mode 100644 index 0000000..8cdcd55 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Devin/DevinProviderDescriptor.swift @@ -0,0 +1,92 @@ +import Foundation + +public enum DevinProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .devin, + metadata: ProviderMetadata( + id: .devin, + displayName: "Devin", + sessionLabel: "Daily", + weeklyLabel: "Weekly", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Devin usage", + cliName: "devin", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.devinCookieImportOrder, + dashboardURL: "https://app.devin.ai", + subscriptionDashboardURL: "https://app.devin.ai/settings/usage", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .devin, + iconResourceName: "ProviderIcon-devin", + color: ProviderColor(red: 70 / 255, green: 180 / 255, blue: 130 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Devin cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [DevinWebFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "devin", + versionDetector: nil)) + } +} + +struct DevinWebFetchStrategy: ProviderFetchStrategy { + let id: String = "devin.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + let settings = context.settings?.devin + let source = settings?.cookieSource ?? .auto + guard source != .off else { return false } + if source == .manual { + return DevinUsageFetcher.manualAuth(from: Self.bearerTokenOverride(context: context)) != nil + } + #if os(macOS) + return true + #else + return false + #endif + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let fetcher = DevinUsageFetcher(browserDetection: context.browserDetection) + let settings = context.settings?.devin + let logger: ((String) -> Void)? = context.verbose + ? { msg in CodexBarLog.logger(LogCategories.devin).verbose(msg) } + : nil + let snapshot = try await fetcher.fetch( + bearerTokenOverride: settings?.cookieSource == .manual ? Self.bearerTokenOverride(context: context) : nil, + organizationOverride: Self.organizationOverride(context: context), + timeout: context.webTimeout, + logger: logger) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "web") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func bearerTokenOverride(context: ProviderFetchContext) -> String? { + context.env["DEVIN_BEARER_TOKEN"] + ?? context.env["DEVIN_AUTHORIZATION"] + ?? context.settings?.devin?.manualBearerToken + } + + private static func organizationOverride(context: ProviderFetchContext) -> String? { + context.env["DEVIN_ORGANIZATION"] + ?? context.env["DEVIN_ORG"] + ?? context.settings?.devin?.organization + } +} diff --git a/Sources/CodexBarCore/Providers/Devin/DevinSessionImporter.swift b/Sources/CodexBarCore/Providers/Devin/DevinSessionImporter.swift new file mode 100644 index 0000000..102ed65 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Devin/DevinSessionImporter.swift @@ -0,0 +1,459 @@ +import Foundation +#if os(macOS) +import SweetCookieKit +#endif + +#if os(macOS) +enum DevinSessionImporter { + nonisolated(unsafe) static var importSessionOverrideForTesting: + ((BrowserDetection, String?, ((String) -> Void)?) -> SessionInfo?)? + + private static let storageOrigin = "https://app.devin.ai" + private static let externalOrgPrefix = "last-internal-org-for-external-org-v1-" + + struct SessionInfo: Equatable { + let accessToken: String + let organization: String? + let internalOrganizationID: String? + let sourceLabel: String + } + + struct LocalStorageCandidate { + let label: String + let url: URL + } + + static func importSession( + browserDetection: BrowserDetection, + organizationOverride: String? = nil, + logger: ((String) -> Void)? = nil) -> SessionInfo? + { + if let override = self.importSessionOverrideForTesting { + return override(browserDetection, organizationOverride, logger) + } + + let sessions = self.importSessions( + browserDetection: browserDetection, + organizationOverride: organizationOverride, + logger: logger) + return sessions.first + } + + static func importSessions( + browserDetection: BrowserDetection, + organizationOverride: String? = nil, + logger: ((String) -> Void)? = nil) -> [SessionInfo] + { + if let override = self.importSessionOverrideForTesting { + return override(browserDetection, organizationOverride, logger).map { [$0] } ?? [] + } + + let log: (String) -> Void = { msg in logger?("[devin-storage] \(msg)") } + let candidates = self.chromeLocalStorageCandidates(browserDetection: browserDetection) + if !candidates.isEmpty { + log("Chrome local storage candidates: \(candidates.count)") + } + + var sessions: [SessionInfo] = [] + for candidate in candidates { + let storage = self.readLocalStorage(from: candidate.url, logger: log) + guard let session = self.session( + from: storage, + organizationOverride: organizationOverride, + sourceLabel: candidate.label) + else { + continue + } + log( + "Found Devin session in \(candidate.label); " + + "organization=\(session.organization != nil), internalOrganizationID=" + + "\(session.internalOrganizationID != nil)") + sessions.append(session) + } + sessions = self.rankSessions(self.deduplicateSessions(sessions)) + + if sessions.isEmpty { + log("No Devin session found in browser local storage") + } + return sessions + } + + static func session( + from storage: [String: String], + organizationOverride: String? = nil, + sourceLabel: String) -> SessionInfo? + { + guard let accessToken = self.accessToken(from: storage) else { + return nil + } + let organizationInfo = self.organizationInfo(from: storage, organizationOverride: organizationOverride) + return SessionInfo( + accessToken: accessToken, + organization: organizationInfo.organization, + internalOrganizationID: organizationInfo.internalOrganizationID, + sourceLabel: sourceLabel) + } + + static func accessToken(from storage: [String: String]) -> String? { + for (key, value) in storage where self.isAuth1StorageKey(key) { + guard let json = self.jsonObject(from: value), + let token = self.findAuth1Token(in: json) + else { + continue + } + return token + } + + for (key, value) in storage where self.isAuth0StorageKey(key) { + guard let json = self.jsonObject(from: value), + let token = self.findAccessToken(in: json) + else { + continue + } + return token + } + + for value in storage.values { + guard let json = self.jsonObject(from: value), + let token = self.findAccessToken(in: json) + else { + continue + } + return token + } + + return nil + } + + static func deduplicateSessions(_ sessions: [SessionInfo]) -> [SessionInfo] { + var order: [String] = [] + var bestByToken: [String: SessionInfo] = [:] + for session in sessions { + if let existing = bestByToken[session.accessToken] { + if self.organizationScore(session) > self.organizationScore(existing) { + bestByToken[session.accessToken] = session + } + } else { + order.append(session.accessToken) + bestByToken[session.accessToken] = session + } + } + return order.compactMap { bestByToken[$0] } + } + + static func rankSessions(_ sessions: [SessionInfo]) -> [SessionInfo] { + sessions.enumerated() + .sorted { lhs, rhs in + let lhsScore = self.organizationScore(lhs.element) + let rhsScore = self.organizationScore(rhs.element) + return lhsScore == rhsScore ? lhs.offset < rhs.offset : lhsScore > rhsScore + } + .map(\.element) + } + + private static func organizationScore(_ session: SessionInfo) -> Int { + (session.organization == nil ? 0 : 1) + (session.internalOrganizationID == nil ? 0 : 2) + } + + static func organizationInfo( + from storage: [String: String], + organizationOverride: String?) -> (organization: String?, internalOrganizationID: String?) + { + let override = DevinUsageFetcher.normalizedOrganization(organizationOverride) + let overrideSlug = override.flatMap(self.slug(fromNormalizedOrganization:)) + var firstInternalOrgID: String? + + for (key, value) in storage where self.isExternalOrgStorageKey(key) { + let suffix = self.externalOrgSlug(from: key) + let orgID = self.cleanedOrgID(value) + if firstInternalOrgID == nil { + firstInternalOrgID = orgID + } + if let overrideSlug, suffix == overrideSlug { + return (override, orgID) + } + if override == nil, suffix != "null" { + return ("org/\(suffix)", orgID) + } + } + + if let inferred = self.inferredOrganizationInfo(from: storage, override: override) { + return inferred + } + + if let override { + return (override, firstInternalOrgID ?? self.orgID(fromNormalizedOrganization: override)) + } + + return (firstInternalOrgID.map { "organizations/\($0)" }, firstInternalOrgID) + } + + static func decodedStorageValue(_ value: String) -> String { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return trimmed } + if let data = trimmed.data(using: .utf8), + let decoded = try? JSONDecoder().decode(String.self, from: data) + { + return decoded.trimmingCharacters(in: .whitespacesAndNewlines) + } + return trimmed.trimmingCharacters(in: CharacterSet(charactersIn: "\"")) + .trimmingCharacters(in: .whitespacesAndNewlines) + } + + static func chromeLocalStorageCandidates(browserDetection: BrowserDetection) -> [LocalStorageCandidate] { + let installedBrowsers = self.localStorageBrowsers(browserDetection: browserDetection) + let roots = ChromiumProfileLocator + .roots(for: installedBrowsers, homeDirectories: BrowserCookieClient.defaultHomeDirectories()) + .map { (url: $0.url, labelPrefix: $0.labelPrefix) } + + var candidates: [LocalStorageCandidate] = [] + for root in roots { + candidates.append(contentsOf: self.chromeProfileLocalStorageDirs( + root: root.url, + labelPrefix: root.labelPrefix)) + } + return candidates + } + + static func localStorageBrowsers(browserDetection: BrowserDetection) -> [Browser] { + let order = ProviderDefaults.metadata[.devin]?.browserCookieOrder ?? [.chrome] + return order.browsersWithProfileData(using: browserDetection) + } + + private static func chromeProfileLocalStorageDirs(root: URL, labelPrefix: String) -> [LocalStorageCandidate] { + guard let entries = try? FileManager.default.contentsOfDirectory( + at: root, + includingPropertiesForKeys: [.isDirectoryKey], + options: [.skipsHiddenFiles]) + else { return [] } + + return entries.filter { url in + guard let isDir = (try? url.resourceValues(forKeys: [.isDirectoryKey]).isDirectory), isDir else { + return false + } + let name = url.lastPathComponent + return name == "Default" || name.hasPrefix("Profile ") || name.hasPrefix("user-") + } + .sorted { $0.lastPathComponent < $1.lastPathComponent } + .compactMap { dir in + let levelDBURL = dir.appendingPathComponent("Local Storage").appendingPathComponent("leveldb") + guard FileManager.default.fileExists(atPath: levelDBURL.path) else { return nil } + return LocalStorageCandidate(label: "\(labelPrefix) \(dir.lastPathComponent)", url: levelDBURL) + } + } + + private static func readLocalStorage(from levelDBURL: URL, logger: ((String) -> Void)?) -> [String: String] { + var storage: [String: String] = [:] + let entries = SweetCookieKit.ChromiumLocalStorageReader.readEntries( + for: self.storageOrigin, + in: levelDBURL, + logger: logger) + for entry in entries { + storage[entry.key] = self.decodedStorageValue(entry.value) + } + + let textEntries = SweetCookieKit.ChromiumLocalStorageReader.readTextEntries( + in: levelDBURL, + logger: logger) + for entry in textEntries where storage[entry.key] == nil { + if self.isUsefulStorageKey(entry.key) { + storage[entry.key] = self.decodedStorageValue(entry.value) + } + } + + return storage + } + + private static func jsonObject(from raw: String) -> Any? { + guard let data = raw.data(using: .utf8) else { return nil } + return try? JSONSerialization.jsonObject(with: data) + } + + private static func findAuth1Token(in object: Any) -> String? { + guard let dictionary = object as? [String: Any], + let token = dictionary["token"] as? String + else { + return nil + } + let value = token.trimmingCharacters(in: .whitespacesAndNewlines) + return value.hasPrefix("auth1_") && value.count > 20 ? value : nil + } + + private static func findAccessToken(in object: Any) -> String? { + if let dictionary = object as? [String: Any] { + for key in ["access_token", "accessToken"] { + if let value = dictionary[key] as? String, + self.looksLikeToken(value) + { + return value + } + } + for value in dictionary.values { + if let found = self.findAccessToken(in: value) { + return found + } + } + } + + if let array = object as? [Any] { + for value in array { + if let found = self.findAccessToken(in: value) { + return found + } + } + } + + return nil + } + + private static func looksLikeToken(_ raw: String) -> Bool { + let value = raw.trimmingCharacters(in: .whitespacesAndNewlines) + return value.count > 20 && (value.hasPrefix("eyJ") || value.contains(".")) + } + + private static func isAuth1StorageKey(_ key: String) -> Bool { + key.hasSuffix("auth1_session") + } + + private static func isAuth0StorageKey(_ key: String) -> Bool { + key.contains("auth0spajs@@::") + } + + private static func isExternalOrgStorageKey(_ key: String) -> Bool { + key.contains(self.externalOrgPrefix) + } + + private static func isUsefulStorageKey(_ key: String) -> Bool { + self.isAuth1StorageKey(key) || + self.isAuth0StorageKey(key) || + self.isExternalOrgStorageKey(key) || + key.contains("post-auth-v") || + key.contains("member-info-v") || + key.contains("feature-flags-cache:org-") || + key.contains("feature-flags-cache:org_") + } + + private static func inferredOrganizationInfo( + from storage: [String: String], + override: String?) -> (organization: String?, internalOrganizationID: String?)? + { + let overrideSlug = override.flatMap(self.slug(fromNormalizedOrganization:)) + let overrideOrgID = override.flatMap(self.orgID(fromNormalizedOrganization:)) + var fallbackSlug: String? + var fallbackInternalOrgID: String? + + for (key, value) in storage { + let object = self.jsonObject(from: value) + let internalOrgID = self.cleanedOrgID(self.firstString( + in: object, + matching: ["internalOrgId", "internal_org_id", "org_id", "orgId"])) + ?? self.internalOrgIDFromStorageKey(key) + let slug = self.cleanedSlug( + self.slugFromPostAuthKey(key) ?? + self.firstString(in: object, matching: ["orgName", "org_name", "externalOrgId", "external_org_id"])) + + if let overrideOrgID, internalOrgID == overrideOrgID { + return (override, internalOrgID) + } + if let overrideSlug, slug == overrideSlug { + return (override, internalOrgID) + } + + if fallbackSlug == nil, let slug { + fallbackSlug = slug + } + if fallbackInternalOrgID == nil, let internalOrgID { + fallbackInternalOrgID = internalOrgID + } + } + + if let override, fallbackInternalOrgID != nil { + return (override, fallbackInternalOrgID) + } + + if let fallbackSlug { + return ("org/\(fallbackSlug)", fallbackInternalOrgID) + } + if let fallbackInternalOrgID { + return ("organizations/\(fallbackInternalOrgID)", fallbackInternalOrgID) + } + + return nil + } + + private static func externalOrgSlug(from key: String) -> String { + guard let range = key.range(of: self.externalOrgPrefix) else { return key } + return String(key[range.upperBound...]) + } + + private static func cleanedOrgID(_ raw: String) -> String? { + let value = self.decodedStorageValue(raw) + guard DevinUsageFetcher.isInternalOrganizationID(value) else { return nil } + return value + } + + private static func cleanedOrgID(_ raw: String?) -> String? { + guard let raw else { return nil } + return self.cleanedOrgID(raw) + } + + private static func cleanedSlug(_ raw: String?) -> String? { + guard let raw else { return nil } + let value = self.decodedStorageValue(raw) + .trimmingCharacters(in: CharacterSet(charactersIn: "/")) + .trimmingCharacters(in: .whitespacesAndNewlines) + guard !value.isEmpty, value != "null", !DevinUsageFetcher.isInternalOrganizationID(value) else { + return nil + } + if value.hasPrefix("org/") { + return String(value.dropFirst(4)) + } + return value + } + + private static func slugFromPostAuthKey(_ key: String) -> String? { + guard let range = key.range(of: "-org_name-") else { return nil } + return String(key[range.upperBound...]) + } + + private static func internalOrgIDFromStorageKey(_ key: String) -> String? { + guard let range = key.range(of: #"org[-_][A-Za-z0-9]{8,}"#, options: .regularExpression) else { + return nil + } + return self.cleanedOrgID(String(key[range])) + } + + private static func firstString(in object: Any?, matching keys: Set) -> String? { + if let dictionary = object as? [String: Any] { + for (key, value) in dictionary { + if keys.contains(key), let string = value as? String, !string.isEmpty { + return string + } + if let found = self.firstString(in: value, matching: keys) { + return found + } + } + } + + if let array = object as? [Any] { + for value in array { + if let found = self.firstString(in: value, matching: keys) { + return found + } + } + } + + return nil + } + + private static func slug(fromNormalizedOrganization organization: String) -> String? { + guard organization.hasPrefix("org/") else { return nil } + return String(organization.dropFirst(4)) + } + + private static func orgID(fromNormalizedOrganization organization: String) -> String? { + guard organization.hasPrefix("organizations/") else { return nil } + return String(organization.dropFirst("organizations/".count)) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Devin/DevinUsageFetcher.swift b/Sources/CodexBarCore/Providers/Devin/DevinUsageFetcher.swift new file mode 100644 index 0000000..5ae7657 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Devin/DevinUsageFetcher.swift @@ -0,0 +1,271 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct DevinUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.devin) + private static let baseURL = URL(string: "https://app.devin.ai")! + private static let defaultUserAgent = + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " + + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36" + + public struct RequestAuth: Sendable, Equatable { + public let bearerToken: String + public let organization: String? + public let internalOrganizationID: String? + public let sourceLabel: String + + public init( + bearerToken: String, + organization: String?, + internalOrganizationID: String?, + sourceLabel: String) + { + self.bearerToken = bearerToken + self.organization = organization + self.internalOrganizationID = internalOrganizationID + self.sourceLabel = sourceLabel + } + } + + public let browserDetection: BrowserDetection + + public init(browserDetection: BrowserDetection) { + self.browserDetection = browserDetection + } + + public func fetch( + bearerTokenOverride: String? = nil, + organizationOverride: String? = nil, + timeout: TimeInterval = 15, + logger: ((String) -> Void)? = nil, + now: Date = Date(), + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> DevinUsageSnapshot + { + let auths = try self.resolveAuths( + bearerTokenOverride: bearerTokenOverride, + organizationOverride: organizationOverride, + logger: logger) + var lastError: Error? + for auth in auths { + do { + return try await Self.fetchQuotaUsage( + auth: auth, + organizationOverride: organizationOverride, + timeout: timeout, + logger: logger, + now: now, + transport: transport) + } catch { + lastError = error + logger?("[devin] Session from \(auth.sourceLabel) failed: \(error.localizedDescription)") + if auth.sourceLabel == "manual" || !Self.shouldTryNextSession(after: error) { + throw error + } + } + } + throw lastError ?? DevinUsageError.noSession + } + + public static func fetchQuotaUsage( + auth: RequestAuth, + organizationOverride: String? = nil, + timeout: TimeInterval = 15, + logger: ((String) -> Void)? = nil, + now: Date = Date(), + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> DevinUsageSnapshot + { + let organization = self.normalizedOrganization(organizationOverride) ?? + self.normalizedOrganization(auth.organization) + guard let organization else { + throw DevinUsageError.missingOrganization + } + + var lastError: Error? + for path in self.candidatePaths( + organization: organization, + internalOrganizationID: auth.internalOrganizationID) + { + let data: Data + do { + data = try await self.fetch( + path: path, + auth: auth, + timeout: timeout, + transport: transport) + } catch { + lastError = error + logger?("[devin] /api/\(path) failed: \(error.localizedDescription)") + if case DevinUsageError.invalidCredentials = error { + throw error + } + continue + } + logger?("[devin] Fetched quota usage from /api/\(path)") + return try DevinUsageParser.parse(data, organization: organization, now: now) + } + + throw lastError ?? DevinUsageError.apiError("No Devin quota endpoint succeeded.") + } + + public static func manualAuth(from raw: String?, organization: String? = nil) -> RequestAuth? { + guard var token = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !token.isEmpty else { + return nil + } + if token.lowercased().hasPrefix("authorization:") { + token = token.dropHeaderName().trimmingCharacters(in: .whitespacesAndNewlines) + } + if token.lowercased().hasPrefix("bearer ") { + token = String(token.dropFirst(7)).trimmingCharacters(in: .whitespacesAndNewlines) + } + guard !token.isEmpty else { return nil } + return RequestAuth( + bearerToken: token, + organization: self.normalizedOrganization(organization), + internalOrganizationID: self.internalOrganizationID(from: organization), + sourceLabel: "manual") + } + + private func resolveAuths( + bearerTokenOverride: String?, + organizationOverride: String?, + logger: ((String) -> Void)?) throws -> [RequestAuth] + { + if let manual = Self.manualAuth(from: bearerTokenOverride, organization: organizationOverride) { + logger?("[devin] Using manual Bearer token") + return [manual] + } + + #if os(macOS) + let normalizedOrganizationOverride = Self.normalizedOrganization(organizationOverride) + let sessions = DevinSessionImporter.importSessions( + browserDetection: self.browserDetection, + organizationOverride: normalizedOrganizationOverride, + logger: logger) + guard !sessions.isEmpty else { + throw DevinUsageError.noSession + } + logger?("[devin] Found \(sessions.count) browser session(s)") + return sessions.map { session in + RequestAuth( + bearerToken: session.accessToken, + organization: normalizedOrganizationOverride ?? Self.normalizedOrganization(session.organization), + internalOrganizationID: session.internalOrganizationID, + sourceLabel: session.sourceLabel) + } + #else + throw DevinUsageError.noSession + #endif + } + + static func shouldTryNextSession(after error: Error) -> Bool { + switch error { + case DevinUsageError.invalidCredentials, DevinUsageError.apiError, DevinUsageError.missingOrganization: + true + default: + false + } + } + + private static func fetch( + path: String, + auth: RequestAuth, + timeout: TimeInterval, + transport: any ProviderHTTPTransport) async throws -> Data + { + let url = self.baseURL.appending(path: "api/\(path)") + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.timeoutInterval = timeout + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("en-US,en;q=0.9", forHTTPHeaderField: "Accept-Language") + request.setValue(self.defaultUserAgent, forHTTPHeaderField: "User-Agent") + request.setValue("Bearer \(auth.bearerToken)", forHTTPHeaderField: "Authorization") + if let internalOrganizationID = auth.internalOrganizationID { + request.setValue(internalOrganizationID, forHTTPHeaderField: "x-cog-org-id") + } + let response = try await transport.response(for: request) + guard response.statusCode == 200 else { + let body = String(data: response.data.prefix(200), encoding: .utf8) ?? "" + if response.statusCode == 401 || response.statusCode == 403 { + throw DevinUsageError.invalidCredentials + } + Self.log.error("Devin API returned \(response.statusCode): \(body)") + throw DevinUsageError.apiError("HTTP \(response.statusCode)") + } + return response.data + } + + private static func candidatePaths(organization: String, internalOrganizationID: String?) -> [String] { + var paths: [String] = [] + let normalized = self.normalizedOrganization(organization) ?? organization + if let internalOrganizationID { + paths.append("\(internalOrganizationID)/billing/quota/usage") + } + paths.append("\(normalized)/billing/quota/usage") + if normalized.hasPrefix("org/") { + let slug = String(normalized.dropFirst(4)) + paths.append("\(slug)/billing/quota/usage") + } + if !normalized.hasPrefix("org/"), !normalized.hasPrefix("organizations/") { + paths.append("org/\(normalized)/billing/quota/usage") + } + if let internalOrganizationID { + paths.append("organizations/\(internalOrganizationID)/billing/quota/usage") + } + return paths.removingDuplicates() + } + + public static func normalizedOrganization(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + if let url = URL(string: value), + let host = url.host?.lowercased(), + host == "devin.ai" || host.hasSuffix(".devin.ai") + { + let components = url.path.split(separator: "/").map(String.init) + if components.count >= 2, components[0] == "org" { + value = "org/\(components[1])" + } else if components.count >= 2, components[0] == "organizations" { + value = "organizations/\(components[1])" + } + } + value = value.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + if value.hasPrefix("org/") || value.hasPrefix("organizations/") { + return value + } + if self.isInternalOrganizationID(value) { + return "organizations/\(value)" + } + return "org/\(value)" + } + + private static func internalOrganizationID(from raw: String?) -> String? { + guard let normalized = self.normalizedOrganization(raw), + normalized.hasPrefix("organizations/") + else { + return nil + } + return String(normalized.dropFirst("organizations/".count)) + } + + static func isInternalOrganizationID(_ value: String) -> Bool { + value.hasPrefix("org-") || value.hasPrefix("org_") + } +} + +extension String { + fileprivate func dropHeaderName() -> String { + guard let index = self.firstIndex(of: ":") else { return self } + return String(self[self.index(after: index)...]) + } +} + +extension [String] { + fileprivate func removingDuplicates() -> [String] { + var seen = Set() + return self.filter { seen.insert($0).inserted } + } +} diff --git a/Sources/CodexBarCore/Providers/Devin/DevinUsageSnapshot.swift b/Sources/CodexBarCore/Providers/Devin/DevinUsageSnapshot.swift new file mode 100644 index 0000000..8b7b2b7 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Devin/DevinUsageSnapshot.swift @@ -0,0 +1,345 @@ +import CoreFoundation +import Foundation + +public enum DevinUsageError: LocalizedError, Sendable { + case noSession + case missingOrganization + case invalidCredentials + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .noSession: + "No Devin browser session found. Please log in to app.devin.ai or paste a Bearer token." + case .missingOrganization: + "No Devin organization was found. Open an app.devin.ai/org/... page " + + "or set the organization in Devin settings." + case .invalidCredentials: + "Devin session token is invalid or expired." + case let .apiError(message): + "Devin API error: \(message)" + case let .parseFailed(message): + "Could not parse Devin usage: \(message)" + } + } +} + +public struct DevinQuotaWindow: Sendable, Equatable { + public let usedPercent: Double + public let resetsAt: Date? + + public init(usedPercent: Double, resetsAt: Date? = nil) { + self.usedPercent = min(100, max(0, usedPercent)) + self.resetsAt = resetsAt + } +} + +public struct DevinUsageSnapshot: Sendable, Equatable { + public let daily: DevinQuotaWindow? + public let weekly: DevinQuotaWindow? + public let planName: String? + public let organization: String? + public let updatedAt: Date + public let overageBalance: Double? + + public init( + daily: DevinQuotaWindow?, + weekly: DevinQuotaWindow?, + planName: String?, + organization: String?, + updatedAt: Date, + overageBalance: Double? = nil) + { + self.daily = daily + self.weekly = weekly + self.planName = planName + self.organization = organization + self.updatedAt = updatedAt + self.overageBalance = overageBalance + } + + public func toUsageSnapshot() -> UsageSnapshot { + let primary = self.daily.map { + RateWindow( + usedPercent: $0.usedPercent, + windowMinutes: 24 * 60, + resetsAt: $0.resetsAt, + resetDescription: "Daily") + } + let secondary = self.weekly.map { + RateWindow( + usedPercent: $0.usedPercent, + windowMinutes: 7 * 24 * 60, + resetsAt: $0.resetsAt, + resetDescription: "Weekly") + } + let identity = ProviderIdentitySnapshot( + providerID: .devin, + accountEmail: nil, + accountOrganization: self.organization, + loginMethod: self.planName) + let providerCost: ProviderCostSnapshot? = self.overageBalance.map { + ProviderCostSnapshot( + used: $0, + limit: 0, + currencyCode: "USD", + period: "Extra usage balance", + updatedAt: self.updatedAt) + } + return UsageSnapshot( + primary: primary, + secondary: secondary, + providerCost: providerCost, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public enum DevinUsageParser { + public static func parse(_ data: Data, organization: String?, now: Date = Date()) throws -> DevinUsageSnapshot { + let object = try JSONSerialization.jsonObject(with: data) + return try self.parse(object, organization: organization, now: now) + } + + public static func parse(_ object: Any, organization: String?, now: Date = Date()) throws -> DevinUsageSnapshot { + let current = (object as? [String: Any]).map(self.currentQuotaWindows) + let daily = current?.daily ?? self.findWindow(in: object, matching: self.isDailyKey) + let weekly = current?.weekly ?? self.findWindow(in: object, matching: self.isWeeklyKey) + guard daily != nil || weekly != nil else { + throw DevinUsageError.parseFailed("Missing Devin quota windows.") + } + + return DevinUsageSnapshot( + daily: daily, + weekly: weekly, + planName: self.findPlanName(in: object), + organization: self.displayOrganization(from: organization), + updatedAt: now, + overageBalance: self.findOverageBalance(in: object)) + } + + private static func findOverageBalance(in object: Any) -> Double? { + guard let dictionary = object as? [String: Any] else { return nil } + if let value = self.nonnegativeFiniteDouble(dictionary["overage_balance"]) { return value } + if let cents = self.nonnegativeFiniteDouble(dictionary["overage_balance_cents"]) { return cents / 100.0 } + return nil + } + + private static func nonnegativeFiniteDouble(_ value: Any?) -> Double? { + guard let value = self.double(value), value.isFinite, value >= 0 else { return nil } + return value + } + + private static func currentQuotaWindows(_ dictionary: [String: Any]) + -> (daily: DevinQuotaWindow?, weekly: DevinQuotaWindow?) + { + let daily = self.currentQuotaWindow( + percent: dictionary["daily_percentage"], + resetsAt: dictionary["daily_reset_at"]) + let weekly = self.currentQuotaWindow( + percent: dictionary["weekly_percentage"], + resetsAt: dictionary["weekly_reset_at"]) + return (daily, weekly) + } + + private static func currentQuotaWindow(percent: Any?, resetsAt: Any?) -> DevinQuotaWindow? { + guard let usedPercent = self.double(percent) else { return nil } + return DevinQuotaWindow( + usedPercent: usedPercent < 1 ? usedPercent * 100 : usedPercent, + resetsAt: self.date(from: resetsAt)) + } + + private static func findWindow(in object: Any, matching keyMatches: (String) -> Bool) -> DevinQuotaWindow? { + if let dictionary = object as? [String: Any] { + for (key, value) in dictionary where keyMatches(key) { + if let window = self.window(from: value) { + return window + } + } + for value in dictionary.values { + if let found = self.findWindow(in: value, matching: keyMatches) { + return found + } + } + } + + if let array = object as? [Any] { + for value in array { + if let found = self.findWindow(in: value, matching: keyMatches) { + return found + } + } + } + + return nil + } + + private static func window(from object: Any) -> DevinQuotaWindow? { + guard let dictionary = object as? [String: Any] else { + guard let percent = self.percent(from: object) else { return nil } + return DevinQuotaWindow(usedPercent: percent, resetsAt: nil) + } + + if let percent = self.percent(from: dictionary) { + return DevinQuotaWindow( + usedPercent: percent, + resetsAt: self.findResetDate(in: dictionary)) + } + + if let nested = dictionary.values.lazy.compactMap({ self.window(from: $0) }).first { + return nested + } + + return nil + } + + private static func percent(from object: Any) -> Double? { + if let number = self.double(object) { + return number <= 1 ? number * 100 : number + } + guard let dictionary = object as? [String: Any] else { return nil } + + let directKeys = [ + "used_percent", + "usedPercent", + "usage_percent", + "usagePercent", + "percent_used", + "percentUsed", + "percent", + ] + for key in directKeys { + if let value = self.double(dictionary[key]) { + return value <= 1 ? value * 100 : value + } + } + + let remainingKeys = ["remaining_percent", "remainingPercent", "percent_remaining", "percentRemaining"] + for key in remainingKeys { + if let value = self.double(dictionary[key]) { + let percent = value <= 1 ? value * 100 : value + return 100 - percent + } + } + + let used = self.firstDouble(in: dictionary, keys: ["used", "usage", "used_count", "usedCount", "consumed"]) + let limit = self.firstDouble(in: dictionary, keys: ["limit", "quota", "total", "max", "available"]) + if let used, let limit, limit > 0 { + return used / limit * 100 + } + + let remaining = self.firstDouble(in: dictionary, keys: ["remaining", "left", "available"]) + if let remaining, let limit, limit > 0 { + return (limit - remaining) / limit * 100 + } + + return nil + } + + private static func findPlanName(in object: Any) -> String? { + if let dictionary = object as? [String: Any] { + for key in ["plan_name", "planName", "plan", "tier", "subscription_tier", "subscriptionTier"] { + if let value = dictionary[key] as? String, + let cleaned = self.cleanDisplay(value) + { + return cleaned + } + } + for value in dictionary.values { + if let found = self.findPlanName(in: value) { + return found + } + } + } + + if let array = object as? [Any] { + for value in array { + if let found = self.findPlanName(in: value) { + return found + } + } + } + + return nil + } + + private static func findResetDate(in dictionary: [String: Any]) -> Date? { + for (key, value) in dictionary where key.localizedCaseInsensitiveContains("reset") { + if let date = self.date(from: value) { + return date + } + } + return nil + } + + private static func date(from value: Any?) -> Date? { + if let raw = value as? String { + if let date = ISO8601DateFormatter().date(from: raw) { + return date + } + if let number = Double(raw) { + return self.date(from: number) + } + } + if let number = self.double(value) { + return self.date(from: number) + } + return nil + } + + private static func date(from number: Double) -> Date? { + guard number > 0 else { return nil } + let seconds = number > 10_000_000_000 ? number / 1000 : number + return Date(timeIntervalSince1970: seconds) + } + + private static func firstDouble(in dictionary: [String: Any], keys: [String]) -> Double? { + for key in keys { + if let value = self.double(dictionary[key]) { + return value + } + } + return nil + } + + private static func double(_ value: Any?) -> Double? { + switch value { + case let number as NSNumber: + CFGetTypeID(number) == CFBooleanGetTypeID() ? nil : number.doubleValue + case let string as String: + Double(string.trimmingCharacters(in: .whitespacesAndNewlines)) + default: + nil + } + } + + private static func isDailyKey(_ raw: String) -> Bool { + let key = raw.lowercased() + return !key.contains("hide") && (key.contains("daily") || key.contains("day")) + } + + private static func isWeeklyKey(_ raw: String) -> Bool { + let key = raw.lowercased() + return !key.contains("hide") && (key.contains("weekly") || key.contains("week")) + } + + private static func displayOrganization(from raw: String?) -> String? { + guard let raw = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { return nil } + if raw.hasPrefix("org/") { + return String(raw.dropFirst(4)) + } + if raw.hasPrefix("organizations/") { + return String(raw.dropFirst("organizations/".count)) + } + return raw + } + + private static func cleanDisplay(_ raw: String) -> String? { + let cleaned = raw.trimmingCharacters(in: .whitespacesAndNewlines) + guard !cleaned.isEmpty else { return nil } + return cleaned.split(separator: "_").flatMap { $0.split(separator: "-") }.map { part in + part.prefix(1).uppercased() + String(part.dropFirst()) + }.joined(separator: " ") + } +} diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift new file mode 100644 index 0000000..386c039 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoProviderDescriptor.swift @@ -0,0 +1,110 @@ +import Foundation + +public enum DoubaoProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + public static func primaryLabel(window: RateWindow?) -> String? { + guard window?.windowMinutes == nil, + window?.resetDescription?.localizedCaseInsensitiveContains("request") == true + else { + return nil + } + return "Requests" + } + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .doubao, + metadata: ProviderMetadata( + id: .doubao, + displayName: "Doubao", + sessionLabel: "5-hour", + weeklyLabel: "Weekly", + opusLabel: "Monthly", + supportsOpus: true, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Doubao usage", + cliName: "doubao", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://console.volcengine.com/ark/region:ark+cn-beijing/openManagement?LLM=%7B%7D&advancedActiveKey=subscribe", + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .doubao, + iconResourceName: "ProviderIcon-doubao", + color: ProviderColor(red: 51 / 255, green: 112 / 255, blue: 255 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Doubao cost summary is not available." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in + [DoubaoAPIFetchStrategy()] + })), + cli: ProviderCLIConfig( + name: "doubao", + aliases: ["volcengine", "ark", "bytedance"], + versionDetector: nil)) + } +} + +struct DoubaoAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "doubao.api" + let kind: ProviderFetchKind = .apiToken + private let codingPlanUsageLoader: @Sendable (DoubaoCodingPlanCredentials) async throws -> DoubaoUsageSnapshot + private let arkUsageLoader: @Sendable (String) async throws -> DoubaoUsageSnapshot + + init( + codingPlanUsageLoader: @escaping @Sendable (DoubaoCodingPlanCredentials) async throws + -> DoubaoUsageSnapshot = { credentials in + try await DoubaoUsageFetcher.fetchCodingPlanUsage(credentials: credentials) + }, + arkUsageLoader: @escaping @Sendable (String) async throws -> DoubaoUsageSnapshot = { apiKey in + try await DoubaoUsageFetcher.fetchUsage(apiKey: apiKey) + }) + { + self.codingPlanUsageLoader = codingPlanUsageLoader + self.arkUsageLoader = arkUsageLoader + } + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + DoubaoSettingsReader.codingPlanCredentials(environment: context.env) != nil || + ProviderTokenResolver.doubaoToken(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let apiKey = ProviderTokenResolver.doubaoToken(environment: context.env) + if let credentials = DoubaoSettingsReader.codingPlanCredentials(environment: context.env) { + do { + let usage = try await self.codingPlanUsageLoader(credentials) + return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "api") + } catch { + if Self.isCancellation(error) { + throw error + } + guard let apiKey else { + throw error + } + let usage = try await self.arkUsageLoader(apiKey) + return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "api") + } + } + + guard let apiKey else { + throw DoubaoUsageError.missingCredentials + } + let usage = try await self.arkUsageLoader(apiKey) + return self.makeResult(usage: usage.toUsageSnapshot(), sourceLabel: "api") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func isCancellation(_ error: Error) -> Bool { + error is CancellationError || (error as? URLError)?.code == .cancelled || Task.isCancelled + } +} diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift new file mode 100644 index 0000000..479d281 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoSettingsReader.swift @@ -0,0 +1,86 @@ +import Foundation + +public struct DoubaoSettingsReader: Sendable { + public static let apiKeyEnvironmentKeys = [ + "ARK_API_KEY", + "VOLCENGINE_API_KEY", + "DOUBAO_API_KEY", + ] + public static let accessKeyIDEnvironmentKeys = [ + "VOLCENGINE_ACCESS_KEY_ID", + "VOLCENGINE_ACCESS_KEY", + "VOLC_ACCESSKEY", + "DOUBAO_ACCESS_KEY_ID", + ] + public static let secretAccessKeyEnvironmentKeys = [ + "VOLCENGINE_SECRET_ACCESS_KEY", + "VOLCENGINE_SECRET_KEY", + "VOLCENGINE_ACCESS_KEY_SECRET", + "VOLC_SECRETKEY", + "DOUBAO_SECRET_ACCESS_KEY", + ] + public static let regionEnvironmentKeys = [ + "VOLCENGINE_REGION", + "VOLCENGINE_REGION_ID", + "VOLC_REGION", + "DOUBAO_REGION", + ] + public static let defaultRegion = "cn-beijing" + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.firstValue(in: environment, keys: self.apiKeyEnvironmentKeys) + } + + public static func accessKeyID( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.firstValue(in: environment, keys: self.accessKeyIDEnvironmentKeys) + } + + public static func secretAccessKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.firstValue(in: environment, keys: self.secretAccessKeyEnvironmentKeys) + } + + public static func region(environment: [String: String] = ProcessInfo.processInfo.environment) -> String { + self.firstValue(in: environment, keys: self.regionEnvironmentKeys) ?? self.defaultRegion + } + + public static func codingPlanCredentials( + environment: [String: String] = ProcessInfo.processInfo.environment) -> DoubaoCodingPlanCredentials? + { + guard let accessKeyID = self.accessKeyID(environment: environment), + let secretAccessKey = self.secretAccessKey(environment: environment) + else { + return nil + } + return DoubaoCodingPlanCredentials( + accessKeyID: accessKeyID, + secretAccessKey: secretAccessKey, + region: self.region(environment: environment)) + } + + private static func firstValue(in environment: [String: String], keys: [String]) -> String? { + for key in keys { + guard let cleaned = self.cleaned(environment[key]) else { continue } + return cleaned + } + return nil + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} diff --git a/Sources/CodexBarCore/Providers/Doubao/DoubaoUsageFetcher.swift b/Sources/CodexBarCore/Providers/Doubao/DoubaoUsageFetcher.swift new file mode 100644 index 0000000..cc5b7a0 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Doubao/DoubaoUsageFetcher.swift @@ -0,0 +1,548 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct DoubaoUsageSnapshot: Sendable { + public let remainingRequests: Int + public let limitRequests: Int + public let resetTime: Date? + public let updatedAt: Date + public let apiKeyValid: Bool + public let totalTokens: Int? + public let requestLimitsReliable: Bool + public let codingPlanUsage: DoubaoCodingPlanUsage? + public init( + remainingRequests: Int, + limitRequests: Int, + resetTime: Date?, + updatedAt: Date, + apiKeyValid: Bool = false, + totalTokens: Int? = nil, + requestLimitsReliable: Bool = true, + codingPlanUsage: DoubaoCodingPlanUsage? = nil) + { + self.remainingRequests = remainingRequests + self.limitRequests = limitRequests + self.resetTime = resetTime + self.updatedAt = updatedAt + self.apiKeyValid = apiKeyValid + self.totalTokens = totalTokens + self.requestLimitsReliable = requestLimitsReliable + self.codingPlanUsage = codingPlanUsage + } + + public func toUsageSnapshot() -> UsageSnapshot { + if let codingPlanUsage { + return codingPlanUsage.toUsageSnapshot(updatedAt: self.updatedAt) + } + + let primary: RateWindow? + if self.limitRequests > 0, self.requestLimitsReliable { + let used = max(0, self.limitRequests - self.remainingRequests) + primary = RateWindow( + usedPercent: min(100, max(0, Double(used) / Double(self.limitRequests) * 100)), + windowMinutes: nil, + resetsAt: self.resetTime, + resetDescription: "\(used)/\(self.limitRequests) requests") + } else if self.apiKeyValid { + // Ark can return successful requests without a trustworthy request-limit window. + // Omitting the window prevents the UI from presenting unknown usage as 100% left. + primary = nil + } else { + primary = RateWindow( + usedPercent: 0, + windowMinutes: nil, + resetsAt: self.resetTime, + resetDescription: "No usage data") + } + + let identity = ProviderIdentitySnapshot( + providerID: .doubao, + accountEmail: nil, + accountOrganization: nil, + loginMethod: nil) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + providerCost: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public struct DoubaoCodingPlanUsage: Sendable, Equatable { + public struct Quota: Sendable, Equatable { + public let level: String + public let percent: Double + public let resetTime: Date? + + public init(level: String, percent: Double, resetTime: Date?) { + self.level = level + self.percent = percent + self.resetTime = resetTime + } + } + + public let status: String? + public let updateTime: Date? + public let quotas: [Quota] + + public init(status: String?, updateTime: Date?, quotas: [Quota]) { + self.status = status + self.updateTime = updateTime + self.quotas = quotas + } + + public func toUsageSnapshot(updatedAt: Date) -> UsageSnapshot { + let primary = self.rateWindow(levels: ["session", "5-hour", "five_hour"], minutes: 5 * 60) + let secondary = self.rateWindow(levels: ["weekly", "week"], minutes: 7 * 24 * 60) + let tertiary = self.rateWindow(levels: ["monthly", "month"], minutes: 30 * 24 * 60) + let identity = ProviderIdentitySnapshot( + providerID: .doubao, + accountEmail: nil, + accountOrganization: nil, + loginMethod: self.status) + + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: tertiary, + providerCost: nil, + updatedAt: self.updateTime ?? updatedAt, + identity: identity) + } + + private func rateWindow(levels: Set, minutes: Int) -> RateWindow? { + guard let quota = self.quotas.first(where: { levels.contains($0.level.lowercased()) }) else { + return nil + } + let percent = min(100, max(0, quota.percent)) + return RateWindow( + usedPercent: percent, + windowMinutes: minutes, + resetsAt: quota.resetTime, + resetDescription: nil) + } +} + +public enum DoubaoUsageError: LocalizedError, Sendable { + case missingCredentials + case networkError(String) + case apiError(Int, String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Doubao API key (ARK_API_KEY)." + case let .networkError(message): + "Doubao network error: \(message)" + case let .apiError(code, message): + "Doubao API error (\(code)): \(message)" + case let .parseFailed(message): + "Failed to parse Doubao response: \(message)" + } + } +} + +public struct DoubaoUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.doubaoUsage) + private static let apiURL = URL(string: "https://ark.cn-beijing.volces.com/api/coding/v3/chat/completions")! + private static let codingPlanAPIURL = URL( + string: "https://open.volcengineapi.com/?Action=GetCodingPlanUsage&Version=2024-01-01")! + + /// Models to probe, ordered by likelihood. We try multiple models because + /// different key types may not have access to every model. + private static let probeModels = [ + "doubao-seed-2.0-code", + "doubao-1.5-pro-32k", + "doubao-lite-32k", + ] + + private struct ProbeResult { + let snapshot: DoubaoUsageSnapshot + let statusCode: Int + + var hasAmbiguousZeroRemaining: Bool { + self.statusCode == 200 + && self.snapshot.requestLimitsReliable + && self.snapshot.limitRequests > 0 + && self.snapshot.remainingRequests == 0 + } + } + + public static func fetchUsage( + apiKey: String, + session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) async throws -> DoubaoUsageSnapshot + { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw DoubaoUsageError.missingCredentials + } + + var lastError: Error? + for model in self.probeModels { + do { + let result = try await self.probe(apiKey: apiKey, model: model, transport: transport) + guard result.hasAmbiguousZeroRemaining else { + return result.snapshot + } + + return try await self.confirmAmbiguousZeroRemaining( + initial: result, + apiKey: apiKey, + model: model, + transport: transport) + } catch let error as DoubaoUsageError { + if case let .apiError(code, _) = error, code == 404 || code == 403 { + Self.log.debug("Doubao probe model \(model) unavailable (\(code)), trying next") + lastError = error + continue + } + throw error + } + } + throw lastError ?? DoubaoUsageError.apiError(0, "All probe models failed") + } + + public static func fetchCodingPlanUsage( + credentials: DoubaoCodingPlanCredentials, + session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + date: Date = Date()) async throws -> DoubaoUsageSnapshot + { + guard !credentials.accessKeyID.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty, + !credentials.secretAccessKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + else { + throw DoubaoUsageError.missingCredentials + } + + let body = Data() + var request = URLRequest(url: self.codingPlanAPIURL) + request.httpMethod = "POST" + request.timeoutInterval = 15 + request.httpBody = body + request.setValue("application/json", forHTTPHeaderField: "Accept") + DoubaoVolcengineSigner.sign( + request: &request, + body: body, + credentials: credentials, + date: date) + + let response = try await transport.response(for: request) + guard response.statusCode == 200 else { + let summary = Self.apiErrorSummary(statusCode: response.statusCode, data: response.data) + Self.log.error("Doubao coding plan API returned \(response.statusCode): \(summary)") + throw DoubaoUsageError.apiError(response.statusCode, summary) + } + + let codingPlanUsage = try self.decodeCodingPlanUsage(from: response.data) + return DoubaoUsageSnapshot( + remainingRequests: 0, + limitRequests: 0, + resetTime: nil, + updatedAt: codingPlanUsage.updateTime ?? date, + apiKeyValid: true, + codingPlanUsage: codingPlanUsage) + } + + static func decodeCodingPlanUsage(from data: Data) throws -> DoubaoCodingPlanUsage { + let response: CodingPlanUsageResponse + do { + response = try JSONDecoder().decode(CodingPlanUsageResponse.self, from: data) + } catch { + throw DoubaoUsageError.parseFailed(error.localizedDescription) + } + let usage = response.result + let quotas = usage.quotaUsage.map { quota in + DoubaoCodingPlanUsage.Quota( + level: quota.level, + percent: quota.percent, + resetTime: self.date(fromEpoch: quota.resetTimestamp)) + } + return DoubaoCodingPlanUsage( + status: usage.status, + updateTime: self.date(fromEpoch: usage.updateTimestamp), + quotas: quotas) + } + + private static func date(fromEpoch timestamp: TimeInterval?) -> Date? { + guard let timestamp, timestamp > 0 else { return nil } + return Date(timeIntervalSince1970: timestamp) + } + + private static func confirmAmbiguousZeroRemaining( + initial: ProbeResult, + apiKey: String, + model: String, + transport: any ProviderHTTPTransport) async throws -> DoubaoUsageSnapshot + { + do { + let confirmation = try await self.probe(apiKey: apiKey, model: model, transport: transport) + // This path starts only after a complete HTTP 200 request-limit pair + // reported zero. An immediate 429 confirms that exhausted state even + // when Ark omits the headers from the throttle response. + if confirmation.statusCode == 429 { + return confirmation.snapshot.requestLimitsReliable + ? confirmation.snapshot + : initial.snapshot + } + guard confirmation.hasAmbiguousZeroRemaining else { + return confirmation.snapshot + } + + Self.log.warning( + """ + Doubao Ark returned limit=\(confirmation.snapshot.limitRequests) remaining=0 \ + with HTTP 200 twice; treating request-limit headers as unreliable. + """) + return DoubaoUsageSnapshot( + remainingRequests: confirmation.snapshot.remainingRequests, + limitRequests: confirmation.snapshot.limitRequests, + resetTime: confirmation.snapshot.resetTime, + updatedAt: confirmation.snapshot.updatedAt, + apiKeyValid: confirmation.snapshot.apiKeyValid, + totalTokens: confirmation.snapshot.totalTokens, + requestLimitsReliable: false) + } catch { + if error is CancellationError || (error as? URLError)?.code == .cancelled { + throw error + } + self.log.warning( + """ + Doubao zero-remaining confirmation failed; preserving the initial exhausted state: \ + \(error.localizedDescription) + """) + return initial.snapshot + } + } + + private static func probe( + apiKey: String, + model: String, + transport: any ProviderHTTPTransport) async throws -> ProbeResult + { + var request = URLRequest(url: self.apiURL) + request.httpMethod = "POST" + request.timeoutInterval = 15 + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + + let body: [String: Any] = [ + "model": model, + "max_tokens": 1, + "messages": [ + ["role": "user", "content": "hi"], + ] as [[String: Any]], + ] + + request.httpBody = try JSONSerialization.data(withJSONObject: body) + + let response = try await transport.response(for: request) + let data = response.data + + // Accept both 200 (success) and 429 (rate limited) – both carry rate limit headers. + guard response.statusCode == 200 || response.statusCode == 429 else { + let summary = Self.apiErrorSummary(statusCode: response.statusCode, data: data) + Self.log.error("Doubao API returned \(response.statusCode): \(summary)") + throw DoubaoUsageError.apiError(response.statusCode, summary) + } + + let headers = response.response.allHeaderFields + let remaining = Self.intHeader(headers, "x-ratelimit-remaining-requests") + let limit = Self.intHeader(headers, "x-ratelimit-limit-requests") + let resetString = Self.stringHeader(headers, "x-ratelimit-reset-requests") + + let resetTime: Date? = resetString.flatMap(Self.parseResetTime) + + var totalTokens: Int? + if remaining == nil, limit == nil, + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let usage = json["usage"] as? [String: Any] + { + totalTokens = usage["total_tokens"] as? Int + } + + // 429 means the key is valid but rate-limited; treat it as valid so the UI + // shows "Active" instead of "No usage data" when headers are absent. + let keyValid = response.statusCode == 200 || response.statusCode == 429 + // A request-limit header on 429 identifies request-bucket exhaustion even + // when Ark omits remaining. A bare 429 may describe another throttle. + let requestLimitsReliable = response.statusCode == 429 + ? limit != nil + : limit != nil && remaining != nil + + let snapshot = DoubaoUsageSnapshot( + remainingRequests: remaining ?? 0, + limitRequests: limit ?? 0, + resetTime: resetTime, + updatedAt: Date(), + apiKeyValid: keyValid, + totalTokens: totalTokens, + requestLimitsReliable: requestLimitsReliable) + + Self.log.debug( + """ + Doubao usage parsed remaining=\(snapshot.remainingRequests) \ + limit=\(snapshot.limitRequests) valid=\(snapshot.apiKeyValid) + """) + + return ProbeResult(snapshot: snapshot, statusCode: response.statusCode) + } + + private static func stringHeader(_ headers: [AnyHashable: Any], _ name: String) -> String? { + if let value = headers[name] as? String { return value } + for (key, val) in headers { + if let keyStr = key as? String, + keyStr.caseInsensitiveCompare(name) == .orderedSame, + let valStr = val as? String + { + return valStr + } + } + return nil + } + + private static func intHeader(_ headers: [AnyHashable: Any], _ name: String) -> Int? { + if let value = headers[name] as? String, let int = Int(value) { + return int + } + if let value = headers[name.lowercased()] as? String, let int = Int(value) { + return int + } + for (key, val) in headers { + if let keyStr = key as? String, + keyStr.lowercased() == name.lowercased(), + let valStr = val as? String, + let int = Int(valStr) + { + return int + } + } + return nil + } + + private static func parseResetTime(_ value: String) -> Date? { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.isEmpty { return nil } + + let isoFormatter = ISO8601DateFormatter() + isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = isoFormatter.date(from: trimmed) { return date } + let isoFallback = ISO8601DateFormatter() + isoFallback.formatOptions = [.withInternetDateTime] + if let date = isoFallback.date(from: trimmed) { return date } + + var seconds: TimeInterval = 0 + let pattern = /(\d+)([dhms])/ + for match in trimmed.matches(of: pattern) { + guard let num = Double(match.1) else { continue } + switch match.2 { + case "d": seconds += num * 86400 + case "h": seconds += num * 3600 + case "m": seconds += num * 60 + case "s": seconds += num + default: break + } + } + if seconds > 0 { + return Date().addingTimeInterval(seconds) + } + + if let secs = TimeInterval(trimmed) { + return Date().addingTimeInterval(secs) + } + + return nil + } + + private static func apiErrorSummary(statusCode: Int, data: Data) -> String { + guard let root = try? JSONSerialization.jsonObject(with: data), + let json = root as? [String: Any] + else { + if let text = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines), + !text.isEmpty + { + return self.compactText(text) + } + return "Unexpected response body (\(data.count) bytes)." + } + + // Volcengine Top OpenAPI error shape: { "ResponseMetadata": { "Error": { "Code": ..., "Message": ... } } } + if let metadata = json["ResponseMetadata"] as? [String: Any], + let volcError = metadata["Error"] as? [String: Any] + { + let code = (volcError["Code"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + let message = (volcError["Message"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + switch (code?.isEmpty == false ? code : nil, message?.isEmpty == false ? message : nil) { + case let (code?, message?): + return Self.compactText("\(code): \(message)") + case let (code?, nil): + return Self.compactText(code) + case let (nil, message?): + return Self.compactText(message) + case (nil, nil): + break + } + } + + if let error = json["error"] as? [String: Any], + let message = error["message"] as? String + { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + if let message = json["message"] as? String { + let trimmed = message.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return Self.compactText(trimmed) } + } + + return "HTTP \(statusCode) (\(data.count) bytes)." + } + + private static func compactText(_ text: String, maxLength: Int = 200) -> String { + let collapsed = text + .components(separatedBy: .newlines) + .joined(separator: " ") + .trimmingCharacters(in: .whitespacesAndNewlines) + if collapsed.count <= maxLength { return collapsed } + let limitIndex = collapsed.index(collapsed.startIndex, offsetBy: maxLength) + return "\(collapsed[.. String + { + let dateKey = Self.hmac(key: SymmetricKey(data: Data(secretAccessKey.utf8)), message: dateStamp) + let regionKey = Self.hmac(key: SymmetricKey(data: dateKey), message: region) + let serviceKey = Self.hmac(key: SymmetricKey(data: regionKey), message: Self.service) + let signingKey = Self.hmac(key: SymmetricKey(data: serviceKey), message: Self.terminator) + let signature = HMAC.authenticationCode( + for: Data(stringToSign.utf8), + using: SymmetricKey(data: signingKey)) + return Data(signature).map { String(format: "%02x", $0) }.joined() + } + + private static func hmac(key: SymmetricKey, message: String) -> Data { + let digest = HMAC.authenticationCode(for: Data(message.utf8), using: key) + return Data(digest) + } + + private static func sha256Hex(_ data: Data) -> String { + SHA256.hash(data: data).map { String(format: "%02x", $0) }.joined() + } + + private static func canonicalURI(_ url: URL) -> String { + let path = url.path.isEmpty ? "/" : url.path + return Self.percentEncode(path, encodeSlash: false) + } + + private static func canonicalQueryString(_ url: URL) -> String { + guard let components = URLComponents(url: url, resolvingAgainstBaseURL: false), + let queryItems = components.queryItems, + !queryItems.isEmpty + else { + return "" + } + var pairs: [(key: String, value: String)] = queryItems.map { item in + ( + key: Self.percentEncode(item.name), + value: Self.percentEncode(item.value ?? "")) + } + pairs.sort { lhs, rhs in + lhs.key == rhs.key ? lhs.value < rhs.value : lhs.key < rhs.key + } + return pairs + .map { pair in "\(pair.key)=\(pair.value)" } + .joined(separator: "&") + } + + private static func percentEncode(_ value: String, encodeSlash: Bool = true) -> String { + var allowed = CharacterSet.alphanumerics + allowed.insert(charactersIn: "-_.~") + if !encodeSlash { + allowed.insert("/") + } + return value.addingPercentEncoding(withAllowedCharacters: allowed) ?? value + } + + private static let timestampFormatter: DateFormatter = { + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = TimeZone(secondsFromGMT: 0) + formatter.dateFormat = "yyyyMMdd'T'HHmmss'Z'" + return formatter + }() + + private static let dateFormatter: DateFormatter = { + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = TimeZone(secondsFromGMT: 0) + formatter.dateFormat = "yyyyMMdd" + return formatter + }() +} diff --git a/Sources/CodexBarCore/Providers/ElevenLabs/ElevenLabsProviderDescriptor.swift b/Sources/CodexBarCore/Providers/ElevenLabs/ElevenLabsProviderDescriptor.swift new file mode 100644 index 0000000..1981614 --- /dev/null +++ b/Sources/CodexBarCore/Providers/ElevenLabs/ElevenLabsProviderDescriptor.swift @@ -0,0 +1,49 @@ +import Foundation + +public enum ElevenLabsProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .elevenlabs, + metadata: ProviderMetadata( + id: .elevenlabs, + displayName: "ElevenLabs", + sessionLabel: "Credits", + weeklyLabel: "Voices", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show ElevenLabs usage", + cliName: "elevenlabs", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://elevenlabs.io/app/developers/usage", + subscriptionDashboardURL: "https://elevenlabs.io/app/subscription", + statusPageURL: nil, + statusLinkURL: "https://status.elevenlabs.io"), + branding: ProviderBranding( + iconStyle: .elevenlabs, + iconResourceName: "ProviderIcon-elevenlabs", + color: ProviderColor(red: 0.92, green: 0.92, blue: 0.90)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "ElevenLabs cost history is not available via API yet." }), + fetchPlan: .apiToken( + strategyID: "elevenlabs.api", + resolveToken: { ProviderTokenResolver.elevenLabsToken(environment: $0) }, + missingCredentialsError: { ElevenLabsUsageError.missingCredentials }, + loadUsage: { apiKey, context in + try await ElevenLabsUsageFetcher.fetchUsage( + apiKey: apiKey, + environment: context.env).toUsageSnapshot() + }), + cli: ProviderCLIConfig( + name: "elevenlabs", + aliases: ["11labs", "eleven"], + versionDetector: nil)) + } +} diff --git a/Sources/CodexBarCore/Providers/ElevenLabs/ElevenLabsSettingsReader.swift b/Sources/CodexBarCore/Providers/ElevenLabs/ElevenLabsSettingsReader.swift new file mode 100644 index 0000000..2ffecec --- /dev/null +++ b/Sources/CodexBarCore/Providers/ElevenLabs/ElevenLabsSettingsReader.swift @@ -0,0 +1,62 @@ +import Foundation + +public enum ElevenLabsSettingsReader { + public static let apiKeyEnvironmentKey = "ELEVENLABS_API_KEY" + public static let apiKeyEnvironmentKeys = [ + Self.apiKeyEnvironmentKey, + "XI_API_KEY", + ] + public static let apiURLEnvironmentKey = "ELEVENLABS_API_URL" + + public static func apiKey(environment: [String: String] = ProcessInfo.processInfo.environment) -> String? { + for key in self.apiKeyEnvironmentKeys { + guard let token = self.cleaned(environment[key]) else { continue } + return token + } + return nil + } + + public static func apiURL(environment: [String: String] = ProcessInfo.processInfo.environment) -> URL { + if let override = self.validAPIURL(environment: environment) { + return override + } + return URL(string: "https://api.elevenlabs.io")! + } + + public static func validateEndpointOverrides( + environment: [String: String] = ProcessInfo.processInfo.environment) throws + { + guard let raw = self.cleaned(environment[self.apiURLEnvironmentKey]) else { return } + guard ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) == nil else { return } + throw ElevenLabsSettingsError.invalidEndpointOverride(self.apiURLEnvironmentKey) + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } + + private static func validAPIURL(environment: [String: String]) -> URL? { + guard let raw = self.cleaned(environment[self.apiURLEnvironmentKey]) else { return nil } + return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) + } +} + +public enum ElevenLabsSettingsError: LocalizedError, Sendable, Equatable { + case invalidEndpointOverride(String) + + public var errorDescription: String? { + switch self { + case let .invalidEndpointOverride(key): + "ElevenLabs endpoint override \(key) must use HTTPS or a bare host." + } + } +} diff --git a/Sources/CodexBarCore/Providers/ElevenLabs/ElevenLabsUsageFetcher.swift b/Sources/CodexBarCore/Providers/ElevenLabs/ElevenLabsUsageFetcher.swift new file mode 100644 index 0000000..9befb3c --- /dev/null +++ b/Sources/CodexBarCore/Providers/ElevenLabs/ElevenLabsUsageFetcher.swift @@ -0,0 +1,252 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct ElevenLabsOverage: Codable, Sendable, Equatable { + public let amount: String? + public let currency: String? +} + +public struct ElevenLabsSubscriptionResponse: Decodable, Sendable { + public let tier: String? + public let characterCount: Int + public let characterLimit: Int + public let voiceSlotsUsed: Int? + public let professionalVoiceSlotsUsed: Int? + public let voiceLimit: Int? + public let professionalVoiceLimit: Int? + public let currentOverage: ElevenLabsOverage? + public let status: String? + public let nextCharacterCountResetUnix: Int? + + private enum CodingKeys: String, CodingKey { + case tier + case characterCount = "character_count" + case characterLimit = "character_limit" + case voiceSlotsUsed = "voice_slots_used" + case professionalVoiceSlotsUsed = "professional_voice_slots_used" + case voiceLimit = "voice_limit" + case professionalVoiceLimit = "professional_voice_limit" + case currentOverage = "current_overage" + case status + case nextCharacterCountResetUnix = "next_character_count_reset_unix" + } +} + +public struct ElevenLabsUsageSnapshot: Codable, Sendable, Equatable { + public let tier: String? + public let characterCount: Int + public let characterLimit: Int + public let voiceSlotsUsed: Int? + public let professionalVoiceSlotsUsed: Int? + public let voiceLimit: Int? + public let professionalVoiceLimit: Int? + public let currentOverage: ElevenLabsOverage? + public let status: String? + public let resetsAt: Date? + public let updatedAt: Date + + public init( + tier: String?, + characterCount: Int, + characterLimit: Int, + voiceSlotsUsed: Int?, + professionalVoiceSlotsUsed: Int?, + voiceLimit: Int?, + professionalVoiceLimit: Int?, + currentOverage: ElevenLabsOverage?, + status: String?, + resetsAt: Date?, + updatedAt: Date) + { + self.tier = tier + self.characterCount = characterCount + self.characterLimit = characterLimit + self.voiceSlotsUsed = voiceSlotsUsed + self.professionalVoiceSlotsUsed = professionalVoiceSlotsUsed + self.voiceLimit = voiceLimit + self.professionalVoiceLimit = professionalVoiceLimit + self.currentOverage = currentOverage + self.status = status + self.resetsAt = resetsAt + self.updatedAt = updatedAt + } + + public var usedPercent: Double { + guard self.characterLimit > 0 else { return 0 } + return max(0, Double(self.characterCount) / Double(self.characterLimit) * 100) + } + + public var remainingCharacters: Int { + max(0, self.characterLimit - self.characterCount) + } + + public func toUsageSnapshot() -> UsageSnapshot { + let primary = RateWindow( + usedPercent: self.usedPercent, + windowMinutes: nil, + resetsAt: self.resetsAt, + resetDescription: self.characterSummary) + let extraWindows = self.voiceWindows() + let identity = ProviderIdentitySnapshot( + providerID: .elevenlabs, + accountEmail: nil, + accountOrganization: nil, + loginMethod: self.displayTier) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + extraRateWindows: extraWindows.isEmpty ? nil : extraWindows, + updatedAt: self.updatedAt, + identity: identity) + } + + private var displayTier: String? { + guard let tier = tier?.trimmingCharacters(in: .whitespacesAndNewlines), !tier.isEmpty else { + return status + } + let statusSuffix = if let status, !status.isEmpty, status.lowercased() != "active" { + " · \(status)" + } else { + "" + } + return "\(tier.replacingOccurrences(of: "_", with: " ").capitalized)\(statusSuffix)" + } + + private var characterSummary: String { + "\(Self.formatCount(self.characterCount)) / \(Self.formatCount(self.characterLimit)) credits" + } + + private func voiceWindows() -> [NamedRateWindow] { + var windows: [NamedRateWindow] = [] + if let used = voiceSlotsUsed, let limit = voiceLimit, limit > 0 { + windows.append(NamedRateWindow( + id: "voice-slots", + title: "Voice slots", + window: RateWindow( + usedPercent: Double(used) / Double(limit) * 100, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "\(used) / \(limit)"))) + } + if let used = professionalVoiceSlotsUsed, let limit = professionalVoiceLimit, limit > 0 { + windows.append(NamedRateWindow( + id: "professional-voices", + title: "Professional voices", + window: RateWindow( + usedPercent: Double(used) / Double(limit) * 100, + windowMinutes: nil, + resetsAt: nil, + resetDescription: "\(used) / \(limit)"))) + } + return windows + } + + private static func formatCount(_ value: Int) -> String { + let formatter = NumberFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.numberStyle = .decimal + formatter.usesGroupingSeparator = true + formatter.groupingSeparator = "," + formatter.maximumFractionDigits = 0 + return formatter.string(from: NSNumber(value: value)) ?? "\(value)" + } +} + +public enum ElevenLabsUsageError: LocalizedError, Sendable { + case missingCredentials + case networkError(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing ElevenLabs API key. Set apiKey in ~/.codexbar/config.json or ELEVENLABS_API_KEY." + case let .networkError(message): + "ElevenLabs network error: \(message)" + case let .apiError(message): + "ElevenLabs API error: \(message)" + case let .parseFailed(message): + "Failed to parse ElevenLabs response: \(message)" + } + } +} + +public struct ElevenLabsUsageFetcher: Sendable { + private static let log = CodexBarLog.logger(LogCategories.elevenLabsUsage) + private static let timeoutSeconds: TimeInterval = 15 + + public static func fetchUsage( + apiKey: String, + environment: [String: String] = ProcessInfo.processInfo.environment) async throws -> ElevenLabsUsageSnapshot + { + let trimmed = apiKey.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { + throw ElevenLabsUsageError.missingCredentials + } + try ElevenLabsSettingsReader.validateEndpointOverrides(environment: environment) + + let url = Self.subscriptionURL(baseURL: ElevenLabsSettingsReader.apiURL(environment: environment)) + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.setValue(trimmed, forHTTPHeaderField: "xi-api-key") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.timeoutInterval = Self.timeoutSeconds + + let response = try await ProviderHTTPClient.shared.response(for: request) + switch response.statusCode { + case 200: + return try Self.parseSnapshot(data: response.data, updatedAt: Date()) + case 401, 403: + throw ElevenLabsUsageError.missingCredentials + default: + Self.log.error("ElevenLabs API returned \(response.statusCode)") + throw ElevenLabsUsageError.apiError("HTTP \(response.statusCode)") + } + } + + static func _parseSnapshotForTesting(_ data: Data, updatedAt: Date) throws -> ElevenLabsUsageSnapshot { + try self.parseSnapshot(data: data, updatedAt: updatedAt) + } + + private static func parseSnapshot(data: Data, updatedAt: Date) throws -> ElevenLabsUsageSnapshot { + let decoded: ElevenLabsSubscriptionResponse + do { + decoded = try JSONDecoder().decode(ElevenLabsSubscriptionResponse.self, from: data) + } catch { + throw ElevenLabsUsageError.parseFailed(error.localizedDescription) + } + + let resetsAt = decoded.nextCharacterCountResetUnix.map { + Date(timeIntervalSince1970: TimeInterval($0)) + } + + return ElevenLabsUsageSnapshot( + tier: decoded.tier, + characterCount: decoded.characterCount, + characterLimit: decoded.characterLimit, + voiceSlotsUsed: decoded.voiceSlotsUsed, + professionalVoiceSlotsUsed: decoded.professionalVoiceSlotsUsed, + voiceLimit: decoded.voiceLimit, + professionalVoiceLimit: decoded.professionalVoiceLimit, + currentOverage: decoded.currentOverage, + status: decoded.status, + resetsAt: resetsAt, + updatedAt: updatedAt) + } + + private static func subscriptionURL(baseURL: URL) -> URL { + var url = baseURL + let pathComponents = url.path.split(separator: "/") + if pathComponents.last == "v1" { + url.append(path: "user/subscription") + } else { + url.append(path: "v1/user/subscription") + } + return url + } +} diff --git a/Sources/CodexBarCore/Providers/Factory/FactoryLocalStorageImporter.swift b/Sources/CodexBarCore/Providers/Factory/FactoryLocalStorageImporter.swift new file mode 100644 index 0000000..08f0cde --- /dev/null +++ b/Sources/CodexBarCore/Providers/Factory/FactoryLocalStorageImporter.swift @@ -0,0 +1,383 @@ +import Foundation +#if os(macOS) +import SQLite3 +import SweetCookieKit +#endif + +#if os(macOS) +enum FactoryLocalStorageImporter { + struct TokenInfo { + let refreshToken: String + let accessToken: String? + let organizationID: String? + let sourceLabel: String + } + + static func importWorkOSTokens( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) -> [TokenInfo] + { + let log: (String) -> Void = { msg in logger?("[factory-storage] \(msg)") } + var tokens: [TokenInfo] = [] + + let safariCandidates = self.safariLocalStorageCandidates() + let chromeCandidates = self.chromeLocalStorageCandidates(browserDetection: browserDetection) + if !safariCandidates.isEmpty { + log("Safari local storage candidates: \(safariCandidates.count)") + } + if !chromeCandidates.isEmpty { + log("Chrome local storage candidates: \(chromeCandidates.count)") + } + + let candidates = safariCandidates + chromeCandidates + for candidate in candidates { + let match: WorkOSTokenMatch? = switch candidate.kind { + case let .chromeLevelDB(levelDBURL): + self.readWorkOSToken(from: levelDBURL) + case let .safariSQLite(dbURL): + self.readWorkOSTokenFromSafariSQLite(dbURL: dbURL, logger: log) + } + guard let token = match else { continue } + log("Found WorkOS refresh token in \(candidate.label)") + tokens.append(TokenInfo( + refreshToken: token.refreshToken, + accessToken: token.accessToken, + organizationID: token.organizationID, + sourceLabel: candidate.label)) + } + + if tokens.isEmpty { + log("No WorkOS refresh token found in browser local storage") + } + + return tokens + } + + static func hasSafariWorkOSRefreshToken() -> Bool { + for candidate in self.safariLocalStorageCandidates() { + guard case let .safariSQLite(dbURL) = candidate.kind else { continue } + if self.readWorkOSTokenFromSafariSQLite(dbURL: dbURL) != nil { + return true + } + } + return false + } + + // MARK: - Chrome local storage discovery + + private enum LocalStorageSourceKind { + case chromeLevelDB(URL) + case safariSQLite(URL) + } + + private struct LocalStorageCandidate { + let label: String + let kind: LocalStorageSourceKind + } + + private static func chromeLocalStorageCandidates(browserDetection: BrowserDetection) -> [LocalStorageCandidate] { + let browsers: [Browser] = [ + .chrome, + .chromeBeta, + .chromeCanary, + .arc, + .arcBeta, + .arcCanary, + .dia, + .chatgptAtlas, + .chromium, + .helium, + ] + + // Filter to browsers with profile data to avoid unnecessary filesystem access. + let installedBrowsers = browsers.browsersWithProfileData(using: browserDetection) + + let roots = ChromiumProfileLocator + .roots(for: installedBrowsers, homeDirectories: BrowserCookieClient.defaultHomeDirectories()) + .map { (url: $0.url, labelPrefix: $0.labelPrefix) } + + var candidates: [LocalStorageCandidate] = [] + for root in roots { + candidates.append(contentsOf: self.chromeProfileLocalStorageDirs( + root: root.url, + labelPrefix: root.labelPrefix)) + } + return candidates + } + + private static func chromeProfileLocalStorageDirs(root: URL, labelPrefix: String) -> [LocalStorageCandidate] { + guard let entries = try? FileManager.default.contentsOfDirectory( + at: root, + includingPropertiesForKeys: [.isDirectoryKey], + options: [.skipsHiddenFiles]) + else { return [] } + + let profileDirs = entries.filter { url in + guard let isDir = (try? url.resourceValues(forKeys: [.isDirectoryKey]).isDirectory), isDir else { + return false + } + let name = url.lastPathComponent + return name == "Default" || name.hasPrefix("Profile ") || name.hasPrefix("user-") + } + .sorted { $0.lastPathComponent < $1.lastPathComponent } + + return profileDirs.compactMap { dir in + let levelDBURL = dir.appendingPathComponent("Local Storage").appendingPathComponent("leveldb") + guard FileManager.default.fileExists(atPath: levelDBURL.path) else { return nil } + let label = "\(labelPrefix) \(dir.lastPathComponent)" + return LocalStorageCandidate(label: label, kind: .chromeLevelDB(levelDBURL)) + } + } + + private static func safariLocalStorageCandidates() -> [LocalStorageCandidate] { + let root = FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent("Library") + .appendingPathComponent("Containers") + .appendingPathComponent("com.apple.Safari") + .appendingPathComponent("Data") + .appendingPathComponent("Library") + .appendingPathComponent("WebKit") + .appendingPathComponent("WebsiteData") + .appendingPathComponent("Default") + + guard FileManager.default.fileExists(atPath: root.path) else { return [] } + + let targets = ["app.factory.ai", "auth.factory.ai"] + var candidates: [LocalStorageCandidate] = [] + + guard let enumerator = FileManager.default.enumerator( + at: root, + includingPropertiesForKeys: [.isRegularFileKey], + options: [.skipsHiddenFiles]) + else { return [] } + + for case let fileURL as URL in enumerator { + guard fileURL.lastPathComponent == "origin" else { continue } + guard (try? fileURL.resourceValues(forKeys: [.isRegularFileKey]).isRegularFile) == true else { + continue + } + guard let data = try? Data(contentsOf: fileURL, options: [.mappedIfSafe]) else { continue } + let ascii = String(data: data, encoding: .utf8) ?? String(data: data, encoding: .isoLatin1) ?? "" + guard targets.contains(where: { ascii.contains($0) }) else { continue } + + let storageURL = fileURL.deletingLastPathComponent() + .appendingPathComponent("LocalStorage") + .appendingPathComponent("localstorage.sqlite3") + guard FileManager.default.fileExists(atPath: storageURL.path) else { continue } + let host = self.extractSafariOriginHost(from: ascii) ?? "app.factory.ai" + candidates.append(LocalStorageCandidate(label: "Safari (\(host))", kind: .safariSQLite(storageURL))) + } + + var seen = Set() + return candidates.filter { candidate in + let key: String = switch candidate.kind { + case let .chromeLevelDB(url): + url.path + case let .safariSQLite(url): + url.path + } + if seen.contains(key) { return false } + seen.insert(key) + return true + } + } + + private static func extractSafariOriginHost(from ascii: String) -> String? { + let targets = ["app.factory.ai", "auth.factory.ai", "factory.ai"] + for host in targets where ascii.contains(host) { + return host + } + return nil + } + + // MARK: - Token extraction + + private struct WorkOSTokenMatch { + let refreshToken: String + let accessToken: String? + let organizationID: String? + } + + private static func readWorkOSToken(from levelDBURL: URL) -> WorkOSTokenMatch? { + guard let entries = try? FileManager.default.contentsOfDirectory( + at: levelDBURL, + includingPropertiesForKeys: [.contentModificationDateKey], + options: [.skipsHiddenFiles]) + else { return nil } + + let files = entries.filter { url in + let ext = url.pathExtension.lowercased() + return ext == "ldb" || ext == "log" + } + .sorted { lhs, rhs in + let left = (try? lhs.resourceValues(forKeys: [.contentModificationDateKey]).contentModificationDate) + let right = (try? rhs.resourceValues(forKeys: [.contentModificationDateKey]).contentModificationDate) + return (left ?? .distantPast) > (right ?? .distantPast) + } + + for file in files { + guard let data = try? Data(contentsOf: file, options: [.mappedIfSafe]) else { continue } + if let match = self.extractWorkOSToken(from: data) { + return match + } + } + return nil + } + + private static func extractWorkOSToken(from data: Data) -> WorkOSTokenMatch? { + guard let contents = String(data: data, encoding: .utf8) ?? + String(data: data, encoding: .isoLatin1) + else { return nil } + guard contents.contains("workos:refresh-token") else { return nil } + + let refreshToken = self.matchToken( + in: contents, + pattern: "workos:refresh-token[^A-Za-z0-9_-]*([A-Za-z0-9_-]{20,})") + guard let refreshToken else { return nil } + + let accessToken = self.matchToken( + in: contents, + pattern: "workos:access-token[^A-Za-z0-9_-]*([A-Za-z0-9_-]{20,})") + + let organizationID = self.extractOrganizationID(from: accessToken) + return WorkOSTokenMatch( + refreshToken: refreshToken, + accessToken: accessToken, + organizationID: organizationID) + } + + private static func readWorkOSTokenFromSafariSQLite( + dbURL: URL, + logger: ((String) -> Void)? = nil) -> WorkOSTokenMatch? + { + var db: OpaquePointer? + guard sqlite3_open_v2(dbURL.path, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { + if let c = sqlite3_errmsg(db) { + logger?("Safari local storage open failed: \(String(cString: c))") + } + return nil + } + defer { sqlite3_close(db) } + + sqlite3_busy_timeout(db, 250) + let tables = self.fetchTableNames(db: db, logger: logger) + if tables.isEmpty { + logger?("Safari local storage table lookup returned no tables") + } + let table = tables + .contains("ItemTable") ? "ItemTable" : (tables.contains("localstorage") ? "localstorage" : nil) + guard let table else { + logger?("Safari local storage missing ItemTable/localstorage tables (found: \(tables.sorted()))") + return nil + } + + let refreshToken = self.fetchLocalStorageValue(db: db, table: table, key: "workos:refresh-token") + guard let refreshToken, !refreshToken.isEmpty else { + logger?("Safari local storage missing workos:refresh-token") + return nil + } + let accessToken = self.fetchLocalStorageValue(db: db, table: table, key: "workos:access-token") + + let organizationID = self.extractOrganizationID(from: accessToken) + return WorkOSTokenMatch( + refreshToken: refreshToken, + accessToken: accessToken, + organizationID: organizationID) + } + + private static func extractOrganizationID(from accessToken: String?) -> String? { + guard let accessToken, accessToken.contains(".") else { return nil } + let parts = accessToken.split(separator: ".") + guard parts.count == 3 else { return nil } + let payload = String(parts[1]) + .replacingOccurrences(of: "-", with: "+") + .replacingOccurrences(of: "_", with: "/") + let padded = payload + String(repeating: "=", count: (4 - payload.count % 4) % 4) + guard let data = Data(base64Encoded: padded), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return nil + } + return json["org_id"] as? String + } + + private static func fetchTableNames(db: OpaquePointer?, logger: ((String) -> Void)? = nil) -> Set { + let sql = "SELECT name FROM sqlite_master WHERE type='table'" + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, sql, -1, &stmt, nil) == SQLITE_OK else { + if let c = sqlite3_errmsg(db) { + logger?("Safari local storage table query failed: \(String(cString: c))") + } + return [] + } + defer { sqlite3_finalize(stmt) } + var names = Set() + while true { + let step = sqlite3_step(stmt) + if step == SQLITE_ROW { + if let c = sqlite3_column_text(stmt, 0) { + names.insert(String(cString: c)) + } + } else { + if step != SQLITE_DONE, let c = sqlite3_errmsg(db) { + logger?("Safari local storage table query failed: \(String(cString: c))") + } + break + } + } + return names + } + + private static func fetchLocalStorageValue(db: OpaquePointer?, table: String, key: String) -> String? { + let sql = "SELECT value FROM \(table) WHERE key = ? LIMIT 1" + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, sql, -1, &stmt, nil) == SQLITE_OK else { return nil } + defer { sqlite3_finalize(stmt) } + let transient = unsafeBitCast(-1, to: sqlite3_destructor_type.self) + _ = key.withCString { cString in + sqlite3_bind_text(stmt, 1, cString, -1, transient) + } + guard sqlite3_step(stmt) == SQLITE_ROW else { return nil } + return self.decodeSQLiteValue(stmt: stmt, index: 0) + } + + private static func decodeSQLiteValue(stmt: OpaquePointer?, index: Int32) -> String? { + let type = sqlite3_column_type(stmt, index) + switch type { + case SQLITE_TEXT: + guard let c = sqlite3_column_text(stmt, index) else { return nil } + return String(cString: c) + case SQLITE_BLOB: + guard let bytes = sqlite3_column_blob(stmt, index) else { return nil } + let count = Int(sqlite3_column_bytes(stmt, index)) + let data = Data(bytes: bytes, count: count) + return self.decodeValueData(data) + default: + return nil + } + } + + private static func decodeValueData(_ data: Data) -> String? { + if let decoded = String(data: data, encoding: .utf16LittleEndian) { + return decoded.trimmingCharacters(in: .controlCharacters) + } + if let decoded = String(data: data, encoding: .utf8) { + return decoded.trimmingCharacters(in: .controlCharacters) + } + if let decoded = String(data: data, encoding: .isoLatin1) { + return decoded.trimmingCharacters(in: .controlCharacters) + } + return nil + } + + private static func matchToken(in contents: String, pattern: String) -> String? { + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { return nil } + let range = NSRange(contents.startIndex.. 1, + let tokenRange = Range(match.range(at: 1), in: contents) + else { return nil } + return String(contents[tokenRange]) + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Factory/FactoryManualCredentials.swift b/Sources/CodexBarCore/Providers/Factory/FactoryManualCredentials.swift new file mode 100644 index 0000000..ae4698a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Factory/FactoryManualCredentials.swift @@ -0,0 +1,58 @@ +import Foundation + +extension FactoryStatusProbe { + struct ManualCredentials { + let cookieHeader: String? + let bearerToken: String? + } + + static func manualCredentials(from raw: String?) -> ManualCredentials? { + guard let raw = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { + return nil + } + let normalizedCookieHeader = CookieHeaderNormalizer.normalize(raw) + let cookieHeader = normalizedCookieHeader.flatMap { + CookieHeaderNormalizer.pairs(from: $0).isEmpty ? nil : $0 + } + let bearerToken = self.authorizationBearerToken(from: raw) + ?? normalizedCookieHeader.flatMap(self.bearerToken(fromHeader:)) + ?? self.bareBearerToken(from: raw) + guard cookieHeader != nil || bearerToken != nil else { return nil } + return ManualCredentials(cookieHeader: cookieHeader, bearerToken: bearerToken) + } + + static func bearerToken(fromHeader cookieHeader: String) -> String? { + for pair in CookieHeaderNormalizer.pairs(from: cookieHeader) where pair.name == "access-token" { + let token = pair.value.trimmingCharacters(in: .whitespacesAndNewlines) + if !token.isEmpty { return token } + } + return nil + } + + private static func authorizationBearerToken(from raw: String) -> String? { + let pattern = #"(?i)(?:authorization\s*:\s*)?bearer\s+([A-Za-z0-9._~+/=-]+)"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { return nil } + let range = NSRange(raw.startIndex..= 2, + let tokenRange = Range(match.range(at: 1), in: raw) + else { + return nil + } + let token = raw[tokenRange].trimmingCharacters(in: .whitespacesAndNewlines) + return token.isEmpty ? nil : String(token) + } + + private static func bareBearerToken(from raw: String) -> String? { + let token = raw.trimmingCharacters(in: .whitespacesAndNewlines) + guard !token.contains("="), + !token.contains(";"), + !token.contains(" "), + !token.contains("\n"), + token.count >= 40 || token.split(separator: ".").count >= 3 + else { + return nil + } + return token + } +} diff --git a/Sources/CodexBarCore/Providers/Factory/FactoryProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Factory/FactoryProviderDescriptor.swift new file mode 100644 index 0000000..f3f3892 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Factory/FactoryProviderDescriptor.swift @@ -0,0 +1,145 @@ +import Foundation + +public enum FactoryProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .factory, + metadata: ProviderMetadata( + id: .factory, + displayName: "Droid", + sessionLabel: "Standard", + weeklyLabel: "Premium", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Droid usage", + cliName: "factory", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder, + dashboardURL: "https://app.factory.ai/settings/billing", + statusPageURL: "https://status.factory.ai", + statusLinkURL: nil), + branding: ProviderBranding( + iconStyle: .factory, + iconResourceName: "ProviderIcon-factory", + color: ProviderColor(red: 255 / 255, green: 107 / 255, blue: 53 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Droid cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + // `.cli` remains as an Auto compatibility alias for persisted configs from older builds + // that advertised `[.auto, .cli]` while only implementing the web strategy. + sourceModes: [.auto, .api, .web, .cli], + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), + cli: ProviderCLIConfig( + name: "factory", + versionDetector: nil)) + } + + private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] { + switch context.sourceMode { + case .api: + [FactoryAPIFetchStrategy()] + case .web: + [FactoryStatusFetchStrategy()] + case .auto, .cli: + // Legacy `source: cli` behaves as Auto (API key first, then cookies/WorkOS on macOS). + // Recoverable API failures fall through to web; explicit `.api` does not. + [FactoryAPIFetchStrategy(), FactoryStatusFetchStrategy()] + case .oauth: + [] + } + } +} + +struct FactoryAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "factory.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + // Explicit API mode always runs so a missing key surfaces as FactoryStatusProbeError.missingAPIKey. + // Auto mode only tries API when a key is resolvable, then falls back to cookies/WorkOS + // for missing/unauthorized keys and other recoverable API failures. + context.sourceMode == .api || Self.resolveAPIKey(environment: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard let apiKey = Self.resolveAPIKey(environment: context.env) else { + throw FactoryStatusProbeError.missingAPIKey + } + + let probe = FactoryStatusProbe(browserDetection: context.browserDetection) + do { + let snap = try await probe.fetch(apiKey: apiKey) + return self.makeResult( + usage: snap.toUsageSnapshot(), + sourceLabel: "api") + } catch let error as FactoryStatusProbeError { + throw Self.mapAPIError(error) + } + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + // Explicit API mode stays strict. Auto/cli keep cookies/WorkOS as a recoverable fallback so + // timeouts, DNS/5xx failures, and parse changes do not strand existing web setups. + guard context.sourceMode == .auto || context.sourceMode == .cli else { return false } + if error is CancellationError || (error as? URLError)?.code == .cancelled { + return false + } + return true + } + + private static func resolveAPIKey(environment: [String: String]) -> String? { + FactorySettingsReader.apiKey(environment: environment) + } + + private static func mapAPIError(_ error: FactoryStatusProbeError) -> FactoryStatusProbeError { + switch error { + case .notLoggedIn: + .unauthorizedAPIKey + case let .networkError(message) + where message.contains("HTTP 401") || message.contains("HTTP 403"): + .unauthorizedAPIKey + default: + error + } + } +} + +struct FactoryStatusFetchStrategy: ProviderFetchStrategy { + let id: String = "factory.web" + let kind: ProviderFetchKind = .web + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + #if !os(macOS) + // Cookie/WorkOS import is macOS-only; Linux relies on API-key auth instead. + return false + #else + guard context.settings?.factory?.cookieSource != .off else { return false } + return true + #endif + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let probe = FactoryStatusProbe(browserDetection: context.browserDetection) + let manual = Self.manualCookieHeader(from: context) + let snap = try await probe.fetch(cookieHeaderOverride: manual) + return self.makeResult( + usage: snap.toUsageSnapshot(), + sourceLabel: "web") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } + + private static func manualCookieHeader(from context: ProviderFetchContext) -> String? { + guard context.settings?.factory?.cookieSource == .manual else { return nil } + return context.settings?.factory?.manualCookieHeader?.trimmingCharacters(in: .whitespacesAndNewlines) + } +} diff --git a/Sources/CodexBarCore/Providers/Factory/FactorySessionStore.swift b/Sources/CodexBarCore/Providers/Factory/FactorySessionStore.swift new file mode 100644 index 0000000..1d569cb --- /dev/null +++ b/Sources/CodexBarCore/Providers/Factory/FactorySessionStore.swift @@ -0,0 +1,179 @@ +import Foundation + +#if os(macOS) + +// MARK: - Factory Session Store + +public actor FactorySessionStore { + public static let shared = FactorySessionStore() + + private var sessionCookies: [HTTPCookie] = [] + private var bearerToken: String? + private var refreshToken: String? + private var fileURL: URL + private var didLoadFromDisk = false + + private init() { + let fm = FileManager.default + let appSupport = fm.urls(for: .applicationSupportDirectory, in: .userDomainMask).first + ?? fm.temporaryDirectory + let dir = appSupport.appendingPathComponent("CodexBar", isDirectory: true) + try? fm.createDirectory(at: dir, withIntermediateDirectories: true) + self.fileURL = dir.appendingPathComponent("factory-session.json") + } + + public func setCookies(_ cookies: [HTTPCookie]) { + self.didLoadFromDisk = true + self.sessionCookies = cookies + self.saveToDisk() + } + + public func getCookies() -> [HTTPCookie] { + self.loadFromDiskIfNeeded() + return self.sessionCookies + } + + public func clearCookies() { + self.loadFromDiskIfNeeded() + self.didLoadFromDisk = true + self.sessionCookies = [] + self.saveToDisk() + } + + public func setBearerToken(_ token: String?) { + self.didLoadFromDisk = true + self.bearerToken = token + self.saveToDisk() + } + + public func getBearerToken() -> String? { + self.loadFromDiskIfNeeded() + return self.bearerToken + } + + public func setRefreshToken(_ token: String?) { + self.didLoadFromDisk = true + self.refreshToken = token + self.saveToDisk() + } + + public func getRefreshToken() -> String? { + self.loadFromDiskIfNeeded() + return self.refreshToken + } + + public func clearSession() { + self.didLoadFromDisk = true + self.sessionCookies = [] + self.bearerToken = nil + self.refreshToken = nil + try? FileManager.default.removeItem(at: self.fileURL) + } + + public func hasValidSession() -> Bool { + self.loadFromDiskIfNeeded() + return !self.sessionCookies.isEmpty || self.bearerToken != nil || self.refreshToken != nil + } + + func resetInMemoryForTesting() { + self.sessionCookies = [] + self.bearerToken = nil + self.refreshToken = nil + self.didLoadFromDisk = false + } + + func useFileURLForTesting(_ fileURL: URL) { + self.fileURL = fileURL + self.sessionCookies = [] + self.bearerToken = nil + self.refreshToken = nil + self.didLoadFromDisk = false + try? FileManager.default.removeItem(at: fileURL) + } + + private func saveToDisk() { + let cookieData = self.sessionCookies.compactMap { cookie -> [String: Any]? in + guard let props = cookie.properties else { return nil } + var serializable: [String: Any] = [:] + for (key, value) in props { + let keyString = key.rawValue + if let date = value as? Date { + serializable[keyString] = date.timeIntervalSince1970 + serializable[keyString + "_isDate"] = true + } else if let url = value as? URL { + serializable[keyString] = url.absoluteString + serializable[keyString + "_isURL"] = true + } else if JSONSerialization.isValidJSONObject([value]) || + value is String || + value is Bool || + value is NSNumber + { + serializable[keyString] = value + } + } + return serializable + } + + var payload: [String: Any] = [:] + if !cookieData.isEmpty { + payload["cookies"] = cookieData + } + if let bearerToken { + payload["bearerToken"] = bearerToken + } + if let refreshToken { + payload["refreshToken"] = refreshToken + } + + guard !payload.isEmpty, + let data = try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted]) + else { + try? FileManager.default.removeItem(at: self.fileURL) + return + } + try? data.write(to: self.fileURL) + } + + private func loadFromDisk() { + guard let data = try? Data(contentsOf: self.fileURL), + let json = try? JSONSerialization.jsonObject(with: data) + else { return } + + var cookieArray: [[String: Any]] = [] + if let dict = json as? [String: Any] { + if let stored = dict["cookies"] as? [[String: Any]] { + cookieArray = stored + } + self.bearerToken = dict["bearerToken"] as? String + self.refreshToken = dict["refreshToken"] as? String + } else if let stored = json as? [[String: Any]] { + cookieArray = stored + } + + self.sessionCookies = cookieArray.compactMap { props in + var cookieProps: [HTTPCookiePropertyKey: Any] = [:] + for (key, value) in props { + if key.hasSuffix("_isDate") || key.hasSuffix("_isURL") { continue } + + let propKey = HTTPCookiePropertyKey(key) + + if props[key + "_isDate"] as? Bool == true, let interval = value as? TimeInterval { + cookieProps[propKey] = Date(timeIntervalSince1970: interval) + } else if props[key + "_isURL"] as? Bool == true, let urlString = value as? String { + cookieProps[propKey] = URL(string: urlString) + } else { + cookieProps[propKey] = value + } + } + return HTTPCookie(properties: cookieProps) + } + } + + private func loadFromDiskIfNeeded() { + guard !self.didLoadFromDisk else { return } + self.didLoadFromDisk = true + self.loadFromDisk() + } +} + +#endif diff --git a/Sources/CodexBarCore/Providers/Factory/FactorySettingsReader.swift b/Sources/CodexBarCore/Providers/Factory/FactorySettingsReader.swift new file mode 100644 index 0000000..e200e4d --- /dev/null +++ b/Sources/CodexBarCore/Providers/Factory/FactorySettingsReader.swift @@ -0,0 +1,71 @@ +import Foundation + +public enum FactorySettingsReader { + public static let apiTokenKey = "FACTORY_API_KEY" + + /// Resolves a Factory API key from `FACTORY_API_KEY`, then optional `~/.factory/.env`. + /// Dotenv lookup uses `HOME` from `environment` when present; otherwise the process home directory + /// when `environment` is omitted. Tests can pass an empty env (no `HOME`) to skip dotenv. + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment, + homeDirectory: URL? = nil) -> String? + { + if let fromEnv = self.cleaned(environment[self.apiTokenKey]) { + return fromEnv + } + guard let home = homeDirectory ?? self.homeDirectory(from: environment) else { + return nil + } + return self.apiKeyFromFactoryDotEnv(homeDirectory: home) + } + + public static func apiKeyFromFactoryDotEnv(homeDirectory: URL) -> String? { + let envFile = homeDirectory + .appendingPathComponent(".factory", isDirectory: true) + .appendingPathComponent(".env", isDirectory: false) + guard let contents = try? String(contentsOf: envFile, encoding: .utf8) else { + return nil + } + return self.parseFactoryAPIKey(fromDotEnv: contents) + } + + static func parseFactoryAPIKey(fromDotEnv contents: String) -> String? { + for rawLine in contents.split(whereSeparator: \.isNewline) { + var line = String(rawLine).trimmingCharacters(in: .whitespacesAndNewlines) + guard !line.isEmpty, !line.hasPrefix("#") else { continue } + if line.hasPrefix("export ") { + line = String(line.dropFirst("export ".count)) + .trimmingCharacters(in: .whitespacesAndNewlines) + } + guard let separator = line.firstIndex(of: "=") else { continue } + let key = String(line[.. URL? { + guard let home = self.cleaned(environment["HOME"]) else { + return nil + } + let expanded = NSString(string: home).expandingTildeInPath + return URL(fileURLWithPath: expanded, isDirectory: true) + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } +} diff --git a/Sources/CodexBarCore/Providers/Factory/FactoryStatusProbe+APIKey.swift b/Sources/CodexBarCore/Providers/Factory/FactoryStatusProbe+APIKey.swift new file mode 100644 index 0000000..44c69bc --- /dev/null +++ b/Sources/CodexBarCore/Providers/Factory/FactoryStatusProbe+APIKey.swift @@ -0,0 +1,62 @@ +import Foundation + +extension FactoryStatusProbe { + /// Fetch Factory usage using a Factory API key (`FACTORY_API_KEY` / `fk-…`). + public func fetch( + apiKey: String, + logger: ((String) -> Void)? = nil) async throws -> FactoryStatusSnapshot + { + let trimmed = apiKey.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { + throw FactoryStatusProbeError.missingAPIKey + } + let log: (String) -> Void = { msg in logger?("[factory] \(msg)") } + log("Using Factory API key") + return try await self.fetchWithBearerToken(trimmed, logger: log) + } + + func fetchWithBearerToken( + _ bearerToken: String, + logger: (String) -> Void) async throws -> FactoryStatusSnapshot + { + let candidates = [Self.apiBaseURL, self.baseURL] + var lastError: Error? + var preferredAuthError: FactoryStatusProbeError? + for baseURL in candidates { + if baseURL != Self.apiBaseURL { + logger("Trying Factory bearer base URL: \(baseURL.host ?? baseURL.absoluteString)") + } + do { + return try await self.fetchWithCookieHeader( + "", + bearerToken: bearerToken, + baseURL: baseURL) + } catch { + if preferredAuthError == nil, + let authError = Self.preferredBearerAuthError(from: error) + { + preferredAuthError = authError + } + lastError = error + } + } + // Prefer 401/403 auth failures over later host noise (e.g. 404) so API-key Auto mode + // can map to unauthorizedAPIKey and fall back to cookies/WorkOS. + if let preferredAuthError { throw preferredAuthError } + if let lastError { throw lastError } + throw FactoryStatusProbeError.notLoggedIn + } + + private static func preferredBearerAuthError(from error: Error) -> FactoryStatusProbeError? { + guard let factoryError = error as? FactoryStatusProbeError else { return nil } + switch factoryError { + case .notLoggedIn: + return factoryError + case let .networkError(message) + where message.contains("HTTP 401") || message.contains("HTTP 403"): + return factoryError + default: + return nil + } + } +} diff --git a/Sources/CodexBarCore/Providers/Factory/FactoryStatusProbe.swift b/Sources/CodexBarCore/Providers/Factory/FactoryStatusProbe.swift new file mode 100644 index 0000000..c7ea7bd --- /dev/null +++ b/Sources/CodexBarCore/Providers/Factory/FactoryStatusProbe.swift @@ -0,0 +1,1485 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif +#if os(macOS) +import SweetCookieKit + +private let factoryCookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.factory]?.browserCookieOrder ?? Browser.defaultImportOrder + +// MARK: - Factory Cookie Importer + +/// Imports Factory session cookies from browser cookies. +public enum FactoryCookieImporter { + private static let cookieClient = BrowserCookieClient() + private static let sessionCookieNames: Set = [ + "wos-session", + "__Secure-next-auth.session-token", + "next-auth.session-token", + "__Secure-authjs.session-token", + "__Host-authjs.csrf-token", + "authjs.session-token", + "session", + "access-token", + ] + + private static let authSessionCookieNames: Set = [ + "__Secure-next-auth.session-token", + "next-auth.session-token", + "__Secure-authjs.session-token", + "authjs.session-token", + ] + private static let appBaseURL = URL(string: "https://app.factory.ai")! + private static let authBaseURL = URL(string: "https://auth.factory.ai")! + private static let apiBaseURL = URL(string: "https://api.factory.ai")! + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + } + + /// Returns all Factory sessions across supported browsers. + public static func importSessions( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + let log: (String) -> Void = { msg in logger?("[factory-cookie] \(msg)") } + var sessions: [SessionInfo] = [] + + // Filter to cookie-eligible browsers to avoid unnecessary keychain prompts + let installedBrowsers = factoryCookieImportOrder.cookieImportCandidates(using: browserDetection) + for browserSource in installedBrowsers { + do { + let perSource = try self.importSessions(from: browserSource, logger: logger) + sessions.append(contentsOf: perSource) + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)") + } + } + + guard !sessions.isEmpty else { + throw FactoryStatusProbeError.noSessionCookie + } + return sessions + } + + public static func importSessions( + from browserSource: Browser, + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + let log: (String) -> Void = { msg in logger?("[factory-cookie] \(msg)") } + let cookieDomains = ["factory.ai", "app.factory.ai", "auth.factory.ai"] + let query = BrowserCookieQuery(domains: cookieDomains) + let sources = try Self.cookieClient.codexBarRecords( + matching: query, + in: browserSource, + logger: log) + + var sessions: [SessionInfo] = [] + for source in sources where !source.records.isEmpty { + let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin) + if httpCookies.contains(where: { Self.sessionCookieNames.contains($0.name) }) { + log("Found \(httpCookies.count) Factory cookies in \(source.label)") + log("\(source.label) cookie names: \(self.cookieNames(from: httpCookies))") + if let token = httpCookies.first(where: { $0.name == "access-token" })?.value { + let hint = token.contains(".") ? "jwt" : "opaque" + log("\(source.label) access-token cookie: \(token.count) chars (\(hint))") + } + if let token = httpCookies.first(where: { self.authSessionCookieNames.contains($0.name) })?.value { + let hint = token.contains(".") ? "jwt" : "opaque" + log("\(source.label) session cookie: \(token.count) chars (\(hint))") + } + sessions.append(SessionInfo(cookies: httpCookies, sourceLabel: source.label)) + } else { + log("\(source.label) cookies found, but no Factory session cookie present") + } + } + return sessions + } + + /// Attempts to import Factory cookies using the standard browser import order. + public static func importSession( + browserDetection: BrowserDetection, + logger: ((String) -> Void)? = nil) throws -> SessionInfo + { + let sessions = try self.importSessions(browserDetection: browserDetection, logger: logger) + guard let first = sessions.first else { + throw FactoryStatusProbeError.noSessionCookie + } + return first + } + + /// Check if Factory session cookies are available + public static func hasSession(browserDetection: BrowserDetection, logger: ((String) -> Void)? = nil) -> Bool { + do { + return try !(self.importSessions(browserDetection: browserDetection, logger: logger)).isEmpty + } catch { + return false + } + } + + private static func cookieNames(from cookies: [HTTPCookie]) -> String { + let names = Set(cookies.map { "\($0.name)@\($0.domain)" }).sorted() + return names.joined(separator: ", ") + } +} + +#endif + +// MARK: - Factory API Models + +public struct FactoryAuthResponse: Codable, Sendable { + public let featureFlags: FactoryFeatureFlags? + public let organization: FactoryOrganization? + public let userProfile: FactoryUserProfile? +} + +public struct FactoryUserProfile: Codable, Sendable { + public let id: String? + public let email: String? +} + +public struct FactoryFeatureFlags: Codable, Sendable { + public let flags: [String: Bool]? + public let configs: [String: AnyCodable]? +} + +public struct FactoryOrganization: Codable, Sendable { + public let id: String? + public let name: String? + public let subscription: FactorySubscription? + + enum CodingKeys: String, CodingKey { + case id + case name + case subscription + } +} + +public struct FactorySubscription: Codable, Sendable { + public let factoryTier: String? + public let orbSubscription: FactoryOrbSubscription? +} + +public struct FactoryOrbSubscription: Codable, Sendable { + public let plan: FactoryPlan? + public let status: String? +} + +public struct FactoryPlan: Codable, Sendable { + public let name: String? + public let id: String? +} + +public struct FactoryUsageResponse: Codable, Sendable { + public let usage: FactoryUsageData? + public let source: String? + public let userId: String? +} + +public struct FactoryUsageData: Codable, Sendable { + public let startDate: Int64? + public let endDate: Int64? + public let standard: FactoryTokenUsage? + public let premium: FactoryTokenUsage? +} + +public struct FactoryTokenUsage: Codable, Sendable { + public let userTokens: Int64? + public let orgTotalTokensUsed: Int64? + public let totalAllowance: Int64? + public let usedRatio: Double? + public let orgOverageUsed: Int64? + public let basicAllowance: Int64? + public let orgOverageLimit: Int64? +} + +public struct FactoryBillingLimitsResponse: Codable, Sendable { + public let usesTokenRateLimitsBilling: Bool + public let limits: FactoryTokenRateLimits? + public let extraUsageBalanceCents: Int + public let overagePreference: String? + public let extraUsageAllowed: Bool + public let tokenRateLimitsRolloutEligible: Bool + + enum CodingKeys: String, CodingKey { + case usesTokenRateLimitsBilling + case limits + case extraUsageBalanceCents + case overagePreference + case extraUsageAllowed + case tokenRateLimitsRolloutEligible + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.usesTokenRateLimitsBilling = try container + .decodeIfPresent(Bool.self, forKey: .usesTokenRateLimitsBilling) ?? false + self.limits = try container.decodeIfPresent(FactoryTokenRateLimits.self, forKey: .limits) + self.extraUsageBalanceCents = try container.decodeIfPresent(Int.self, forKey: .extraUsageBalanceCents) ?? 0 + self.overagePreference = try container.decodeIfPresent(String.self, forKey: .overagePreference) + self.extraUsageAllowed = try container.decodeIfPresent(Bool.self, forKey: .extraUsageAllowed) ?? false + self.tokenRateLimitsRolloutEligible = try container + .decodeIfPresent(Bool.self, forKey: .tokenRateLimitsRolloutEligible) ?? false + } +} + +public struct FactoryTokenRateLimits: Codable, Sendable { + public let standard: FactoryLimitPool + public let core: FactoryLimitPool? +} + +public struct FactoryLimitPool: Codable, Sendable { + public let fiveHour: FactoryBillingWindow + public let weekly: FactoryBillingWindow + public let monthly: FactoryBillingWindow + + public var hasUsageData: Bool { + [self.fiveHour, self.weekly, self.monthly].contains { + $0.usedPercent > 0 || $0.windowEnd != nil || $0.secondsRemaining != nil + } + } +} + +public struct FactoryBillingWindow: Codable, Sendable { + public let usedPercent: Double + public let windowEnd: FlexibleFactoryDate? + public let secondsRemaining: Double? + + public func resetAt(now: Date) -> Date? { + if let secondsRemaining, secondsRemaining > 0 { + return now.addingTimeInterval(secondsRemaining) + } + guard let windowEnd = self.windowEnd?.date, windowEnd > now else { + return nil + } + return windowEnd + } + + public func effectiveUsedPercent(now: Date) -> Double { + // Factory can leave stale values after short rolling windows expire. The web UI treats + // that state as reset, so mirror it here instead of showing expired usage. + if self.resetAt(now: now) == nil, self.windowEnd != nil, self.secondsRemaining == nil { + return 0 + } + return min(100, max(0, self.usedPercent)) + } + + public func rateWindow(windowMinutes: Int?, title: String, now: Date) -> RateWindow { + let reset = self.resetAt(now: now) + return RateWindow( + usedPercent: self.effectiveUsedPercent(now: now), + windowMinutes: windowMinutes, + resetsAt: reset, + resetDescription: reset.map { FactoryStatusSnapshot.formatResetDate($0) }) + } +} + +public struct FlexibleFactoryDate: Codable, Sendable { + public let date: Date + + public init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if let seconds = try? container.decode(Double.self) { + self.date = Date(timeIntervalSince1970: seconds > 1e12 ? seconds / 1000.0 : seconds) + return + } + let string = try container.decode(String.self) + if let numeric = Double(string) { + self.date = Date(timeIntervalSince1970: numeric > 1e12 ? numeric / 1000.0 : numeric) + return + } + + let fractional = ISO8601DateFormatter() + fractional.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let parsed = fractional.date(from: string) ?? ISO8601DateFormatter().date(from: string) { + self.date = parsed + return + } + throw DecodingError.dataCorruptedError(in: container, debugDescription: "Invalid Factory date") + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + try container.encode(self.date) + } +} + +/// Helper for encoding arbitrary JSON +public struct AnyCodable: Codable, Sendable { + public init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if container.decodeNil() { + return + } + _ = try? container.decode([String: AnyCodable].self) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + try container.encodeNil() + } +} + +// MARK: - Factory Status Snapshot + +public struct FactoryStatusSnapshot: Sendable { + /// Standard token usage (user) + public let standardUserTokens: Int64 + /// Standard token usage (org total) + public let standardOrgTokens: Int64 + /// Standard token allowance + public let standardAllowance: Int64 + /// Standard usage ratio from API (0.0-1.0), preferred over manual calculation + /// Falls back to percent-scale (0.0-100.0) when allowance is unavailable. + public let standardUsedRatio: Double? + /// Premium token usage (user) + public let premiumUserTokens: Int64 + /// Premium token usage (org total) + public let premiumOrgTokens: Int64 + /// Premium token allowance + public let premiumAllowance: Int64 + /// Premium usage ratio from API (0.0-1.0), preferred over manual calculation + /// Falls back to percent-scale (0.0-100.0) when allowance is unavailable. + public let premiumUsedRatio: Double? + /// Billing period start + public let periodStart: Date? + /// Billing period end + public let periodEnd: Date? + /// Plan name + public let planName: String? + /// Factory tier (enterprise, team, etc.) + public let tier: String? + /// Organization name + public let organizationName: String? + /// User email + public let accountEmail: String? + /// User ID + public let userId: String? + /// Raw JSON for debugging + public let rawJSON: String? + /// New Factory token-rate-limits billing payload, when enabled for the account. + public let tokenRateLimits: FactoryTokenRateLimits? + public let extraUsageBalanceCents: Int? + public let overagePreference: String? + + public init( + standardUserTokens: Int64, + standardOrgTokens: Int64, + standardAllowance: Int64, + standardUsedRatio: Double? = nil, + premiumUserTokens: Int64, + premiumOrgTokens: Int64, + premiumAllowance: Int64, + premiumUsedRatio: Double? = nil, + periodStart: Date?, + periodEnd: Date?, + planName: String?, + tier: String?, + organizationName: String?, + accountEmail: String?, + userId: String?, + rawJSON: String?, + tokenRateLimits: FactoryTokenRateLimits? = nil, + extraUsageBalanceCents: Int? = nil, + overagePreference: String? = nil) + { + self.standardUserTokens = standardUserTokens + self.standardOrgTokens = standardOrgTokens + self.standardAllowance = standardAllowance + self.standardUsedRatio = standardUsedRatio + self.premiumUserTokens = premiumUserTokens + self.premiumOrgTokens = premiumOrgTokens + self.premiumAllowance = premiumAllowance + self.premiumUsedRatio = premiumUsedRatio + self.periodStart = periodStart + self.periodEnd = periodEnd + self.planName = planName + self.tier = tier + self.organizationName = organizationName + self.accountEmail = accountEmail + self.userId = userId + self.rawJSON = rawJSON + self.tokenRateLimits = tokenRateLimits + self.extraUsageBalanceCents = extraUsageBalanceCents + self.overagePreference = overagePreference + } + + /// Convert to UsageSnapshot for the common provider interface + public func toUsageSnapshot() -> UsageSnapshot { + if let tokenRateLimits { + return self.tokenRateLimitsUsageSnapshot(from: tokenRateLimits) + } + + // Primary: Standard tokens used (as percentage of allowance, capped reasonably) + let standardPercent = self.calculateUsagePercent( + used: self.standardUserTokens, + allowance: self.standardAllowance, + apiRatio: self.standardUsedRatio) + + let primary = RateWindow( + usedPercent: standardPercent, + windowMinutes: nil, + resetsAt: self.periodEnd, + resetDescription: self.periodEnd.map { Self.formatResetDate($0) }) + + // Secondary: Premium tokens used + let premiumPercent = self.calculateUsagePercent( + used: self.premiumUserTokens, + allowance: self.premiumAllowance, + apiRatio: self.premiumUsedRatio) + + let secondary = RateWindow( + usedPercent: premiumPercent, + windowMinutes: nil, + resetsAt: self.periodEnd, + resetDescription: self.periodEnd.map { Self.formatResetDate($0) }) + + // Format login method as tier + plan + let loginMethod: String? = { + var parts: [String] = [] + if let tier = self.tier, !tier.isEmpty { + parts.append("Factory \(tier.capitalized)") + } + if let plan = self.planName, !plan.isEmpty, !plan.lowercased().contains("factory") { + parts.append(plan) + } + return parts.isEmpty ? nil : parts.joined(separator: " - ") + }() + + let identity = ProviderIdentitySnapshot( + providerID: .factory, + accountEmail: self.accountEmail, + accountOrganization: self.organizationName, + loginMethod: loginMethod) + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: nil, + providerCost: nil, + updatedAt: Date(), + identity: identity) + } + + private func tokenRateLimitsUsageSnapshot(from limits: FactoryTokenRateLimits) -> UsageSnapshot { + let now = Date() + let primary = limits.standard.fiveHour.rateWindow(windowMinutes: 5 * 60, title: "5h", now: now) + let secondary = limits.standard.weekly.rateWindow(windowMinutes: 7 * 24 * 60, title: "7-day", now: now) + let tertiary = limits.standard.monthly.rateWindow(windowMinutes: nil, title: "Monthly", now: now) + + let coreWindows: [NamedRateWindow]? = if let core = limits.core, core.hasUsageData { + [ + NamedRateWindow( + id: "factory-core-5h", + title: "Core 5h", + window: core.fiveHour.rateWindow(windowMinutes: 5 * 60, title: "Core 5h", now: now)), + NamedRateWindow( + id: "factory-core-7d", + title: "Core 7-day", + window: core.weekly.rateWindow(windowMinutes: 7 * 24 * 60, title: "Core 7-day", now: now)), + NamedRateWindow( + id: "factory-core-monthly", + title: "Core Monthly", + window: core.monthly.rateWindow(windowMinutes: nil, title: "Core Monthly", now: now)), + ] + } else { + nil + } + + let loginMethod: String? = { + var parts: [String] = [] + if let tier = self.tier, !tier.isEmpty { + parts.append("Factory \(tier.capitalized)") + } + if let plan = self.planName, !plan.isEmpty, !plan.lowercased().contains("factory") { + parts.append(plan) + } + if let overagePreference, !overagePreference.isEmpty { + parts.append("Fallback: \(overagePreference)") + } + return parts.isEmpty ? nil : parts.joined(separator: " - ") + }() + + let identity = ProviderIdentitySnapshot( + providerID: .factory, + accountEmail: self.accountEmail, + accountOrganization: self.organizationName, + loginMethod: loginMethod) + let providerCost = self.extraUsageBalanceCents.map { + ProviderCostSnapshot( + used: Double($0) / 100.0, + limit: 0, + currencyCode: "USD", + period: "Extra usage balance", + updatedAt: now) + } + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: tertiary, + extraRateWindows: coreWindows, + providerCost: providerCost, + updatedAt: now, + identity: identity) + } + + private func calculateUsagePercent(used: Int64, allowance: Int64, apiRatio: Double?) -> Double { + // Prefer API-provided ratio when available and valid. + // This handles plan-specific limits correctly on the server side, + // avoiding issues with missing/sentinel values in totalAllowance. + let unlimitedThreshold: Int64 = 1_000_000_000_000 + if let ratio = apiRatio, + !(ratio == 0 && used > 0 && allowance > 0 && allowance <= unlimitedThreshold), + let percent = Self.percentFromAPIRatio(ratio, allowance: allowance, unlimitedThreshold: unlimitedThreshold) + { + return percent + } + + // Fallback: calculate from used/allowance. + // Treat very large allowances (> 1 trillion) as unlimited. + if allowance > unlimitedThreshold { + // For unlimited, show a token count-based pseudo-percentage (capped at 100%). + // Use 100M tokens as a reference point for "100%". + let referenceTokens: Double = 100_000_000 + return min(100, Double(used) / referenceTokens * 100) + } + guard allowance > 0 else { return 0 } + return min(100, Double(used) / Double(allowance) * 100) + } + + private static func percentFromAPIRatio( + _ ratio: Double, + allowance: Int64, + unlimitedThreshold: Int64) -> Double? + { + guard ratio.isFinite else { return nil } + + // Primary: ratio scale (0.0 - 1.0). Clamp to account for rounding. + if ratio >= -0.001, ratio <= 1.001 { + return min(100, max(0, ratio * 100)) + } + + // TODO: Confirm usedRatio contract (0.0-1.0 vs 0.0-100.0) and tighten this fallback. + // Secondary: percent scale (0.0 - 100.0), only when allowance is missing/unreliable. + // This avoids misinterpreting slightly-over-1 ratios when we can calculate locally. + let allowanceIsReliable = allowance > 0 && allowance <= unlimitedThreshold + if !allowanceIsReliable, ratio >= -0.1, ratio <= 100.1 { + return min(100, max(0, ratio)) + } + + return nil + } + + static func formatResetDate(_ date: Date) -> String { + let formatter = DateFormatter() + formatter.dateFormat = "MMM d 'at' h:mma" + formatter.locale = Locale(identifier: "en_US_POSIX") + return "Resets " + formatter.string(from: date) + } +} + +// MARK: - Factory Status Probe + +public struct FactoryStatusProbe: Sendable { + public let baseURL: URL + public var timeout: TimeInterval = 15.0 + private static let staleTokenCookieNames: Set = [ + "access-token", + "__recent_auth", + ] + private static let sessionCookieNames: Set = [ + "session", + "wos-session", + ] + private static let authSessionCookieNames: Set = [ + "__Secure-next-auth.session-token", + "next-auth.session-token", + "__Secure-authjs.session-token", + "authjs.session-token", + ] + static let appBaseURL = URL(string: "https://app.factory.ai")! + static let authBaseURL = URL(string: "https://auth.factory.ai")! + static let apiBaseURL = URL(string: "https://api.factory.ai")! + private static let workosClientIDs = [ + "client_01HXRMBQ9BJ3E7QSTQ9X2PHVB7", + "client_01HNM792M5G5G1A2THWPXKFMXB", + ] + + private struct WorkOSAuthResponse: Decodable { + let access_token: String + let refresh_token: String? + let organization_id: String? + } + + private let browserDetection: BrowserDetection + private let transport: any ProviderHTTPTransport + + public init( + baseURL: URL = URL(string: "https://app.factory.ai")!, + timeout: TimeInterval = 15.0, + browserDetection: BrowserDetection, + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) + { + self.baseURL = baseURL + self.timeout = timeout + self.browserDetection = browserDetection + self.transport = transport + } + + /// Fetch Factory usage using browser cookies with fallback to stored session. + public func fetch( + cookieHeaderOverride: String? = nil, + logger: ((String) -> Void)? = nil) async throws -> FactoryStatusSnapshot + { + #if os(macOS) + let log: (String) -> Void = { msg in logger?("[factory] \(msg)") } + var lastError: Error? + + let manualOverride = cookieHeaderOverride?.trimmingCharacters(in: .whitespacesAndNewlines) + if manualOverride?.isEmpty == false { + guard let override = Self.manualCredentials(from: manualOverride) else { + throw FactoryStatusProbeError.noSessionCookie + } + if let cookieHeader = override.cookieHeader { + log("Using manual cookie header") + let candidates = [ + self.baseURL, + Self.authBaseURL, + Self.apiBaseURL, + ] + for baseURL in candidates { + do { + return try await self.fetchWithCookieHeader( + cookieHeader, + bearerToken: override.bearerToken, + baseURL: baseURL) + } catch { + lastError = error + } + } + } + if let bearerToken = override.bearerToken { + log("Using manual Factory bearer token") + return try await self.fetchWithBearerToken(bearerToken, logger: log) + } + if let lastError { throw lastError } + throw FactoryStatusProbeError.noSessionCookie + } + + if let cached = CookieHeaderCache.load(provider: .factory), + !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + { + log("Using cached cookie header from \(cached.sourceLabel)") + let bearer = Self.bearerToken(fromHeader: cached.cookieHeader) + do { + return try await self.fetchWithCookieHeader( + cached.cookieHeader, + bearerToken: bearer, + baseURL: self.baseURL) + } catch { + if case FactoryStatusProbeError.notLoggedIn = error { + CookieHeaderCache.clear(provider: .factory) + } + lastError = error + } + } + + // Filter to only installed browsers to avoid unnecessary keychain prompts + let installedChromiumAndFirefox = [.chrome, .firefox].cookieImportCandidates(using: self.browserDetection) + + let attempts: [() async -> FetchAttemptResult] = [ + { await self.attemptStoredCookies(logger: log) }, + { await self.attemptStoredBearer(logger: log) }, + { await self.attemptStoredRefreshToken(logger: log) }, + { await self.attemptLocalStorageTokens(logger: log) }, + { await self.attemptBrowserCookies(logger: log, sources: [.safari]) }, + { await self.attemptWorkOSCookies(logger: log, sources: [.safari]) }, + { await self.attemptBrowserCookies(logger: log, sources: installedChromiumAndFirefox) }, + { await self.attemptWorkOSCookies(logger: log, sources: installedChromiumAndFirefox) }, + ] + + for attempt in attempts { + switch await attempt() { + case let .success(snapshot): + return snapshot + case let .failure(error): + lastError = error + case .skipped: + continue + } + } + + if let lastError { throw lastError } + throw FactoryStatusProbeError.noSessionCookie + #else + _ = cookieHeaderOverride + _ = logger + throw FactoryStatusProbeError.notSupported + #endif + } + + #if os(macOS) + private enum FetchAttemptResult { + case success(FactoryStatusSnapshot) + case failure(Error) + case skipped + } + + private func attemptBrowserCookies( + logger: @escaping (String) -> Void, + sources: [Browser]) async -> FetchAttemptResult + { + do { + var lastError: Error? + for browserSource in sources { + let sessions = try FactoryCookieImporter.importSessions(from: browserSource, logger: logger) + for session in sessions { + logger("Using cookies from \(session.sourceLabel)") + do { + let snapshot = try await self.fetchWithCookies(session.cookies, logger: logger) + await FactorySessionStore.shared.setCookies(session.cookies) + CookieHeaderCache.store( + provider: .factory, + cookieHeader: session.cookieHeader, + sourceLabel: session.sourceLabel) + return .success(snapshot) + } catch { + lastError = error + logger("Browser session fetch failed for \(session.sourceLabel): \(error.localizedDescription)") + } + } + } + if let lastError { return .failure(lastError) } + return .skipped + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + logger("Browser cookie import failed: \(error.localizedDescription)") + return .failure(error) + } + } + + private func attemptStoredCookies(logger: (String) -> Void) async -> FetchAttemptResult { + let storedCookies = await FactorySessionStore.shared.getCookies() + guard !storedCookies.isEmpty else { return .skipped } + logger("Using stored session cookies") + do { + return try await .success(self.fetchWithCookies(storedCookies, logger: logger)) + } catch { + if case FactoryStatusProbeError.notLoggedIn = error { + await FactorySessionStore.shared.clearCookies() + logger("Stored session cookies invalid, cleared") + } else { + logger("Stored session failed: \(error.localizedDescription)") + } + return .failure(error) + } + } + + private func attemptStoredBearer(logger: (String) -> Void) async -> FetchAttemptResult { + guard let bearerToken = await FactorySessionStore.shared.getBearerToken() else { return .skipped } + logger("Using stored Factory bearer token") + do { + return try await .success(self.fetchWithBearerToken(bearerToken, logger: logger)) + } catch { + return .failure(error) + } + } + + private func attemptStoredRefreshToken(logger: (String) -> Void) async -> FetchAttemptResult { + guard let refreshToken = await FactorySessionStore.shared.getRefreshToken(), + !refreshToken.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + else { + return .skipped + } + logger("Using stored WorkOS refresh token") + do { + return try await .success(self.fetchWithWorkOSRefreshToken( + refreshToken, + organizationID: nil, + logger: logger)) + } catch { + if self.isInvalidGrant(error) { + await FactorySessionStore.shared.setRefreshToken(nil) + } else if case FactoryStatusProbeError.noSessionCookie = error { + await FactorySessionStore.shared.setRefreshToken(nil) + } + return .failure(error) + } + } + + private func attemptLocalStorageTokens(logger: @escaping (String) -> Void) async -> FetchAttemptResult { + let workosTokens = FactoryLocalStorageImporter.importWorkOSTokens( + browserDetection: self.browserDetection, + logger: logger) + guard !workosTokens.isEmpty else { return .skipped } + var lastError: Error? + for token in workosTokens { + guard !token.refreshToken.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + continue + } + logger("Using WorkOS refresh token from \(token.sourceLabel)") + if let accessToken = token.accessToken { + do { + await FactorySessionStore.shared.setBearerToken(accessToken) + return try await .success(self.fetchWithBearerToken(accessToken, logger: logger)) + } catch { + lastError = error + } + } + do { + return try await .success(self.fetchWithWorkOSRefreshToken( + token.refreshToken, + organizationID: token.organizationID, + logger: logger)) + } catch { + if self.isInvalidGrant(error) { + await FactorySessionStore.shared.setRefreshToken(nil) + } + lastError = error + } + } + if let lastError { return .failure(lastError) } + return .skipped + } + + private func attemptWorkOSCookies( + logger: @escaping (String) -> Void, + sources: [Browser]) async -> FetchAttemptResult + { + let log: (String) -> Void = { msg in logger("[factory-workos] \(msg)") } + var lastError: Error? + + for browserSource in sources { + do { + let query = BrowserCookieQuery(domains: ["workos.com"]) + let sources = try BrowserCookieClient().codexBarRecords( + matching: query, + in: browserSource, + logger: log) + for source in sources where !source.records.isEmpty { + let cookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin) + guard !cookies.isEmpty else { continue } + log("Using WorkOS cookies from \(source.label)") + do { + let auth = try await self.fetchWorkOSAccessTokenWithCookies( + cookies: cookies, + logger: logger) + await FactorySessionStore.shared.setBearerToken(auth.access_token) + if let refreshToken = auth.refresh_token { + await FactorySessionStore.shared.setRefreshToken(refreshToken) + } + return try await .success(self.fetchWithBearerToken(auth.access_token, logger: logger)) + } catch { + lastError = error + log("WorkOS cookie auth failed for \(source.label): \(error.localizedDescription)") + } + } + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + log("\(browserSource.displayName) WorkOS cookie import failed: \(error.localizedDescription)") + lastError = error + } + } + + if let lastError { return .failure(lastError) } + return .skipped + } + + private func fetchWithWorkOSRefreshToken( + _ refreshToken: String, + organizationID: String?, + logger: (String) -> Void) async throws -> FactoryStatusSnapshot + { + let auth = try await self.fetchWorkOSAccessToken( + refreshToken: refreshToken, + organizationID: organizationID) + await FactorySessionStore.shared.setBearerToken(auth.access_token) + if let newRefresh = auth.refresh_token { + await FactorySessionStore.shared.setRefreshToken(newRefresh) + } + return try await self.fetchWithBearerToken(auth.access_token, logger: logger) + } + + private func fetchWithCookies( + _ cookies: [HTTPCookie], + logger: (String) -> Void) async throws -> FactoryStatusSnapshot + { + let candidates = Self.baseURLCandidates(default: self.baseURL, cookies: cookies) + var lastError: Error? + + for baseURL in candidates { + if baseURL != self.baseURL { + logger("Trying Factory base URL: \(baseURL.host ?? baseURL.absoluteString)") + } + do { + return try await self.fetchWithCookies(cookies, baseURL: baseURL, logger: logger) + } catch { + lastError = error + } + } + + if let lastError { throw lastError } + throw FactoryStatusProbeError.noSessionCookie + } + + private func fetchWithCookies( + _ cookies: [HTTPCookie], + baseURL: URL, + logger: (String) -> Void) async throws -> FactoryStatusSnapshot + { + let header = Self.cookieHeader(from: cookies) + let bearerToken = Self.bearerToken(from: cookies) + do { + return try await self.fetchWithCookieHeader(header, bearerToken: bearerToken, baseURL: baseURL) + } catch let error as FactoryStatusProbeError { + if case .notLoggedIn = error, bearerToken != nil { + logger("Retrying without Authorization header") + return try await self.fetchWithCookieHeader(header, bearerToken: nil, baseURL: baseURL) + } + guard case let .networkError(message) = error, + message.contains("HTTP 409") + else { + throw error + } + + var lastError: Error? = error + if bearerToken != nil { + logger("Retrying without Authorization header (HTTP 409)") + do { + return try await self.fetchWithCookieHeader(header, bearerToken: nil, baseURL: baseURL) + } catch { + lastError = error + } + } + + let retries: [(String, (HTTPCookie) -> Bool)] = [ + ("Retrying without access-token cookies", { !Self.staleTokenCookieNames.contains($0.name) }), + ("Retrying without session cookies", { !Self.sessionCookieNames.contains($0.name) }), + ("Retrying without access-token/session cookies", { + !Self.staleTokenCookieNames.contains($0.name) && !Self.sessionCookieNames.contains($0.name) + }), + ] + + for (label, predicate) in retries { + let filtered = cookies.filter(predicate) + guard filtered.count < cookies.count else { continue } + logger(label) + do { + let filteredBearer = Self.bearerToken(from: filtered) + return try await self.fetchWithCookieHeader( + Self.cookieHeader(from: filtered), + bearerToken: filteredBearer, + baseURL: baseURL) + } catch let retryError as FactoryStatusProbeError { + switch retryError { + case let .networkError(retryMessage) + where retryMessage.contains("HTTP 409") && + retryMessage.localizedCaseInsensitiveContains("stale token"): + lastError = retryError + continue + case .notLoggedIn: + lastError = retryError + continue + default: + throw retryError + } + } + } + + let authOnly = cookies.filter { + Self.authSessionCookieNames.contains($0.name) || $0.name == "__Host-authjs.csrf-token" + } + if !authOnly.isEmpty, authOnly.count < cookies.count { + logger("Retrying with auth session cookies only") + do { + return try await self.fetchWithCookieHeader( + Self.cookieHeader(from: authOnly), + bearerToken: Self.bearerToken(from: authOnly), + baseURL: baseURL) + } catch let retryError as FactoryStatusProbeError { + lastError = retryError + } + } + + if let lastError { throw lastError } + throw error + } catch { + throw error + } + } + + private static func cookieHeader(from cookies: [HTTPCookie]) -> String { + cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + + #endif + + func fetchWithCookieHeader( + _ cookieHeader: String, + bearerToken: String?, + baseURL: URL) async throws -> FactoryStatusSnapshot + { + // First fetch auth info to get user ID and org info + let authInfo = try await self.fetchAuthInfo( + cookieHeader: cookieHeader, + bearerToken: bearerToken, + baseURL: baseURL) + + let userId = factoryUserIdFromAuth(authInfo) + ?? factoryUserIdFromBearerToken(bearerToken) + + if let billingLimits = try await self.fetchBillingLimitsIfAvailable( + cookieHeader: cookieHeader, + bearerToken: bearerToken), + billingLimits.usesTokenRateLimitsBilling, + let tokenRateLimits = billingLimits.limits + { + return self.buildTokenRateLimitsSnapshot( + authInfo: authInfo, + billingLimits: billingLimits, + tokenRateLimits: tokenRateLimits, + userId: userId) + } + + // Fetch usage data + let usageData = try await self.fetchUsage( + cookieHeader: cookieHeader, + bearerToken: bearerToken, + userId: userId, + baseURL: baseURL) + + return self.buildSnapshot(authInfo: authInfo, usageData: usageData, userId: userId) + } + + private func fetchBillingLimitsIfAvailable( + cookieHeader: String, + bearerToken: String?) async throws -> FactoryBillingLimitsResponse? + { + let url = Self.apiBaseURL.appendingPathComponent("/api/billing/limits") + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.httpMethod = "GET" + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("https://app.factory.ai", forHTTPHeaderField: "Origin") + request.setValue("https://app.factory.ai/", forHTTPHeaderField: "Referer") + request.setValue("web-app", forHTTPHeaderField: "x-factory-client") + if !cookieHeader.isEmpty { + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + } + if let bearerToken { + request.setValue("Bearer \(bearerToken)", forHTTPHeaderField: "Authorization") + } + + let data: Data + let response: URLResponse + do { + (data, response) = try await self.transport.data(for: request) + } catch { + return nil + } + + guard let httpResponse = response as? HTTPURLResponse else { + return nil + } + + guard httpResponse.statusCode == 200 else { + return nil + } + + return try? JSONDecoder().decode(FactoryBillingLimitsResponse.self, from: data) + } + + private func fetchAuthInfo( + cookieHeader: String, + bearerToken: String?, + baseURL: URL) async throws -> FactoryAuthResponse + { + let url = baseURL.appendingPathComponent("/api/app/auth/me") + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.httpMethod = "GET" + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("https://app.factory.ai", forHTTPHeaderField: "Origin") + request.setValue("https://app.factory.ai/", forHTTPHeaderField: "Referer") + if !cookieHeader.isEmpty { + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + } + request.setValue("web-app", forHTTPHeaderField: "x-factory-client") + if let bearerToken { + request.setValue("Bearer \(bearerToken)", forHTTPHeaderField: "Authorization") + } + + let (data, response) = try await self.transport.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw FactoryStatusProbeError.networkError("Invalid response") + } + + if httpResponse.statusCode == 401 { + throw FactoryStatusProbeError.notLoggedIn + } + + if httpResponse.statusCode == 403 { + let body = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + let snippet = body.isEmpty ? "" : ": \(body.prefix(200))" + throw FactoryStatusProbeError.networkError("HTTP 403 Forbidden\(snippet)") + } + + guard httpResponse.statusCode == 200 else { + let body = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + let snippet = body.isEmpty ? "" : ": \(body.prefix(200))" + throw FactoryStatusProbeError.networkError("HTTP \(httpResponse.statusCode)\(snippet)") + } + + do { + let decoder = JSONDecoder() + return try decoder.decode(FactoryAuthResponse.self, from: data) + } catch { + let rawJSON = String(data: data, encoding: .utf8) ?? "" + throw FactoryStatusProbeError + .parseFailed("Auth decode failed: \(error.localizedDescription). Raw: \(rawJSON.prefix(200))") + } + } + + private func fetchUsage( + cookieHeader: String, + bearerToken: String?, + userId: String?, + baseURL: URL) async throws -> FactoryUsageResponse + { + var components = URLComponents( + url: baseURL.appendingPathComponent("/api/organization/subscription/usage"), + resolvingAgainstBaseURL: false) + components?.queryItems = [ + URLQueryItem(name: "useCache", value: "true"), + ] + if let userId { + components?.queryItems?.append(URLQueryItem(name: "userId", value: userId)) + } + let url = components?.url ?? baseURL.appendingPathComponent("/api/organization/subscription/usage") + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.httpMethod = "GET" + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("https://app.factory.ai", forHTTPHeaderField: "Origin") + request.setValue("https://app.factory.ai/", forHTTPHeaderField: "Referer") + if !cookieHeader.isEmpty { + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + } + request.setValue("web-app", forHTTPHeaderField: "x-factory-client") + if let bearerToken { + request.setValue("Bearer \(bearerToken)", forHTTPHeaderField: "Authorization") + } + + let (data, response) = try await self.transport.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw FactoryStatusProbeError.networkError("Invalid response") + } + + if httpResponse.statusCode == 401 || httpResponse.statusCode == 403 { + throw FactoryStatusProbeError.notLoggedIn + } + + guard httpResponse.statusCode == 200 else { + let body = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + let snippet = body.isEmpty ? "" : ": \(body.prefix(200))" + throw FactoryStatusProbeError.networkError("HTTP \(httpResponse.statusCode)\(snippet)") + } + + do { + let decoder = JSONDecoder() + return try decoder.decode(FactoryUsageResponse.self, from: data) + } catch { + let rawJSON = String(data: data, encoding: .utf8) ?? "" + throw FactoryStatusProbeError + .parseFailed("Usage decode failed: \(error.localizedDescription). Raw: \(rawJSON.prefix(200))") + } + } + + #if os(macOS) + private static func baseURLCandidates(default baseURL: URL, cookies: [HTTPCookie]) -> [URL] { + let cookieDomains = Set( + cookies.map { + $0.domain.trimmingCharacters(in: CharacterSet(charactersIn: ".")) + }) + + var candidates: [URL] = [] + if cookieDomains.contains("auth.factory.ai") { + candidates.append(Self.authBaseURL) + } + candidates.append(Self.apiBaseURL) + candidates.append(Self.appBaseURL) + candidates.append(baseURL) + + var seen = Set() + return candidates.filter { url in + let key = url.absoluteString + if seen.contains(key) { return false } + seen.insert(key) + return true + } + } + + private static func bearerToken(from cookies: [HTTPCookie]) -> String? { + let accessToken = cookies.first(where: { $0.name == "access-token" })?.value + let sessionToken = cookies.first(where: { Self.authSessionCookieNames.contains($0.name) })?.value + let legacySession = cookies.first(where: { $0.name == "session" })?.value + + if let accessToken, accessToken.contains(".") { + return accessToken + } + if let sessionToken, sessionToken.contains(".") { + return sessionToken + } + if let legacySession, legacySession.contains(".") { + return legacySession + } + return accessToken ?? sessionToken + } + + private func fetchWorkOSAccessToken( + refreshToken: String, + organizationID: String?) async throws -> WorkOSAuthResponse + { + var lastError: Error? + for clientID in Self.workosClientIDs { + do { + return try await self.fetchWorkOSAccessToken( + refreshToken: refreshToken, + organizationID: organizationID, + clientID: clientID) + } catch { + lastError = error + } + } + if let lastError { throw lastError } + throw FactoryStatusProbeError.networkError("WorkOS auth failed") + } + + private func fetchWorkOSAccessToken( + refreshToken: String, + organizationID: String?, + clientID: String) async throws -> WorkOSAuthResponse + { + guard let url = URL(string: "https://api.workos.com/user_management/authenticate") else { + throw FactoryStatusProbeError.networkError("WorkOS auth URL unavailable") + } + + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + + var body: [String: Any] = [ + "client_id": clientID, + "grant_type": "refresh_token", + "refresh_token": refreshToken, + ] + if let organizationID { + body["organization_id"] = organizationID + } + request.httpBody = try JSONSerialization.data(withJSONObject: body, options: []) + + let (data, response) = try await self.transport.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw FactoryStatusProbeError.networkError("Invalid WorkOS response") + } + + guard httpResponse.statusCode == 200 else { + if httpResponse.statusCode == 400, Self.isMissingWorkOSRefreshToken(data) { + throw FactoryStatusProbeError.noSessionCookie + } + let body = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + let snippet = body.isEmpty ? "" : ": \(body.prefix(200))" + throw FactoryStatusProbeError.networkError("WorkOS HTTP \(httpResponse.statusCode)\(snippet)") + } + + let decoder = JSONDecoder() + return try decoder.decode(WorkOSAuthResponse.self, from: data) + } + + private func fetchWorkOSAccessTokenWithCookies( + cookies: [HTTPCookie], + logger: (String) -> Void) async throws -> WorkOSAuthResponse + { + let cookieHeader = Self.cookieHeader(from: cookies) + guard !cookieHeader.isEmpty else { + throw FactoryStatusProbeError.networkError("Missing WorkOS cookies") + } + + var lastError: Error? + for clientID in Self.workosClientIDs { + do { + return try await self.fetchWorkOSAccessTokenWithCookies( + cookieHeader: cookieHeader, + organizationID: nil, + clientID: clientID) + } catch { + lastError = error + logger("WorkOS cookie auth failed for client \(clientID): \(error.localizedDescription)") + } + } + if let lastError { throw lastError } + throw FactoryStatusProbeError.networkError("WorkOS cookie auth failed") + } + + private func fetchWorkOSAccessTokenWithCookies( + cookieHeader: String, + organizationID: String?, + clientID: String) async throws -> WorkOSAuthResponse + { + guard let url = URL(string: "https://api.workos.com/user_management/authenticate") else { + throw FactoryStatusProbeError.networkError("WorkOS auth URL unavailable") + } + + var request = URLRequest(url: url) + request.timeoutInterval = self.timeout + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + + var body: [String: Any] = [ + "client_id": clientID, + "grant_type": "refresh_token", + "useCookie": true, + ] + if let organizationID { + body["organization_id"] = organizationID + } + request.httpBody = try JSONSerialization.data(withJSONObject: body, options: []) + + let (data, response) = try await self.transport.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw FactoryStatusProbeError.networkError("Invalid WorkOS response") + } + + guard httpResponse.statusCode == 200 else { + if httpResponse.statusCode == 400, Self.isMissingWorkOSRefreshToken(data) { + throw FactoryStatusProbeError.noSessionCookie + } + let bodyText = String(data: data, encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + let snippet = bodyText.isEmpty ? "" : ": \(bodyText.prefix(200))" + throw FactoryStatusProbeError.networkError("WorkOS HTTP \(httpResponse.statusCode)\(snippet)") + } + + let decoder = JSONDecoder() + return try decoder.decode(WorkOSAuthResponse.self, from: data) + } + + private func isInvalidGrant(_ error: Error) -> Bool { + guard case let FactoryStatusProbeError.networkError(message) = error else { + return false + } + return message.localizedCaseInsensitiveContains("invalid_grant") + } + + static func isMissingWorkOSRefreshToken(_ data: Data) -> Bool { + guard let json = try? JSONSerialization.jsonObject(with: data, options: []) as? [String: Any] else { + return false + } + guard let description = json["error_description"] as? String else { return false } + return description.localizedCaseInsensitiveContains("missing refresh token") + } + + #endif + + private func buildSnapshot( + authInfo: FactoryAuthResponse, + usageData: FactoryUsageResponse, + userId: String?) -> FactoryStatusSnapshot + { + let usage = usageData.usage + + let periodStart: Date? = usage?.startDate.map { Date(timeIntervalSince1970: TimeInterval($0) / 1000) } + let periodEnd: Date? = usage?.endDate.map { Date(timeIntervalSince1970: TimeInterval($0) / 1000) } + + return FactoryStatusSnapshot( + standardUserTokens: usage?.standard?.userTokens ?? 0, + standardOrgTokens: usage?.standard?.orgTotalTokensUsed ?? 0, + standardAllowance: usage?.standard?.totalAllowance ?? 0, + standardUsedRatio: usage?.standard?.usedRatio, + premiumUserTokens: usage?.premium?.userTokens ?? 0, + premiumOrgTokens: usage?.premium?.orgTotalTokensUsed ?? 0, + premiumAllowance: usage?.premium?.totalAllowance ?? 0, + premiumUsedRatio: usage?.premium?.usedRatio, + periodStart: periodStart, + periodEnd: periodEnd, + planName: authInfo.organization?.subscription?.orbSubscription?.plan?.name, + tier: authInfo.organization?.subscription?.factoryTier, + organizationName: authInfo.organization?.name, + accountEmail: nil, // Email is in JWT, not in auth response body + userId: userId ?? usageData.userId, + rawJSON: nil) + } + + private func buildTokenRateLimitsSnapshot( + authInfo: FactoryAuthResponse, + billingLimits: FactoryBillingLimitsResponse, + tokenRateLimits: FactoryTokenRateLimits, + userId: String?) -> FactoryStatusSnapshot + { + FactoryStatusSnapshot( + standardUserTokens: 0, + standardOrgTokens: 0, + standardAllowance: 0, + standardUsedRatio: nil, + premiumUserTokens: 0, + premiumOrgTokens: 0, + premiumAllowance: 0, + premiumUsedRatio: nil, + periodStart: nil, + periodEnd: nil, + planName: authInfo.organization?.subscription?.orbSubscription?.plan?.name, + tier: authInfo.organization?.subscription?.factoryTier, + organizationName: authInfo.organization?.name, + accountEmail: nil, + userId: userId, + rawJSON: nil, + tokenRateLimits: tokenRateLimits, + extraUsageBalanceCents: billingLimits.extraUsageBalanceCents, + overagePreference: billingLimits.overagePreference) + } +} + +private func factoryUserIdFromAuth(_ auth: FactoryAuthResponse) -> String? { + factoryNormalizedString(auth.userProfile?.id) +} + +private func factoryUserIdFromBearerToken(_ token: String?) -> String? { + guard let token, + let claims = UsageFetcher.parseJWT(token), + let subject = claims["sub"] as? String + else { + return nil + } + return factoryNormalizedString(subject) +} + +private func factoryNormalizedString(_ value: String?) -> String? { + guard let value = value?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + return value +} diff --git a/Sources/CodexBarCore/Providers/Factory/FactoryStatusProbeError.swift b/Sources/CodexBarCore/Providers/Factory/FactoryStatusProbeError.swift new file mode 100644 index 0000000..a486eec --- /dev/null +++ b/Sources/CodexBarCore/Providers/Factory/FactoryStatusProbeError.swift @@ -0,0 +1,47 @@ +import Foundation +#if os(macOS) +import SweetCookieKit + +private let factoryAPIKeyCookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.factory]?.browserCookieOrder ?? Browser.defaultImportOrder +#endif + +public enum FactoryStatusProbeError: LocalizedError, Sendable, Equatable { + case notSupported + case notLoggedIn + case missingAPIKey + case unauthorizedAPIKey + case networkError(String) + case parseFailed(String) + case noSessionCookie + + public var errorDescription: String? { + switch self { + case .notSupported: + "Factory browser-cookie auth is only supported on macOS. Use FACTORY_API_KEY / --source api on Linux." + case .notLoggedIn: + #if os(macOS) + "No usable Droid session found. Log in to app.factory.ai in \(factoryAPIKeyCookieImportOrder.loginHint), " + + "then refresh Droid." + #else + "No usable Droid session found. Log in to app.factory.ai, then refresh Droid." + #endif + case .missingAPIKey: + "Droid API key missing. Set FACTORY_API_KEY, add providers[].apiKey for factory in " + + "~/.codexbar/config.json, or run `codexbar config set-api-key --provider factory`." + case .unauthorizedAPIKey: + "Droid API authentication failed (401/403). Refresh FACTORY_API_KEY or regenerate a key at " + + "app.factory.ai/settings/api-keys." + case let .networkError(msg): + "Factory API error: \(msg)" + case let .parseFailed(msg): + "Could not parse Factory usage: \(msg)" + case .noSessionCookie: + #if os(macOS) + "No Factory session found. Please log in to app.factory.ai in \(factoryAPIKeyCookieImportOrder.loginHint)." + #else + "No Factory session found. Please log in to app.factory.ai." + #endif + } + } +} diff --git a/Sources/CodexBarCore/Providers/Gemini/GeminiConsumerTierMigration.swift b/Sources/CodexBarCore/Providers/Gemini/GeminiConsumerTierMigration.swift new file mode 100644 index 0000000..36ee22b --- /dev/null +++ b/Sources/CodexBarCore/Providers/Gemini/GeminiConsumerTierMigration.swift @@ -0,0 +1,16 @@ +import Foundation + +/// Shared copy for Google's June 2026 Gemini CLI consumer-tier shutdown. +public enum GeminiConsumerTierMigration { + public static let deprecationError = """ + Google no longer supports Gemini CLI OAuth for individual, AI Pro, or Ultra accounts. \ + Enable CodexBar's Antigravity provider, sign in to Antigravity or run `agy`, then refresh. + """ + + public static let oauthRecoveryError = """ + Could not refresh Gemini OAuth credentials. Reinstall or update Gemini CLI, or set \ + GEMINI_OAUTH_CLIENT_ID and GEMINI_OAUTH_CLIENT_SECRET. Consumer Google AI Pro/Ultra \ + accounts blocked by the June 2026 Gemini CLI shutdown should use CodexBar's Antigravity \ + provider instead. Workspace and education accounts should keep using Gemini. + """ +} diff --git a/Sources/CodexBarCore/Providers/Gemini/GeminiOAuthConfig.swift b/Sources/CodexBarCore/Providers/Gemini/GeminiOAuthConfig.swift new file mode 100644 index 0000000..4463648 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Gemini/GeminiOAuthConfig.swift @@ -0,0 +1,77 @@ +import Foundation + +/// Gemini CLI OAuth client resolution helpers for token refresh. +/// Mirrors Antigravity's env override pattern. +public enum GeminiOAuthConfig: Sendable { + public struct ClientCredentials: Sendable, Equatable { + public let clientID: String + public let clientSecret: String + + public init(clientID: String, clientSecret: String) { + self.clientID = clientID + self.clientSecret = clientSecret + } + } + + /// Test-injectable environment view so suites can override OAuth knobs without + /// mutating process-wide env (which races parallel Gemini suites). + public struct EnvironmentValues: Sendable, Equatable { + public var clientID: String? + public var clientSecret: String? + public var oauth2JSPath: String? + + public init(clientID: String? = nil, clientSecret: String? = nil, oauth2JSPath: String? = nil) { + self.clientID = clientID + self.clientSecret = clientSecret + self.oauth2JSPath = oauth2JSPath + } + + public static func fromProcessEnvironment( + _ environment: [String: String] = ProcessInfo.processInfo.environment) -> Self + { + Self( + clientID: environment["GEMINI_OAUTH_CLIENT_ID"], + clientSecret: environment["GEMINI_OAUTH_CLIENT_SECRET"], + oauth2JSPath: environment["GEMINI_OAUTH2_JS_PATH"]) + } + } + + @TaskLocal public static var environmentOverride: EnvironmentValues? + + public static var currentEnvironment: EnvironmentValues { + self.environmentOverride ?? .fromProcessEnvironment() + } + + public static var configuredClientID: String? { + self.currentEnvironment.clientID? + .trimmingCharacters(in: .whitespacesAndNewlines) + .nilIfEmpty + } + + public static var configuredClientSecret: String? { + self.currentEnvironment.clientSecret? + .trimmingCharacters(in: .whitespacesAndNewlines) + .nilIfEmpty + } + + public static var configuredOAuth2JSPath: String? { + self.currentEnvironment.oauth2JSPath? + .trimmingCharacters(in: .whitespacesAndNewlines) + .nilIfEmpty + } + + public static func environmentClient() -> ClientCredentials? { + guard let clientID = configuredClientID, + let clientSecret = configuredClientSecret + else { + return nil + } + return ClientCredentials(clientID: clientID, clientSecret: clientSecret) + } +} + +extension String { + fileprivate var nilIfEmpty: String? { + self.isEmpty ? nil : self + } +} diff --git a/Sources/CodexBarCore/Providers/Gemini/GeminiProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Gemini/GeminiProviderDescriptor.swift new file mode 100644 index 0000000..322ad1f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Gemini/GeminiProviderDescriptor.swift @@ -0,0 +1,65 @@ +import Foundation + +public enum GeminiProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .gemini, + metadata: ProviderMetadata( + id: .gemini, + displayName: "Gemini", + sessionLabel: "Pro", + weeklyLabel: "Flash", + opusLabel: "Flash Lite", + supportsOpus: true, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Gemini usage", + cliName: "gemini", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + dashboardURL: "https://gemini.google.com", + changelogURL: "https://github.com/google-gemini/gemini-cli/releases", + statusPageURL: nil, + statusLinkURL: "https://www.google.com/appsstatus/dashboard/products/npdyhgECDJ6tB66MxXyo/history", + statusWorkspaceProductID: "npdyhgECDJ6tB66MxXyo"), + branding: ProviderBranding( + iconStyle: .gemini, + iconResourceName: "ProviderIcon-gemini", + color: ProviderColor(red: 171 / 255, green: 135 / 255, blue: 234 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Gemini cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .api], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [GeminiStatusFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "gemini", + versionDetector: { _ in ProviderVersionDetector.geminiVersion() })) + } +} + +struct GeminiStatusFetchStrategy: ProviderFetchStrategy { + static let sourceLabel = "oauth-api" + + let id: String = "gemini.api" + let kind: ProviderFetchKind = .apiToken + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { + let probe = GeminiStatusProbe() + let snap = try await probe.fetch() + return self.makeResult( + usage: snap.toUsageSnapshot(), + sourceLabel: Self.sourceLabel) + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} diff --git a/Sources/CodexBarCore/Providers/Gemini/GeminiStatusProbe+DataLoader.swift b/Sources/CodexBarCore/Providers/Gemini/GeminiStatusProbe+DataLoader.swift new file mode 100644 index 0000000..502f496 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Gemini/GeminiStatusProbe+DataLoader.swift @@ -0,0 +1,131 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +extension GeminiStatusProbe { + public static func defaultDataLoader(for request: URLRequest) async throws -> (Data, URLResponse) { + let loader = Self.dataLoaderWithCurlFallback( + primary: { request in + try await ProviderHTTPClient.shared.data(for: request) + }, + fallback: Self.curlDataLoader) + return try await loader(request) + } + + public static func dataLoaderWithCurlFallback( + primary: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse), + fallback: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) + -> @Sendable (URLRequest) async throws -> (Data, URLResponse) + { + { request in + do { + return try await primary(request) + } catch { + guard Self.isURLSessionTimeout(error) else { + throw error + } + CodexBarLog.logger(LogCategories.geminiProbe) + .warning("Gemini URLSession timed out; retrying with curl") + return try await fallback(request) + } + } + } + + private static func isURLSessionTimeout(_ error: Error) -> Bool { + if let urlError = error as? URLError { + return urlError.code == .timedOut + } + + let nsError = error as NSError + return nsError.domain == NSURLErrorDomain && nsError.code == NSURLErrorTimedOut + } + + private static func curlDataLoader(for request: URLRequest) async throws -> (Data, URLResponse) { + guard let url = request.url else { + throw URLError(.badURL) + } + + let fileManager = FileManager.default + let tempDir = fileManager.temporaryDirectory + .appendingPathComponent("codexbar-gemini-curl-\(UUID().uuidString)", isDirectory: true) + try fileManager.createDirectory( + at: tempDir, + withIntermediateDirectories: true, + attributes: [.posixPermissions: 0o700]) + defer { + try? fileManager.removeItem(at: tempDir) + } + + let configURL = tempDir.appendingPathComponent("curl.conf") + var config = [ + "silent", + "show-error", + "location", + "url = \(Self.curlConfigQuote(url.absoluteString))", + "max-time = \(max(1, Int(ceil(request.timeoutInterval))))", + ] + + if let method = request.httpMethod, !method.isEmpty { + config.append("request = \(Self.curlConfigQuote(method))") + } + + for (name, value) in (request.allHTTPHeaderFields ?? [:]).sorted(by: { $0.key < $1.key }) { + let header = "\(name): \(value)" + guard !header.contains("\n"), !header.contains("\r") else { + throw GeminiStatusProbeError.apiError("Invalid request header") + } + config.append("header = \(Self.curlConfigQuote(header))") + } + + if let body = request.httpBody { + let bodyURL = tempDir.appendingPathComponent("body") + try body.write(to: bodyURL, options: .atomic) + try fileManager.setAttributes([.posixPermissions: 0o600], ofItemAtPath: bodyURL.path) + config.append("data-binary = \(Self.curlConfigQuote("@\(bodyURL.path)"))") + } + + config.append("write-out = \(Self.curlConfigQuote(Self.curlHTTPStatusMarker + "%{http_code}"))") + try config.joined(separator: "\n").write(to: configURL, atomically: true, encoding: .utf8) + try fileManager.setAttributes([.posixPermissions: 0o600], ofItemAtPath: configURL.path) + + let result = try await SubprocessRunner.run( + binary: "/usr/bin/curl", + arguments: ["--config", configURL.path], + environment: TTYCommandRunner.enrichedEnvironment(), + timeout: max(5, request.timeoutInterval + 2), + label: "gemini-api-curl") + + return try Self.parseCurlDataLoaderResult(result.stdout, url: url) + } + + private static let curlHTTPStatusMarker = "__CODEXBAR_HTTP_STATUS__:" + + private static func curlConfigQuote(_ value: String) -> String { + let escaped = value + .replacingOccurrences(of: "\\", with: "\\\\") + .replacingOccurrences(of: "\"", with: "\\\"") + return "\"\(escaped)\"" + } + + private static func parseCurlDataLoaderResult(_ output: String, url: URL) throws -> (Data, URLResponse) { + guard let markerRange = output.range(of: curlHTTPStatusMarker, options: .backwards) else { + throw GeminiStatusProbeError.apiError("curl response missing HTTP status") + } + + let body = String(output[.. 0, + let response = HTTPURLResponse( + url: url, + statusCode: statusCode, + httpVersion: nil, + headerFields: nil) + else { + throw GeminiStatusProbeError.apiError("curl response had invalid HTTP status") + } + + return (Data(body.utf8), response) + } +} diff --git a/Sources/CodexBarCore/Providers/Gemini/GeminiStatusProbe.swift b/Sources/CodexBarCore/Providers/Gemini/GeminiStatusProbe.swift new file mode 100644 index 0000000..27c46d8 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Gemini/GeminiStatusProbe.swift @@ -0,0 +1,1397 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#elseif canImport(Musl) +import Musl +#endif + +public struct GeminiModelQuota: Sendable { + public let modelId: String + public let percentLeft: Double + public let resetTime: Date? + public let resetDescription: String? +} + +public struct GeminiStatusSnapshot: Sendable { + public let modelQuotas: [GeminiModelQuota] + public let rawText: String + public let accountEmail: String? + public let accountPlan: String? + + // Convenience: lowest quota across all models (for icon display) + public var lowestPercentLeft: Double? { + self.modelQuotas.min(by: { $0.percentLeft < $1.percentLeft })?.percentLeft + } + + /// Legacy compatibility + public var dailyPercentLeft: Double? { + self.lowestPercentLeft + } + + public var resetDescription: String? { + self.modelQuotas.min(by: { $0.percentLeft < $1.percentLeft })?.resetDescription + } + + /// Converts Gemini quotas to a unified UsageSnapshot. + /// Groups quotas by tier: Pro (24h window) as primary, Flash (24h window) as secondary, + /// Flash Lite (24h window) as tertiary. + public func toUsageSnapshot() -> UsageSnapshot { + let lower = self.modelQuotas.map { ($0.modelId.lowercased(), $0) } + let flashLiteQuotas = lower.filter { Self.isFlashLiteModel(id: $0.0) }.map(\.1) + let flashQuotas = lower.filter { Self.isFlashModel(id: $0.0) }.map(\.1) + let proQuotas = lower.filter { Self.isProModel(id: $0.0) }.map(\.1) + + let flashLiteMin = flashLiteQuotas.min(by: { $0.percentLeft < $1.percentLeft }) + let flashMin = flashQuotas.min(by: { $0.percentLeft < $1.percentLeft }) + let proMin = proQuotas.min(by: { $0.percentLeft < $1.percentLeft }) + + // Keep missing tiers nil so downstream selectors can fall back to a quota the account actually reports. + let primary: RateWindow? = proMin.map { + RateWindow( + usedPercent: 100 - $0.percentLeft, + windowMinutes: 1440, + resetsAt: $0.resetTime, + resetDescription: $0.resetDescription) + } + + let secondary: RateWindow? = flashMin.map { + RateWindow( + usedPercent: 100 - $0.percentLeft, + windowMinutes: 1440, + resetsAt: $0.resetTime, + resetDescription: $0.resetDescription) + } + let tertiary: RateWindow? = flashLiteMin.map { + RateWindow( + usedPercent: 100 - $0.percentLeft, + windowMinutes: 1440, + resetsAt: $0.resetTime, + resetDescription: $0.resetDescription) + } + + let identity = ProviderIdentitySnapshot( + providerID: .gemini, + accountEmail: self.accountEmail, + accountOrganization: nil, + loginMethod: self.accountPlan) + return UsageSnapshot( + primary: primary, + secondary: secondary, + tertiary: tertiary, + updatedAt: Date(), + identity: identity) + } + + private static func isFlashLiteModel(id: String) -> Bool { + id.contains("flash-lite") + } + + private static func isFlashModel(id: String) -> Bool { + id.contains("flash") && !self.isFlashLiteModel(id: id) + } + + private static func isProModel(id: String) -> Bool { + id.contains("pro") + } +} + +public enum GeminiStatusProbeError: LocalizedError, Sendable, Equatable { + case geminiNotInstalled + case notLoggedIn + case unsupportedAuthType(String) + case consumerTierDeprecated + case parseFailed(String) + case timedOut + case apiError(String) + + public var errorDescription: String? { + switch self { + case .geminiNotInstalled: + "Gemini CLI is not installed or not on PATH." + case .notLoggedIn: + "Not logged in to Gemini. Run 'gemini' in Terminal to authenticate." + case let .unsupportedAuthType(authType): + "Gemini \(authType) auth not supported. Use Google account (OAuth) instead." + case .consumerTierDeprecated: + GeminiConsumerTierMigration.deprecationError + case let .parseFailed(msg): + "Could not parse Gemini usage: \(msg)" + case .timedOut: + "Gemini quota API request timed out." + case let .apiError(msg): + "Gemini API error: \(msg)" + } + } + + /// Detects Google's consumer-tier Gemini CLI shutdown responses. + public static func isConsumerTierDeprecationSignal(_ text: String) -> Bool { + let normalized = text.lowercased() + if normalized.contains("unsupported_client") { + return true + } + if normalized.contains("ineligibletiererror") { + return true + } + if normalized.contains("no longer supported"), normalized.contains("gemini code assist") { + return true + } + if normalized.contains("migrate"), normalized.contains("antigravity"), normalized.contains("gemini") { + return true + } + return false + } + + public static func throwIfConsumerTierDeprecated(data: Data) throws { + guard let text = String(data: data, encoding: .utf8), + isConsumerTierDeprecationSignal(text) + else { + return + } + throw GeminiStatusProbeError.consumerTierDeprecated + } +} + +public enum GeminiAuthType: String, Sendable { + case oauthPersonal = "oauth-personal" + case apiKey = "gemini-api-key" + case vertexAI = "vertex-ai" + case unknown +} + +/// User tier IDs returned from the Cloud Code Private API (loadCodeAssist). +/// Maps to: google3/cloud/developer_experience/cloudcode/pa/service/usertier.go +public enum GeminiUserTierId: String, Sendable { + case free = "free-tier" + case legacy = "legacy-tier" + case standard = "standard-tier" +} + +public struct GeminiStatusProbe: Sendable { + public var timeout: TimeInterval = 10.0 + public var homeDirectory: String + public var dataLoader: @Sendable (URLRequest) async throws -> (Data, URLResponse) + private static let log = CodexBarLog.logger(LogCategories.geminiProbe) + private static let quotaEndpoint = "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota" + private static let loadCodeAssistEndpoint = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" + private static let projectsEndpoint = "https://cloudresourcemanager.googleapis.com/v1/projects" + private static let credentialsPath = "/.gemini/oauth_creds.json" + private static let settingsPath = "/.gemini/settings.json" + private static let tokenRefreshEndpoint = "https://oauth2.googleapis.com/token" + + public init( + timeout: TimeInterval = 10.0, + homeDirectory: String = NSHomeDirectory(), + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse) = Self.defaultDataLoader) + { + self.timeout = timeout + self.homeDirectory = homeDirectory + self.dataLoader = dataLoader + } + + /// Reads the current Gemini auth type from settings.json + public static func currentAuthType(homeDirectory: String = NSHomeDirectory()) -> GeminiAuthType { + let settingsURL = URL(fileURLWithPath: homeDirectory + Self.settingsPath) + + guard let data = try? Data(contentsOf: settingsURL), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let security = json["security"] as? [String: Any], + let auth = security["auth"] as? [String: Any], + let selectedType = auth["selectedType"] as? String + else { + return .unknown + } + + if selectedType == "api-key" { + return .apiKey + } + return GeminiAuthType(rawValue: selectedType) ?? .unknown + } + + public func fetch() async throws -> GeminiStatusSnapshot { + // Block explicitly unsupported auth types; allow unknown to try OAuth creds + let authType = Self.currentAuthType(homeDirectory: self.homeDirectory) + switch authType { + case .apiKey: + throw GeminiStatusProbeError.unsupportedAuthType("API key") + case .vertexAI: + throw GeminiStatusProbeError.unsupportedAuthType("Vertex AI") + case .oauthPersonal, .unknown: + break + } + + let snap = try await Self.fetchViaAPI( + timeout: self.timeout, + homeDirectory: self.homeDirectory, + dataLoader: self.dataLoader) + + Self.log.info("Gemini API fetch ok", metadata: [ + "dailyPercentLeft": "\(snap.dailyPercentLeft ?? -1)", + ]) + return snap + } + + // MARK: - Direct API approach + + private static func fetchViaAPI( + timeout: TimeInterval, + homeDirectory: String, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> GeminiStatusSnapshot + { + let creds = try Self.loadCredentials(homeDirectory: homeDirectory) + + let expiryStr = creds.expiryDate.map { "\($0)" } ?? "nil" + let hasRefresh = creds.refreshToken != nil + Self.log.debug("Token check", metadata: [ + "expiry": expiryStr, + "hasRefresh": hasRefresh ? "1" : "0", + "now": "\(Date())", + ]) + + var accessToken = creds.accessToken?.isEmpty == false ? creds.accessToken : nil + var idToken = creds.idToken + let needsRefresh = accessToken == nil || creds.expiryDate.map { $0 < Date() } == true + if needsRefresh { + if accessToken == nil { + Self.log.info("No access token found; attempting refresh from stored Gemini credentials") + } else if let expiry = creds.expiryDate { + Self.log.info("Token expired; attempting refresh", metadata: [ + "expiry": "\(expiry)", + ]) + } + + guard let refreshToken = creds.refreshToken, !refreshToken.isEmpty else { + Self.log.error("No refresh token available") + throw GeminiStatusProbeError.notLoggedIn + } + accessToken = try await Self.refreshAccessToken( + refreshToken: refreshToken, + timeout: timeout, + homeDirectory: homeDirectory, + dataLoader: dataLoader) + idToken = (try? Self.loadCredentials(homeDirectory: homeDirectory).idToken) ?? idToken + } + guard let accessToken else { + Self.log.error("No access token found") + throw GeminiStatusProbeError.notLoggedIn + } + + // Extract account info from JWT + let claims = Self.extractClaimsFromToken(idToken) + + // Load Code Assist status to get project ID and tier (aligned with CLI setupUser logic) + let caStatus = try await Self.loadCodeAssistStatus( + accessToken: accessToken, + timeout: timeout, + dataLoader: dataLoader) + + // Determine the project ID to use for quota fetching. + // Priority: + // 1. Project ID returned by loadCodeAssist (e.g. managed project for free tier) + // 2. Discovered project ID from cloud resource manager (e.g. user's own project) + var projectId = caStatus.projectId + if projectId == nil { + projectId = try? await Self.discoverGeminiProjectId( + accessToken: accessToken, + timeout: timeout, + dataLoader: dataLoader) + } + + guard let url = URL(string: Self.quotaEndpoint) else { + throw GeminiStatusProbeError.apiError("Invalid endpoint URL") + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + + // Include project ID for accurate quota + if let projectId { + request.httpBody = Data("{\"project\": \"\(projectId)\"}".utf8) + } else { + request.httpBody = Data("{}".utf8) + } + request.timeoutInterval = timeout + + let (data, response) = try await dataLoader(request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw GeminiStatusProbeError.apiError("Invalid response") + } + + if httpResponse.statusCode == 401 { + try GeminiStatusProbeError.throwIfConsumerTierDeprecated(data: data) + throw GeminiStatusProbeError.notLoggedIn + } + + guard httpResponse.statusCode == 200 else { + try GeminiStatusProbeError.throwIfConsumerTierDeprecated(data: data) + throw GeminiStatusProbeError.apiError("HTTP \(httpResponse.statusCode)") + } + + let snapshot = try Self.parseAPIResponse(data, email: claims.email) + + let plan = Self.resolveAccountPlan( + tier: caStatus.tier, + hostedDomain: claims.hostedDomain, + paidTierName: caStatus.paidTierName) + + return GeminiStatusSnapshot( + modelQuotas: snapshot.modelQuotas, + rawText: snapshot.rawText, + accountEmail: snapshot.accountEmail, + accountPlan: plan) + } + + private static func discoverGeminiProjectId( + accessToken: String, + timeout: TimeInterval, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> String? + { + guard let url = URL(string: projectsEndpoint) else { return nil } + + var request = URLRequest(url: url) + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.timeoutInterval = timeout + + let (data, response) = try await dataLoader(request) + + guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else { + return nil + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let projects = json["projects"] as? [[String: Any]] + else { + return nil + } + + // Look for Gemini API project (has "generative-language" label or "gen-lang-client" prefix) + for project in projects { + guard let projectId = project["projectId"] as? String else { continue } + + // Check for gen-lang-client prefix (Gemini CLI projects) + if projectId.hasPrefix("gen-lang-client") { + return projectId + } + + // Check for generative-language label + if let labels = project["labels"] as? [String: String], + labels["generative-language"] != nil + { + return projectId + } + } + + return nil + } + + private struct CodeAssistStatus { + let tier: GeminiUserTierId? + let projectId: String? + let paidTierName: String? + + static let empty = CodeAssistStatus(tier: nil, projectId: nil, paidTierName: nil) + } + + private static func loadCodeAssistStatus( + accessToken: String, + timeout: TimeInterval, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> CodeAssistStatus + { + guard let url = URL(string: loadCodeAssistEndpoint) else { + self.log.warning("loadCodeAssist: invalid endpoint URL") + return .empty + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.httpBody = Data("{\"metadata\":{\"ideType\":\"GEMINI_CLI\",\"pluginType\":\"GEMINI\"}}".utf8) + request.timeoutInterval = timeout + + let data: Data + let response: URLResponse + do { + (data, response) = try await dataLoader(request) + } catch { + Self.log.warning("loadCodeAssist: request failed", metadata: ["error": "\(error)"]) + return .empty + } + + guard let httpResponse = response as? HTTPURLResponse else { + Self.log.warning("loadCodeAssist: invalid response type") + return .empty + } + + guard httpResponse.statusCode == 200 else { + try GeminiStatusProbeError.throwIfConsumerTierDeprecated(data: data) + Self.log.warning("loadCodeAssist: HTTP error", metadata: [ + "statusCode": "\(httpResponse.statusCode)", + "body": String(data: data, encoding: .utf8) ?? "", + ]) + return .empty + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + Self.log.warning("loadCodeAssist: failed to parse JSON", metadata: [ + "body": String(data: data, encoding: .utf8) ?? "", + ]) + return .empty + } + + let rawProjectId: String? = { + if let project = json["cloudaicompanionProject"] as? String { + return project + } + if let project = json["cloudaicompanionProject"] as? [String: Any] { + if let projectId = project["id"] as? String { + return projectId + } + if let projectId = project["projectId"] as? String { + return projectId + } + } + return nil + }() + let trimmedProjectId = rawProjectId?.trimmingCharacters(in: .whitespacesAndNewlines) + let projectId = trimmedProjectId?.isEmpty == true ? nil : trimmedProjectId + if let projectId { + Self.log.info("loadCodeAssist: project detected", metadata: ["projectId": projectId]) + } + + let tierId = (json["currentTier"] as? [String: Any])?["id"] as? String + let paidTierName = Self.parsePaidTierName(from: json) + + guard let tierId else { + Self.log.warning("loadCodeAssist: no currentTier.id in response", metadata: [ + "json": "\(json)", + ]) + return CodeAssistStatus(tier: nil, projectId: projectId, paidTierName: paidTierName) + } + + guard let tier = GeminiUserTierId(rawValue: tierId) else { + Self.log.warning("loadCodeAssist: unknown tier ID", metadata: ["tierId": tierId]) + return CodeAssistStatus(tier: nil, projectId: projectId, paidTierName: paidTierName) + } + + Self.log.info("loadCodeAssist: success", metadata: [ + "tier": tierId, + "projectId": projectId ?? "nil", + "paidTierName": paidTierName ?? "nil", + ]) + return CodeAssistStatus(tier: tier, projectId: projectId, paidTierName: paidTierName) + } + + private struct OAuthCredentials { + let accessToken: String? + let idToken: String? + let refreshToken: String? + let expiryDate: Date? + } + + fileprivate struct OAuthClientCredentials { + let clientId: String + let clientSecret: String + } + + private static func isLikelyFnmManagedPath(_ path: String) -> Bool { + let normalized = path.replacingOccurrences(of: "\\", with: "/") + return normalized.contains("/fnm_multishells/") + || (normalized.contains("/node-versions/") && normalized.contains("/fnm/")) + } + + private static func resolveExecutableOnEnvironmentPath( + named executable: String, + environment: [String: String]) -> String? + { + guard let path = environment["PATH"] else { return nil } + for directory in path.split(separator: ":") where !directory.isEmpty { + let candidate = URL(fileURLWithPath: String(directory), isDirectory: true) + .appendingPathComponent(executable) + .path + if FileManager.default.isExecutableFile(atPath: candidate) { + return candidate + } + } + return nil + } + + private static func resolveGeminiPackageRootViaFnm( + fnmPath: String, + environment: [String: String]) -> String? + { + guard let currentVersion = runProcess( + executable: fnmPath, + arguments: ["current"], + environment: environment, + timeout: 2.0), + !currentVersion.isEmpty + else { + return nil + } + + // Prefer npm root -g because require.resolve searches from the current + // working directory and often fails for globally-installed packages. + if let npmRoot = runProcess( + executable: fnmPath, + arguments: [ + "exec", + "--using", + currentVersion, + "npm", + "root", + "-g", + ], + environment: environment, + timeout: 4.0), + !npmRoot.isEmpty + { + let packageRoot = "\(npmRoot)/@google/gemini-cli" + let packageJSONPath = "\(packageRoot)/package.json" + if FileManager.default.fileExists(atPath: packageJSONPath) { + return packageRoot + } + } + + // Fallback for non-npm global installations. + if let packageJSONPath = runProcess( + executable: fnmPath, + arguments: [ + "exec", + "--using", + currentVersion, + "node", + "-p", + "require.resolve('@google/gemini-cli/package.json')", + ], + environment: environment, + timeout: 4.0), + !packageJSONPath.isEmpty + { + return (packageJSONPath as NSString).deletingLastPathComponent + } + + return nil + } + + private static func findGeminiPackageRoot(startingAt path: String) -> String? { + let fileManager = FileManager.default + var currentURL = URL(fileURLWithPath: path).standardizedFileURL + + var isDirectory: ObjCBool = false + if !fileManager.fileExists(atPath: currentURL.path, isDirectory: &isDirectory) || !isDirectory.boolValue { + currentURL.deleteLastPathComponent() + } + + // Bound the walk so an unrelated Gemini install elsewhere on the host + // (e.g. a global npm/brew install unrelated to the resolved binary) can't + // contaminate discovery started from the actual binary path. + let maxAscents = 8 + for _ in 0...maxAscents { + let packageJSONURL = currentURL.appendingPathComponent("package.json") + if let data = try? Data(contentsOf: packageJSONURL), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + json["name"] as? String == "@google/gemini-cli" + { + return currentURL.path + } + + // Also check for a global Node installation layout: + // /lib/node_modules/@google/gemini-cli/package.json + let globalPackageJSONURL = currentURL + .appendingPathComponent("lib") + .appendingPathComponent("node_modules") + .appendingPathComponent("@google") + .appendingPathComponent("gemini-cli") + .appendingPathComponent("package.json") + if let data = try? Data(contentsOf: globalPackageJSONURL), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + json["name"] as? String == "@google/gemini-cli" + { + return globalPackageJSONURL.deletingLastPathComponent().path + } + + // Homebrew layout: + // /libexec/lib/node_modules/@google/gemini-cli/package.json + let homebrewPackageJSONURL = currentURL + .appendingPathComponent("libexec") + .appendingPathComponent("lib") + .appendingPathComponent("node_modules") + .appendingPathComponent("@google") + .appendingPathComponent("gemini-cli") + .appendingPathComponent("package.json") + if let data = try? Data(contentsOf: homebrewPackageJSONURL), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + json["name"] as? String == "@google/gemini-cli" + { + return homebrewPackageJSONURL.deletingLastPathComponent().path + } + + let parentURL = currentURL.deletingLastPathComponent() + if parentURL.path == currentURL.path { + return nil + } + currentURL = parentURL + } + + return nil + } + + private static func extractOAuthCredentials(fromGeminiPackageRoot packageRoot: String) -> OAuthClientCredentials? { + // Check the standard distributed file first, then any sibling core package + let oauthFile = "dist/src/code_assist/oauth2.js" + let candidatePaths = [ + "\(packageRoot)/\(oauthFile)", + "\(packageRoot)/node_modules/@google/gemini-cli-core/\(oauthFile)", + ] + + for path in candidatePaths { + if let content = try? String(contentsOfFile: path, encoding: .utf8), + let credentials = Self.parseOAuthCredentials(from: content) + { + return credentials + } + } + + return Self.extractOAuthCredentialsFromBundle(packageRoot: packageRoot) + } + + private static func extractOAuthCredentialsFromBundle(packageRoot: String) -> OAuthClientCredentials? { + let bundleRoot = URL(fileURLWithPath: packageRoot).appendingPathComponent("bundle", isDirectory: true) + let entryURL = bundleRoot.appendingPathComponent("gemini.js") + + guard FileManager.default.fileExists(atPath: entryURL.path) else { + return nil + } + + var pendingURLs = [entryURL] + var visitedPaths = Set() + + while !pendingURLs.isEmpty { + let currentURL = pendingURLs.removeFirst() + let standardizedPath = currentURL.standardizedFileURL.path + guard visitedPaths.insert(standardizedPath).inserted, + let content = try? String(contentsOf: currentURL, encoding: .utf8) + else { + continue + } + + if let credentials = Self.parseOAuthCredentials(from: content) { + return credentials + } + + let imports = Self.extractRelativeJavaScriptImports(from: content) + for importPath in imports { + let nextURL = URL(fileURLWithPath: importPath, relativeTo: currentURL.deletingLastPathComponent()) + .standardizedFileURL + guard nextURL.path.hasPrefix(bundleRoot.path) else { continue } + pendingURLs.append(nextURL) + } + } + + guard let bundleFiles = try? FileManager.default.contentsOfDirectory( + at: bundleRoot, + includingPropertiesForKeys: nil, + options: [.skipsHiddenFiles]) + else { + return nil + } + + for url in bundleFiles where url.pathExtension == "js" && !visitedPaths.contains(url.standardizedFileURL.path) { + guard let content = try? String(contentsOf: url, encoding: .utf8) else { continue } + if let credentials = Self.parseOAuthCredentials(from: content) { + return credentials + } + } + + return nil + } + + private static func extractRelativeJavaScriptImports(from content: String) -> [String] { + let patterns = [ + #"(?:import|export)\s+(?:[^;]*?\s+from\s+)?[\"'](\./[^\"']+\.js)[\"']"#, + #"import\(\s*[\"'](\./[^\"']+\.js)[\"']\s*\)"#, + ] + + var discoveredPaths: [String] = [] + var seen = Set() + let fullRange = NSRange(content.startIndex..., in: content) + + for pattern in patterns { + guard let regex = try? NSRegularExpression(pattern: pattern) else { continue } + for match in regex.matches(in: content, range: fullRange) { + guard let range = Range(match.range(at: 1), in: content) else { continue } + let path = String(content[range]) + if seen.insert(path).inserted { + discoveredPaths.append(path) + } + } + } + + return discoveredPaths + } + + private static func extractOAuthCredentialsFromLegacyPaths(realGeminiPath: String) -> OAuthClientCredentials? { + let binDir = (realGeminiPath as NSString).deletingLastPathComponent + let baseDir = (binDir as NSString).deletingLastPathComponent + + let oauthSubpath = + "node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/code_assist/oauth2.js" + let nixShareSubpath = + "share/gemini-cli/node_modules/@google/gemini-cli-core/dist/src/code_assist/oauth2.js" + let oauthFile = "dist/src/code_assist/oauth2.js" + let possiblePaths = [ + // Homebrew nested structure + "\(baseDir)/libexec/lib/\(oauthSubpath)", + "\(baseDir)/lib/\(oauthSubpath)", + // Nix package layout + "\(baseDir)/\(nixShareSubpath)", + // Bun/npm sibling structure: gemini-cli-core is a sibling to gemini-cli + "\(baseDir)/../gemini-cli-core/\(oauthFile)", + // npm nested inside gemini-cli + "\(baseDir)/node_modules/@google/gemini-cli-core/\(oauthFile)", + ] + + for path in possiblePaths { + if let content = try? String(contentsOfFile: path, encoding: .utf8), + let credentials = Self.parseOAuthCredentials(from: content) + { + return credentials + } + } + + return nil + } + + private static func parseOAuthCredentials(from content: String) -> OAuthClientCredentials? { + // Match: const/let/var OAUTH_CLIENT_ID = '...'; + let clientIdPattern = #"(?:const|let|var)?\s*OAUTH_CLIENT_ID\s*=\s*['"]([\w\-\.]+)['"]\s*;"# + let secretPattern = #"(?:const|let|var)?\s*OAUTH_CLIENT_SECRET\s*=\s*['"]([\w\-]+)['"]\s*;"# + + guard let clientIdRegex = try? NSRegularExpression(pattern: clientIdPattern), + let secretRegex = try? NSRegularExpression(pattern: secretPattern) + else { + return nil + } + + let range = NSRange(content.startIndex..., in: content) + + guard let clientIdMatch = clientIdRegex.firstMatch(in: content, range: range), + let clientIdRange = Range(clientIdMatch.range(at: 1), in: content), + let secretMatch = secretRegex.firstMatch(in: content, range: range), + let secretRange = Range(secretMatch.range(at: 1), in: content) + else { + return nil + } + + let clientId = String(content[clientIdRange]) + let clientSecret = String(content[secretRange]) + + return OAuthClientCredentials(clientId: clientId, clientSecret: clientSecret) + } + + private static func refreshAccessToken( + refreshToken: String, + timeout: TimeInterval, + homeDirectory: String, + dataLoader: @escaping @Sendable (URLRequest) async throws -> (Data, URLResponse)) async throws + -> String + { + guard let url = URL(string: tokenRefreshEndpoint) else { + throw GeminiStatusProbeError.apiError("Invalid token refresh URL") + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.timeoutInterval = timeout + + guard let oauthCreds = Self.extractOAuthCredentials() else { + Self.log.error("Could not extract OAuth credentials from Gemini CLI") + throw GeminiStatusProbeError.apiError(GeminiConsumerTierMigration.oauthRecoveryError) + } + + let body = [ + "client_id=\(oauthCreds.clientId)", + "client_secret=\(oauthCreds.clientSecret)", + "refresh_token=\(refreshToken)", + "grant_type=refresh_token", + ].joined(separator: "&") + request.httpBody = body.data(using: .utf8) + + let (data, response) = try await dataLoader(request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw GeminiStatusProbeError.apiError("Invalid refresh response") + } + + guard httpResponse.statusCode == 200 else { + try GeminiStatusProbeError.throwIfConsumerTierDeprecated(data: data) + Self.log.error("Token refresh failed", metadata: [ + "statusCode": "\(httpResponse.statusCode)", + ]) + throw GeminiStatusProbeError.notLoggedIn + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let newAccessToken = json["access_token"] as? String + else { + throw GeminiStatusProbeError.parseFailed("Could not parse refresh response") + } + + // Update stored credentials with new token + try Self.updateStoredCredentials(json, homeDirectory: homeDirectory) + + Self.log.info("Token refreshed successfully") + return newAccessToken + } + + private static func updateStoredCredentials(_ refreshResponse: [String: Any], homeDirectory: String) throws { + let credsURL = URL(fileURLWithPath: homeDirectory + Self.credentialsPath) + + guard let existingCreds = try? Data(contentsOf: credsURL), + var json = try? JSONSerialization.jsonObject(with: existingCreds) as? [String: Any] + else { + return + } + + // Update with new values from refresh response + if let accessToken = refreshResponse["access_token"] { + json["access_token"] = accessToken + } + if let expiresIn = refreshResponse["expires_in"] as? Double { + json["expiry_date"] = (Date().timeIntervalSince1970 + expiresIn) * 1000 + } + if let idToken = refreshResponse["id_token"] { + json["id_token"] = idToken + } + + let updatedData = try JSONSerialization.data(withJSONObject: json, options: [.prettyPrinted]) + try updatedData.write(to: credsURL, options: .atomic) + } + + private static func loadCredentials(homeDirectory: String) throws -> OAuthCredentials { + let credsURL = URL(fileURLWithPath: homeDirectory + Self.credentialsPath) + + guard FileManager.default.fileExists(atPath: credsURL.path) else { + throw GeminiStatusProbeError.notLoggedIn + } + + let data = try Data(contentsOf: credsURL) + guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw GeminiStatusProbeError.parseFailed("Invalid credentials file") + } + + let accessToken = json["access_token"] as? String + let idToken = json["id_token"] as? String + let refreshToken = json["refresh_token"] as? String + + var expiryDate: Date? + if let expiryMs = json["expiry_date"] as? Double { + expiryDate = Date(timeIntervalSince1970: expiryMs / 1000) + } + + return OAuthCredentials( + accessToken: accessToken, + idToken: idToken, + refreshToken: refreshToken, + expiryDate: expiryDate) + } + + private struct TokenClaims { + let email: String? + let hostedDomain: String? + } + + private static func extractClaimsFromToken(_ idToken: String?) -> TokenClaims { + guard let token = idToken else { return TokenClaims(email: nil, hostedDomain: nil) } + + let parts = token.components(separatedBy: ".") + guard parts.count >= 2 else { return TokenClaims(email: nil, hostedDomain: nil) } + + var payload = parts[1] + .replacingOccurrences(of: "-", with: "+") + .replacingOccurrences(of: "_", with: "/") + + let remainder = payload.count % 4 + if remainder > 0 { + payload += String(repeating: "=", count: 4 - remainder) + } + + guard let data = Data(base64Encoded: payload, options: .ignoreUnknownCharacters), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return TokenClaims(email: nil, hostedDomain: nil) + } + + return TokenClaims( + email: json["email"] as? String, + hostedDomain: json["hd"] as? String) + } + + private static func extractEmailFromToken(_ idToken: String?) -> String? { + self.extractClaimsFromToken(idToken).email + } + + private struct QuotaBucket: Decodable { + let remainingFraction: Double? + let resetTime: String? + let modelId: String? + let tokenType: String? + } + + private struct QuotaResponse: Decodable { + let buckets: [QuotaBucket]? + } + + private static func parseAPIResponse(_ data: Data, email: String?) throws -> GeminiStatusSnapshot { + let decoder = JSONDecoder() + let response = try decoder.decode(QuotaResponse.self, from: data) + + guard let buckets = response.buckets, !buckets.isEmpty else { + throw GeminiStatusProbeError.parseFailed("No quota buckets in response") + } + + // Group quotas by model, keeping lowest per model (input tokens usually) + var modelQuotaMap: [String: (fraction: Double, resetString: String?)] = [:] + + for bucket in buckets { + guard let modelId = bucket.modelId, let fraction = bucket.remainingFraction else { continue } + + if let existing = modelQuotaMap[modelId] { + if fraction < existing.fraction { + modelQuotaMap[modelId] = (fraction, bucket.resetTime) + } + } else { + modelQuotaMap[modelId] = (fraction, bucket.resetTime) + } + } + + // Convert to sorted array (by model name for consistent ordering) + let quotas = modelQuotaMap + .sorted { $0.key < $1.key } + .map { modelId, info in + let resetDate = info.resetString.flatMap { Self.parseResetTime($0) } + return GeminiModelQuota( + modelId: modelId, + percentLeft: info.fraction * 100, + resetTime: resetDate, + resetDescription: info.resetString.flatMap { Self.formatResetTime($0) }) + } + + let rawText = String(data: data, encoding: .utf8) ?? "" + + return GeminiStatusSnapshot( + modelQuotas: quotas, + rawText: rawText, + accountEmail: email, + accountPlan: nil) + } + + private static func parseResetTime(_ isoString: String) -> Date? { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + + if let date = formatter.date(from: isoString) { + return date + } + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: isoString) + } + + private static func formatResetTime(_ isoString: String) -> String { + guard let resetDate = parseResetTime(isoString) else { + return "Resets soon" + } + + let now = Date() + let interval = resetDate.timeIntervalSince(now) + + if interval <= 0 { + return "Resets soon" + } + + let hours = Int(interval / 3600) + let minutes = Int((interval.truncatingRemainder(dividingBy: 3600)) / 60) + + if hours > 0 { + return "Resets in \(hours)h \(minutes)m" + } else { + return "Resets in \(minutes)m" + } + } + + // MARK: - Legacy CLI parsing (kept for fallback) + + public static func parse(text: String) throws -> GeminiStatusSnapshot { + let clean = TextParsing.stripANSICodes(text) + guard !clean.isEmpty else { throw GeminiStatusProbeError.timedOut } + + let quotas = Self.parseModelUsageTable(clean) + + if quotas.isEmpty { + if clean.contains("Login with Google") || clean.contains("Use Gemini API key") { + throw GeminiStatusProbeError.notLoggedIn + } + if clean.contains("Waiting for auth"), !clean.contains("Usage") { + throw GeminiStatusProbeError.notLoggedIn + } + throw GeminiStatusProbeError.parseFailed("No usage data found in /stats output") + } + + return GeminiStatusSnapshot( + modelQuotas: quotas, + rawText: text, + accountEmail: nil, + accountPlan: nil) + } + + private static func parseModelUsageTable(_ text: String) -> [GeminiModelQuota] { + let lines = text.components(separatedBy: .newlines) + var quotas: [GeminiModelQuota] = [] + + let pattern = #"(gemini[-\w.]+)\s+[\d-]+\s+([0-9]+(?:\.[0-9]+)?)\s*%\s*\(([^)]+)\)"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: [.caseInsensitive]) else { + return [] + } + + for line in lines { + let cleanLine = line.replacingOccurrences(of: "│", with: " ") + let range = NSRange(cleanLine.startIndex..= 4 else { continue } + + guard let modelRange = Range(match.range(at: 1), in: cleanLine), + let pctRange = Range(match.range(at: 2), in: cleanLine), + let pct = Double(cleanLine[pctRange]) + else { continue } + + let modelId = String(cleanLine[modelRange]) + var resetDesc: String? + if let resetRange = Range(match.range(at: 3), in: cleanLine) { + resetDesc = String(cleanLine[resetRange]).trimmingCharacters(in: .whitespaces) + } + + quotas.append(GeminiModelQuota( + modelId: modelId, + percentLeft: pct, + resetTime: nil, + resetDescription: resetDesc)) + } + + return quotas + } +} + +extension GeminiStatusProbe { + fileprivate static func extractOAuthCredentials() -> OAuthClientCredentials? { + if let resolved = GeminiOAuthConfig.environmentClient() { + return OAuthClientCredentials(clientId: resolved.clientID, clientSecret: resolved.clientSecret) + } + if let path = GeminiOAuthConfig.configuredOAuth2JSPath, + let credentials = Self.parseOAuthCredentials(fromFile: path) + { + return credentials + } + if let credentials = Self.discoverOAuthCredentialsFromInstalledCLI() { + return OAuthClientCredentials(clientId: credentials.clientID, clientSecret: credentials.clientSecret) + } + if let credentials = Self.discoverOAuthCredentialsFromKnownInstallPaths() { + return credentials + } + return nil + } + + /// Optional Homebrew/npm prefixes for last-resort discovery when the gemini + /// binary cannot be resolved (GUI PATH / sandbox). Tests inject fake Cellar trees. + @TaskLocal public static var knownInstallPrefixesForTesting: [String]? + + fileprivate static func discoverOAuthCredentialsFromKnownInstallPaths() -> OAuthClientCredentials? { + let oauthFile = "dist/src/code_assist/oauth2.js" + let nestedOAuthFile = + "node_modules/@google/gemini-cli/node_modules/@google/gemini-cli-core/\(oauthFile)" + let home = FileManager.default.homeDirectoryForCurrentUser.path + + // When tests inject Homebrew prefixes, skip the hard-coded host paths so + // fixture Cellar trees are not shadowed by the developer's real install. + if Self.knownInstallPrefixesForTesting == nil { + let possiblePaths = [ + "/opt/homebrew/lib/node_modules/@google/gemini-cli-core/\(oauthFile)", + "/opt/homebrew/lib/\(nestedOAuthFile)", + "/usr/local/lib/node_modules/@google/gemini-cli-core/\(oauthFile)", + "/usr/local/lib/\(nestedOAuthFile)", + "\(home)/.npm-global/lib/node_modules/@google/gemini-cli-core/\(oauthFile)", + "\(home)/.npm-global/lib/\(nestedOAuthFile)", + ] + + for path in possiblePaths { + if let credentials = Self.parseOAuthCredentials(fromFile: path) { + return credentials + } + } + } + + // Homebrew Cellar/opt libexec layouts keep credentials inside the package + // bundle when GUI PATH cannot resolve `/opt/homebrew/bin/gemini`. + for packageRoot in Self.knownHomebrewGeminiPackageRoots() { + if let credentials = Self.extractOAuthCredentials(fromGeminiPackageRoot: packageRoot) { + return credentials + } + } + return nil + } + + fileprivate static func knownHomebrewPrefixes() -> [String] { + if let overrides = self.knownInstallPrefixesForTesting, !overrides.isEmpty { + return overrides + } + return ["/opt/homebrew", "/usr/local"] + } + + fileprivate static func knownHomebrewGeminiPackageRoots( + fileManager: FileManager = .default) -> [String] + { + var roots: [String] = [] + var seen = Set() + + func appendPackageRoot(under prefixRoot: String) { + let packageRoot = URL(fileURLWithPath: prefixRoot) + .appendingPathComponent("libexec") + .appendingPathComponent("lib") + .appendingPathComponent("node_modules") + .appendingPathComponent("@google") + .appendingPathComponent("gemini-cli") + .path + guard fileManager.fileExists(atPath: packageRoot), + seen.insert(packageRoot).inserted + else { + return + } + roots.append(packageRoot) + } + + for prefix in Self.knownHomebrewPrefixes() { + let optRoot = "\(prefix)/opt/gemini-cli" + if fileManager.fileExists(atPath: optRoot) { + appendPackageRoot(under: optRoot) + } + + let cellarRoot = "\(prefix)/Cellar/gemini-cli" + guard let versions = try? fileManager.contentsOfDirectory(atPath: cellarRoot) else { + continue + } + for version in versions.sorted() { + appendPackageRoot(under: "\(cellarRoot)/\(version)") + } + } + + return roots + } + + fileprivate static func parseOAuthCredentials(fromFile path: String) -> OAuthClientCredentials? { + guard let content = try? String(contentsOfFile: path, encoding: .utf8) else { + return nil + } + return Self.parseOAuthCredentials(from: content) + } + + fileprivate static func discoverOAuthCredentialsFromInstalledCLI() -> GeminiOAuthConfig.ClientCredentials? { + let env = ProcessInfo.processInfo.environment + + guard let geminiPath = BinaryLocator.resolveGeminiBinary( + env: env, + loginPATH: LoginShellPathCache.shared.current) + ?? TTYCommandRunner.which("gemini") + else { + return nil + } + + let resolvedGeminiPath = URL(fileURLWithPath: geminiPath).resolvingSymlinksInPath().path + + if let credentials = Self.extractOAuthCredentialsFromLegacyPaths(realGeminiPath: resolvedGeminiPath) { + return GeminiOAuthConfig.ClientCredentials( + clientID: credentials.clientId, + clientSecret: credentials.clientSecret) + } + + if Self.isLikelyFnmManagedPath(geminiPath) || Self.isLikelyFnmManagedPath(resolvedGeminiPath), + let fnmPath = Self.resolveExecutableOnEnvironmentPath(named: "fnm", environment: env) + ?? TTYCommandRunner.which("fnm"), + let packageRoot = Self.resolveGeminiPackageRootViaFnm(fnmPath: fnmPath, environment: env), + let credentials = Self.extractOAuthCredentials(fromGeminiPackageRoot: packageRoot) + { + return GeminiOAuthConfig.ClientCredentials( + clientID: credentials.clientId, + clientSecret: credentials.clientSecret) + } + + if let packageRoot = Self.findGeminiPackageRoot(startingAt: resolvedGeminiPath), + let credentials = Self.extractOAuthCredentials(fromGeminiPackageRoot: packageRoot) + { + return GeminiOAuthConfig.ClientCredentials( + clientID: credentials.clientId, + clientSecret: credentials.clientSecret) + } + + return nil + } + + /// Plan display strings with tier mapping: + /// - paidTier.name: Most-specific paid subscription label from Google, regardless of currentTier + /// - standard-tier: Paid subscription fallback (Code Assist Standard/Enterprise, Developer Program Premium) + /// - free-tier + hd claim: Workspace account (Gemini included free since Jan 2025) + /// - free-tier: Personal free account + /// - legacy-tier: Unknown legacy/grandfathered tier + /// - nil (API failed): Leave blank (no display) + fileprivate static func resolveAccountPlan( + tier: GeminiUserTierId?, + hostedDomain: String?, + paidTierName: String?) -> String? + { + // Match Gemini CLI's contract: a named paid tier is the most specific plan signal, + // even when currentTier is missing, unknown, or still reports free-tier. + if let paidTierName { + self.log.info("Paid tier detected", metadata: [ + "tier": tier?.rawValue ?? "unknown", + "plan": paidTierName, + ]) + return paidTierName + } + + switch (tier, hostedDomain) { + case (.standard, _): + return "Paid" + case let (.free, .some(domain)): + Self.log.info("Workspace account detected", metadata: ["domain": domain]) + return "Workspace" + case (.free, .none): + Self.log.info("Personal free account") + return "Free" + case (.legacy, _): + return "Legacy" + case (.none, _): + self.log.info("Tier detection failed, leaving plan blank") + return nil + } + } + + private static func parsePaidTierName(from json: [String: Any]) -> String? { + guard let paidTier = json["paidTier"] as? [String: Any], + let rawName = paidTier["name"] as? String + else { + return nil + } + let trimmed = rawName.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + } +} + +extension GeminiStatusProbe { + private static let processTimeoutQueue = DispatchQueue( + label: "com.steipete.codexbar.gemini-process-timeout", + qos: .utility, + attributes: .concurrent) + + package static func runProcess( + executable: String, + arguments: [String], + environment: [String: String], + timeout: TimeInterval) -> String? + { + let process = Process() + process.executableURL = URL(fileURLWithPath: executable) + process.arguments = arguments + + var mergedEnvironment = environment + mergedEnvironment["PATH"] = PathBuilder.effectivePATH( + purposes: [.tty, .nodeTooling], + env: environment, + loginPATH: LoginShellPathCache.shared.current) + process.environment = mergedEnvironment + + let stdout = Pipe() + let stderr = Pipe() + process.standardOutput = stdout + process.standardError = stderr + process.standardInput = nil + + let stdoutCapture = ProcessPipeCapture(pipe: stdout) + let stderrCapture = ProcessPipeCapture(pipe: stderr) + + do { + try process.run() + } catch { + stdoutCapture.stop() + stdout.fileHandleForWriting.closeFile() + stderrCapture.stop() + stderr.fileHandleForWriting.closeFile() + return nil + } + // Process has duplicated these descriptors. Closing the parent's copies lets readers observe EOF as + // soon as the helper exits instead of spending the full drain grace period on every successful call. + stdout.fileHandleForWriting.closeFile() + stderr.fileHandleForWriting.closeFile() + stdoutCapture.start() + stderrCapture.start() + let pid = process.processIdentifier + let processGroup: pid_t? = setpgid(pid, pid) == 0 ? pid : nil + + // Wait synchronously for Foundation to reap the helper. A separate timer owns timeout termination, + // so neither termination-handler scheduling nor stale Process.isRunning state controls normal exits. + let timeoutState = GeminiProcessTimeoutState() + let timeoutTimer = DispatchSource.makeTimerSource(queue: Self.processTimeoutQueue) + timeoutTimer.schedule(deadline: .now() + max(0, timeout)) + timeoutTimer.setEventHandler { + guard process.isRunning else { return } + timeoutState.markTimedOut() + SubprocessRunner.terminateProcess(process, processGroup: processGroup) + } + timeoutTimer.resume() + process.waitUntilExit() + timeoutTimer.cancel() + + if timeoutState.didTimeOut { + stdoutCapture.stop() + stderrCapture.stop() + return nil + } + + let data = stdoutCapture.finishFirstLineSynchronously(timeout: 1) + stderrCapture.stop() + let output = ProcessPipeCapture.decodeUTF8(data) + .trimmingCharacters(in: .whitespacesAndNewlines) + guard process.terminationStatus == 0, !output.isEmpty else { + return nil + } + + return output.components(separatedBy: .newlines).first? + .trimmingCharacters(in: .whitespacesAndNewlines) + } +} + +private final class GeminiProcessTimeoutState: @unchecked Sendable { + private let lock = NSLock() + private var timedOut = false + + var didTimeOut: Bool { + self.lock.lock() + defer { self.lock.unlock() } + return self.timedOut + } + + func markTimedOut() { + self.lock.lock() + self.timedOut = true + self.lock.unlock() + } +} diff --git a/Sources/CodexBarCore/Providers/Grok/GrokAuth.swift b/Sources/CodexBarCore/Providers/Grok/GrokAuth.swift new file mode 100644 index 0000000..9544808 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Grok/GrokAuth.swift @@ -0,0 +1,188 @@ +import Foundation + +public struct GrokCredentials: Sendable { + public let accessToken: String + public let refreshToken: String? + public let scope: String + public let authMode: String? + public let userId: String? + public let email: String? + public let firstName: String? + public let lastName: String? + public let teamId: String? + public let oidcIssuer: String? + public let oidcClientId: String? + public let expiresAt: Date? + public let createTime: Date? + + public init( + accessToken: String, + refreshToken: String?, + scope: String, + authMode: String?, + userId: String?, + email: String?, + firstName: String?, + lastName: String?, + teamId: String?, + oidcIssuer: String?, + oidcClientId: String?, + expiresAt: Date?, + createTime: Date?) + { + self.accessToken = accessToken + self.refreshToken = refreshToken + self.scope = scope + self.authMode = authMode + self.userId = userId + self.email = email + self.firstName = firstName + self.lastName = lastName + self.teamId = teamId + self.oidcIssuer = oidcIssuer + self.oidcClientId = oidcClientId + self.expiresAt = expiresAt + self.createTime = createTime + } + + public var displayName: String? { + let parts = [self.firstName, self.lastName].compactMap { $0?.nilIfEmpty } + guard !parts.isEmpty else { return nil } + return parts.joined(separator: " ") + } + + public var isExpired: Bool { + guard let expiresAt else { return false } + return Date() >= expiresAt + } + + public var loginMethod: String? { + switch self.authMode?.lowercased() { + case "oidc": "SuperGrok" + case "session": "session" + case nil: nil + default: self.authMode + } + } +} + +public enum GrokCredentialsError: LocalizedError, Sendable { + case notFound + case decodeFailed(String) + case missingTokens + + public var errorDescription: String? { + switch self { + case .notFound: + "Grok auth.json not found. Run `grok login` to authenticate." + case let .decodeFailed(message): + "Failed to decode Grok credentials: \(message)" + case .missingTokens: + "Grok auth.json exists but contains no access tokens." + } + } +} + +public enum GrokCredentialsStore { + /// Top-level OIDC scope used by `grok login` for SuperGrok subscribers. + public static let oidcScopePrefix = "https://auth.x.ai::" + /// Legacy/session scope used by older `grok login` flows. + public static let legacySessionScope = "https://accounts.x.ai/sign-in" + + public static func grokHomeURL( + env: [String: String] = ProcessInfo.processInfo.environment, + fileManager: FileManager = .default) -> URL + { + if let custom = env["GROK_HOME"]?.nilIfEmpty { + return URL(fileURLWithPath: (custom as NSString).expandingTildeInPath) + } + let home = fileManager.homeDirectoryForCurrentUser + return home.appendingPathComponent(".grok", isDirectory: true) + } + + public static func authFileURL( + env: [String: String] = ProcessInfo.processInfo.environment, + fileManager: FileManager = .default) -> URL + { + self.grokHomeURL(env: env, fileManager: fileManager).appendingPathComponent("auth.json") + } + + public static func load(env: [String: String] = ProcessInfo.processInfo.environment) throws -> GrokCredentials { + let url = self.authFileURL(env: env) + guard FileManager.default.fileExists(atPath: url.path) else { + throw GrokCredentialsError.notFound + } + let data = try Data(contentsOf: url) + return try self.parse(data: data) + } + + public static func parse(data: Data) throws -> GrokCredentials { + let raw: Any + do { + raw = try JSONSerialization.jsonObject(with: data) + } catch { + throw GrokCredentialsError.decodeFailed(error.localizedDescription) + } + guard let root = raw as? [String: Any] else { + throw GrokCredentialsError.decodeFailed("Invalid JSON (expected object at root)") + } + + // `auth.json` is a map keyed by scope URL. Prefer the OIDC scope (SuperGrok), + // fall back to the legacy session scope. + let preferredEntry = Self.selectPreferredEntry(in: root) + guard let (scope, entry) = preferredEntry else { + throw GrokCredentialsError.missingTokens + } + guard let key = entry["key"] as? String, !key.isEmpty else { + throw GrokCredentialsError.missingTokens + } + + return GrokCredentials( + accessToken: key, + refreshToken: (entry["refresh_token"] as? String)?.nilIfEmpty, + scope: scope, + authMode: (entry["auth_mode"] as? String)?.nilIfEmpty, + userId: (entry["user_id"] as? String)?.nilIfEmpty, + email: (entry["email"] as? String)?.nilIfEmpty, + firstName: (entry["first_name"] as? String)?.nilIfEmpty, + lastName: (entry["last_name"] as? String)?.nilIfEmpty, + teamId: (entry["team_id"] as? String)?.nilIfEmpty, + oidcIssuer: (entry["oidc_issuer"] as? String)?.nilIfEmpty, + oidcClientId: (entry["oidc_client_id"] as? String)?.nilIfEmpty, + expiresAt: Self.parseDate(entry["expires_at"]), + createTime: Self.parseDate(entry["create_time"])) + } + + private static func selectPreferredEntry(in root: [String: Any]) -> (scope: String, entry: [String: Any])? { + var oidcCandidate: (String, [String: Any])? + var legacyCandidate: (String, [String: Any])? + for (scope, value) in root { + guard let entry = value as? [String: Any] else { continue } + // Only accept entries that actually carry a usable bearer token. A + // stale/partial OIDC record (key missing or empty) must not shadow a + // healthy legacy session entry. + guard let key = entry["key"] as? String, !key.isEmpty else { continue } + if scope.hasPrefix(self.oidcScopePrefix) { + oidcCandidate = (scope, entry) + } else if scope == self.legacySessionScope || scope.contains("/sign-in") { + legacyCandidate = (scope, entry) + } + } + return oidcCandidate ?? legacyCandidate + } + + private static func parseDate(_ raw: Any?) -> Date? { + guard let value = raw as? String, !value.isEmpty else { return nil } + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = formatter.date(from: value) { return date } + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: value) + } +} + +extension String { + fileprivate var nilIfEmpty: String? { + self.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty ? nil : self + } +} diff --git a/Sources/CodexBarCore/Providers/Grok/GrokCookieImporter.swift b/Sources/CodexBarCore/Providers/Grok/GrokCookieImporter.swift new file mode 100644 index 0000000..2efd4ee --- /dev/null +++ b/Sources/CodexBarCore/Providers/Grok/GrokCookieImporter.swift @@ -0,0 +1,213 @@ +import Foundation + +#if os(macOS) +import SweetCookieKit + +public enum GrokCookieImporter { + private static let importSessionCacheTTL: TimeInterval = 5 + private static let importSessionCache = ImportSessionCache(ttl: importSessionCacheTTL) + private static let log = CodexBarLog.logger(LogCategories.providers) + private static let cookieClient = BrowserCookieClient() + private static let cookieDomains = ["grok.com"] + private static let cookieImportOrder: BrowserCookieImportOrder = + ProviderDefaults.metadata[.grok]?.browserCookieOrder ?? Browser.defaultImportOrder + + public struct SessionInfo: Sendable { + public let cookies: [HTTPCookie] + public let sourceLabel: String + + public init(cookies: [HTTPCookie], sourceLabel: String) { + self.cookies = cookies + self.sourceLabel = sourceLabel + } + + public var cookieHeader: String { + self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ") + } + } + + public static func importSessions( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + if let cached = self.cachedImportSessions() { + return cached + } + + var sessions: [SessionInfo] = [] + let candidates = self.cookieImportOrder.cookieImportCandidates(using: browserDetection) + for browserSource in candidates { + do { + let perSource = try self.importSessions(from: browserSource, logger: logger) + sessions.append(contentsOf: perSource) + } catch { + BrowserCookieAccessGate.recordIfNeeded(error) + self.emit( + "\(browserSource.displayName) cookie import failed: \(error.localizedDescription)", + logger: logger) + } + } + + guard !sessions.isEmpty else { throw GrokWebBillingError.missingCredentials } + self.storeImportSessions(sessions) + return sessions + } + + public static func importSessions( + from browserSource: Browser, + logger: ((String) -> Void)? = nil) throws -> [SessionInfo] + { + let query = BrowserCookieQuery(domains: self.cookieDomains) + let log: (String) -> Void = { msg in self.emit(msg, logger: logger) } + let sources = try Self.cookieClient.codexBarRecords( + matching: query, + in: browserSource, + logger: log) + + var sessions: [SessionInfo] = [] + let grouped = Dictionary(grouping: sources, by: { $0.store.profile.id }) + let sortedGroups = grouped.values.sorted { lhs, rhs in + self.mergedLabel(for: lhs) < self.mergedLabel(for: rhs) + } + + for group in sortedGroups where !group.isEmpty { + let label = self.mergedLabel(for: group) + let mergedRecords = self.mergeRecords(group) + guard mergedRecords.contains(where: { $0.name == "sso" || $0.name == "sso-rw" }) else { continue } + let httpCookies = BrowserCookieClient.makeHTTPCookies(mergedRecords, origin: query.origin) + guard !httpCookies.isEmpty else { continue } + log("Found Grok session cookies in \(label)") + sessions.append(SessionInfo(cookies: httpCookies, sourceLabel: label)) + } + return sessions + } + + public static func importSession( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) throws -> SessionInfo + { + let sessions = try self.importSessions(browserDetection: browserDetection, logger: logger) + guard let first = sessions.first else { throw GrokWebBillingError.missingCredentials } + return first + } + + public static func hasSession( + browserDetection: BrowserDetection = BrowserDetection(), + logger: ((String) -> Void)? = nil) -> Bool + { + do { + _ = try self.importSession(browserDetection: browserDetection, logger: logger) + return true + } catch { + return false + } + } + + static func invalidateImportSessionCache() { + self.importSessionCache.invalidate() + } + + private static func emit(_ message: String, logger: ((String) -> Void)?) { + logger?("[grok-cookie] \(message)") + self.log.debug("\(message)") + } + + private static func cachedImportSessions(now: Date = Date()) -> [SessionInfo]? { + self.importSessionCache.load(now: now) + } + + private static func storeImportSessions(_ sessions: [SessionInfo], now: Date = Date()) { + self.importSessionCache.store(sessions, now: now) + } + + private static func mergedLabel(for sources: [BrowserCookieStoreRecords]) -> String { + guard let base = sources.map(\.label).min() else { return "Unknown" } + if base.hasSuffix(" (Network)") { + return String(base.dropLast(" (Network)".count)) + } + return base + } + + private static func mergeRecords(_ sources: [BrowserCookieStoreRecords]) -> [BrowserCookieRecord] { + let sortedSources = sources.sorted { lhs, rhs in + self.storePriority(lhs.store.kind) < self.storePriority(rhs.store.kind) + } + var mergedByKey: [String: BrowserCookieRecord] = [:] + for source in sortedSources { + for record in source.records { + let key = self.recordKey(record) + if let existing = mergedByKey[key] { + if self.shouldReplace(existing: existing, candidate: record) { + mergedByKey[key] = record + } + } else { + mergedByKey[key] = record + } + } + } + return Array(mergedByKey.values) + } + + private static func storePriority(_ kind: BrowserCookieStoreKind) -> Int { + switch kind { + case .network: 0 + case .primary: 1 + case .safari: 2 + } + } + + private static func recordKey(_ record: BrowserCookieRecord) -> String { + "\(record.name)|\(record.domain)|\(record.path)" + } + + private static func shouldReplace(existing: BrowserCookieRecord, candidate: BrowserCookieRecord) -> Bool { + switch (existing.expires, candidate.expires) { + case let (lhs?, rhs?): rhs > lhs + case (nil, .some): true + case (.some, nil): false + case (nil, nil): false + } + } + + private final class ImportSessionCache: @unchecked Sendable { + private let ttl: TimeInterval + private let lock = NSLock() + private var entry: (sessions: [SessionInfo], expiresAt: Date)? + + init(ttl: TimeInterval) { + self.ttl = ttl + } + + func load(now: Date) -> [SessionInfo]? { + self.lock.lock() + defer { self.lock.unlock() } + guard let entry = self.entry, entry.expiresAt > now else { + self.entry = nil + return nil + } + return entry.sessions + } + + func store(_ sessions: [SessionInfo], now: Date) { + self.lock.lock() + self.entry = (sessions, now.addingTimeInterval(self.ttl)) + self.lock.unlock() + } + + func invalidate() { + self.lock.lock() + self.entry = nil + self.lock.unlock() + } + } +} +#else +public enum GrokCookieImporter { + public static func hasSession( + browserDetection _: BrowserDetection = BrowserDetection(), + logger _: ((String) -> Void)? = nil) -> Bool + { + false + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Grok/GrokLocalSessionScanner.swift b/Sources/CodexBarCore/Providers/Grok/GrokLocalSessionScanner.swift new file mode 100644 index 0000000..c5b1afd --- /dev/null +++ b/Sources/CodexBarCore/Providers/Grok/GrokLocalSessionScanner.swift @@ -0,0 +1,100 @@ +import Foundation + +/// Aggregated stats from local `~/.grok/sessions/**/signals.json` files. +/// Used as a local fallback view when the JSON-RPC billing call is unavailable. +public struct GrokLocalSessionSummary: Sendable { + public let sessionCount: Int + public let totalTokens: Int + public let lastSessionAt: Date? + public let primaryModel: String? + public let models: [String] + + public init( + sessionCount: Int, + totalTokens: Int, + lastSessionAt: Date?, + primaryModel: String?, + models: [String]) + { + self.sessionCount = sessionCount + self.totalTokens = totalTokens + self.lastSessionAt = lastSessionAt + self.primaryModel = primaryModel + self.models = models + } +} + +public enum GrokLocalSessionScanner { + public static let defaultLookbackDays = 30 + + /// Walk `~/.grok/sessions///signals.json` and aggregate stats. + public static func summarize( + env: [String: String] = ProcessInfo.processInfo.environment, + fileManager: FileManager = .default, + lookbackDays: Int = defaultLookbackDays, + now: Date = .init()) -> GrokLocalSessionSummary + { + let root = GrokCredentialsStore.grokHomeURL(env: env, fileManager: fileManager) + .appendingPathComponent("sessions", isDirectory: true) + guard let rootEnum = fileManager.enumerator( + at: root, + includingPropertiesForKeys: [.contentModificationDateKey, .isDirectoryKey], + options: [.skipsHiddenFiles]) + else { + return GrokLocalSessionSummary( + sessionCount: 0, + totalTokens: 0, + lastSessionAt: nil, + primaryModel: nil, + models: []) + } + + let lookbackCutoff = Calendar.current.date(byAdding: .day, value: -lookbackDays, to: now) ?? now + var sessionCount = 0 + var totalTokens = 0 + var lastSessionAt: Date? + var modelCounts: [String: Int] = [:] + + while let url = rootEnum.nextObject() as? URL { + guard url.lastPathComponent == "signals.json" else { continue } + let attrs = try? url.resourceValues(forKeys: [.contentModificationDateKey]) + let mtime = attrs?.contentModificationDate ?? Date.distantPast + guard mtime >= lookbackCutoff else { continue } + + guard let data = try? Data(contentsOf: url), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { continue } + + sessionCount += 1 + let beforeCompaction = (json["totalTokensBeforeCompaction"] as? Int) ?? 0 + let contextUsed = (json["contextTokensUsed"] as? Int) ?? 0 + totalTokens += beforeCompaction + contextUsed + + if mtime > (lastSessionAt ?? Date.distantPast) { + lastSessionAt = mtime + } + + if let primary = (json["primaryModelId"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines), + !primary.isEmpty + { + modelCounts[primary, default: 0] += 1 + } + if let models = json["modelsUsed"] as? [String] { + for model in models { + let trimmed = model.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { + modelCounts[trimmed, default: 0] += 1 + } + } + } + } + + let sortedModels = modelCounts.sorted { $0.value > $1.value }.map(\.key) + return GrokLocalSessionSummary( + sessionCount: sessionCount, + totalTokens: totalTokens, + lastSessionAt: lastSessionAt, + primaryModel: sortedModels.first, + models: sortedModels) + } +} diff --git a/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift new file mode 100644 index 0000000..6aaf76f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift @@ -0,0 +1,223 @@ +import Foundation + +public enum GrokProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .grok, + metadata: ProviderMetadata( + id: .grok, + displayName: "Grok", + sessionLabel: "Credits", + weeklyLabel: "On-demand", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Grok usage", + cliName: "grok", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: ProviderBrowserCookieDefaults.grokCookieImportOrder, + dashboardURL: "https://grok.com/?_s=usage", + changelogURL: "https://x.ai/news", + statusPageURL: nil, + statusLinkURL: "https://status.x.ai"), + branding: ProviderBranding( + iconStyle: .grok, + iconResourceName: "ProviderIcon-grok", + color: ProviderColor(red: 16 / 255, green: 163 / 255, blue: 127 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Grok cost summary is not supported yet." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .cli, .web], + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), + cli: ProviderCLIConfig( + name: "grok", + versionDetector: { _ in GrokStatusProbe.detectVersion() })) + } + + private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] { + switch context.sourceMode { + case .auto: + [GrokCLIFetchStrategy(), GrokWebFetchStrategy()] + case .cli: + [GrokCLIFetchStrategy()] + case .web: + [GrokWebFetchStrategy()] + case .api, .oauth: + [] + } + } + + /// Returns a contextual label for Grok's primary usage bar ("Weekly" or "Monthly"). + /// Prefer the billing period duration when available; fall back to reset distance for + /// web billing payloads that expose only a reset timestamp. + public static func primaryLabel(window: RateWindow?, now: Date = .now) -> String? { + if let minutes = window?.windowMinutes { + return self.primaryLabel(duration: TimeInterval(minutes) * 60) + } + return self.primaryLabel(resetsAt: window?.resetsAt, now: now) + } + + public static func primaryLabel(resetsAt: Date?, now: Date = .now) -> String? { + guard let resetsAt else { return nil } + return self.primaryLabel(duration: resetsAt.timeIntervalSince(now)) + } + + private static func primaryLabel(duration seconds: TimeInterval) -> String? { + guard seconds > 3600 else { return nil } + let days = Int((seconds / 86400).rounded(.toNearestOrAwayFromZero)) + if (4...12).contains(days) { return "Weekly" } + if (20...45).contains(days) { return "Monthly" } + return nil + } +} + +struct GrokCLIFetchStrategy: ProviderFetchStrategy { + let id: String = "grok.cli" + let kind: ProviderFetchKind = .cli + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + BinaryLocator.resolveGrokBinary(env: context.env) != nil + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let probe = GrokStatusProbe() + let snap = try await probe.fetch(env: context.env) + return self.makeResult( + usage: snap.toUsageSnapshot(), + sourceLabel: "grok-cli") + } + + func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool { + context.sourceMode == .auto + } +} + +struct GrokWebFetchStrategy: ProviderFetchStrategy { + let id: String = "grok.web" + let kind: ProviderFetchKind = .web + + static func canImportBrowserCookies(runtime: ProviderRuntime, env: [String: String]) -> Bool { + runtime == .app || env["CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT"] == "1" + } + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + #if os(macOS) + if Self.canImportBrowserCookies(runtime: context.runtime, env: context.env), + GrokCookieImporter.hasSession(browserDetection: context.browserDetection) + { + return true + } + #endif + return FileManager.default.fileExists(atPath: GrokCredentialsStore.authFileURL(env: context.env).path) + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let (webBilling, sourceLabel, authenticatedByAuthFile) = try await self.fetchWebBilling(context: context) + let credentials = Self.credentialsForWebBillingSnapshot( + credentials: try? GrokCredentialsStore.load(env: context.env), + authenticatedByAuthFile: authenticatedByAuthFile) + let snapshot = GrokUsageSnapshot( + billing: nil, + webBilling: webBilling, + credentials: GrokStatusProbe.credentialsForSnapshot( + credentials: credentials, + billing: nil, + webBilling: webBilling), + localSummary: GrokLocalSessionScanner.summarize(env: context.env), + cliVersion: GrokStatusProbe.detectVersion(env: context.env), + updatedAt: Date()) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: sourceLabel) + } + + private func fetchWebBilling(context: ProviderFetchContext) async throws -> ( + snapshot: GrokWebBillingSnapshot, + sourceLabel: String, + authenticatedByAuthFile: Bool) + { + let credentialsResult: Result = Result { + try GrokCredentialsStore.load(env: context.env) + } + let browserCredentials = try? credentialsResult.get() + + #if os(macOS) + if Self.canImportBrowserCookies(runtime: context.runtime, env: context.env) { + var lastCookieError: Error? + do { + let sessions = try GrokCookieImporter.importSessions(browserDetection: context.browserDetection) + let (snapshot, sourceLabel) = try await Self.fetchFirstValidCookieSession( + sessions, + credentials: browserCredentials) + return (snapshot, sourceLabel, false) + } catch { + lastCookieError = error + } + if browserCredentials == nil { + if FileManager.default.fileExists( + atPath: GrokCredentialsStore.authFileURL(env: context.env).path) + { + _ = try credentialsResult.get() + } + throw lastCookieError ?? GrokWebBillingError.missingCredentials + } + } + #endif + + let authCredentials = try credentialsResult.get() + guard !authCredentials.isExpired else { + throw GrokWebBillingError.missingCredentials + } + let snapshot = try await GrokWebBillingFetcher.fetch(credentials: authCredentials) + return (snapshot, "grok-web", true) + } + + static func credentialsForWebBillingSnapshot( + credentials: GrokCredentials?, + authenticatedByAuthFile: Bool) -> GrokCredentials? + { + authenticatedByAuthFile ? credentials : nil + } + + #if os(macOS) + static func fetchFirstValidCookieSession( + _ sessions: [GrokCookieImporter.SessionInfo], + credentials: GrokCredentials? = nil, + fetch: ((String, GrokCredentials?) async throws -> GrokWebBillingSnapshot)? = nil) async throws + -> (GrokWebBillingSnapshot, String) + { + let fetchSnapshot = fetch ?? { cookieHeader, credentials in + try await GrokWebBillingFetcher.fetch( + cookieHeader: cookieHeader, + credentials: credentials) + } + var lastError: Error? + for session in sessions { + for authCredentials in Self.cookieAuthAttempts(credentials: credentials) { + do { + let snapshot = try await fetchSnapshot(session.cookieHeader, authCredentials) + return (snapshot, session.sourceLabel) + } catch { + lastError = error + } + } + } + throw lastError ?? GrokWebBillingError.missingCredentials + } + + static func cookieAuthAttempts(credentials: GrokCredentials?) -> [GrokCredentials?] { + guard let credentials, !credentials.isExpired else { return [nil] } + return [credentials, nil] + } + #endif + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} diff --git a/Sources/CodexBarCore/Providers/Grok/GrokRPCClient.swift b/Sources/CodexBarCore/Providers/Grok/GrokRPCClient.swift new file mode 100644 index 0000000..a0feea7 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Grok/GrokRPCClient.swift @@ -0,0 +1,371 @@ +import Foundation + +/// JSON-RPC client for `grok agent stdio` (ACP protocol). +/// +/// The protocol mirrors Codex's app-server (newline-delimited JSON-RPC 2.0 over stdin/stdout), +/// but uses `protocolVersion`/`clientCapabilities` for the `initialize` call instead of +/// `clientInfo`. Billing is fetched via the `x.ai/billing` extension method. +final class GrokRPCClient: @unchecked Sendable { + private static let log = CodexBarLog.logger(LogCategories.grok) + + private let process = Process() + private let stdinPipe = Pipe() + private let stdoutPipe = Pipe() + private let stderrPipe = Pipe() + private let initializeTimeoutSeconds: TimeInterval + private let requestTimeoutSeconds: TimeInterval + private var nextID: Int = 1 + private let stdoutLineStream: AsyncStream + private let stdoutLineContinuation: AsyncStream.Continuation + + init( + executable: String = "grok", + arguments: [String] = ["agent", "stdio"], + environment: [String: String] = ProcessInfo.processInfo.environment, + initializeTimeoutSeconds: TimeInterval = 4.0, + requestTimeoutSeconds: TimeInterval = 3.0) throws + { + self.initializeTimeoutSeconds = initializeTimeoutSeconds + self.requestTimeoutSeconds = requestTimeoutSeconds + var stdoutContinuation: AsyncStream.Continuation! + self.stdoutLineStream = AsyncStream { continuation in + stdoutContinuation = continuation + } + self.stdoutLineContinuation = stdoutContinuation + + let resolvedExec = BinaryLocator.resolveGrokBinary(env: environment) + ?? TTYCommandRunner.which(executable) + + guard let resolvedExec else { + Self.log.warning("Grok RPC binary not found", metadata: ["binary": executable]) + throw GrokRPCError.binaryNotFound + } + + var env = environment + env["PATH"] = PathBuilder.effectivePATH(purposes: [.rpc], env: env) + + self.process.environment = env + self.process.executableURL = URL(fileURLWithPath: "/usr/bin/env") + self.process.arguments = [resolvedExec] + arguments + self.process.standardInput = self.stdinPipe + self.process.standardOutput = self.stdoutPipe + self.process.standardError = self.stderrPipe + + do { + try self.process.run() + Self.log.debug("Grok RPC started", metadata: ["binary": resolvedExec]) + } catch { + Self.log.warning("Grok RPC failed to start", metadata: ["error": error.localizedDescription]) + throw GrokRPCError.startFailed(error.localizedDescription) + } + + let stdoutHandle = self.stdoutPipe.fileHandleForReading + let stdoutLineContinuation = self.stdoutLineContinuation + let stdoutBuffer = LineBuffer() + stdoutHandle.readabilityHandler = { handle in + let data = handle.availableData + if data.isEmpty { + handle.readabilityHandler = nil + stdoutLineContinuation.finish() + return + } + let lines = stdoutBuffer.appendAndDrainLines(data) + for lineData in lines { + stdoutLineContinuation.yield(lineData) + } + } + + let stderrHandle = self.stderrPipe.fileHandleForReading + stderrHandle.readabilityHandler = { handle in + let data = handle.availableData + if data.isEmpty { + handle.readabilityHandler = nil + return + } + guard let text = String(data: data, encoding: .utf8), !text.isEmpty else { return } + for line in text.split(whereSeparator: \.isNewline) { + #if !os(Linux) + fputs("[grok stderr] \(line)\n", stderr) + #endif + } + } + } + + deinit { + self.shutdown() + } + + func initialize() async throws { + let params: [String: Any] = [ + "protocolVersion": "1", + "clientCapabilities": [ + "fs": ["readTextFile": false, "writeTextFile": false], + "terminal": false, + ], + ] + _ = try await self.request( + method: "initialize", + params: params, + timeout: self.initializeTimeoutSeconds) + } + + /// Calls `x.ai/billing` and returns the decoded response. + func fetchBilling() async throws -> GrokBillingResponse { + let message = try await self.request(method: "x.ai/billing", params: [:]) + return try self.decodeResult(from: message) + } + + func shutdown() { + if self.process.isRunning { + Self.log.debug("Grok RPC stopping") + self.process.terminate() + } + } + + // MARK: - JSON-RPC plumbing (mirrors CodexRPCClient) + + private struct SendableJSONMessage: @unchecked Sendable { + let value: [String: Any] + } + + private func request( + method: String, + params: [String: Any]? = nil, + timeout: TimeInterval? = nil) async throws -> [String: Any] + { + let id = self.nextID + self.nextID += 1 + try self.sendRequest(id: id, method: method, params: params) + + let resolvedTimeout = timeout ?? self.requestTimeoutSeconds + let wrapped = try await self.withTimeout(seconds: resolvedTimeout, method: method) { + while true { + let message = try await self.readNextMessage() + // Skip notifications (no id) or unrelated responses. + if message["id"] == nil { continue } + guard let messageID = self.jsonID(message["id"]), messageID == id else { continue } + if let error = message["error"] as? [String: Any] { + let messageText = (error["message"] as? String) ?? "unknown JSON-RPC error" + throw GrokRPCError.requestFailed(messageText) + } + return SendableJSONMessage(value: message) + } + } + return wrapped.value + } + + private func withTimeout( + seconds: TimeInterval, + method: String, + body: @escaping @Sendable () async throws -> T) async throws -> T + { + try await withThrowingTaskGroup(of: T.self) { group in + group.addTask { try await body() } + group.addTask { [weak self] in + try await Task.sleep(for: .seconds(seconds)) + self?.terminateProcessForTimeout(method: method) + throw GrokRPCError.timeout(method: method) + } + do { + guard let result = try await group.next() else { + throw GrokRPCError.timeout(method: method) + } + group.cancelAll() + return result + } catch { + group.cancelAll() + throw error + } + } + } + + private func terminateProcessForTimeout(method: String) { + if self.process.isRunning { + Self.log.warning("Grok RPC timed out on `\(method)`; terminating process") + self.process.terminate() + } + } + + private func sendRequest(id: Int, method: String, params: [String: Any]?) throws { + let paramsValue: Any = params ?? [:] + let payload: [String: Any] = [ + "jsonrpc": "2.0", + "id": id, + "method": method, + "params": paramsValue, + ] + try self.sendPayload(payload) + } + + private func sendPayload(_ payload: [String: Any]) throws { + let raw = try JSONSerialization.data(withJSONObject: payload) + // Foundation's JSONSerialization escapes "/" as "\/" by default. Grok's + // ACP server treats the escaped form as a *different* method name (it does + // not unescape before lookup), so `x.ai/billing` becomes "Method not found" + // when sent as `x.ai\/billing`. Re-encode without slash escapes to match + // the on-the-wire shape the grok agent expects. + let unescaped = String(data: raw, encoding: .utf8)? + .replacingOccurrences(of: "\\/", with: "/") + let data = unescaped.flatMap { $0.data(using: .utf8) } ?? raw + if let preview = String(data: data.prefix(200), encoding: .utf8) { + Self.log.debug("grok rpc -> \(preview)") + } + self.stdinPipe.fileHandleForWriting.write(data) + self.stdinPipe.fileHandleForWriting.write(Data([0x0A])) + } + + private func readNextMessage() async throws -> [String: Any] { + for await lineData in self.stdoutLineStream { + if lineData.isEmpty { continue } + if let preview = String(data: lineData.prefix(300), encoding: .utf8) { + Self.log.debug("grok rpc <- \(preview)") + } + if let json = try? JSONSerialization.jsonObject(with: lineData) as? [String: Any] { + return json + } + } + throw GrokRPCError.malformed("grok agent stdio closed stdout") + } + + private func decodeResult(from message: [String: Any]) throws -> T { + guard let result = message["result"] else { + throw GrokRPCError.malformed("missing result field") + } + let data = try JSONSerialization.data(withJSONObject: result) + let decoder = JSONDecoder() + return try decoder.decode(T.self, from: data) + } + + private func jsonID(_ value: Any?) -> Int? { + switch value { + case let int as Int: int + case let number as NSNumber: number.intValue + default: nil + } + } + + private final class LineBuffer: @unchecked Sendable { + private var buffer = Data() + private let lock = NSLock() + + func appendAndDrainLines(_ data: Data) -> [Data] { + self.lock.lock() + defer { lock.unlock() } + self.buffer.append(data) + var out: [Data] = [] + while let newline = self.buffer.firstIndex(of: 0x0A) { + let lineData = Data(self.buffer[.. }` in the wire format. +public struct GrokBillingResponse: Codable, Sendable { + public let billingCycle: GrokBillingCycle? + public let monthlyLimit: GrokCent? + public let onDemandCap: GrokCent? + public let onDemandEnabled: Bool? + public let disabledByConfig: Bool? + public let usage: GrokBillingUsage? + + private enum CodingKeys: String, CodingKey { + case billingCycle + case monthlyLimit + case onDemandCap + case onDemandEnabled = "on_demand_enabled" + case disabledByConfig + case usage + } +} + +public struct GrokBillingCycle: Codable, Sendable { + public let billingPeriodStart: String? + public let billingPeriodEnd: String? +} + +public struct GrokBillingUsage: Codable, Sendable { + public let includedUsed: GrokCent? + public let onDemandUsed: GrokCent? + public let totalUsed: GrokCent? +} + +public struct GrokCent: Codable, Sendable { + public let val: Int? +} + +extension GrokBillingResponse { + /// Convenience accessor: monthly usage as a 0–100 percent. + public var monthlyUsedPercent: Double? { + guard let limit = self.monthlyLimit?.val, limit > 0, + let used = self.usage?.totalUsed?.val + else { return nil } + return min(100.0, max(0.0, Double(used) / Double(limit) * 100.0)) + } + + public var billingPeriodEndDate: Date? { + guard let raw = self.billingCycle?.billingPeriodEnd else { return nil } + return GrokBillingResponse.parseISO8601(raw) + } + + public var billingPeriodStartDate: Date? { + guard let raw = self.billingCycle?.billingPeriodStart else { return nil } + return GrokBillingResponse.parseISO8601(raw) + } + + public var billingPeriodMinutes: Int? { + guard let start = self.billingPeriodStartDate, + let end = self.billingPeriodEndDate, + end > start + else { return nil } + return Int(end.timeIntervalSince(start) / 60) + } + + private static func parseISO8601(_ raw: String) -> Date? { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = formatter.date(from: raw) { return date } + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: raw) + } +} diff --git a/Sources/CodexBarCore/Providers/Grok/GrokStatusProbe.swift b/Sources/CodexBarCore/Providers/Grok/GrokStatusProbe.swift new file mode 100644 index 0000000..5473d24 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Grok/GrokStatusProbe.swift @@ -0,0 +1,147 @@ +import Foundation + +public struct GrokUsageSnapshot: Sendable { + public let billing: GrokBillingResponse? + public let webBilling: GrokWebBillingSnapshot? + public let credentials: GrokCredentials? + public let localSummary: GrokLocalSessionSummary? + public let cliVersion: String? + public let updatedAt: Date + + public init( + billing: GrokBillingResponse?, + webBilling: GrokWebBillingSnapshot? = nil, + credentials: GrokCredentials?, + localSummary: GrokLocalSessionSummary?, + cliVersion: String?, + updatedAt: Date) + { + self.billing = billing + self.webBilling = webBilling + self.credentials = credentials + self.localSummary = localSummary + self.cliVersion = cliVersion + self.updatedAt = updatedAt + } + + public func toUsageSnapshot() -> UsageSnapshot { + // Primary window: credit usage (against included limit) from the CLI RPC, + // falling back to the web billing RPC used by grok.com when the agent surface lacks billing. + var primary: RateWindow? + if let billing, + let percent = billing.monthlyUsedPercent + { + primary = RateWindow( + usedPercent: percent, + windowMinutes: billing.billingPeriodMinutes, + resetsAt: billing.billingPeriodEndDate, + resetDescription: nil) + } else if let webBilling, + let percent = webBilling.usedPercent + { + primary = RateWindow( + usedPercent: percent, + windowMinutes: nil, + resetsAt: webBilling.resetsAt, + resetDescription: nil) + } + + let identity = ProviderIdentitySnapshot( + providerID: .grok, + accountEmail: self.credentials?.email, + accountOrganization: self.credentials?.teamId, + loginMethod: self.credentials?.loginMethod) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + updatedAt: self.updatedAt, + identity: identity) + } +} + +public struct GrokStatusProbe: Sendable { + public init() {} + + public static func detectVersion(env: [String: String] = ProcessInfo.processInfo.environment) -> String? { + guard let binary = BinaryLocator.resolveGrokBinary(env: env) else { return nil } + guard let output = ProviderVersionDetector.run( + path: binary, + args: ["--version"], + environment: env, + mergeStandardError: true) + else { return nil } + // Output is like "grok 0.1.210 (8b63e9068c)" — strip the leading "grok " so + // callers can prefix the CLI name themselves without duplicating it. + let withoutPrefix = output.replacingOccurrences( + of: #"^grok\s+"#, + with: "", + options: [.regularExpression]) + .trimmingCharacters(in: .whitespacesAndNewlines) + return withoutPrefix.isEmpty ? nil : withoutPrefix + } + + public func fetch(env: [String: String] = ProcessInfo.processInfo.environment) async throws -> GrokUsageSnapshot { + // Credentials are optional: we still show identity-less state if the user + // hasn't logged in, with a clear hint via the RPC error. + let credentials = try? GrokCredentialsStore.load(env: env) + + var billing: GrokBillingResponse? + var rpcError: Error? + do { + let client = try GrokRPCClient(environment: env) + defer { client.shutdown() } + try await client.initialize() + billing = try await client.fetchBilling() + } catch { + rpcError = error + } + + // Local fallback summary always succeeds (empty if no sessions yet). + let localSummary = GrokLocalSessionScanner.summarize(env: env) + let cliVersion = Self.detectVersion(env: env) + + // `localSummary` is *not* currently projected into a visible RateWindow or + // identity field, so a stale `~/.grok/sessions/` directory must not + // suppress the auth-required hint. CLI-only fetches need a billing + // response; the provider pipeline owns the separate web fallback. + if billing == nil { + throw rpcError ?? GrokRPCError.notAuthenticated + } + + return GrokUsageSnapshot( + billing: billing, + webBilling: nil, + credentials: Self.credentialsForSnapshot( + credentials: credentials, + billing: billing, + webBilling: nil), + localSummary: localSummary, + cliVersion: cliVersion, + updatedAt: Date()) + } + + static func credentialsForSnapshot( + credentials: GrokCredentials?, + billing: GrokBillingResponse?, + webBilling: GrokWebBillingSnapshot? = nil) -> GrokCredentials? + { + // If remote usage succeeded, xAI accepted auth and the local + // identity is still useful even when the persisted expires_at is stale. + if billing != nil || webBilling != nil { return credentials } + return credentials.flatMap { $0.isExpired ? nil : $0 } + } + + static func shouldSurfaceRemoteAuthError(_ error: Error?) -> Bool { + guard let error = error as? GrokWebBillingError else { return false } + switch error { + case let .requestFailed(status, _): + return status == 401 || status == 403 + case let .rpcFailed(status, _): + return status == 16 + case .missingCredentials, .emptyResponse, .invalidResponse, .parseFailed: + return false + } + } +} diff --git a/Sources/CodexBarCore/Providers/Grok/GrokWebBillingFetcher.swift b/Sources/CodexBarCore/Providers/Grok/GrokWebBillingFetcher.swift new file mode 100644 index 0000000..4ab67ac --- /dev/null +++ b/Sources/CodexBarCore/Providers/Grok/GrokWebBillingFetcher.swift @@ -0,0 +1,433 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct GrokWebBillingSnapshot: Sendable, Equatable { + public let usedPercent: Double? + public let resetsAt: Date? + + public init(usedPercent: Double?, resetsAt: Date?) { + self.usedPercent = usedPercent + self.resetsAt = resetsAt + } +} + +public enum GrokWebBillingError: LocalizedError, Sendable { + case missingCredentials + case emptyResponse + case invalidResponse + case requestFailed(Int, String) + case rpcFailed(Int, String) + case parseFailed + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Grok web billing requires a signed-in grok.com browser session or `grok login`." + case .emptyResponse: + "Grok web billing returned no protobuf payload." + case .invalidResponse: + "Grok web billing returned an invalid response." + case let .requestFailed(status, body): + if status == 401 || status == 403 { + Self.reauthMessage + } else { + "Grok web billing request failed with HTTP \(status): \(body)" + } + case let .rpcFailed(status, message): + if Self.isAuthenticationFailure(status: status, message: message) { + Self.reauthMessage + } else { + "Grok web billing RPC failed with status \(status): \(message)" + } + case .parseFailed: + "Could not parse Grok web billing usage." + } + } + + private static let reauthMessage = + "Grok web billing rejected credentials. Sign in to grok.com in Chrome or run `grok login` to refresh xAI auth." + + static func isAuthenticationFailure(status: Int, message: String) -> Bool { + if status == 16 { return true } + guard status == 7 else { return false } + let lower = message.lowercased() + return lower.contains("bad-credentials") || + lower.contains("unauthenticated") || + (lower.contains("oauth2") && lower.contains("could not be validated")) || + (lower.contains("access token") && + (lower.contains("invalid") || + lower.contains("expired") || + lower.contains("could not be validated"))) + } +} + +public enum GrokWebBillingFetcher { + public static let defaultEndpoint = + URL(string: "https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")! + private static let requestTimeoutSeconds: TimeInterval = 15 + + public static func fetch( + credentials: GrokCredentials, + session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + endpoint: URL = Self.defaultEndpoint) async throws -> GrokWebBillingSnapshot + { + try await self.fetch( + authorizationHeader: "Bearer \(credentials.accessToken)", + cookieHeader: nil, + transport: transport, + endpoint: endpoint) + } + + public static func fetch( + cookieHeader: String, + session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + endpoint: URL = Self.defaultEndpoint) async throws -> GrokWebBillingSnapshot + { + try await self.fetch( + cookieHeader: cookieHeader, + credentials: nil, + session: transport, + endpoint: endpoint) + } + + public static func fetch( + cookieHeader: String, + credentials: GrokCredentials?, + session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + endpoint: URL = Self.defaultEndpoint) async throws -> GrokWebBillingSnapshot + { + let authorizationHeader = credentials.flatMap { credential in + credential.isExpired ? nil : "Bearer \(credential.accessToken)" + } + return try await self.fetch( + authorizationHeader: authorizationHeader, + cookieHeader: cookieHeader, + transport: transport, + endpoint: endpoint) + } + + private static func fetch( + authorizationHeader: String?, + cookieHeader: String?, + transport: any ProviderHTTPTransport, + endpoint: URL) async throws -> GrokWebBillingSnapshot + { + do { + return try await self.fetchOnce( + authorizationHeader: authorizationHeader, + cookieHeader: cookieHeader, + transport: transport, + endpoint: endpoint) + } catch where self.shouldRetry(error) { + return try await self.fetchOnce( + authorizationHeader: authorizationHeader, + cookieHeader: cookieHeader, + transport: transport, + endpoint: endpoint) + } + } + + private static func fetchOnce( + authorizationHeader: String?, + cookieHeader: String?, + transport: any ProviderHTTPTransport, + endpoint: URL) async throws -> GrokWebBillingSnapshot + { + var request = URLRequest(url: endpoint) + request.httpMethod = "POST" + request.timeoutInterval = Self.requestTimeoutSeconds + request.httpBody = Data([0x00, 0x00, 0x00, 0x00, 0x00]) + if let authorizationHeader { + request.setValue(authorizationHeader, forHTTPHeaderField: "Authorization") + } + if let cookieHeader { + request.setValue(cookieHeader, forHTTPHeaderField: "Cookie") + } + request.setValue("https://grok.com", forHTTPHeaderField: "Origin") + request.setValue("https://grok.com/?_s=usage", forHTTPHeaderField: "Referer") + request.setValue("*/*", forHTTPHeaderField: "Accept") + request.setValue("application/grpc-web+proto", forHTTPHeaderField: "Content-Type") + request.setValue("1", forHTTPHeaderField: "x-grpc-web") + request.setValue("connect-es/2.1.1", forHTTPHeaderField: "x-user-agent") + request.setValue("CodexBar", forHTTPHeaderField: "User-Agent") + + let response: ProviderHTTPResponse + do { + response = try await transport.response(for: request) + } catch let error as URLError where error.code == .badServerResponse { + throw GrokWebBillingError.invalidResponse + } catch { + throw error + } + guard response.statusCode == 200 else { + let body = String(data: response.data.prefix(400), encoding: .utf8) ?? "" + throw GrokWebBillingError.requestFailed(response.statusCode, body) + } + try Self.validateGRPCStatusFields(Self.grpcHeaderFields(from: response.response.allHeaderFields)) + try Self.validateGRPCWebTrailers(response.data) + + return try Self.parseGRPCWebResponse(response.data) + } + + private static func shouldRetry(_ error: Error) -> Bool { + if let urlError = error as? URLError { + return urlError.code == .timedOut || urlError.code == .networkConnectionLost + } + if case let GrokWebBillingError.requestFailed(status, body) = error { + if [408, 502, 503, 504].contains(status) { return true } + return body.localizedCaseInsensitiveContains("timeout") + || body.localizedCaseInsensitiveContains("deadline") + } + guard case let GrokWebBillingError.rpcFailed(status, message) = error else { return false } + if status == 4 { return true } + guard status == 1 else { return false } + return message.localizedCaseInsensitiveContains("timeout") + || message.localizedCaseInsensitiveContains("deadline") + || message.localizedCaseInsensitiveContains("expired") + } + + static func parseGRPCWebResponse(_ data: Data, now: Date = Date()) throws -> GrokWebBillingSnapshot { + var payloads = Self.grpcWebDataFrames(from: data) + if payloads.isEmpty, Self.looksLikeProtobufPayload(data) { + payloads = [data] + } + guard !payloads.isEmpty else { throw GrokWebBillingError.emptyResponse } + + var scan = ProtobufScan() + for payload in payloads { + scan.merge(Self.scanProtobuf(payload, depth: 0)) + } + + let parsedPercent = scan.fixed32Fields + .filter { field in + field.path.last == 1 && field.value.isFinite && field.value >= 0 && field.value <= 100 + } + .min { lhs, rhs in + lhs.path.count == rhs.path.count ? lhs.order < rhs.order : lhs.path.count < rhs.path.count + } + .map { Double($0.value) } + + let resetFields = scan.varintFields.compactMap { field -> (path: [UInt64], date: Date)? in + let raw = field.value + guard raw >= 1_700_000_000, raw <= 2_100_000_000 else { return nil } + return (field.path, Date(timeIntervalSince1970: TimeInterval(raw))) + } + let futureResetFields = resetFields.filter { $0.date > now } + let reset = futureResetFields + .filter { $0.path == [1, 5, 1] } + .map(\.date) + .min() ?? futureResetFields + .map(\.date) + .min() + + let hasUsagePeriod = scan.varintFields.contains { field in + field.path.starts(with: [1, 6]) || + (field.path == [1, 8, 1] && (field.value == 1 || field.value == 2)) + } + let noUsageYet = parsedPercent == nil && + scan.fixed32Fields.isEmpty && + reset != nil && + hasUsagePeriod + guard let percent = parsedPercent ?? (noUsageYet ? 0 : nil) else { + throw GrokWebBillingError.parseFailed + } + return GrokWebBillingSnapshot(usedPercent: percent, resetsAt: reset) + } + + static func looksLikeProtobufPayload(_ data: Data) -> Bool { + guard let first = data.first else { return false } + let fieldNumber = first >> 3 + let wireType = first & 0x07 + return fieldNumber > 0 && (wireType == 0 || wireType == 1 || wireType == 2 || wireType == 5) + } + + static func grpcWebDataFrames(from data: Data) -> [Data] { + let bytes = [UInt8](data) + var frames: [Data] = [] + var index = 0 + while index < bytes.count { + guard index + 5 <= bytes.count else { return [] } + let flags = bytes[index] + let length = (Int(bytes[index + 1]) << 24) + | (Int(bytes[index + 2]) << 16) + | (Int(bytes[index + 3]) << 8) + | Int(bytes[index + 4]) + let start = index + 5 + let end = start + length + guard length >= 0, end <= bytes.count else { return [] } + if flags & 0x80 == 0 { + frames.append(Data(bytes[start.. [String: String] { + var fields: [String: String] = [:] + for (key, value) in headers { + let normalizedKey = String(describing: key) + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + guard normalizedKey.hasPrefix("grpc-") else { continue } + fields[normalizedKey] = String(describing: value) + .trimmingCharacters(in: .whitespacesAndNewlines) + .removingPercentEncoding ?? "" + } + return fields + } + + static func grpcWebTrailerFields(from data: Data) -> [String: String] { + let bytes = [UInt8](data) + var fields: [String: String] = [:] + var index = 0 + while index + 5 <= bytes.count { + let flags = bytes[index] + let length = (Int(bytes[index + 1]) << 24) + | (Int(bytes[index + 2]) << 16) + | (Int(bytes[index + 3]) << 8) + | Int(bytes[index + 4]) + let start = index + 5 + let end = start + length + guard length >= 0, end <= bytes.count else { break } + if flags & 0x80 != 0, let text = String(data: Data(bytes[start.. ProtobufScan { + self.scanProtobuf(data, depth: depth, path: [], order: 0).scan + } + + private static func scanProtobuf( + _ data: Data, + depth: Int, + path: [UInt64], + order: Int) -> (scan: ProtobufScan, order: Int) + { + let bytes = [UInt8](data) + var scan = ProtobufScan() + var index = 0 + var nextOrder = order + + while index < bytes.count { + let fieldStart = index + guard let key = Self.readVarint(bytes, index: &index), key != 0 else { + index = fieldStart + 1 + continue + } + let fieldNumber = key >> 3 + let wireType = key & 0x07 + let fieldPath = path + [fieldNumber] + + switch wireType { + case 0: + if let value = Self.readVarint(bytes, index: &index) { + scan.varintFields.append(ProtobufScan.VarintField(path: fieldPath, value: value)) + } else { + index = fieldStart + 1 + } + case 1: + guard index + 8 <= bytes.count else { return (scan, nextOrder) } + index += 8 + case 2: + guard let length = Self.readVarint(bytes, index: &index), + length <= UInt64(bytes.count - index) + else { + index = fieldStart + 1 + continue + } + let start = index + let end = index + Int(length) + if depth < 4 { + let nested = Self.scanProtobuf( + Data(bytes[start.. UInt64? { + var value: UInt64 = 0 + var shift: UInt64 = 0 + while index < bytes.count, shift < 64 { + let byte = bytes[index] + index += 1 + value |= UInt64(byte & 0x7F) << shift + if byte & 0x80 == 0 { return value } + shift += 7 + } + return nil + } +} diff --git a/Sources/CodexBarCore/Providers/Groq/GroqProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Groq/GroqProviderDescriptor.swift new file mode 100644 index 0000000..f4214da --- /dev/null +++ b/Sources/CodexBarCore/Providers/Groq/GroqProviderDescriptor.swift @@ -0,0 +1,49 @@ +import Foundation + +public enum GroqProviderDescriptor { + public static let descriptor: ProviderDescriptor = Self.makeDescriptor() + + static func makeDescriptor() -> ProviderDescriptor { + ProviderDescriptor( + id: .groq, + metadata: ProviderMetadata( + id: .groq, + displayName: "Groq", + sessionLabel: "Requests", + weeklyLabel: "Tokens", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show Groq usage", + cliName: "groqcloud", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + browserCookieOrder: nil, + dashboardURL: "https://console.groq.com/dashboard/metrics", + statusPageURL: nil, + statusLinkURL: "https://status.groq.com"), + branding: ProviderBranding( + iconStyle: .groq, + iconResourceName: "ProviderIcon-groq", + color: ProviderColor(red: 245 / 255, green: 104 / 255, blue: 68 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "Groq cost history is not available via the metrics API." }), + fetchPlan: .apiToken( + strategyID: "groq.api", + sourceLabel: "metrics", + resolveToken: { ProviderTokenResolver.groqToken(environment: $0) }, + missingCredentialsError: { GroqUsageError.missingCredentials }, + loadUsage: { apiKey, context in + try await GroqUsageFetcher.fetchUsage( + apiKey: apiKey, + environment: context.env).toUsageSnapshot() + }), + cli: ProviderCLIConfig( + name: "groqcloud", + aliases: ["groq", "groq-api"], + versionDetector: nil)) + } +} diff --git a/Sources/CodexBarCore/Providers/Groq/GroqSettingsReader.swift b/Sources/CodexBarCore/Providers/Groq/GroqSettingsReader.swift new file mode 100644 index 0000000..65c8a4a --- /dev/null +++ b/Sources/CodexBarCore/Providers/Groq/GroqSettingsReader.swift @@ -0,0 +1,58 @@ +import Foundation + +public enum GroqSettingsReader { + public static let apiKeyEnvironmentKey = "GROQ_API_KEY" + public static let apiURLEnvironmentKey = "GROQ_API_URL" + + public static func apiKey( + environment: [String: String] = ProcessInfo.processInfo.environment) -> String? + { + self.cleaned(environment[self.apiKeyEnvironmentKey]) + } + + public static func apiURL( + environment: [String: String] = ProcessInfo.processInfo.environment) -> URL + { + if let override = self.validAPIURL(environment: environment) { + return override + } + return URL(string: "https://api.groq.com/v1")! + } + + public static func validateEndpointOverrides( + environment: [String: String] = ProcessInfo.processInfo.environment) throws + { + guard let raw = self.cleaned(environment[self.apiURLEnvironmentKey]) else { return } + guard ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) == nil else { return } + throw GroqSettingsError.invalidEndpointOverride(self.apiURLEnvironmentKey) + } + + static func cleaned(_ raw: String?) -> String? { + guard var value = raw?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { + return nil + } + if (value.hasPrefix("\"") && value.hasSuffix("\"")) || + (value.hasPrefix("'") && value.hasSuffix("'")) + { + value = String(value.dropFirst().dropLast()) + } + value = value.trimmingCharacters(in: .whitespacesAndNewlines) + return value.isEmpty ? nil : value + } + + private static func validAPIURL(environment: [String: String]) -> URL? { + guard let raw = self.cleaned(environment[self.apiURLEnvironmentKey]) else { return nil } + return ProviderEndpointOverrideValidator.normalizedHTTPSURL(from: raw) + } +} + +public enum GroqSettingsError: LocalizedError, Sendable, Equatable { + case invalidEndpointOverride(String) + + public var errorDescription: String? { + switch self { + case let .invalidEndpointOverride(key): + "Groq endpoint override \(key) must use HTTPS or a bare host." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Groq/GroqUsageFetcher.swift b/Sources/CodexBarCore/Providers/Groq/GroqUsageFetcher.swift new file mode 100644 index 0000000..bff9a1b --- /dev/null +++ b/Sources/CodexBarCore/Providers/Groq/GroqUsageFetcher.swift @@ -0,0 +1,235 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public enum GroqUsageError: LocalizedError, Sendable { + case missingCredentials + case invalidURL + case accessDenied(String) + case apiError(String) + case parseFailed(String) + + public var errorDescription: String? { + switch self { + case .missingCredentials: + "Missing Groq API key. Set apiKey in ~/.codexbar/config.json or GROQ_API_KEY." + case .invalidURL: + "Groq metrics URL is invalid." + case let .accessDenied(message): + "Groq metrics access denied: \(message)" + case let .apiError(message): + "Groq metrics API error: \(message)" + case let .parseFailed(message): + "Groq metrics parse error: \(message)" + } + } +} + +public struct GroqUsageSnapshot: Codable, Sendable, Equatable { + public let requestRatePerSecond: Double + public let inputTokenRatePerSecond: Double + public let outputTokenRatePerSecond: Double + public let promptCacheHitRatePerSecond: Double + public let updatedAt: Date + + public init( + requestRatePerSecond: Double, + inputTokenRatePerSecond: Double, + outputTokenRatePerSecond: Double, + promptCacheHitRatePerSecond: Double = 0, + updatedAt: Date) + { + self.requestRatePerSecond = requestRatePerSecond + self.inputTokenRatePerSecond = inputTokenRatePerSecond + self.outputTokenRatePerSecond = outputTokenRatePerSecond + self.promptCacheHitRatePerSecond = promptCacheHitRatePerSecond + self.updatedAt = updatedAt + } + + public var requestsPerMinute: Double { + self.requestRatePerSecond * 60 + } + + public var tokensPerMinute: Double { + (self.inputTokenRatePerSecond + self.outputTokenRatePerSecond) * 60 + } + + public var cacheHitsPerMinute: Double { + self.promptCacheHitRatePerSecond * 60 + } + + public func toUsageSnapshot() -> UsageSnapshot { + UsageSnapshot( + primary: RateWindow( + usedPercent: 0, + windowMinutes: 5, + resetsAt: nil, + resetDescription: "\(Self.formatDecimal(self.requestsPerMinute)) req/min"), + secondary: RateWindow( + usedPercent: 0, + windowMinutes: 5, + resetsAt: nil, + resetDescription: "\(Self.formatDecimal(self.tokensPerMinute)) tok/min"), + tertiary: self.promptCacheHitRatePerSecond > 0 ? RateWindow( + usedPercent: 0, + windowMinutes: 5, + resetsAt: nil, + resetDescription: "\(Self.formatDecimal(self.cacheHitsPerMinute)) cache/min") : nil, + updatedAt: self.updatedAt, + identity: ProviderIdentitySnapshot( + providerID: .groq, + accountEmail: nil, + accountOrganization: nil, + loginMethod: "Prometheus metrics")) + } + + static func formatDecimal(_ value: Double) -> String { + if value >= 100 { + return String(format: "%.0f", value) + } + if value >= 10 { + return String(format: "%.1f", value) + } + return String(format: "%.2f", value) + } +} + +private struct GroqPrometheusResponse: Decodable { + struct Payload: Decodable { + let result: [Series] + } + + struct Series: Decodable { + let value: [PrometheusValue]? + } + + enum PrometheusValue: Decodable { + case number(Double) + case string(String) + + init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if let number = try? container.decode(Double.self) { + self = .number(number) + return + } + self = try .string(container.decode(String.self)) + } + + var doubleValue: Double? { + switch self { + case let .number(number): + number + case let .string(text): + Double(text) + } + } + } + + let status: String + let data: Payload? + let error: String? +} + +public struct GroqUsageFetcher: Sendable { + public init() {} + + public static func fetchUsage( + apiKey: String, + environment: [String: String] = ProcessInfo.processInfo.environment, + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared, + updatedAt: Date = Date()) async throws -> GroqUsageSnapshot + { + guard !apiKey.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw GroqUsageError.missingCredentials + } + try GroqSettingsReader.validateEndpointOverrides(environment: environment) + let baseURL = GroqSettingsReader.apiURL(environment: environment) + .appendingPathComponent("metrics") + .appendingPathComponent("prometheus") + + async let requests = Self.queryScalar( + query: "sum(model_project_id_status_code:requests:rate5m)", + apiKey: apiKey, + baseURL: baseURL, + transport: transport) + async let inputTokens = Self.queryScalar( + query: "sum(model_project_id:tokens_in:rate5m)", + apiKey: apiKey, + baseURL: baseURL, + transport: transport) + async let outputTokens = Self.queryScalar( + query: "sum(model_project_id:tokens_out:rate5m)", + apiKey: apiKey, + baseURL: baseURL, + transport: transport) + async let cacheHits = Self.queryScalar( + query: "sum(model_project_id:prompt_cache_hits:rate5m)", + apiKey: apiKey, + baseURL: baseURL, + transport: transport) + + return try await GroqUsageSnapshot( + requestRatePerSecond: requests, + inputTokenRatePerSecond: inputTokens, + outputTokenRatePerSecond: outputTokens, + promptCacheHitRatePerSecond: cacheHits, + updatedAt: updatedAt) + } + + public static func _parseScalarForTesting(_ data: Data) throws -> Double { + try self.parseScalar(data: data) + } + + private static func queryScalar( + query: String, + apiKey: String, + baseURL: URL, + transport: any ProviderHTTPTransport) async throws -> Double + { + guard var components = URLComponents( + url: baseURL.appendingPathComponent("api/v1/query"), + resolvingAgainstBaseURL: false) + else { throw GroqUsageError.invalidURL } + components.queryItems = [URLQueryItem(name: "query", value: query)] + guard let url = components.url else { throw GroqUsageError.invalidURL } + + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization") + request.setValue("application/json", forHTTPHeaderField: "Accept") + + let response = try await transport.response(for: request) + guard (200..<300).contains(response.statusCode) else { + let summary = Self.responseSummary(response.data) + if response.statusCode == 401 || response.statusCode == 403 { + throw GroqUsageError.accessDenied(summary) + } + throw GroqUsageError.apiError("HTTP \(response.statusCode): \(summary)") + } + return try self.parseScalar(data: response.data) + } + + private static func parseScalar(data: Data) throws -> Double { + do { + let decoded = try JSONDecoder().decode(GroqPrometheusResponse.self, from: data) + guard decoded.status == "success" else { + throw GroqUsageError.apiError(decoded.error ?? "query failed") + } + return decoded.data?.result.compactMap { series in + series.value?.last?.doubleValue + }.reduce(0, +) ?? 0 + } catch let error as GroqUsageError { + throw error + } catch { + throw GroqUsageError.parseFailed(error.localizedDescription) + } + } + + private static func responseSummary(_ data: Data) -> String { + String(bytes: data.prefix(500), encoding: .utf8)? + .trimmingCharacters(in: .whitespacesAndNewlines) + ?? "" + } +} diff --git a/Sources/CodexBarCore/Providers/JetBrains/JetBrainsIDEDetector.swift b/Sources/CodexBarCore/Providers/JetBrains/JetBrainsIDEDetector.swift new file mode 100644 index 0000000..c1850f5 --- /dev/null +++ b/Sources/CodexBarCore/Providers/JetBrains/JetBrainsIDEDetector.swift @@ -0,0 +1,150 @@ +import Foundation + +public struct JetBrainsIDEInfo: Sendable, Equatable, Hashable { + public let name: String + public let version: String + public let basePath: String + public let quotaFilePath: String + + public init(name: String, version: String, basePath: String, quotaFilePath: String) { + self.name = name + self.version = version + self.basePath = basePath + self.quotaFilePath = quotaFilePath + } + + public var displayName: String { + "\(self.name) \(self.version)" + } +} + +public enum JetBrainsIDEDetector { + private static let idePatterns: [(prefix: String, displayName: String)] = [ + ("IntelliJIdea", "IntelliJ IDEA"), + ("PyCharm", "PyCharm"), + ("WebStorm", "WebStorm"), + ("GoLand", "GoLand"), + ("CLion", "CLion"), + ("DataGrip", "DataGrip"), + ("RubyMine", "RubyMine"), + ("Rider", "Rider"), + ("PhpStorm", "PhpStorm"), + ("AppCode", "AppCode"), + ("Fleet", "Fleet"), + ("AndroidStudio", "Android Studio"), + ("RustRover", "RustRover"), + ("Aqua", "Aqua"), + ("DataSpell", "DataSpell"), + ] + + private static let quotaFileName = "AIAssistantQuotaManager2.xml" + + public static func detectInstalledIDEs(includeMissingQuota: Bool = false) -> [JetBrainsIDEInfo] { + let basePaths = self.jetBrainsConfigBasePaths() + var detectedIDEs: [JetBrainsIDEInfo] = [] + + let fileManager = FileManager.default + for basePath in basePaths { + guard fileManager.fileExists(atPath: basePath) else { continue } + guard let contents = try? fileManager.contentsOfDirectory(atPath: basePath) else { continue } + + for dirname in contents { + guard let ideInfo = parseIDEDirectory(dirname: dirname, basePath: basePath) else { continue } + if includeMissingQuota || fileManager.fileExists(atPath: ideInfo.quotaFilePath) { + detectedIDEs.append(ideInfo) + } + } + } + + return detectedIDEs.sorted { lhs, rhs in + if lhs.name == rhs.name { + return self.compareVersions(lhs.version, rhs.version) > 0 + } + return lhs.name < rhs.name + } + } + + public static func detectLatestIDE() -> JetBrainsIDEInfo? { + let ides = self.detectInstalledIDEs() + guard !ides.isEmpty else { return nil } + + let fileManager = FileManager.default + var latestIDE: JetBrainsIDEInfo? + var latestModificationDate: Date? + + for ide in ides { + guard let attrs = try? fileManager.attributesOfItem(atPath: ide.quotaFilePath), + let modDate = attrs[.modificationDate] as? Date + else { continue } + + if latestModificationDate == nil || modDate > latestModificationDate! { + latestModificationDate = modDate + latestIDE = ide + } + } + + return latestIDE ?? ides.first + } + + private static func jetBrainsConfigBasePaths() -> [String] { + let homeDir = FileManager.default.homeDirectoryForCurrentUser.path + #if os(macOS) + return [ + "\(homeDir)/Library/Application Support/JetBrains", + "\(homeDir)/Library/Application Support/Google", + ] + #else + return [ + "\(homeDir)/.config/JetBrains", + "\(homeDir)/.local/share/JetBrains", + "\(homeDir)/.config/Google", + ] + #endif + } + + private static func parseIDEDirectory(dirname: String, basePath: String) -> JetBrainsIDEInfo? { + let dirnameLower = dirname.lowercased() + for (prefix, displayName) in self.idePatterns { + let prefixLower = prefix.lowercased() + guard dirnameLower.hasPrefix(prefixLower) else { continue } + let versionPart = String(dirname.dropFirst(prefix.count)) + let version = versionPart.isEmpty ? "Unknown" : versionPart + let idePath = "\(basePath)/\(dirname)" + let quotaFilePath = quotaFilePath(for: idePath) + return JetBrainsIDEInfo( + name: displayName, + version: version, + basePath: idePath, + quotaFilePath: quotaFilePath) + } + return nil + } + + #if DEBUG + + // MARK: - Test hooks (DEBUG-only) + + public static func _parseIDEDirectoryForTesting(dirname: String, basePath: String) -> JetBrainsIDEInfo? { + self.parseIDEDirectory(dirname: dirname, basePath: basePath) + } + #endif + + public static func quotaFilePath(for ideBasePath: String) -> String { + "\(ideBasePath)/options/\(self.quotaFileName)" + } + + private static func compareVersions(_ v1: String, _ v2: String) -> Int { + let parts1 = v1.split(separator: ".").compactMap { Int($0) } + let parts2 = v2.split(separator: ".").compactMap { Int($0) } + + let maxLen = max(parts1.count, parts2.count) + for i in 0.. ProviderDescriptor { + ProviderDescriptor( + id: .jetbrains, + metadata: ProviderMetadata( + id: .jetbrains, + displayName: "JetBrains AI", + sessionLabel: "Current", + weeklyLabel: "Refill", + opusLabel: nil, + supportsOpus: false, + supportsCredits: false, + creditsHint: "", + toggleTitle: "Show JetBrains AI usage", + cliName: "jetbrains", + defaultEnabled: false, + isPrimaryProvider: false, + usesAccountFallback: false, + dashboardURL: nil, + statusPageURL: nil), + branding: ProviderBranding( + iconStyle: .jetbrains, + iconResourceName: "ProviderIcon-jetbrains", + color: ProviderColor(red: 255 / 255, green: 51 / 255, blue: 153 / 255)), + tokenCost: ProviderTokenCostConfig( + supportsTokenCost: false, + noDataMessage: { "JetBrains AI cost summary is not supported." }), + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .cli], + pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [JetBrainsStatusFetchStrategy()] })), + cli: ProviderCLIConfig( + name: "jetbrains", + versionDetector: nil)) + } +} + +struct JetBrainsStatusFetchStrategy: ProviderFetchStrategy { + let id: String = "jetbrains.local" + let kind: ProviderFetchKind = .localProbe + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let probe = JetBrainsStatusProbe(settings: context.settings) + let snap = try await probe.fetch() + let usage = try snap.toUsageSnapshot() + return self.makeResult( + usage: usage, + sourceLabel: "local") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} diff --git a/Sources/CodexBarCore/Providers/JetBrains/JetBrainsStatusProbe.swift b/Sources/CodexBarCore/Providers/JetBrains/JetBrainsStatusProbe.swift new file mode 100644 index 0000000..198d830 --- /dev/null +++ b/Sources/CodexBarCore/Providers/JetBrains/JetBrainsStatusProbe.swift @@ -0,0 +1,355 @@ +import Foundation +#if os(macOS) && canImport(FoundationXML) +import FoundationXML +#endif + +public struct JetBrainsQuotaInfo: Sendable, Equatable { + public let type: String? + public let used: Double + public let maximum: Double + public let available: Double + public let until: Date? + + public init(type: String?, used: Double, maximum: Double, available: Double?, until: Date?) { + self.type = type + self.used = used + self.maximum = maximum + // Use available if provided, otherwise calculate from maximum - used + self.available = available ?? max(0, maximum - used) + self.until = until + } + + /// Percentage of quota that has been used (0-100) + public var usedPercent: Double { + guard self.maximum > 0 else { return 0 } + return min(100, max(0, (self.used / self.maximum) * 100)) + } + + /// Percentage of quota remaining (0-100), based on available value + public var remainingPercent: Double { + guard self.maximum > 0 else { return 100 } + return min(100, max(0, (self.available / self.maximum) * 100)) + } +} + +public struct JetBrainsRefillInfo: Sendable, Equatable { + public let type: String? + public let next: Date? + public let amount: Double? + public let duration: String? + + public init(type: String?, next: Date?, amount: Double?, duration: String?) { + self.type = type + self.next = next + self.amount = amount + self.duration = duration + } +} + +public struct JetBrainsStatusSnapshot: Sendable { + public let quotaInfo: JetBrainsQuotaInfo + public let refillInfo: JetBrainsRefillInfo? + public let detectedIDE: JetBrainsIDEInfo? + + public init(quotaInfo: JetBrainsQuotaInfo, refillInfo: JetBrainsRefillInfo?, detectedIDE: JetBrainsIDEInfo?) { + self.quotaInfo = quotaInfo + self.refillInfo = refillInfo + self.detectedIDE = detectedIDE + } + + public func toUsageSnapshot() throws -> UsageSnapshot { + // Primary shows monthly credits usage with next refill date + // IDE displays: "今月のクレジット残り X / Y" with "Z月D日に更新されます" + let refillDate = self.refillInfo?.next + let primary = RateWindow( + usedPercent: self.quotaInfo.usedPercent, + windowMinutes: nil, + resetsAt: refillDate, + resetDescription: Self.formatResetDescription(refillDate)) + + let identity = ProviderIdentitySnapshot( + providerID: .jetbrains, + accountEmail: nil, + accountOrganization: self.detectedIDE?.displayName, + loginMethod: self.quotaInfo.type) + + return UsageSnapshot( + primary: primary, + secondary: nil, + tertiary: nil, + updatedAt: Date(), + identity: identity) + } + + private static func formatResetDescription(_ date: Date?) -> String? { + guard let date else { return nil } + let now = Date() + let interval = date.timeIntervalSince(now) + guard interval > 0 else { return "Expired" } + + let hours = Int(interval / 3600) + let minutes = Int((interval.truncatingRemainder(dividingBy: 3600)) / 60) + + if hours > 24 { + let days = hours / 24 + let remainingHours = hours % 24 + return "Resets in \(days)d \(remainingHours)h" + } else if hours > 0 { + return "Resets in \(hours)h \(minutes)m" + } else { + return "Resets in \(minutes)m" + } + } +} + +public enum JetBrainsStatusProbeError: LocalizedError, Sendable, Equatable { + case noIDEDetected + case quotaFileNotFound(String) + case parseError(String) + case noQuotaInfo + + public var errorDescription: String? { + switch self { + case .noIDEDetected: + "No JetBrains IDE with AI Assistant detected. Install a JetBrains IDE and enable AI Assistant." + case let .quotaFileNotFound(path): + "JetBrains AI quota file not found at \(path). Enable AI Assistant in your IDE." + case let .parseError(message): + "Could not parse JetBrains AI quota: \(message)" + case .noQuotaInfo: + "No quota information found in the JetBrains AI configuration." + } + } +} + +public struct JetBrainsStatusProbe: Sendable { + private let settings: ProviderSettingsSnapshot? + + public init(settings: ProviderSettingsSnapshot? = nil) { + self.settings = settings + } + + public func fetch() async throws -> JetBrainsStatusSnapshot { + let (quotaFilePath, detectedIDE) = try self.resolveQuotaFilePath() + return try Self.parseQuotaFile(at: quotaFilePath, detectedIDE: detectedIDE) + } + + private func resolveQuotaFilePath() throws -> (String, JetBrainsIDEInfo?) { + if let customPath = self.settings?.jetbrainsIDEBasePath?.trimmingCharacters(in: .whitespacesAndNewlines), + !customPath.isEmpty + { + let expandedBasePath = (customPath as NSString).expandingTildeInPath + let quotaPath = JetBrainsIDEDetector.quotaFilePath(for: expandedBasePath) + return (quotaPath, nil) + } + + guard let detectedIDE = JetBrainsIDEDetector.detectLatestIDE() else { + throw JetBrainsStatusProbeError.noIDEDetected + } + return (detectedIDE.quotaFilePath, detectedIDE) + } + + public static func parseQuotaFile( + at path: String, + detectedIDE: JetBrainsIDEInfo?) throws -> JetBrainsStatusSnapshot + { + guard FileManager.default.fileExists(atPath: path) else { + throw JetBrainsStatusProbeError.quotaFileNotFound(path) + } + + let xmlData: Data + do { + xmlData = try Data(contentsOf: URL(fileURLWithPath: path)) + } catch { + throw JetBrainsStatusProbeError.parseError("Failed to read file: \(error.localizedDescription)") + } + + return try Self.parseXMLData(xmlData, detectedIDE: detectedIDE) + } + + public static func parseXMLData(_ data: Data, detectedIDE: JetBrainsIDEInfo?) throws -> JetBrainsStatusSnapshot { + #if os(macOS) + let document: XMLDocument + do { + document = try XMLDocument(data: data) + } catch { + throw JetBrainsStatusProbeError.parseError("Invalid XML: \(error.localizedDescription)") + } + + let quotaInfoRaw = try? document + .nodes(forXPath: "//component[@name='AIAssistantQuotaManager2']/option[@name='quotaInfo']/@value") + .first? + .stringValue + let nextRefillRaw = try? document + .nodes(forXPath: "//component[@name='AIAssistantQuotaManager2']/option[@name='nextRefill']/@value") + .first? + .stringValue + #else + let parseResult = JetBrainsXMLParser.parse(data: data) + let quotaInfoRaw = parseResult.quotaInfo + let nextRefillRaw = parseResult.nextRefill + #endif + + guard let quotaInfoRaw, !quotaInfoRaw.isEmpty else { + throw JetBrainsStatusProbeError.noQuotaInfo + } + + let quotaInfoDecoded = Self.decodeHTMLEntities(quotaInfoRaw) + let quotaInfo = try Self.parseQuotaInfoJSON(quotaInfoDecoded) + + var refillInfo: JetBrainsRefillInfo? + if let nextRefillRaw, !nextRefillRaw.isEmpty { + let nextRefillDecoded = Self.decodeHTMLEntities(nextRefillRaw) + refillInfo = try? Self.parseRefillInfoJSON(nextRefillDecoded) + } + + return JetBrainsStatusSnapshot( + quotaInfo: quotaInfo, + refillInfo: refillInfo, + detectedIDE: detectedIDE) + } + + private static func decodeHTMLEntities(_ string: String) -> String { + string + .replacingOccurrences(of: " ", with: "\n") + .replacingOccurrences(of: """, with: "\"") + .replacingOccurrences(of: "&", with: "&") + .replacingOccurrences(of: "<", with: "<") + .replacingOccurrences(of: ">", with: ">") + .replacingOccurrences(of: "'", with: "'") + } + + private static func parseQuotaInfoJSON(_ jsonString: String) throws -> JetBrainsQuotaInfo { + guard let data = jsonString.data(using: .utf8) else { + throw JetBrainsStatusProbeError.parseError("Invalid JSON encoding") + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw JetBrainsStatusProbeError.parseError("Invalid JSON format") + } + + let type = json["type"] as? String + let currentStr = json["current"] as? String + let maximumStr = json["maximum"] as? String + let untilStr = json["until"] as? String + + // tariffQuota contains the actual available credits + let tariffQuota = json["tariffQuota"] as? [String: Any] + let availableStr = tariffQuota?["available"] as? String + + let used = currentStr.flatMap { Double($0) } ?? 0 + let maximum = maximumStr.flatMap { Double($0) } ?? 0 + let available = availableStr.flatMap { Double($0) } + let until = untilStr.flatMap { Self.parseDate($0) } + + return JetBrainsQuotaInfo(type: type, used: used, maximum: maximum, available: available, until: until) + } + + private static func parseRefillInfoJSON(_ jsonString: String) throws -> JetBrainsRefillInfo { + guard let data = jsonString.data(using: .utf8) else { + throw JetBrainsStatusProbeError.parseError("Invalid JSON encoding") + } + + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw JetBrainsStatusProbeError.parseError("Invalid JSON format") + } + + let type = json["type"] as? String + let nextStr = json["next"] as? String + let amountStr = json["amount"] as? String + let duration = json["duration"] as? String + + let next = nextStr.flatMap { Self.parseDate($0) } + let amount = amountStr.flatMap { Double($0) } + + let tariff = json["tariff"] as? [String: Any] + let tariffAmountStr = tariff?["amount"] as? String + let tariffDuration = tariff?["duration"] as? String + let finalAmount = amount ?? tariffAmountStr.flatMap { Double($0) } + let finalDuration = duration ?? tariffDuration + + return JetBrainsRefillInfo(type: type, next: next, amount: finalAmount, duration: finalDuration) + } + + private static func parseDate(_ string: String) -> Date? { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + if let date = formatter.date(from: string) { + return date + } + + formatter.formatOptions = [.withInternetDateTime] + return formatter.date(from: string) + } +} + +#if !os(macOS) +/// Simple regex-based XML parser to avoid libxml2 dependency on Linux. +/// Only extracts quotaInfo and nextRefill values from AIAssistantQuotaManager2 component. +private enum JetBrainsXMLParser { + struct ParseResult { + let quotaInfo: String? + let nextRefill: String? + } + + static func parse(data: Data) -> ParseResult { + guard let content = String(data: data, encoding: .utf8) else { + return ParseResult(quotaInfo: nil, nextRefill: nil) + } + + // Find the AIAssistantQuotaManager2 component block + guard let componentRange = self.findComponentRange(in: content) else { + return ParseResult(quotaInfo: nil, nextRefill: nil) + } + + let componentContent = String(content[componentRange]) + + let quotaInfo = self.extractOptionValue(named: "quotaInfo", from: componentContent) + let nextRefill = self.extractOptionValue(named: "nextRefill", from: componentContent) + + return ParseResult(quotaInfo: quotaInfo, nextRefill: nextRefill) + } + + private static func findComponentRange(in content: String) -> Range? { + // Match ... + let pattern = #"]*name\s*=\s*["']AIAssistantQuotaManager2["'][^>]*>[\s\S]*?"# + guard let regex = try? NSRegularExpression(pattern: pattern, options: []), + let match = regex.firstMatch( + in: content, + options: [], + range: NSRange(content.startIndex..., in: content)), + let range = Range(match.range, in: content) + else { + return nil + } + return range + } + + private static func extractOptionValue(named name: String, from content: String) -> String? { + // Match