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 e811cbe..9e28a00 100644 --- a/bot_bottle/backend/docker/launch.py +++ b/bot_bottle/backend/docker/launch.py @@ -42,11 +42,7 @@ from contextlib import ExitStack, contextmanager from pathlib import Path from typing import Callable, Generator -from ...codex_auth import codex_host_access_token -from ...egress import ( - CODEX_HOST_CREDENTIAL_TOKEN_REF, - egress_resolve_token_values, -) +from ...egress import egress_resolve_token_values from ...log import info from . import network as network_mod from . import util as docker_mod @@ -180,18 +176,10 @@ 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. - token_values: dict[str, str] = {} - if plan.egress_plan.routes: - token_values = egress_resolve_token_values( - plan.egress_plan.token_env_map, dict(os.environ), - ) - if plan.spec.manifest.bottle_for( - plan.spec.agent_name, - ).agent_provider.forward_host_credentials: - access_token = codex_host_access_token(dict(os.environ)) - for token_env, token_ref in plan.egress_plan.token_env_map.items(): - if token_ref == CODEX_HOST_CREDENTIAL_TOKEN_REF: - token_values[token_env] = access_token + 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, **plan.forwarded_env, diff --git a/bot_bottle/backend/smolmachines/launch.py b/bot_bottle/backend/smolmachines/launch.py index da73dae..14d5a33 100644 --- a/bot_bottle/backend/smolmachines/launch.py +++ b/bot_bottle/backend/smolmachines/launch.py @@ -26,9 +26,7 @@ from contextlib import ExitStack, contextmanager from pathlib import Path from typing import Callable, Generator -from ...codex_auth import codex_host_access_token from ...egress import ( - CODEX_HOST_CREDENTIAL_TOKEN_REF, EGRESS_ROUTES_IN_CONTAINER, egress_resolve_token_values, ) @@ -146,7 +144,7 @@ def launch( # spec's ports_to_publish list expands depending on which # daemons the agent needs to reach from the smolvm guest. bundle_spec = _bundle_launch_spec(plan, network, loopback_ip) - token_env = _resolve_token_env(plan, os.environ) + token_env = _resolve_token_env(plan, dict(os.environ)) _bundle.ensure_bundle_image(bundle_spec.image) _bundle.start_bundle(bundle_spec, env={**os.environ, **token_env}) stack.callback(_bundle.stop_bundle, plan.slug) @@ -420,24 +418,13 @@ def _bundle_launch_spec( def _resolve_token_env( - plan: SmolmachinesBottlePlan, host_env: object + plan: SmolmachinesBottlePlan, host_env: dict[str, str], ) -> dict[str, str]: """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.""" - ep = plan.egress_plan - if not ep.routes: - return {} - env = dict(host_env) - token_values = egress_resolve_token_values(ep.token_env_map, env) - if plan.spec.manifest.bottle_for( - plan.spec.agent_name, - ).agent_provider.forward_host_credentials: - access_token = codex_host_access_token(env) - for token_env, token_ref in ep.token_env_map.items(): - if token_ref == CODEX_HOST_CREDENTIAL_TOKEN_REF: - token_values[token_env] = access_token - return token_values + 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 000d5e1..0e1bc0d 100644 --- a/bot_bottle/egress.py +++ b/bot_bottle/egress.py @@ -341,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( diff --git a/docs/prds/0030-deduplicate-egress-token-resolution.md b/docs/prds/0030-deduplicate-egress-token-resolution.md new file mode 100644 index 0000000..2067347 --- /dev/null +++ b/docs/prds/0030-deduplicate-egress-token-resolution.md @@ -0,0 +1,121 @@ +# PRD 0030: Deduplicate egress token resolution across backends + +- **Status:** Active +- **Author:** didericis-claude +- **Created:** 2026-06-02 +- **Issue:** #118 + +## Summary + +Eliminate the duplicated egress token resolution block — which resolved +manifest-declared tokens and the Codex host credential — by moving +provider-specific token reading into `AgentProvisionPlan.provisioned_env` +at prepare time, and having both backends merge that map into `os.environ` +before calling the now-generic `egress_resolve_token_values`. + +## Problem + +The same logic block appeared in two places: + +- `bot_bottle/backend/docker/launch.py` (~line 183): inline inside + `launch()`, building `token_values` before `compose_up`. +- `bot_bottle/backend/smolmachines/launch.py` (~line 422): the private + `_resolve_token_env` helper, called before `_bundle.start_bundle`. + +Both blocks: +1. Short-circuit to `{}` when there are no egress routes. +2. Call `egress_resolve_token_values(token_env_map, host_env)` to + resolve ordinary manifest-declared token refs. +3. Check `agent_provider.forward_host_credentials` and, when true, + call `codex_host_access_token` and slot the result into any + `token_env` whose `token_ref` is `CODEX_HOST_CREDENTIAL_TOKEN_REF`. + +The duplication means any change to step 3 must be applied twice. +PRD 0029, which introduced `forward_host_credentials`, already had to +wire both backends; the next change would too. This is a near-certain +future sync bug. + +`egress_resolve_token_values` also carried a sentinel `continue` skip +for `CODEX_HOST_CREDENTIAL_TOKEN_REF`, which tied an otherwise generic +egress helper to a Codex-specific contract. + +## Goals / Success Criteria + +- The `forward_host_credentials` resolution logic exists in exactly one + place in the codebase. +- Both `docker/launch.py` and `smolmachines/launch.py` call + `egress_resolve_token_values` with no provider-specific arguments. +- `egress_resolve_token_values` is fully generic — it neither knows nor + cares about provider identity or the `CODEX_HOST_CREDENTIAL_TOKEN_REF` + sentinel. +- No behaviour change for either backend. + +## Non-goals + +- Changes to token resolution semantics. This is a pure refactor. +- Adding support for any new credential type or provider. +- Consolidating any other backend differences beyond this one block. + +## Design + +### `AgentProvisionPlan.provisioned_env` + +Add `provisioned_env: dict[str, str]` (default empty) to +`AgentProvisionPlan`. This map holds host-side secrets that the +provisioning stage resolved and that egress needs injected into the +sidecar environ. + +When `forward_host_credentials=True` for Codex, `agent_provision_plan` +calls `codex_host_access_token(host_env)` and stores the result under +`CODEX_HOST_CREDENTIAL_TOKEN_REF`. This is already the prepare-time +stage where `write_codex_dummy_auth_file` runs, so the access-token +read is colocated with all other Codex-specific provisioning. + +### Backend call sites + +Both backends merge `provisioned_env` over `os.environ` before calling +`egress_resolve_token_values`: + +```python +effective_env = {**dict(os.environ), **plan.agent_provision.provisioned_env} +token_values = egress_resolve_token_values(plan.egress_plan.token_env_map, effective_env) +``` + +`provisioned_env` values take precedence over same-named host env vars, +so the Codex token slot resolves from the map written at prepare time +rather than from a raw `os.environ` lookup. + +### `egress_resolve_token_values` + +Remove the `CODEX_HOST_CREDENTIAL_TOKEN_REF` sentinel `continue` skip. +The function now resolves every slot in `token_env_map` from `host_env` +without special-casing any key name. The `CODEX_HOST_CREDENTIAL_TOKEN_REF` +key is present in `effective_env` (injected by `provisioned_env`) exactly +when it is needed. + +## Implementation + +- `bot_bottle/agent_provider.py`: add `provisioned_env` field; populate + it for Codex when `forward_host_credentials=True`. +- `bot_bottle/egress.py`: remove sentinel skip; remove + `egress_resolve_token_values_with_provider` (the intermediate function + that was introduced and then superseded by this design); drop the + `codex_auth` import. +- `bot_bottle/backend/docker/launch.py`: replace the provider-aware + resolution block with the `effective_env` merge. +- `bot_bottle/backend/smolmachines/launch.py`: same replacement in + `_resolve_token_env`. +- `tests/unit/test_agent_provider.py`: add tests verifying + `provisioned_env` is populated (or empty) for Codex with and without + `forward_host_credentials`. +- `tests/unit/test_egress.py`: remove `TestResolveTokenValuesWithProvider`; + replace the sentinel-skip test with a test verifying the Codex token + ref resolves normally when present in `host_env`. + +## References + +- Issue #118: Deduplicate egress token resolution across backends. +- Issue #117: Complexity hotspots in launch, egress, and auth paths + (source of the finding). +- PRD 0029: Provider auth credentials through egress (introduced the + duplicated block). 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 7d7107c..f1129cc 100644 --- a/tests/unit/test_egress.py +++ b/tests/unit/test_egress.py @@ -341,12 +341,12 @@ 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) + self.assertEqual({"EGRESS_TOKEN_0": "codex-access-token"}, out) if __name__ == "__main__":