diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f5a5859..3ee6c29 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -62,6 +62,10 @@ jobs: - name: Swarm agent launcher (plan/limit/key-resolve) (#16) run: python scripts/test-agent-launcher.py + - name: Model-key injection + revocation handshake (#16/#60) + env: { REDIS_FAKE: "1" } + run: python scripts/test-key-injection-contract.py + - name: Manager event-contract test env: { REDIS_FAKE: "1" } run: python scripts/test-contract-events.py diff --git a/CLIENT_GUIDE.md b/CLIENT_GUIDE.md index deb157f..f729041 100644 --- a/CLIENT_GUIDE.md +++ b/CLIENT_GUIDE.md @@ -172,7 +172,7 @@ python scripts/stub_agent.py ▼ [HM 查询面/回调摄入] 持久化事件 + 透传 sequence ▼ -[桌面客户端] ⑥ 经 HM 拉取/SSE：GET /{id}（状态）、/events?after=（事件流）、 +[桌面客户端] ⑥ 经 HM 拉取/SSE：GET /{id}（状态）、/events?after=（事件流；游标为 HM 不透明分页游标 next_after，非 per-swarm sequence——sequence 仅客户端去重/排序）、 /tasks（DAG）、/workflow、/metrics、/audit、**/result**（综述+deliverable+artifacts） 终态事件 swarm.completed 已直接带 summary+deliverable diff --git a/docs/integration/event-schema.md b/docs/integration/event-schema.md index 275f2d4..d4aa976 100644 --- a/docs/integration/event-schema.md +++ b/docs/integration/event-schema.md @@ -93,9 +93,12 @@ | `handoff.requested` | `task_id`, `from_role`, `to_role` | swarm_task_flow | ✅ | | `timeline.updated` | `title` | timeline | ✅（emit 时回填 `title`=`summary`） | | `budget.alert` | `threshold_pct` | budget | ✅（emit 时回填 `threshold_pct`） | +| `swarm.pool_terminated` | `user_id`, `secret_ref` | lifecycle | ✅（HM #60 模型 key 吊销握手；**非**客户端事件，见下） | | `phase.changed` / `agent.*` / `sk_tool.*` | — | ordinary_sub / sk | ❌（本运行时不产生） | > 说明：评审/重做循环复用 `task.retried` + `timeline.updated`（`Review cycle N`）表达，无单独 review 事件类型；蜂群内部遥测（提案/竞价/交叉评审/收敛/health）**不经回调外发**（见各协议文档「事件不进 Manager 流」），故不在本表。`swarm.completed/failed/stopped` 与 `approval.approved/rejected`、`handoff.created` 为本次冻结**新增**，HM 注册表需据此登记（agent_swarm#15.2）。 +> +> **`swarm.pool_terminated`（HM #60 模型 key 吊销握手）**：HM 控制面生命周期事件，**非**客户端 cockpit 事件、**不**入 `FROZEN_CLIENT_EVENT_TYPES`、不渲染到驾驶舱。当某用户**所有 swarm run 均被 stop**（completed/failed 可经 `…/input` 重开故仍保 key）时，运行时发恰好一次，载 `user_id` + `secret_ref`（`azkv://` 引用，非明文 key）。HM 需在 `subscribed_events` + `agent_callback` 登记，收到即吊销该用户 per-user `sk-` + 清 KV。口径见 [runtime-contract §3.3.1](./runtime-contract.md)。 ## 5. HM 响应 diff --git a/docs/integration/runtime-contract.md b/docs/integration/runtime-contract.md index e2bfd3e..3f65e32 100644 --- a/docs/integration/runtime-contract.md +++ b/docs/integration/runtime-contract.md @@ -86,9 +86,24 @@ Swarm 拉起每个专家 agent 时注入以下进程环境变量（`agent/main.p **约束**： - 模型 key 由 **Swarm 服务端从 `billing_context.secret_ref`（`azkv://`）解析**（`resolve_model_key`：override → azkv（部署 SecretResolver / dev `HEICODE_SECRET_`）→ 编排器 `OPENAI_API_KEY` 兜底）后注入被拉起 agent 的 env。key **不入** create 请求体 / 回调 / 日志 / argv（满足 §3.1）；解析不到则 agent keyless 启动并明确报错（不伪造）。`billing_context` 的 `newapi_user_ref`/`quota_ref` 仍仅为归因元数据。见 [usage-billing-schema.md §2](./usage-billing-schema.md)。 -- 同一用户的专家 agent 注入**同一把** `sk-` + `OPENAI_API_BASE=HM/v1`，计费归一到发起用户；`task_executor` 随模型请求带 `X-Agent-*` 归因头供 HM/NewAPI 关联。token 由 HM 按 agent 生命周期吊销（stop/delete）。 +- 同一用户的专家 agent 注入**同一把** `sk-` + `OPENAI_API_BASE=HM/v1`，计费归一到发起用户；`task_executor` 随模型请求带 `X-Agent-*` 归因头供 HM/NewAPI 关联。token 吊销走 §3.3.1 A.5 的 `swarm.pool_terminated` 握手（per-user，全部 run stop 后吊销）。 - **运行时事件前置**：只有 (i) HM 真把 create 派发到编排器（`SWARM_RUNTIME_ENABLED=true`，非 manager-local 适配器）且 (ii) Swarm 已按上表拉起 agent 并连入，编排器才回推 `task.*`/`swarm.*` 运行时事件（否则 events feed 仅有 HM 控制面 `deployment.status_changed`，见 #39）。 +#### 3.3.1 模型 key 注入对接参数（HM #60，Swarm 侧已定死） + +HM 实现 #60（计费 key 注入）前需 Swarm 定死的参数，逐条口径如下（实现：`agent_launcher.resolve_model_key` / `_resolve_secret_ref`，`swarm_runtime` 吊销握手；测试 `scripts/test-key-injection-contract.py`）： + +| # | 参数 | Swarm 口径（已定） | +|---|---|---| +| A.1 | `sk-` 粒度 + mint 时机 | **每用户一把**（首次蜂群开通 mint、跨该用户所有 run 复用）；同一用户的专家 agent 注入同一把 `sk-`、计费归一到发起用户。 | +| A.1 | KV secret 命名约定 | `swarm-model-key-`（Swarm 只取 `secret_ref` 末段做名，不强约束路径；命名供 HM 定位）。 | +| A.2 | KV secret **value 格式** | **JSON `{"openai_api_key":"sk-..."}`**（对齐 callback 签名密钥的 `{"callback_signing_secret":"..."}` 约定，可扩展）。Swarm 解析字段名 = `openai_api_key`；裸 `sk-` 字符串亦兼容（`_extract_model_key`）；解析不到/字段缺失 → 不伪造、agent keyless 明确报错。 | +| A.3 | Swarm 读 KV 身份 / RBAC | **⚠ 待定（联调阻塞前置）**：Swarm 运行时需对 `heicode-kv`（`https://heicode-kv.vault.azure.net`）有 `Key Vault Secrets User`（经 Workload Identity / MI）。该 RBAC 授权归属（运维 / Swarm 部署侧）**尚未敲定**，须先确认——不通则 `secret_ref` 解不出。生产 KV 适配器（`SECRET_RESOLVER`）在仓外经部署接线；dev/CI 用 `HEICODE_SECRET_` 环境映射。 | +| A.4 | `OPENAI_API_BASE` | **Swarm 部署常量**（`AGENT_OPENAI_API_BASE`/`OPENAI_API_BASE`，设为 HM 网关 `https://code.xinghanlab.com/v1`），**不**经 create 的 `billing_context` 下发。 | +| A.5 | 吊销信号 | **事件驱动（方案 A）**：`sk-` per-user 长存；`stop` 是唯一**终态**（`completed`/`failed` 可经 `POST …/input` 重开，故仍保留 key）。当某用户**所有 run 均被 stop**（retained 集清空）时，运行时发**恰好一次** `swarm.pool_terminated{user_id, secret_ref}`，HM 收到即吊销 `sk-` + 清 KV。单 run 的 `swarm.stopped` **不**触发吊销（key per-user 复用）；账户停用 / HM 主动 delete 走同一 stop 路径。 | + +> `swarm.pool_terminated` 是 **HM 控制面生命周期事件**（非客户端 cockpit 事件，**不**入 `FROZEN_CLIENT_EVENT_TYPES`）；HM 需在 `subscribed_events` + `agent_callback` 登记以收吊销信号。payload `secret_ref` 为 `azkv://` 引用（非明文 key）。 + ## 4. 状态机 部署状态：`waiting_approval` → `running` →（`blocked` ⇄ `running`）→ 终态 `completed` / `failed` / `stopped`。 diff --git a/docs/integration/security-boundary.md b/docs/integration/security-boundary.md index 429dd5a..2b302df 100644 --- a/docs/integration/security-boundary.md +++ b/docs/integration/security-boundary.md @@ -59,7 +59,9 @@ - **Swarm 模型**：Agent 主动出站连编排器 WebSocket（`/ws/{agent_id}`），**不**对公网暴露每 Agent 子域名。AM 单 Agent 模型里的「客户端↔agent 直连 + `AGENT_ACCESS_TOKEN` 本地校验」**不适用于** swarm（无直连回路）。 - 服务间鉴权：HM→Swarm 用 `AGENT_RUNTIME_SERVICE_TOKEN`（Bearer）；回调 HMAC 签名。 - **每用户并发 Agent 配额**：一个 `user_id` 同时连接的 Agent 数上限为 `MAX_AGENTS_PER_USER`（env，默认 10）。注册（WS `register` 消息携带 `user_id`）超额即被拒绝（回 `registration_rejected` 并关闭，code 1008），断开后释放名额。归因主轴仍为 `user.id`/`channelId`。未带 `user_id` 的 Agent 为 unbound，不计入该配额。实现：`ConnectionManager.can_bind_user/bind_user/unbind` + 注册处强制；测试 `scripts/test-max-agents-per-user.py`。 -- **Swarm 拉起 agent + 服务端解析 key（team 决议，runtime-contract §3.3）**：由 **Swarm 运行时**（`orchestrator/agent_launcher.py`）拉起专家 agent 池（拉起数 `min(池大小, MAX_AGENTS_PER_USER − 已连)`，与上面的注册兜底一致）。模型 key 由 **Swarm 从 `billing_context.secret_ref`（`azkv://`）服务端解析**后注入被拉起 agent 的 env——**不入** create 请求体 / 回调 / 日志 / argv（`command` 后端的密钥经进程 env 传入，不上命令行）。azkv 真实解析为部署侧 SecretResolver；dev/CI 用 `HEICODE_SECRET_`。解析不到即 keyless 启动并明确报错（不伪造）。 +- **Swarm 拉起 agent + 服务端解析 key（team 决议，runtime-contract §3.3）**：由 **Swarm 运行时**（`orchestrator/agent_launcher.py`）拉起专家 agent 池（拉起数 `min(池大小, MAX_AGENTS_PER_USER − 已连)`，与上面的注册兜底一致）。模型 key 由 **Swarm 从 `billing_context.secret_ref`（`azkv://`）服务端解析**后注入被拉起 agent 的 env——**不入** create 请求体 / 回调 / 日志 / argv（`command` 后端的密钥经进程 env 传入，不上命令行）。azkv 真实解析为部署侧 SecretResolver；dev/CI 用 `HEICODE_SECRET_`。解析不到即 keyless 启动并明确报错（不伪造）。KV secret value 为 JSON `{"openai_api_key":"sk-..."}`（HM #60；裸串兼容），解析字段 `openai_api_key`。 +- **per-user key 吊销握手（HM #60，runtime-contract §3.3.1 A.5）**：`sk-` per-user 长存；`stop` 为唯一终态（completed/failed 经 `…/input` 可重开故保 key）。某用户全部 run 被 stop 时运行时发一次 `swarm.pool_terminated{user_id, secret_ref}`（HM 控制面事件，非客户端、非冻结集），HM 据此吊销 `sk-` + 清 KV，杜绝永不吊销 token。 +- 🟡 **待定（联调阻塞前置）**：Swarm 运行时读 `heicode-kv` 的身份/RBAC（Workload Identity / MI + `Key Vault Secrets User`）**归属未敲定**（运维 / Swarm 部署侧），须先确认——不通则 `secret_ref` 解不出。 - 🟡 待接入：多租户运行时隔离（命名空间/网络/配额）由 Agent 平台（AKS Workload Identity）承载，非本仓编排器；归因主轴为 `user.id`/`channelId`（见 `usage-billing-schema.md`），不引入 tenant 概念。 ## 7. 外部 API 与传输 diff --git a/orchestrator/agent_launcher.py b/orchestrator/agent_launcher.py index e9521a1..a341a7f 100644 --- a/orchestrator/agent_launcher.py +++ b/orchestrator/agent_launcher.py @@ -25,6 +25,7 @@ it. Pure planning (`plan_launch_specs`) is separated from side-effecting launch from __future__ import annotations import asyncio +import json import logging import os import shlex @@ -87,15 +88,43 @@ def resolve_model_key(body: Dict[str, Any]) -> Optional[str]: return os.getenv("OPENAI_API_KEY") or None -def _resolve_secret_ref(secret_ref: str) -> Optional[str]: - """Resolve an azkv:// ref to its secret value. +#: Key Vault secret-value JSON field carrying the model key (HM #60 contract). HM writes the KV +#: secret as JSON ``{"openai_api_key": "sk-..."}`` — mirroring the callback-secret convention +#: (``{"callback_signing_secret": "..."}``) and leaving room for sibling fields later. A bare +#: ``sk-`` string value is still accepted (back-compat / hand-set dev secrets). +KV_MODEL_KEY_FIELD = "openai_api_key" - Production: an Azure Key Vault adapter (out of this repo; wire via deployment). Dev/CI: read - env `HEICODE_SECRET_` where is the trailing path segment of the ref. Returns None - when unavailable — we never fabricate a key. + +def _extract_model_key(secret_value: str) -> Optional[str]: + """Pull the model key out of a resolved KV secret value. + + Accepts the contract JSON ``{"openai_api_key": "sk-..."}`` (HM #60) and, as a fallback, a bare + ``sk-`` string. Never fabricates: unparseable/missing field → None. + """ + value = (secret_value or "").strip() + if not value: + return None + if value.startswith("{"): + try: + obj = json.loads(value) + except ValueError: + return None + key = obj.get(KV_MODEL_KEY_FIELD) if isinstance(obj, dict) else None + return str(key).strip() or None if key else None + return value # bare string secret value + + +def _resolve_secret_ref(secret_ref: str) -> Optional[str]: + """Resolve an azkv:// ref to the model key. + + Production: an Azure Key Vault adapter (out of this repo; wire via deployment) reads the secret + named by the ref's trailing segment and returns its value. Dev/CI: read env + `HEICODE_SECRET_`. Either way the value is the HM #60 JSON ``{"openai_api_key": "sk-..."}`` + (bare string also accepted). Returns None when unavailable — we never fabricate a key. """ name = secret_ref.rstrip("/").rsplit("/", 1)[-1] - return os.getenv(f"HEICODE_SECRET_{name}") + raw = os.getenv(f"HEICODE_SECRET_{name}") + return _extract_model_key(raw) if raw is not None else None def model_api_base() -> str: diff --git a/orchestrator/swarm_runtime.py b/orchestrator/swarm_runtime.py index f183e36..1fca836 100644 --- a/orchestrator/swarm_runtime.py +++ b/orchestrator/swarm_runtime.py @@ -116,6 +116,13 @@ class SwarmRuntime: IDEMPOTENCY_KEY_PREFIX = "swarm_idempotency:" EVENT_KEY_PREFIX = "swarm_events:" EVENT_SEQ_KEY_PREFIX = "swarm_event_seq:" + # Per-user "retained" run set (agent_swarm#16/#60 — model-key revocation handshake). Holds the + # swarm_ids whose per-user `sk-` must stay valid: a run is retained from create until it is + # *stopped*. completed/failed runs stay retained because they are reopenable via POST …/input + # (their relaunched agents reuse the same per-user key). Only when ALL of a user's runs are + # stopped does the runtime emit `swarm.pool_terminated` so HM can revoke the `sk-` + clear KV. + USER_RETAINED_KEY_PREFIX = "swarm_user_retained:" + POOL_TERMINATED_FLAG_PREFIX = "swarm_user_pool_terminated:" def __init__(self): self.callback_service_token = ( @@ -339,6 +346,10 @@ class SwarmRuntime: for approval in run.approvals.values(): await self.emit_event(run, "approval.requested", payload=approval) + # Mark the per-user `sk-` as in-use for the model-key revocation handshake (#16/#60). Held + # until every run of this user is stopped (completed/failed stay retained — reopenable). + await self.retain_run_for_user(run) + return run, True async def save_run(self, run: SwarmRun): @@ -426,6 +437,75 @@ class SwarmRuntime: await redis_client.set(f"{self.TASK_RUN_KEY_PREFIX}{task_id}", run.swarm_id) await self.save_run(run) + # ------------------------------------------------------------------------------------------ + # Per-user model-key revocation handshake (agent_swarm#16 / HM #60) + # ------------------------------------------------------------------------------------------ + @staticmethod + def _run_user_id(run: SwarmRun) -> Optional[str]: + """The initiating user for a run (from the create `x-user-id`, passed via runtime_headers).""" + headers = (run.metadata or {}).get("runtime_headers") or {} + uid = headers.get("x_user_id") or headers.get("x-user-id") + return str(uid) if uid else None + + @staticmethod + def _run_secret_ref(run: SwarmRun) -> Optional[str]: + """The azkv:// secret_ref the per-user `sk-` was resolved from (for HM to target the KV + secret on revocation). This is a reference, never a plaintext key (see _redact_sensitive).""" + ref = ((run.request_body or {}).get("billing_context") or {}).get("secret_ref") + return ref or None + + async def _read_retained(self, user_id: str) -> List[str]: + raw = await redis_client.get(f"{self.USER_RETAINED_KEY_PREFIX}{user_id}") + if not raw: + return [] + try: + return list(json.loads(raw)) + except (ValueError, TypeError): + return [] + + async def _write_retained(self, user_id: str, swarm_ids: List[str]) -> None: + key = f"{self.USER_RETAINED_KEY_PREFIX}{user_id}" + if swarm_ids: + await redis_client.set(key, json.dumps(sorted(set(swarm_ids)))) + else: + await redis_client.delete(key) + + async def retain_run_for_user(self, run: SwarmRun) -> None: + """Mark a run's per-user `sk-` as in-use. Idempotent. Clears any prior pool-terminated flag + so a future drain can signal again (the user has live work after a previous full teardown).""" + user_id = self._run_user_id(run) + if not user_id: + return + retained = await self._read_retained(user_id) + if run.swarm_id not in retained: + retained.append(run.swarm_id) + await self._write_retained(user_id, retained) + await redis_client.delete(f"{self.POOL_TERMINATED_FLAG_PREFIX}{user_id}") + + async def release_run_and_maybe_terminate_pool(self, run: SwarmRun, reason: str) -> None: + """Drop a *stopped* run from its user's retained set; if that user now has no retained runs + (all stopped — completed/failed stay retained, being reopenable), emit `swarm.pool_terminated` + once so HM revokes the per-user `sk-` and clears the KV secret. Best-effort; never raises.""" + user_id = self._run_user_id(run) + if not user_id: + return + try: + retained = [sid for sid in await self._read_retained(user_id) if sid != run.swarm_id] + await self._write_retained(user_id, retained) + if retained: + return # user still has live/reopenable runs — keep the key valid + flag_key = f"{self.POOL_TERMINATED_FLAG_PREFIX}{user_id}" + if await redis_client.get(flag_key): + return # already signalled since the last active run — don't double-revoke + await redis_client.set(flag_key, "1", ex=7 * 24 * 3600) + payload = {"user_id": user_id, "reason": reason or "user swarm pool torn down"} + secret_ref = self._run_secret_ref(run) + if secret_ref: + payload["secret_ref"] = secret_ref # azkv ref (not the key) — HM clears this KV secret + await self.emit_event(run, "swarm.pool_terminated", payload=payload) + except Exception as exc: # never break stop on the revocation signal + logger.warning("pool_terminated signal failed for user %s: %s", user_id, exc) + async def stop_run(self, deployment_id: str, reason: str = "") -> Optional[SwarmRun]: """Stop a run by runtime deployment id or swarm id.""" run = await self.find_run_by_deployment_id(deployment_id) @@ -449,6 +529,9 @@ class SwarmRuntime: "status": "stopped", "reason": reason or "Heicode Manager requested stop", }) + # Model-key revocation handshake (#16/#60): a stop is the only *final* state (completed/ + # failed are reopenable via …/input). If this drains the user's retained set, signal HM. + await self.release_run_and_maybe_terminate_pool(run, reason or "Heicode Manager requested stop") return run async def find_run_by_deployment_id(self, deployment_id: str) -> Optional[SwarmRun]: diff --git a/scripts/test-key-injection-contract.py b/scripts/test-key-injection-contract.py new file mode 100644 index 0000000..9e77e07 --- /dev/null +++ b/scripts/test-key-injection-contract.py @@ -0,0 +1,157 @@ +"""Model-key injection + revocation handshake contract tests (agent_swarm#16 / HM #60). + +Pins the two Swarm-side parameters HM locked for #60 (calling-key injection) before HM writes its +mint→KV→revoke side: + + A.2 — KV secret VALUE format: Swarm resolves the per-user `sk-` from the HM #60 JSON + ``{"openai_api_key": "sk-..."}`` (mirrors the callback-secret convention); a bare `sk-` + string is still accepted; malformed / missing field never fabricates a key. + A.5 — revocation signal: a `stop` is the only *final* state (completed/failed are reopenable via + POST …/input, so they keep the key). When a stop drains a user's retained-run set, the + runtime emits exactly one `swarm.pool_terminated{user_id, secret_ref}` so HM revokes the + `sk-` + clears the KV secret. No double-revoke; still-live runs keep the key. + +Hermetic: REDIS_FAKE, no model key, no real Manager callback (callback url empty). + +Run from agent_swarm_v6 (install deps first — needs fakeredis): + pip install -r orchestrator/requirements.txt + REDIS_FAKE=1 python scripts/test-key-injection-contract.py +""" +import asyncio +import json +import os +import sys +from pathlib import Path + +os.environ["REDIS_FAKE"] = "1" +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from orchestrator.redis_client import redis_client +from orchestrator import agent_launcher +from orchestrator.swarm_runtime import swarm_runtime, FROZEN_CLIENT_EVENT_TYPES + +failures = [] + + +def check(name, cond): + print(("PASS" if cond else "FAIL"), "-", name) + if not cond: + failures.append(name) + + +async def stored_events(swarm_id): + raw = await redis_client.lrange(f"{swarm_runtime.EVENT_KEY_PREFIX}{swarm_id}", 0, -1) + return [json.loads(r) for r in raw] + + +async def new_run(user_id, secret_ref="azkv://heicode-kv/secrets/swarm-model-key-u1"): + body = { + "mode": "swarm", + "orchestration_plan": {"objective": "key-contract test"}, + "billing_context": {"secret_ref": secret_ref}, + "callback": {"url": "", "subscribed_events": []}, + "metadata": {"manager_deployment_id": "m-key", "runtime_headers": {"x_user_id": user_id}}, + } + run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="k") + return run + + +# --------------------------------------------------------------------------------------------- +# A.2 — KV secret value format +# --------------------------------------------------------------------------------------------- +def test_kv_value_format(): + check("JSON {openai_api_key} → key extracted", + agent_launcher._extract_model_key('{"openai_api_key": "sk-CONTRACT"}') == "sk-CONTRACT") + check("JSON with sibling fields → still extracts the key", + agent_launcher._extract_model_key('{"openai_api_key":"sk-X","note":"per-user"}') == "sk-X") + check("bare sk- string value still accepted (back-compat)", + agent_launcher._extract_model_key("sk-BARE") == "sk-BARE") + check("malformed JSON → None (never fabricated)", + agent_launcher._extract_model_key("{not json") is None) + check("JSON missing the key field → None", + agent_launcher._extract_model_key('{"callback_signing_secret":"x"}') is None) + check("empty value → None", agent_launcher._extract_model_key("") is None) + + # End-to-end via resolve_model_key: azkv ref → dev env map → JSON value. + os.environ.pop("AGENT_LAUNCH_MODEL_KEY", None) + os.environ.pop("OPENAI_API_KEY", None) + os.environ["HEICODE_SECRET_swarm-model-key-u1"] = '{"openai_api_key": "sk-RESOLVED"}' + body = {"billing_context": {"secret_ref": "azkv://heicode-kv/secrets/swarm-model-key-u1"}} + check("resolve_model_key resolves the azkv JSON secret to the bare sk-", + agent_launcher.resolve_model_key(body) == "sk-RESOLVED") + os.environ.pop("HEICODE_SECRET_swarm-model-key-u1", None) + + +# --------------------------------------------------------------------------------------------- +# A.5 — revocation handshake (swarm.pool_terminated on retained-set drain) +# --------------------------------------------------------------------------------------------- +async def pool_terminated_events(swarm_id): + return [e for e in await stored_events(swarm_id) if e["event_type"] == "swarm.pool_terminated"] + + +async def test_single_run_stop_signals(): + run = await new_run("user-A") + # not signalled before stop + check("no pool_terminated before stop", not await pool_terminated_events(run.swarm_id)) + await swarm_runtime.stop_run(run.deployment_id, reason="manager stop") + evs = await pool_terminated_events(run.swarm_id) + check("pool_terminated emitted when user's only run is stopped", len(evs) == 1) + if evs: + p = evs[0]["payload"] + check("pool_terminated carries user_id", p.get("user_id") == "user-A") + check("pool_terminated carries azkv secret_ref (not a plaintext key)", + str(p.get("secret_ref", "")).startswith("azkv://")) + # It is an HM control-plane lifecycle signal, NOT a client task-cockpit event. + check("pool_terminated is NOT in the frozen client event set", + "swarm.pool_terminated" not in FROZEN_CLIENT_EVENT_TYPES) + + +async def test_multi_run_keeps_key_until_last_stop(): + a = await new_run("user-B") + b = await new_run("user-B") + await swarm_runtime.stop_run(a.deployment_id, reason="stop a") + # B still retained → no signal on either run's stream yet + sig_a = await pool_terminated_events(a.swarm_id) + sig_b = await pool_terminated_events(b.swarm_id) + check("stopping one of two runs does NOT revoke (other still live)", not sig_a and not sig_b) + await swarm_runtime.stop_run(b.deployment_id, reason="stop b") + check("pool_terminated fires only when the user's LAST run is stopped", + len(await pool_terminated_events(b.swarm_id)) == 1) + + +async def test_no_double_revoke(): + run = await new_run("user-C") + await swarm_runtime.stop_run(run.deployment_id, reason="stop 1") + # a redundant second stop of the same (already drained) user must not signal again + await swarm_runtime.stop_run(run.deployment_id, reason="stop 2") + check("no double pool_terminated on a repeated stop", + len(await pool_terminated_events(run.swarm_id)) == 1) + + +async def test_relife_after_new_run(): + # After a full teardown, a brand-new run re-arms the signal (clears the terminated flag). + run1 = await new_run("user-D") + await swarm_runtime.stop_run(run1.deployment_id, reason="stop d1") + check("first teardown signals", len(await pool_terminated_events(run1.swarm_id)) == 1) + run2 = await new_run("user-D") + await swarm_runtime.stop_run(run2.deployment_id, reason="stop d2") + check("a new run after teardown re-arms a second pool_terminated", + len(await pool_terminated_events(run2.swarm_id)) == 1) + + +async def main(): + await redis_client.connect() + test_kv_value_format() + await test_single_run_stop_signals() + await test_multi_run_keeps_key_until_last_stop() + await test_no_double_revoke() + await test_relife_after_new_run() + print() + if failures: + print(f"{len(failures)} key-injection-contract check(s) FAILED: {failures}") + sys.exit(1) + print("all key-injection-contract checks passed") + + +if __name__ == "__main__": + asyncio.run(main())