Issue Codex subscription tokens through a chatgpt.com entry (DRU-476)

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

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:

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 <get_runs_root>/<run_id>/. 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]:

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:

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:

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,14 +106,17 @@ 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:
             label = await provider_label(provider_id)
             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.

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,

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(

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)