回应 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:
co-authored by
Claude Opus 4.8
parent
ad343e6404
commit
0cbab3f750
@@ -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
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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 响应
|
||||||
|
|
||||||
|
|||||||
@@ -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`。
|
||||||
|
|||||||
@@ -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 与传输
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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]:
|
||||||
|
|||||||
@@ -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())
|
||||||
Reference in New Issue
Block a user