diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..e67ff070a --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,3 @@ +# CLAUDE Notes + +- 2026-03-12: `refactor.md` on disk was empty, while the active collaboration context contained the full v4 refactor design. Treat the conversation-approved v4 design as source of truth unless a newer committed document replaces it. diff --git a/pyproject.toml b/pyproject.toml index 3c173e398..c8e559f43 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,7 @@ +[build-system] +requires = ["setuptools>=80", "wheel"] +build-backend = "setuptools.build_meta" + [project] name = "astrbot-sdk" version = "0.1.0" @@ -19,3 +23,9 @@ dependencies = [ [project.scripts] astr = "astrbot_sdk.cli:cli" + +[tool.setuptools] +package-dir = {"" = "src-new"} + +[tool.setuptools.packages.find] +where = ["src-new"] diff --git a/refactor.md b/refactor.md new file mode 100644 index 000000000..cc27256fc --- /dev/null +++ b/refactor.md @@ -0,0 +1,1385 @@ +# AstrBot SDK 重构架构设计 v4 + +--- + +## 一、全局架构图 + +``` +╔══════════════════════════════════════════════════════════════════════╗ +║ 插件作者的世界 ║ +║ ║ +║ class MyPlugin(Star): ║ +║ @on_command("hello") ║ +║ async def hello(self, event: MessageEvent, ctx: Context): ║ +║ reply = await ctx.llm.chat(event.text) ║ +║ await event.reply(reply) ║ +║ ║ +╚══════════════════════════════════════════════════════════════════════╝ + │ Star / 装饰器 / Event │ Context / Clients + ▼ ▼ +┌─────────────────────┐ ┌────────────────────────────────┐ +│ Handler 系统 │ │ Capability 调用系统 │ +│ │ │ │ +│ HandlerDescriptor │ │ ctx.llm.chat() │ +│ HandlerDispatcher │ │ ctx.memory.search() │ +│ │ │ ctx.db.get() │ +│ 插件 → 主进程 │ │ ctx.platform.send() │ +│ "我能响应这些事件" │ │ │ +│ │ │ 插件 → 主进程 │ +│ │ │ "帮我调用这个能力" │ +└──────────┬──────────┘ └──────────────┬─────────────────┘ + │ │ + └─────────────────┬────────────────┘ + ▼ +┌──────────────────────────────────────────────────────────────────┐ +│ 通信层 │ +│ │ +│ 所有消息统一使用 id 字段关联请求与响应 │ +│ │ +│ Peer.initialize(handlers=[...]) │ +│ Peer.invoke("llm.chat", input, stream=false) → result │ +│ Peer.invoke("llm.stream_chat", input, stream=true) → event* │ +│ Peer.invoke("handler.invoke", {handler_id, event}) │ +│ │ +│ Transport: StdioTransport / WebSocketTransport │ +└──────────────────────────────────────────────────────────────────┘ + │ JSON 消息流 + ▼ +┌──────────────────────────────────────────────────────────────────┐ +│ 主进程（AstrBot Core） │ +│ │ +│ CapabilityRouter ──► "llm.chat" ──► LLM Service │ +│ ──► "db.get" ──► Storage │ +│ ──► "handler.invoke" ──► 转发给插件 │ +│ │ +│ HandlerDispatcher ◄── 外部消息 ──► 匹配订阅 ──► 回调插件 │ +└──────────────────────────────────────────────────────────────────┘ + + ┌─────────────────────┐ + │ compat.py（旁路） │ ← 不是核心层 + │ 旧 API → 转发新 API │ 新代码不感知它 + └─────────────────────┘ +``` + +--- + +## 二、两个核心概念的区分 + +``` +┌──────────────────────────────────────────────────────────────┐ +│ HandlerDescriptor │ +│ 方向：插件 ──► 主进程（initialize 时声明） │ +│ 含义：插件订阅"我能响应哪些事件" │ +│ 例子：@on_command("hello") → 订阅 /hello 命令 │ +└──────────────────────────────────────────────────────────────┘ + +┌──────────────────────────────────────────────────────────────┐ +│ CapabilityInvocation │ +│ 方向：插件 ──► 主进程（运行时按需调用） │ +│ 含义：插件请求"帮我执行这个能力" │ +│ 例子：ctx.llm.chat() → invoke "llm.chat" │ +└──────────────────────────────────────────────────────────────┘ + +┌──────────────────────────────────────────────────────────────┐ +│ CapabilityDescriptor │ +│ 方向：主进程 ──► 插件（initialize_result 时返回） │ +│ 含义：主进程声明"我提供哪些能力" │ +│ 例子：{ name: "llm.chat", supports_stream: false, ... } │ +└──────────────────────────────────────────────────────────────┘ +``` + +| | HandlerDescriptor | CapabilityDescriptor | CapabilityInvocation | +|---|---|---|---| +| 谁发 | 插件 | 主进程 | 插件 | +| 何时 | initialize 时 | initialize_result 时 | 运行时 | +| 主进程动作 | 注册订阅 | 告知可用能力 | 执行并返回结果 | + +--- + +## 三、分层职责 + +``` +┌─────────────────────────────────────────────────────┐ +│ Layer 1：用户层 │ +│ Star / 装饰器 / MessageEvent │ +│ 插件作者只接触这一层 │ +│ 不知道：RPC、进程、序列化、订阅协议 │ +└──────────────────────────┬──────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ Layer 2：API 层 │ +│ Context / LLMClient / DBClient / MemoryClient │ +│ PlatformClient │ +│ 把能力包装成类型化 API │ +│ 不知道：JSON 格式、id、transport │ +└──────────────────────────┬──────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ Layer 3：翻译层 │ +│ CapabilityProxy │ +│ API 调用 → Peer.invoke(name, input, stream) │ +│ output dict → 返回类型 │ +│ 无业务逻辑，一一对应 │ +└──────────────────────────┬──────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ Layer 4：通信层 │ +│ Peer / Transport / Protocol Messages │ +│ 可靠收发消息 │ +│ 不知道业务，只知道消息格式 │ +└─────────────────────────────────────────────────────┘ + + ※ compat.py 不是第五层，是用户层和 API 层的旁路入口。 + 新代码不 import 它，可整体删除。 +``` + +--- + +## 四、目录结构 + +``` +astrbot_sdk/ +│ +├── star.py +├── context.py +├── decorators.py +├── events.py +├── errors.py +├── compat.py ← 旁路，不是核心层 +│ +├── clients/ +│ ├── llm.py +│ ├── memory.py +│ ├── db.py +│ └── platform.py +│ +├── runtime/ +│ ├── peer.py +│ ├── transport.py +│ ├── capability_router.py +│ ├── handler_dispatcher.py +│ ├── loader.py +│ └── bootstrap.py +│ +└── protocol/ + ├── messages.py ← 所有协议消息类型 + ├── descriptors.py ← HandlerDescriptor / CapabilityDescriptor + └── legacy_adapter.py ← 旧线协议翻译，只做翻译无业务逻辑 +``` + +--- + +## 五、协议消息定义（完整版） + +### 五条硬规则 + +**规则一：统一使用 `id` 字段关联所有请求与响应** +``` +所有消息只用一个关联字段：id +不区分 request_id / invocation_id，全部统一成 id。 +发送方生成 id，接收方响应时原样带回，双方按 id 配对。 +``` + +**规则二：event 只用于 stream=true 的调用** +``` +stream=false 的调用只能以单个 result 结束。 +stream=true 的调用只能以 event 序列结束。 +stream=false 的调用不得发送 event(started/delta/completed/failed)。 +违反此规则的实现视为协议错误。 +``` + +**规则三：插件 handler 回调走统一 invoke，不新增消息类型** +``` +主进程触发插件处理器时： + capability: "handler.invoke" + input: { handler_id: str, event: { 纯数据 } } + +ctx 不通过线协议传输。 +ctx 由插件进程本地重建并注入处理器。 +看到处理器签名有 ctx 参数，不要误以为需要从主进程发过来。 +``` + +**规则四：cancel 是"请求停止"，不是"立即停止"** +``` +收到 cancel 后： + 若调用已结束 → 忽略，不报错 + 若调用仍在执行 → 尽力中断，发送统一终止态 + +统一终止态： + stream=true: event { phase: "failed", error: { code: "cancelled" } } + stream=false: result { success: false, error: { code: "cancelled" } } + +调用方收到 cancel 后必须等待终止态，不能认为发完 cancel 就已结束。 +``` + +**规则五：initialize 失败后连接进入不可用状态** +``` +initialize 失败（协议版本不兼容 / handlers 非法 / 元信息缺失）时： + 返回 result { kind: "initialize_result", success: false, error: {...} } + 连接进入不可用状态 + 除关闭连接外，不得继续发送普通 invoke + 对端收到失败的 initialize_result 后应立即关闭连接 +``` + +--- + +### 消息格式 + +**initialize** +```json +{ + "type": "initialize", + "id": "msg_001", + "protocol_version": "1.0", + "peer": { + "name": "my-plugin", + "role": "plugin", + "version": "1.2.0" + }, + "handlers": [ "HandlerDescriptor ..." ], + "metadata": {} +} +``` + +**initialize_result（成功）** +```json +{ + "type": "result", + "id": "msg_001", + "kind": "initialize_result", + "success": true, + "output": { + "peer": { "name": "astrbot-core", "role": "core" }, + "capabilities": [ "CapabilityDescriptor ..." ], + "metadata": {} + } +} +``` + +**initialize_result（失败）** +```json +{ + "type": "result", + "id": "msg_001", + "kind": "initialize_result", + "success": false, + "error": { + "code": "protocol_version_mismatch", + "message": "服务端支持协议版本 1.0，客户端请求版本 2.0", + "hint": "请升级 astrbot_sdk 至最新版本", + "retryable": false + } +} +``` +※ 失败后连接进入不可用状态，对端应立即关闭连接。 + +**invoke（普通能力）** +```json +{ + "type": "invoke", + "id": "msg_002", + "capability": "llm.chat", + "input": { "prompt": "hi", "system": null }, + "stream": false +} +``` + +**invoke（流式能力）** +```json +{ + "type": "invoke", + "id": "msg_003", + "capability": "llm.stream_chat", + "input": { "prompt": "hi" }, + "stream": true +} +``` + +**invoke（handler 回调）** +```json +{ + "type": "invoke", + "id": "msg_010", + "capability": "handler.invoke", + "input": { + "handler_id": "handler_abc123", + "event": { + "text": "/hello", + "user_id": "u_001", + "group_id": null, + "platform": "qq" + } + }, + "stream": false +} +``` +※ input.event 只含纯数据字段。ctx 由插件进程本地构建并注入，不经过线协议传输。 + +**result（成功）** +```json +{ + "type": "result", + "id": "msg_002", + "success": true, + "output": { "text": "你好！" } +} +``` + +**result（失败）** +```json +{ + "type": "result", + "id": "msg_002", + "success": false, + "error": { + "code": "llm_not_configured", + "message": "未找到可用的大模型配置", + "hint": "请在管理面板的「模型管理」中添加模型", + "retryable": false + } +} +``` + +**event 序列（stream=true 专用）** +```json +{ "type": "event", "id": "msg_003", "phase": "started" } +{ "type": "event", "id": "msg_003", "phase": "delta", "data": { "text": "你" } } +{ "type": "event", "id": "msg_003", "phase": "delta", "data": { "text": "好" } } +{ "type": "event", "id": "msg_003", "phase": "completed", "output": { "text": "你好" } } +``` + +**event（取消终止态）** +```json +{ + "type": "event", + "id": "msg_003", + "phase": "failed", + "error": { + "code": "cancelled", + "message": "调用被取消", + "hint": "", + "retryable": false + } +} +``` + +**cancel** +```json +{ + "type": "cancel", + "id": "msg_003", + "reason": "user_cancelled" +} +``` + +--- + +## 六、描述符定义 + +### HandlerDescriptor + +``` +HandlerDescriptor +{ + id: str 唯一标识，主进程回调时填入 handler_id + trigger: CommandTrigger + | MessageTrigger + | EventTrigger + | ScheduleTrigger + priority: int 默认 0，越大越先执行 + permissions: { + require_admin: bool + level: int + } +} +``` + +trigger 判别联合：不同 type 只允许对应字段出现，其他字段必须省略。 + +``` +CommandTrigger +{ + type: "command" + command: str 必填 + aliases: [str] 可选，默认 [] + description: str 可选 +} + +MessageTrigger +{ + type: "message" + regex: str | null 可选 + keywords: [str] 可选，默认 [] + platforms: [str] 可选，默认 []（空表示所有平台） +} + +EventTrigger +{ + type: "event" + event_type: str 必填 +} + +ScheduleTrigger +{ + type: "schedule" + cron: str | null + interval_seconds: int | null +} +规则：cron 和 interval_seconds 必须且只能有一个非 null +``` + +### CapabilityDescriptor + +主进程在 initialize_result.output.capabilities 中返回。 + +``` +CapabilityDescriptor +{ + name: str capability name，如 "llm.chat" + description: str 一句话说明 + input_schema: JSONSchema | null 输入结构定义 + output_schema: JSONSchema | null 输出结构定义 + supports_stream: bool 是否支持 stream=true 调用 + cancelable: bool 是否支持 cancel +} + +schema 治理规则： + 内建核心 capability（llm.* / db.* / memory.* / platform.*） + 必须提供 input_schema 和 output_schema + 兼容期或动态注册的 capability + 允许为 null，但应在路线图中补全 + 不得以"动态能力"为由长期保持 null +``` + +示例： +```json +{ + "name": "llm.chat", + "description": "发送对话请求，返回模型回复文本", + "input_schema": { + "type": "object", + "properties": { + "prompt": { "type": "string" }, + "system": { "type": "string" }, + "model": { "type": "string" }, + "temperature": { "type": "number" } + }, + "required": ["prompt"] + }, + "output_schema": { + "type": "object", + "properties": { + "text": { "type": "string" } + }, + "required": ["text"] + }, + "supports_stream": false, + "cancelable": false +} +``` + +--- + +## 七、Capability Name 约定 + +``` +格式：{namespace}.{method} + +内建 capability 列表： + llm.chat ctx.llm.chat() + llm.chat_raw ctx.llm.chat_raw() + llm.stream_chat ctx.llm.stream_chat() + memory.search ctx.memory.search() + memory.save ctx.memory.save() + memory.delete ctx.memory.delete() + db.get ctx.db.get() + db.set ctx.db.set() + db.delete ctx.db.delete() + db.list ctx.db.list() + platform.send ctx.platform.send() + platform.send_image ctx.platform.send_image() + platform.get_members ctx.platform.get_members() + +保留命名空间（插件不可使用这些前缀）： + handler.* 框架内部：处理器回调 + system.* 框架内部：系统级操作 + internal.* 框架内部：保留扩展 + +命名规则： + 全小写，点分隔命名空间，下划线分隔单词，不用驼峰 + capability name 是协议约定，手写定义，不自动从方法名推导 + 方法名重构不影响协议；协议变更需同步更新方法名和文档 +``` + +--- + +## 八、错误模型 + +```python +@dataclass +class AstrBotError(Exception): + code: str # 机器可读，如 "llm_not_configured" + message: str # 发生了什么 + hint: str # 用户怎么修 + retryable: bool # True = 可重试（超时、网络抖动、临时不可用） + # False = 重试无意义（权限不足、能力不存在、配置缺失） +``` + +``` +可重试（retryable=true） 不可重试（retryable=false） +───────────────────────── ──────────────────────────── +CapabilityTimeout LLMNotConfigured +NetworkError CapabilityNotFound +LLMTemporaryError PermissionDenied + LLMError（模型返回错误） + InvalidInput + Cancelled + ProtocolVersionMismatch +``` + +**Star.on_error 默认兜底：** +``` +AstrBotError retryable=true → 回复"请求失败，请稍后重试" +AstrBotError retryable=false → 回复 error.hint +其他异常 → 回复"出了点问题，请联系插件作者" +所有情况均打完整 traceback 日志 +插件作者覆盖 on_error 可完全自定义 +``` + +--- + +## 九、Context 设计规则 + +```python +class Context: + + # 第一类：插件常用能力 Client（稳定，只扩展不删除） + llm: LLMClient + memory: MemoryClient + db: DBClient + platform: PlatformClient + + # 第二类：少量基础运行时信息 + plugin_id: str + logger: Logger # 自动带插件名前缀 + cancel_token: ... # 取消当前调用 + + # ❌ 不直接挂顶层： + # tools / runtime / scheduler / http / + # storage / persona / workflow / config + # 有需要时设计专属 Client 后再加 +``` + +--- + +## 十、LLM Client 分层 + +```python +class LLMClient: + + async def chat( + self, + prompt: str, + *, + system: str | None = None, + history: list[ChatMessage] | None = None, + model: str | None = None, + temperature: float | None = None, + ) -> str: + """发送对话请求，返回回复文本。爱好者场景首选。""" + + async def chat_raw( + self, + prompt: str, + **kwargs, + ) -> LLMResponse: + """返回完整响应，含 usage / finish_reason / tool_calls。""" + + async def stream_chat( + self, + prompt: str, + *, + system: str | None = None, + history: list[ChatMessage] | None = None, + ) -> AsyncGenerator[str, None]: + """流式对话，逐字返回文本片段。""" +``` + +chat() 和 chat_raw() 是唯二入口，不再增加第三种变体。 + +--- + +## 十一、关键数据流 + +### 11.1 插件加载与握手 + +``` +框架启动 + → loader.py 扫描目录，发现 Star 子类 + → 收集 __handlers__，转成 HandlerDescriptor 列表 + → Peer 发送 initialize { id: "msg_001", handlers: [...] } + → 主进程注册事件订阅 + → 主进程返回 initialize_result { id: "msg_001", + success: true, + capabilities: [...] } + → 插件 CapabilityProxy 缓存 capabilities + → 插件就绪 + +握手失败时： + → 主进程返回 initialize_result { success: false, error: {...} } + → 连接进入不可用状态 + → 插件进程关闭连接，打错误日志 + → 不发送任何 invoke +``` + +### 11.2 外部消息触发处理器 + +``` +外部用户发送 /hello + → 主进程 HandlerDispatcher 匹配订阅 + → 主进程发送 invoke { + id: "msg_010", + capability: "handler.invoke", + input: { + handler_id: "handler_abc", + event: { text: "/hello", user_id: "u_001", ... } + }, + stream: false + } + → 插件 handler_dispatcher 找到处理器方法 + → 本地构建 ctx，注入 event 和 ctx，执行处理器 + → 处理器内调用 ctx.llm.chat()（进入 11.3） +``` + +注：ctx 在插件进程本地构建，不经过线协议传输。 + +### 11.3 非流式能力调用 + +``` +ctx.llm.chat("hi") + → CapabilityProxy 构造 input + → Peer.invoke("llm.chat", {prompt:"hi"}, stream=false) id="msg_020" + → 发送 { type:"invoke", id:"msg_020", capability:"llm.chat", + input:{...}, stream:false } + ← 收到 { type:"result", id:"msg_020", success:true, + output:{text:"你好"} } + → 解包 output.text → 返回 str + +※ stream=false 不会收到任何 event 消息 +``` + +### 11.4 流式能力调用 + +``` +async for chunk in ctx.llm.stream_chat("hi"): + → Peer.invoke("llm.stream_chat", {...}, stream=true) id="msg_030" + ← event { id:"msg_030", phase:"started" } + ← event { id:"msg_030", phase:"delta", data:{text:"你"} } → yield "你" + ← event { id:"msg_030", phase:"delta", data:{text:"好"} } → yield "好" + ← event { id:"msg_030", phase:"completed", output:{text:"你好"} } + → 生成器结束 + +※ stream=true 不会收到 result 消息，只收到 event 序列 +``` + +--- + +## 十二、兼容层 + +``` +compat.py 三条铁律： + 1. 新代码不 import compat.py + 2. compat.py 只 import 新代码 + 3. compat.py 里只有"转发"，无业务逻辑 + +旧 API 映射： + +旧写法 新写法 +────────────────────────────────────────────────────────────── +CommandComponent → Star +context.llm_generate(prompt) → ctx.llm.chat(prompt) +context.tool_loop_agent(...) → ctx.llm.chat_raw(...) 含 tools +context.send_message(session, mc) → ctx.platform.send(session, text) +context.put_kv_data(key, value) → ctx.db.set(key, value) +context.get_kv_data(key) → ctx.db.get(key) +@filter.command("hello") → @on_command("hello") +@filter.regex("pattern") → @on_message(regex="pattern") +@filter.permission(ADMIN) → @require_admin +yield event.plain_result("hi") → await event.reply("hi") + +deprecated warning 格式（每个方法只打一次）： + [AstrBot] 警告：context.llm_generate() 已过时。 + 请替换为：ctx.llm.chat(prompt) + 迁移文档：https://docs.astrbot.app/migration/v3 + +流式兼容： + 旧 yield 写法只在 compat 层兜底处理 + 新 API 只推荐 AsyncGenerator，不双轨并行 + +legacy_adapter.py 职责边界： + 只翻译旧线协议消息 ↔ 新线协议消息 + 不含业务逻辑，不被新代码 import + 生命周期结束时整个删掉，新代码零修改 +``` + +--- + +## 十三、迁移计划 + +``` +阶段 0：立骨架（当前可开始） +────────────────────────────────────────────────────── + 做什么： + ✦ 新建 star / context / decorators / events / errors / clients + ✦ protocol/descriptors.py 写清 HandlerDescriptor（判别联合） + 和 CapabilityDescriptor（含 schema 治理规则） + ✦ Peer 用 mock 占位（invoke 返回假数据） + ✦ 写 compat.py + + 验收： + ✦ 旧插件加载不报错 + ✦ 新写法能跑通基本流程 + ✦ IDE 对 ctx.llm / ctx.db 有完整补全 + + +阶段 1：接通信层 +────────────────────────────────────────────────────── + 做什么： + ✦ 实现 Transport / Peer（统一 id 字段） + ✦ 实现 capability_router + handler_dispatcher + ✦ 实现 legacy_adapter（旧协议翻译） + ✦ clients/ 接上真实 capability 调用 + ✦ 实现 cancel 语义（请求停止，等终止态） + ✦ 实现 initialize 失败处理（连接不可用 + 关闭） + + 验收： + ✦ 端到端调用成功 + ✦ 流式响应正常，stream=false 不出现 event 消息 + ✦ initialize 失败时连接正确关闭 + ✦ retryable 错误触发自动提示，不可重试触发 hint + + +阶段 2：清理旧实现 +────────────────────────────────────────────────────── + 做什么： + ✦ 删除 api/star/context.py（旧 Context） + ✦ 删除 runtime/rpc/ 旧角色划分 + ✦ 删除 runtime/stars/filter/ 旧装饰器实现 + ✦ deprecated warning 升级为更显眼提示 + + 验收： + ✦ 旧插件仍通过 compat.py 运行 + ✦ 核心路径无旧抽象引用 + + +阶段 3：废弃旧 API（下一大版本） +────────────────────────────────────────────────────── + 做什么： + ✦ deprecated warning 变启动报错 + ✦ 生态迁移完成后删除 compat.py 和 legacy_adapter.py + + 验收： + ✦ 删除 compat.py 后新代码零修改 +``` + +--- + +## 十四、设计决策记录 + +| 问题 | 决策 | 理由 | +|------|------|------| +| 关联字段用什么名 | 统一 `id` | 防止 request_id / invocation_id 在 initialize_result 处产生歧义 | +| event 能用于非流式吗 | 不能，硬规则 | 防止"非流式先发 started 再发 result"污染处理逻辑 | +| initialize 失败后能继续发 invoke 吗 | 不能，连接进入不可用状态 | 防止在无效连接上堆积调用 | +| handler 回调走什么机制 | handler.invoke，不新增消息类型 | 协议保持五种消息 | +| ctx 从哪来 | 插件进程本地构建，不经线协议 | ctx 含运行时状态不可序列化，且无需传输 | +| HandlerDescriptor trigger 结构 | 判别联合 | 防止大量可空字段，方便校验和类型推导 | +| CapabilityDescriptor schema 是否可为 null | 可以，但内建 capability 必须提供 | 防止所有人偷懒填 null 导致 schema 形同虚设 | +| 保留命名空间 | handler.* / system.* / internal.* | 集中声明，防止插件误用或冲突 | +| 错误模型 | code + message + hint + retryable | retryable 区分策略差异，hint 直接告诉用户怎么修 | +| cancel 语义 | 请求停止，等终止态 | 避免实现侧歧义，调用方行为确定 | +| compat 定位 | 旁路入口，不是核心层 | 新代码不感知，可整体删除 | +| Context 扩展规则 | 只放常用能力 Client + 少量运行时信息 | 防止变成圣诞树 | +| chat() 返回类型 | str；进阶用 chat_raw() | 爱好者不拆包装，进阶有专用入口，两个定死 | +| 序列化 | 默认 JSON，不用 pickle | 跨语言，安全，可观测 | + +--- + +*本文档是代码的源头。Python SDK、通信层、主进程三端有分歧时，以本文档为准。* + +*v4 修正：补充 event 只用于 stream=true 的硬规则；initialize 失败场景和连接不可用状态；ctx 不经线协议传输的明确说明；CapabilityDescriptor schema 治理规则；保留命名空间集中声明（handler.* / system.* / internal.*）；initialize_result 失败示例。* +# AstrBot SDK 重构架构设计 v4 + +--- + +## 一、全局架构图 + +``` +╔══════════════════════════════════════════════════════════════════════╗ +║ 插件作者的世界 ║ +║ ║ +║ class MyPlugin(Star): ║ +║ @on_command("hello") ║ +║ async def hello(self, event: MessageEvent, ctx: Context): ║ +║ reply = await ctx.llm.chat(event.text) ║ +║ await event.reply(reply) ║ +║ ║ +╚══════════════════════════════════════════════════════════════════════╝ + │ Star / 装饰器 / Event │ Context / Clients + ▼ ▼ +┌─────────────────────┐ ┌────────────────────────────────┐ +│ Handler 系统 │ │ Capability 调用系统 │ +│ │ │ │ +│ HandlerDescriptor │ │ ctx.llm.chat() │ +│ HandlerDispatcher │ │ ctx.memory.search() │ +│ │ │ ctx.db.get() │ +│ 插件 → 主进程 │ │ ctx.platform.send() │ +│ "我能响应这些事件" │ │ │ +│ │ │ 插件 → 主进程 │ +│ │ │ "帮我调用这个能力" │ +└──────────┬──────────┘ └──────────────┬─────────────────┘ + │ │ + └─────────────────┬────────────────┘ + ▼ +┌──────────────────────────────────────────────────────────────────┐ +│ 通信层 │ +│ │ +│ 所有消息统一使用 id 字段关联请求与响应 │ +│ │ +│ Peer.initialize(handlers=[...]) │ +│ Peer.invoke("llm.chat", input, stream=false) → result │ +│ Peer.invoke("llm.stream_chat", input, stream=true) → event* │ +│ Peer.invoke("handler.invoke", {handler_id, event}) │ +│ │ +│ Transport: StdioTransport / WebSocketTransport │ +└──────────────────────────────────────────────────────────────────┘ + │ JSON 消息流 + ▼ +┌──────────────────────────────────────────────────────────────────┐ +│ 主进程（AstrBot Core） │ +│ │ +│ CapabilityRouter ──► "llm.chat" ──► LLM Service │ +│ ──► "db.get" ──► Storage │ +│ ──► "handler.invoke" ──► 转发给插件 │ +│ │ +│ HandlerDispatcher ◄── 外部消息 ──► 匹配订阅 ──► 回调插件 │ +└──────────────────────────────────────────────────────────────────┘ + + ┌─────────────────────┐ + │ compat.py（旁路） │ ← 不是核心层 + │ 旧 API → 转发新 API │ 新代码不感知它 + └─────────────────────┘ +``` + +--- + +## 二、两个核心概念的区分 + +``` +┌──────────────────────────────────────────────────────────────┐ +│ HandlerDescriptor │ +│ 方向：插件 ──► 主进程（initialize 时声明） │ +│ 含义：插件订阅"我能响应哪些事件" │ +│ 例子：@on_command("hello") → 订阅 /hello 命令 │ +└──────────────────────────────────────────────────────────────┘ + +┌──────────────────────────────────────────────────────────────┐ +│ CapabilityInvocation │ +│ 方向：插件 ──► 主进程（运行时按需调用） │ +│ 含义：插件请求"帮我执行这个能力" │ +│ 例子：ctx.llm.chat() → invoke "llm.chat" │ +└──────────────────────────────────────────────────────────────┘ + +┌──────────────────────────────────────────────────────────────┐ +│ CapabilityDescriptor │ +│ 方向：主进程 ──► 插件（initialize_result 时返回） │ +│ 含义：主进程声明"我提供哪些能力" │ +│ 例子：{ name: "llm.chat", supports_stream: false, ... } │ +└──────────────────────────────────────────────────────────────┘ +``` + +| | HandlerDescriptor | CapabilityDescriptor | CapabilityInvocation | +|---|---|---|---| +| 谁发 | 插件 | 主进程 | 插件 | +| 何时 | initialize 时 | initialize_result 时 | 运行时 | +| 主进程动作 | 注册订阅 | 告知可用能力 | 执行并返回结果 | + +--- + +## 三、分层职责 + +``` +┌─────────────────────────────────────────────────────┐ +│ Layer 1：用户层 │ +│ Star / 装饰器 / MessageEvent │ +│ 插件作者只接触这一层 │ +│ 不知道：RPC、进程、序列化、订阅协议 │ +└──────────────────────────┬──────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ Layer 2：API 层 │ +│ Context / LLMClient / DBClient / MemoryClient │ +│ PlatformClient │ +│ 把能力包装成类型化 API │ +│ 不知道：JSON 格式、id、transport │ +└──────────────────────────┬──────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ Layer 3：翻译层 │ +│ CapabilityProxy │ +│ API 调用 → Peer.invoke(name, input, stream) │ +│ output dict → 返回类型 │ +│ 无业务逻辑，一一对应 │ +└──────────────────────────┬──────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ Layer 4：通信层 │ +│ Peer / Transport / Protocol Messages │ +│ 可靠收发消息 │ +│ 不知道业务，只知道消息格式 │ +└─────────────────────────────────────────────────────┘ + + ※ compat.py 不是第五层，是用户层和 API 层的旁路入口。 + 新代码不 import 它，可整体删除。 +``` + +--- + +## 四、目录结构 + +``` +astrbot_sdk/ +│ +├── star.py +├── context.py +├── decorators.py +├── events.py +├── errors.py +├── compat.py ← 旁路，不是核心层 +│ +├── clients/ +│ ├── llm.py +│ ├── memory.py +│ ├── db.py +│ └── platform.py +│ +├── runtime/ +│ ├── peer.py +│ ├── transport.py +│ ├── capability_router.py +│ ├── handler_dispatcher.py +│ ├── loader.py +│ └── bootstrap.py +│ +└── protocol/ + ├── messages.py + ├── descriptors.py + └── legacy_adapter.py +``` + +--- + +## 五、协议消息定义（完整版） + +### 五条硬规则 + +**规则一：统一使用 `id` 字段关联所有请求与响应** +``` +所有消息只用一个关联字段：id +不区分 request_id / invocation_id，全部统一成 id。 +发送方生成 id，接收方响应时原样带回，双方按 id 配对。 +``` + +**规则二：event 只用于 stream=true 的调用** +``` +stream=false 的调用只能以单个 result 结束。 +stream=true 的调用只能以 event 序列结束。 +stream=false 的调用不得发送 event(started/delta/completed/failed)。 +违反此规则的实现视为协议错误。 +``` + +**规则三：插件 handler 回调走统一 invoke，不新增消息类型** +``` +主进程触发插件处理器时： + capability: "handler.invoke" + input: { handler_id: str, event: { 纯数据 } } + +ctx 不通过线协议传输。 +ctx 由插件进程本地重建并注入处理器。 +看到处理器签名有 ctx 参数，不要误以为需要从主进程发过来。 +``` + +**规则四：cancel 是"请求停止"，不是"立即停止"** +``` +收到 cancel 后： + 若调用已结束 → 忽略，不报错 + 若调用仍在执行 → 尽力中断，发送统一终止态 + +统一终止态： + stream=true: event { phase: "failed", error: { code: "cancelled" } } + stream=false: result { success: false, error: { code: "cancelled" } } + +调用方收到 cancel 后必须等待终止态，不能认为发完 cancel 就已结束。 +``` + +**规则五：initialize 失败后连接进入不可用状态** +``` +initialize 失败（协议版本不兼容 / handlers 非法 / 元信息缺失）时： + 返回 result { kind: "initialize_result", success: false, error: {...} } + 连接进入不可用状态 + 除关闭连接外，不得继续发送普通 invoke + 对端收到失败的 initialize_result 后应立即关闭连接 +``` + +--- + +### 消息格式 + +**initialize** +```json +{ + "type": "initialize", + "id": "msg_001", + "protocol_version": "1.0", + "peer": { + "name": "my-plugin", + "role": "plugin", + "version": "1.2.0" + }, + "handlers": [ "HandlerDescriptor ..." ], + "metadata": {} +} +``` + +**initialize_result（成功）** +```json +{ + "type": "result", + "id": "msg_001", + "kind": "initialize_result", + "success": true, + "output": { + "peer": { "name": "astrbot-core", "role": "core" }, + "capabilities": [ "CapabilityDescriptor ..." ], + "metadata": {} + } +} +``` + +**initialize_result（失败）** +```json +{ + "type": "result", + "id": "msg_001", + "kind": "initialize_result", + "success": false, + "error": { + "code": "protocol_version_mismatch", + "message": "服务端支持协议版本 1.0，客户端请求版本 2.0", + "hint": "请升级 astrbot_sdk 至最新版本", + "retryable": false + } +} +``` +※ 失败后连接进入不可用状态，对端应立即关闭连接。 + +**invoke（普通能力）** +```json +{ + "type": "invoke", + "id": "msg_002", + "capability": "llm.chat", + "input": { "prompt": "hi", "system": null }, + "stream": false +} +``` + +**invoke（流式能力）** +```json +{ + "type": "invoke", + "id": "msg_003", + "capability": "llm.stream_chat", + "input": { "prompt": "hi" }, + "stream": true +} +``` + +**invoke（handler 回调）** +```json +{ + "type": "invoke", + "id": "msg_010", + "capability": "handler.invoke", + "input": { + "handler_id": "handler_abc123", + "event": { + "text": "/hello", + "user_id": "u_001", + "group_id": null, + "platform": "qq" + } + }, + "stream": false +} +``` +※ input.event 只含纯数据字段。ctx 由插件进程本地构建并注入，不经过线协议传输。 + +**result（成功）** +```json +{ + "type": "result", + "id": "msg_002", + "success": true, + "output": { "text": "你好！" } +} +``` + +**result（失败）** +```json +{ + "type": "result", + "id": "msg_002", + "success": false, + "error": { + "code": "llm_not_configured", + "message": "未找到可用的大模型配置", + "hint": "请在管理面板的「模型管理」中添加模型", + "retryable": false + } +} +``` + +**event 序列（stream=true 专用）** +```json +{ "type": "event", "id": "msg_003", "phase": "started" } +{ "type": "event", "id": "msg_003", "phase": "delta", "data": { "text": "你" } } +{ "type": "event", "id": "msg_003", "phase": "delta", "data": { "text": "好" } } +{ "type": "event", "id": "msg_003", "phase": "completed", "output": { "text": "你好" } } +``` + +**event（取消终止态）** +```json +{ + "type": "event", + "id": "msg_003", + "phase": "failed", + "error": { + "code": "cancelled", + "message": "调用被取消", + "hint": "", + "retryable": false + } +} +``` + +**cancel** +```json +{ + "type": "cancel", + "id": "msg_003", + "reason": "user_cancelled" +} +``` + +--- + +## 六、描述符定义 + +### HandlerDescriptor + +``` +HandlerDescriptor +{ + id: str + trigger: CommandTrigger + | MessageTrigger + | EventTrigger + | ScheduleTrigger + priority: int + permissions: { + require_admin: bool + level: int + } +} +``` + +trigger 判别联合：不同 type 只允许对应字段出现，其他字段必须省略。 + +``` +CommandTrigger +{ + type: "command" + command: str + aliases: [str] + description: str +} + +MessageTrigger +{ + type: "message" + regex: str | null + keywords: [str] + platforms: [str] +} + +EventTrigger +{ + type: "event" + event_type: str +} + +ScheduleTrigger +{ + type: "schedule" + cron: str | null + interval_seconds: int | null +} +``` + +### CapabilityDescriptor + +``` +CapabilityDescriptor +{ + name: str + description: str + input_schema: JSONSchema | null + output_schema: JSONSchema | null + supports_stream: bool + cancelable: bool +} +``` + +schema 治理规则： +``` +内建核心 capability（llm.* / db.* / memory.* / platform.*）必须提供 input_schema 和 output_schema +兼容期或动态注册的 capability 允许为 null，但应在路线图中补全 +不得以“动态能力”为由长期保持 null +``` + +--- + +## 七、Capability Name 约定 + +``` +格式：{namespace}.{method} + +内建 capability 列表： + llm.chat + llm.chat_raw + llm.stream_chat + memory.search + memory.save + memory.delete + db.get + db.set + db.delete + db.list + platform.send + platform.send_image + platform.get_members + +保留命名空间： + handler.* + system.* + internal.* +``` + +--- + +## 八、错误模型 + +```python +@dataclass +class AstrBotError(Exception): + code: str + message: str + hint: str + retryable: bool +``` + +--- + +## 九、Context 设计规则 + +```python +class Context: + llm: LLMClient + memory: MemoryClient + db: DBClient + platform: PlatformClient + plugin_id: str + logger: Logger + cancel_token: ... +``` + +--- + +## 十、LLM Client 分层 + +```python +class LLMClient: + async def chat(...) -> str: ... + async def chat_raw(...) -> LLMResponse: ... + async def stream_chat(...) -> AsyncGenerator[str, None]: ... +``` + +--- + +## 十一、关键数据流 + +- 插件启动后发送 `initialize` +- 主进程返回 `initialize_result` +- 外部消息触发 `handler.invoke` +- 非流式能力调用只返回 `result` +- 流式能力调用只返回 `event` 序列 + +--- + +## 十二、兼容层 + +``` +compat.py 三条铁律： + 1. 新代码不 import compat.py + 2. compat.py 只 import 新代码 + 3. compat.py 里只有转发，无业务逻辑 +``` + +--- + +## 十三、迁移计划 + +- 阶段 0：立骨架 +- 阶段 1：接通信层 +- 阶段 2：清理旧实现 +- 阶段 3：废弃旧 API + +--- + +## 十四、设计决策记录 + +- 统一 `id` +- `event` 只用于 `stream=true` +- initialize 失败后连接不可用 +- handler 回调走 `handler.invoke` +- ctx 在插件进程本地构建 +- HandlerDescriptor 使用判别联合 +- 内建 capability schema 必填 +- 保留命名空间 `handler.* / system.* / internal.*` +- 错误模型固定为 `code + message + hint + retryable` +- cancel 为“请求停止，等待终止态” +- compat 是旁路，不是核心层 +- Context 只挂常用 client 与少量运行时信息 +- `chat()` 返回 `str`，进阶使用 `chat_raw()` +- 序列化使用 JSON，不用 pickle + +--- + +*本文档是代码的源头。Python SDK、通信层、主进程三端有分歧时，以本文档为准。* diff --git a/src-new.rar b/src-new.rar new file mode 100644 index 000000000..f9dc50ca2 Binary files /dev/null and b/src-new.rar differ diff --git a/src-new/astrbot_sdk/__init__.py b/src-new/astrbot_sdk/__init__.py new file mode 100644 index 000000000..b9e2d98df --- /dev/null +++ b/src-new/astrbot_sdk/__init__.py @@ -0,0 +1,17 @@ +from .context import Context +from .decorators import on_command, on_event, on_message, on_schedule, require_admin +from .errors import AstrBotError +from .events import MessageEvent +from .star import Star + +__all__ = [ + "AstrBotError", + "Context", + "MessageEvent", + "Star", + "on_command", + "on_event", + "on_message", + "on_schedule", + "require_admin", +] diff --git a/src-new/astrbot_sdk/__main__.py b/src-new/astrbot_sdk/__main__.py new file mode 100644 index 000000000..491a4d136 --- /dev/null +++ b/src-new/astrbot_sdk/__main__.py @@ -0,0 +1,5 @@ +from .cli import cli + + +if __name__ == "__main__": + cli() diff --git a/src-new/astrbot_sdk/api/__init__.py b/src-new/astrbot_sdk/api/__init__.py new file mode 100644 index 000000000..c9c2ef67b --- /dev/null +++ b/src-new/astrbot_sdk/api/__init__.py @@ -0,0 +1 @@ +__all__: list[str] = [] diff --git a/src-new/astrbot_sdk/api/components/__init__.py b/src-new/astrbot_sdk/api/components/__init__.py new file mode 100644 index 000000000..9d075b559 --- /dev/null +++ b/src-new/astrbot_sdk/api/components/__init__.py @@ -0,0 +1,3 @@ +from .command import CommandComponent + +__all__ = ["CommandComponent"] diff --git a/src-new/astrbot_sdk/api/components/command.py b/src-new/astrbot_sdk/api/components/command.py new file mode 100644 index 000000000..487fa288b --- /dev/null +++ b/src-new/astrbot_sdk/api/components/command.py @@ -0,0 +1,3 @@ +from ...compat import CommandComponent + +__all__ = ["CommandComponent"] diff --git a/src-new/astrbot_sdk/api/event/__init__.py b/src-new/astrbot_sdk/api/event/__init__.py new file mode 100644 index 000000000..611042749 --- /dev/null +++ b/src-new/astrbot_sdk/api/event/__init__.py @@ -0,0 +1,4 @@ +from ...events import MessageEvent as AstrMessageEvent +from .filter import ADMIN, filter + +__all__ = ["ADMIN", "AstrMessageEvent", "filter"] diff --git a/src-new/astrbot_sdk/api/event/filter.py b/src-new/astrbot_sdk/api/event/filter.py new file mode 100644 index 000000000..dc6b70bf2 --- /dev/null +++ b/src-new/astrbot_sdk/api/event/filter.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from ...decorators import on_command, on_message, require_admin + +ADMIN = "admin" + + +def command(name: str): + return on_command(name) + + +def regex(pattern: str): + return on_message(regex=pattern) + + +def permission(level): + if level == ADMIN: + return require_admin + + def decorator(func): + return func + + return decorator + + +class _FilterNamespace: + command = staticmethod(command) + regex = staticmethod(regex) + permission = staticmethod(permission) + + +filter = _FilterNamespace() diff --git a/src-new/astrbot_sdk/api/star/__init__.py b/src-new/astrbot_sdk/api/star/__init__.py new file mode 100644 index 000000000..43a016780 --- /dev/null +++ b/src-new/astrbot_sdk/api/star/__init__.py @@ -0,0 +1,3 @@ +from .context import Context + +__all__ = ["Context"] diff --git a/src-new/astrbot_sdk/api/star/context.py b/src-new/astrbot_sdk/api/star/context.py new file mode 100644 index 000000000..4b29df5c9 --- /dev/null +++ b/src-new/astrbot_sdk/api/star/context.py @@ -0,0 +1,3 @@ +from ...compat import Context + +__all__ = ["Context"] diff --git a/src-new/astrbot_sdk/cli.py b/src-new/astrbot_sdk/cli.py new file mode 100644 index 000000000..01748a64d --- /dev/null +++ b/src-new/astrbot_sdk/cli.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path + +import click +from loguru import logger + +from .runtime.bootstrap import run_plugin_worker, run_supervisor, run_websocket_server + + +def setup_logger(verbose: bool = False) -> None: + logger.remove() + logger.add( + sys.stderr, + format="{time:HH:mm:ss} | {level: <8} | {message}", + level="DEBUG" if verbose else "INFO", + colorize=True, + ) + + +@click.group() +@click.option("-v", "--verbose", is_flag=True, help="Enable verbose output") +@click.pass_context +def cli(ctx, verbose: bool) -> None: + ctx.ensure_object(dict) + ctx.obj["verbose"] = verbose + setup_logger(verbose) + + +@cli.command() +@click.option( + "--plugins-dir", + default="plugins", + type=click.Path(file_okay=False, dir_okay=True, path_type=Path), + help="Directory containing plugin folders", +) +def run(plugins_dir: Path) -> None: + asyncio.run(run_supervisor(plugins_dir=plugins_dir)) + + +@cli.command(hidden=True) +@click.option( + "--plugin-dir", + required=True, + type=click.Path(file_okay=False, dir_okay=True, path_type=Path), +) +def worker(plugin_dir: Path) -> None: + asyncio.run(run_plugin_worker(plugin_dir=plugin_dir)) + + +@cli.command(hidden=True) +@click.option("--port", default=8765, type=int) +def websocket(port: int) -> None: + asyncio.run(run_websocket_server(port=port)) diff --git a/src-new/astrbot_sdk/clients/__init__.py b/src-new/astrbot_sdk/clients/__init__.py new file mode 100644 index 000000000..0d66078df --- /dev/null +++ b/src-new/astrbot_sdk/clients/__init__.py @@ -0,0 +1,13 @@ +from .db import DBClient +from .llm import ChatMessage, LLMClient, LLMResponse +from .memory import MemoryClient +from .platform import PlatformClient + +__all__ = [ + "ChatMessage", + "DBClient", + "LLMClient", + "LLMResponse", + "MemoryClient", + "PlatformClient", +] diff --git a/src-new/astrbot_sdk/clients/_proxy.py b/src-new/astrbot_sdk/clients/_proxy.py new file mode 100644 index 000000000..3c3f58531 --- /dev/null +++ b/src-new/astrbot_sdk/clients/_proxy.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from typing import Any + +from ..errors import AstrBotError + + +class CapabilityProxy: + def __init__(self, peer) -> None: + self._peer = peer + + def _get_descriptor(self, name: str): + return self._peer.remote_capability_map.get(name) + + def _ensure_available(self, name: str, *, stream: bool) -> None: + descriptor = self._get_descriptor(name) + if descriptor is None: + if self._peer.remote_capability_map: + raise AstrBotError.capability_not_found(name) + return + if stream and not descriptor.supports_stream: + raise AstrBotError.invalid_input(f"{name} 不支持 stream=true") + if not stream and descriptor.supports_stream is False: + return + + async def call(self, name: str, payload: dict[str, Any]) -> dict[str, Any]: + self._ensure_available(name, stream=False) + return await self._peer.invoke(name, payload, stream=False) + + async def stream( + self, + name: str, + payload: dict[str, Any], + ) -> AsyncIterator[dict[str, Any]]: + self._ensure_available(name, stream=True) + event_stream = await self._peer.invoke_stream(name, payload) + async for event in event_stream: + if event.phase == "delta": + yield event.data diff --git a/src-new/astrbot_sdk/clients/db.py b/src-new/astrbot_sdk/clients/db.py new file mode 100644 index 000000000..98e85a967 --- /dev/null +++ b/src-new/astrbot_sdk/clients/db.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from typing import Any + +from ._proxy import CapabilityProxy + + +class DBClient: + def __init__(self, proxy: CapabilityProxy) -> None: + self._proxy = proxy + + async def get(self, key: str) -> dict[str, Any] | None: + output = await self._proxy.call("db.get", {"key": key}) + value = output.get("value") + return value if isinstance(value, dict) else None + + async def set(self, key: str, value: dict[str, Any]) -> None: + await self._proxy.call("db.set", {"key": key, "value": value}) + + async def delete(self, key: str) -> None: + await self._proxy.call("db.delete", {"key": key}) + + async def list(self, prefix: str | None = None) -> list[str]: + output = await self._proxy.call("db.list", {"prefix": prefix}) + return [str(item) for item in output.get("keys", [])] diff --git a/src-new/astrbot_sdk/clients/llm.py b/src-new/astrbot_sdk/clients/llm.py new file mode 100644 index 000000000..ac20c3487 --- /dev/null +++ b/src-new/astrbot_sdk/clients/llm.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +from collections.abc import AsyncGenerator +from typing import Any + +from pydantic import BaseModel, Field + +from ._proxy import CapabilityProxy + + +class ChatMessage(BaseModel): + role: str + content: str + + +class LLMResponse(BaseModel): + text: str + usage: dict[str, Any] | None = None + finish_reason: str | None = None + tool_calls: list[dict[str, Any]] = Field(default_factory=list) + + +class LLMClient: + def __init__(self, proxy: CapabilityProxy) -> None: + self._proxy = proxy + + async def chat( + self, + prompt: str, + *, + system: str | None = None, + history: list[ChatMessage] | None = None, + model: str | None = None, + temperature: float | None = None, + ) -> str: + output = await self._proxy.call( + "llm.chat", + { + "prompt": prompt, + "system": system, + "history": [item.model_dump() for item in history or []], + "model": model, + "temperature": temperature, + }, + ) + return str(output.get("text", "")) + + async def chat_raw( + self, + prompt: str, + **kwargs: Any, + ) -> LLMResponse: + output = await self._proxy.call( + "llm.chat_raw", + { + "prompt": prompt, + **kwargs, + }, + ) + return LLMResponse.model_validate(output) + + async def stream_chat( + self, + prompt: str, + *, + system: str | None = None, + history: list[ChatMessage] | None = None, + ) -> AsyncGenerator[str, None]: + async for data in self._proxy.stream( + "llm.stream_chat", + { + "prompt": prompt, + "system": system, + "history": [item.model_dump() for item in history or []], + }, + ): + yield str(data.get("text", "")) diff --git a/src-new/astrbot_sdk/clients/memory.py b/src-new/astrbot_sdk/clients/memory.py new file mode 100644 index 000000000..f8a5528b8 --- /dev/null +++ b/src-new/astrbot_sdk/clients/memory.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from typing import Any + +from ._proxy import CapabilityProxy + + +class MemoryClient: + def __init__(self, proxy: CapabilityProxy) -> None: + self._proxy = proxy + + async def search(self, query: str) -> list[dict[str, Any]]: + output = await self._proxy.call("memory.search", {"query": query}) + return list(output.get("items", [])) + + async def save(self, key: str, value: dict[str, Any]) -> None: + await self._proxy.call("memory.save", {"key": key, "value": value}) + + async def delete(self, key: str) -> None: + await self._proxy.call("memory.delete", {"key": key}) diff --git a/src-new/astrbot_sdk/clients/platform.py b/src-new/astrbot_sdk/clients/platform.py new file mode 100644 index 000000000..d5c5b2b75 --- /dev/null +++ b/src-new/astrbot_sdk/clients/platform.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from typing import Any + +from ._proxy import CapabilityProxy + + +class PlatformClient: + def __init__(self, proxy: CapabilityProxy) -> None: + self._proxy = proxy + + async def send(self, session: str, text: str) -> dict[str, Any]: + return await self._proxy.call( + "platform.send", + {"session": session, "text": text}, + ) + + async def send_image(self, session: str, image_url: str) -> dict[str, Any]: + return await self._proxy.call( + "platform.send_image", + {"session": session, "image_url": image_url}, + ) + + async def get_members(self, session: str) -> list[dict[str, Any]]: + output = await self._proxy.call("platform.get_members", {"session": session}) + return list(output.get("members", [])) diff --git a/src-new/astrbot_sdk/compat.py b/src-new/astrbot_sdk/compat.py new file mode 100644 index 000000000..3d70f3141 --- /dev/null +++ b/src-new/astrbot_sdk/compat.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +from collections import defaultdict +from typing import Any + +from loguru import logger + +from .clients.llm import LLMResponse +from .context import Context as NewContext + +_warned_methods: set[str] = set() + + +def _warn_once(old_name: str, replacement: str) -> None: + if old_name in _warned_methods: + return + _warned_methods.add(old_name) + logger.warning( + "[AstrBot] 警告：{} 已过时。请替换为：{}", + old_name, + replacement, + ) + + +class LegacyConversationManager: + def __init__(self, parent: "LegacyContext") -> None: + self._parent = parent + self._counters: defaultdict[str, int] = defaultdict(int) + + def _ctx(self) -> NewContext: + return self._parent.require_runtime_context() + + async def new_conversation( + self, + unified_msg_origin: str, + platform_id: str | None = None, + content: list[dict] | None = None, + title: str | None = None, + persona_id: str | None = None, + ) -> str: + ctx = self._ctx() + self._counters[unified_msg_origin] += 1 + conversation_id = f"{ctx.plugin_id}-conv-{self._counters[unified_msg_origin]}" + stored = await ctx.db.get("__compat_conversations__") or {} + stored[conversation_id] = { + "unified_msg_origin": unified_msg_origin, + "platform_id": platform_id, + "content": content or [], + "title": title, + "persona_id": persona_id, + } + await ctx.db.set("__compat_conversations__", stored) + return conversation_id + + +class LegacyContext: + def __init__(self, plugin_id: str) -> None: + self.plugin_id = plugin_id + self._runtime_context: NewContext | None = None + self.conversation_manager = LegacyConversationManager(self) + + def bind_runtime_context(self, runtime_context: NewContext) -> None: + self._runtime_context = runtime_context + + def require_runtime_context(self) -> NewContext: + if self._runtime_context is None: + raise RuntimeError("LegacyContext 尚未绑定运行时 Context") + return self._runtime_context + + async def llm_generate( + self, + chat_provider_id: str, + prompt: str | None = None, + image_urls: list[str] | None = None, + tools: Any | None = None, + system_prompt: str | None = None, + contexts: list[dict] | None = None, + **kwargs: Any, + ) -> LLMResponse: + _warn_once("context.llm_generate()", "ctx.llm.chat(prompt)") + ctx = self.require_runtime_context() + return await ctx.llm.chat_raw( + prompt or "", + system=system_prompt, + history=contexts or [], + image_urls=image_urls or [], + tools=tools, + **kwargs, + ) + + async def tool_loop_agent( + self, + chat_provider_id: str, + prompt: str | None = None, + image_urls: list[str] | None = None, + tools: Any | None = None, + system_prompt: str | None = None, + contexts: list[dict] | None = None, + max_steps: int = 30, + **kwargs: Any, + ) -> LLMResponse: + _warn_once("context.tool_loop_agent()", "ctx.llm.chat_raw(...)") + ctx = self.require_runtime_context() + return await ctx.llm.chat_raw( + prompt or "", + system=system_prompt, + history=contexts or [], + image_urls=image_urls or [], + tools=tools, + max_steps=max_steps, + **kwargs, + ) + + async def send_message(self, session: str, message_chain: Any) -> None: + _warn_once("context.send_message()", "ctx.platform.send(session, text)") + ctx = self.require_runtime_context() + await ctx.platform.send(session, str(message_chain)) + + async def add_llm_tools(self, *tools: Any) -> None: + _warn_once("context.add_llm_tools()", "ctx.llm.chat_raw(..., tools=...)") + return None + + async def put_kv_data(self, key: str, value: dict[str, Any]) -> None: + _warn_once("context.put_kv_data()", "ctx.db.set(key, value)") + ctx = self.require_runtime_context() + await ctx.db.set(key, value) + + async def get_kv_data(self, key: str) -> dict[str, Any] | None: + _warn_once("context.get_kv_data()", "ctx.db.get(key)") + ctx = self.require_runtime_context() + return await ctx.db.get(key) + + async def delete_kv_data(self, key: str) -> None: + _warn_once("context.delete_kv_data()", "ctx.db.delete(key)") + ctx = self.require_runtime_context() + await ctx.db.delete(key) + + +class CommandComponent: + pass + + +Context = LegacyContext diff --git a/src-new/astrbot_sdk/context.py b/src-new/astrbot_sdk/context.py new file mode 100644 index 000000000..216caec23 --- /dev/null +++ b/src-new/astrbot_sdk/context.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Any + +from loguru import logger as base_logger + +from .clients import DBClient, LLMClient, MemoryClient, PlatformClient +from .clients._proxy import CapabilityProxy + + +@dataclass(slots=True) +class CancelToken: + _cancelled: asyncio.Event + + def __init__(self) -> None: + self._cancelled = asyncio.Event() + + def cancel(self) -> None: + self._cancelled.set() + + @property + def cancelled(self) -> bool: + return self._cancelled.is_set() + + async def wait(self) -> None: + await self._cancelled.wait() + + def raise_if_cancelled(self) -> None: + if self.cancelled: + raise asyncio.CancelledError + + +class Context: + def __init__( + self, + *, + peer, + plugin_id: str, + cancel_token: CancelToken | None = None, + logger: Any | None = None, + ) -> None: + proxy = CapabilityProxy(peer) + self.llm = LLMClient(proxy) + self.memory = MemoryClient(proxy) + self.db = DBClient(proxy) + self.platform = PlatformClient(proxy) + self.plugin_id = plugin_id + self.logger = logger or base_logger.bind(plugin_id=plugin_id) + self.cancel_token = cancel_token or CancelToken() diff --git a/src-new/astrbot_sdk/decorators.py b/src-new/astrbot_sdk/decorators.py new file mode 100644 index 000000000..58e59f640 --- /dev/null +++ b/src-new/astrbot_sdk/decorators.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable + +from .protocol.descriptors import ( + CommandTrigger, + EventTrigger, + MessageTrigger, + Permissions, + ScheduleTrigger, +) + +HandlerCallable = Callable[..., Any] +HANDLER_META_ATTR = "__astrbot_handler_meta__" + + +@dataclass(slots=True) +class HandlerMeta: + trigger: CommandTrigger | MessageTrigger | EventTrigger | ScheduleTrigger | None = None + priority: int = 0 + permissions: Permissions = field(default_factory=Permissions) + + +def _get_or_create_meta(func: HandlerCallable) -> HandlerMeta: + meta = getattr(func, HANDLER_META_ATTR, None) + if meta is None: + meta = HandlerMeta() + setattr(func, HANDLER_META_ATTR, meta) + return meta + + +def get_handler_meta(func: HandlerCallable) -> HandlerMeta | None: + return getattr(func, HANDLER_META_ATTR, None) + + +def on_command( + command: str, + *, + aliases: list[str] | None = None, + description: str | None = None, +) -> Callable[[HandlerCallable], HandlerCallable]: + def decorator(func: HandlerCallable) -> HandlerCallable: + meta = _get_or_create_meta(func) + meta.trigger = CommandTrigger( + command=command, + aliases=aliases or [], + description=description, + ) + return func + + return decorator + + +def on_message( + *, + regex: str | None = None, + keywords: list[str] | None = None, + platforms: list[str] | None = None, +) -> Callable[[HandlerCallable], HandlerCallable]: + def decorator(func: HandlerCallable) -> HandlerCallable: + meta = _get_or_create_meta(func) + meta.trigger = MessageTrigger( + regex=regex, + keywords=keywords or [], + platforms=platforms or [], + ) + return func + + return decorator + + +def on_event(event_type: str) -> Callable[[HandlerCallable], HandlerCallable]: + def decorator(func: HandlerCallable) -> HandlerCallable: + meta = _get_or_create_meta(func) + meta.trigger = EventTrigger(event_type=event_type) + return func + + return decorator + + +def on_schedule( + *, + cron: str | None = None, + interval_seconds: int | None = None, +) -> Callable[[HandlerCallable], HandlerCallable]: + def decorator(func: HandlerCallable) -> HandlerCallable: + meta = _get_or_create_meta(func) + meta.trigger = ScheduleTrigger(cron=cron, interval_seconds=interval_seconds) + return func + + return decorator + + +def require_admin(func: HandlerCallable) -> HandlerCallable: + meta = _get_or_create_meta(func) + meta.permissions.require_admin = True + return func diff --git a/src-new/astrbot_sdk/errors.py b/src-new/astrbot_sdk/errors.py new file mode 100644 index 000000000..a8629d722 --- /dev/null +++ b/src-new/astrbot_sdk/errors.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(slots=True) +class AstrBotError(Exception): + code: str + message: str + hint: str = "" + retryable: bool = False + + def __str__(self) -> str: + return self.message + + @classmethod + def cancelled(cls, message: str = "调用被取消") -> "AstrBotError": + return cls( + code="cancelled", + message=message, + hint="", + retryable=False, + ) + + @classmethod + def capability_not_found(cls, name: str) -> "AstrBotError": + return cls( + code="capability_not_found", + message=f"未找到能力：{name}", + hint="请确认 AstrBot Core 是否已注册该 capability", + retryable=False, + ) + + @classmethod + def invalid_input(cls, message: str) -> "AstrBotError": + return cls( + code="invalid_input", + message=message, + hint="请检查调用参数", + retryable=False, + ) + + @classmethod + def protocol_version_mismatch(cls, message: str) -> "AstrBotError": + return cls( + code="protocol_version_mismatch", + message=message, + hint="请升级 astrbot_sdk 至最新版本", + retryable=False, + ) + + @classmethod + def protocol_error(cls, message: str) -> "AstrBotError": + return cls( + code="protocol_error", + message=message, + hint="请检查通信双方的协议实现", + retryable=False, + ) + + @classmethod + def internal_error(cls, message: str) -> "AstrBotError": + return cls( + code="internal_error", + message=message, + hint="请联系插件作者", + retryable=False, + ) + + def to_payload(self) -> dict[str, object]: + return { + "code": self.code, + "message": self.message, + "hint": self.hint, + "retryable": self.retryable, + } + + @classmethod + def from_payload(cls, payload: dict[str, object]) -> "AstrBotError": + return cls( + code=str(payload.get("code", "unknown_error")), + message=str(payload.get("message", "未知错误")), + hint=str(payload.get("hint", "")), + retryable=bool(payload.get("retryable", False)), + ) diff --git a/src-new/astrbot_sdk/events.py b/src-new/astrbot_sdk/events.py new file mode 100644 index 000000000..9965b7188 --- /dev/null +++ b/src-new/astrbot_sdk/events.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .context import Context + + +@dataclass(slots=True) +class PlainTextResult: + text: str + + +class MessageEvent: + def __init__( + self, + *, + text: str = "", + user_id: str | None = None, + group_id: str | None = None, + platform: str | None = None, + session_id: str | None = None, + raw: dict[str, Any] | None = None, + context: "Context | None" = None, + ) -> None: + self.text = text + self.user_id = user_id + self.group_id = group_id + self.platform = platform + self.session_id = session_id or group_id or user_id or "" + self.raw = raw or {} + self._context = context + + @classmethod + def from_payload( + cls, + payload: dict[str, Any], + *, + context: "Context | None" = None, + ) -> "MessageEvent": + return cls( + text=str(payload.get("text", "")), + user_id=payload.get("user_id"), + group_id=payload.get("group_id"), + platform=payload.get("platform"), + session_id=payload.get("session_id"), + raw=payload, + context=context, + ) + + def to_payload(self) -> dict[str, Any]: + return { + "text": self.text, + "user_id": self.user_id, + "group_id": self.group_id, + "platform": self.platform, + "session_id": self.session_id, + } + + async def reply(self, text: str) -> None: + if self._context is None: + raise RuntimeError("MessageEvent 未绑定 Context，无法 reply") + await self._context.platform.send(self.session_id, text) + + def plain_result(self, text: str) -> PlainTextResult: + return PlainTextResult(text=text) diff --git a/src-new/astrbot_sdk/protocol/__init__.py b/src-new/astrbot_sdk/protocol/__init__.py new file mode 100644 index 000000000..882adb401 --- /dev/null +++ b/src-new/astrbot_sdk/protocol/__init__.py @@ -0,0 +1,21 @@ +from .descriptors import CapabilityDescriptor, HandlerDescriptor, Permissions +from .messages import ( + CancelMessage, + EventMessage, + InitializeMessage, + InitializeOutput, + InvokeMessage, + ResultMessage, +) + +__all__ = [ + "CapabilityDescriptor", + "CancelMessage", + "EventMessage", + "HandlerDescriptor", + "InitializeMessage", + "InitializeOutput", + "InvokeMessage", + "Permissions", + "ResultMessage", +] diff --git a/src-new/astrbot_sdk/protocol/descriptors.py b/src-new/astrbot_sdk/protocol/descriptors.py new file mode 100644 index 000000000..1481c5132 --- /dev/null +++ b/src-new/astrbot_sdk/protocol/descriptors.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from typing import Annotated, Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class _DescriptorBase(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class Permissions(_DescriptorBase): + require_admin: bool = False + level: int = 0 + + +class CommandTrigger(_DescriptorBase): + type: Literal["command"] = "command" + command: str + aliases: list[str] = Field(default_factory=list) + description: str | None = None + + +class MessageTrigger(_DescriptorBase): + type: Literal["message"] = "message" + regex: str | None = None + keywords: list[str] = Field(default_factory=list) + platforms: list[str] = Field(default_factory=list) + + +class EventTrigger(_DescriptorBase): + type: Literal["event"] = "event" + event_type: str + + +class ScheduleTrigger(_DescriptorBase): + type: Literal["schedule"] = "schedule" + cron: str | None = None + interval_seconds: int | None = None + + @model_validator(mode="after") + def validate_schedule(self) -> "ScheduleTrigger": + has_cron = self.cron is not None + has_interval = self.interval_seconds is not None + if has_cron == has_interval: + raise ValueError("cron 和 interval_seconds 必须且只能有一个非 null") + return self + + +Trigger = Annotated[ + CommandTrigger | MessageTrigger | EventTrigger | ScheduleTrigger, + Field(discriminator="type"), +] + + +class HandlerDescriptor(_DescriptorBase): + id: str + trigger: Trigger + priority: int = 0 + permissions: Permissions = Field(default_factory=Permissions) + + +class CapabilityDescriptor(_DescriptorBase): + name: str + description: str + input_schema: dict[str, Any] | None = None + output_schema: dict[str, Any] | None = None + supports_stream: bool = False + cancelable: bool = False diff --git a/src-new/astrbot_sdk/protocol/legacy_adapter.py b/src-new/astrbot_sdk/protocol/legacy_adapter.py new file mode 100644 index 000000000..de63424bc --- /dev/null +++ b/src-new/astrbot_sdk/protocol/legacy_adapter.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from typing import Any + +from .messages import CancelMessage, EventMessage, InvokeMessage, ResultMessage + + +def legacy_request_to_invoke(payload: dict[str, Any]) -> InvokeMessage: + method = str(payload.get("method", "")) + params = payload.get("params") or {} + request_id = str(payload.get("id", "")) + if method == "call_handler": + return InvokeMessage( + id=request_id, + capability="handler.invoke", + input={ + "handler_id": params.get("handler_full_name", ""), + "event": params.get("event", {}), + "args": params.get("args", {}), + }, + stream=False, + ) + return InvokeMessage( + id=request_id, + capability=method, + input=params if isinstance(params, dict) else {}, + stream=False, + ) + + +def invoke_to_legacy_request(message: InvokeMessage) -> dict[str, Any]: + if message.capability == "handler.invoke": + return { + "jsonrpc": "2.0", + "id": message.id, + "method": "call_handler", + "params": { + "handler_full_name": message.input.get("handler_id"), + "event": message.input.get("event", {}), + "args": message.input.get("args", {}), + }, + } + return { + "jsonrpc": "2.0", + "id": message.id, + "method": message.capability, + "params": message.input, + } + + +def result_to_legacy_response(message: ResultMessage) -> dict[str, Any]: + if message.success: + return { + "jsonrpc": "2.0", + "id": message.id, + "result": message.output, + } + return { + "jsonrpc": "2.0", + "id": message.id, + "error": { + "code": -32000, + "message": message.error.message if message.error else "unknown error", + "data": message.error.model_dump() if message.error else None, + }, + } + + +def event_to_legacy_notification(message: EventMessage) -> dict[str, Any]: + method = { + "started": "handler_stream_start", + "delta": "handler_stream_update", + "completed": "handler_stream_end", + "failed": "handler_stream_end", + }[message.phase] + params: dict[str, Any] = {"id": message.id} + if message.phase == "delta": + params["data"] = message.data + if message.phase == "failed" and message.error is not None: + params["error"] = message.error.model_dump() + return {"jsonrpc": "2.0", "method": method, "params": params} + + +def cancel_to_legacy_request(message: CancelMessage) -> dict[str, Any]: + return { + "jsonrpc": "2.0", + "id": message.id, + "method": "cancel", + "params": {"reason": message.reason}, + } diff --git a/src-new/astrbot_sdk/protocol/messages.py b/src-new/astrbot_sdk/protocol/messages.py new file mode 100644 index 000000000..1f25775cd --- /dev/null +++ b/src-new/astrbot_sdk/protocol/messages.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import json +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field + +from .descriptors import CapabilityDescriptor, HandlerDescriptor + + +class _MessageBase(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class ErrorPayload(_MessageBase): + code: str + message: str + hint: str = "" + retryable: bool = False + + +class PeerInfo(_MessageBase): + name: str + role: Literal["plugin", "core"] + version: str | None = None + + +class InitializeMessage(_MessageBase): + type: Literal["initialize"] = "initialize" + id: str + protocol_version: str + peer: PeerInfo + handlers: list[HandlerDescriptor] = Field(default_factory=list) + metadata: dict[str, Any] = Field(default_factory=dict) + + +class InitializeOutput(_MessageBase): + peer: PeerInfo + capabilities: list[CapabilityDescriptor] = Field(default_factory=list) + metadata: dict[str, Any] = Field(default_factory=dict) + + +class ResultMessage(_MessageBase): + type: Literal["result"] = "result" + id: str + kind: str | None = None + success: bool + output: dict[str, Any] = Field(default_factory=dict) + error: ErrorPayload | None = None + + +class InvokeMessage(_MessageBase): + type: Literal["invoke"] = "invoke" + id: str + capability: str + input: dict[str, Any] = Field(default_factory=dict) + stream: bool = False + + +class EventMessage(_MessageBase): + type: Literal["event"] = "event" + id: str + phase: Literal["started", "delta", "completed", "failed"] + data: dict[str, Any] = Field(default_factory=dict) + output: dict[str, Any] = Field(default_factory=dict) + error: ErrorPayload | None = None + + +class CancelMessage(_MessageBase): + type: Literal["cancel"] = "cancel" + id: str + reason: str = "user_cancelled" + + +ProtocolMessage = ( + InitializeMessage | ResultMessage | InvokeMessage | EventMessage | CancelMessage +) + + +def parse_message(payload: str | bytes | dict[str, Any]) -> ProtocolMessage: + if isinstance(payload, bytes): + payload = payload.decode("utf-8") + if isinstance(payload, str): + payload = json.loads(payload) + message_type = payload.get("type") + if message_type == "initialize": + return InitializeMessage.model_validate(payload) + if message_type == "result": + return ResultMessage.model_validate(payload) + if message_type == "invoke": + return InvokeMessage.model_validate(payload) + if message_type == "event": + return EventMessage.model_validate(payload) + if message_type == "cancel": + return CancelMessage.model_validate(payload) + raise ValueError(f"未知消息类型：{message_type}") diff --git a/src-new/astrbot_sdk/runtime/__init__.py b/src-new/astrbot_sdk/runtime/__init__.py new file mode 100644 index 000000000..c9c2ef67b --- /dev/null +++ b/src-new/astrbot_sdk/runtime/__init__.py @@ -0,0 +1 @@ +__all__: list[str] = [] diff --git a/src-new/astrbot_sdk/runtime/bootstrap.py b/src-new/astrbot_sdk/runtime/bootstrap.py new file mode 100644 index 000000000..08b2204ec --- /dev/null +++ b/src-new/astrbot_sdk/runtime/bootstrap.py @@ -0,0 +1,352 @@ +from __future__ import annotations + +import asyncio +import os +import signal +import sys +from pathlib import Path +from typing import IO, Any + +from loguru import logger + +from ..errors import AstrBotError +from ..protocol.messages import InitializeOutput, PeerInfo +from .capability_router import CapabilityRouter +from .handler_dispatcher import HandlerDispatcher +from .loader import ( + PluginEnvironmentManager, + PluginSpec, + discover_plugins, + load_plugin, + load_plugin_spec, +) +from .peer import Peer +from .transport import StdioTransport, WebSocketServerTransport + + +def _install_signal_handlers(stop_event: asyncio.Event) -> None: + loop = asyncio.get_running_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + try: + loop.add_signal_handler(sig, stop_event.set) + except NotImplementedError: + logger.debug("Signal handlers are not supported for {}", sig) + + +def _prepare_stdio_transport( + stdin: IO[str] | None, + stdout: IO[str] | None, +) -> tuple[IO[str], IO[str], IO[str] | None]: + if stdin is not None and stdout is not None: + return stdin, stdout, None + transport_stdin = stdin or sys.stdin + transport_stdout = stdout or sys.stdout + original_stdout = sys.stdout + sys.stdout = sys.stderr + return transport_stdin, transport_stdout, original_stdout + + +async def _wait_for_shutdown(peer: Peer, stop_event: asyncio.Event) -> None: + stop_waiter = asyncio.create_task(stop_event.wait()) + transport_waiter = asyncio.create_task(peer.wait_closed()) + done, pending = await asyncio.wait( + {stop_waiter, transport_waiter}, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + for task in done: + if not task.cancelled(): + task.result() + + +class WorkerSession: + def __init__( + self, + *, + plugin: PluginSpec, + repo_root: Path, + env_manager: PluginEnvironmentManager, + capability_router: CapabilityRouter, + ) -> None: + self.plugin = plugin + self.repo_root = repo_root.resolve() + self.env_manager = env_manager + self.capability_router = capability_router + self.peer: Peer | None = None + self.handlers = [] + + async def start(self) -> None: + python_path = self.env_manager.prepare_environment(self.plugin) + repo_src_dir = str(self.repo_root / "src-new") + env = os.environ.copy() + existing_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + f"{repo_src_dir}{os.pathsep}{existing_pythonpath}" + if existing_pythonpath + else repo_src_dir + ) + + transport = StdioTransport( + command=[ + str(python_path), + "-m", + "astrbot_sdk", + "worker", + "--plugin-dir", + str(self.plugin.plugin_dir), + ], + cwd=str(self.plugin.plugin_dir), + env=env, + ) + self.peer = Peer( + transport=transport, + peer_info=PeerInfo(name="astrbot-core", role="core", version="v4"), + ) + self.peer.set_initialize_handler(self._handle_initialize) + self.peer.set_invoke_handler(self._handle_capability_invoke) + try: + await self.peer.start() + await self.peer.wait_until_remote_initialized() + self.handlers = list(self.peer.remote_handlers) + except Exception: + await self.stop() + raise + + async def stop(self) -> None: + if self.peer is not None: + await self.peer.stop() + + async def invoke_handler( + self, + handler_id: str, + event_payload: dict[str, Any], + *, + request_id: str, + ) -> dict[str, Any]: + if self.peer is None: + raise RuntimeError("worker session is not running") + return await self.peer.invoke( + "handler.invoke", + { + "handler_id": handler_id, + "event": event_payload, + }, + request_id=request_id, + ) + + async def cancel(self, request_id: str) -> None: + if self.peer is None: + return + await self.peer.cancel(request_id) + + async def _handle_initialize(self, _message) -> InitializeOutput: + return InitializeOutput( + peer=PeerInfo(name="astrbot-supervisor", role="core", version="v4"), + capabilities=self.capability_router.descriptors(), + metadata={"plugin": self.plugin.name}, + ) + + async def _handle_capability_invoke(self, message, cancel_token): + return await self.capability_router.execute( + message.capability, + message.input, + stream=message.stream, + cancel_token=cancel_token, + ) + + +class SupervisorRuntime: + def __init__( + self, + *, + transport, + plugins_dir: Path, + env_manager: PluginEnvironmentManager | None = None, + ) -> None: + self.transport = transport + self.plugins_dir = plugins_dir.resolve() + self.repo_root = Path(__file__).resolve().parents[3] + self.env_manager = env_manager or PluginEnvironmentManager(self.repo_root) + self.capability_router = CapabilityRouter() + self.peer = Peer( + transport=self.transport, + peer_info=PeerInfo(name="astrbot-supervisor", role="plugin", version="v4"), + ) + self.peer.set_invoke_handler(self._handle_upstream_invoke) + self.peer.set_cancel_handler(self._handle_upstream_cancel) + self.worker_sessions: dict[str, WorkerSession] = {} + self.handler_to_worker: dict[str, WorkerSession] = {} + self.active_requests: dict[str, WorkerSession] = {} + self.loaded_plugins: list[str] = [] + self.skipped_plugins: dict[str, str] = {} + + async def start(self) -> None: + discovery = discover_plugins(self.plugins_dir) + self.skipped_plugins = dict(discovery.skipped_plugins) + try: + for plugin in discovery.plugins: + session = WorkerSession( + plugin=plugin, + repo_root=self.repo_root, + env_manager=self.env_manager, + capability_router=self.capability_router, + ) + try: + await session.start() + except Exception as exc: + self.skipped_plugins[plugin.name] = str(exc) + await session.stop() + continue + self.worker_sessions[plugin.name] = session + self.loaded_plugins.append(plugin.name) + for handler in session.handlers: + self.handler_to_worker[handler.id] = session + + aggregated_handlers = list(self.handler_to_worker.keys()) + logger.info("Loaded plugins: {}", ", ".join(sorted(self.loaded_plugins)) or "none") + + await self.peer.start() + await self.peer.initialize( + [handler for session in self.worker_sessions.values() for handler in session.handlers], + metadata={ + "plugins": sorted(self.loaded_plugins), + "skipped_plugins": self.skipped_plugins, + "aggregated_handler_ids": aggregated_handlers, + }, + ) + except Exception: + await self.stop() + raise + + async def stop(self) -> None: + for session in list(self.worker_sessions.values()): + await session.stop() + await self.peer.stop() + + async def _handle_upstream_invoke(self, message, _cancel_token): + if message.capability != "handler.invoke": + raise AstrBotError.capability_not_found(message.capability) + handler_id = str(message.input.get("handler_id", "")) + session = self.handler_to_worker.get(handler_id) + if session is None: + raise AstrBotError.invalid_input(f"handler not found: {handler_id}") + self.active_requests[message.id] = session + try: + return await session.invoke_handler( + handler_id, + message.input.get("event", {}), + request_id=message.id, + ) + finally: + self.active_requests.pop(message.id, None) + + async def _handle_upstream_cancel(self, request_id: str) -> None: + session = self.active_requests.get(request_id) + if session is not None: + await session.cancel(request_id) + + +class PluginWorkerRuntime: + def __init__(self, *, plugin_dir: Path, transport) -> None: + self.plugin = load_plugin_spec(plugin_dir) + self.transport = transport + self.loaded_plugin = load_plugin(self.plugin) + self.peer = Peer( + transport=self.transport, + peer_info=PeerInfo(name=self.plugin.name, role="plugin", version="v4"), + ) + self.dispatcher = HandlerDispatcher( + plugin_id=self.plugin.name, + peer=self.peer, + handlers=self.loaded_plugin.handlers, + ) + self.peer.set_invoke_handler(self._handle_invoke) + self.peer.set_cancel_handler(self.dispatcher.cancel) + + async def start(self) -> None: + await self.peer.start() + await self.peer.initialize( + [item.descriptor for item in self.loaded_plugin.handlers], + metadata={"plugin_id": self.plugin.name}, + ) + + async def stop(self) -> None: + await self.peer.stop() + + async def _handle_invoke(self, message, cancel_token): + if message.capability != "handler.invoke": + raise AstrBotError.capability_not_found(message.capability) + return await self.dispatcher.invoke(message, cancel_token) + + +async def run_supervisor( + *, + plugins_dir: Path = Path("plugins"), + stdin: IO[str] | None = None, + stdout: IO[str] | None = None, + env_manager: PluginEnvironmentManager | None = None, +) -> None: + transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport( + stdin, + stdout, + ) + transport = StdioTransport(stdin=transport_stdin, stdout=transport_stdout) + runtime = SupervisorRuntime( + transport=transport, + plugins_dir=plugins_dir, + env_manager=env_manager, + ) + + try: + await runtime.start() + stop_event = asyncio.Event() + _install_signal_handlers(stop_event) + await _wait_for_shutdown(runtime.peer, stop_event) + finally: + await runtime.stop() + if original_stdout is not None: + sys.stdout = original_stdout + + +async def run_plugin_worker( + *, + plugin_dir: Path, + stdin: IO[str] | None = None, + stdout: IO[str] | None = None, +) -> None: + transport_stdin, transport_stdout, original_stdout = _prepare_stdio_transport( + stdin, + stdout, + ) + transport = StdioTransport(stdin=transport_stdin, stdout=transport_stdout) + runtime = PluginWorkerRuntime(plugin_dir=plugin_dir, transport=transport) + try: + await runtime.start() + stop_event = asyncio.Event() + _install_signal_handlers(stop_event) + await _wait_for_shutdown(runtime.peer, stop_event) + finally: + await runtime.stop() + if original_stdout is not None: + sys.stdout = original_stdout + + +async def run_websocket_server( + *, + host: str = "127.0.0.1", + port: int = 8765, + path: str = "/", + plugin_dir: Path | None = None, +) -> None: + runtime = PluginWorkerRuntime( + plugin_dir=plugin_dir or Path.cwd(), + transport=WebSocketServerTransport(host=host, port=port, path=path), + ) + try: + await runtime.start() + stop_event = asyncio.Event() + _install_signal_handlers(stop_event) + await _wait_for_shutdown(runtime.peer, stop_event) + finally: + await runtime.stop() diff --git a/src-new/astrbot_sdk/runtime/capability_router.py b/src-new/astrbot_sdk/runtime/capability_router.py new file mode 100644 index 000000000..70782b766 --- /dev/null +++ b/src-new/astrbot_sdk/runtime/capability_router.py @@ -0,0 +1,330 @@ +from __future__ import annotations + +import asyncio +import json +from collections.abc import AsyncIterator, Awaitable, Callable +from dataclasses import dataclass +from typing import Any + +from ..errors import AstrBotError +from ..protocol.descriptors import CapabilityDescriptor + +CallHandler = Callable[[dict[str, Any], object], Awaitable[dict[str, Any]]] +StreamHandler = Callable[[dict[str, Any], object], AsyncIterator[dict[str, Any]]] +FinalizeHandler = Callable[[list[dict[str, Any]]], dict[str, Any]] + + +@dataclass(slots=True) +class StreamExecution: + iterator: AsyncIterator[dict[str, Any]] + finalize: FinalizeHandler + + +@dataclass(slots=True) +class _CapabilityRegistration: + descriptor: CapabilityDescriptor + call_handler: CallHandler | None = None + stream_handler: StreamHandler | None = None + finalize: FinalizeHandler | None = None + + +class CapabilityRouter: + def __init__(self) -> None: + self._registrations: dict[str, _CapabilityRegistration] = {} + self.db_store: dict[str, dict[str, Any]] = {} + self.memory_store: dict[str, dict[str, Any]] = {} + self.sent_messages: list[dict[str, Any]] = [] + self._register_builtin_capabilities() + + def descriptors(self) -> list[CapabilityDescriptor]: + return [entry.descriptor for entry in self._registrations.values()] + + def register( + self, + descriptor: CapabilityDescriptor, + *, + call_handler: CallHandler | None = None, + stream_handler: StreamHandler | None = None, + finalize: FinalizeHandler | None = None, + ) -> None: + self._registrations[descriptor.name] = _CapabilityRegistration( + descriptor=descriptor, + call_handler=call_handler, + stream_handler=stream_handler, + finalize=finalize, + ) + + async def execute( + self, + capability: str, + payload: dict[str, Any], + *, + stream: bool, + cancel_token, + ) -> dict[str, Any] | StreamExecution: + registration = self._registrations.get(capability) + if registration is None: + raise AstrBotError.capability_not_found(capability) + + self._validate_schema(registration.descriptor.input_schema, payload) + if stream: + if registration.stream_handler is None: + raise AstrBotError.invalid_input(f"{capability} 不支持 stream=true") + finalize = registration.finalize or (lambda chunks: {"items": chunks}) + return StreamExecution( + iterator=registration.stream_handler(payload, cancel_token), + finalize=finalize, + ) + + if registration.call_handler is None: + raise AstrBotError.invalid_input(f"{capability} 只能以 stream=true 调用") + output = await registration.call_handler(payload, cancel_token) + self._validate_schema(registration.descriptor.output_schema, output) + return output + + def _register_builtin_capabilities(self) -> None: + def obj_schema(required: list[str], **properties: Any) -> dict[str, Any]: + return { + "type": "object", + "properties": properties, + "required": required, + } + + async def llm_chat(payload: dict[str, Any], _token) -> dict[str, Any]: + prompt = str(payload.get("prompt", "")) + return {"text": f"Echo: {prompt}"} + + async def llm_chat_raw(payload: dict[str, Any], _token) -> dict[str, Any]: + prompt = str(payload.get("prompt", "")) + text = f"Echo: {prompt}" + return { + "text": text, + "usage": { + "input_tokens": len(prompt), + "output_tokens": len(text), + }, + "finish_reason": "stop", + "tool_calls": [], + } + + async def llm_stream( + payload: dict[str, Any], + token, + ) -> AsyncIterator[dict[str, Any]]: + text = f"Echo: {str(payload.get('prompt', ''))}" + for char in text: + token.raise_if_cancelled() + await asyncio.sleep(0) + yield {"text": char} + + async def memory_search(payload: dict[str, Any], _token) -> dict[str, Any]: + query = str(payload.get("query", "")) + items = [ + {"key": key, "value": value} + for key, value in self.memory_store.items() + if query in key or query in json.dumps(value, ensure_ascii=False) + ] + return {"items": items} + + async def memory_save(payload: dict[str, Any], _token) -> dict[str, Any]: + key = str(payload.get("key", "")) + value = payload.get("value") + if not isinstance(value, dict): + raise AstrBotError.invalid_input("memory.save 的 value 必须是 object") + self.memory_store[key] = value + return {} + + async def memory_delete(payload: dict[str, Any], _token) -> dict[str, Any]: + self.memory_store.pop(str(payload.get("key", "")), None) + return {} + + async def db_get(payload: dict[str, Any], _token) -> dict[str, Any]: + return {"value": self.db_store.get(str(payload.get("key", "")))} + + async def db_set(payload: dict[str, Any], _token) -> dict[str, Any]: + key = str(payload.get("key", "")) + value = payload.get("value") + if not isinstance(value, dict): + raise AstrBotError.invalid_input("db.set 的 value 必须是 object") + self.db_store[key] = value + return {} + + async def db_delete(payload: dict[str, Any], _token) -> dict[str, Any]: + self.db_store.pop(str(payload.get("key", "")), None) + return {} + + async def db_list(payload: dict[str, Any], _token) -> dict[str, Any]: + prefix = payload.get("prefix") + keys = sorted(self.db_store.keys()) + if isinstance(prefix, str): + keys = [item for item in keys if item.startswith(prefix)] + return {"keys": keys} + + async def platform_send(payload: dict[str, Any], _token) -> dict[str, Any]: + session = str(payload.get("session", "")) + text = str(payload.get("text", "")) + message_id = f"msg_{len(self.sent_messages) + 1}" + self.sent_messages.append( + { + "message_id": message_id, + "session": session, + "text": text, + } + ) + return {"message_id": message_id} + + async def platform_send_image(payload: dict[str, Any], _token) -> dict[str, Any]: + session = str(payload.get("session", "")) + image_url = str(payload.get("image_url", "")) + message_id = f"img_{len(self.sent_messages) + 1}" + self.sent_messages.append( + { + "message_id": message_id, + "session": session, + "image_url": image_url, + } + ) + return {"message_id": message_id} + + async def platform_get_members(payload: dict[str, Any], _token) -> dict[str, Any]: + session = str(payload.get("session", "")) + return { + "members": [ + {"user_id": f"{session}:member-1", "nickname": "Member 1"}, + {"user_id": f"{session}:member-2", "nickname": "Member 2"}, + ] + } + + self.register( + CapabilityDescriptor( + name="llm.chat", + description="发送对话请求，返回文本", + input_schema=obj_schema(["prompt"], prompt={"type": "string"}), + output_schema=obj_schema(["text"], text={"type": "string"}), + ), + call_handler=llm_chat, + ) + self.register( + CapabilityDescriptor( + name="llm.chat_raw", + description="发送对话请求，返回完整响应", + input_schema=obj_schema(["prompt"], prompt={"type": "string"}), + output_schema=obj_schema(["text"], text={"type": "string"}), + ), + call_handler=llm_chat_raw, + ) + self.register( + CapabilityDescriptor( + name="llm.stream_chat", + description="流式对话", + input_schema=obj_schema(["prompt"], prompt={"type": "string"}), + output_schema=obj_schema(["text"], text={"type": "string"}), + supports_stream=True, + cancelable=True, + ), + stream_handler=llm_stream, + finalize=lambda chunks: {"text": "".join(item.get("text", "") for item in chunks)}, + ) + self.register( + CapabilityDescriptor( + name="memory.search", + description="搜索记忆", + input_schema=obj_schema(["query"], query={"type": "string"}), + output_schema=obj_schema(["items"], items={"type": "array"}), + ), + call_handler=memory_search, + ) + self.register( + CapabilityDescriptor( + name="memory.save", + description="保存记忆", + input_schema=obj_schema(["key", "value"], key={"type": "string"}, value={"type": "object"}), + output_schema=obj_schema([]), + ), + call_handler=memory_save, + ) + self.register( + CapabilityDescriptor( + name="memory.delete", + description="删除记忆", + input_schema=obj_schema(["key"], key={"type": "string"}), + output_schema=obj_schema([]), + ), + call_handler=memory_delete, + ) + self.register( + CapabilityDescriptor( + name="db.get", + description="读取 KV", + input_schema=obj_schema(["key"], key={"type": "string"}), + output_schema=obj_schema([], value={"type": "object"}), + ), + call_handler=db_get, + ) + self.register( + CapabilityDescriptor( + name="db.set", + description="写入 KV", + input_schema=obj_schema(["key", "value"], key={"type": "string"}, value={"type": "object"}), + output_schema=obj_schema([]), + ), + call_handler=db_set, + ) + self.register( + CapabilityDescriptor( + name="db.delete", + description="删除 KV", + input_schema=obj_schema(["key"], key={"type": "string"}), + output_schema=obj_schema([]), + ), + call_handler=db_delete, + ) + self.register( + CapabilityDescriptor( + name="db.list", + description="列出 KV", + input_schema=obj_schema([], prefix={"type": "string"}), + output_schema=obj_schema(["keys"], keys={"type": "array"}), + ), + call_handler=db_list, + ) + self.register( + CapabilityDescriptor( + name="platform.send", + description="发送消息", + input_schema=obj_schema(["session", "text"], session={"type": "string"}, text={"type": "string"}), + output_schema=obj_schema(["message_id"], message_id={"type": "string"}), + ), + call_handler=platform_send, + ) + self.register( + CapabilityDescriptor( + name="platform.send_image", + description="发送图片", + input_schema=obj_schema(["session", "image_url"], session={"type": "string"}, image_url={"type": "string"}), + output_schema=obj_schema(["message_id"], message_id={"type": "string"}), + ), + call_handler=platform_send_image, + ) + self.register( + CapabilityDescriptor( + name="platform.get_members", + description="获取群成员", + input_schema=obj_schema(["session"], session={"type": "string"}), + output_schema=obj_schema(["members"], members={"type": "array"}), + ), + call_handler=platform_get_members, + ) + + def _validate_schema( + self, + schema: dict[str, Any] | None, + payload: dict[str, Any], + ) -> None: + if schema is None: + return + if schema.get("type") == "object" and not isinstance(payload, dict): + raise AstrBotError.invalid_input("输入必须是 object") + for field_name in schema.get("required", []): + if field_name not in payload or payload[field_name] is None: + raise AstrBotError.invalid_input(f"缺少必填字段：{field_name}") diff --git a/src-new/astrbot_sdk/runtime/handler_dispatcher.py b/src-new/astrbot_sdk/runtime/handler_dispatcher.py new file mode 100644 index 000000000..032cb7f85 --- /dev/null +++ b/src-new/astrbot_sdk/runtime/handler_dispatcher.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import asyncio +import inspect +from typing import Any + +from ..context import CancelToken, Context +from ..events import MessageEvent, PlainTextResult +from ..star import Star +from .loader import LoadedHandler + + +class HandlerDispatcher: + def __init__(self, *, plugin_id: str, peer, handlers: list[LoadedHandler]) -> None: + self._plugin_id = plugin_id + self._peer = peer + self._handlers = {item.descriptor.id: item for item in handlers} + self._active: dict[str, tuple[asyncio.Task[Any], CancelToken]] = {} + + async def invoke(self, message, cancel_token: CancelToken) -> dict[str, Any]: + handler_id = str(message.input.get("handler_id", "")) + loaded = self._handlers.get(handler_id) + if loaded is None: + raise LookupError(f"handler not found: {handler_id}") + + ctx = Context(peer=self._peer, plugin_id=self._plugin_id, cancel_token=cancel_token) + event = MessageEvent.from_payload(message.input.get("event", {}), context=ctx) + if loaded.legacy_context is not None: + loaded.legacy_context.bind_runtime_context(ctx) + + task = asyncio.create_task(self._run_handler(loaded, event, ctx)) + self._active[message.id] = (task, cancel_token) + try: + await task + return {} + finally: + self._active.pop(message.id, None) + + async def cancel(self, request_id: str) -> None: + active = self._active.get(request_id) + if active is None: + return + task, cancel_token = active + cancel_token.cancel() + task.cancel() + + async def _run_handler( + self, + loaded: LoadedHandler, + event: MessageEvent, + ctx: Context, + ) -> None: + try: + result = loaded.callable(*self._build_args(loaded.callable, event, ctx)) + if inspect.isasyncgen(result): + async for item in result: + await self._consume_legacy_result(item, event) + return + if inspect.isawaitable(result): + result = await result + if result is not None: + await self._consume_legacy_result(result, event) + except Exception as exc: + await self._handle_error(loaded.owner, exc, event, ctx) + raise + + def _build_args(self, handler, event: MessageEvent, ctx: Context) -> list[Any]: + signature = inspect.signature(handler) + args: list[Any] = [] + for parameter in signature.parameters.values(): + if parameter.kind not in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ): + continue + if parameter.name == "event": + args.append(event) + elif parameter.name in {"ctx", "context"}: + args.append(ctx) + return args + + async def _consume_legacy_result(self, item: Any, event: MessageEvent) -> None: + if isinstance(item, PlainTextResult): + await event.reply(item.text) + return + if isinstance(item, str): + await event.reply(item) + return + if isinstance(item, dict) and "text" in item: + await event.reply(str(item["text"])) + + async def _handle_error( + self, + owner: Any, + exc: Exception, + event: MessageEvent, + ctx: Context, + ) -> None: + if hasattr(owner, "on_error") and callable(owner.on_error): + await owner.on_error(exc, event, ctx) + return + await Star().on_error(exc, event, ctx) diff --git a/src-new/astrbot_sdk/runtime/loader.py b/src-new/astrbot_sdk/runtime/loader.py new file mode 100644 index 000000000..270524a6c --- /dev/null +++ b/src-new/astrbot_sdk/runtime/loader.py @@ -0,0 +1,307 @@ +from __future__ import annotations + +import json +import os +import re +import shutil +import subprocess +import sys +from dataclasses import dataclass +from importlib import import_module +from pathlib import Path +from typing import Any + +import yaml + +from ..compat import LegacyContext +from ..decorators import get_handler_meta +from ..protocol.descriptors import HandlerDescriptor +from ..star import Star + +STATE_FILE_NAME = ".astrbot-worker-state.json" + + +def _venv_python_path(venv_dir: Path) -> Path: + if os.name == "nt": + return venv_dir / "Scripts" / "python.exe" + return venv_dir / "bin" / "python" + + +@dataclass(slots=True) +class PluginSpec: + name: str + plugin_dir: Path + manifest_path: Path + requirements_path: Path + python_version: str + manifest_data: dict[str, Any] + + +@dataclass(slots=True) +class PluginDiscoveryResult: + plugins: list[PluginSpec] + skipped_plugins: dict[str, str] + + +@dataclass(slots=True) +class LoadedHandler: + descriptor: HandlerDescriptor + callable: Any + owner: Any + legacy_context: LegacyContext | None = None + + +@dataclass(slots=True) +class LoadedPlugin: + plugin: PluginSpec + handlers: list[LoadedHandler] + instances: list[Any] + + +def load_plugin_spec(plugin_dir: Path) -> PluginSpec: + plugin_dir = plugin_dir.resolve() + manifest_path = plugin_dir / "plugin.yaml" + requirements_path = plugin_dir / "requirements.txt" + manifest_data = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {} + runtime = manifest_data.get("runtime") or {} + python_version = runtime.get("python") or f"{sys.version_info.major}.{sys.version_info.minor}" + return PluginSpec( + name=str(manifest_data.get("name") or plugin_dir.name), + plugin_dir=plugin_dir, + manifest_path=manifest_path, + requirements_path=requirements_path, + python_version=str(python_version), + manifest_data=manifest_data, + ) + + +def discover_plugins(plugins_dir: Path) -> PluginDiscoveryResult: + plugins_root = plugins_dir.resolve() + skipped_plugins: dict[str, str] = {} + plugins: list[PluginSpec] = [] + seen_names: set[str] = set() + + if not plugins_root.exists(): + return PluginDiscoveryResult([], {}) + + for entry in sorted(plugins_root.iterdir()): + if not entry.is_dir() or entry.name.startswith("."): + continue + manifest_path = entry / "plugin.yaml" + requirements_path = entry / "requirements.txt" + if not manifest_path.exists(): + continue + if not requirements_path.exists(): + skipped_plugins[entry.name] = "missing requirements.txt" + continue + try: + manifest_data = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {} + except Exception as exc: + skipped_plugins[entry.name] = f"failed to parse plugin.yaml: {exc}" + continue + plugin_name = manifest_data.get("name") + runtime = manifest_data.get("runtime") or {} + python_version = runtime.get("python") + components = manifest_data.get("components") + if not isinstance(plugin_name, str) or not plugin_name: + skipped_plugins[entry.name] = "plugin name is required" + continue + if plugin_name in seen_names: + skipped_plugins[plugin_name] = "duplicate plugin name" + continue + if not isinstance(components, list) or not components: + skipped_plugins[plugin_name] = "components must be a non-empty list" + continue + if not isinstance(python_version, str) or not python_version: + skipped_plugins[plugin_name] = "runtime.python is required" + continue + seen_names.add(plugin_name) + plugins.append( + PluginSpec( + name=plugin_name, + plugin_dir=entry.resolve(), + manifest_path=manifest_path.resolve(), + requirements_path=requirements_path.resolve(), + python_version=python_version, + manifest_data=manifest_data, + ) + ) + return PluginDiscoveryResult(plugins=plugins, skipped_plugins=skipped_plugins) + + +class PluginEnvironmentManager: + def __init__(self, repo_root: Path, uv_binary: str | None = None) -> None: + self.repo_root = repo_root.resolve() + self.uv_binary = uv_binary or shutil.which("uv") + self.cache_dir = self.repo_root / ".uv-cache" + + def prepare_environment(self, plugin: PluginSpec) -> Path: + if not self.uv_binary: + raise RuntimeError("uv executable not found") + state_path = plugin.plugin_dir / STATE_FILE_NAME + venv_dir = plugin.plugin_dir / ".venv" + python_path = _venv_python_path(venv_dir) + fingerprint = self._fingerprint(plugin) + state = self._load_state(state_path) + if ( + not python_path.exists() + or not self._matches_python_version(venv_dir, plugin.python_version) + or state.get("fingerprint") != fingerprint + ): + self._rebuild(plugin, venv_dir, python_path) + self._write_state(state_path, plugin, fingerprint) + return python_path + + def _rebuild(self, plugin: PluginSpec, venv_dir: Path, python_path: Path) -> None: + if venv_dir.exists(): + shutil.rmtree(venv_dir) + self._run_command( + [ + self.uv_binary, + "venv", + "--python", + plugin.python_version, + "--system-site-packages", + "--no-python-downloads", + "--no-managed-python", + str(venv_dir), + ], + cwd=self.repo_root, + command_name=f"create venv for {plugin.name}", + ) + requirements_text = plugin.requirements_path.read_text(encoding="utf-8").strip() + if not requirements_text: + return + self._run_command( + [ + self.uv_binary, + "pip", + "install", + "--python", + str(python_path), + "-r", + str(plugin.requirements_path), + ], + cwd=plugin.plugin_dir, + command_name=f"install requirements for {plugin.name}", + ) + + def _run_command( + self, + command: list[str], + *, + cwd: Path, + command_name: str, + ) -> None: + process = subprocess.run( + command, + cwd=str(cwd), + env={**os.environ, "UV_CACHE_DIR": str(self.cache_dir)}, + capture_output=True, + text=True, + check=False, + ) + if process.returncode != 0: + raise RuntimeError( + f"{command_name} failed with exit code {process.returncode}: " + f"{process.stderr.strip() or process.stdout.strip()}" + ) + + @staticmethod + def _fingerprint(plugin: PluginSpec) -> str: + requirements = plugin.requirements_path.read_text(encoding="utf-8") + payload = { + "python_version": plugin.python_version, + "requirements": requirements, + } + return json.dumps(payload, ensure_ascii=True, sort_keys=True) + + @staticmethod + def _load_state(state_path: Path) -> dict[str, Any]: + if not state_path.exists(): + return {} + try: + data = json.loads(state_path.read_text(encoding="utf-8")) + except Exception: + return {} + return data if isinstance(data, dict) else {} + + @staticmethod + def _write_state(state_path: Path, plugin: PluginSpec, fingerprint: str) -> None: + state_path.write_text( + json.dumps( + { + "plugin": plugin.name, + "python_version": plugin.python_version, + "fingerprint": fingerprint, + }, + ensure_ascii=True, + indent=2, + sort_keys=True, + ), + encoding="utf-8", + ) + + @staticmethod + def _matches_python_version(venv_dir: Path, version: str) -> bool: + pyvenv_cfg = venv_dir / "pyvenv.cfg" + if not pyvenv_cfg.exists(): + return False + content = pyvenv_cfg.read_text(encoding="utf-8") + match = re.search(r"version\s*=\s*(\d+\.\d+)\.\d+", content, re.IGNORECASE) + return match is not None and match.group(1) == version + + +def load_plugin(plugin: PluginSpec) -> LoadedPlugin: + plugin_path = str(plugin.plugin_dir) + if plugin_path not in sys.path: + sys.path.insert(0, plugin_path) + + instances: list[Any] = [] + handlers: list[LoadedHandler] = [] + for component in plugin.manifest_data.get("components", []): + class_path = component.get("class") + if not isinstance(class_path, str) or ":" not in class_path: + continue + component_cls = import_string(class_path) + legacy_context: LegacyContext | None = None + if isinstance(component_cls, type) and issubclass(component_cls, Star): + instance = component_cls() + else: + legacy_context = LegacyContext(plugin.name) + try: + instance = component_cls(legacy_context) + except TypeError: + instance = component_cls() + if getattr(instance, "context", None) is None: + setattr(instance, "context", legacy_context) + instances.append(instance) + for name in dir(instance): + bound = getattr(instance, name) + func = getattr(bound, "__func__", bound) + meta = get_handler_meta(func) + if meta is None or meta.trigger is None: + continue + handler_id = ( + f"{plugin.name}:{instance.__class__.__module__}.{instance.__class__.__name__}.{name}" + ) + handlers.append( + LoadedHandler( + descriptor=HandlerDescriptor( + id=handler_id, + trigger=meta.trigger, + priority=meta.priority, + permissions=meta.permissions.model_copy(deep=True), + ), + callable=bound, + owner=instance, + legacy_context=legacy_context, + ) + ) + return LoadedPlugin(plugin=plugin, handlers=handlers, instances=instances) + + +def import_string(path: str) -> Any: + module_name, attr = path.split(":", 1) + module = import_module(module_name) + return getattr(module, attr) diff --git a/src-new/astrbot_sdk/runtime/peer.py b/src-new/astrbot_sdk/runtime/peer.py new file mode 100644 index 000000000..aa6cc2d2d --- /dev/null +++ b/src-new/astrbot_sdk/runtime/peer.py @@ -0,0 +1,346 @@ +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Awaitable, Callable +from typing import Any + +from ..context import CancelToken +from ..errors import AstrBotError +from ..protocol.messages import ( + CancelMessage, + ErrorPayload, + EventMessage, + InitializeMessage, + InitializeOutput, + InvokeMessage, + PeerInfo, + ResultMessage, + parse_message, +) +from .capability_router import StreamExecution + +InitializeHandler = Callable[[InitializeMessage], Awaitable[InitializeOutput]] +InvokeHandler = Callable[[InvokeMessage, CancelToken], Awaitable[dict[str, Any] | StreamExecution]] +CancelHandler = Callable[[str], Awaitable[None]] + + +class Peer: + def __init__( + self, + *, + transport, + peer_info: PeerInfo, + protocol_version: str = "1.0", + ) -> None: + self.transport = transport + self.peer_info = peer_info + self.protocol_version = protocol_version + self.remote_peer: PeerInfo | None = None + self.remote_handlers = [] + self.remote_capabilities = [] + self.remote_capability_map: dict[str, Any] = {} + self.remote_metadata: dict[str, Any] = {} + + self._initialize_handler: InitializeHandler | None = None + self._invoke_handler: InvokeHandler | None = None + self._cancel_handler: CancelHandler | None = None + self._counter = 0 + self._closed = False + self._unusable = False + self._pending_results: dict[str, asyncio.Future[ResultMessage]] = {} + self._pending_streams: dict[str, asyncio.Queue[Any]] = {} + self._inbound_tasks: dict[str, tuple[asyncio.Task[None], CancelToken]] = {} + self._remote_initialized = asyncio.Event() + + def set_initialize_handler(self, handler: InitializeHandler) -> None: + self._initialize_handler = handler + + def set_invoke_handler(self, handler: InvokeHandler) -> None: + self._invoke_handler = handler + + def set_cancel_handler(self, handler: CancelHandler) -> None: + self._cancel_handler = handler + + async def start(self) -> None: + self.transport.set_message_handler(self._handle_raw_message) + await self.transport.start() + + async def stop(self) -> None: + self._closed = True + await self.transport.stop() + + async def wait_closed(self) -> None: + await self.transport.wait_closed() + + async def wait_until_remote_initialized(self, timeout: float | None = 30.0) -> None: + if timeout is None: + await self._remote_initialized.wait() + return + await asyncio.wait_for(self._remote_initialized.wait(), timeout=timeout) + + async def initialize( + self, + handlers, + *, + metadata: dict[str, Any] | None = None, + ) -> InitializeOutput: + self._ensure_usable() + request_id = self._next_id() + future: asyncio.Future[ResultMessage] = asyncio.get_running_loop().create_future() + self._pending_results[request_id] = future + await self._send( + InitializeMessage( + id=request_id, + protocol_version=self.protocol_version, + peer=self.peer_info, + handlers=list(handlers), + metadata=metadata or {}, + ) + ) + result = await future + if result.kind != "initialize_result": + raise AstrBotError.protocol_error("initialize 必须收到 initialize_result") + if not result.success: + self._unusable = True + raise AstrBotError.from_payload(result.error.model_dump() if result.error else {}) + output = InitializeOutput.model_validate(result.output) + self.remote_peer = output.peer + self.remote_capabilities = output.capabilities + self.remote_capability_map = {item.name: item for item in output.capabilities} + self.remote_metadata = output.metadata + return output + + async def invoke( + self, + capability: str, + payload: dict[str, Any], + *, + stream: bool = False, + request_id: str | None = None, + ) -> dict[str, Any]: + self._ensure_usable() + if stream: + raise ValueError("stream=True 请使用 invoke_stream()") + request_id = request_id or self._next_id() + future: asyncio.Future[ResultMessage] = asyncio.get_running_loop().create_future() + self._pending_results[request_id] = future + await self._send( + InvokeMessage( + id=request_id, + capability=capability, + input=payload, + stream=False, + ) + ) + result = await future + if not result.success: + raise AstrBotError.from_payload(result.error.model_dump() if result.error else {}) + return result.output + + async def invoke_stream( + self, + capability: str, + payload: dict[str, Any], + *, + request_id: str | None = None, + ) -> AsyncIterator[EventMessage]: + self._ensure_usable() + request_id = request_id or self._next_id() + queue: asyncio.Queue[Any] = asyncio.Queue() + self._pending_streams[request_id] = queue + await self._send( + InvokeMessage( + id=request_id, + capability=capability, + input=payload, + stream=True, + ) + ) + + async def iterator() -> AsyncIterator[EventMessage]: + try: + while True: + item = await queue.get() + if isinstance(item, Exception): + raise item + if not isinstance(item, EventMessage): + raise AstrBotError.protocol_error("流式调用收到非法事件") + if item.phase == "started": + continue + if item.phase == "delta": + yield item + continue + if item.phase == "completed": + break + if item.phase == "failed": + raise AstrBotError.from_payload(item.error.model_dump() if item.error else {}) + finally: + self._pending_streams.pop(request_id, None) + + return iterator() + + async def cancel(self, request_id: str, reason: str = "user_cancelled") -> None: + await self._send(CancelMessage(id=request_id, reason=reason)) + + def _next_id(self) -> str: + self._counter += 1 + return f"msg_{self._counter:04d}" + + def _ensure_usable(self) -> None: + if self._unusable: + raise AstrBotError.protocol_error("连接已进入不可用状态") + + async def _handle_raw_message(self, payload: str) -> None: + message = parse_message(payload) + if isinstance(message, ResultMessage): + await self._handle_result(message) + return + if isinstance(message, EventMessage): + await self._handle_event(message) + return + if isinstance(message, InitializeMessage): + await self._handle_initialize(message) + return + if isinstance(message, InvokeMessage): + task = asyncio.create_task(self._handle_invoke(message)) + token = CancelToken() + self._inbound_tasks[message.id] = (task, token) + task.add_done_callback(lambda _task, request_id=message.id: self._inbound_tasks.pop(request_id, None)) + return + if isinstance(message, CancelMessage): + await self._handle_cancel(message) + return + + async def _handle_initialize(self, message: InitializeMessage) -> None: + self.remote_peer = message.peer + self.remote_handlers = message.handlers + self.remote_metadata = message.metadata + if self._initialize_handler is None: + error = AstrBotError.protocol_error("对端不接受 initialize") + await self._send( + ResultMessage( + id=message.id, + kind="initialize_result", + success=False, + error=ErrorPayload.model_validate(error.to_payload()), + ) + ) + self._unusable = True + self._remote_initialized.set() + return + + if message.protocol_version != self.protocol_version: + error = AstrBotError.protocol_version_mismatch( + f"服务端支持协议版本 {self.protocol_version}，客户端请求版本 {message.protocol_version}" + ) + await self._send( + ResultMessage( + id=message.id, + kind="initialize_result", + success=False, + error=ErrorPayload.model_validate(error.to_payload()), + ) + ) + self._unusable = True + self._remote_initialized.set() + return + + output = await self._initialize_handler(message) + await self._send( + ResultMessage( + id=message.id, + kind="initialize_result", + success=True, + output=output.model_dump(), + ) + ) + self._remote_initialized.set() + + async def _handle_invoke(self, message: InvokeMessage) -> None: + active = self._inbound_tasks.get(message.id) + token = active[1] if active is not None else CancelToken() + try: + if self._invoke_handler is None: + raise AstrBotError.capability_not_found(message.capability) + execution = await self._invoke_handler(message, token) + if message.stream: + if not isinstance(execution, StreamExecution): + raise AstrBotError.protocol_error("stream=true 必须返回 StreamExecution") + await self._send(EventMessage(id=message.id, phase="started")) + chunks: list[dict[str, Any]] = [] + async for chunk in execution.iterator: + chunks.append(chunk) + await self._send(EventMessage(id=message.id, phase="delta", data=chunk)) + await self._send( + EventMessage( + id=message.id, + phase="completed", + output=execution.finalize(chunks), + ) + ) + return + if isinstance(execution, StreamExecution): + raise AstrBotError.protocol_error("stream=false 不能返回流式执行对象") + await self._send(ResultMessage(id=message.id, success=True, output=execution)) + except asyncio.CancelledError: + await self._send_cancelled_termination(message) + except LookupError as exc: + error = AstrBotError.invalid_input(str(exc)) + await self._send_error_result(message, error) + except AstrBotError as exc: + await self._send_error_result(message, exc) + except Exception as exc: + await self._send_error_result(message, AstrBotError.internal_error(str(exc))) + + async def _handle_cancel(self, message: CancelMessage) -> None: + inbound = self._inbound_tasks.get(message.id) + if inbound is None: + return + task, token = inbound + token.cancel() + if self._cancel_handler is not None: + await self._cancel_handler(message.id) + task.cancel() + + async def _handle_result(self, message: ResultMessage) -> None: + future = self._pending_results.pop(message.id, None) + if future is None: + queue = self._pending_streams.get(message.id) + if queue is not None: + await queue.put(AstrBotError.protocol_error("stream=true 调用不应收到 result")) + return + future.set_result(message) + + async def _handle_event(self, message: EventMessage) -> None: + queue = self._pending_streams.get(message.id) + if queue is None: + future = self._pending_results.get(message.id) + if future is not None: + future.set_exception(AstrBotError.protocol_error("stream=false 调用不应收到 event")) + return + await queue.put(message) + + async def _send_error_result(self, message: InvokeMessage, error: AstrBotError) -> None: + if message.stream: + await self._send( + EventMessage( + id=message.id, + phase="failed", + error=ErrorPayload.model_validate(error.to_payload()), + ) + ) + return + await self._send( + ResultMessage( + id=message.id, + success=False, + error=ErrorPayload.model_validate(error.to_payload()), + ) + ) + + async def _send_cancelled_termination(self, message: InvokeMessage) -> None: + error = AstrBotError.cancelled() + await self._send_error_result(message, error) + + async def _send(self, message) -> None: + await self.transport.send(message.model_dump_json(exclude_none=True)) diff --git a/src-new/astrbot_sdk/runtime/transport.py b/src-new/astrbot_sdk/runtime/transport.py new file mode 100644 index 000000000..dc246d295 --- /dev/null +++ b/src-new/astrbot_sdk/runtime/transport.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import asyncio +import sys +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable, Sequence +from typing import IO + +import aiohttp +from aiohttp import web +from loguru import logger + +MessageHandler = Callable[[str], Awaitable[None]] + + +class Transport(ABC): + def __init__(self) -> None: + self._handler: MessageHandler | None = None + self._closed = asyncio.Event() + + def set_message_handler(self, handler: MessageHandler) -> None: + self._handler = handler + + @abstractmethod + async def start(self) -> None: + raise NotImplementedError + + @abstractmethod + async def stop(self) -> None: + raise NotImplementedError + + @abstractmethod + async def send(self, payload: str) -> None: + raise NotImplementedError + + async def wait_closed(self) -> None: + await self._closed.wait() + + async def _dispatch(self, payload: str) -> None: + if self._handler is not None: + await self._handler(payload) + + +class StdioTransport(Transport): + def __init__( + self, + *, + stdin: IO[str] | None = None, + stdout: IO[str] | None = None, + command: Sequence[str] | None = None, + cwd: str | None = None, + env: dict[str, str] | None = None, + ) -> None: + super().__init__() + self._stdin = stdin + self._stdout = stdout + self._command = list(command) if command is not None else None + self._cwd = cwd + self._env = env + self._process: asyncio.subprocess.Process | None = None + self._reader_task: asyncio.Task[None] | None = None + + async def start(self) -> None: + self._closed.clear() + if self._command is not None: + self._process = await asyncio.create_subprocess_exec( + *self._command, + cwd=self._cwd, + env=self._env, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=sys.stderr, + ) + self._reader_task = asyncio.create_task(self._read_process_loop()) + return + + self._stdin = self._stdin or sys.stdin + self._stdout = self._stdout or sys.stdout + self._reader_task = asyncio.create_task(self._read_file_loop()) + + async def stop(self) -> None: + if self._reader_task is not None: + self._reader_task.cancel() + try: + await self._reader_task + except asyncio.CancelledError: + pass + self._reader_task = None + + if self._process is not None: + if self._process.returncode is None: + self._process.terminate() + try: + await asyncio.wait_for(self._process.wait(), timeout=5) + except asyncio.TimeoutError: + self._process.kill() + await self._process.wait() + self._process = None + self._closed.set() + + async def send(self, payload: str) -> None: + line = payload if payload.endswith("\n") else f"{payload}\n" + if self._process is not None: + if self._process.stdin is None: + raise RuntimeError("STDIO subprocess stdin 不可用") + self._process.stdin.write(line.encode("utf-8")) + await self._process.stdin.drain() + return + + if self._stdout is None: + raise RuntimeError("STDIO stdout 不可用") + + def _write() -> None: + assert self._stdout is not None + self._stdout.write(line) + self._stdout.flush() + + await asyncio.to_thread(_write) + + async def _read_process_loop(self) -> None: + assert self._process is not None + assert self._process.stdout is not None + try: + while True: + raw = await self._process.stdout.readline() + if not raw: + break + await self._dispatch(raw.decode("utf-8").rstrip("\r\n")) + finally: + self._closed.set() + + async def _read_file_loop(self) -> None: + assert self._stdin is not None + try: + while True: + raw = await asyncio.to_thread(self._stdin.readline) + if not raw: + break + await self._dispatch(raw.rstrip("\r\n")) + finally: + self._closed.set() + + +class WebSocketServerTransport(Transport): + def __init__( + self, + *, + host: str = "127.0.0.1", + port: int = 8765, + path: str = "/", + heartbeat: float = 30.0, + ) -> None: + super().__init__() + self._host = host + self._port = port + self._actual_port: int | None = None + self._path = path + self._heartbeat = heartbeat + self._app: web.Application | None = None + self._runner: web.AppRunner | None = None + self._site: web.TCPSite | None = None + self._ws: web.WebSocketResponse | None = None + self._write_lock = asyncio.Lock() + + async def start(self) -> None: + self._closed.clear() + self._app = web.Application() + self._app.router.add_get(self._path, self._handle_socket) + self._runner = web.AppRunner(self._app) + await self._runner.setup() + self._site = web.TCPSite(self._runner, self._host, self._port) + await self._site.start() + if self._site._server and getattr(self._site._server, "sockets", None): + socket = self._site._server.sockets[0] + self._actual_port = socket.getsockname()[1] + + async def stop(self) -> None: + if self._ws is not None and not self._ws.closed: + await self._ws.close() + if self._site is not None: + await self._site.stop() + self._site = None + if self._runner is not None: + await self._runner.cleanup() + self._runner = None + self._closed.set() + + async def send(self, payload: str) -> None: + if self._ws is None or self._ws.closed: + raise RuntimeError("WebSocket 尚未连接") + async with self._write_lock: + await self._ws.send_str(payload) + + async def _handle_socket(self, request: web.Request) -> web.WebSocketResponse: + if self._ws is not None and not self._ws.closed: + ws = web.WebSocketResponse() + await ws.prepare(request) + await ws.close(code=1008, message=b"only one websocket connection allowed") + return ws + + ws = web.WebSocketResponse( + heartbeat=self._heartbeat if self._heartbeat > 0 else None + ) + await ws.prepare(request) + self._ws = ws + try: + async for msg in ws: + if msg.type == aiohttp.WSMsgType.TEXT: + await self._dispatch(msg.data) + elif msg.type == aiohttp.WSMsgType.BINARY: + await self._dispatch(msg.data.decode("utf-8")) + elif msg.type == aiohttp.WSMsgType.ERROR: + logger.error("websocket server error: {}", ws.exception()) + break + finally: + self._closed.set() + self._ws = None + return ws + + @property + def port(self) -> int: + return self._actual_port or self._port + + @property + def url(self) -> str: + return f"ws://{self._host}:{self.port}{self._path}" + + +class WebSocketClientTransport(Transport): + def __init__( + self, + *, + url: str, + heartbeat: float = 30.0, + ) -> None: + super().__init__() + self._url = url + self._heartbeat = heartbeat + self._session: aiohttp.ClientSession | None = None + self._ws: aiohttp.ClientWebSocketResponse | None = None + self._reader_task: asyncio.Task[None] | None = None + + async def start(self) -> None: + self._closed.clear() + self._session = aiohttp.ClientSession() + self._ws = await self._session.ws_connect( + self._url, + heartbeat=self._heartbeat if self._heartbeat > 0 else None, + ) + self._reader_task = asyncio.create_task(self._read_loop()) + + async def stop(self) -> None: + if self._reader_task is not None: + self._reader_task.cancel() + try: + await self._reader_task + except asyncio.CancelledError: + pass + self._reader_task = None + if self._ws is not None and not self._ws.closed: + await self._ws.close() + if self._session is not None: + await self._session.close() + self._ws = None + self._session = None + self._closed.set() + + async def send(self, payload: str) -> None: + if self._ws is None or self._ws.closed: + raise RuntimeError("WebSocket client 尚未连接") + await self._ws.send_str(payload) + + async def _read_loop(self) -> None: + assert self._ws is not None + try: + async for msg in self._ws: + if msg.type == aiohttp.WSMsgType.TEXT: + await self._dispatch(msg.data) + elif msg.type == aiohttp.WSMsgType.BINARY: + await self._dispatch(msg.data.decode("utf-8")) + elif msg.type == aiohttp.WSMsgType.ERROR: + logger.error("websocket client error: {}", self._ws.exception()) + break + finally: + self._closed.set() diff --git a/src-new/astrbot_sdk/star.py b/src-new/astrbot_sdk/star.py new file mode 100644 index 000000000..4cdc80506 --- /dev/null +++ b/src-new/astrbot_sdk/star.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +import traceback +from typing import Any + +from loguru import logger + +from .errors import AstrBotError + + +class Star: + async def on_error(self, error: Exception, event, ctx) -> None: + if isinstance(error, AstrBotError): + if error.retryable: + await event.reply("请求失败，请稍后重试") + elif error.hint: + await event.reply(error.hint) + else: + await event.reply(error.message) + else: + await event.reply("出了点问题，请联系插件作者") + logger.error("handler 执行失败\n{}", traceback.format_exc()) + + @classmethod + def __astrbot_is_new_star__(cls) -> bool: + return True diff --git a/tests_v4/__init__.py b/tests_v4/__init__.py new file mode 100644 index 000000000..c9c2ef67b --- /dev/null +++ b/tests_v4/__init__.py @@ -0,0 +1 @@ +__all__: list[str] = [] diff --git a/tests_v4/helpers.py b/tests_v4/helpers.py new file mode 100644 index 000000000..528231f69 --- /dev/null +++ b/tests_v4/helpers.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path + +from astrbot_sdk.runtime.transport import Transport + + +class MemoryTransport(Transport): + def __init__(self) -> None: + super().__init__() + self.partner: "MemoryTransport | None" = None + + async def start(self) -> None: + self._closed.clear() + + async def stop(self) -> None: + self._closed.set() + + async def send(self, payload: str) -> None: + if self.partner is None: + raise RuntimeError("MemoryTransport 未连接 partner") + await self.partner._dispatch(payload) + + +def make_transport_pair() -> tuple[MemoryTransport, MemoryTransport]: + left = MemoryTransport() + right = MemoryTransport() + left.partner = right + right.partner = left + return left, right + + +class FakeEnvManager: + def prepare_environment(self, _plugin) -> Path: + return Path(__import__("sys").executable) + + +async def drain_loop() -> None: + await asyncio.sleep(0.05) diff --git a/tests_v4/test_entrypoints.py b/tests_v4/test_entrypoints.py new file mode 100644 index 000000000..46f7470d4 --- /dev/null +++ b/tests_v4/test_entrypoints.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import subprocess +import sys +import unittest + + +class EntryPointTest(unittest.TestCase): + def test_import_package(self) -> None: + process = subprocess.run( + [sys.executable, "-c", "import astrbot_sdk; print(astrbot_sdk.__name__)"], + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(process.returncode, 0, process.stderr) + self.assertIn("astrbot_sdk", process.stdout) + + def test_module_help(self) -> None: + process = subprocess.run( + [sys.executable, "-m", "astrbot_sdk", "--help"], + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(process.returncode, 0, process.stderr) + self.assertIn("Usage", process.stdout) + + def test_run_help(self) -> None: + process = subprocess.run( + [sys.executable, "-m", "astrbot_sdk", "run", "--help"], + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(process.returncode, 0, process.stderr) + self.assertIn("--plugins-dir", process.stdout) diff --git a/tests_v4/test_peer.py b/tests_v4/test_peer.py new file mode 100644 index 000000000..e307c8e1b --- /dev/null +++ b/tests_v4/test_peer.py @@ -0,0 +1,189 @@ +from __future__ import annotations + +import asyncio +import unittest + +from astrbot_sdk.context import CancelToken +from astrbot_sdk.errors import AstrBotError +from astrbot_sdk.protocol.descriptors import CapabilityDescriptor +from astrbot_sdk.protocol.messages import EventMessage, InitializeOutput, PeerInfo, ResultMessage +from astrbot_sdk.runtime.capability_router import CapabilityRouter, StreamExecution +from astrbot_sdk.runtime.peer import Peer +from astrbot_sdk.runtime.transport import WebSocketClientTransport, WebSocketServerTransport + +from tests_v4.helpers import make_transport_pair + + +class PeerRuntimeTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self) -> None: + self.left, self.right = make_transport_pair() + + async def test_initialize_and_call_builtin_capabilities(self) -> None: + router = CapabilityRouter() + core = Peer( + transport=self.left, + peer_info=PeerInfo(name="core", role="core", version="v4"), + ) + core.set_initialize_handler( + lambda _message: asyncio.sleep(0, result=InitializeOutput( + peer=PeerInfo(name="core", role="core", version="v4"), + capabilities=router.descriptors(), + metadata={}, + )) + ) + core.set_invoke_handler( + lambda message, token: router.execute( + message.capability, + message.input, + stream=message.stream, + cancel_token=token, + ) + ) + + plugin = Peer( + transport=self.right, + peer_info=PeerInfo(name="plugin", role="plugin", version="v4"), + ) + + await core.start() + await plugin.start() + await plugin.initialize([]) + + result = await plugin.invoke("llm.chat", {"prompt": "hello"}) + self.assertEqual(result["text"], "Echo: hello") + + stream = await plugin.invoke_stream("llm.stream_chat", {"prompt": "hi"}) + chunks = [event.data["text"] async for event in stream] + self.assertEqual("".join(chunks), "Echo: hi") + + await plugin.stop() + await core.stop() + + async def test_stream_false_receiving_event_is_protocol_error(self) -> None: + plugin = Peer( + transport=self.right, + peer_info=PeerInfo(name="plugin", role="plugin", version="v4"), + ) + await self.left.start() + await plugin.start() + + task = asyncio.create_task( + plugin.invoke("llm.chat", {"prompt": "bad"}, request_id="req-1") + ) + await asyncio.sleep(0) + await self.left.send(EventMessage(id="req-1", phase="started").model_dump_json()) + + with self.assertRaises(AstrBotError) as raised: + await task + self.assertEqual(raised.exception.code, "protocol_error") + await plugin.stop() + await self.left.stop() + + async def test_stream_true_receiving_result_is_protocol_error(self) -> None: + plugin = Peer( + transport=self.right, + peer_info=PeerInfo(name="plugin", role="plugin", version="v4"), + ) + await self.left.start() + await plugin.start() + + stream = await plugin.invoke_stream("llm.stream_chat", {"prompt": "bad"}, request_id="stream-1") + await self.left.send(ResultMessage(id="stream-1", success=True, output={}).model_dump_json()) + + with self.assertRaises(AstrBotError) as raised: + async for _ in stream: + pass + self.assertEqual(raised.exception.code, "protocol_error") + await plugin.stop() + await self.left.stop() + + async def test_cancel_waits_for_failed_terminal_event(self) -> None: + descriptor = CapabilityDescriptor( + name="slow.stream", + description="slow stream", + input_schema={"type": "object", "properties": {}, "required": []}, + output_schema={"type": "object", "properties": {"count": {"type": "number"}}, "required": ["count"]}, + supports_stream=True, + cancelable=True, + ) + + async def init_handler(_message): + return InitializeOutput( + peer=PeerInfo(name="core", role="core", version="v4"), + capabilities=[descriptor], + metadata={}, + ) + + async def invoke_handler(message, token: CancelToken): + async def iterator(): + while True: + token.raise_if_cancelled() + await asyncio.sleep(0.01) + yield {"text": "x"} + + return StreamExecution( + iterator=iterator(), + finalize=lambda chunks: {"count": len(chunks)}, + ) + + core = Peer( + transport=self.left, + peer_info=PeerInfo(name="core", role="core", version="v4"), + ) + core.set_initialize_handler(init_handler) + core.set_invoke_handler(invoke_handler) + plugin = Peer( + transport=self.right, + peer_info=PeerInfo(name="plugin", role="plugin", version="v4"), + ) + + await core.start() + await plugin.start() + await plugin.initialize([]) + + stream = await plugin.invoke_stream("slow.stream", {}, request_id="cancel-1") + first = await anext(stream) + self.assertEqual(first.data["text"], "x") + await plugin.cancel("cancel-1") + + with self.assertRaises(AstrBotError) as raised: + await anext(stream) + self.assertEqual(raised.exception.code, "cancelled") + await plugin.stop() + await core.stop() + + async def test_websocket_transport_smoke(self) -> None: + router = CapabilityRouter() + server_transport = WebSocketServerTransport(port=0) + core = Peer( + transport=server_transport, + peer_info=PeerInfo(name="core", role="core", version="v4"), + ) + core.set_initialize_handler( + lambda _message: asyncio.sleep(0, result=InitializeOutput( + peer=PeerInfo(name="core", role="core", version="v4"), + capabilities=router.descriptors(), + metadata={}, + )) + ) + core.set_invoke_handler( + lambda message, token: router.execute( + message.capability, + message.input, + stream=message.stream, + cancel_token=token, + ) + ) + + await core.start() + client_transport = WebSocketClientTransport(url=server_transport.url) + plugin = Peer( + transport=client_transport, + peer_info=PeerInfo(name="plugin", role="plugin", version="v4"), + ) + await plugin.start() + await plugin.initialize([]) + result = await plugin.invoke("llm.chat", {"prompt": "ws"}) + self.assertEqual(result["text"], "Echo: ws") + await plugin.stop() + await core.stop() diff --git a/tests_v4/test_protocol.py b/tests_v4/test_protocol.py new file mode 100644 index 000000000..8d828df71 --- /dev/null +++ b/tests_v4/test_protocol.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import unittest + +from astrbot_sdk.protocol.descriptors import ( + CommandTrigger, + HandlerDescriptor, + ScheduleTrigger, +) +from astrbot_sdk.protocol.messages import ( + CancelMessage, + EventMessage, + InitializeMessage, + InvokeMessage, + PeerInfo, + ResultMessage, + parse_message, +) + + +class ProtocolModelsTest(unittest.TestCase): + def test_parse_message_roundtrip(self) -> None: + message = InitializeMessage( + id="msg_001", + protocol_version="1.0", + peer=PeerInfo(name="plugin", role="plugin", version="1.0.0"), + handlers=[ + HandlerDescriptor( + id="handler_1", + trigger=CommandTrigger(command="hello"), + ) + ], + metadata={"a": 1}, + ) + parsed = parse_message(message.model_dump_json()) + self.assertEqual(parsed.id, "msg_001") + self.assertEqual(parsed.peer.name, "plugin") + + for sample in [ + InvokeMessage(id="msg_002", capability="llm.chat", input={"prompt": "hi"}), + ResultMessage(id="msg_002", success=True, output={"text": "ok"}), + EventMessage(id="msg_003", phase="started"), + CancelMessage(id="msg_003"), + ]: + self.assertEqual(parse_message(sample.model_dump()).type, sample.type) + + def test_schedule_trigger_requires_exactly_one_strategy(self) -> None: + with self.assertRaises(ValueError): + ScheduleTrigger() + with self.assertRaises(ValueError): + ScheduleTrigger(cron="* * * * *", interval_seconds=10) + trigger = ScheduleTrigger(interval_seconds=30) + self.assertEqual(trigger.interval_seconds, 30) diff --git a/tests_v4/test_runtime.py b/tests_v4/test_runtime.py new file mode 100644 index 000000000..1ac536d62 --- /dev/null +++ b/tests_v4/test_runtime.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +import asyncio +import shutil +import sys +import tempfile +import textwrap +import unittest +from pathlib import Path + +from astrbot_sdk.protocol.messages import InitializeOutput, PeerInfo +from astrbot_sdk.runtime.bootstrap import SupervisorRuntime +from astrbot_sdk.runtime.peer import Peer + +from tests_v4.helpers import FakeEnvManager, make_transport_pair + + +def write_new_plugin(plugin_root: Path) -> None: + (plugin_root / "commands").mkdir(parents=True, exist_ok=True) + (plugin_root / "commands" / "__init__.py").write_text("", encoding="utf-8") + (plugin_root / "requirements.txt").write_text("", encoding="utf-8") + (plugin_root / "plugin.yaml").write_text( + textwrap.dedent( + f"""\ + _schema_version: 2 + name: v4_plugin + display_name: V4 Plugin + desc: test + author: tester + version: 0.1.0 + runtime: + python: "{sys.version_info.major}.{sys.version_info.minor}" + components: + - class: commands.sample:MyPlugin + type: command + name: hello + description: hello + """ + ), + encoding="utf-8", + ) + (plugin_root / "commands" / "sample.py").write_text( + textwrap.dedent( + """\ + from astrbot_sdk import Context, MessageEvent, Star, on_command + + + class MyPlugin(Star): + @on_command("hello") + async def hello(self, event: MessageEvent, ctx: Context): + reply = await ctx.llm.chat(event.text) + await event.reply(reply) + chunks = [] + async for chunk in ctx.llm.stream_chat("stream"): + chunks.append(chunk) + await event.reply("".join(chunks)) + """ + ), + encoding="utf-8", + ) + + +class RuntimeIntegrationTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self) -> None: + self.left, self.right = make_transport_pair() + self.core = Peer( + transport=self.left, + peer_info=PeerInfo(name="outer-core", role="core", version="v4"), + ) + self.core.set_initialize_handler( + lambda _message: asyncio.sleep(0, result=InitializeOutput( + peer=PeerInfo(name="outer-core", role="core", version="v4"), + capabilities=[], + metadata={}, + )) + ) + await self.core.start() + + async def asyncTearDown(self) -> None: + await self.core.stop() + + async def test_supervisor_runs_v4_plugin_over_stdio_worker(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + plugins_root = Path(temp_dir) / "plugins" + plugin_root = plugins_root / "v4_plugin" + write_new_plugin(plugin_root) + + runtime = SupervisorRuntime( + transport=self.right, + plugins_dir=plugins_root, + env_manager=FakeEnvManager(), + ) + try: + await runtime.start() + await self.core.wait_until_remote_initialized() + handler_id = self.core.remote_handlers[0].id + + await self.core.invoke( + "handler.invoke", + { + "handler_id": handler_id, + "event": { + "text": "hello", + "session_id": "session-1", + "user_id": "user-1", + "platform": "test", + }, + }, + request_id="call-v4", + ) + texts = [item.get("text") for item in runtime.capability_router.sent_messages] + self.assertEqual(texts, ["Echo: hello", "Echo: stream"]) + finally: + await runtime.stop() + + async def test_supervisor_runs_compat_plugin(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + plugins_root = Path(temp_dir) / "plugins" + plugin_root = plugins_root / "compat_plugin" + shutil.copytree(Path.cwd() / "test_plugin", plugin_root) + + runtime = SupervisorRuntime( + transport=self.right, + plugins_dir=plugins_root, + env_manager=FakeEnvManager(), + ) + try: + await runtime.start() + await self.core.wait_until_remote_initialized() + handler_id = self.core.remote_handlers[0].id + + await self.core.invoke( + "handler.invoke", + { + "handler_id": handler_id, + "event": { + "text": "/hello", + "session_id": "session-compat", + "user_id": "user-1", + "platform": "test", + }, + }, + request_id="call-compat", + ) + texts = [item.get("text") for item in runtime.capability_router.sent_messages] + self.assertEqual(len(texts), 4) + self.assertIn("Created conversation ID", texts[0]) + finally: + await runtime.stop()