fz-mirrors/AstrBot
1
0
Fork 0
mirror of https://github.com/AstrBotDevs/AstrBot synced 2026-07-15 17:30:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: Implement Peer and Transport layers for asynchronous communication

Browse Source 
- Added Peer class for managing peer-to-peer communication, including initialization, invocation, and cancellation of requests.
- Introduced Transport abstract base class with StdioTransport and WebSocketTransport implementations for message handling.
- Created Star class for error handling in event-driven architecture.
- Developed MemoryTransport for testing purposes, allowing in-memory communication between peers.
- Implemented unit tests for Peer functionality, protocol message parsing, and runtime integration with plugins.
- Established entry point tests to ensure package import and command-line interface functionality.
This commit is contained in:
whatevertogo
2026-03-12 21:42:46 +08:00
parent b7d4b647be
commit 9d04cec8a4
43 changed files with 4664 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
CLAUDE.md Normal file
View File

@@ -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.

10
pyproject.toml
View File

@@ -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"]

1385
refactor.md Normal file
View File

File diff suppressed because it is too large Load Diff

BIN
src-new.rar Normal file
View File

Binary file not shown.

17
src-new/astrbot_sdk/__init__.py Normal file
View File

@@ -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",
]

5
src-new/astrbot_sdk/__main__.py Normal file
View File

@@ -0,0 +1,5 @@
from .cli import cli
if __name__ == "__main__":
cli()

1
src-new/astrbot_sdk/api/__init__.py Normal file
View File

@@ -0,0 +1 @@
__all__: list[str] = []

3
src-new/astrbot_sdk/api/components/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
from .command import CommandComponent
__all__ = ["CommandComponent"]

3
src-new/astrbot_sdk/api/components/command.py Normal file
View File

@@ -0,0 +1,3 @@
from ...compat import CommandComponent
__all__ = ["CommandComponent"]

4
src-new/astrbot_sdk/api/event/__init__.py Normal file
View File

@@ -0,0 +1,4 @@
from ...events import MessageEvent as AstrMessageEvent
from .filter import ADMIN, filter
__all__ = ["ADMIN", "AstrMessageEvent", "filter"]

32
src-new/astrbot_sdk/api/event/filter.py Normal file
View File

@@ -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()

3
src-new/astrbot_sdk/api/star/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
from .context import Context
__all__ = ["Context"]

3
src-new/astrbot_sdk/api/star/context.py Normal file
View File

@@ -0,0 +1,3 @@
from ...compat import Context
__all__ = ["Context"]

56
src-new/astrbot_sdk/cli.py Normal file
View File

@@ -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="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
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))

13
src-new/astrbot_sdk/clients/__init__.py Normal file
View File

@@ -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",
]

40
src-new/astrbot_sdk/clients/_proxy.py Normal file
View File

@@ -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

25
src-new/astrbot_sdk/clients/db.py Normal file
View File

@@ -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", [])]

77
src-new/astrbot_sdk/clients/llm.py Normal file
View File

@@ -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", ""))

20
src-new/astrbot_sdk/clients/memory.py Normal file
View File

@@ -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})

26
src-new/astrbot_sdk/clients/platform.py Normal file
View File

@@ -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", []))

143
src-new/astrbot_sdk/compat.py Normal file
View File

@@ -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

51
src-new/astrbot_sdk/context.py Normal file
View File

@@ -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()

98
src-new/astrbot_sdk/decorators.py Normal file
View File

@@ -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

85
src-new/astrbot_sdk/errors.py Normal file
View File

@@ -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)),
)

67
src-new/astrbot_sdk/events.py Normal file
View File

@@ -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)

21
src-new/astrbot_sdk/protocol/__init__.py Normal file
View File

@@ -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",
]

69
src-new/astrbot_sdk/protocol/descriptors.py Normal file
View File

@@ -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

90
src-new/astrbot_sdk/protocol/legacy_adapter.py Normal file
View File

@@ -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},
}

96
src-new/astrbot_sdk/protocol/messages.py Normal file
View File

@@ -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}")

1
src-new/astrbot_sdk/runtime/__init__.py Normal file
View File

@@ -0,0 +1 @@
__all__: list[str] = []

352
src-new/astrbot_sdk/runtime/bootstrap.py Normal file
View File

@@ -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()

330
src-new/astrbot_sdk/runtime/capability_router.py Normal file
View File

@@ -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}")

102
src-new/astrbot_sdk/runtime/handler_dispatcher.py Normal file
View File

@@ -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)

307
src-new/astrbot_sdk/runtime/loader.py Normal file
View File

@@ -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)

346
src-new/astrbot_sdk/runtime/peer.py Normal file
View File

@@ -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))

285
src-new/astrbot_sdk/runtime/transport.py Normal file
View File

@@ -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()

26
src-new/astrbot_sdk/star.py Normal file
View File

@@ -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

1
tests_v4/__init__.py Normal file
View File

@@ -0,0 +1 @@
__all__: list[str] = []

40
tests_v4/helpers.py Normal file
View File

@@ -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)

37
tests_v4/test_entrypoints.py Normal file
View File

@@ -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)

189
tests_v4/test_peer.py Normal file
View File

@@ -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()

53
tests_v4/test_protocol.py Normal file
View File

@@ -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)

149
tests_v4/test_runtime.py Normal file
View File

@@ -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()