模型 key 注入对接:定死 HM #60 参数(KV value JSON + per-user 吊销握手)(Refs #16 #60)

回应 HM「#60 落地前需 Swarm 定死的参数清单」。Swarm 侧逐条定死并落实现:

- A.2 KV secret value 格式:JSON {"openai_api_key":"sk-..."}(对齐 callback 签名密钥
  约定),解析字段 openai_api_key;裸 sk- 串兼容;解析不到/字段缺失不伪造。实现
  agent_launcher._extract_model_key + _resolve_secret_ref。
- A.5 吊销信号(事件驱动):sk- per-user 长存;stop 为唯一终态(completed/failed 经
  …/input 可重开故保 key)。某用户全部 run 被 stop(retained 集清空)时,运行时发
  恰好一次 swarm.pool_terminated{user_id, secret_ref},HM 据此吊销 sk- + 清 KV。
  单 run swarm.stopped 不触发吊销。实现 swarm_runtime.retain_run_for_user /
  release_run_and_maybe_terminate_pool(per-user retained 集 + 一次性 flag)。
- A.1 粒度:每用户一把、跨 run 复用;KV 命名 swarm-model-key-<user_id>(文档)。
- A.4 OPENAI_API_BASE:Swarm 部署常量(已实现),不经 create 下发(文档确认)。
- A.3 KV 读 RBAC:⚠ 待定(联调阻塞前置)——如实标注归属未敲定,不伪造已就绪。
- B.1:CLIENT_GUIDE §9 /events 游标改 after=<next_after>(HM 不透明游标,非 sequence)。

swarm.pool_terminated 为 HM 控制面生命周期事件,不入 FROZEN_CLIENT_EVENT_TYPES、
不渲染驾驶舱;payload secret_ref 为 azkv:// 引用(非明文 key)。

文档:runtime-contract §3.3.1(参数表)、event-schema(注册 + 说明)、
security-boundary §6(吊销握手 + RBAC 待定)。测试 scripts/test-key-injection-contract.py
(KV 格式解析 + 吊销握手:单 run/多 run 保留/不重复吊销/teardown 后重新武装)+ CI 步。

