From 0e29bcc8297f2ebcb89121f7bd70c562119621f2 Mon Sep 17 00:00:00 2001 From: claude Date: Tue, 2 Jun 2026 04:53:23 +0000 Subject: [PATCH] refactor(egress): use provisioned_env instead of sentinel for Codex token (PRD 0030) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `provisioned_env: dict[str, str]` to `AgentProvisionPlan`. When `forward_host_credentials=True`, `agent_provision_plan` reads the host Codex access token at prepare time and stores it under `CODEX_HOST_CREDENTIAL_TOKEN_REF`. Both backends merge `provisioned_env` over `os.environ` before calling `egress_resolve_token_values`, so the token slot resolves like any other manifest-declared token ref. Removes `egress_resolve_token_values_with_provider` and the sentinel `continue` skip from `egress_resolve_token_values`. The function is now fully generic — it neither knows nor cares about provider identity. --- bot_bottle/agent_provider.py | 11 +++- bot_bottle/backend/docker/launch.py | 10 ++-- bot_bottle/backend/smolmachines/launch.py | 10 ++-- bot_bottle/egress.py | 29 ---------- tests/unit/test_agent_provider.py | 31 +++++++++++ tests/unit/test_egress.py | 65 +---------------------- 6 files changed, 49 insertions(+), 107 deletions(-) diff --git a/bot_bottle/agent_provider.py b/bot_bottle/agent_provider.py index cd271dc..c610116 100644 --- a/bot_bottle/agent_provider.py +++ b/bot_bottle/agent_provider.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import Literal -from .codex_auth import write_codex_dummy_auth_file +from .codex_auth import codex_host_access_token, write_codex_dummy_auth_file from .egress import CODEX_HOST_CREDENTIAL_TOKEN_REF, EgressRoute @@ -92,6 +92,7 @@ class AgentProvisionPlan: verify: tuple[AgentProvisionCommand, ...] = () egress_routes: tuple[EgressRoute, ...] = () hidden_env_names: frozenset[str] = field(default_factory=frozenset) + provisioned_env: dict[str, str] = field(default_factory=dict) _REPO_ROOT = Path(__file__).resolve().parent.parent @@ -139,6 +140,7 @@ def agent_provision_plan( runtime = runtime_for(template) resolved_guest_env = dict(guest_env or {}) env_vars: dict[str, str] = {} + provisioned_env: dict[str, str] = {} dirs: list[AgentProvisionDir] = [] files: list[AgentProvisionFile] = [] pre_copy: list[AgentProvisionCommand] = [] @@ -169,8 +171,12 @@ def agent_provision_plan( tls_passthrough=True, )) if forward_host_credentials: + _host_env = host_env or dict(os.environ) + provisioned_env[CODEX_HOST_CREDENTIAL_TOKEN_REF] = codex_host_access_token( + _host_env, + ) auth_file = state_dir / "codex-auth.json" - write_codex_dummy_auth_file(auth_file, host_env or dict(os.environ)) + write_codex_dummy_auth_file(auth_file, _host_env) files.append(AgentProvisionFile(auth_file, f"{auth_dir}/auth.json")) pre_copy.append(AgentProvisionCommand(( "find", auth_dir, @@ -220,6 +226,7 @@ def agent_provision_plan( verify=tuple(verify), egress_routes=tuple(egress_routes), hidden_env_names=hidden_env_names, + provisioned_env=provisioned_env, ) diff --git a/bot_bottle/backend/docker/launch.py b/bot_bottle/backend/docker/launch.py index a1a6f5e..9e28a00 100644 --- a/bot_bottle/backend/docker/launch.py +++ b/bot_bottle/backend/docker/launch.py @@ -42,7 +42,7 @@ from contextlib import ExitStack, contextmanager from pathlib import Path from typing import Callable, Generator -from ...egress import egress_resolve_token_values_with_provider +from ...egress import egress_resolve_token_values from ...log import info from . import network as network_mod from . import util as docker_mod @@ -176,11 +176,9 @@ def launch( # Step 7: compose up. Token values + the OAuth placeholder # flow through subprocess env; the compose file holds only # bare names for the secret-carrying entries. - bottle = plan.spec.manifest.bottle_for(plan.spec.agent_name) - token_values = egress_resolve_token_values_with_provider( - plan.egress_plan.token_env_map, - bottle.agent_provider.forward_host_credentials, - dict(os.environ), + effective_env = {**dict(os.environ), **plan.agent_provision.provisioned_env} + token_values = egress_resolve_token_values( + plan.egress_plan.token_env_map, effective_env, ) compose_env: dict[str, str] = { **os.environ, diff --git a/bot_bottle/backend/smolmachines/launch.py b/bot_bottle/backend/smolmachines/launch.py index 629eb2f..14d5a33 100644 --- a/bot_bottle/backend/smolmachines/launch.py +++ b/bot_bottle/backend/smolmachines/launch.py @@ -28,7 +28,7 @@ from typing import Callable, Generator from ...egress import ( EGRESS_ROUTES_IN_CONTAINER, - egress_resolve_token_values_with_provider, + egress_resolve_token_values, ) from ...pipelock import ( PIPELOCK_CA_CERT_IN_CONTAINER, @@ -423,12 +423,8 @@ def _resolve_token_env( """Resolve the egress token env-var values from the host's environ so they reach the bundle's process env via docker's `-e NAME` inheritance. Empty when no routes declare auth.""" - bottle = plan.spec.manifest.bottle_for(plan.spec.agent_name) - return egress_resolve_token_values_with_provider( - plan.egress_plan.token_env_map, - bottle.agent_provider.forward_host_credentials, - host_env, - ) + effective_env = {**host_env, **plan.agent_provision.provisioned_env} + return egress_resolve_token_values(plan.egress_plan.token_env_map, effective_env) def _ensure_smolmachine(image_ref: str, *, dockerfile: str = "") -> Path: diff --git a/bot_bottle/egress.py b/bot_bottle/egress.py index 19988d4..0e1bc0d 100644 --- a/bot_bottle/egress.py +++ b/bot_bottle/egress.py @@ -29,7 +29,6 @@ from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING -from .codex_auth import codex_host_access_token from .log import die if TYPE_CHECKING: @@ -342,8 +341,6 @@ def egress_resolve_token_values( a sealed mapping without touching `os.environ`.""" out: dict[str, str] = {} for token_env, token_ref in token_env_map.items(): - if token_ref == CODEX_HOST_CREDENTIAL_TOKEN_REF: - continue value = host_env.get(token_ref) if value is None: die( @@ -361,31 +358,6 @@ def egress_resolve_token_values( return out -def egress_resolve_token_values_with_provider( - token_env_map: dict[str, str], - forward_host_credentials: bool, - host_env: dict[str, str], -) -> dict[str, str]: - """Resolve all egress token env-var values, including the optional - Codex host credential slot. - - Combines `egress_resolve_token_values` (manifest-declared token refs) - with the `forward_host_credentials` path (Codex ChatGPT bearer). - Returns an empty dict when `token_env_map` is empty. - - Pure function: `host_env` is passed in so tests can use a sealed - mapping without touching `os.environ`.""" - if not token_env_map: - return {} - token_values = egress_resolve_token_values(token_env_map, host_env) - if forward_host_credentials: - access_token = codex_host_access_token(host_env) - for token_env, token_ref in token_env_map.items(): - if token_ref == CODEX_HOST_CREDENTIAL_TOKEN_REF: - token_values[token_env] = access_token - return token_values - - class Egress(ABC): """The per-bottle egress proxy. Encapsulates the host-side prepare (route lift + routes.yaml render + token-env-map derivation); the @@ -429,7 +401,6 @@ __all__ = [ "egress_manifest_routes", "egress_render_routes", "egress_resolve_token_values", - "egress_resolve_token_values_with_provider", "egress_routes_for_bottle", "egress_token_env_map", ] diff --git a/tests/unit/test_agent_provider.py b/tests/unit/test_agent_provider.py index 8d4f4b6..3f90db5 100644 --- a/tests/unit/test_agent_provider.py +++ b/tests/unit/test_agent_provider.py @@ -145,6 +145,37 @@ class TestAgentProviderRuntime(unittest.TestCase): self.assertNotIn("CLAUDE_CODE_OAUTH_TOKEN", plan.env_vars) self.assertEqual(frozenset(), plan.hidden_env_names) + def test_codex_forward_host_credentials_populates_provisioned_env(self): + access = _jwt(2000000000) + with tempfile.TemporaryDirectory(prefix="bb-provider.") as tmp: + home = Path(tmp) / "host-codex" + home.mkdir() + (home / "auth.json").write_text(json.dumps({ + "auth_mode": "chatgpt", + "tokens": {"access_token": access}, + })) + plan = agent_provision_plan( + template="codex", + dockerfile="", + state_dir=Path(tmp), + forward_host_credentials=True, + host_env={"CODEX_HOME": str(home)}, + ) + self.assertEqual( + {CODEX_HOST_CREDENTIAL_TOKEN_REF: access}, + plan.provisioned_env, + ) + + def test_codex_without_forward_host_credentials_has_empty_provisioned_env(self): + with tempfile.TemporaryDirectory(prefix="bb-provider.") as tmp: + plan = agent_provision_plan( + template="codex", + dockerfile="", + state_dir=Path(tmp), + forward_host_credentials=False, + ) + self.assertEqual({}, plan.provisioned_env) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_egress.py b/tests/unit/test_egress.py index 55b89c4..f1129cc 100644 --- a/tests/unit/test_egress.py +++ b/tests/unit/test_egress.py @@ -2,7 +2,6 @@ resolution (PRD 0017).""" import unittest -import unittest.mock from bot_bottle.egress import ( CODEX_HOST_CREDENTIAL_TOKEN_REF, @@ -10,7 +9,6 @@ from bot_bottle.egress import ( egress_manifest_routes, egress_render_routes, egress_resolve_token_values, - egress_resolve_token_values_with_provider, egress_routes_for_bottle, egress_token_env_map, ) @@ -343,72 +341,13 @@ class TestResolveTokenValues(unittest.TestCase): {"GH_PAT": ""}, ) - def test_codex_host_credential_ref_is_resolved_by_launch(self): + def test_codex_host_credential_ref_resolved_via_provisioned_env(self): out = egress_resolve_token_values( {"EGRESS_TOKEN_0": CODEX_HOST_CREDENTIAL_TOKEN_REF}, - {}, + {CODEX_HOST_CREDENTIAL_TOKEN_REF: "codex-access-token"}, ) - self.assertEqual({}, out) - - -class TestResolveTokenValuesWithProvider(unittest.TestCase): - def test_empty_map_returns_empty(self): - out = egress_resolve_token_values_with_provider({}, False, {}) - self.assertEqual({}, out) - - def test_empty_map_with_forward_credentials_returns_empty(self): - # forward_host_credentials=True but no slots → no codex call needed. - out = egress_resolve_token_values_with_provider({}, True, {}) - self.assertEqual({}, out) - - def test_manifest_tokens_resolved_without_forward_credentials(self): - out = egress_resolve_token_values_with_provider( - {"EGRESS_TOKEN_0": "GH_PAT"}, - False, - {"GH_PAT": "ghp_secret"}, - ) - self.assertEqual({"EGRESS_TOKEN_0": "ghp_secret"}, out) - - def test_codex_token_slotted_in_when_forward_credentials_and_matching_ref(self): - with unittest.mock.patch( - "bot_bottle.egress.codex_host_access_token", - return_value="codex-access-token", - ): - out = egress_resolve_token_values_with_provider( - {"EGRESS_TOKEN_0": CODEX_HOST_CREDENTIAL_TOKEN_REF}, - True, - {}, - ) self.assertEqual({"EGRESS_TOKEN_0": "codex-access-token"}, out) - def test_codex_token_not_slotted_when_no_matching_ref(self): - # forward_host_credentials=True but no CODEX_HOST_CREDENTIAL_TOKEN_REF - # slot in the map → manifest tokens only; Codex token is fetched but - # nothing to slot it into. - with unittest.mock.patch( - "bot_bottle.egress.codex_host_access_token", - return_value="codex-access-token", - ): - out = egress_resolve_token_values_with_provider( - {"EGRESS_TOKEN_0": "GH_PAT"}, - True, - {"GH_PAT": "ghp_secret"}, - ) - self.assertEqual({"EGRESS_TOKEN_0": "ghp_secret"}, out) - - def test_codex_not_called_when_forward_credentials_false(self): - called = [] - with unittest.mock.patch( - "bot_bottle.egress.codex_host_access_token", - side_effect=lambda *_: called.append(1) or "tok", - ): - egress_resolve_token_values_with_provider( - {"EGRESS_TOKEN_0": "GH_PAT"}, - False, - {"GH_PAT": "ghp_secret"}, - ) - self.assertEqual([], called) - if __name__ == "__main__": unittest.main()