import asyncio from collections.abc import Awaitable, Callable from contextlib import contextmanager from contextvars import ContextVar import dataclasses import inspect from typing import TYPE_CHECKING, Any, Generic, cast, get_origin from typing_extensions import TypeVar import uuid if TYPE_CHECKING: from zhenxun.services.ai.context.memory.types import SessionMetadata from nonebot.adapters import Bot, Event from nonebot.matcher import Matcher, current_bot, current_event, current_matcher from pydantic import BaseModel, ConfigDict, Field from zhenxun.services.ai.capabilities.base import AbstractCapability from zhenxun.services.ai.core.engine.append_only import AppendOnlyContextManager from zhenxun.services.ai.core.messages import AgentEvent, AgentMessage from zhenxun.services.ai.core.models import CancellationToken from zhenxun.services.ai.core.protocols.tool import ToolExecutable from zhenxun.services.ai.core.stream_events import AgentStreamEvent, EventBus from zhenxun.services.ai.utils import ContextUtils from zhenxun.services.scheduler.types import ScheduleContext from zhenxun.utils.platform import PlatformUtils from zhenxun.utils.utils import infer_plugin_namespace from .blackboard import BlackboardManager AgentDepsT = TypeVar("AgentDepsT", default=Any) """泛型类型变量：外部环境依赖对象 (Agent Dependencies)。""" ProviderFunc = Callable[["RunContext"], Any | Awaitable[Any]] """函数签名类型别名：依赖提供者函数 (Dependency Provider)。""" ToolsPrepareFunc = Callable[["RunContext[AgentDepsT]", list[Any]], Any | Awaitable[Any]] """全局/Agent 级动态工具干预函数类型""" class NoneBotDeps(BaseModel): """NoneBot 环境下的标准依赖容器。""" model_config = ConfigDict(arbitrary_types_allowed=True) bot: Bot | None = Field(default=None) """当前触发事件的 Bot 实例""" event: Event | None = Field(default=None) """当前触发的事件实例""" matcher: Matcher | None = Field(default=None) """当前处理该事件的 Matcher 实例""" @classmethod def get_current(cls) -> "NoneBotDeps | None": """利用 NoneBot 原生魔法，基于 ContextVars 隐式提取当前执行上下文""" bot = current_bot.get(None) event = current_event.get(None) matcher = current_matcher.get(None) if bot or event: return cls(bot=bot, event=event, matcher=matcher) return None class ScheduledDeps(BaseModel): """后台/定时任务环境下的标准依赖容器。""" model_config = ConfigDict(arbitrary_types_allowed=True) bot: Bot | None = Field(default=None) """机器人实例""" group_id: str | None = Field(default=None) """群组/频道 ID""" user_id: str | None = Field(default=None) """用户 ID""" schedule_id: int | str | None = Field(default=None) """定时任务/调度任务 ID""" platform: str | None = Field(default=None) """适配器平台标识""" @classmethod def from_schedule_context( cls, bot: Bot, context: ScheduleContext ) -> "ScheduledDeps": """从调度器上下文中快速构造并提取依赖""" return cls( bot=bot, group_id=getattr(context, "group_id", None), user_id=getattr(context, "user_id", None), schedule_id=getattr(context, "schedule_id", None), platform=getattr(context, "platform_scope", None) or PlatformUtils.get_platform_scope(bot), ) @dataclasses.dataclass class SessionContext(Generic[AgentDepsT]): """ 生命周期：Session（会话）层。 跨越多次对话轮次，负责保存长线状态与物理隔离信息。 """ session_id: str """核心会话标识符，用于区分不同用户或群组的上下文隔离。""" deps: AgentDepsT """强类型的外部依赖注入对象（如 Bot, Event），供跨工具共享。""" shared_state: dict[str, Any] = dataclasses.field(default_factory=dict) """共享状态字典：全局引用穿透，用于主智能体与嵌套子智能体之间的数据通信。""" blackboard: BlackboardManager | None = None """结构化黑板管理器，作为共享状态的高级替代方案，提供并发锁和强类型校验。""" namespace: str = "global" """触发事件的插件命名空间""" append_only_manager: Any = dataclasses.field(default=None) """用于大模型前缀缓存命中优化的追加写入管理器。""" session_meta: "SessionMetadata | None" = dataclasses.field(default=None) """隔离会话的元信息(Session ID, 命名空间, 权限等)，在状态流转中生成""" @dataclasses.dataclass class AgentRunContext(Generic[AgentDepsT]): """ 生命周期：Run（运行）层。 伴随 Agent 的单次执行 (run_stream)，保存大模型推理时的状态与原生消息历史。 """ session: SessionContext[AgentDepsT] """指向所属 Session 层上下文的引用。""" state: dict[str, Any] = dataclasses.field(default_factory=dict) """状态字典：用于在当前 Agent 执行轮次、工具和中间件中透传动态变量。""" agent_name: str | None = None """当前正在执行的 Agent 名称。""" current_model: str | None = None """当前实际调用的底层大模型名称 (Provider/Model)。""" user_input: str | None = None """当前轮次用户的原始文本输入。""" messages: list[AgentMessage] = dataclasses.field(default_factory=list) """大模型原生上下文 (LLMMessage 列表)，与执行器中的执行历史保持内存引用同步。""" hitl_locks: dict[str, asyncio.Lock] = dataclasses.field(default_factory=dict) """人机交互 (HITL) 并发锁，防止同群组内并发审批冲突。""" delegate_depth: int = 0 """子智能体委派深度标记，用于防范无限递归嵌套。""" tool_retries: dict[str, int] = dataclasses.field(default_factory=dict) """记录当前轮次内各个工具的累积失败重试次数，用于系统熔断。""" cancellation_token: CancellationToken | None = None """全局级联取消令牌，用于跨 Agent 的协程挂起中断。""" event_bus: EventBus | None = None """底层的事件流发射器 (EventBus)，由执行引擎在运行时挂载。""" dynamic_prompts: dict[str, str] = dataclasses.field(default_factory=dict) """动态提示词字典（保持插入顺序并去重）。 仅在 HTTP 请求前 JIT 渲染，不会污染持久化的上下文对话历史。""" def add_system_prompt(self, prompt: str, key: str | None = None) -> None: """动态追加临时系统提示词到大模型上下文中（实时生效且不污染历史）。""" dict_key = key or prompt if prompt: self.dynamic_prompts[dict_key] = prompt def add_event(self, event: AgentEvent) -> None: """向当前运行上下文中安全追加业务事件""" self.messages.append(event) async def emit(self, event: AgentStreamEvent) -> None: """安全地向事件总线发射流式事件（内置判空逻辑，消解业务层的样板代码）""" if self.event_bus: await self.event_bus.emit(event) @dataclasses.dataclass class ToolCallContext(Generic[AgentDepsT]): """ 生命周期：Call（工具调用）层。 单次工具调用分配的绝对私有状态，彻底消灭并发调用时的属性污染。 """ run: AgentRunContext[AgentDepsT] """指向所属 Run 层上下文的引用。""" tool_call_id: str = "unknown" """大模型为本次工具调用分配的唯一 ID。""" tool_name: str = "unknown" """本次调用的工具名称。""" retry_count: int = 0 """当前工具调用的重试序号 (第几次重试)。""" current_tool: ToolExecutable | None = None """当前工具的可执行实例 (ToolExecutable)。""" @dataclasses.dataclass class RunContext(Generic[AgentDepsT]): """ 依赖注入容器（DI Container），保留原有上下文信息的同时提升获取类型的能力。 """ session_id: str | None = None """当前运行所在的会话ID，用于区分不同用户的独立上下文""" di_cache: dict[str, Any] = dataclasses.field(default_factory=dict) """依赖注入引擎的缓存容器，支持父子层级浅拷贝隔离""" state: dict[str, Any] = dataclasses.field(default_factory=dict) """状态字典：用于在会话轮次、工具和中间件中透传动态变量""" shared_state: dict[str, Any] = dataclasses.field(default_factory=dict) """共享状态字典：全局引用穿透，用于主智能体与嵌套子智能体之间的 黑板模式 (Blackboard) 数据通信""" upstream_results: dict[str, Any] = dataclasses.field(default_factory=dict) """前置节点产出字典：标准化的数据流载荷契约，键为 Agent Name，值为输出内容""" capabilities: list[AbstractCapability] = dataclasses.field(default_factory=list) """当前上下文绑定的拦截器 (Capabilities) 链，用于生命周期拦截""" deps: AgentDepsT = dataclasses.field(default=cast(AgentDepsT, None)) """强类型的外部依赖注入对象，用于跨工具共享业务状态或配置""" session: SessionContext[AgentDepsT] = dataclasses.field( init=False, repr=False, compare=False ) """会话层上下文：承载跨轮次共享的依赖与共享状态引用。""" run: AgentRunContext[AgentDepsT] = dataclasses.field( init=False, repr=False, compare=False ) """运行层上下文：承载当前 Agent 执行轮次的模型状态与运行时元信息。""" call: ToolCallContext[AgentDepsT] = dataclasses.field( init=False, repr=False, compare=False ) """调用层上下文：承载单次工具调用的私有状态与执行引用。""" _is_auto_session_id: bool = dataclasses.field( default=False, init=False, repr=False, compare=False ) """标记 session_id 是否为框架隐式生成的。""" @classmethod def from_schedule( cls, bot: Bot, context: ScheduleContext, **kwargs ) -> "RunContext[ScheduledDeps]": """极简构造语法糖：从定时调度任务上下文中直接生成 RunContext""" deps = ScheduledDeps.from_schedule_context(bot, context) return cast("RunContext[ScheduledDeps]", cls(deps=cast(Any, deps), **kwargs)) def get_bot(self) -> Bot | None: """强类型安全地提取 Bot 实例""" if not self.deps: return None bot = getattr(self.deps, "bot", None) return bot if isinstance(bot, Bot) else None def get_event(self) -> Event | None: """强类型安全地提取 Event 实例""" if not self.deps: return None event = getattr(self.deps, "event", None) return event if isinstance(event, Event) else None def get_matcher(self) -> Matcher | None: """强类型安全地提取 Matcher 实例""" if not self.deps: return None matcher = getattr(self.deps, "matcher", None) return matcher if isinstance(matcher, Matcher) else None def get_user_id(self) -> str | None: """安全提取当前触发任务的用户 ID""" return ContextUtils.extract_user_id(self.deps) def get_group_id(self) -> str | None: """安全提取当前触发任务的群组 ID（私聊则为 None）""" return ContextUtils.extract_group_id(self.deps) def get_platform(self) -> str: """安全提取当前连接的适配器平台标识""" return ContextUtils.extract_platform(self.deps) def __post_init__(self): if self.deps is None: self.deps = cast(AgentDepsT, NoneBotDeps.get_current()) if not self.session_id and self.deps: from zhenxun.services.ai.context.memory.types import ( Isolation, ) bot = self.get_bot() event = self.get_event() if bot: meta = ContextUtils.generate_session_meta( bot=bot, event=event, deps=self.deps, scope_builder=Isolation.AGENT_USER(), ) self.session_id = meta.session_id self._is_auto_session_id = True ns = "global" if self.deps: ns = getattr(self.deps, "namespace", None) or ( self.deps.get("namespace") if isinstance(self.deps, dict) else None ) if not ns: ns = infer_plugin_namespace() self.session = SessionContext( session_id=self.session_id or "default_session", deps=self.deps, shared_state=self.shared_state, namespace=ns, ) self.session.append_only_manager = AppendOnlyContextManager() self.run = AgentRunContext(session=self.session, state=self.state) self.call = ToolCallContext(run=self.run) def clone_for_execution(self, **kwargs) -> "RunContext": new_state = self.state.copy() changes = {k: v for k, v in kwargs.items() if hasattr(self, k)} if "state" not in changes: changes["state"] = new_state if "shared_state" not in changes: changes["shared_state"] = self.shared_state new_ctx = dataclasses.replace(cast(Any, self), **changes) new_ctx.upstream_results = self.upstream_results.copy() new_ctx.di_cache = self.di_cache.copy() new_ctx.session = self.session new_ctx.run = AgentRunContext( session=new_ctx.session, state=new_ctx.state, agent_name=self.run.agent_name, current_model=self.run.current_model, user_input=self.run.user_input, messages=list(self.run.messages), hitl_locks=self.run.hitl_locks, delegate_depth=self.run.delegate_depth, tool_retries=self.run.tool_retries, cancellation_token=self.run.cancellation_token, event_bus=self.run.event_bus, dynamic_prompts=self.run.dynamic_prompts.copy(), ) new_ctx.call = ToolCallContext(run=new_ctx.run) return new_ctx def clone_for_tool_call(self, tool_call_id: str, tool_name: str) -> "RunContext": """ 为每个并发的工具调用派生绝对独立的 ToolCallContext。消除多工具并行时的属性污染。 """ new_ctx = dataclasses.replace(cast(Any, self)) new_ctx.di_cache = self.di_cache.copy() new_ctx.session = self.session new_ctx.run = self.run new_ctx.call = ToolCallContext( run=new_ctx.run, tool_call_id=tool_call_id, tool_name=tool_name, retry_count=new_ctx.run.tool_retries.get(tool_name, 0), ) return new_ctx def clone_for_member(self, member_name: str = "unknown") -> "RunContext": """ 为团队子成员克隆上下文。 强制清空消息历史并分配独立 SessionID，实现绝对的记忆沙箱物理隔离。 """ new_ctx = self.clone_for_execution() new_ctx.run.delegate_depth = 0 new_ctx.run.messages = [] new_ctx.capabilities = [] base_sid = self.session_id or "default" new_ctx.session_id = f"{base_sid}/sub_{member_name}_{uuid.uuid4().hex[:6]}" new_ctx.session.session_id = new_ctx.session_id return new_ctx _CURRENT_RUN_CONTEXT: ContextVar[RunContext | None] = ContextVar( "current_run_context", default=None ) def get_current_run_context() -> RunContext | None: """ [全局逃生舱] 获取当前运行中的上下文对象。 适用于深层嵌套业务逻辑，无需层层透传 context 参数。 """ return _CURRENT_RUN_CONTEXT.get() @contextmanager def set_run_context(ctx: RunContext): """[内部 API] 挂载当前上下文至全局""" token = _CURRENT_RUN_CONTEXT.set(ctx) try: yield finally: _CURRENT_RUN_CONTEXT.reset(token) def _is_run_context_type(annotation: Any) -> bool: """判断类型标注是否为 RunContext 类型或其子类。""" if annotation is RunContext: return True origin = get_origin(annotation) if origin is RunContext: return True if inspect.isclass(origin) and issubclass(origin, RunContext): return True if inspect.isclass(annotation) and issubclass(annotation, RunContext): return True if "RunContext" in str(annotation): return True return False __all__ = [ "AgentDepsT", "AgentRunContext", "NoneBotDeps", "RunContext", "ScheduledDeps", "SessionContext", "ToolCallContext", "ToolsPrepareFunc", "get_current_run_context", "set_run_context", ]