不动 Manager 面接口、HMAC 回调、审批链、计费/审计字段语义。

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Songhaoz666
2026-06-11 17:57:15 +08:00
co-authored by Claude Opus 4.8
parent ad343e6404
commit 0cbab3f750
8 changed files with 302 additions and 9 deletions
+4
View File
@@ -62,6 +62,10 @@ jobs:
- name: Swarm agent launcher (plan/limit/key-resolve) (#16) - name: Swarm agent launcher (plan/limit/key-resolve) (#16)
run: python scripts/test-agent-launcher.py 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 - name: Manager event-contract test
env: { REDIS_FAKE: "1" } env: { REDIS_FAKE: "1" }
run: python scripts/test-contract-events.py run: python scripts/test-contract-events.py
+1 -1
View File
@@ -172,7 +172,7 @@ python scripts/stub_agent.py
▼ ▼
[HM 查询面/回调摄入] 持久化事件 + 透传 sequence [HM 查询面/回调摄入] 持久化事件 + 透传 sequence
▼ ▼
[桌面客户端] ⑥ 经 HM 拉取/SSE:GET /{id}(状态)、/events?after=<sequence>(事件流)、 [桌面客户端] ⑥ 经 HM 拉取/SSE:GET /{id}(状态)、/events?after=<next_after>(事件流;游标为 HM 不透明分页游标 next_after,非 per-swarm sequence——sequence 仅客户端去重/排序)、
/tasks(DAG)、/workflow、/metrics、/audit、**/result**(综述+deliverable+artifacts) /tasks(DAG)、/workflow、/metrics、/audit、**/result**(综述+deliverable+artifacts)
终态事件 swarm.completed 已直接带 summary+deliverable 终态事件 swarm.completed 已直接带 summary+deliverable
+3
View File
@@ -93,9 +93,12 @@
| `handoff.requested` | `task_id`, `from_role`, `to_role` | swarm_task_flow | ✅ | | `handoff.requested` | `task_id`, `from_role`, `to_role` | swarm_task_flow | ✅ |
| `timeline.updated` | `title` | timeline | ✅(emit 时回填 `title`=`summary`) | | `timeline.updated` | `title` | timeline | ✅(emit 时回填 `title`=`summary`) |
| `budget.alert` | `threshold_pct` | budget | ✅(emit 时回填 `threshold_pct`) | | `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 | ❌(本运行时不产生) | | `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)。 > 说明:评审/重做循环复用 `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 响应 ## 5. HM 响应
+16 -1
View File
@@ -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_<name>`)→ 编排器 `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)。 - 模型 key 由 **Swarm 服务端从 `billing_context.secret_ref`(`azkv://`)解析**(`resolve_model_key`:override → azkv(部署 SecretResolver / dev `HEICODE_SECRET_<name>`)→ 编排器 `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)。 - **运行时事件前置**:只有 (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-<user_id>`(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_<name>` 环境映射。 |
| 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. 状态机 ## 4. 状态机
部署状态:`waiting_approval` → `running` →(`blocked` ⇄ `running`)→ 终态 `completed` / `failed` / `stopped`。 部署状态:`waiting_approval` → `running` →(`blocked` ⇄ `running`)→ 终态 `completed` / `failed` / `stopped`。
+3 -1
View File
@@ -59,7 +59,9 @@
- **Swarm 模型**:Agent 主动出站连编排器 WebSocket(`/ws/{agent_id}`),**不**对公网暴露每 Agent 子域名。AM 单 Agent 模型里的「客户端↔agent 直连 + `AGENT_ACCESS_TOKEN` 本地校验」**不适用于** swarm(无直连回路)。 - **Swarm 模型**:Agent 主动出站连编排器 WebSocket(`/ws/{agent_id}`),**不**对公网暴露每 Agent 子域名。AM 单 Agent 模型里的「客户端↔agent 直连 + `AGENT_ACCESS_TOKEN` 本地校验」**不适用于** swarm(无直连回路)。
- 服务间鉴权:HM→Swarm 用 `AGENT_RUNTIME_SERVICE_TOKEN`(Bearer);回调 HMAC 签名。 - 服务间鉴权: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`。 - **每用户并发 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_<name>`。解析不到即 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_<name>`。解析不到即 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 概念。 - 🟡 待接入:多租户运行时隔离(命名空间/网络/配额)由 Agent 平台(AKS Workload Identity)承载,非本仓编排器;归因主轴为 `user.id`/`channelId`(见 `usage-billing-schema.md`),不引入 tenant 概念。
## 7. 外部 API 与传输 ## 7. 外部 API 与传输
+35 -6
View File
@@ -25,6 +25,7 @@ it. Pure planning (`plan_launch_specs`) is separated from side-effecting launch
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import json
import logging import logging
import os import os
import shlex import shlex
@@ -87,15 +88,43 @@ def resolve_model_key(body: Dict[str, Any]) -> Optional[str]:
return os.getenv("OPENAI_API_KEY") or None return os.getenv("OPENAI_API_KEY") or None
def _resolve_secret_ref(secret_ref: str) -> Optional[str]: #: Key Vault secret-value JSON field carrying the model key (HM #60 contract). HM writes the KV
"""Resolve an azkv:// ref to its secret value. #: 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_<name>` where <name> is the trailing path segment of the ref. Returns None def _extract_model_key(secret_value: str) -> Optional[str]:
when unavailable — we never fabricate a key. """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_<name>`. 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] 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: def model_api_base() -> str:
+83
View File
@@ -116,6 +116,13 @@ class SwarmRuntime:
IDEMPOTENCY_KEY_PREFIX = "swarm_idempotency:" IDEMPOTENCY_KEY_PREFIX = "swarm_idempotency:"
EVENT_KEY_PREFIX = "swarm_events:" EVENT_KEY_PREFIX = "swarm_events:"
EVENT_SEQ_KEY_PREFIX = "swarm_event_seq:" 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): def __init__(self):
self.callback_service_token = ( self.callback_service_token = (
@@ -339,6 +346,10 @@ class SwarmRuntime:
for approval in run.approvals.values(): for approval in run.approvals.values():
await self.emit_event(run, "approval.requested", payload=approval) 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 return run, True
async def save_run(self, run: SwarmRun): 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 redis_client.set(f"{self.TASK_RUN_KEY_PREFIX}{task_id}", run.swarm_id)
await self.save_run(run) 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]: async def stop_run(self, deployment_id: str, reason: str = "") -> Optional[SwarmRun]:
"""Stop a run by runtime deployment id or swarm id.""" """Stop a run by runtime deployment id or swarm id."""
run = await self.find_run_by_deployment_id(deployment_id) run = await self.find_run_by_deployment_id(deployment_id)
@@ -449,6 +529,9 @@ class SwarmRuntime:
"status": "stopped", "status": "stopped",
"reason": reason or "Heicode Manager requested stop", "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 return run
async def find_run_by_deployment_id(self, deployment_id: str) -> Optional[SwarmRun]: async def find_run_by_deployment_id(self, deployment_id: str) -> Optional[SwarmRun]:
+157
View File
@@ -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())