cloud: production control plane (API, auth/RBAC, queue, sandbox, memory/KB, fallback/cost, webhooks)
This commit is contained in:
1 parent
b97c34465f
commit
b74d99e959
17 files changed
+3420
No files matched your search
+191
@@ -0,0 +1,191 @@
|
||||
# 云端控制面(cloud/)
|
||||
|
||||
本目录是 host.py 本地编排之上的云端架构实施:控制面 server +
|
||||
已就绪的模块(认证、队列、沙箱、记忆、知识库、模型路由、成本、
|
||||
webhook)。零第三方依赖,只用 Python 标准库。
|
||||
|
||||
> 诚实边界:这是一个本地可跑的控制面演示,不是已建成的真实云服务。
|
||||
> 本 server 不跑模型回合、不连真实 GitHub/Slack/云服务。生产前仍缺
|
||||
> 的外部条件见文末“生产前仍缺什么”一节,不要绕过去。
|
||||
|
||||
## 架构:谁负责什么
|
||||
|
||||
```
|
||||
客户端 / GitHub / Slack
|
||||
│ HTTP + X-API-Key(webhook 靠签名,见下)
|
||||
▼
|
||||
cloud/server.py(控制面,ThreadingHTTPServer)
|
||||
├─ 项目 / 任务:建、查、取消、心跳、stale 回收
|
||||
├─ 认证 / 权限 / 审计:cloud/auth.py(AuthStore)
|
||||
└─ webhook 入站:cloud/webhooks.py(验签 + 解析,只解析不执行)
|
||||
│ (后续由 worker 接线)
|
||||
▼
|
||||
worker(尚未接成 HTTP 端点,本版需另行调用模块)
|
||||
├─ 派发顺序:cloud/queue.py(order_tasks / ready_tasks /
|
||||
│ budget_exceeded / TaskQueue;状态机与 host.py 对齐)
|
||||
├─ 执行:cloud/sandbox.py(LocalExecutor 调 host.py CLI 做只读
|
||||
│ 封装;DockerExecutor / CloudExecutor 仍是占位,未接通)
|
||||
├─ 记忆:cloud/memory.py(MemoryStore:任务摘要,供 prompt 上下文)
|
||||
├─ 知识库:cloud/kb.py(KBStore:按项目隔离的文档检索)
|
||||
├─ 模型路由:cloud/models.py(ModelRouter / record_fallback:
|
||||
│ 只决定 fallback,不跑模型回合)
|
||||
├─ 成本:cloud/cost.py(CostStore / fetch_account_usage;
|
||||
│ usage 来自调用方注入的 client,本模块不联网)
|
||||
└─ 通知:cloud/webhooks.py(NotificationService:webhook + Slack
|
||||
占位分发;Web Push / APNs / FCM 明确未接,需要证书才做)
|
||||
```
|
||||
|
||||
状态库:单一 SQLite 文件。AuthStore 用 `users` / `project_members` /
|
||||
`audit_log` 三张表;server 用 `projects` / `tasks` 两张表,表名不冲突。
|
||||
同一个 db_path 传给两边即可。不要再用共享 JSON 文件当多用户状态库。
|
||||
|
||||
并发:ThreadingHTTPServer 多线程处理请求,每个请求内新建 SQLite
|
||||
连接(server 表用新连接,AuthStore 每次请求新建一个实例),绝不在
|
||||
多线程 handler 间共用一个连接。注意:AuthStore 自身持有一个连接,
|
||||
server 因此不在 handler 间共享 AuthStore 实例。
|
||||
|
||||
任务状态机(与 host.py 对齐):
|
||||
|
||||
```
|
||||
queued → running → completed / failed / interrupted / merged / merge-failed
|
||||
```
|
||||
|
||||
本版 server 只流转 queued → running(心跳)→ interrupted(取消 /
|
||||
sweep 回收)。`completed` / `merged` 由 worker / 控制面后续调用方在
|
||||
跑完模型回合、合完分支后回写,本 server 不提供“跑模型”端点,也不
|
||||
要假装一次 POST 就能让任务变 completed。
|
||||
|
||||
## 启动方式
|
||||
|
||||
```bash
|
||||
cd ~/workspace/codex-app-server
|
||||
python3 -m cloud.server [--host 127.0.0.1] [--port 8765] [--db <路径>]
|
||||
# 默认:host 127.0.0.1,port 8765,db 为当前目录下的 cloud-control-plane.db
|
||||
```
|
||||
|
||||
测试用(不自动 serve):
|
||||
|
||||
```python
|
||||
from cloud.server import create_server
|
||||
server = create_server(db_path, host="127.0.0.1", port=8765)
|
||||
server.serve_forever()
|
||||
```
|
||||
|
||||
用户与 API key 由调用方先用 AuthStore 建(明文 key 只返回一次,
|
||||
库里只存 sha256 哈希,绝不把 key 写进文件或聊天记录):
|
||||
|
||||
```python
|
||||
from cloud.auth import AuthStore
|
||||
store = AuthStore(db_path)
|
||||
user_id, api_key_plain = store.create_user("alice")
|
||||
```
|
||||
|
||||
## API 列表
|
||||
|
||||
全部返回 JSON,`Content-Type: application/json`。除 `GET /health`
|
||||
与 `/webhooks/*` 外,所有端点要求请求头 `X-API-Key`,无 key 或 key
|
||||
错返回 `401 {"error": "unauthorized"}`。`/webhooks/*` 不走 X-API-Key,
|
||||
因为外部服务无法持本站 API key,改靠各自的签名认证(见下)。
|
||||
|
||||
| 方法与路径 | 权限 | 说明 |
|
||||
| --- | --- | --- |
|
||||
| `GET /health` | 无需认证 | `200 {"ok": true, "service": "cloud-control-plane"}` |
|
||||
| `POST /projects` | 认证即可 | body `{name, id?}`;无 id 则生成 uuid;调用者记为 owner(调 `AuthStore.create_project`,同时写 projects 表并记审计 `create_project`);返回 `201` + 项目 |
|
||||
| `GET /projects` | 认证即可 | 只返回调用者是成员的项目:`{"projects": [...]}` |
|
||||
| `POST /tasks` | `create_task` | body `{project_id, prompt, depends_on?, priority?, idempotency_key?}`;无权限 `403`;同 `project_id + idempotency_key` 已存在则返回已有任务(`200`,不建重复);否则建 `queued` 任务并记审计 `create_task`,返回 `201` |
|
||||
| `GET /tasks?project_id=...` | 该项目 `read` | 返回该项目任务列表:`{"tasks": [...]}`;无权限 `403` |
|
||||
| `GET /tasks/{id}` | 该任务所属项目 `read` | 无权限 `403`;不存在 `404` |
|
||||
| `POST /tasks/{id}/cancel` | 该项目 `create_task` | `queued` / `running` 改为 `interrupted` 并记审计 `cancel_task`;已是终态(completed / failed / interrupted / merged / merge-failed)返回 `400`;不存在 `404` |
|
||||
| `POST /tasks/{id}/heartbeat` | 该项目 `create_task` | 更新 `last_heartbeat = now`;状态为 `queued` 则改为 `running`;返回任务 |
|
||||
| `POST /admin/sweep` | 认证即可 | 调 `sweep_stale_tasks(db_path, threshold_seconds=300)`(body 可传 `threshold_seconds` 覆盖默认值);把 `running` 且 `last_heartbeat`(无心跳则用 `updated_at`)早于 `now - threshold` 的任务改为 `interrupted` 并记审计 `sweep_interrupted`;返回 `{"swept": [ids]}` |
|
||||
| `POST /webhooks/github` | GitHub 签名 | 见“webhook 端点” |
|
||||
| `POST /webhooks/slack` | Slack 签名 | 见“webhook 端点” |
|
||||
|
||||
单个项目 / 任务以裸对象返回;列表以 `{"projects": [...]}` /
|
||||
`{"tasks": [...]}` 包装。任务对象字段:`id`、`project_id`、`prompt`、
|
||||
`status`、`idempotency_key`、`depends_on`(list)、`priority`、
|
||||
`created_by`、`created_at`、`updated_at`、`last_heartbeat`。
|
||||
|
||||
函数 `sweep_stale_tasks(db_path, threshold_seconds=300) -> list` 独立
|
||||
可单测,不依赖 HTTP;回收动作的审计 `user_id` 记为 `"system"`。
|
||||
|
||||
### webhook 端点(加分项,已实现;只解析,不执行任务)
|
||||
|
||||
- `POST /webhooks/github`:用 `webhooks.verify_github_signature`
|
||||
验签(请求头 `X-Hub-Signature-256`),secret 从环境变量
|
||||
`GITHUB_WEBHOOK_SECRET` 读。未设 secret 返回 `503`;验签失败
|
||||
返回 `401`;成功返回 `parse_github_pr_comment` 的结果(PR 评论含
|
||||
`/fix-review` 时返回 fix_review 触发点,否则 `{"result": null}`)。
|
||||
真正执行 fix-review 由控制面后续调 host.py 完成,本端点不执行。
|
||||
- `POST /webhooks/slack`:用 `verify_slack_signature` 验签(请求头
|
||||
`X-Slack-Signature` / `X-Slack-Request-Timestamp`,含 300 秒防重放),
|
||||
secret 从环境变量 `SLACK_SIGNING_SECRET` 读。未设返回 `503`;
|
||||
验签失败返回 `401`;成功返回 `parse_slack_command` 的结果,不执行任务。
|
||||
|
||||
## 权限表
|
||||
|
||||
角色高低:viewer < contributor < maintainer < owner。未加入项目一律
|
||||
无权限。动作的最低角色(`AuthStore.check_permission` 的实际规则):
|
||||
|
||||
| 动作 | viewer | contributor | maintainer | owner |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `read`(读项目任务) | ✅ | ✅ | ✅ | ✅ |
|
||||
| `create_task`(建任务 / 取消 / 心跳) | ❌ | ✅ | ✅ | ✅ |
|
||||
| `approve_merge`(批合并,worker 侧用) | ❌ | ❌ | ✅ | ✅ |
|
||||
| `manage_project`(仅 owner,改成员等) | ❌ | ❌ | ❌ | ✅ |
|
||||
|
||||
说明:`manage_project` 的最低角色是 owner 本人,不是 maintainer;
|
||||
改成员(`AuthStore.set_member`)只有 owner 能做,否则抛
|
||||
`PermissionError`。审计(`audit` / `list_audit`)只追加、只读,本模块
|
||||
不提供删除或修改审计记录的方法。
|
||||
|
||||
注意一个实际签名细节:`AuthStore.create_project(project_id,
|
||||
owner_user_id)` 只写 `project_members` 的 owner 成员记录,不建项目
|
||||
记录本身;项目记录(name 等)由 server 写 `projects` 表。两者配合
|
||||
才算一个完整项目,不要只调一处。
|
||||
|
||||
## 测试
|
||||
|
||||
```bash
|
||||
cd ~/workspace/codex-app-server
|
||||
python3 -m unittest discover -s cloud/tests -v # 模块单测 + 集成单测
|
||||
python3 -m unittest discover -s tests -v # 原有 host.py 测试
|
||||
```
|
||||
|
||||
`cloud/tests/test_server.py` 用 http.client 对本地临时端口的 server
|
||||
做真实 HTTP 测试(临时 SQLite 文件与临时端口,跑完关闭 server),
|
||||
覆盖:无 key / 错误 key 401、非成员读项目 403、viewer 建任务 403、
|
||||
幂等键第二次返回同一任务且任务数不增加、cancel 把 running 任务变
|
||||
interrupted 且终态任务 cancel 返回 400、心跳后 sweep 不扫刚心跳的
|
||||
任务、sweep 扫掉心跳过期(手动改旧 last_heartbeat 或 threshold=0)
|
||||
的 running 任务、GET /health 不认证也通、GET /projects 只返回成员项目,
|
||||
以及两个 webhook 端点的 503 / 401 / 验签成功解析。
|
||||
|
||||
## 生产前仍缺什么
|
||||
|
||||
以下每一项都是外部条件,缺了就不能算生产可用,不要绕过去:
|
||||
|
||||
- **常驻机器 / 托管服务**:本 VM 会被平台定时回收,跑在这里只能算
|
||||
演示。生产需要一台常驻机器或托管服务(以及相应的进程守护与
|
||||
重启策略),否则控制面会随 VM 回收一起消失。
|
||||
- **GitHub 远端 + PAT**:codex-app-server 目前仍无 git 远端,`gh`
|
||||
也未登录;真实 PR 流程(host.py 的 `pr` / `pr-checks` / `pr-merge`)
|
||||
需要先建远端并用 Secure Vault 供给 PAT,PAT 绝不写文件。
|
||||
- **Slack App**:需要真实 Slack App(签名密钥、斜杠命令、事件订阅)
|
||||
才能接通 `SLACK_SIGNING_SECRET` 背后的入站;目前只有验签与解析,
|
||||
且 `NotificationService` 的 Slack 发送端仍是占位注入。
|
||||
- **云容器账号**:`CloudExecutor` 仍是占位(无账号即抛错),Docker
|
||||
执行也未接通容器挂载 / 网络隔离 / 资源限额;接通前不要宣称任务
|
||||
已在云端或容器里跑。
|
||||
- **模型 key 供给方式**:这是当前最实在的卡点。上次
|
||||
real-nightly-tests automation 就卡在这里:定时进程手里没有中转站
|
||||
key,key 从不存文件、只在进程环境里 transient 供给。控制面要跑
|
||||
真实模型回合,必须先由用户定下供给方式(例如常驻进程的环境、
|
||||
Secure Vault 注入、或其它经用户明确同意的渠道),而不是把 key
|
||||
写进配置文件、数据库或代码里凑合。定下之前,定时自动化与 worker
|
||||
的模型回合都只能停在“失败路径可用、模型跑不起来”的状态。
|
||||
- **其它待接线(非外部条件,但同样未完成)**:memory / KB / 成本 /
|
||||
模型路由尚未接成 HTTP 端点;OAuth(GitHub/Google)未接,首版只有
|
||||
API key + 哈希认证;Web Push / APNs / FCM 推送未接,需要证书才做;
|
||||
单机 SQLite 在高并发写下仍有 busy 锁上限,生产量上来后需评估
|
||||
换库或单写者队列。
|
||||
@@ -0,0 +1,5 @@
|
||||
"""云端控制面模块包。
|
||||
|
||||
只含 docstring,无副作用。子模块(auth、queue、sandbox、memory、kb、
|
||||
models、cost、webhooks、server)均零第三方依赖,只用 Python 标准库。
|
||||
"""
|
||||
+178
@@ -0,0 +1,178 @@
|
||||
"""认证与权限模块(云端架构,首版)。
|
||||
|
||||
认证方式:首版只用 API key + 哈希。
|
||||
|
||||
- 创建用户时用 ``secrets.token_urlsafe(32)`` 生成明文 key,仅返回给调用方一次;
|
||||
库里只存 ``hashlib.sha256(api_key_plain.encode()).hexdigest()``,绝不存明文。
|
||||
- 认证时对输入做同样的 sha256 再查库比对。
|
||||
|
||||
OAuth(GitHub/Google)是下一步接线:本模块不提供任何 OAuth 流程、
|
||||
不编造假 OAuth 端点或假 token 校验,接线时再由控制面按真实
|
||||
OAuth 协议接入,本模块的 API key 认证保持独立可用。
|
||||
|
||||
权限模型:项目内角色 viewer < contributor < maintainer < owner。
|
||||
动作最低角色:read -> viewer,create_task -> contributor,
|
||||
approve_merge -> maintainer,manage_project -> owner(仅 owner)。
|
||||
未加入项目一律无权限。
|
||||
|
||||
审计:audit() 只追加,list_audit() 只读;本模块不提供删除或修改
|
||||
审计记录的方法。
|
||||
|
||||
零第三方依赖,只用 Python 标准库。
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import secrets
|
||||
import sqlite3
|
||||
import time
|
||||
import uuid
|
||||
|
||||
ROLES = ["viewer", "contributor", "maintainer", "owner"]
|
||||
|
||||
_ROLE_LEVEL = {role: i for i, role in enumerate(ROLES)}
|
||||
|
||||
_ACTION_MIN_ROLE = {
|
||||
"read": "viewer",
|
||||
"create_task": "contributor",
|
||||
"approve_merge": "maintainer",
|
||||
"manage_project": "owner",
|
||||
}
|
||||
|
||||
|
||||
class AuthStore:
|
||||
"""基于 SQLite 的用户 / 项目成员 / 审计存储。"""
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = str(db_path)
|
||||
self.conn = sqlite3.connect(self.db_path)
|
||||
self.conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS users("
|
||||
"id TEXT PRIMARY KEY, "
|
||||
"username TEXT UNIQUE, "
|
||||
"api_key_hash TEXT, "
|
||||
"created_at REAL)"
|
||||
)
|
||||
self.conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS project_members("
|
||||
"project_id TEXT, "
|
||||
"user_id TEXT, "
|
||||
"role TEXT, "
|
||||
"PRIMARY KEY(project_id, user_id))"
|
||||
)
|
||||
self.conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS audit_log("
|
||||
"id INTEGER PRIMARY KEY AUTOINCREMENT, "
|
||||
"user_id TEXT, "
|
||||
"project_id TEXT, "
|
||||
"action TEXT, "
|
||||
"detail TEXT, "
|
||||
"created_at REAL)"
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def create_user(self, username):
|
||||
"""创建用户,返回 (user_id, api_key_plain),明文仅返回这一次。"""
|
||||
user_id = uuid.uuid4().hex
|
||||
api_key_plain = secrets.token_urlsafe(32)
|
||||
api_key_hash = hashlib.sha256(api_key_plain.encode()).hexdigest()
|
||||
self.conn.execute(
|
||||
"INSERT INTO users(id, username, api_key_hash, created_at)"
|
||||
" VALUES(?, ?, ?, ?)",
|
||||
(user_id, username, api_key_hash, time.time()),
|
||||
)
|
||||
self.conn.commit()
|
||||
return user_id, api_key_plain
|
||||
|
||||
def authenticate(self, api_key):
|
||||
"""用 API key 认证,成功返回 user_id,失败返回 None。"""
|
||||
if not api_key:
|
||||
return None
|
||||
api_key_hash = hashlib.sha256(api_key.encode()).hexdigest()
|
||||
row = self.conn.execute(
|
||||
"SELECT id FROM users WHERE api_key_hash = ?", (api_key_hash,)
|
||||
).fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
def create_project(self, project_id, owner_user_id):
|
||||
"""创建项目成员记录:指定用户为该项目的 owner。"""
|
||||
self.conn.execute(
|
||||
"INSERT OR REPLACE INTO project_members(project_id, user_id, role)"
|
||||
" VALUES(?, ?, 'owner')",
|
||||
(project_id, owner_user_id),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def _member_role(self, project_id, user_id):
|
||||
row = self.conn.execute(
|
||||
"SELECT role FROM project_members"
|
||||
" WHERE project_id = ? AND user_id = ?",
|
||||
(project_id, user_id),
|
||||
).fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
def set_member(self, project_id, actor_user_id, target_user_id, role):
|
||||
"""改成员角色:只有 owner 能改,否则 raise PermissionError;同时记审计。"""
|
||||
if role not in ROLES:
|
||||
raise ValueError("unknown role: %r" % (role,))
|
||||
if self._member_role(project_id, actor_user_id) != "owner":
|
||||
raise PermissionError(
|
||||
"only project owner can change members"
|
||||
)
|
||||
self.conn.execute(
|
||||
"INSERT OR REPLACE INTO project_members(project_id, user_id, role)"
|
||||
" VALUES(?, ?, ?)",
|
||||
(project_id, target_user_id, role),
|
||||
)
|
||||
self.conn.commit()
|
||||
self.audit(
|
||||
actor_user_id,
|
||||
project_id,
|
||||
"set_member",
|
||||
"%s:%s" % (target_user_id, role),
|
||||
)
|
||||
|
||||
def check_permission(self, user_id, project_id, action):
|
||||
"""检查权限:满足最低角色返回 True,否则 False(含未加入项目)。"""
|
||||
min_role = _ACTION_MIN_ROLE.get(action)
|
||||
if min_role is None:
|
||||
return False
|
||||
role = self._member_role(project_id, user_id)
|
||||
if role is None:
|
||||
return False
|
||||
return _ROLE_LEVEL[role] >= _ROLE_LEVEL[min_role]
|
||||
|
||||
def require_permission(self, user_id, project_id, action):
|
||||
"""不满足权限时 raise PermissionError。"""
|
||||
if not self.check_permission(user_id, project_id, action):
|
||||
raise PermissionError(
|
||||
"user %s lacks %s on project %s" % (user_id, action, project_id)
|
||||
)
|
||||
|
||||
def audit(self, user_id, project_id, action, detail=""):
|
||||
"""只追加一条审计记录。"""
|
||||
self.conn.execute(
|
||||
"INSERT INTO audit_log(user_id, project_id, action, detail, created_at)"
|
||||
" VALUES(?, ?, ?, ?, ?)",
|
||||
(user_id, project_id, action, detail, time.time()),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def list_audit(self, project_id, limit=100):
|
||||
"""按 created_at 倒序返回该项目的审计列表(dict)。"""
|
||||
rows = self.conn.execute(
|
||||
"SELECT id, user_id, project_id, action, detail, created_at"
|
||||
" FROM audit_log WHERE project_id = ?"
|
||||
" ORDER BY created_at DESC, id DESC LIMIT ?",
|
||||
(project_id, limit),
|
||||
).fetchall()
|
||||
return [
|
||||
{
|
||||
"id": row[0],
|
||||
"user_id": row[1],
|
||||
"project_id": row[2],
|
||||
"action": row[3],
|
||||
"detail": row[4],
|
||||
"created_at": row[5],
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
+158
@@ -0,0 +1,158 @@
|
||||
"""Cost estimation and cost recording.
|
||||
|
||||
Zero third-party dependencies; standard library only.
|
||||
|
||||
Actual token counts come from the model/provider response when
|
||||
available; when they are not, token counts are estimates (see
|
||||
``estimate_tokens``) and are stored separately as ``source="estimated"``
|
||||
so estimates are never mixed into actual totals.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
def estimate_tokens(prompt_text, output_text) -> int:
|
||||
"""Estimate total tokens as (chars of prompt + output) // 4."""
|
||||
prompt_text = prompt_text or ""
|
||||
output_text = output_text or ""
|
||||
return (len(prompt_text) + len(output_text)) // 4
|
||||
|
||||
|
||||
def fetch_account_usage(usage_client) -> dict:
|
||||
"""Fetch account usage through an injected ``usage_client`` callable.
|
||||
|
||||
This is a call wrapper around the app-server account/usage/read
|
||||
concept: the real ``usage_client`` is injected by the control plane
|
||||
once the billing route is confirmed available. This function itself
|
||||
never runs a model turn and never touches the network.
|
||||
|
||||
On success it returns ``{"available": True, "data": <raw data>}``.
|
||||
If ``usage_client`` raises, or returns None / an empty dict,
|
||||
it returns ``{"available": False, "data": None}`` and never raises
|
||||
to the caller. On a relay, usage being unreadable is a normal
|
||||
outcome (available=False), not a bug.
|
||||
"""
|
||||
try:
|
||||
data = usage_client()
|
||||
except Exception:
|
||||
return {"available": False, "data": None}
|
||||
if data is None:
|
||||
return {"available": False, "data": None}
|
||||
if isinstance(data, dict) and not data:
|
||||
return {"available": False, "data": None}
|
||||
return {"available": True, "data": data}
|
||||
|
||||
|
||||
class CostStore:
|
||||
"""SQLite store for per-task / per-project / per-day cost records."""
|
||||
|
||||
_SOURCES = ("actual", "estimated")
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = db_path
|
||||
if db_path != ":memory:":
|
||||
parent = os.path.dirname(os.path.abspath(db_path))
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
self._conn = sqlite3.connect(db_path)
|
||||
self._conn.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS cost_records (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
task_id TEXT,
|
||||
project_id TEXT,
|
||||
model TEXT,
|
||||
input_tokens INTEGER,
|
||||
output_tokens INTEGER,
|
||||
total_tokens INTEGER,
|
||||
usd_micros INTEGER,
|
||||
source TEXT,
|
||||
fallback_count INTEGER DEFAULT 0,
|
||||
day TEXT,
|
||||
created_at REAL
|
||||
)"""
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def record(self, task_id, project_id, model, input_tokens=0,
|
||||
output_tokens=0, usd_micros=None, source="actual",
|
||||
fallback_count=0):
|
||||
"""Insert one cost record; returns its row id.
|
||||
|
||||
``total_tokens`` is computed as input + output, and ``day`` is
|
||||
the UTC date (YYYY-MM-DD) of ``created_at``.
|
||||
"""
|
||||
if source not in self._SOURCES:
|
||||
raise ValueError(
|
||||
"source must be one of %r, got %r" % (self._SOURCES, source)
|
||||
)
|
||||
input_tokens = int(input_tokens or 0)
|
||||
output_tokens = int(output_tokens or 0)
|
||||
total_tokens = input_tokens + output_tokens
|
||||
created_at = time.time()
|
||||
day = datetime.fromtimestamp(created_at, tz=timezone.utc).strftime("%Y-%m-%d")
|
||||
cursor = self._conn.execute(
|
||||
"""INSERT INTO cost_records
|
||||
(task_id, project_id, model, input_tokens, output_tokens,
|
||||
total_tokens, usd_micros, source, fallback_count, day,
|
||||
created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
(task_id, project_id, model, input_tokens, output_tokens,
|
||||
total_tokens, usd_micros, source, int(fallback_count or 0),
|
||||
day, created_at),
|
||||
)
|
||||
self._conn.commit()
|
||||
return cursor.lastrowid
|
||||
|
||||
def record_estimated(self, task_id, project_id, model, prompt_text,
|
||||
output_text):
|
||||
"""Record an estimated cost row; returns its row id.
|
||||
|
||||
The token split cannot be recovered from a character estimate,
|
||||
so ``input_tokens`` holds the whole estimate and
|
||||
``output_tokens`` is 0; ``source`` is ``"estimated"`` and
|
||||
``usd_micros`` is None.
|
||||
"""
|
||||
total = estimate_tokens(prompt_text, output_text)
|
||||
return self.record(
|
||||
task_id, project_id, model,
|
||||
input_tokens=total, output_tokens=0,
|
||||
usd_micros=None, source="estimated", fallback_count=0,
|
||||
)
|
||||
|
||||
def _report(self, where, params):
|
||||
row = self._conn.execute(
|
||||
"""SELECT
|
||||
COALESCE(SUM(total_tokens), 0),
|
||||
COALESCE(SUM(usd_micros), 0),
|
||||
COUNT(*),
|
||||
COALESCE(SUM(fallback_count), 0),
|
||||
COALESCE(SUM(CASE WHEN source = 'actual'
|
||||
THEN total_tokens ELSE 0 END), 0),
|
||||
COALESCE(SUM(CASE WHEN source = 'estimated'
|
||||
THEN total_tokens ELSE 0 END), 0)
|
||||
FROM cost_records WHERE %s""" % where,
|
||||
params,
|
||||
).fetchone()
|
||||
total_tokens, usd_micros, turns, fallback_count, actual, estimated = row
|
||||
return {
|
||||
"total_tokens": int(total_tokens),
|
||||
"usd_micros": int(usd_micros),
|
||||
"turns": int(turns),
|
||||
"fallback_count": int(fallback_count),
|
||||
"actual_tokens": int(actual),
|
||||
"estimated_tokens": int(estimated),
|
||||
}
|
||||
|
||||
def report_by_task(self, task_id) -> dict:
|
||||
"""Aggregate cost records for one task."""
|
||||
return self._report("task_id = ?", (task_id,))
|
||||
|
||||
def report_by_project(self, project_id) -> dict:
|
||||
"""Aggregate cost records for one project."""
|
||||
return self._report("project_id = ?", (project_id,))
|
||||
|
||||
def report_by_day(self, day) -> dict:
|
||||
"""Aggregate cost records for one UTC day (YYYY-MM-DD)."""
|
||||
return self._report("day = ?", (day,))
|
||||
+169
@@ -0,0 +1,169 @@
|
||||
"""Project knowledge-base store: indexed docs searchable per project.
|
||||
|
||||
Zero third-party dependencies; Python standard library only.
|
||||
|
||||
Design notes:
|
||||
- Every document belongs to exactly one ``project_id``. Every search
|
||||
filters on ``project_id`` first: results from another project are
|
||||
never returned. This project isolation is a hard rule, not a default.
|
||||
- Full-text search uses an FTS5 virtual table ``kb_fts(content)`` when
|
||||
the local sqlite build supports it; otherwise search falls back to
|
||||
``LIKE``. The fallback is chosen once at construction time (try to
|
||||
create the FTS5 table, record whether it worked) and the search path
|
||||
follows that branch, so a build without FTS5 never crashes.
|
||||
- Re-indexing the same ``project_id`` + ``path`` replaces the old
|
||||
document (old row and its FTS row are deleted first), so repeated
|
||||
indexing never creates duplicates.
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
import time
|
||||
|
||||
_KINDS = ("file", "readme", "doc", "task", "review")
|
||||
|
||||
|
||||
class KBStore:
|
||||
"""SQLite-backed, per-project knowledge base with keyword search."""
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = db_path
|
||||
self.conn = sqlite3.connect(db_path)
|
||||
self.conn.row_factory = sqlite3.Row
|
||||
self.conn.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS kb_docs(
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
project_id TEXT,
|
||||
path TEXT,
|
||||
content TEXT,
|
||||
kind TEXT,
|
||||
created_at REAL
|
||||
)"""
|
||||
)
|
||||
# Try FTS5; fall back to LIKE if this sqlite build lacks it.
|
||||
self._fts = False
|
||||
try:
|
||||
self.conn.execute(
|
||||
"CREATE VIRTUAL TABLE IF NOT EXISTS kb_fts USING fts5(content)"
|
||||
)
|
||||
self._fts = True
|
||||
except sqlite3.OperationalError:
|
||||
self._fts = False
|
||||
self.conn.commit()
|
||||
|
||||
def index_file(self, project_id, path, content, kind="file"):
|
||||
"""Index one document; returns the new row id.
|
||||
|
||||
If the same ``project_id`` + ``path`` was already indexed, the
|
||||
old document (and its FTS row, when FTS5 is in use) is deleted
|
||||
first, so re-indexing replaces rather than duplicates.
|
||||
"""
|
||||
if kind not in _KINDS:
|
||||
raise ValueError("kind must be one of %s, got %r" % (_KINDS, kind))
|
||||
old_ids = [
|
||||
row["id"]
|
||||
for row in self.conn.execute(
|
||||
"SELECT id FROM kb_docs WHERE project_id = ? AND path = ?",
|
||||
(project_id, path),
|
||||
).fetchall()
|
||||
]
|
||||
for old_id in old_ids:
|
||||
if self._fts:
|
||||
self.conn.execute(
|
||||
"DELETE FROM kb_fts WHERE rowid = ?", (old_id,)
|
||||
)
|
||||
self.conn.execute("DELETE FROM kb_docs WHERE id = ?", (old_id,))
|
||||
cur = self.conn.execute(
|
||||
"""INSERT INTO kb_docs
|
||||
(project_id, path, content, kind, created_at)
|
||||
VALUES (?, ?, ?, ?, ?)""",
|
||||
(project_id, path, content, kind, time.time()),
|
||||
)
|
||||
doc_id = cur.lastrowid
|
||||
if self._fts:
|
||||
self.conn.execute(
|
||||
"INSERT INTO kb_fts(rowid, content) VALUES (?, ?)",
|
||||
(doc_id, content),
|
||||
)
|
||||
self.conn.commit()
|
||||
return doc_id
|
||||
|
||||
def index_task(self, project_id, task_id, prompt, conclusion, kind="task"):
|
||||
"""Index a task's prompt + conclusion as one document.
|
||||
|
||||
The document's path is ``"task:<task_id>"`` and its content is
|
||||
the prompt and conclusion joined together.
|
||||
"""
|
||||
content = "%s\n%s" % (prompt, conclusion)
|
||||
return self.index_file(project_id, "task:%s" % task_id, content, kind=kind)
|
||||
|
||||
def search(self, project_id, query, top_k=5):
|
||||
"""Search one project's docs by keyword; returns list[dict].
|
||||
|
||||
Each dict has ``path``, ``kind`` and ``snippet`` (first 200
|
||||
characters of content), plus ``id`` and ``project_id``.
|
||||
The ``project_id`` filter is always applied: docs from other
|
||||
projects are never returned.
|
||||
"""
|
||||
if not query:
|
||||
return []
|
||||
if self._fts:
|
||||
try:
|
||||
rows = self.conn.execute(
|
||||
"""SELECT d.id AS id, d.project_id AS project_id,
|
||||
d.path AS path, d.kind AS kind,
|
||||
d.content AS content
|
||||
FROM kb_docs d JOIN kb_fts f ON f.rowid = d.id
|
||||
WHERE d.project_id = ? AND kb_fts MATCH ?
|
||||
ORDER BY d.created_at DESC, d.id DESC
|
||||
LIMIT ?""",
|
||||
(project_id, self._fts_query(query), top_k),
|
||||
).fetchall()
|
||||
if rows:
|
||||
return [self._row_to_result(row) for row in rows]
|
||||
# Empty FTS result falls through to LIKE below. This is
|
||||
# not just a safety net: FTS5's default unicode61
|
||||
# tokenizer does not split CJK text into searchable
|
||||
# words, so a Chinese keyword that is present verbatim
|
||||
# in a document (verified: "登录" in a doc containing
|
||||
# "用户登录流程") returns zero FTS hits. LIKE substring
|
||||
# search recovers those cases; the project_id filter
|
||||
# still applies, so isolation is unchanged.
|
||||
except sqlite3.OperationalError:
|
||||
# Query syntax FTS5 cannot parse, or a broken FTS table:
|
||||
# fall through to the LIKE path rather than crashing.
|
||||
pass
|
||||
# LIKE fallback (also the path when FTS5 is unavailable).
|
||||
pattern = "%" + query.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + "%"
|
||||
rows = self.conn.execute(
|
||||
"""SELECT id, project_id, path, kind, content FROM kb_docs
|
||||
WHERE project_id = ? AND content LIKE ? ESCAPE '\\'
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT ?""",
|
||||
(project_id, pattern, top_k),
|
||||
).fetchall()
|
||||
return [self._row_to_result(row) for row in rows]
|
||||
|
||||
def list_docs(self, project_id):
|
||||
"""List all documents of one project, newest first."""
|
||||
rows = self.conn.execute(
|
||||
"""SELECT id, project_id, path, kind, content, created_at
|
||||
FROM kb_docs WHERE project_id = ?
|
||||
ORDER BY created_at DESC, id DESC""",
|
||||
(project_id,),
|
||||
).fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
@staticmethod
|
||||
def _fts_query(query):
|
||||
"""Turn a raw keyword query into a safe FTS5 phrase query."""
|
||||
# Phrase-quoting escapes special FTS5 syntax in the user query;
|
||||
# a multi-word query is searched as one phrase, which is the
|
||||
# conservative reading for keyword lookup.
|
||||
return '"%s"' % query.replace('"', '""')
|
||||
|
||||
@staticmethod
|
||||
def _row_to_result(row):
|
||||
d = dict(row)
|
||||
content = d.pop("content", "") or ""
|
||||
d["snippet"] = content[:200]
|
||||
return d
|
||||
+144
@@ -0,0 +1,144 @@
|
||||
"""Project memory store: short task summaries for prompt context.
|
||||
|
||||
Zero third-party dependencies; Python standard library only.
|
||||
|
||||
Design notes:
|
||||
- Summaries are the durable, compacted memory of a project. They are
|
||||
deliberately short (goal / conclusion / files changed / open issues)
|
||||
so that a prompt can carry recent context without replaying history.
|
||||
- ``build_prompt_context`` never returns the whole history: it returns
|
||||
at most ``top_k`` recent active summaries, formatted as a prompt prefix.
|
||||
- ``deprecate`` is a soft mark (status change, not deletion); ``delete``
|
||||
is a physical removal reserved for user-initiated correction/forgetting.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import time
|
||||
|
||||
|
||||
class MemoryStore:
|
||||
"""SQLite-backed store for per-project task summaries."""
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = db_path
|
||||
self.conn = sqlite3.connect(db_path)
|
||||
self.conn.row_factory = sqlite3.Row
|
||||
self.conn.execute(
|
||||
"""CREATE TABLE IF NOT EXISTS summaries(
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
project_id TEXT,
|
||||
thread_id TEXT,
|
||||
goal TEXT,
|
||||
conclusion TEXT,
|
||||
files_changed TEXT,
|
||||
open_issues TEXT,
|
||||
status TEXT DEFAULT 'active',
|
||||
created_at REAL
|
||||
)"""
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def add_summary(self, project_id, thread_id, goal, conclusion,
|
||||
files_changed=None, open_issues=None):
|
||||
"""Insert a summary; returns the new row id."""
|
||||
cur = self.conn.execute(
|
||||
"""INSERT INTO summaries
|
||||
(project_id, thread_id, goal, conclusion,
|
||||
files_changed, open_issues, status, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, 'active', ?)""",
|
||||
(
|
||||
project_id,
|
||||
thread_id,
|
||||
goal,
|
||||
conclusion,
|
||||
json.dumps(list(files_changed) if files_changed is not None else []),
|
||||
json.dumps(list(open_issues) if open_issues is not None else []),
|
||||
time.time(),
|
||||
),
|
||||
)
|
||||
self.conn.commit()
|
||||
return cur.lastrowid
|
||||
|
||||
def list_summaries(self, project_id, include_deprecated=False):
|
||||
"""List summaries for one project, newest first.
|
||||
|
||||
By default only ``status='active'`` rows are returned; pass
|
||||
``include_deprecated=True`` to also see deprecated ones.
|
||||
"""
|
||||
if include_deprecated:
|
||||
rows = self.conn.execute(
|
||||
"""SELECT * FROM summaries WHERE project_id = ?
|
||||
ORDER BY created_at DESC, id DESC""",
|
||||
(project_id,),
|
||||
).fetchall()
|
||||
else:
|
||||
rows = self.conn.execute(
|
||||
"""SELECT * FROM summaries WHERE project_id = ?
|
||||
AND status = 'active'
|
||||
ORDER BY created_at DESC, id DESC""",
|
||||
(project_id,),
|
||||
).fetchall()
|
||||
return [self._row_to_dict(row) for row in rows]
|
||||
|
||||
def deprecate(self, summary_id):
|
||||
"""Mark a summary deprecated (soft mark, not a physical delete)."""
|
||||
self.conn.execute(
|
||||
"UPDATE summaries SET status = 'deprecated' WHERE id = ?",
|
||||
(summary_id,),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def delete(self, summary_id):
|
||||
"""Physically delete a summary.
|
||||
|
||||
Deletion is a user-initiated action (correction / forgetting):
|
||||
it is the only way a summary is removed from storage, as opposed
|
||||
to :meth:`deprecate`, which merely hides it from default listings
|
||||
and prompt context while keeping the record.
|
||||
"""
|
||||
self.conn.execute(
|
||||
"DELETE FROM summaries WHERE id = ?",
|
||||
(summary_id,),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def build_prompt_context(self, project_id, top_k=5):
|
||||
"""Build a prompt prefix from the newest ``top_k`` active summaries.
|
||||
|
||||
Each summary becomes one line::
|
||||
|
||||
- [thread_id] 目标:...;结论:...;改动文件:...;遗留:...
|
||||
|
||||
Returns an empty string when the project has no active summaries.
|
||||
Never returns the full history: only the newest ``top_k`` rows
|
||||
(see :meth:`list_summaries`) are included.
|
||||
"""
|
||||
summaries = self.list_summaries(project_id)[:top_k]
|
||||
if not summaries:
|
||||
return ""
|
||||
lines = []
|
||||
for s in summaries:
|
||||
files = "、".join(s["files_changed"]) if s["files_changed"] else "无"
|
||||
issues = "、".join(s["open_issues"]) if s["open_issues"] else "无"
|
||||
lines.append(
|
||||
"- [{thread_id}] 目标:{goal};结论:{conclusion};"
|
||||
"改动文件:{files};遗留:{issues}".format(
|
||||
thread_id=s["thread_id"],
|
||||
goal=s["goal"],
|
||||
conclusion=s["conclusion"],
|
||||
files=files,
|
||||
issues=issues,
|
||||
)
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def _row_to_dict(row):
|
||||
d = dict(row)
|
||||
for key in ("files_changed", "open_issues"):
|
||||
try:
|
||||
d[key] = json.loads(d[key]) if d[key] else []
|
||||
except (TypeError, ValueError):
|
||||
d[key] = []
|
||||
return d
|
||||
+125
@@ -0,0 +1,125 @@
|
||||
"""Model fallback routing.
|
||||
|
||||
Zero third-party dependencies; standard library only.
|
||||
|
||||
This module decides *whether* a failed model attempt may fall back to
|
||||
the next model in a priority list, and *which* model to try next. It
|
||||
never runs a model turn itself and never talks to any provider.
|
||||
"""
|
||||
|
||||
FALLBACKABLE = {404, 429, 500, 502, 503, 504, "timeout"}
|
||||
|
||||
|
||||
def _is_fallbackable_code(value):
|
||||
"""Return True for a single status code / error token."""
|
||||
if isinstance(value, bool):
|
||||
return False
|
||||
if isinstance(value, int):
|
||||
if value in FALLBACKABLE:
|
||||
return True
|
||||
# Any other 5xx is a server-side failure of the same kind.
|
||||
return 500 <= value <= 599
|
||||
if isinstance(value, str):
|
||||
text = value.strip().lower()
|
||||
if text in FALLBACKABLE:
|
||||
return True
|
||||
if text.isdigit():
|
||||
return _is_fallbackable_code(int(text))
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def is_fallbackable(error) -> bool:
|
||||
"""Return True if ``error`` may trigger a model fallback.
|
||||
|
||||
``error`` may be:
|
||||
|
||||
* an int HTTP status code (404, 429, 5xx -> True; 401, 400 -> False);
|
||||
* the string ``"timeout"``;
|
||||
* an object / exception carrying a ``status`` or ``code`` attribute
|
||||
holding one of the above.
|
||||
|
||||
Authentication failures (401) and bad-request errors (400) are
|
||||
deliberately *not* fallbackable: retrying them on another model
|
||||
cannot fix bad credentials or a bad request.
|
||||
"""
|
||||
if error is None:
|
||||
return False
|
||||
if isinstance(error, (int, str)) and not isinstance(error, bool):
|
||||
return _is_fallbackable_code(error)
|
||||
for attr in ("status", "code"):
|
||||
value = getattr(error, attr, None)
|
||||
if value is not None:
|
||||
return _is_fallbackable_code(value)
|
||||
return False
|
||||
|
||||
|
||||
class ModelRouter:
|
||||
"""Walk a priority list of models, one fallback at a time."""
|
||||
|
||||
def __init__(self, priority_list):
|
||||
if not priority_list:
|
||||
raise ValueError("priority_list must not be empty")
|
||||
self.priority_list = list(priority_list)
|
||||
|
||||
def first(self):
|
||||
"""Return the first (highest-priority) model."""
|
||||
return self.priority_list[0]
|
||||
|
||||
def next_after(self, current_model, error):
|
||||
"""Return the model after ``current_model``, or None.
|
||||
|
||||
Returns None when ``error`` is not fallbackable, when
|
||||
``current_model`` is not in the priority list, or when there
|
||||
is no model left after ``current_model``.
|
||||
"""
|
||||
if not is_fallbackable(error):
|
||||
return None
|
||||
try:
|
||||
index = self.priority_list.index(current_model)
|
||||
except ValueError:
|
||||
return None
|
||||
if index + 1 < len(self.priority_list):
|
||||
return self.priority_list[index + 1]
|
||||
return None
|
||||
|
||||
def route(self, attempts):
|
||||
"""Return the next model to try given ``attempts``, or None.
|
||||
|
||||
``attempts`` is a list of ``(model, error)`` pairs in the order
|
||||
they were tried. With no attempts yet, the next model to try
|
||||
is the first in the priority list. A non-fallbackable error
|
||||
anywhere in the chain stops routing (returns None), as does
|
||||
exhausting the priority list.
|
||||
"""
|
||||
if not attempts:
|
||||
return self.first()
|
||||
last_model = None
|
||||
for model, error in attempts:
|
||||
if not is_fallbackable(error):
|
||||
return None
|
||||
last_model = model
|
||||
return self.next_after(last_model, attempts[-1][1])
|
||||
|
||||
|
||||
def record_fallback(task_record, from_model, to_model, reason) -> dict:
|
||||
"""Record a model fallback on ``task_record`` and return the event.
|
||||
|
||||
Fallback must be recorded as an event, never silently downgraded:
|
||||
the returned event is appended to ``task_record["events"]``
|
||||
(created if absent) so the downgrade is always visible in the
|
||||
task's history.
|
||||
"""
|
||||
event = {
|
||||
"type": "model_fallback",
|
||||
"task_id": task_record.get("id"),
|
||||
"from_model": from_model,
|
||||
"to_model": to_model,
|
||||
"reason": reason,
|
||||
}
|
||||
events = task_record.get("events")
|
||||
if not isinstance(events, list):
|
||||
events = []
|
||||
task_record["events"] = events
|
||||
events.append(event)
|
||||
return event
|
||||
+163
@@ -0,0 +1,163 @@
|
||||
"""派发队列:依赖拓扑排序、就绪筛选、预算判断与内存任务队列。
|
||||
|
||||
零第三方依赖,只用 Python 标准库;不依赖 host.py,纯函数 + 内存
|
||||
状态实现。状态机与 host.py 对齐:
|
||||
|
||||
queued / running / completed / failed / interrupted / merged
|
||||
/ merge-failed
|
||||
|
||||
注意:本模块只决定“按什么顺序派发”,不跑模型回合、不调 agent。
|
||||
"""
|
||||
|
||||
VALID_STATUSES = frozenset({
|
||||
"queued",
|
||||
"running",
|
||||
"completed",
|
||||
"failed",
|
||||
"interrupted",
|
||||
"merged",
|
||||
"merge-failed",
|
||||
})
|
||||
|
||||
|
||||
def _task_id(task):
|
||||
task_id = task.get("id")
|
||||
if task_id is None:
|
||||
raise ValueError("任务缺少 id")
|
||||
return task_id
|
||||
|
||||
|
||||
def _depends_on(task):
|
||||
deps = task.get("depends_on") or []
|
||||
if isinstance(deps, str):
|
||||
deps = [deps]
|
||||
return list(deps)
|
||||
|
||||
|
||||
def _priority(task):
|
||||
return task.get("priority") or 0
|
||||
|
||||
|
||||
def _topo_sort(tasks, initial_done=frozenset(), strict_missing=True):
|
||||
"""按波次拓扑排序:依赖先排,同一波内 priority 大的先,
|
||||
priority 相同时按 id 升序(保证确定性)。
|
||||
|
||||
initial_done:已视为完成的 id 集合(ready_tasks 用 completed_ids
|
||||
填充它,这样就绪子集里的依赖不再被当成缺失依赖)。
|
||||
strict_missing=True 时,依赖不在任务集且不在 initial_done 里
|
||||
即 raise ValueError;成环同样 raise ValueError。
|
||||
"""
|
||||
by_id = {}
|
||||
for task in tasks:
|
||||
task_id = _task_id(task)
|
||||
if task_id in by_id:
|
||||
raise ValueError(f"任务 id 重复:{task_id}")
|
||||
by_id[task_id] = task
|
||||
for task in tasks:
|
||||
for dep in _depends_on(task):
|
||||
if dep == _task_id(task):
|
||||
raise ValueError(f"任务 {_task_id(task)} 不能依赖自己")
|
||||
if strict_missing and dep not in by_id and dep not in initial_done:
|
||||
raise ValueError(
|
||||
f"任务 {_task_id(task)} 依赖不存在的任务 {dep}")
|
||||
|
||||
ordered = []
|
||||
done = set(initial_done)
|
||||
remaining = list(tasks)
|
||||
while remaining:
|
||||
ready = [t for t in remaining
|
||||
if all(d in done for d in _depends_on(t))]
|
||||
if not ready:
|
||||
names = ", ".join(str(_task_id(t)) for t in remaining)
|
||||
raise ValueError(f"依赖成环,无法排序:{names}")
|
||||
ready.sort(key=lambda t: (-_priority(t), str(_task_id(t))))
|
||||
for task in ready:
|
||||
ordered.append(task)
|
||||
done.add(_task_id(task))
|
||||
remaining = [t for t in remaining if _task_id(t) not in done]
|
||||
return ordered
|
||||
|
||||
|
||||
def order_tasks(tasks) -> list:
|
||||
"""按依赖拓扑排序,同一波内 priority 大的先。
|
||||
|
||||
tasks 是 dict 列表,每项含 id、depends_on(list,可缺省)、
|
||||
priority(int,可缺省 0)。缺失依赖或成环 raise ValueError。
|
||||
不修改输入列表,返回新列表。
|
||||
"""
|
||||
return _topo_sort(list(tasks))
|
||||
|
||||
|
||||
def ready_tasks(tasks, completed_ids) -> list:
|
||||
"""返回依赖全部在 completed_ids 里、且状态为 queued 的任务,
|
||||
再按 order_tasks 的规则(priority 降序、同波拓扑)排序。
|
||||
|
||||
未写 status 字段的任务按 queued 计;running/completed/failed 等
|
||||
状态的任务一律不返回。
|
||||
"""
|
||||
completed = set(completed_ids or [])
|
||||
ready = [t for t in tasks
|
||||
if t.get("status", "queued") == "queued"
|
||||
and all(d in completed for d in _depends_on(t))]
|
||||
return _topo_sort(ready, initial_done=completed, strict_missing=False)
|
||||
|
||||
|
||||
def budget_exceeded(usage_turns, max_turns, usage_cost_usd,
|
||||
max_cost_usd) -> bool:
|
||||
"""任一上限非 None 且用量 >= 上限即 True,否则 False。
|
||||
|
||||
上限为 None 表示该项不限;用量为 None 按 0 计。
|
||||
"""
|
||||
if max_turns is not None and (usage_turns or 0) >= max_turns:
|
||||
return True
|
||||
if max_cost_usd is not None and (usage_cost_usd or 0) >= max_cost_usd:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class TaskQueue:
|
||||
"""内存队列:add / mark_running / mark_done / next_ready。
|
||||
|
||||
内部存任务 dict 的浅拷贝,add 时缺省状态为 queued;status
|
||||
只能是 VALID_STATUSES 中的值,否则 raise ValueError。
|
||||
"""
|
||||
|
||||
def __init__(self, tasks=None):
|
||||
self._tasks = {}
|
||||
for task in tasks or []:
|
||||
self.add(task)
|
||||
|
||||
def add(self, task):
|
||||
task_id = _task_id(task)
|
||||
if task_id in self._tasks:
|
||||
raise ValueError(f"任务 id 重复:{task_id}")
|
||||
stored = dict(task)
|
||||
status = stored.get("status", "queued")
|
||||
if status not in VALID_STATUSES:
|
||||
raise ValueError(f"非法任务状态:{status}")
|
||||
stored["status"] = status
|
||||
self._tasks[task_id] = stored
|
||||
return stored
|
||||
|
||||
def mark_running(self, task_id):
|
||||
task = self._get(task_id)
|
||||
task["status"] = "running"
|
||||
return task
|
||||
|
||||
def mark_done(self, task_id, status):
|
||||
if status not in VALID_STATUSES:
|
||||
raise ValueError(f"非法任务状态:{status}")
|
||||
task = self._get(task_id)
|
||||
task["status"] = status
|
||||
return task
|
||||
|
||||
def next_ready(self, completed_ids):
|
||||
"""按 ready_tasks 顺序的第一个可派任务,或 None。"""
|
||||
ready = ready_tasks(list(self._tasks.values()), completed_ids)
|
||||
return ready[0] if ready else None
|
||||
|
||||
def _get(self, task_id):
|
||||
try:
|
||||
return self._tasks[task_id]
|
||||
except KeyError:
|
||||
raise ValueError(f"未知任务:{task_id}") from None
|
||||
@@ -0,0 +1,122 @@
|
||||
"""沙箱执行器:可插拔的任务执行接口。
|
||||
|
||||
零第三方依赖,只用 Python 标准库;不重写 agent 逻辑。
|
||||
|
||||
- Executor:基类,子类必须实现 execute(task)。
|
||||
- LocalExecutor:先行实现,复用 host.py 现有 worktree + codex 沙箱,
|
||||
通过 host.py CLI 子进程做只读查询封装;不直接跑模型。
|
||||
- DockerExecutor:占位,未接通容器执行。
|
||||
- CloudExecutor:占位,未接通云端执行。
|
||||
"""
|
||||
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
MAX_STDOUT_CHARS = 2000
|
||||
|
||||
|
||||
class Executor:
|
||||
"""执行器基类。
|
||||
|
||||
task 是 dict,含 id、project_id、repo、prompt、branch(可选)。
|
||||
execute(task) -> dict,至少含 executor、task_id、status、stdout。
|
||||
"""
|
||||
|
||||
def execute(self, task) -> dict:
|
||||
raise NotImplementedError("子类必须实现 execute(task)")
|
||||
|
||||
|
||||
def _truncate_stdout(text):
|
||||
if text is None:
|
||||
return ""
|
||||
return str(text)[:MAX_STDOUT_CHARS]
|
||||
|
||||
|
||||
def _default_runner(cmd):
|
||||
"""subprocess.run 的薄封装:捕获输出、文本模式。"""
|
||||
return subprocess.run(cmd, capture_output=True, text=True)
|
||||
|
||||
|
||||
class LocalExecutor(Executor):
|
||||
"""本地执行器:调 host.py CLI 子进程做执行封装。
|
||||
|
||||
真正跑模型回合的 dispatch/run 由控制面在确认模型 key 供给后才
|
||||
调用,本模块只做执行器骨架和 CLI 调用封装。因此 execute() 在此
|
||||
只发 host.py 的只读命令(例如 status),绝不在 execute 里直接
|
||||
跑模型回合;需要跑模型时,控制面另行调用 host.py dispatch/run。
|
||||
|
||||
构造函数:LocalExecutor(host_py_path, config_path=None, runner=None)。
|
||||
runner 默认是 subprocess.run 的薄封装;单测可注入假 runner,假
|
||||
runner 返回 CompletedProcess 风格的 result 对象(含 returncode、
|
||||
stdout),本类只读这两个字段,不依赖其它字段。
|
||||
"""
|
||||
|
||||
def __init__(self, host_py_path, config_path=None, runner=None):
|
||||
self.host_py_path = host_py_path
|
||||
self.config_path = config_path
|
||||
self.runner = runner if runner is not None else _default_runner
|
||||
|
||||
def _command(self, task_id):
|
||||
cmd = [sys.executable, self.host_py_path]
|
||||
if self.config_path is not None:
|
||||
cmd += ["--config", self.config_path]
|
||||
cmd += ["status", "--task", str(task_id)]
|
||||
return cmd
|
||||
|
||||
def execute(self, task) -> dict:
|
||||
task_id = task.get("id") if isinstance(task, dict) else None
|
||||
if task_id is None:
|
||||
raise ValueError("任务缺少 id")
|
||||
result = self.runner(self._command(task_id))
|
||||
returncode = getattr(result, "returncode", 1)
|
||||
stdout = _truncate_stdout(getattr(result, "stdout", ""))
|
||||
return {
|
||||
"executor": "local",
|
||||
"task_id": task_id,
|
||||
"status": "completed" if returncode == 0 else "failed",
|
||||
"stdout": stdout,
|
||||
"returncode": returncode,
|
||||
}
|
||||
|
||||
|
||||
class DockerExecutor(Executor):
|
||||
"""容器执行器(占位,未接通)。
|
||||
|
||||
本机无 docker(shutil.which("docker") 为 None)时 raise
|
||||
RuntimeError("docker not available"),绝不静默降级为本地执行。
|
||||
有 docker 时也只返回占位失败,不真跑容器。
|
||||
|
||||
后续接线:容器挂载(把任务 worktree 以只读/受限写挂载进容器)、
|
||||
网络隔离(默认无外网,仅放行模型 API 所需出口)、资源限额与
|
||||
容器回收,届时再替换本占位实现。
|
||||
"""
|
||||
|
||||
def execute(self, task) -> dict:
|
||||
task_id = task.get("id") if isinstance(task, dict) else None
|
||||
if task_id is None:
|
||||
raise ValueError("任务缺少 id")
|
||||
if shutil.which("docker") is None:
|
||||
raise RuntimeError("docker not available")
|
||||
return {
|
||||
"executor": "docker",
|
||||
"task_id": task_id,
|
||||
"status": "failed",
|
||||
"stdout": "docker executor not implemented: "
|
||||
"container mount and network isolation pending",
|
||||
"returncode": None,
|
||||
}
|
||||
|
||||
|
||||
class CloudExecutor(Executor):
|
||||
"""云端执行器(占位,未接通)。
|
||||
|
||||
execute(task) 一律 raise RuntimeError,绝不假装已接通云端。
|
||||
接通前提:真实云账号 + 预算上限配置,届时由控制面注入凭证与
|
||||
限额后再替换本占位实现。
|
||||
"""
|
||||
|
||||
def execute(self, task) -> dict:
|
||||
raise RuntimeError(
|
||||
"cloud executor not connected: missing cloud account "
|
||||
"and budget limit")
|
||||
+644
@@ -0,0 +1,644 @@
|
||||
"""云端控制面 HTTP server(首版)。
|
||||
|
||||
零第三方依赖,只用 http.server + sqlite3 + 标准库。
|
||||
|
||||
分工与边界(以实际代码为准,不夸大):
|
||||
- 本 server 只做控制面:项目/任务的创建、查询、取消、心跳、
|
||||
stale 任务回收,以及 GitHub/Slack webhook 的验签与解析。
|
||||
- 本版不直接跑模型回合。任务状态机与 host.py 对齐:
|
||||
queued -> running -> completed/failed/interrupted/merged/merge-failed,
|
||||
其中 completed/merged 由 worker/控制面后续调用方更新(例如 worker
|
||||
调 host.py 跑完回合后回写),本 server 不提供“跑模型”端点。
|
||||
- 其它已就绪模块(queue 的 order_tasks/ready_tasks、sandbox 的执行器、
|
||||
memory 的 MemoryStore、kb 的 KBStore、models 的 ModelRouter、
|
||||
cost 的 CostStore、webhooks 的 NotificationService)本版尚未接成
|
||||
HTTP 端点,留给 worker 接线时调用;本 server 只接入其中两处:
|
||||
AuthStore(认证/权限/审计)与 webhooks 的验签/解析函数。
|
||||
- DockerExecutor/CloudExecutor 目前仍是占位(见 sandbox.py),
|
||||
本 server 不假装已接通容器或云端执行。
|
||||
|
||||
导入规则:只用绝对包导入(from cloud.auth import ...)。绝不要把
|
||||
cloud/ 目录加进 sys.path:标准库已有 queue 模块,会和 cloud/queue.py
|
||||
撞名,导致 stdlib queue 被错误模块覆盖。
|
||||
|
||||
状态库:单一 SQLite 文件。AuthStore 用 users/project_members/
|
||||
audit_log 三张表;本 server 用 projects/tasks 两张表,表名不冲突。
|
||||
并发:ThreadingHTTPServer 多线程处理请求,每个请求内新建 SQLite
|
||||
连接(server 表用新连接,AuthStore 每次请求新建一个实例),绝不在
|
||||
多线程 handler 间共用一个连接。
|
||||
|
||||
启动:
|
||||
python3 -m cloud.server [--host 127.0.0.1] [--port 8765] [--db <路径>]
|
||||
|
||||
测试用 create_server(db_path, host, port) 只返回 ThreadingHTTPServer
|
||||
实例,不自动 serve。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.webhooks import (
|
||||
parse_github_pr_comment,
|
||||
parse_slack_command,
|
||||
verify_github_signature,
|
||||
verify_slack_signature,
|
||||
)
|
||||
|
||||
TERMINAL_STATUSES = frozenset(
|
||||
{"completed", "failed", "interrupted", "merged", "merge-failed"}
|
||||
)
|
||||
DEFAULT_SWEEP_THRESHOLD_SECONDS = 300
|
||||
|
||||
|
||||
def _connect(db_path):
|
||||
"""每个请求新建一个连接;busy timeout 给并发写留重试窗口。"""
|
||||
return sqlite3.connect(str(db_path), timeout=30)
|
||||
|
||||
|
||||
def init_db(db_path):
|
||||
"""创建 server 自己的 projects/tasks 表(AuthStore 的表由其自建)。"""
|
||||
conn = _connect(db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS projects("
|
||||
"id TEXT PRIMARY KEY, "
|
||||
"name TEXT, "
|
||||
"owner_id TEXT, "
|
||||
"created_at REAL)"
|
||||
)
|
||||
conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS tasks("
|
||||
"id TEXT PRIMARY KEY, "
|
||||
"project_id TEXT, "
|
||||
"prompt TEXT, "
|
||||
"status TEXT, "
|
||||
"idempotency_key TEXT, "
|
||||
"depends_on TEXT, "
|
||||
"priority INTEGER, "
|
||||
"created_by TEXT, "
|
||||
"created_at REAL, "
|
||||
"updated_at REAL, "
|
||||
"last_heartbeat REAL, "
|
||||
"UNIQUE(project_id, idempotency_key))"
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _auth_store(db_path):
|
||||
"""每请求一个 AuthStore 实例(自带新连接),用完即关闭。"""
|
||||
store = AuthStore(db_path)
|
||||
try:
|
||||
yield store
|
||||
finally:
|
||||
try:
|
||||
store.conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _task_to_dict(row):
|
||||
depends_on = row["depends_on"]
|
||||
try:
|
||||
depends_on = json.loads(depends_on) if depends_on else []
|
||||
except (TypeError, ValueError):
|
||||
depends_on = []
|
||||
if not isinstance(depends_on, list):
|
||||
depends_on = [depends_on]
|
||||
return {
|
||||
"id": row["id"],
|
||||
"project_id": row["project_id"],
|
||||
"prompt": row["prompt"],
|
||||
"status": row["status"],
|
||||
"idempotency_key": row["idempotency_key"],
|
||||
"depends_on": depends_on,
|
||||
"priority": row["priority"],
|
||||
"created_by": row["created_by"],
|
||||
"created_at": row["created_at"],
|
||||
"updated_at": row["updated_at"],
|
||||
"last_heartbeat": row["last_heartbeat"],
|
||||
}
|
||||
|
||||
|
||||
def _project_to_dict(row):
|
||||
return {
|
||||
"id": row["id"],
|
||||
"name": row["name"],
|
||||
"owner_id": row["owner_id"],
|
||||
"created_at": row["created_at"],
|
||||
}
|
||||
|
||||
|
||||
def sweep_stale_tasks(db_path, threshold_seconds=300):
|
||||
"""把 running 且心跳过期的任务改为 interrupted,返回被扫任务 id 列表。
|
||||
|
||||
心跳取 last_heartbeat,无心跳则用 updated_at;早于
|
||||
now - threshold_seconds 即过期。审计记一条 "sweep_interrupted"
|
||||
(user_id 记 "system",因为回收不是某个用户的操作)。
|
||||
独立可单测,不依赖 HTTP。
|
||||
"""
|
||||
now = time.time()
|
||||
cutoff = now - threshold_seconds
|
||||
conn = _connect(db_path)
|
||||
try:
|
||||
conn.row_factory = sqlite3.Row
|
||||
rows = conn.execute(
|
||||
"SELECT id, project_id FROM tasks WHERE status = 'running'"
|
||||
" AND COALESCE(last_heartbeat, updated_at) < ?",
|
||||
(cutoff,),
|
||||
).fetchall()
|
||||
swept = []
|
||||
for row in rows:
|
||||
cur = conn.execute(
|
||||
"UPDATE tasks SET status = 'interrupted', updated_at = ?"
|
||||
" WHERE id = ? AND status = 'running'",
|
||||
(now, row["id"]),
|
||||
)
|
||||
if cur.rowcount:
|
||||
swept.append(row["id"])
|
||||
conn.commit()
|
||||
stale = [(row["id"], row["project_id"]) for row in rows
|
||||
if row["id"] in swept]
|
||||
finally:
|
||||
conn.close()
|
||||
with _auth_store(db_path) as store:
|
||||
for task_id, project_id in stale:
|
||||
store.audit("system", project_id, "sweep_interrupted", task_id)
|
||||
return swept
|
||||
|
||||
|
||||
class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
"""控制面请求处理:除 /health 与 /webhooks/* 外一律要求 X-API-Key。"""
|
||||
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
# -- 通用工具 ------------------------------------------------------
|
||||
def log_message(self, *args): # 测试输出保持干净
|
||||
pass
|
||||
|
||||
def _send_json(self, status, obj):
|
||||
body = json.dumps(obj).encode("utf-8")
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def _read_body(self):
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length") or 0)
|
||||
except ValueError:
|
||||
length = 0
|
||||
if length <= 0:
|
||||
return b""
|
||||
return self.rfile.read(length)
|
||||
|
||||
def _read_json(self):
|
||||
"""返回 (obj, error):空 body 视为 {};非法 JSON 返回 error。"""
|
||||
raw = self._read_body()
|
||||
if not raw:
|
||||
return {}, None
|
||||
try:
|
||||
obj = json.loads(raw.decode("utf-8"))
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
return None, "invalid json"
|
||||
if not isinstance(obj, dict):
|
||||
return None, "body must be a JSON object"
|
||||
return obj, None
|
||||
|
||||
def _db(self):
|
||||
conn = _connect(self.server.db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
def _authenticate(self):
|
||||
"""X-API-Key -> user_id;无 key 或 key 错返回 None。"""
|
||||
api_key = self.headers.get("X-API-Key")
|
||||
if not api_key:
|
||||
return None
|
||||
with _auth_store(self.server.db_path) as store:
|
||||
return store.authenticate(api_key)
|
||||
|
||||
# -- 入口 ----------------------------------------------------------
|
||||
def do_GET(self):
|
||||
self._dispatch("GET")
|
||||
|
||||
def do_POST(self):
|
||||
self._dispatch("POST")
|
||||
|
||||
def _dispatch(self, method):
|
||||
parsed = urlparse(self.path)
|
||||
path = parsed.path
|
||||
if len(path) > 1 and path.endswith("/"):
|
||||
path = path.rstrip("/")
|
||||
query = parse_qs(parsed.query)
|
||||
|
||||
if path == "/health" and method == "GET":
|
||||
self._send_json(200, {"ok": True, "service": "cloud-control-plane"})
|
||||
return
|
||||
if path.startswith("/webhooks/"):
|
||||
# webhook 靠签名认证(外部服务无法持 X-API-Key),不走 API key。
|
||||
if method != "POST":
|
||||
self._send_json(404, {"error": "not found"})
|
||||
return
|
||||
if path == "/webhooks/github":
|
||||
self._handle_github_webhook()
|
||||
elif path == "/webhooks/slack":
|
||||
self._handle_slack_webhook()
|
||||
else:
|
||||
self._send_json(404, {"error": "not found"})
|
||||
return
|
||||
|
||||
user_id = self._authenticate()
|
||||
if user_id is None:
|
||||
self._send_json(401, {"error": "unauthorized"})
|
||||
return
|
||||
|
||||
if path == "/projects" and method == "POST":
|
||||
self._create_project(user_id)
|
||||
elif path == "/projects" and method == "GET":
|
||||
self._list_projects(user_id)
|
||||
elif path == "/tasks" and method == "POST":
|
||||
self._create_task(user_id)
|
||||
elif path == "/tasks" and method == "GET":
|
||||
self._list_tasks(user_id, query)
|
||||
elif path == "/admin/sweep" and method == "POST":
|
||||
self._admin_sweep()
|
||||
elif path.startswith("/tasks/"):
|
||||
rest = path[len("/tasks/"):]
|
||||
if "/" not in rest and method == "GET":
|
||||
self._get_task(user_id, rest)
|
||||
elif rest.endswith("/cancel") and method == "POST":
|
||||
self._cancel_task(user_id, rest[: -len("/cancel")])
|
||||
elif rest.endswith("/heartbeat") and method == "POST":
|
||||
self._heartbeat_task(user_id, rest[: -len("/heartbeat")])
|
||||
else:
|
||||
self._send_json(404, {"error": "not found"})
|
||||
else:
|
||||
self._send_json(404, {"error": "not found"})
|
||||
|
||||
# -- 项目 ----------------------------------------------------------
|
||||
def _create_project(self, user_id):
|
||||
body, err = self._read_json()
|
||||
if err:
|
||||
self._send_json(400, {"error": err})
|
||||
return
|
||||
name = body.get("name")
|
||||
if not name or not isinstance(name, str):
|
||||
self._send_json(400, {"error": "name required"})
|
||||
return
|
||||
project_id = body.get("id") or uuid.uuid4().hex
|
||||
if not isinstance(project_id, str):
|
||||
self._send_json(400, {"error": "id must be a string"})
|
||||
return
|
||||
now = time.time()
|
||||
conn = self._db()
|
||||
try:
|
||||
existing = conn.execute(
|
||||
"SELECT id FROM projects WHERE id = ?", (project_id,)
|
||||
).fetchone()
|
||||
if existing:
|
||||
self._send_json(409, {"error": "project exists"})
|
||||
return
|
||||
conn.execute(
|
||||
"INSERT INTO projects(id, name, owner_id, created_at)"
|
||||
" VALUES(?, ?, ?, ?)",
|
||||
(project_id, name, user_id, now),
|
||||
)
|
||||
conn.commit()
|
||||
project = {
|
||||
"id": project_id,
|
||||
"name": name,
|
||||
"owner_id": user_id,
|
||||
"created_at": now,
|
||||
}
|
||||
except sqlite3.IntegrityError:
|
||||
self._send_json(409, {"error": "project exists"})
|
||||
return
|
||||
finally:
|
||||
conn.close()
|
||||
# AuthStore.create_project 只写成员表(owner 角色),项目记录
|
||||
# 本身在上面的 projects 表;两者配合才算一个完整项目。
|
||||
with _auth_store(self.server.db_path) as store:
|
||||
store.create_project(project_id, user_id)
|
||||
store.audit(user_id, project_id, "create_project", name)
|
||||
self._send_json(201, project)
|
||||
|
||||
def _list_projects(self, user_id):
|
||||
conn = self._db()
|
||||
try:
|
||||
member_ids = {
|
||||
row["project_id"]
|
||||
for row in conn.execute(
|
||||
"SELECT project_id FROM project_members WHERE user_id = ?",
|
||||
(user_id,),
|
||||
).fetchall()
|
||||
}
|
||||
rows = conn.execute(
|
||||
"SELECT id, name, owner_id, created_at FROM projects"
|
||||
" ORDER BY created_at, id"
|
||||
).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
projects = [_project_to_dict(r) for r in rows
|
||||
if r["id"] in member_ids]
|
||||
self._send_json(200, {"projects": projects})
|
||||
|
||||
# -- 任务 ----------------------------------------------------------
|
||||
def _create_task(self, user_id):
|
||||
body, err = self._read_json()
|
||||
if err:
|
||||
self._send_json(400, {"error": err})
|
||||
return
|
||||
project_id = body.get("project_id")
|
||||
prompt = body.get("prompt")
|
||||
if not project_id or not isinstance(project_id, str):
|
||||
self._send_json(400, {"error": "project_id required"})
|
||||
return
|
||||
if not prompt or not isinstance(prompt, str):
|
||||
self._send_json(400, {"error": "prompt required"})
|
||||
return
|
||||
depends_on = body.get("depends_on") or []
|
||||
if isinstance(depends_on, str):
|
||||
depends_on = [depends_on]
|
||||
if not isinstance(depends_on, list):
|
||||
self._send_json(400, {"error": "depends_on must be a list"})
|
||||
return
|
||||
try:
|
||||
priority = int(body.get("priority", 0) or 0)
|
||||
except (TypeError, ValueError):
|
||||
self._send_json(400, {"error": "priority must be an integer"})
|
||||
return
|
||||
idempotency_key = body.get("idempotency_key")
|
||||
if idempotency_key is not None and not isinstance(idempotency_key, str):
|
||||
self._send_json(400, {"error": "idempotency_key must be a string"})
|
||||
return
|
||||
with _auth_store(self.server.db_path) as store:
|
||||
if not store.check_permission(user_id, project_id, "create_task"):
|
||||
self._send_json(403, {"error": "forbidden"})
|
||||
return
|
||||
now = time.time()
|
||||
conn = self._db()
|
||||
try:
|
||||
if idempotency_key is not None:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM tasks WHERE project_id = ?"
|
||||
" AND idempotency_key = ?",
|
||||
(project_id, idempotency_key),
|
||||
).fetchone()
|
||||
if row:
|
||||
self._send_json(200, _task_to_dict(row))
|
||||
return
|
||||
task_id = uuid.uuid4().hex
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||||
" idempotency_key, depends_on, priority, created_by,"
|
||||
" created_at, updated_at, last_heartbeat)"
|
||||
" VALUES(?, ?, ?, 'queued', ?, ?, ?, ?, ?, ?, NULL)",
|
||||
(task_id, project_id, prompt, idempotency_key,
|
||||
json.dumps(depends_on), priority, user_id, now, now),
|
||||
)
|
||||
conn.commit()
|
||||
except sqlite3.IntegrityError:
|
||||
# 并发下同幂等键抢建:回读已有任务,不建重复。
|
||||
row = conn.execute(
|
||||
"SELECT * FROM tasks WHERE project_id = ?"
|
||||
" AND idempotency_key = ?",
|
||||
(project_id, idempotency_key),
|
||||
).fetchone()
|
||||
if row:
|
||||
self._send_json(200, _task_to_dict(row))
|
||||
return
|
||||
raise
|
||||
row = conn.execute(
|
||||
"SELECT * FROM tasks WHERE id = ?", (task_id,)
|
||||
).fetchone()
|
||||
task = _task_to_dict(row)
|
||||
finally:
|
||||
conn.close()
|
||||
with _auth_store(self.server.db_path) as store:
|
||||
store.audit(user_id, project_id, "create_task", task_id)
|
||||
self._send_json(201, task)
|
||||
|
||||
def _list_tasks(self, user_id, query):
|
||||
project_ids = query.get("project_id")
|
||||
project_id = project_ids[0] if project_ids else None
|
||||
if not project_id:
|
||||
self._send_json(400, {"error": "project_id required"})
|
||||
return
|
||||
with _auth_store(self.server.db_path) as store:
|
||||
if not store.check_permission(user_id, project_id, "read"):
|
||||
self._send_json(403, {"error": "forbidden"})
|
||||
return
|
||||
conn = self._db()
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM tasks WHERE project_id = ?"
|
||||
" ORDER BY created_at, id",
|
||||
(project_id,),
|
||||
).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
self._send_json(200, {"tasks": [_task_to_dict(r) for r in rows]})
|
||||
|
||||
def _fetch_task(self, task_id):
|
||||
conn = self._db()
|
||||
try:
|
||||
return conn.execute(
|
||||
"SELECT * FROM tasks WHERE id = ?", (task_id,)
|
||||
).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _get_task(self, user_id, task_id):
|
||||
row = self._fetch_task(task_id)
|
||||
if row is None:
|
||||
self._send_json(404, {"error": "task not found"})
|
||||
return
|
||||
task = _task_to_dict(row)
|
||||
with _auth_store(self.server.db_path) as store:
|
||||
if not store.check_permission(user_id, task["project_id"], "read"):
|
||||
self._send_json(403, {"error": "forbidden"})
|
||||
return
|
||||
self._send_json(200, task)
|
||||
|
||||
def _cancel_task(self, user_id, task_id):
|
||||
row = self._fetch_task(task_id)
|
||||
if row is None:
|
||||
self._send_json(404, {"error": "task not found"})
|
||||
return
|
||||
task = _task_to_dict(row)
|
||||
with _auth_store(self.server.db_path) as store:
|
||||
if not store.check_permission(
|
||||
user_id, task["project_id"], "create_task"):
|
||||
self._send_json(403, {"error": "forbidden"})
|
||||
return
|
||||
if task["status"] in TERMINAL_STATUSES:
|
||||
self._send_json(400, {"error": "task already terminal"})
|
||||
return
|
||||
now = time.time()
|
||||
conn = self._db()
|
||||
try:
|
||||
conn.execute(
|
||||
"UPDATE tasks SET status = 'interrupted', updated_at = ?"
|
||||
" WHERE id = ?",
|
||||
(now, task_id),
|
||||
)
|
||||
conn.commit()
|
||||
task = _task_to_dict(
|
||||
conn.execute(
|
||||
"SELECT * FROM tasks WHERE id = ?", (task_id,)
|
||||
).fetchone()
|
||||
)
|
||||
finally:
|
||||
conn.close()
|
||||
with _auth_store(self.server.db_path) as store:
|
||||
store.audit(user_id, task["project_id"], "cancel_task", task_id)
|
||||
self._send_json(200, task)
|
||||
|
||||
def _heartbeat_task(self, user_id, task_id):
|
||||
row = self._fetch_task(task_id)
|
||||
if row is None:
|
||||
self._send_json(404, {"error": "task not found"})
|
||||
return
|
||||
task = _task_to_dict(row)
|
||||
with _auth_store(self.server.db_path) as store:
|
||||
if not store.check_permission(
|
||||
user_id, task["project_id"], "create_task"):
|
||||
self._send_json(403, {"error": "forbidden"})
|
||||
return
|
||||
now = time.time()
|
||||
new_status = "running" if task["status"] == "queued" else task["status"]
|
||||
conn = self._db()
|
||||
try:
|
||||
conn.execute(
|
||||
"UPDATE tasks SET status = ?, last_heartbeat = ?,"
|
||||
" updated_at = ? WHERE id = ?",
|
||||
(new_status, now, now, task_id),
|
||||
)
|
||||
conn.commit()
|
||||
task = _task_to_dict(
|
||||
conn.execute(
|
||||
"SELECT * FROM tasks WHERE id = ?", (task_id,)
|
||||
).fetchone()
|
||||
)
|
||||
finally:
|
||||
conn.close()
|
||||
self._send_json(200, task)
|
||||
|
||||
# -- 管理 ----------------------------------------------------------
|
||||
def _admin_sweep(self):
|
||||
body, err = self._read_json()
|
||||
if err:
|
||||
self._send_json(400, {"error": err})
|
||||
return
|
||||
threshold = body.get(
|
||||
"threshold_seconds", DEFAULT_SWEEP_THRESHOLD_SECONDS)
|
||||
try:
|
||||
threshold = int(threshold)
|
||||
except (TypeError, ValueError):
|
||||
self._send_json(400, {"error": "threshold_seconds must be an integer"})
|
||||
return
|
||||
swept = sweep_stale_tasks(self.server.db_path, threshold)
|
||||
self._send_json(200, {"swept": swept})
|
||||
|
||||
# -- webhooks(签名认证,不走 X-API-Key)---------------------------
|
||||
def _handle_github_webhook(self):
|
||||
raw = self._read_body()
|
||||
secret = os.environ.get("GITHUB_WEBHOOK_SECRET")
|
||||
if not secret:
|
||||
self._send_json(503, {"error": "github webhook secret not configured"})
|
||||
return
|
||||
signature = self.headers.get("X-Hub-Signature-256")
|
||||
if not verify_github_signature(secret, raw, signature):
|
||||
self._send_json(401, {"error": "unauthorized"})
|
||||
return
|
||||
try:
|
||||
payload = json.loads(raw.decode("utf-8")) if raw else {}
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
self._send_json(400, {"error": "invalid json"})
|
||||
return
|
||||
if not isinstance(payload, dict):
|
||||
self._send_json(400, {"error": "body must be a JSON object"})
|
||||
return
|
||||
# 只解析不执行:/fix-review 的真正执行由控制面后续调 host.py 完成。
|
||||
result = parse_github_pr_comment(payload)
|
||||
if result is None:
|
||||
self._send_json(200, {"result": None})
|
||||
else:
|
||||
self._send_json(200, result)
|
||||
|
||||
def _handle_slack_webhook(self):
|
||||
raw = self._read_body()
|
||||
secret = os.environ.get("SLACK_SIGNING_SECRET")
|
||||
if not secret:
|
||||
self._send_json(503, {"error": "slack webhook secret not configured"})
|
||||
return
|
||||
signature = self.headers.get("X-Slack-Signature")
|
||||
timestamp = self.headers.get("X-Slack-Request-Timestamp")
|
||||
if not verify_slack_signature(secret, timestamp, raw, signature):
|
||||
self._send_json(401, {"error": "unauthorized"})
|
||||
return
|
||||
# 只解析不执行任务:Slack 斜杠命令的执行由控制面后续决定。
|
||||
content_type = self.headers.get("Content-Type") or ""
|
||||
text = ""
|
||||
if "application/x-www-form-urlencoded" in content_type:
|
||||
parsed = parse_qs(raw.decode("utf-8", "replace"))
|
||||
values = parsed.get("text")
|
||||
text = values[0] if values else ""
|
||||
elif raw:
|
||||
try:
|
||||
payload = json.loads(raw.decode("utf-8"))
|
||||
if isinstance(payload, dict):
|
||||
text = payload.get("text", "") or ""
|
||||
else:
|
||||
text = raw.decode("utf-8", "replace")
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
text = raw.decode("utf-8", "replace")
|
||||
self._send_json(200, parse_slack_command(text))
|
||||
|
||||
|
||||
def create_server(db_path, host="127.0.0.1", port=8765):
|
||||
"""建 server 实例(不自动 serve,供测试与 __main__ 共用)。"""
|
||||
db_path = str(db_path)
|
||||
init_db(db_path)
|
||||
# 确保 AuthStore 的三张表也存在,后续请求才能直接查 project_members。
|
||||
with _auth_store(db_path):
|
||||
pass
|
||||
server = ThreadingHTTPServer((host, port), CloudControlPlaneHandler)
|
||||
server.db_path = db_path
|
||||
server.daemon_threads = True
|
||||
return server
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
parser = argparse.ArgumentParser(description="cloud control plane")
|
||||
parser.add_argument("--host", default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8765)
|
||||
parser.add_argument("--db", default="cloud-control-plane.db")
|
||||
args = parser.parse_args(argv)
|
||||
server = create_server(args.db, args.host, args.port)
|
||||
print(
|
||||
"cloud-control-plane listening on %s:%d db=%s"
|
||||
% (args.host, server.server_address[1], args.db),
|
||||
flush=True,
|
||||
)
|
||||
try:
|
||||
server.serve_forever()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,149 @@
|
||||
"""auth 模块单测:API key 认证、角色权限、审计只增不减。"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.auth import ROLES, AuthStore
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import ROLES, AuthStore
|
||||
|
||||
|
||||
class TestAuth(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.store = AuthStore(os.path.join(self.tmp.name, "auth.db"))
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _user(self, username):
|
||||
return self.store.create_user(username)
|
||||
|
||||
def _project_with_roles(self):
|
||||
"""owner + viewer + contributor + maintainer + 一个非成员,返回其 id。"""
|
||||
owner_id, _ = self._user("owner")
|
||||
viewer_id, _ = self._user("viewer")
|
||||
contributor_id, _ = self._user("contributor")
|
||||
maintainer_id, _ = self._user("maintainer")
|
||||
outsider_id, _ = self._user("outsider")
|
||||
project_id = "proj-1"
|
||||
self.store.create_project(project_id, owner_id)
|
||||
self.store.set_member(project_id, owner_id, viewer_id, "viewer")
|
||||
self.store.set_member(project_id, owner_id, contributor_id, "contributor")
|
||||
self.store.set_member(project_id, owner_id, maintainer_id, "maintainer")
|
||||
return project_id, owner_id, viewer_id, contributor_id, maintainer_id, outsider_id
|
||||
|
||||
def test_create_user_plain_key_authenticates(self):
|
||||
user_id, api_key_plain = self._user("alice")
|
||||
self.assertEqual(self.store.authenticate(api_key_plain), user_id)
|
||||
|
||||
def test_wrong_key_rejected(self):
|
||||
user_id, api_key_plain = self._user("alice")
|
||||
self.assertIsNone(self.store.authenticate(api_key_plain + "x"))
|
||||
self.assertIsNone(self.store.authenticate("not-a-real-key"))
|
||||
self.assertIsNone(self.store.authenticate(""))
|
||||
self.assertIsNone(self.store.authenticate(None))
|
||||
|
||||
def test_hash_stored_not_plaintext(self):
|
||||
_user_id, api_key_plain = self._user("alice")
|
||||
db_path = os.path.join(self.tmp.name, "auth.db")
|
||||
conn = sqlite3.connect(db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT api_key_hash FROM users WHERE username = ?", ("alice",)
|
||||
).fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
self.assertIsNotNone(row)
|
||||
stored = row[0]
|
||||
self.assertNotEqual(stored, api_key_plain)
|
||||
self.assertNotIn(api_key_plain, stored)
|
||||
# 库里只有哈希列,没有明文列可查
|
||||
columns = [c[1] for c in sqlite3.connect(db_path).execute(
|
||||
"PRAGMA table_info(users)").fetchall()]
|
||||
self.assertNotIn("api_key_plain", columns)
|
||||
self.assertNotIn("api_key", columns)
|
||||
|
||||
def test_permission_denials(self):
|
||||
project_id, _o, viewer_id, contributor_id, _m, outsider_id = (
|
||||
self._project_with_roles()
|
||||
)
|
||||
# viewer 不能 create_task
|
||||
self.assertFalse(
|
||||
self.store.check_permission(viewer_id, project_id, "create_task"))
|
||||
with self.assertRaises(PermissionError):
|
||||
self.store.require_permission(viewer_id, project_id, "create_task")
|
||||
# contributor 不能 approve_merge
|
||||
self.assertFalse(
|
||||
self.store.check_permission(contributor_id, project_id, "approve_merge"))
|
||||
with self.assertRaises(PermissionError):
|
||||
self.store.require_permission(
|
||||
contributor_id, project_id, "approve_merge")
|
||||
# 非成员读不到
|
||||
self.assertFalse(
|
||||
self.store.check_permission(outsider_id, project_id, "read"))
|
||||
self.assertFalse(
|
||||
self.store.check_permission(outsider_id, project_id, "create_task"))
|
||||
with self.assertRaises(PermissionError):
|
||||
self.store.require_permission(outsider_id, project_id, "read")
|
||||
|
||||
def test_permission_grants(self):
|
||||
project_id, owner_id, viewer_id, contributor_id, maintainer_id, _ = (
|
||||
self._project_with_roles()
|
||||
)
|
||||
self.assertTrue(self.store.check_permission(viewer_id, project_id, "read"))
|
||||
self.assertTrue(
|
||||
self.store.check_permission(contributor_id, project_id, "create_task"))
|
||||
self.assertTrue(
|
||||
self.store.check_permission(contributor_id, project_id, "read"))
|
||||
self.assertTrue(
|
||||
self.store.check_permission(maintainer_id, project_id, "approve_merge"))
|
||||
self.assertTrue(
|
||||
self.store.check_permission(owner_id, project_id, "manage_project"))
|
||||
self.assertFalse(
|
||||
self.store.check_permission(maintainer_id, project_id, "manage_project"))
|
||||
self.assertEqual(
|
||||
ROLES, ["viewer", "contributor", "maintainer", "owner"])
|
||||
|
||||
def test_owner_can_set_member_maintainer_cannot(self):
|
||||
project_id, owner_id, _v, _c, maintainer_id, _o = (
|
||||
self._project_with_roles()
|
||||
)
|
||||
target_id, _ = self._user("target")
|
||||
# owner 改成员成功
|
||||
self.store.set_member(project_id, owner_id, target_id, "viewer")
|
||||
self.assertTrue(self.store.check_permission(target_id, project_id, "read"))
|
||||
self.assertFalse(
|
||||
self.store.check_permission(target_id, project_id, "create_task"))
|
||||
# maintainer 改成员被拒,且原角色不变
|
||||
with self.assertRaises(PermissionError):
|
||||
self.store.set_member(
|
||||
project_id, maintainer_id, target_id, "contributor")
|
||||
self.assertFalse(
|
||||
self.store.check_permission(target_id, project_id, "create_task"))
|
||||
|
||||
def test_audit_append_only(self):
|
||||
project_id, owner_id, _v, _c, _m, _o = self._project_with_roles()
|
||||
self.assertFalse(hasattr(AuthStore, "delete_audit"))
|
||||
self.assertFalse(hasattr(AuthStore, "update_audit"))
|
||||
before = len(self.store.list_audit(project_id))
|
||||
self.store.audit(owner_id, project_id, "create_task", "task-1")
|
||||
self.store.audit(owner_id, project_id, "approve_merge", "task-1")
|
||||
entries = self.store.list_audit(project_id)
|
||||
self.assertEqual(len(entries), before + 2)
|
||||
# 最新两条在前(倒序),且均为本次追加的动作
|
||||
actions = [e["action"] for e in entries[:2]]
|
||||
self.assertEqual(actions, ["approve_merge", "create_task"])
|
||||
# 审计条目为 dict 且含关键字段
|
||||
self.assertIn("user_id", entries[0])
|
||||
self.assertIn("created_at", entries[0])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,182 @@
|
||||
"""Tests for cloud/memory.py (MemoryStore) and cloud/kb.py (KBStore).
|
||||
|
||||
Run from ~/workspace/codex-app-server:
|
||||
python3 -m unittest discover -s cloud/tests -p "test_memory_kb.py" -v
|
||||
|
||||
No network, no real model turns, no secrets. Each test uses a fresh
|
||||
temporary SQLite file.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.normpath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
|
||||
from kb import KBStore # noqa: E402
|
||||
from memory import MemoryStore # noqa: E402
|
||||
|
||||
|
||||
def _tmp_db():
|
||||
fd, path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
os.unlink(path) # let sqlite create it fresh
|
||||
return path
|
||||
|
||||
|
||||
class MemoryStoreTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.db = _tmp_db()
|
||||
self.store = MemoryStore(self.db)
|
||||
|
||||
def tearDown(self):
|
||||
try:
|
||||
self.store.conn.close()
|
||||
finally:
|
||||
if os.path.exists(self.db):
|
||||
os.unlink(self.db)
|
||||
|
||||
def test_add_then_list(self):
|
||||
sid = self.store.add_summary(
|
||||
"proj-a", "thread-1", "add login page", "login page done",
|
||||
files_changed=["app.py"], open_issues=["tests pending"],
|
||||
)
|
||||
self.assertIsInstance(sid, int)
|
||||
rows = self.store.list_summaries("proj-a")
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertEqual(rows[0]["thread_id"], "thread-1")
|
||||
self.assertEqual(rows[0]["goal"], "add login page")
|
||||
self.assertEqual(rows[0]["files_changed"], ["app.py"])
|
||||
self.assertEqual(rows[0]["open_issues"], ["tests pending"])
|
||||
self.assertEqual(rows[0]["status"], "active")
|
||||
|
||||
def test_deprecate_hidden_by_default_visible_with_flag(self):
|
||||
sid = self.store.add_summary("proj-a", "t1", "g1", "c1")
|
||||
self.store.add_summary("proj-a", "t2", "g2", "c2")
|
||||
self.store.deprecate(sid)
|
||||
default_rows = self.store.list_summaries("proj-a")
|
||||
self.assertEqual([r["thread_id"] for r in default_rows], ["t2"])
|
||||
all_rows = self.store.list_summaries("proj-a", include_deprecated=True)
|
||||
self.assertEqual(len(all_rows), 2)
|
||||
deprecated = [r for r in all_rows if r["id"] == sid]
|
||||
self.assertEqual(len(deprecated), 1)
|
||||
self.assertEqual(deprecated[0]["status"], "deprecated")
|
||||
|
||||
def test_delete_removes(self):
|
||||
sid = self.store.add_summary("proj-a", "t1", "g1", "c1")
|
||||
self.store.delete(sid)
|
||||
self.assertEqual(self.store.list_summaries("proj-a"), [])
|
||||
self.assertEqual(
|
||||
self.store.list_summaries("proj-a", include_deprecated=True), []
|
||||
)
|
||||
|
||||
def test_build_prompt_context_top_k_and_excludes_deprecated(self):
|
||||
# 7 active summaries; only the newest 5 may appear.
|
||||
for i in range(7):
|
||||
self.store.add_summary("proj-a", "t%d" % i, "goal-%d" % i, "done-%d" % i)
|
||||
ctx = self.store.build_prompt_context("proj-a", top_k=5)
|
||||
self.assertEqual(ctx.count("\n") + 1, 5)
|
||||
for i in range(2, 7):
|
||||
self.assertIn("goal-%d" % i, ctx)
|
||||
for i in range(2):
|
||||
self.assertNotIn("goal-%d" % i, ctx)
|
||||
# Format spot-check on one line.
|
||||
self.assertIn("- [t6] 目标:goal-6;结论:done-6;", ctx)
|
||||
# A deprecated newest summary must not appear even within top_k.
|
||||
sid = self.store.add_summary("proj-a", "t-new", "goal-new", "done-new")
|
||||
self.store.deprecate(sid)
|
||||
ctx2 = self.store.build_prompt_context("proj-a", top_k=5)
|
||||
self.assertNotIn("goal-new", ctx2)
|
||||
self.assertIn("goal-6", ctx2)
|
||||
|
||||
def test_build_prompt_context_empty(self):
|
||||
self.assertEqual(self.store.build_prompt_context("proj-empty"), "")
|
||||
|
||||
|
||||
class KBStoreTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.db = _tmp_db()
|
||||
self.store = KBStore(self.db)
|
||||
|
||||
def tearDown(self):
|
||||
try:
|
||||
self.store.conn.close()
|
||||
finally:
|
||||
if os.path.exists(self.db):
|
||||
os.unlink(self.db)
|
||||
|
||||
def test_search_finds_file_by_keyword(self):
|
||||
self.store.index_file(
|
||||
"proj-a", "docs/auth.md",
|
||||
"authentication flow uses token refresh and session cookies",
|
||||
kind="doc",
|
||||
)
|
||||
self.store.index_file(
|
||||
"proj-a", "docs/other.md", "unrelated gardening notes", kind="doc"
|
||||
)
|
||||
results = self.store.search("proj-a", "authentication")
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0]["path"], "docs/auth.md")
|
||||
self.assertEqual(results[0]["kind"], "doc")
|
||||
self.assertIn("authentication", results[0]["snippet"])
|
||||
self.assertLessEqual(len(results[0]["snippet"]), 200)
|
||||
|
||||
def test_search_project_isolation(self):
|
||||
self.store.index_file(
|
||||
"proj-a", "docs/secret.md", "zebra secret keyword", kind="doc"
|
||||
)
|
||||
self.store.index_file(
|
||||
"proj-b", "docs/plain.md", "plain notes with no keyword", kind="doc"
|
||||
)
|
||||
# A's doc must never leak into B's search, and vice versa.
|
||||
self.assertEqual(self.store.search("proj-b", "zebra"), [])
|
||||
self.assertEqual(self.store.search("proj-a", "zebra")[0]["path"], "docs/secret.md")
|
||||
|
||||
def test_reindex_same_path_no_duplicates(self):
|
||||
self.store.index_file("proj-a", "README.md", "version one content", kind="readme")
|
||||
self.store.index_file("proj-a", "README.md", "version two content", kind="readme")
|
||||
docs = self.store.list_docs("proj-a")
|
||||
self.assertEqual(len(docs), 1)
|
||||
self.assertEqual(docs[0]["path"], "README.md")
|
||||
self.assertEqual(docs[0]["content"], "version two content")
|
||||
# Search must also see only the new version, once.
|
||||
self.assertEqual(len(self.store.search("proj-a", "version")), 1)
|
||||
self.assertEqual(self.store.search("proj-a", "version")[0]["snippet"],
|
||||
"version two content")
|
||||
self.assertEqual(self.store.search("proj-a", "one"), [])
|
||||
|
||||
def test_index_task_and_list_docs(self):
|
||||
tid = self.store.index_task(
|
||||
"proj-a", "task-42", "fix the login bug", "fixed by resetting tokens"
|
||||
)
|
||||
self.assertIsInstance(tid, int)
|
||||
docs = self.store.list_docs("proj-a")
|
||||
self.assertEqual(len(docs), 1)
|
||||
self.assertEqual(docs[0]["path"], "task:task-42")
|
||||
self.assertEqual(docs[0]["kind"], "task")
|
||||
results = self.store.search("proj-a", "login")
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0]["path"], "task:task-42")
|
||||
|
||||
def test_search_chinese_keyword(self):
|
||||
# FTS5 unicode61 does not tokenize CJK words, so a verbatim
|
||||
# Chinese keyword must still be found via the LIKE fallback.
|
||||
self.store.index_file(
|
||||
"proj-a", "docs/login.md", "用户登录流程使用令牌刷新", kind="doc"
|
||||
)
|
||||
results = self.store.search("proj-a", "登录")
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(results[0]["path"], "docs/login.md")
|
||||
|
||||
def test_search_top_k_and_empty_query(self):
|
||||
for i in range(4):
|
||||
self.store.index_file(
|
||||
"proj-a", "doc%d.md" % i, "shared keyword doc %d" % i, kind="doc"
|
||||
)
|
||||
self.assertEqual(len(self.store.search("proj-a", "shared", top_k=2)), 2)
|
||||
self.assertEqual(self.store.search("proj-a", ""), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Tests for cloud.models (fallback routing) and cloud.cost.
|
||||
|
||||
No real model turns, no network, no secrets: usage clients are fakes
|
||||
and the CostStore uses an in-memory SQLite database.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from datetime import datetime, timezone
|
||||
|
||||
# cloud/ has no __init__.py by design (the coordinator creates it), so
|
||||
# make the project root importable regardless of the discover top level.
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(
|
||||
os.path.dirname(os.path.abspath(__file__)))))
|
||||
|
||||
from cloud.cost import CostStore, estimate_tokens, fetch_account_usage
|
||||
from cloud.models import FALLBACKABLE, ModelRouter, is_fallbackable, record_fallback
|
||||
|
||||
|
||||
class FakeError(Exception):
|
||||
def __init__(self, status=None, code=None):
|
||||
super().__init__(status or code)
|
||||
self.status = status
|
||||
self.code = code
|
||||
|
||||
|
||||
class TestIsFallbackable(unittest.TestCase):
|
||||
def test_fallbackable_status_codes(self):
|
||||
for code in (404, 429, 500, 502, 503, 504):
|
||||
self.assertTrue(is_fallbackable(code), code)
|
||||
self.assertIn(404, FALLBACKABLE)
|
||||
self.assertIn(429, FALLBACKABLE)
|
||||
self.assertIn("timeout", FALLBACKABLE)
|
||||
|
||||
def test_timeout_string(self):
|
||||
self.assertTrue(is_fallbackable("timeout"))
|
||||
|
||||
def test_not_fallbackable(self):
|
||||
for code in (400, 401, 403, 200):
|
||||
self.assertFalse(is_fallbackable(code), code)
|
||||
self.assertFalse(is_fallbackable("other"))
|
||||
|
||||
def test_object_with_status_or_code(self):
|
||||
self.assertTrue(is_fallbackable(FakeError(status=429)))
|
||||
self.assertTrue(is_fallbackable(FakeError(code=503)))
|
||||
self.assertTrue(is_fallbackable(FakeError(status="timeout")))
|
||||
self.assertFalse(is_fallbackable(FakeError(status=401)))
|
||||
self.assertFalse(is_fallbackable(FakeError(code=400)))
|
||||
|
||||
|
||||
class TestModelRouter(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.router = ModelRouter(["gpt-6.1-sol", "gpt-6-sol"])
|
||||
|
||||
def test_empty_priority_list_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
ModelRouter([])
|
||||
|
||||
def test_first(self):
|
||||
self.assertEqual(self.router.first(), "gpt-6.1-sol")
|
||||
|
||||
def test_next_after_in_priority_order(self):
|
||||
self.assertEqual(
|
||||
self.router.next_after("gpt-6.1-sol", 404), "gpt-6-sol")
|
||||
self.assertEqual(
|
||||
self.router.next_after("gpt-6.1-sol", "timeout"), "gpt-6-sol")
|
||||
|
||||
def test_next_after_exhausted_returns_none(self):
|
||||
self.assertIsNone(self.router.next_after("gpt-6-sol", 500))
|
||||
self.assertIsNone(self.router.next_after("gpt-6-sol", 429))
|
||||
|
||||
def test_next_after_non_fallbackable_returns_none(self):
|
||||
self.assertIsNone(self.router.next_after("gpt-6.1-sol", 401))
|
||||
self.assertIsNone(self.router.next_after("gpt-6.1-sol", 400))
|
||||
|
||||
def test_route_walks_attempts(self):
|
||||
self.assertEqual(
|
||||
self.router.route([("gpt-6.1-sol", 404)]), "gpt-6-sol")
|
||||
self.assertIsNone(self.router.route(
|
||||
[("gpt-6.1-sol", 404), ("gpt-6-sol", 500)]))
|
||||
self.assertIsNone(self.router.route([("gpt-6.1-sol", 401)]))
|
||||
self.assertEqual(self.router.route([]), "gpt-6.1-sol")
|
||||
|
||||
|
||||
class TestRecordFallback(unittest.TestCase):
|
||||
def test_event_recorded_on_task_record(self):
|
||||
task = {"id": "task-1"}
|
||||
event = record_fallback(task, "gpt-6.1-sol", "gpt-6-sol", "429")
|
||||
self.assertEqual(event, {
|
||||
"type": "model_fallback",
|
||||
"task_id": "task-1",
|
||||
"from_model": "gpt-6.1-sol",
|
||||
"to_model": "gpt-6-sol",
|
||||
"reason": "429",
|
||||
})
|
||||
self.assertIn(event, task["events"])
|
||||
|
||||
def test_event_appended_to_existing_events(self):
|
||||
task = {"id": "task-2", "events": [{"type": "started"}]}
|
||||
record_fallback(task, "a", "b", "500")
|
||||
self.assertEqual(len(task["events"]), 2)
|
||||
self.assertEqual(task["events"][1]["type"], "model_fallback")
|
||||
|
||||
|
||||
class TestFetchAccountUsage(unittest.TestCase):
|
||||
def test_success_returns_available_true(self):
|
||||
data = {"balance": 123, "currency": "USD"}
|
||||
result = fetch_account_usage(lambda: data)
|
||||
self.assertEqual(result, {"available": True, "data": data})
|
||||
|
||||
def test_raising_client_returns_unavailable_not_raise(self):
|
||||
def boom():
|
||||
raise RuntimeError("no route")
|
||||
self.assertEqual(
|
||||
fetch_account_usage(boom), {"available": False, "data": None})
|
||||
|
||||
def test_none_and_empty_dict_return_unavailable(self):
|
||||
self.assertEqual(
|
||||
fetch_account_usage(lambda: None),
|
||||
{"available": False, "data": None})
|
||||
self.assertEqual(
|
||||
fetch_account_usage(lambda: {}),
|
||||
{"available": False, "data": None})
|
||||
|
||||
|
||||
class TestEstimateTokens(unittest.TestCase):
|
||||
def test_chars_divided_by_four(self):
|
||||
self.assertEqual(estimate_tokens("abcd", "efgh"), 2)
|
||||
self.assertEqual(estimate_tokens("a" * 10, "b" * 7), 4) # 17 // 4
|
||||
self.assertEqual(estimate_tokens("", ""), 0)
|
||||
|
||||
|
||||
class TestCostStore(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.store = CostStore(":memory:")
|
||||
|
||||
def test_record_computes_total_and_day(self):
|
||||
row_id = self.store.record(
|
||||
"t1", "p1", "gpt-6.1-sol", input_tokens=100, output_tokens=50,
|
||||
usd_micros=2000, source="actual", fallback_count=1)
|
||||
self.assertIsInstance(row_id, int)
|
||||
report = self.store.report_by_task("t1")
|
||||
self.assertEqual(report["total_tokens"], 150)
|
||||
self.assertEqual(report["usd_micros"], 2000)
|
||||
self.assertEqual(report["turns"], 1)
|
||||
self.assertEqual(report["fallback_count"], 1)
|
||||
self.assertEqual(report["actual_tokens"], 150)
|
||||
self.assertEqual(report["estimated_tokens"], 0)
|
||||
|
||||
def test_invalid_source_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.record("t1", "p1", "m", source="guess")
|
||||
|
||||
def test_actual_and_estimated_kept_separate(self):
|
||||
self.store.record("t1", "p1", "m", input_tokens=100,
|
||||
output_tokens=0, source="actual")
|
||||
self.store.record_estimated("t1", "p1", "m",
|
||||
"x" * 40, "y" * 40) # 80 // 4 = 20
|
||||
report = self.store.report_by_task("t1")
|
||||
self.assertEqual(report["total_tokens"], 120)
|
||||
self.assertEqual(report["actual_tokens"], 100)
|
||||
self.assertEqual(report["estimated_tokens"], 20)
|
||||
self.assertEqual(report["turns"], 2)
|
||||
|
||||
def test_record_estimated_shape(self):
|
||||
self.store.record_estimated("t9", "p9", "m", "abcd", "efgh")
|
||||
report = self.store.report_by_task("t9")
|
||||
self.assertEqual(report["estimated_tokens"], 2)
|
||||
self.assertEqual(report["actual_tokens"], 0)
|
||||
self.assertEqual(report["usd_micros"], 0)
|
||||
|
||||
def test_report_by_project_and_day(self):
|
||||
self.store.record("t1", "p1", "m", input_tokens=10, output_tokens=0)
|
||||
self.store.record("t2", "p1", "m", input_tokens=30, output_tokens=0)
|
||||
self.store.record("t3", "p2", "m", input_tokens=999, output_tokens=0)
|
||||
project = self.store.report_by_project("p1")
|
||||
self.assertEqual(project["total_tokens"], 40)
|
||||
self.assertEqual(project["turns"], 2)
|
||||
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||||
day = self.store.report_by_day(today)
|
||||
self.assertEqual(day["total_tokens"], 10 + 30 + 999)
|
||||
self.assertEqual(day["turns"], 3)
|
||||
self.assertEqual(self.store.report_by_day("2000-01-01")["turns"], 0)
|
||||
|
||||
def test_empty_report_is_zeroed(self):
|
||||
report = self.store.report_by_task("missing")
|
||||
self.assertEqual(report, {
|
||||
"total_tokens": 0,
|
||||
"usd_micros": 0,
|
||||
"turns": 0,
|
||||
"fallback_count": 0,
|
||||
"actual_tokens": 0,
|
||||
"estimated_tokens": 0,
|
||||
})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,206 @@
|
||||
"""派发队列与沙箱执行器的单测。
|
||||
|
||||
运行(在 ~/workspace/codex-app-server 下):
|
||||
python3 -m unittest discover -s cloud/tests -p "test_queue_sandbox.py" -v
|
||||
|
||||
说明:cloud/ 没有 __init__.py(协调员统一创建),因此这里把 cloud/
|
||||
目录插到 sys.path 前面再导入 queue / sandbox,与 discover 的
|
||||
top-level 目录无关。
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(
|
||||
os.path.abspath(__file__))))
|
||||
|
||||
from queue import TaskQueue, budget_exceeded, order_tasks, ready_tasks
|
||||
from sandbox import CloudExecutor, DockerExecutor, Executor, LocalExecutor
|
||||
|
||||
|
||||
def _task(task_id, depends_on=None, priority=0, status="queued"):
|
||||
task = {"id": task_id, "priority": priority, "status": status}
|
||||
if depends_on is not None:
|
||||
task["depends_on"] = depends_on
|
||||
return task
|
||||
|
||||
|
||||
class TestOrderTasks(unittest.TestCase):
|
||||
def test_dependency_order(self):
|
||||
# A -> B 时 B 在 A 后
|
||||
ordered = order_tasks([_task("B", depends_on=["A"]),
|
||||
_task("A")])
|
||||
self.assertEqual([t["id"] for t in ordered], ["A", "B"])
|
||||
|
||||
def test_dependency_chain(self):
|
||||
ordered = order_tasks([_task("C", depends_on=["B"]),
|
||||
_task("B", depends_on=["A"]),
|
||||
_task("A")])
|
||||
self.assertEqual([t["id"] for t in ordered], ["A", "B", "C"])
|
||||
|
||||
def test_priority_within_wave(self):
|
||||
# 同一波(均无依赖)按 priority 大的先
|
||||
ordered = order_tasks([_task("low", priority=1),
|
||||
_task("high", priority=10),
|
||||
_task("mid", priority=5)])
|
||||
self.assertEqual([t["id"] for t in ordered],
|
||||
["high", "mid", "low"])
|
||||
|
||||
def test_priority_does_not_jump_dependency(self):
|
||||
# 高优先级也不能越过依赖:B 依赖 A,即使 B 优先级更高也排 A 后
|
||||
ordered = order_tasks([_task("B", depends_on=["A"], priority=99),
|
||||
_task("A", priority=0)])
|
||||
self.assertEqual([t["id"] for t in ordered], ["A", "B"])
|
||||
|
||||
def test_missing_dependency_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
order_tasks([_task("B", depends_on=["NOPE"])])
|
||||
|
||||
def test_cycle_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
order_tasks([_task("A", depends_on=["B"]),
|
||||
_task("B", depends_on=["A"])])
|
||||
|
||||
def test_self_dependency_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
order_tasks([_task("A", depends_on=["A"])])
|
||||
|
||||
def test_defaults(self):
|
||||
# depends_on / priority 缺省时按 0 处理、不报错
|
||||
ordered = order_tasks([{"id": "A"}, {"id": "B"}])
|
||||
self.assertEqual([t["id"] for t in ordered], ["A", "B"])
|
||||
|
||||
|
||||
class TestReadyTasks(unittest.TestCase):
|
||||
def test_only_dependency_satisfied_queued(self):
|
||||
tasks = [_task("A"),
|
||||
_task("B", depends_on=["A"]),
|
||||
_task("C", depends_on=["A"], status="running"),
|
||||
_task("D", depends_on=["MISSING"])]
|
||||
ready = ready_tasks(tasks, completed_ids=[])
|
||||
self.assertEqual([t["id"] for t in ready], ["A"])
|
||||
|
||||
# A 已完成(状态随之改为 completed)后,B 依赖满足且 queued;
|
||||
# C 是 running 不返回;D 依赖未满足不返回
|
||||
done_tasks = [_task("A", status="completed"),
|
||||
_task("B", depends_on=["A"]),
|
||||
_task("C", depends_on=["A"], status="running"),
|
||||
_task("D", depends_on=["MISSING"])]
|
||||
ready = ready_tasks(done_tasks, completed_ids=["A"])
|
||||
self.assertEqual([t["id"] for t in ready], ["B"])
|
||||
|
||||
def test_ready_ordered_by_priority(self):
|
||||
tasks = [_task("low", priority=1), _task("high", priority=9)]
|
||||
ready = ready_tasks(tasks, completed_ids=[])
|
||||
self.assertEqual([t["id"] for t in ready], ["high", "low"])
|
||||
|
||||
|
||||
class TestBudgetExceeded(unittest.TestCase):
|
||||
def test_below_limits(self):
|
||||
self.assertFalse(budget_exceeded(5, 10, 0.5, 1.0))
|
||||
|
||||
def test_at_limit_counts_as_exceeded(self):
|
||||
self.assertTrue(budget_exceeded(10, 10, 0.5, 1.0))
|
||||
self.assertTrue(budget_exceeded(5, 10, 1.0, 1.0))
|
||||
|
||||
def test_none_limit_means_unlimited(self):
|
||||
self.assertFalse(budget_exceeded(10 ** 9, None, 10 ** 9, None))
|
||||
self.assertTrue(budget_exceeded(10 ** 9, None, 10 ** 9, 1.0))
|
||||
|
||||
|
||||
class TestTaskQueue(unittest.TestCase):
|
||||
def test_next_ready_flow(self):
|
||||
queue = TaskQueue()
|
||||
queue.add(_task("A"))
|
||||
queue.add(_task("B", depends_on=["A"]))
|
||||
self.assertEqual(queue.next_ready([])["id"], "A")
|
||||
queue.mark_running("A")
|
||||
# A 在 running,不再是可派任务;B 依赖未完成,也不可派
|
||||
self.assertIsNone(queue.next_ready([]))
|
||||
queue.mark_done("A", "completed")
|
||||
self.assertEqual(queue.next_ready(["A"])["id"], "B")
|
||||
|
||||
def test_mark_done_validates_status(self):
|
||||
queue = TaskQueue([_task("A")])
|
||||
with self.assertRaises(ValueError):
|
||||
queue.mark_done("A", "bogus-status")
|
||||
with self.assertRaises(ValueError):
|
||||
queue.add(_task("A")) # 重复 id
|
||||
|
||||
def test_all_host_statuses_accepted(self):
|
||||
queue = TaskQueue()
|
||||
for status in ("running", "completed", "failed",
|
||||
"interrupted", "merged", "merge-failed"):
|
||||
queue.add(_task("t-" + status, status=status))
|
||||
# 非 queued 任务一律不可派
|
||||
self.assertIsNone(queue.next_ready([]))
|
||||
|
||||
|
||||
class TestExecutors(unittest.TestCase):
|
||||
def test_base_executor_must_implement(self):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
Executor().execute({"id": "x"})
|
||||
|
||||
def test_local_executor_success(self):
|
||||
seen = []
|
||||
|
||||
def fake_runner(cmd):
|
||||
seen.append(cmd)
|
||||
return SimpleNamespace(returncode=0, stdout="ok-status")
|
||||
|
||||
executor = LocalExecutor("/fake/host.py", config_path="/fake/cfg.json",
|
||||
runner=fake_runner)
|
||||
result = executor.execute({"id": "t1", "project_id": "p",
|
||||
"repo": "/r", "prompt": "hi"})
|
||||
self.assertEqual(result["executor"], "local")
|
||||
self.assertEqual(result["task_id"], "t1")
|
||||
self.assertEqual(result["status"], "completed")
|
||||
self.assertEqual(result["stdout"], "ok-status")
|
||||
# 只发只读 status 命令,不发 dispatch/run(不直接跑模型)
|
||||
self.assertIn("status", seen[0])
|
||||
self.assertNotIn("dispatch", seen[0])
|
||||
self.assertNotIn("run", seen[0])
|
||||
self.assertIn("--config", seen[0])
|
||||
|
||||
def test_local_executor_failure(self):
|
||||
def fake_runner(cmd):
|
||||
return SimpleNamespace(returncode=1, stdout="boom")
|
||||
|
||||
result = LocalExecutor("/fake/host.py",
|
||||
runner=fake_runner).execute({"id": "t2"})
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertEqual(result["stdout"], "boom")
|
||||
|
||||
def test_local_executor_truncates_stdout(self):
|
||||
def fake_runner(cmd):
|
||||
return SimpleNamespace(returncode=0, stdout="x" * 5000)
|
||||
|
||||
result = LocalExecutor("/fake/host.py",
|
||||
runner=fake_runner).execute({"id": "t3"})
|
||||
self.assertEqual(len(result["stdout"]), 2000)
|
||||
|
||||
def test_docker_executor_no_docker_raises_not_local(self):
|
||||
with mock.patch("sandbox.shutil.which", return_value=None):
|
||||
with self.assertRaises(RuntimeError) as ctx:
|
||||
DockerExecutor().execute({"id": "t1"})
|
||||
self.assertIn("docker not available", str(ctx.exception))
|
||||
# 报错而不是返回 local 结果:绝不静默降级为本地执行
|
||||
|
||||
def test_docker_executor_with_docker_is_placeholder_failure(self):
|
||||
with mock.patch("sandbox.shutil.which", return_value="/usr/bin/docker"):
|
||||
result = DockerExecutor().execute({"id": "t1"})
|
||||
self.assertEqual(result["executor"], "docker")
|
||||
self.assertEqual(result["status"], "failed")
|
||||
self.assertNotEqual(result["executor"], "local")
|
||||
|
||||
def test_cloud_executor_raises_with_cloud_account(self):
|
||||
with self.assertRaises(RuntimeError) as ctx:
|
||||
CloudExecutor().execute({"id": "t1"})
|
||||
self.assertIn("cloud account", str(ctx.exception))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,311 @@
|
||||
"""控制面 server 集成测试:对本地临时端口做真实 HTTP 请求。
|
||||
|
||||
只用标准库(http.client),不跑真实模型回合,不联网连真实
|
||||
GitHub/Slack(webhook 测试用的 secret 均为单测内临时值)。
|
||||
每个测试用独立的临时 SQLite 文件与临时端口,跑完关闭 server。
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
try:
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.server import create_server, sweep_stale_tasks
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.server import create_server, sweep_stale_tasks
|
||||
|
||||
|
||||
class ServerTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
self.auth = AuthStore(self.db_path)
|
||||
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
||||
self.viewer_id, self.viewer_key = self.auth.create_user("viewer")
|
||||
self.outsider_id, self.outsider_key = self.auth.create_user("outsider")
|
||||
self.server = create_server(self.db_path, "127.0.0.1", 0)
|
||||
self.port = self.server.server_address[1]
|
||||
self.thread = threading.Thread(
|
||||
target=self.server.serve_forever, daemon=True)
|
||||
self.thread.start()
|
||||
# owner 建一个项目,viewer 设为该项目的 viewer 角色。
|
||||
status, self.project = self._req(
|
||||
"POST", "/projects", {"name": "demo"}, self.owner_key)
|
||||
self.assertEqual(status, 201)
|
||||
self.auth.set_member(
|
||||
self.project["id"], self.owner_id, self.viewer_id, "viewer")
|
||||
|
||||
def tearDown(self):
|
||||
self.server.shutdown()
|
||||
self.server.server_close()
|
||||
self.thread.join(timeout=5)
|
||||
try:
|
||||
self.auth.conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
self.tmp.cleanup()
|
||||
|
||||
# -- HTTP 工具 -----------------------------------------------------
|
||||
def _req(self, method, path, body=None, key=None, headers=None,
|
||||
raw_body=None):
|
||||
conn = http.client.HTTPConnection(
|
||||
"127.0.0.1", self.port, timeout=10)
|
||||
hdrs = dict(headers or {})
|
||||
if raw_body is not None:
|
||||
payload = raw_body
|
||||
elif body is not None:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
hdrs.setdefault("Content-Type", "application/json")
|
||||
else:
|
||||
payload = None
|
||||
if key is not None:
|
||||
hdrs["X-API-Key"] = key
|
||||
try:
|
||||
conn.request(method, path, body=payload, headers=hdrs)
|
||||
resp = conn.getresponse()
|
||||
data = resp.read()
|
||||
finally:
|
||||
conn.close()
|
||||
try:
|
||||
parsed = json.loads(data.decode("utf-8")) if data else None
|
||||
except ValueError:
|
||||
parsed = None
|
||||
return resp.status, parsed
|
||||
|
||||
def _create_task(self, key, **extra):
|
||||
body = {"project_id": self.project["id"], "prompt": "做一个加法函数"}
|
||||
body.update(extra)
|
||||
status, task = self._req("POST", "/tasks", body, key)
|
||||
self.assertEqual(status, 201, task)
|
||||
return task
|
||||
|
||||
# -- 健康检查与认证 -------------------------------------------------
|
||||
def test_health_no_auth(self):
|
||||
status, body = self._req("GET", "/health")
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(
|
||||
body, {"ok": True, "service": "cloud-control-plane"})
|
||||
|
||||
def test_no_key_401(self):
|
||||
status, body = self._req("GET", "/projects")
|
||||
self.assertEqual(status, 401)
|
||||
self.assertEqual(body, {"error": "unauthorized"})
|
||||
|
||||
def test_wrong_key_401(self):
|
||||
status, body = self._req("GET", "/projects", key="wrong-key")
|
||||
self.assertEqual(status, 401)
|
||||
self.assertEqual(body, {"error": "unauthorized"})
|
||||
|
||||
# -- 项目可见性与权限 -----------------------------------------------
|
||||
def test_projects_only_member_projects(self):
|
||||
# outsider 建自己的项目;双方互不可见对方的项目。
|
||||
status, other = self._req(
|
||||
"POST", "/projects", {"name": "other"}, self.outsider_key)
|
||||
self.assertEqual(status, 201)
|
||||
status, body = self._req("GET", "/projects", key=self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
ids = [p["id"] for p in body["projects"]]
|
||||
self.assertIn(self.project["id"], ids)
|
||||
self.assertNotIn(other["id"], ids)
|
||||
status, body = self._req("GET", "/projects", key=self.outsider_key)
|
||||
self.assertEqual(status, 200)
|
||||
ids = [p["id"] for p in body["projects"]]
|
||||
self.assertIn(other["id"], ids)
|
||||
self.assertNotIn(self.project["id"], ids)
|
||||
|
||||
def test_nonmember_read_project_tasks_403(self):
|
||||
status, body = self._req(
|
||||
"GET", "/tasks?project_id=%s" % self.project["id"],
|
||||
key=self.outsider_key)
|
||||
self.assertEqual(status, 403)
|
||||
|
||||
def test_viewer_create_task_403(self):
|
||||
status, body = self._req(
|
||||
"POST", "/tasks",
|
||||
{"project_id": self.project["id"], "prompt": "x"},
|
||||
self.viewer_key)
|
||||
self.assertEqual(status, 403)
|
||||
|
||||
def test_viewer_can_read_tasks(self):
|
||||
self._create_task(self.owner_key)
|
||||
status, body = self._req(
|
||||
"GET", "/tasks?project_id=%s" % self.project["id"],
|
||||
key=self.viewer_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(len(body["tasks"]), 1)
|
||||
|
||||
# -- 任务 CRUD、幂等、取消、心跳、回收 -------------------------------
|
||||
def test_create_and_get_task(self):
|
||||
task = self._create_task(self.owner_key, priority=3)
|
||||
self.assertEqual(task["status"], "queued")
|
||||
self.assertEqual(task["priority"], 3)
|
||||
status, got = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(got["id"], task["id"])
|
||||
status, _ = self._req("GET", "/tasks/nope", key=self.owner_key)
|
||||
self.assertEqual(status, 404)
|
||||
status, _ = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.outsider_key)
|
||||
self.assertEqual(status, 403)
|
||||
|
||||
def test_idempotent_key_returns_same_task(self):
|
||||
first = self._create_task(
|
||||
self.owner_key, idempotency_key="key-1")
|
||||
status, second = self._req(
|
||||
"POST", "/tasks",
|
||||
{"project_id": self.project["id"], "prompt": "另一个 prompt",
|
||||
"idempotency_key": "key-1"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(second["id"], first["id"])
|
||||
self.assertEqual(second["prompt"], first["prompt"])
|
||||
status, body = self._req(
|
||||
"GET", "/tasks?project_id=%s" % self.project["id"],
|
||||
key=self.owner_key)
|
||||
self.assertEqual(len(body["tasks"]), 1)
|
||||
|
||||
def test_cancel_running_then_terminal_cancel_400(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
status, running = self._req(
|
||||
"POST", "/tasks/%s/heartbeat" % task["id"], {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(running["status"], "running")
|
||||
status, cancelled = self._req(
|
||||
"POST", "/tasks/%s/cancel" % task["id"], {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(cancelled["status"], "interrupted")
|
||||
status, body = self._req(
|
||||
"POST", "/tasks/%s/cancel" % task["id"], {}, self.owner_key)
|
||||
self.assertEqual(status, 400)
|
||||
|
||||
def test_cancel_queued_task(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
status, body = self._req(
|
||||
"POST", "/tasks/%s/cancel" % task["id"], {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(body["status"], "interrupted")
|
||||
|
||||
def test_heartbeat_after_sweep_keeps_fresh_task(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
self._req("POST", "/tasks/%s/heartbeat" % task["id"], {},
|
||||
self.owner_key)
|
||||
status, body = self._req(
|
||||
"POST", "/admin/sweep", {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertNotIn(task["id"], body["swept"])
|
||||
status, got = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.owner_key)
|
||||
self.assertEqual(got["status"], "running")
|
||||
|
||||
def test_sweep_stale_with_threshold_zero(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
self._req("POST", "/tasks/%s/heartbeat" % task["id"], {},
|
||||
self.owner_key)
|
||||
swept = sweep_stale_tasks(self.db_path, threshold_seconds=0)
|
||||
self.assertIn(task["id"], swept)
|
||||
status, got = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.owner_key)
|
||||
self.assertEqual(got["status"], "interrupted")
|
||||
|
||||
def test_sweep_stale_manually_aged_heartbeat(self):
|
||||
task = self._create_task(self.owner_key)
|
||||
self._req("POST", "/tasks/%s/heartbeat" % task["id"], {},
|
||||
self.owner_key)
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
conn.execute(
|
||||
"UPDATE tasks SET last_heartbeat = ? WHERE id = ?",
|
||||
(time.time() - 3600, task["id"]))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
status, body = self._req(
|
||||
"POST", "/admin/sweep", {}, self.owner_key)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertIn(task["id"], body["swept"])
|
||||
status, got = self._req(
|
||||
"GET", "/tasks/%s" % task["id"], key=self.owner_key)
|
||||
self.assertEqual(got["status"], "interrupted")
|
||||
|
||||
def test_sweep_skips_queued_task(self):
|
||||
task = self._create_task(self.owner_key) # 仍是 queued
|
||||
swept = sweep_stale_tasks(self.db_path, threshold_seconds=0)
|
||||
self.assertNotIn(task["id"], swept)
|
||||
|
||||
# -- webhooks(验签,解析但不执行)----------------------------------
|
||||
def test_github_webhook_secret_missing_503(self):
|
||||
with mock.patch.dict(os.environ, {}, clear=False):
|
||||
os.environ.pop("GITHUB_WEBHOOK_SECRET", None)
|
||||
status, _ = self._req(
|
||||
"POST", "/webhooks/github", {"action": "created"},
|
||||
key=None)
|
||||
self.assertEqual(status, 503)
|
||||
|
||||
def test_github_webhook_bad_signature_401(self):
|
||||
payload = {"action": "created", "comment": {"body": "/fix-review"}}
|
||||
raw = json.dumps(payload).encode("utf-8")
|
||||
with mock.patch.dict(
|
||||
os.environ, {"GITHUB_WEBHOOK_SECRET": "test-secret"}):
|
||||
status, body = self._req(
|
||||
"POST", "/webhooks/github", raw_body=raw,
|
||||
headers={"Content-Type": "application/json",
|
||||
"X-Hub-Signature-256": "sha256=deadbeef"})
|
||||
self.assertEqual(status, 401)
|
||||
|
||||
def test_github_webhook_good_signature_parses(self):
|
||||
payload = {
|
||||
"action": "created",
|
||||
"comment": {"body": "请 /fix-review 处理", "user": {"login": "alice"}},
|
||||
"repository": {"full_name": "org/repo"},
|
||||
"issue": {"number": 7},
|
||||
}
|
||||
raw = json.dumps(payload).encode("utf-8")
|
||||
sig = "sha256=" + hmac.new(
|
||||
b"test-secret", raw, hashlib.sha256).hexdigest()
|
||||
with mock.patch.dict(
|
||||
os.environ, {"GITHUB_WEBHOOK_SECRET": "test-secret"}):
|
||||
status, body = self._req(
|
||||
"POST", "/webhooks/github", raw_body=raw,
|
||||
headers={"Content-Type": "application/json",
|
||||
"X-Hub-Signature-256": sig})
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(body["type"], "fix_review")
|
||||
self.assertEqual(body["pr_number"], 7)
|
||||
|
||||
def test_slack_webhook_bad_and_good_signature(self):
|
||||
raw = b"text=/codex status my-task"
|
||||
ts = str(int(time.time()))
|
||||
good_sig = "v0=" + hmac.new(
|
||||
b"slack-secret",
|
||||
b"v0:" + ts.encode() + b":" + raw, hashlib.sha256).hexdigest()
|
||||
headers = {"Content-Type": "application/x-www-form-urlencoded",
|
||||
"X-Slack-Signature": "v0=deadbeef",
|
||||
"X-Slack-Request-Timestamp": ts}
|
||||
with mock.patch.dict(
|
||||
os.environ, {"SLACK_SIGNING_SECRET": "slack-secret"}):
|
||||
status, _ = self._req(
|
||||
"POST", "/webhooks/slack", raw_body=raw, headers=headers)
|
||||
self.assertEqual(status, 401)
|
||||
headers["X-Slack-Signature"] = good_sig
|
||||
status, body = self._req(
|
||||
"POST", "/webhooks/slack", raw_body=raw, headers=headers)
|
||||
self.assertEqual(status, 200)
|
||||
self.assertEqual(body, {"command": "status", "task": "my-task"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,236 @@
|
||||
"""webhooks.py 单测:签名校验、事件解析、命令派发、用户映射、通知门禁。
|
||||
|
||||
只用标准库,不联网、不连真实 GitHub/Slack、不跑真实模型回合。
|
||||
签名测试用的 secret 均为单测内临时值,不是生产密钥。
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
|
||||
# 本文件位于 cloud/tests/,被测模块位于 cloud/;cloud/ 目前不是包
|
||||
# (协调员统一创建 __init__.py),故把 cloud/ 加到 sys.path 后按
|
||||
# 顶层模块导入 webhooks。
|
||||
sys.path.insert(
|
||||
0, os.path.normpath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
)
|
||||
|
||||
import webhooks # noqa: E402
|
||||
from webhooks import ( # noqa: E402
|
||||
ALLOWED_NOTIFICATION_EVENTS,
|
||||
NotificationService,
|
||||
SlackUserMap,
|
||||
parse_github_pr_comment,
|
||||
parse_slack_command,
|
||||
verify_github_signature,
|
||||
verify_slack_signature,
|
||||
)
|
||||
|
||||
|
||||
def _github_signature(secret, payload):
|
||||
return "sha256=" + hmac.new(
|
||||
secret.encode("utf-8"), payload, hashlib.sha256
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def _slack_signature(secret, timestamp, body):
|
||||
basestring = f"v0:{timestamp}:{body}"
|
||||
return "v0=" + hmac.new(
|
||||
secret.encode("utf-8"), basestring.encode("utf-8"), hashlib.sha256
|
||||
).hexdigest()
|
||||
|
||||
|
||||
class TestGithubSignature(unittest.TestCase):
|
||||
def test_correct_signature_passes(self):
|
||||
secret = "test-github-secret"
|
||||
payload = b'{"action":"created"}'
|
||||
header = _github_signature(secret, payload)
|
||||
self.assertTrue(verify_github_signature(secret, payload, header))
|
||||
|
||||
def test_wrong_signature_rejected(self):
|
||||
secret = "test-github-secret"
|
||||
payload = b'{"action":"created"}'
|
||||
header = _github_signature("wrong-secret", payload)
|
||||
self.assertFalse(verify_github_signature(secret, payload, header))
|
||||
|
||||
def test_tampered_payload_rejected(self):
|
||||
secret = "test-github-secret"
|
||||
header = _github_signature(secret, b'{"action":"created"}')
|
||||
self.assertFalse(
|
||||
verify_github_signature(secret, b'{"action":"deleted"}', header)
|
||||
)
|
||||
|
||||
def test_missing_header_rejected(self):
|
||||
self.assertFalse(
|
||||
verify_github_signature("test-github-secret", b"{}", None)
|
||||
)
|
||||
self.assertFalse(verify_github_signature("test-github-secret", b"{}", ""))
|
||||
|
||||
def test_missing_secret_rejected(self):
|
||||
payload = b"{}"
|
||||
self.assertFalse(verify_github_signature("", payload, "sha256=abc"))
|
||||
self.assertFalse(verify_github_signature(None, payload, "sha256=abc"))
|
||||
|
||||
|
||||
class TestParseGithubPrComment(unittest.TestCase):
|
||||
def _payload(self, body, action="created"):
|
||||
return {
|
||||
"action": action,
|
||||
"repository": {"full_name": "owner/repo"},
|
||||
"issue": {"number": 42},
|
||||
"comment": {"body": body, "user": {"login": "alice"}},
|
||||
}
|
||||
|
||||
def test_fix_review_comment_parsed(self):
|
||||
result = parse_github_pr_comment(
|
||||
self._payload("请帮忙 /fix-review 看一下")
|
||||
)
|
||||
self.assertEqual(
|
||||
result,
|
||||
{
|
||||
"type": "fix_review",
|
||||
"repo": "owner/repo",
|
||||
"pr_number": 42,
|
||||
"comment": "请帮忙 /fix-review 看一下",
|
||||
"author": "alice",
|
||||
},
|
||||
)
|
||||
|
||||
def test_comment_without_trigger_returns_none(self):
|
||||
self.assertIsNone(parse_github_pr_comment(self._payload("好看,LGTM")))
|
||||
|
||||
def test_non_created_action_returns_none(self):
|
||||
self.assertIsNone(
|
||||
parse_github_pr_comment(
|
||||
self._payload("有 /fix-review 但是已编辑", action="edited")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class TestSlackSignature(unittest.TestCase):
|
||||
def test_correct_signature_passes(self):
|
||||
secret = "test-slack-secret"
|
||||
timestamp = str(int(time.time()))
|
||||
body = "command=/codex&text=status"
|
||||
header = _slack_signature(secret, timestamp, body)
|
||||
self.assertTrue(
|
||||
verify_slack_signature(secret, timestamp, body, header)
|
||||
)
|
||||
|
||||
def test_wrong_signature_rejected(self):
|
||||
secret = "test-slack-secret"
|
||||
timestamp = str(int(time.time()))
|
||||
body = "command=/codex&text=status"
|
||||
header = _slack_signature("wrong-secret", timestamp, body)
|
||||
self.assertFalse(
|
||||
verify_slack_signature(secret, timestamp, body, header)
|
||||
)
|
||||
|
||||
def test_expired_timestamp_rejected(self):
|
||||
# 超过 300 秒:用正确密钥签一个 301 秒前的时间戳,仍应拒绝。
|
||||
secret = "test-slack-secret"
|
||||
timestamp = str(int(time.time()) - 301)
|
||||
body = "command=/codex&text=status"
|
||||
header = _slack_signature(secret, timestamp, body)
|
||||
self.assertFalse(
|
||||
verify_slack_signature(secret, timestamp, body, header)
|
||||
)
|
||||
|
||||
def test_missing_params_rejected(self):
|
||||
timestamp = str(int(time.time()))
|
||||
self.assertFalse(
|
||||
verify_slack_signature("", timestamp, "body", "v0=abc")
|
||||
)
|
||||
self.assertFalse(
|
||||
verify_slack_signature("test-slack-secret", timestamp, "body", None)
|
||||
)
|
||||
|
||||
|
||||
class TestParseSlackCommand(unittest.TestCase):
|
||||
def test_dispatch_prompt(self):
|
||||
self.assertEqual(
|
||||
parse_slack_command("/codex 做一个加法函数"),
|
||||
{"command": "codex", "prompt": "做一个加法函数"},
|
||||
)
|
||||
|
||||
def test_dispatch_status(self):
|
||||
self.assertEqual(
|
||||
parse_slack_command("/codex status my-task"),
|
||||
{"command": "status", "task": "my-task"},
|
||||
)
|
||||
|
||||
def test_empty_or_non_codex_returns_none_command(self):
|
||||
self.assertEqual(
|
||||
parse_slack_command(""), {"command": None, "prompt": ""}
|
||||
)
|
||||
self.assertEqual(
|
||||
parse_slack_command("hello"),
|
||||
{"command": None, "prompt": "hello"},
|
||||
)
|
||||
self.assertEqual(
|
||||
parse_slack_command("/other do something"),
|
||||
{"command": None, "prompt": "/other do something"},
|
||||
)
|
||||
|
||||
|
||||
class TestSlackUserMap(unittest.TestCase):
|
||||
def test_map_and_lookup(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
db_path = os.path.join(tmpdir, "slack_users.db")
|
||||
user_map = SlackUserMap(db_path)
|
||||
try:
|
||||
self.assertIsNone(user_map.lookup("U123"))
|
||||
user_map.map("U123", "user-1")
|
||||
self.assertEqual(user_map.lookup("U123"), "user-1")
|
||||
# 覆盖映射
|
||||
user_map.map("U123", "user-2")
|
||||
self.assertEqual(user_map.lookup("U123"), "user-2")
|
||||
# 未映射的 id 仍返回 None
|
||||
self.assertIsNone(user_map.lookup("U999"))
|
||||
finally:
|
||||
user_map.close()
|
||||
|
||||
|
||||
class TestNotificationService(unittest.TestCase):
|
||||
def test_allowed_event_sent_and_senders_called(self):
|
||||
calls = {"webhook": [], "slack": []}
|
||||
service = NotificationService(
|
||||
webhook_sender=lambda payload: calls["webhook"].append(payload),
|
||||
slack_sender=lambda payload: calls["slack"].append(payload),
|
||||
)
|
||||
result = service.notify("completed", "task-1", "proj-1", "done")
|
||||
self.assertEqual(result, {"sent": True, "channels": ["webhook", "slack"]})
|
||||
self.assertEqual(len(calls["webhook"]), 1)
|
||||
self.assertEqual(len(calls["slack"]), 1)
|
||||
self.assertEqual(calls["webhook"][0]["event_type"], "completed")
|
||||
self.assertEqual(calls["webhook"][0]["task_id"], "task-1")
|
||||
|
||||
def test_disallowed_event_not_sent_and_senders_not_called(self):
|
||||
calls = []
|
||||
service = NotificationService(
|
||||
webhook_sender=lambda payload: calls.append(("webhook", payload)),
|
||||
slack_sender=lambda payload: calls.append(("slack", payload)),
|
||||
)
|
||||
result = service.notify("running", "task-1", "proj-1")
|
||||
self.assertEqual(result, {"sent": False, "reason": "event not allowed"})
|
||||
self.assertEqual(calls, [])
|
||||
|
||||
def test_no_senders_does_not_error(self):
|
||||
service = NotificationService()
|
||||
result = service.notify("completed", "task-1", "proj-1")
|
||||
self.assertEqual(result, {"sent": True, "channels": []})
|
||||
|
||||
def test_allowed_events_set(self):
|
||||
self.assertEqual(
|
||||
ALLOWED_NOTIFICATION_EVENTS,
|
||||
{"completed", "failed", "goal-unmet", "merge-failed",
|
||||
"review-requested"},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,238 @@
|
||||
"""外部集成 webhook 与推送通知模块(云端架构实施)。
|
||||
|
||||
零第三方依赖,只用 Python 标准库。
|
||||
|
||||
本模块只做两件事:
|
||||
|
||||
1. 入站 webhook 的签名校验与事件解析(GitHub、Slack)。
|
||||
2. 出站状态通知的占位分发(webhook + Slack)。
|
||||
|
||||
边界(不要超出本模块宣称的范围):
|
||||
|
||||
- Web Push / APNs / FCM 明确是后续需要证书的一步,本版只有
|
||||
webhook + Slack 占位,不要编造已接通手机推送。
|
||||
- 密钥走环境变量是调用方的事:本模块的校验函数只接受参数,
|
||||
不自己读生产密钥,也不把任何密钥写入文件。
|
||||
- 本模块不联网、不调用真实 GitHub/Slack、不跑真实模型回合。
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import sqlite3
|
||||
import time
|
||||
|
||||
__all__ = [
|
||||
"verify_github_signature",
|
||||
"parse_github_pr_comment",
|
||||
"verify_slack_signature",
|
||||
"parse_slack_command",
|
||||
"SlackUserMap",
|
||||
"ALLOWED_NOTIFICATION_EVENTS",
|
||||
"NotificationService",
|
||||
]
|
||||
|
||||
|
||||
def verify_github_signature(secret, payload_bytes, signature_header):
|
||||
"""校验 GitHub webhook 签名。
|
||||
|
||||
signature_header 形如 ``"sha256=<hex>"``,用
|
||||
``hmac.new(secret.encode(), payload_bytes, hashlib.sha256).hexdigest()``
|
||||
比对,比较用 ``hmac.compare_digest``。缺 secret 或缺 header 返回 False。
|
||||
"""
|
||||
if not secret or not signature_header:
|
||||
return False
|
||||
if isinstance(payload_bytes, str):
|
||||
payload_bytes = payload_bytes.encode("utf-8")
|
||||
if not isinstance(payload_bytes, (bytes, bytearray)):
|
||||
return False
|
||||
if not isinstance(signature_header, str):
|
||||
return False
|
||||
try:
|
||||
expected = "sha256=" + hmac.new(
|
||||
secret.encode("utf-8"), bytes(payload_bytes), hashlib.sha256
|
||||
).hexdigest()
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
return False
|
||||
return hmac.compare_digest(expected, signature_header)
|
||||
|
||||
|
||||
def parse_github_pr_comment(payload):
|
||||
"""解析 GitHub issue_comment 事件,返回 fix-review 触发点或 None。
|
||||
|
||||
这是 PR 评论 → fix-review 的触发点,真正执行 fix-review 由控制面调
|
||||
host.py 完成,本函数只解析不执行。
|
||||
|
||||
只有当 action == "created" 且 comment body 里含 "/fix-review" 时返回::
|
||||
|
||||
{"type": "fix_review", "repo": ..., "pr_number": ...,
|
||||
"comment": ..., "author": ...}
|
||||
|
||||
其中 repo 取 payload repository full_name,pr_number 取 issue number,
|
||||
comment 取 body,author 取 comment user login。其他情况返回 None。
|
||||
"""
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
if payload.get("action") != "created":
|
||||
return None
|
||||
comment = payload.get("comment")
|
||||
if not isinstance(comment, dict):
|
||||
return None
|
||||
body = comment.get("body")
|
||||
if not isinstance(body, str) or "/fix-review" not in body:
|
||||
return None
|
||||
repository = payload.get("repository") or {}
|
||||
issue = payload.get("issue") or {}
|
||||
user = comment.get("user") or {}
|
||||
return {
|
||||
"type": "fix_review",
|
||||
"repo": repository.get("full_name"),
|
||||
"pr_number": issue.get("number"),
|
||||
"comment": body,
|
||||
"author": user.get("login"),
|
||||
}
|
||||
|
||||
|
||||
def verify_slack_signature(signing_secret, timestamp, body, signature_header):
|
||||
"""校验 Slack 请求签名(防重放)。
|
||||
|
||||
Slack 规则:basestring = ``f"v0:{timestamp}:{body}"``,hmac sha256,
|
||||
header 形如 ``"v0=<hex>"``。时间戳与当前时间差超过 300 秒返回 False
|
||||
(防重放);缺参返回 False。
|
||||
"""
|
||||
if not signing_secret or timestamp is None or body is None:
|
||||
return False
|
||||
if not signature_header:
|
||||
return False
|
||||
try:
|
||||
ts_int = int(timestamp)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
if abs(time.time() - ts_int) > 300:
|
||||
return False
|
||||
try:
|
||||
if isinstance(body, (bytes, bytearray)):
|
||||
basestring = b"v0:" + str(ts_int).encode("utf-8") + b":" + bytes(body)
|
||||
expected = "v0=" + hmac.new(
|
||||
signing_secret.encode("utf-8"), basestring, hashlib.sha256
|
||||
).hexdigest()
|
||||
else:
|
||||
basestring = f"v0:{timestamp}:{body}"
|
||||
expected = "v0=" + hmac.new(
|
||||
signing_secret.encode("utf-8"),
|
||||
basestring.encode("utf-8"),
|
||||
hashlib.sha256,
|
||||
).hexdigest()
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
return False
|
||||
if not isinstance(signature_header, str):
|
||||
return False
|
||||
return hmac.compare_digest(expected, signature_header)
|
||||
|
||||
|
||||
def parse_slack_command(text):
|
||||
"""解析 Slack 斜杠命令。
|
||||
|
||||
- ``"/codex 做一个加法函数"`` ->
|
||||
``{"command": "codex", "prompt": "做一个加法函数"}``
|
||||
- ``"/codex status <任务名>"`` ->
|
||||
``{"command": "status", "task": "<任务名>"}``
|
||||
- 空或非 ``/codex`` 开头 -> ``{"command": None, "prompt": text}``
|
||||
"""
|
||||
if not text or not isinstance(text, str):
|
||||
return {"command": None, "prompt": text}
|
||||
stripped = text.strip()
|
||||
if stripped == "/codex" or stripped.startswith("/codex "):
|
||||
rest = stripped[len("/codex"):].strip()
|
||||
if rest == "status" or rest.startswith("status "):
|
||||
task = rest[len("status"):].strip()
|
||||
return {"command": "status", "task": task}
|
||||
return {"command": "codex", "prompt": rest}
|
||||
return {"command": None, "prompt": text}
|
||||
|
||||
|
||||
class SlackUserMap:
|
||||
"""Slack 用户 ID 到内部 user_id 的映射(SQLite)。"""
|
||||
|
||||
def __init__(self, db_path):
|
||||
self.db_path = db_path
|
||||
self._conn = sqlite3.connect(db_path)
|
||||
self._conn.execute(
|
||||
"CREATE TABLE IF NOT EXISTS slack_users ("
|
||||
"slack_id TEXT PRIMARY KEY, "
|
||||
"user_id TEXT, "
|
||||
"created_at REAL)"
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def map(self, slack_id, user_id):
|
||||
"""把 slack_id 映射到 user_id(已存在则覆盖)。"""
|
||||
self._conn.execute(
|
||||
"INSERT OR REPLACE INTO slack_users "
|
||||
"(slack_id, user_id, created_at) VALUES (?, ?, ?)",
|
||||
(slack_id, user_id, time.time()),
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def lookup(self, slack_id):
|
||||
"""返回 slack_id 对应的 user_id,未映射返回 None。"""
|
||||
row = self._conn.execute(
|
||||
"SELECT user_id FROM slack_users WHERE slack_id = ?",
|
||||
(slack_id,),
|
||||
).fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
def close(self):
|
||||
"""关闭数据库连接(可选,单测/退出时用)。"""
|
||||
self._conn.close()
|
||||
|
||||
|
||||
# 只发状态变化的事件集合:进行中的状态(如 running)不触发通知。
|
||||
ALLOWED_NOTIFICATION_EVENTS = {
|
||||
"completed",
|
||||
"failed",
|
||||
"goal-unmet",
|
||||
"merge-failed",
|
||||
"review-requested",
|
||||
}
|
||||
|
||||
|
||||
class NotificationService:
|
||||
"""状态通知分发(webhook + Slack 占位)。
|
||||
|
||||
Web Push / APNs / FCM 明确是后续需要证书的一步,本版只有 webhook +
|
||||
Slack 占位,不要编造已接通手机推送。
|
||||
|
||||
sender 都是可调用对象,单测可注入假 sender,默认 None 时不发、
|
||||
不报错。sender 被调用时收到一个 dict payload::
|
||||
|
||||
{"event_type": ..., "task_id": ..., "project_id": ..., "detail": ...}
|
||||
"""
|
||||
|
||||
def __init__(self, webhook_sender=None, slack_sender=None):
|
||||
self.webhook_sender = webhook_sender
|
||||
self.slack_sender = slack_sender
|
||||
|
||||
def notify(self, event_type, task_id, project_id, detail=""):
|
||||
"""发送一条状态通知。
|
||||
|
||||
event_type 不在 ALLOWED_NOTIFICATION_EVENTS 里时返回
|
||||
``{"sent": False, "reason": "event not allowed"}`` 且不调用 sender;
|
||||
否则依次调 webhook_sender 和 slack_sender(不为 None 的才调),
|
||||
返回 ``{"sent": True, "channels": [...已调的通道名...]}``。
|
||||
"""
|
||||
if event_type not in ALLOWED_NOTIFICATION_EVENTS:
|
||||
return {"sent": False, "reason": "event not allowed"}
|
||||
payload = {
|
||||
"event_type": event_type,
|
||||
"task_id": task_id,
|
||||
"project_id": project_id,
|
||||
"detail": detail,
|
||||
}
|
||||
channels = []
|
||||
if self.webhook_sender is not None:
|
||||
self.webhook_sender(payload)
|
||||
channels.append("webhook")
|
||||
if self.slack_sender is not None:
|
||||
self.slack_sender(payload)
|
||||
channels.append("slack")
|
||||
return {"sent": True, "channels": channels}
|
||||
Reference in new issue
Block a user