diff --git a/backend/druks/harnesses/base.py b/backend/druks/harnesses/base.py index 13c38648..62421e6a 100644 --- a/backend/druks/harnesses/base.py +++ b/backend/druks/harnesses/base.py @@ -120,8 +120,8 @@ def get_secrets(cls, provider: str, key: str) -> dict[str, Secret]: @classmethod def get_secret_refs(cls, subscription: VaultSecret) -> list[SecretRef]: - """The secrets a box fetches for the subscription, by Drukbox catalog - name. Empty for a CLI that reads its credential from a file.""" + """The secrets a box fetches for the subscription, under the box's + name for each. Empty for a key-only CLI.""" return [] @property diff --git a/backend/druks/harnesses/claude.py b/backend/druks/harnesses/claude.py index fdb882e8..4851c1f8 100644 --- a/backend/druks/harnesses/claude.py +++ b/backend/druks/harnesses/claude.py @@ -67,9 +67,7 @@ async def build_invocation( skills: tuple[str, ...] = (), extra_env: dict[str, str] | None = None, mcp_servers: tuple[McpServer, ...] = (), - # Accepted for signature parity. The sandbox holds a placeholder for - # the subscription token or the key. Drukbox delivers it. - subscription: VaultSecret | None = None, + identity: dict | None = None, timeout: int = Harness.default_timeout, ) -> AgentInvocation: if not self.sandbox: diff --git a/backend/druks/harnesses/codex.py b/backend/druks/harnesses/codex.py index ef1a9c67..68d2ca9f 100644 --- a/backend/druks/harnesses/codex.py +++ b/backend/druks/harnesses/codex.py @@ -1,3 +1,4 @@ +import base64 import json import logging import os @@ -16,10 +17,10 @@ Credentials, HarnessRunResult, HomeCopy, - HomeFile, McpServer, ) from druks.sandbox.layout import get_runs_root, get_work_root +from druks.sandbox.models import SecretRef from druks.secrets.models import VaultSecret from druks.skills.models import Skill @@ -33,12 +34,14 @@ HarnessRateLimitError, HarnessUsageLimitError, ) -from .providers import OpenAiProvider +from .providers import OPENAI_AUTH_CLAIM, OpenAiProvider from .subprocess import read_result_json logger = logging.getLogger(__name__) _TOKEN_COUNT_MARKERS = ('"type":"token_count"', '"type": "token_count"') +_CHATGPT_HOST = "chatgpt.com" +_SUBSCRIPTION_TOKEN = "codex_subscription_token" @dataclass(frozen=True) @@ -138,6 +141,35 @@ def _with_final_message_note(prompt: str) -> str: ) +def _auth_file(identity: dict) -> str: + """The ``auth.json`` Codex reads in subscription mode. Codex needs the id + token and a refresh token key present, else it switches to API-key mode.""" + claims = {"chatgpt_account_id": identity["account_id"]} + if "plan" in identity: + claims["chatgpt_plan_type"] = identity["plan"] + # Codex decodes the id token without a signature check. + header = {"alg": "none", "typ": "JWT"} + payload = {OPENAI_AUTH_CLAIM: claims, "email": identity["email"]} + segments = [ + base64.urlsafe_b64encode(json.dumps(part).encode()).rstrip(b"=").decode() + for part in (header, payload) + ] + return json.dumps( + { + "OPENAI_API_KEY": None, + "tokens": { + "id_token": ".".join((*segments, "unsigned")), + "access_token": f"${_SUBSCRIPTION_TOKEN.upper()}", + # The one refresh Codex attempts after a 401 fails fast on this + # value, and the turn ends. Druks refreshes; the box never does. + "refresh_token": "druks-placeholder", + "account_id": identity["account_id"], + }, + "last_refresh": datetime.now(UTC).isoformat().replace("+00:00", "Z"), + } + ) + + def _parse_token_count_events( path: Path, window_start: datetime, @@ -291,6 +323,7 @@ def _build_codex_wrapper( run_id: str, codex_flags: tuple[str, ...], cwd: str, + identity: dict | None, ) -> list[str]: # In-VM paths under //. Schema is inlined via # printf (~few KB); the prompt rides as stdin via the helper. @@ -329,6 +362,16 @@ def _build_codex_wrapper( codex_cmdline = " ".join(shlex.quote(a) for a in codex_argv) marker_q = shlex.quote(in_vm_marker) session_q = shlex.quote(in_vm_session) + # A subscription run writes its login before the command. Drukbox hands + # the placeholder to the box only, so the box fills the file itself. + login = "" + if identity: + auth_file = _auth_file(identity).replace('"', '\\"') + login = ( + 'mkdir -p "$HOME/.codex" && ' + f'printf %s "{auth_file}" > "$HOME/.codex/auth.json" && ' + 'chmod 600 "$HOME/.codex/auth.json" && ' + ) # Codex runs against its real ``~/.codex`` because CODEX_HOME re-homes # all home-resolved state: auth, skills, and future features. A marker # file plus ``-newer`` identifies the session JSONL. More than one match @@ -337,6 +380,7 @@ def _build_codex_wrapper( wrapper = ( f"mkdir -p {shlex.quote(in_vm_run_dir)} && " f"printf %s {shlex.quote(schema_body)} > {shlex.quote(in_vm_schema)} && " + f"{login}" f"touch {marker_q} && " f"{codex_cmdline}; " "ec=$?; " @@ -367,7 +411,7 @@ async def build_invocation( skills: tuple[str, ...] = (), extra_env: dict[str, str] | None = None, mcp_servers: tuple[McpServer, ...] = (), - subscription: VaultSecret | None = None, + identity: dict | None = None, timeout: int = Harness.default_timeout, ) -> AgentInvocation: sandbox = self.sandbox @@ -383,6 +427,7 @@ async def build_invocation( run_id=run_id, codex_flags=(*self._prompt_flags(), *self._mcp_flags(mcp_servers)), cwd=get_work_root(ssh_username), + identity=identity, ) # --output-schema constrains EVERY agent_message mechanically (the # "final response shape" in its docs is inaccurate — verified by @@ -394,11 +439,7 @@ async def build_invocation( name=self.name, args=tuple(cmd), stdin=_with_final_message_note(prompt).encode("utf-8"), - credentials=await self._get_credentials( - sandbox, - skills=skills, - subscription=subscription, - ), + credentials=await self._get_credentials(sandbox, skills=skills), env=extra_env, extra_artifact_filenames=("output.json", "session.jsonl"), ) @@ -473,29 +514,25 @@ def _prompt_flags(self) -> tuple[str, ...]: return args async def _get_credentials( - self, - sandbox: SandboxSettings, - *, - skills: tuple[str, ...] = (), - subscription: VaultSecret | None, + self, sandbox: SandboxSettings, *, skills: tuple[str, ...] = () ) -> Credentials: config_dir = sandbox.harness_config_root / self.name - home: list[HomeFile | HomeCopy] = [] - if subscription: - home.append(self.auth_file(subscription)) - home += [ - HomeCopy(".codex/config.toml", config_dir / "config.toml"), - HomeCopy(".codex/AGENTS.md", config_dir / "AGENTS.md"), - ] skills_dir = sandbox.skills_dir or config_dir / "skills" - home.append( - HomeCopy(".codex/skills", skills_dir, excludes=await Skill.delivery_excludes(skills)) + return Credentials( + home=( + HomeCopy(".codex/config.toml", config_dir / "config.toml"), + HomeCopy(".codex/AGENTS.md", config_dir / "AGENTS.md"), + HomeCopy( + ".codex/skills", skills_dir, excludes=await Skill.delivery_excludes(skills) + ), + ) ) - return Credentials(home=tuple(home)) @classmethod - def auth_file(cls, subscription: VaultSecret) -> HomeFile: - return HomeFile(".codex/auth.json", json.dumps(dict(subscription.secrets))) + def get_secret_refs(cls, subscription: VaultSecret) -> list[SecretRef]: + # A custom entry: the proxy swaps the placeholder on every chatgpt.com + # request. CODEX_API_KEY would put codex exec in API-key mode. + return [SecretRef(name=_SUBSCRIPTION_TOKEN, secret_id=subscription.id, host=_CHATGPT_HOST)] @classmethod def get_secrets(cls, provider: str, key: str) -> dict[str, Secret]: diff --git a/backend/druks/harnesses/opencode.py b/backend/druks/harnesses/opencode.py index f08c2dac..dcfbaf0b 100644 --- a/backend/druks/harnesses/opencode.py +++ b/backend/druks/harnesses/opencode.py @@ -11,7 +11,6 @@ McpServer, ) from druks.sandbox.layout import get_runs_root, get_work_root -from druks.secrets.models import VaultSecret from . import exceptions from .artifacts import call_dir, write_cost @@ -49,9 +48,7 @@ async def build_invocation( skills: tuple[str, ...] = (), extra_env: dict[str, str] | None = None, mcp_servers: tuple[McpServer, ...] = (), - # Accepted for signature parity; opencode runs on an API key only, and - # the sandbox holds it as a placeholder in the provider's variable. - subscription: VaultSecret | None = None, + identity: dict | None = None, timeout: int = Harness.default_timeout, ) -> AgentInvocation: if not self.sandbox: diff --git a/backend/druks/harnesses/pi.py b/backend/druks/harnesses/pi.py index e7e7114e..1351bfa8 100644 --- a/backend/druks/harnesses/pi.py +++ b/backend/druks/harnesses/pi.py @@ -11,7 +11,6 @@ McpServer, ) from druks.sandbox.layout import get_runs_root -from druks.secrets.models import VaultSecret from . import exceptions from .artifacts import write_cost @@ -51,9 +50,7 @@ async def build_invocation( skills: tuple[str, ...] = (), extra_env: dict[str, str] | None = None, mcp_servers: tuple[McpServer, ...] = (), - # Accepted for signature parity; pi runs on an API key only, and the - # sandbox holds it as a placeholder in the provider's variable. - subscription: VaultSecret | None = None, + identity: dict | None = None, timeout: int = Harness.default_timeout, ) -> AgentInvocation: if not self.sandbox: diff --git a/backend/druks/harnesses/profiles.py b/backend/druks/harnesses/profiles.py index 95e259d2..edadae48 100644 --- a/backend/druks/harnesses/profiles.py +++ b/backend/druks/harnesses/profiles.py @@ -29,6 +29,8 @@ class Profile: secrets: dict[str, Secret] # The secrets a box fetches through the issuer, beyond its pasted key. secret_refs: list[SecretRef] + # The subscription's non-secret facts, for the login a box sees. Empty for a key. + identity: dict billing: str effort: str timeout: int @@ -104,6 +106,7 @@ async def get_profile(agent_name: str, account_id: str | None) -> Profile: provider_key = None secrets: dict[str, Secret] = {} secret_refs: list[SecretRef] = [] + identity: dict = {} if billing == "api_key": provider_key = await VaultSecret.lookup(SecretKind.STATIC, Audience.provider(provider_id)) if not provider_key: @@ -111,7 +114,9 @@ async def get_profile(agent_name: str, account_id: str | None) -> Profile: raise HarnessNotConnectedError(f"add the {label} API key in Settings → Providers.") secrets = harness_class.get_secrets(provider_id, provider_key.secrets["value"]) else: - subscription = await get_provider(provider_id).get_subscription(account_id) + provider = get_provider(provider_id) + subscription = await provider.get_subscription(account_id) + identity = provider.get_identity(subscription) secret_refs = harness_class.get_secret_refs(subscription) timeout = ( await SettingsOverride.agent_timeout(agent_name, agent.timeout, settings=settings) @@ -123,6 +128,7 @@ async def get_profile(agent_name: str, account_id: str | None) -> Profile: api_key=provider_key, secrets=secrets, secret_refs=secret_refs, + identity=identity, billing=billing, effort=(await SettingsOverride.agent_effort(agent_name, settings=settings)).value, # Capped so a single call always fits inside a fresh sandbox lease. diff --git a/backend/druks/harnesses/providers.py b/backend/druks/harnesses/providers.py index ba334000..aa232935 100644 --- a/backend/druks/harnesses/providers.py +++ b/backend/druks/harnesses/providers.py @@ -176,6 +176,11 @@ async def get_subscription( f"connect your {cls.label} subscription in Settings → Providers." ) + @classmethod + def get_identity(cls, subscription: VaultSecret) -> dict: + """The subscription's non-secret facts, for the login a box sees.""" + return dict(subscription.identity) + @classmethod def load_token(cls, subscription: VaultSecret, *, now: datetime | None = None) -> Token: """Read + validate ``subscription``'s access token, or raise @@ -883,8 +888,8 @@ def _parse_iso(value: object) -> datetime | None: return ensure_utc(parsed) -# Namespaced claims OpenAI packs into the Codex access-token JWT. -_OPENAI_AUTH_CLAIM = "https://api.openai.com/auth" +# Namespaced claims OpenAI packs into the Codex access and id token JWTs. +OPENAI_AUTH_CLAIM = "https://api.openai.com/auth" _OPENAI_PROFILE_CLAIM = "https://api.openai.com/profile" # ChatGPT subscription usage endpoint — the standalone fetch the codex CLI's @@ -919,6 +924,15 @@ def get_secret(cls, key: str) -> Secret: # The catalog entry: OPENAI_API_KEY as a bearer on api.openai.com. return Secret(key) + @classmethod + def get_identity(cls, subscription: VaultSecret) -> dict: + tokens = subscription.secrets["tokens"] + auth = (jwt_claims(tokens.get("id_token") or "") or {}).get(OPENAI_AUTH_CLAIM) or {} + identity = {**super().get_identity(subscription), "account_id": tokens.get("account_id")} + if plan := auth.get("chatgpt_plan_type"): + identity["plan"] = plan + return identity + @classmethod def _token_from_credentials(cls, data: dict) -> CodexToken: tokens = data.get("tokens") if isinstance(data.get("tokens"), dict) else {} @@ -975,7 +989,7 @@ async def exchange(cls, *, code: str, verifier: str) -> tuple[dict, str | None]: ) access = grant["access_token"] claims = jwt_claims(access) or {} - auth = claims.get(_OPENAI_AUTH_CLAIM) or {} + auth = claims.get(OPENAI_AUTH_CLAIM) or {} profile = claims.get(_OPENAI_PROFILE_CLAIM) or {} payload = { "OPENAI_API_KEY": None, diff --git a/backend/druks/sandbox/host.py b/backend/druks/sandbox/host.py index 7695ff0b..b8f3e032 100644 --- a/backend/druks/sandbox/host.py +++ b/backend/druks/sandbox/host.py @@ -39,7 +39,6 @@ if TYPE_CHECKING: from druks.harnesses.base import Harness from druks.harnesses.profiles import Profile - from druks.secrets.models import VaultSecret from .runner import Exec @@ -253,7 +252,7 @@ async def run_agent( extra_env=extra_env, mcp_servers=mcp_servers, call_id=run_id, - subscription=profile.subscription, + identity=profile.identity, ) except HarnessError as exc: error = exc @@ -291,7 +290,7 @@ async def run_prompt( extra_env: dict[str, str] | None = None, mcp_servers: tuple[McpServer, ...] = (), call_id: str | None = None, - subscription: "VaultSecret | None" = None, + identity: dict | None = None, ) -> Any: """Drive one prompt through ``harness`` on this VM: the harness builds the invocation and parses the result; this sandbox executes it.""" @@ -314,7 +313,7 @@ async def run_prompt( skills=skills, extra_env=extra_env, mcp_servers=mcp_servers, - subscription=subscription, + identity=identity, timeout=timeout, ) result = await self._exec( diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index 8572144c..97a3bfc7 100644 --- a/backend/tests/conftest.py +++ b/backend/tests/conftest.py @@ -1,3 +1,5 @@ +import base64 +import json from pathlib import Path from unittest import mock @@ -208,6 +210,13 @@ def bind_ambient_session(session) -> None: db_session.registry.set(session) +def make_jwt(claims: dict) -> str: + """An unsigned JWT carrying ``claims``; the providers read claims without a signature check.""" + header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode() + payload = base64.urlsafe_b64encode(json.dumps(claims).encode()).rstrip(b"=").decode() + return f"{header}.{payload}.sig" + + async def connect_provider(provider_cls, payload: dict, *, provider_email: str = "op@example.com"): """Seed the vault row a finished OAuth connect flow would leave for a subscription.""" account = await Account.get_or_create(provider_email) diff --git a/backend/tests/test_harness_auth.py b/backend/tests/test_harness_auth.py index b2624e48..68d62732 100644 --- a/backend/tests/test_harness_auth.py +++ b/backend/tests/test_harness_auth.py @@ -1,21 +1,26 @@ import base64 import json +import shlex from pathlib import Path from unittest.mock import AsyncMock import pytest -from conftest import connect_provider +from conftest import connect_provider, make_jwt from drukbox_sdk import Secret from druks.accounts.models import Account +from druks.database import db_session from druks.harnesses.claude import ClaudeHarness, _get_credentials from druks.harnesses.codex import CodexHarness from druks.harnesses.datastructures import SandboxSettings from druks.harnesses.exceptions import HarnessNotConnectedError, ProfileSettingsError from druks.harnesses.opencode import OpenCodeHarness from druks.harnesses.pi import PiHarness -from druks.harnesses.providers import AnthropicProvider, OpenAiProvider +from druks.harnesses.providers import AnthropicProvider, OpenAiProvider, jwt_claims from druks.sandbox.datastructures import HomeCopy, HomeFile +from druks.sandbox.models import SandboxIdentity from druks.secrets.models import VaultSecret +from druks.testing import seed_run +from druks_field_notes.workflows import Summarize async def _seed_claude( @@ -81,26 +86,31 @@ async def test_the_operators_claude_config_reaches_the_box_without_its_mcp_serve assert "lin_secret" not in repr(bundle) -async def test_credentials_builders_read_their_harness_config_directories(druks_db): - far_future_expiration = 4_102_444_800 - jwt_header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode() - jwt_payload = ( - base64.urlsafe_b64encode(json.dumps({"exp": far_future_expiration}).encode()) - .rstrip(b"=") - .decode() +async def _seed_codex() -> VaultSecret: + id_token = make_jwt( + { + "https://api.openai.com/auth": { + "chatgpt_account_id": "acc-1", + "chatgpt_plan_type": "pro", + }, + "email": "op@example.com", + } ) - codex_subscription = await connect_provider( + return await connect_provider( OpenAiProvider, { - "auth_mode": "chatgpt", "OPENAI_API_KEY": None, "tokens": { - "access_token": f"{jwt_header}.{jwt_payload}.sig", - "refresh_token": "R0", + "access_token": make_jwt({"exp": 4_102_444_800}), + "refresh_token": "rt-secret", + "id_token": id_token, "account_id": "acc-1", }, }, ) + + +async def test_credentials_builders_read_their_harness_config_directories(druks_db): config_root = Path("/harnesses") sandbox = SandboxSettings( service_url="x", @@ -116,9 +126,10 @@ async def test_credentials_builders_read_their_harness_config_directories(druks_ fast_mode=False, effort=None, sandbox=sandbox, - )._get_credentials(sandbox, subscription=codex_subscription) + )._get_credentials(sandbox) - assert codex_bundle.home[0].path == ".codex/auth.json" + # No credential file: each CLI reads a placeholder the box holds. + assert not any(type(entry) is HomeFile for entry in (*claude_bundle.home, *codex_bundle.home)) assert HomeCopy(".claude/settings.json", config_root / "claude/settings.json") in ( claude_bundle.home ) @@ -204,6 +215,71 @@ async def test_config_delivery_does_not_copy_host_provider_credentials( assert not invocation.env +async def test_a_codex_subscription_binds_a_custom_entry_on_chatgpt(druks_db): + subscription = await _seed_codex() + [ref] = CodexHarness.get_secret_refs(subscription) + await seed_run(db_session(), kind=Summarize.kind, run_id="run-1") + + identity, entries = await SandboxIdentity.create( + run_id="run-1", scoped_to="workflow", secret_refs=[ref] + ) + entry = entries["codex_subscription_token"].entry() + + assert ref.key == ("codex_subscription_token", subscription.id, "", "chatgpt.com") + assert (entry["host"], entry["auth_variable"], entry["auth_header"], entry["auth_prefix"]) == ( + "chatgpt.com", + "CODEX_SUBSCRIPTION_TOKEN", + "Authorization", + "Bearer ", + ) + assert entry["issuer"]["refresh"] == "1h" + assert entry["issuer"]["url"].endswith(f"/api/secrets/{identity.id}/codex_subscription_token") + + +async def test_the_codex_wrapper_writes_its_login_around_the_placeholder(druks_db): + subscription = await _seed_codex() + tokens = subscription.secrets["tokens"] + sandbox = SandboxSettings( + service_url="x", + service_token="x", + service_timeout=30.0, + image="x", + harness_config_root=Path("/harnesses"), + ) + + invocation = await CodexHarness( + model=CodexHarness.default_model, fast_mode=False, effort=None, sandbox=sandbox + ).build_invocation( + prompt="hello", + schema={"type": "object"}, + run_id="run-1", + ssh_username="druks", + identity=OpenAiProvider.get_identity(subscription), + ) + + wrapper = invocation.args[2] + # The file rides in a double-quoted shell word, so the box expands the variable. + [auth] = [json.loads(word) for word in shlex.split(wrapper) if word.startswith('{"OPENAI')] + assert auth["OPENAI_API_KEY"] is None + assert auth["tokens"]["access_token"] == "$CODEX_SUBSCRIPTION_TOKEN" + assert auth["tokens"]["refresh_token"] == "druks-placeholder" + assert auth["tokens"]["account_id"] == "acc-1" + assert auth["last_refresh"].endswith("Z") + header, _, _ = auth["tokens"]["id_token"].split(".") + assert json.loads(base64.urlsafe_b64decode(header + "=" * (-len(header) % 4))) == { + "alg": "none", + "typ": "JWT", + } + assert jwt_claims(auth["tokens"]["id_token"]) == { + "https://api.openai.com/auth": {"chatgpt_account_id": "acc-1", "chatgpt_plan_type": "pro"}, + "email": "op@example.com", + } + assert wrapper.index('chmod 600 "$HOME/.codex/auth.json"') < wrapper.index("codex exec") + for secret in (tokens["access_token"], tokens["refresh_token"], tokens["id_token"]): + assert secret not in wrapper + assert not any(type(entry) is HomeFile for entry in invocation.credentials.home) + + _ANTHROPIC_ENTRY = Secret( "sk-1", host="api.anthropic.com", diff --git a/backend/tests/test_harness_reasoning_flags.py b/backend/tests/test_harness_reasoning_flags.py index b296e7a6..1c89b890 100644 --- a/backend/tests/test_harness_reasoning_flags.py +++ b/backend/tests/test_harness_reasoning_flags.py @@ -2,30 +2,15 @@ import shlex from pathlib import Path -import pytest -from conftest import connect_provider from drukbox_sdk import Secret -from druks.accounts.models import Account from druks.harnesses.claude import ClaudeHarness from druks.harnesses.codex import CodexHarness from druks.harnesses.datastructures import SandboxSettings -from druks.harnesses.providers import AnthropicProvider, OpenAiProvider from druks.sandbox.datastructures import McpServer -from druks.secrets.datastructures import Audience -from druks.secrets.enums import SecretKind -from druks.secrets.models import VaultSecret _CODEX_MODEL = CodexHarness.default_model -@pytest.fixture(autouse=True) -async def _connected_harnesses(druks_db): - # build_invocation renders each subscription bundle from the DB row and - # raises when that harness isn't connected. - await connect_provider(AnthropicProvider, {"claudeAiOauth": {"accessToken": "t"}}) - await connect_provider(OpenAiProvider, {"tokens": {"access_token": "t"}}) - - def _sandbox_config(): return SandboxSettings( service_url="https://sb.test", @@ -46,11 +31,7 @@ async def test_claude_build_invocation_carries_every_flag(): effort="high", sandbox=_sandbox_config(), ).build_invocation( - subscription=await VaultSecret.lookup( - SecretKind.SUBSCRIPTION, - Audience.provider("anthropic"), - (await Account.get_default()).id, - ), + identity={"email": "op@example.com"}, prompt="hello", schema=schema, run_id="run-1", @@ -106,9 +87,7 @@ async def test_codex_build_invocation_carries_every_flag(): effort="high", sandbox=_sandbox_config(), ).build_invocation( - subscription=await VaultSecret.lookup( - SecretKind.SUBSCRIPTION, Audience.provider("openai"), (await Account.get_default()).id - ), + identity={"email": "op@example.com", "account_id": "acc-1"}, prompt="hello", schema={"type": "object"}, run_id="run-1", diff --git a/backend/tests/test_profiles.py b/backend/tests/test_profiles.py index 70cc6d66..12c3daf8 100644 --- a/backend/tests/test_profiles.py +++ b/backend/tests/test_profiles.py @@ -1,7 +1,13 @@ from types import SimpleNamespace import pytest -from conftest import PROFILE_PROBE, ProfileOutput, connect_anthropic_subscription +from conftest import ( + PROFILE_PROBE, + ProfileOutput, + connect_anthropic_subscription, + connect_provider, + make_jwt, +) from drukbox_sdk import Secret from druks import agents from druks.accounts.models import Account @@ -10,10 +16,12 @@ from druks.database import db_session from druks.durable.models import AgentCall from druks.harnesses.claude import ClaudeHarness +from druks.harnesses.codex import CodexHarness from druks.harnesses.exceptions import HarnessNotConnectedError, ProfileSettingsError from druks.harnesses.models import ProviderCatalog from druks.harnesses.opencode import OpenCodeHarness from druks.harnesses.profiles import check_profile, get_profile +from druks.harnesses.providers import OpenAiProvider from druks.sandbox.constants import MAX_AGENT_TIMEOUT_SECONDS from druks.secrets.datastructures import Audience from druks.secrets.enums import SecretKind @@ -134,6 +142,7 @@ async def test_a_subscription_agent_runs_as_its_actor_or_the_default_account(dru assert as_actor.subscription.id == actor.id assert as_actor.charged_account_id == actor.account_id assert unattended.subscription.id == default_subscription.id + assert as_actor.identity == {"email": "b@example.com"} assert (as_actor.secrets, as_actor.secrets_id) == ({}, actor.id) [secret] = as_actor.secret_refs assert secret.key == ("anthropic", actor.id, "", "") @@ -142,6 +151,43 @@ async def test_a_subscription_agent_runs_as_its_actor_or_the_default_account(dru assert (as_actor.effort, as_actor.timeout, as_actor.fast_mode) == ("high", 1800, False) +async def test_a_codex_subscription_profile_carries_its_login_facts_and_its_ref(druks_db): + id_token = make_jwt( + { + "https://api.openai.com/auth": { + "chatgpt_account_id": "acc-1", + "chatgpt_plan_type": "pro", + }, + "email": "a@example.com", + } + ) + subscription = await connect_provider( + OpenAiProvider, + { + "OPENAI_API_KEY": None, + "tokens": { + "access_token": make_jwt({"exp": 4_102_444_800}), + "refresh_token": "R0", + "id_token": id_token, + "account_id": "acc-1", + }, + }, + provider_email="a@example.com", + ) + await SettingsOverride.set_agent_harness(PROFILE_PROBE.id, "codex") + await SettingsOverride.set_agent_model(PROFILE_PROBE.id, "openai/gpt-5.5") + + profile = await get_profile(PROFILE_PROBE.id, subscription.account_id) + + assert profile.harness_class is CodexHarness + # The facts the box's login names come from the row and its id token; the + # tokens stay on the server. + assert profile.identity == {"email": "a@example.com", "account_id": "acc-1", "plan": "pro"} + [ref] = profile.secret_refs + assert ref.key == ("codex_subscription_token", subscription.id, "", "chatgpt.com") + assert profile.secrets_id == subscription.id + + async def test_a_subscription_agent_refuses_without_the_actors_own_subscription(druks_db): await connect_anthropic_subscription("a@example.com") await _key() @@ -162,7 +208,7 @@ async def test_a_key_agent_runs_on_the_installations_key_for_anyone(druks_db): # Claude reads the key from a placeholder in the VM, never from its invocation. assert (as_actor.secrets, as_actor.subscription) == ({"anthropic": _SHARED_ENTRY}, None) - assert unattended.secrets == {"anthropic": _SHARED_ENTRY} + assert (unattended.secrets, unattended.identity) == ({"anthropic": _SHARED_ENTRY}, {}) # The entries' identity is the pasted key, with no secret material. assert as_actor.secrets_id == f"anthropic.{pasted.updated_at:%Y%m%dT%H%M%S}" assert "sk-shared" not in as_actor.secrets_id diff --git a/backend/tests/test_provider_auth.py b/backend/tests/test_provider_auth.py index 623b925d..026f6e0b 100644 --- a/backend/tests/test_provider_auth.py +++ b/backend/tests/test_provider_auth.py @@ -574,12 +574,14 @@ class MinimalProvider(pbase.Provider): assert calls == [] -async def _bound_identity(subscription, *, host_id: str, run_id: str) -> SandboxIdentity: +async def _bound_identity( + subscription, *, host_id: str, run_id: str, name: str = "anthropic", host: str = "" +) -> SandboxIdentity: await seed_run(db_session(), kind=Summarize.kind, run_id=run_id) identity, _ = await SandboxIdentity.create( run_id=run_id, scoped_to="workflow", - secret_refs=[SecretRef(name="anthropic", secret_id=subscription.id)], + secret_refs=[SecretRef(name=name, secret_id=subscription.id, host=host)], ) await identity.bind(host_id) return identity @@ -801,3 +803,133 @@ async def fake_get(self, url, *, headers=None, **_kwargs): AnthropicProvider._TOKEN_URL, _REFRESH_URL.format(host_id="host-a"), ] + + +def _jwt_in(delta: timedelta) -> str: + return _jwt(int(_in(delta).timestamp())) + + +def _codex_refreshed() -> dict: + return {"access_token": _jwt_in(timedelta(days=9)), "refresh_token": "R1", "id_token": "id-1"} + + +_CODEX_REFRESH_URL = "http://127.0.0.1:8781/refresh/{host_id}/codex_subscription_token" + + +async def test_a_codex_fetch_answers_a_fresh_token_with_its_exp_without_a_provider_call_or_the_gate( + monkeypatch, druks_db +): + # Above the 24-hour margin. The answer carries the JWT exp as its expiry. + exp = int(_in(timedelta(hours=48)).timestamp()) + connection = await _seed_codex(access=_jwt(exp)) + calls = _mock_post(monkeypatch, _resp(200, _codex_refreshed())) + monkeypatch.setattr(pbase.gate, "shut", _no_gate) + + token = await OpenAiProvider.issue_token(connection.id) + + assert (token.access_token, token.expires_at) == ( + _jwt(exp), + datetime.fromtimestamp(exp, tz=UTC), + ) + assert calls == [] + + +async def test_two_codex_fetches_inside_the_margin_rotate_once_and_request_once( + monkeypatch, druks_db +): + connection = await _seed_codex(access=_jwt_in(timedelta(hours=12)), refresh="R0") + await _bound_identity( + connection, + host_id="host-other", + run_id="run-other", + name="codex_subscription_token", + host="chatgpt.com", + ) + refreshed = _codex_refreshed() + calls = _mock_post(monkeypatch, _resp(200, refreshed)) + + first = await OpenAiProvider.issue_token(connection.id, except_host_id="host-mine") + second = await OpenAiProvider.issue_token(connection.id, except_host_id="host-mine") + + assert first.access_token == second.access_token == refreshed["access_token"] + assert calls[0]["json"]["refresh_token"] == "R0" + assert [call["url"] for call in calls] == [ + OpenAiProvider._TOKEN_URL, + _CODEX_REFRESH_URL.format(host_id="host-other"), + ] + + +async def test_a_codex_fetch_on_a_busy_subscription_answers_the_current_token( + monkeypatch, druks_db +): + # Inside the margin, above the call horizon: the call in flight keeps its token. + current = _jwt_in(timedelta(hours=12)) + connection = await _seed_codex(access=current, refresh="R0") + calls = _mock_post(monkeypatch, _resp(200, _codex_refreshed())) + + async with gate.use(connection.id, "call-1"): + token = await OpenAiProvider.issue_token(connection.id) + + assert token.access_token == current + assert calls == [] + + +async def test_a_codex_fetch_rotates_a_busy_subscription_once_urgent(monkeypatch, druks_db): + connection = await _seed_codex(access=_jwt_in(timedelta(minutes=30)), refresh="R0") + refreshed = _codex_refreshed() + calls = _mock_post(monkeypatch, _resp(200, refreshed)) + + async with gate.use(connection.id, "call-1"): + token = await OpenAiProvider.issue_token(connection.id) + + assert token.access_token == refreshed["access_token"] + assert calls[0]["json"]["refresh_token"] == "R0" + + +async def test_a_codex_fetch_waits_out_a_shut_gate_then_answers_the_stored_token( + monkeypatch, druks_db +): + monkeypatch.setattr(gate, "_POLL", 0.01) + connection = await _seed_codex(access=_jwt_in(timedelta(hours=12)), refresh="R0") + calls = _mock_post(monkeypatch, _resp(200, _codex_refreshed())) + client = druks.redis.get_client() + rotating = f"druks:sandbox:rotating:{connection.id}" + await client.set(rotating, "1", ex=60) + session = db_session() + stored = _jwt_in(timedelta(hours=48)) + + async def other_rotator() -> None: + # The holder advances the row, then reopens the gate. + await asyncio.sleep(0.03) + await session.execute( + update(VaultSecret) + .where(VaultSecret.id == connection.id) + .values(secrets=_codex_payload(access=stored, refresh="R1")) + ) + await client.delete(rotating) + + holder = asyncio.create_task(other_rotator()) + token = await OpenAiProvider.issue_token(connection.id) + await holder + + assert token.access_token == stored + assert calls == [] + + +async def test_a_failed_codex_rotation_answers_the_live_token(monkeypatch, druks_db): + live = _jwt_in(timedelta(hours=12)) + connection = await _seed_codex(access=live, refresh="R0") + _mock_post(monkeypatch, httpx.ConnectError("boom")) + + token = await OpenAiProvider.issue_token(connection.id) + + assert token.access_token == live + + +async def test_an_expired_codex_token_with_a_failed_rotation_answers_nothing(monkeypatch, druks_db): + connection = await _seed_codex(access=_jwt_in(-timedelta(minutes=1)), refresh="R0") + _mock_post(monkeypatch, httpx.ConnectError("boom")) + + with pytest.raises(OAuthTokenError) as error: + await OpenAiProvider.issue_token(connection.id) + assert error.value.tag == "token_expired" diff --git a/backend/tests/test_sandboxed_harness.py b/backend/tests/test_sandboxed_harness.py index c92d3ffb..3d3090a5 100644 --- a/backend/tests/test_sandboxed_harness.py +++ b/backend/tests/test_sandboxed_harness.py @@ -9,7 +9,7 @@ from typing import Any import pytest -from conftest import PROFILE_PROBE, connect_provider, installation_key +from conftest import PROFILE_PROBE, connect_provider, installation_key, make_jwt from druks.durable.enums import AgentCallStatus from druks.harnesses.base import Harness from druks.harnesses.claude import ClaudeHarness @@ -26,7 +26,7 @@ Retry, ) from druks.harnesses.profiles import Profile, get_profile -from druks.harnesses.providers import AnthropicProvider +from druks.harnesses.providers import AnthropicProvider, OpenAiProvider from druks.sandbox.datastructures import ( AgentInvocation, AgentResult, @@ -363,7 +363,7 @@ def parse(self, result: Any, *, artifact_dir: Path, run_id: str) -> Any: artifact_dir=ctx.artifact_dir, timeout=60, call_id="call-7", - subscription=SimpleNamespace(provider="anthropic"), + identity={"email": "op@example.com"}, extra_env={"GITHUB_MCP_TOKEN": "ghs_x"}, ) @@ -534,6 +534,7 @@ def agent_profile(): api_key=None, secrets={}, secret_refs=[], + identity={"email": "op@example.com"}, billing="subscription", effort="high", timeout=60, @@ -699,3 +700,62 @@ async def test_claude_subscription_token_stays_on_the_server( assert secret not in start.kwargs["stdin_data"].decode() for artifact in (ctx.artifact_dir / "call-9").iterdir(): assert secret not in artifact.read_text() + + +async def test_codex_subscription_token_stays_on_the_server( + ctx: SimpleNamespace, druks_db, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + """Under subscription billing the VM holds a placeholder in CODEX_SUBSCRIPTION_TOKEN, + and the run wrapper writes it into auth.json. The token, its refresh token, and its + id token reach no invocation, VM file, artifact, or result.""" + tokens = { + "access_token": make_jwt({"exp": int((datetime.now(UTC) + timedelta(days=9)).timestamp())}), + "refresh_token": "rt-secret", + "id_token": make_jwt({"email": "op@example.com"}), + "account_id": "acc-1", + } + await connect_provider(OpenAiProvider, {"OPENAI_API_KEY": None, "tokens": tokens}) + await SettingsOverride.set_agent_harness(PROFILE_PROBE.id, "codex") + await SettingsOverride.set_agent_model(PROFILE_PROBE.id, "openai/gpt-5.5") + await SettingsOverride.set_agent_billing(PROFILE_PROBE.id, "subscription") + profile = await get_profile(PROFILE_PROBE.id, None) + # Codex leaves its result in the box; the fake download pulls nothing, so + # the file is in place before the run. + (ctx.artifact_dir / "call-9").mkdir() + (ctx.artifact_dir / "call-9" / "output.json").write_text('{"ok": true}') + run = _FakeRun(stdout_chunks=[b'{"type":"thread.started"}\n']) + sandbox = _fake_sandbox(run) + sandbox.run_prompt = functools.partial(Host.run_prompt, sandbox) + sandbox._exec = functools.partial(Host._exec, sandbox) + settings = SimpleNamespace( + sandbox=SimpleNamespace(service_url="x", service_token="x", timeout=30.0, image="x"), + harness_config_root=tmp_path / "harnesses", + skills_dir=None, + ) + monkeypatch.setattr("druks.sandbox.host.load_settings", lambda: settings) + + result = await Host.run_agent( + sandbox, + agent="evaluate", + profile=profile, + prompt="p", + schema={"type": "object"}, + artifact_dir=ctx.artifact_dir, + call_id="call-9", + ) + + assert result.status is AgentCallStatus.SUCCEEDED + assert result.output == {"ok": True} + [secret] = profile.secret_refs + assert secret.key == ("codex_subscription_token", profile.subscription.id, "", "chatgpt.com") + [start] = sandbox.calls + assert not start.kwargs["extra_env"] + bundle = start.kwargs["credentials_bundle"] + assert not any(type(entry) is HomeFile for entry in bundle.home) + assert "$CODEX_SUBSCRIPTION_TOKEN" in " ".join(start.kwargs["cmd"]) + for secret in (tokens["access_token"], tokens["refresh_token"], tokens["id_token"]): + assert secret not in " ".join(start.kwargs["cmd"]) + assert secret not in start.kwargs["stdin_data"].decode() + for artifact in (ctx.artifact_dir / "call-9").iterdir(): + assert secret not in artifact.read_text() + assert secret not in repr(result) and secret not in repr(profile) diff --git a/docs/concepts.md b/docs/concepts.md index 00dd0544..ce51a2af 100644 --- a/docs/concepts.md +++ b/docs/concepts.md @@ -194,9 +194,13 @@ holds: the Drukbox name, the vault row, and the resource. The row can be the GitHub App key, a pasted token, or a provider subscription. A replay after a crash finds the sandbox through the run's live identity with the same scope. Drukbox holds the bearer in the sandbox's issuer entry and fetches the token -from the Druks issuer, `GET /api/secrets//`. The issuer -answers a fresh token at once. A token inside its refresh margin rotates first, while the -subscription is idle or the token is urgent. One rotator runs at a time, and +from the Druks issuer, `GET /api/secrets//`. The sandbox +sees a placeholder in the variable the entry names. Claude reads it from +`ANTHROPIC_AUTH_TOKEN`. The Codex run wrapper writes it into +`~/.codex/auth.json` beside a sentinel refresh token, so Codex never refreshes +inside the sandbox. The issuer answers a fresh token at once. A token inside +its refresh margin rotates first, while the subscription is idle or the token +is urgent. One rotator runs at a time, and new calls wait for it. After a rotation, Druks requests a refresh from the secrets exchange for every live sandbox on that subscription. A provider can revoke the previous token at the rotation. Druks revokes the identity when it diff --git a/docs/configuration.md b/docs/configuration.md index d88c2aa8..5215b278 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -309,28 +309,42 @@ optional. It reports pending setup if the selected tracker lacks a connection. Druks registers two subscription providers, `anthropic` and `openai`. Each also accepts an API key. Both connect from **Settings → Providers**. The connection flow stores each credential in Postgres. Druks refreshes a -subscription token on a schedule. A `claude` sandbox holds a placeholder in -`ANTHROPIC_AUTH_TOKEN` and never the token. Drukbox fetches the token from the +subscription token on a schedule. A sandbox holds a placeholder for the +subscription token and never the token. Drukbox fetches the token from the Druks issuer through the sandbox's identity, and the secrets proxy swaps the -placeholder on each request. See +placeholder on each request to the entry's host. See [sandbox identities and the issuer](concepts.md#agents-harnesses-workspaces-and-sandboxes). -Codex still receives its subscription file inside the sandbox. Druks does -not copy a host login. This is a capability connection for the requesting -account. In a fresh `none`-mode install, the first completed subscription -connection also creates the operator account. See +Druks does not copy a host login. This is a capability connection for the +requesting account. In a fresh `none`-mode install, the first completed +subscription connection also creates the operator account. See [access control](#public-urls-and-access-control). +A `claude` sandbox reads its placeholder from `ANTHROPIC_AUTH_TOKEN`. A +`codex` sandbox reads its placeholder from `CODEX_SUBSCRIPTION_TOKEN`. The +Codex run wrapper writes `~/.codex/auth.json` from that variable before the +command. The file carries the account id, the sentinel refresh token +`druks-placeholder`, and an unsigned id token with the account id, the plan, +and the email. Codex sends the placeholder to `chatgpt.com` on every request. +Codex never refreshes it: the one refresh it attempts after a 401 fails on +the sentinel, and the turn ends. The real refresh token and id token stay in +Postgres. + An API key never enters the sandbox. Druks gives the key to Drukbox as a secret entry when it creates the sandbox. The sandbox holds a placeholder in the variable the entry names, and the CLI reads it from the environment. The Drukbox secrets proxy swaps the placeholder for the key in the entry's header on each request to the entry's host. -| Harness | Provider | Variable | Host | Header | +| Harness | Credential | Variable | Host | Header | | --- | --- | --- | --- | --- | -| `claude`, `pi`, `opencode` | Anthropic | `ANTHROPIC_API_KEY` | `api.anthropic.com` | `x-api-key` | -| `pi`, `opencode` | OpenAI | `OPENAI_API_KEY` | `api.openai.com` | `Authorization: Bearer` | -| `codex` | OpenAI | `CODEX_API_KEY` | `api.openai.com` | `Authorization: Bearer` | +| `claude` | Anthropic subscription | `ANTHROPIC_AUTH_TOKEN` | `api.anthropic.com` | `Authorization: Bearer` | +| `codex` | OpenAI subscription | `CODEX_SUBSCRIPTION_TOKEN` | `chatgpt.com` | `Authorization: Bearer` | +| `claude`, `pi`, `opencode` | Anthropic API key | `ANTHROPIC_API_KEY` | `api.anthropic.com` | `x-api-key` | +| `pi`, `opencode` | OpenAI API key | `OPENAI_API_KEY` | `api.openai.com` | `Authorization: Bearer` | +| `codex` | OpenAI API key | `CODEX_API_KEY` | `api.openai.com` | `Authorization: Bearer` | + +The `ANTHROPIC_AUTH_TOKEN` and `OPENAI_API_KEY` entries come from the Drukbox +catalog. Druks declares the other entries with their host and header. The Compose stack runs the secrets proxy on every provider but docker-sbx. See [the secrets exchange and the secrets proxy](deployment.md#the-secrets-exchange-and-the-secrets-proxy).