fix(secrets): recover encrypted tokens on all backends

bot_bottle/backend/consolidated_util.py

@@ -26,6 +26,7 @@ def provision_bottle(
     *,
     image_ref: str = "",
     tokens: dict[str, str] | None = None,
+    env_var_secret: str | None = None,
 ) -> RegisteredBottle:
     """Register the bottle and provision its git-gate state. Rolls back the
     registration if provisioning fails so no orphan is left.
@@ -35,7 +36,7 @@ def provision_bottle(
     ``RegisteredBottle`` so callers can inject it into the agent container's
     environment."""
     inputs = registration_inputs(egress_plan)
-    env_var_secret = new_env_var_secret()
+    env_var_secret = env_var_secret or new_env_var_secret()
     reg = client.register_bottle(
         source_ip, image_ref=image_ref, policy=inputs.policy,
         metadata=inputs.metadata, tokens=tokens, env_var_secret=env_var_secret,

bot_bottle/backend/docker/consolidated_launch.py

@@ -23,6 +23,7 @@ from ...orchestrator.client import OrchestratorClient
 from ...orchestrator.gateway import GATEWAY_NETWORK
 from ...orchestrator.lifecycle import INFRA_NAME, OrchestratorService
 from ...orchestrator.secret_store import ENV_VAR_SECRET_NAME
+from ...orchestrator.reprovision import reprovision_bottles
 from ..consolidated_util import provision_bottle
 from ..consolidated_util import teardown_consolidated as _teardown_util
 from .gateway_provision import DockerGatewayTransport
@@ -102,17 +103,17 @@ def _reprovision_running_bottles(
     when the orchestrator already has all tokens loaded. Best-effort: a single
     container exec failure never blocks a new bottle launch."""
     client = OrchestratorClient(orchestrator_url)
-    bottles = client.list_bottles()
-    if not bottles:
-        return
-
     # Build {source_ip: container_name} from live containers on the gateway
     # network, excluding the infra container itself.
-    proc = run_docker([
-        "docker", "network", "inspect",
-        "--format", "{{range .Containers}}{{.Name}} {{.IPv4Address}}\n{{end}}",
-        network,
-    ])
+    try:
+        proc = run_docker([
+            "docker", "network", "inspect",
+            "--format", "{{range .Containers}}{{.Name}} {{.IPv4Address}}\n{{end}}",
+            network,
+        ])
+    except OSError as exc:
+        log.info(f"egress token reprovision skipped: {exc}")
+        return
     ip_to_container: dict[str, str] = {}
     for line in proc.stdout.splitlines():
         parts = line.strip().split()
@@ -121,26 +122,15 @@ def _reprovision_running_bottles(
             if ip:
                 ip_to_container[ip] = parts[0]
 
-    reprovisioned = 0
-    for bottle in bottles:
-        bottle_id = bottle.get("bottle_id")
-        source_ip = bottle.get("source_ip")
-        if not isinstance(bottle_id, str) or not isinstance(source_ip, str):
-            continue
-        container_name = ip_to_container.get(source_ip)
-        if not container_name:
-            continue
+    secrets_by_ip: dict[str, str] = {}
+    for source_ip, container_name in ip_to_container.items():
         proc = run_docker(
             ["docker", "exec", container_name, "printenv", ENV_VAR_SECRET_NAME]
         )
-        if proc.returncode != 0 or not proc.stdout.strip():
-            continue
-        try:
-            if client.reprovision_gateway(bottle_id, proc.stdout.strip()):
-                reprovisioned += 1
-        except Exception:  # noqa: BLE001 — best-effort, never block a launch
-            pass
+        if proc.returncode == 0 and proc.stdout.strip():
+            secrets_by_ip[source_ip] = proc.stdout.strip()
 
+    reprovisioned = reprovision_bottles(client, secrets_by_ip)
     if reprovisioned:
         log.info(
             "reprovisioned egress tokens",

bot_bottle/backend/firecracker/bottle_plan.py

@@ -22,6 +22,9 @@ class FirecrackerBottlePlan(BottlePlan):
     # (egress proxy credentials, git-gate/supervise headers); set by launch
     # from the orchestrator registration. Empty pre-registration.
     identity_token: str = ""
+    # Applied to every agent SSH exec and mirrored into /run inside the VM so
+    # the host can recover it after the infra VM restarts.
+    env_var_secret: str = ""
 
     @property
     def container_name(self) -> str:

bot_bottle/backend/firecracker/cleanup.py

@@ -88,6 +88,12 @@ def _scan_processes(run_root: Path) -> tuple[set[str], list[int]]:
     return live, orphan_pids
 
 
+def live_run_dirs() -> tuple[Path, ...]:
+    """Run directories backed by currently running agent microVMs."""
+    live, _ = _scan_processes(_run_root())
+    return tuple(Path(path) for path in sorted(live))
+
+
 def _orphan_run_dirs(run_root: Path, live: set[str]) -> list[str]:
     """Run dirs with no live VM behind them — the leaked ones to remove."""
     if not run_root.is_dir():

bot_bottle/backend/firecracker/consolidated_launch.py

@@ -25,16 +25,27 @@ The TAP slot allocation, rootfs build, and VM boot are the caller's job.
 
 from __future__ import annotations
 
+import json
+import subprocess
 from dataclasses import dataclass
+from pathlib import Path
 
 from ...egress import EgressPlan
 from ...git_gate import GitGatePlan
-from ...orchestrator.client import OrchestratorClient
+from ...log import info
+from ...orchestrator.client import OrchestratorClient, OrchestratorClientError
 from ...orchestrator.lifecycle import (
     OrchestratorStartError,  # re-exported so callers can catch it
 )
-from ..consolidated_util import provision_bottle, teardown_consolidated as _teardown_util
-from . import infra_vm
+from ...orchestrator.reprovision import reprovision_bottles
+from ...orchestrator.secret_store import ENV_VAR_SECRET_NAME
+from ..consolidated_util import (
+    provision_bottle,
+    teardown_consolidated as _teardown_util,
+)
+from . import cleanup, infra_vm, util
+
+_ENV_VAR_SECRET_PATH = "/run/bot-bottle/env-var-secret"
 
 
 class ConsolidatedLaunchError(RuntimeError):
@@ -53,6 +64,54 @@ class LaunchContext:
     env_var_secret: str = ""  # encryption key injected into the agent's env
 
 
+def _guest_ip_from_config(config_path: Path) -> str:
+    """Read the kernel's configured guest IP from a Firecracker config."""
+    try:
+        config = json.loads(config_path.read_text())
+        args = config["boot-source"]["boot_args"]
+        ip_arg = next(part for part in args.split() if part.startswith("ip="))
+        return ip_arg.removeprefix("ip=").split(":", 1)[0]
+    except (OSError, ValueError, KeyError, TypeError, StopIteration):
+        return ""
+
+
+def persist_env_var_secret(private_key: Path, guest_ip: str, secret: str) -> None:
+    """Mirror the exec-time key into guest tmpfs for restart recovery."""
+    proc = subprocess.run(
+        util.ssh_base_argv(private_key, guest_ip)
+        + [f"umask 077; mkdir -p /run/bot-bottle; cat > {_ENV_VAR_SECRET_PATH}"],
+        input=secret, capture_output=True, text=True, check=False,
+    )
+    if proc.returncode != 0:
+        raise ConsolidatedLaunchError(
+            f"failed to persist {ENV_VAR_SECRET_NAME} in agent VM: "
+            f"{proc.stderr.strip() or '<no stderr>'}"
+        )
+
+
+def _reprovision_running_bottles(client: OrchestratorClient) -> None:
+    """Read keys from live agent VMs and restore the restarted gateway."""
+    try:
+        secrets_by_ip: dict[str, str] = {}
+        for run_dir in cleanup.live_run_dirs():
+            guest_ip = _guest_ip_from_config(run_dir / "config.json")
+            private_key = run_dir / "bottle_id_ed25519"
+            if not guest_ip or not private_key.is_file():
+                continue
+            proc = subprocess.run(
+                util.ssh_base_argv(private_key, guest_ip)
+                + [f"cat {_ENV_VAR_SECRET_PATH}"],
+                capture_output=True, text=True, check=False,
+            )
+            if proc.returncode == 0 and proc.stdout.strip():
+                secrets_by_ip[guest_ip] = proc.stdout.strip()
+        count = reprovision_bottles(client, secrets_by_ip)
+        if count:
+            info(f"reprovisioned egress tokens for {count} Firecracker bottle(s)")
+    except (OSError, OrchestratorClientError) as exc:
+        info(f"egress token reprovision skipped: {exc}")
+
+
 def launch_consolidated(
     egress_plan: EgressPlan,
     git_gate_plan: GitGatePlan,
@@ -67,6 +126,7 @@ def launch_consolidated(
     infra = infra_vm.ensure_running()
     url = infra.control_plane_url
     client = OrchestratorClient(url)
+    _reprovision_running_bottles(client)
 
     transport = infra_vm.gateway_transport()
     reg = provision_bottle(

bot_bottle/backend/firecracker/launch.py

@@ -55,8 +55,10 @@ from . import firecracker_vm, image_builder, isolation_probe, netpool, util
 from .bottle import FirecrackerBottle
 from .bottle_plan import FirecrackerBottlePlan
 from ...orchestrator.config_store import resolve_teardown_timeout
+from ...orchestrator.secret_store import ENV_VAR_SECRET_NAME
 from .consolidated_launch import (
     launch_consolidated,
+    persist_env_var_secret,
     teardown_consolidated,
 )
 
@@ -153,6 +155,7 @@ def launch(
             git_gate_plan=git_gate_plan,
             egress_plan=egress_plan,
             identity_token=ctx.identity_token,
+            env_var_secret=ctx.env_var_secret,
             # Deliver the identity token as egress proxy credentials — clients
             # honor `HTTPS_PROXY=http://id:token@gw` without app changes; the
             # gateway reads Proxy-Authorization, validates the (source_ip,
@@ -187,6 +190,7 @@ def launch(
         )
         stack.callback(vm.terminate)
         firecracker_vm.wait_for_ssh(vm, private_key)
+        persist_env_var_secret(private_key, slot.guest_ip, ctx.env_var_secret)
 
         # Authoritative fail-closed egress-boundary check, before the agent
         # runs: prove the VM cannot reach the host directly.
@@ -281,6 +285,8 @@ def _agent_guest_env(plan: FirecrackerBottlePlan, host_ip: str) -> dict[str, str
         env["GIT_GATE_URL"] = plan.agent_git_gate_url
     if plan.agent_supervise_url:
         env["MCP_SUPERVISE_URL"] = plan.agent_supervise_url
+    if plan.env_var_secret:
+        env[ENV_VAR_SECRET_NAME] = plan.env_var_secret
     for entry in egress_agent_env_entries(plan.egress_plan):
         key, _, value = entry.partition("=")
         env[key] = value

bot_bottle/backend/firecracker/util.py

@@ -399,6 +399,9 @@ fi
 chown -R 0:0 /root 2>/dev/null || true
 
 mkdir -p /etc/dropbear /run
+# Keep restart-recovery key material memory-backed, separate from both the
+# agent rootfs and the infra VM's persistent registry volume.
+mount -t tmpfs -o mode=0755 tmpfs /run 2>/dev/null || true
 # -R: generate host keys on demand. -E: log auth failures to stderr,
 # captured in the host-side console.log for debugging.
 /bb-dropbear -R -E -p 22 &

bot_bottle/backend/macos_container/bottle_plan.py

@@ -20,6 +20,9 @@ class MacosContainerBottlePlan(BottlePlan):
     # bottle is registered. See launch.py's stamp for why it lives here and not
     # only in the exec-time proxy env.
     identity_token: str = ""
+    # Generated before `container run` so it becomes part of the container's
+    # configured environment and can be read back after an infra restart.
+    env_var_secret: str = ""
 
     @property
     def container_name(self) -> str:

bot_bottle/backend/macos_container/consolidated_launch.py

@@ -38,7 +38,12 @@ from ...egress import EgressPlan
 from ...git_gate import GitGatePlan
 from ...log import info
 from ...orchestrator.client import OrchestratorClient, OrchestratorClientError
-from ..consolidated_util import provision_bottle, teardown_consolidated as _teardown_util
+from ...orchestrator.reprovision import reprovision_bottles
+from ...orchestrator.secret_store import ENV_VAR_SECRET_NAME
+from ..consolidated_util import (
+    provision_bottle,
+    teardown_consolidated as _teardown_util,
+)
 from . import util as container_mod
 from .enumerate import CONTAINER_NAME_PREFIX, EnumerationError, enumerate_active
 from .gateway import GATEWAY_NETWORK
@@ -84,12 +89,35 @@ def ensure_gateway(
     needs `gateway_ip` at run time."""
     service = service or MacosInfraService()
     infra = service.ensure_running()
-    return GatewayEndpoint(
+    endpoint = GatewayEndpoint(
         orchestrator_url=infra.control_plane_url,
         gateway_ip=infra.gateway_ip,
         gateway_ca_pem=service.ca_cert_pem(),
         network=service.network,
     )
+    _reprovision_running_bottles(endpoint)
+    return endpoint
+
+
+def _reprovision_running_bottles(endpoint: GatewayEndpoint) -> None:
+    """Recover keys from live Apple containers and restore gateway tokens."""
+    try:
+        secrets_by_ip: dict[str, str] = {}
+        for agent in enumerate_active():
+            name = f"{CONTAINER_NAME_PREFIX}{agent.slug}"
+            source_ip = container_mod.inspect_container_network_ip(name, endpoint.network)
+            if not source_ip:
+                continue
+            secret = container_mod.read_container_env(name, ENV_VAR_SECRET_NAME)
+            if secret:
+                secrets_by_ip[source_ip] = secret
+        count = reprovision_bottles(
+            OrchestratorClient(endpoint.orchestrator_url), secrets_by_ip,
+        )
+        if count:
+            info(f"reprovisioned egress tokens for {count} macOS bottle(s)")
+    except (OrchestratorClientError, EnumerationError, OSError) as exc:
+        info(f"egress token reprovision skipped: {exc}")
 
 
 def live_source_ips(network: str) -> list[str]:
@@ -126,6 +154,7 @@ def register_agent(
     endpoint: GatewayEndpoint,
     image_ref: str = "",
     tokens: dict[str, str] | None = None,
+    env_var_secret: str | None = None,
 ) -> LaunchContext:
     """Register the (already running) agent by its address and provision its
     git-gate state into the gateway. `source_ip` must be read from the live
@@ -145,6 +174,7 @@ def register_agent(
     reg = provision_bottle(
         client, source_ip, egress_plan, git_gate_plan, AppleGatewayTransport(),
         image_ref=image_ref, tokens=tokens,
+        env_var_secret=env_var_secret,
     )
     return LaunchContext(
         bottle_id=reg.bottle_id,

bot_bottle/backend/macos_container/launch.py

@@ -68,6 +68,7 @@ from .gateway_hosts import (
 )
 from .bottle_plan import MacosContainerBottlePlan
 from ...orchestrator.config_store import resolve_teardown_timeout
+from ...orchestrator.secret_store import ENV_VAR_SECRET_NAME, new_env_var_secret
 from .consolidated_launch import (
     GatewayEndpoint,
     ensure_gateway,
@@ -142,6 +143,7 @@ def launch(
         plan = _provision_git_gate_keys(plan)
         plan = _install_gateway_ca(plan, endpoint)
         plan = _stamp_agent_urls(plan, endpoint)
+        plan = dataclasses.replace(plan, env_var_secret=new_env_var_secret())
 
         # Step 3: run the agent. It has no identity token yet — registration
         # needs the address this run assigns.
@@ -176,6 +178,7 @@ def launch(
             endpoint=endpoint,
             image_ref=plan.image,
             tokens=token_values,
+            env_var_secret=plan.env_var_secret,
         )
         stack.callback(
             teardown_consolidated, ctx.bottle_id,
@@ -406,6 +409,8 @@ def _agent_env_entries(
         env.append(f"GIT_GATE_URL={plan.agent_git_gate_url}")
     if plan.agent_supervise_url:
         env.append(f"MCP_SUPERVISE_URL={plan.agent_supervise_url}")
+    if getattr(plan, "env_var_secret", ""):
+        env.append(f"{ENV_VAR_SECRET_NAME}={plan.env_var_secret}")
     for name, value in sorted(plan.agent_provision.guest_env.items()):
         env.append(f"{name}={value}")
     # Forwarded vars: bare name → inherits from the `container run` process env

bot_bottle/backend/macos_container/util.py

@@ -361,6 +361,12 @@ def exec_container(name: str, argv: list[str]) -> None:
         )
 
 
+def read_container_env(name: str, env_name: str) -> str:
+    """Read one configured env value from a running container, or ``""``."""
+    result = _run_container_op([_CONTAINER, "exec", name, "printenv", env_name])
+    return result.stdout.strip() if result.returncode == 0 else ""
+
+
 def exec_container_as_root(name: str, argv: list[str]) -> None:
     """`exec_container`, but as uid 0 inside the container.