"""Centralized runtime loading for Beaver agents.""" from __future__ import annotations import asyncio import os from dataclasses import dataclass, field from pathlib import Path from typing import Callable from beaver.coordinator.registry import AgentRegistry from beaver.engine.context import ContextBuilder from beaver.engine.session import SessionManager from beaver.foundation.config import BeaverConfig, load_config from beaver.integrations.mcp import MCPConnectionManager from beaver.memory.curated.store import MemoryStore from beaver.memory.runs import RunMemoryStore from beaver.memory.skills import SkillLearningStore from beaver.services.memory_service import MemoryService from beaver.skills.drafts import DraftService from beaver.skills.learning import EvidenceSelector, SkillDraftSynthesizer, SkillLearningPipelineService, SkillLearningService from beaver.skills.learning.safety import SkillDraftSafetyChecker from beaver.skills.learning.eval import SkillDraftEvaluator from beaver.skills.publisher import SkillPublisher from beaver.skills.reviews import ReviewService from beaver.skills.specs import SkillSpecStore from beaver.tasks import TaskExecutionPlanner, TaskService from beaver.tasks.skill_resolver import TaskSkillResolver from beaver.skills import SkillAssembler, SkillsLoader from beaver.tools import ObjectBackedTool, ToolAssembler, ToolExecutor, ToolRegistry from beaver.tools.builtins import ( ClarifyTool, CronTool, DelegateTool, EchoTool, ExecuteCodeTool, ListDirectoryTool, MemoryTool, PatchFileTool, ProcessTool, ReadFileTool, SearchFilesTool, SendMessageTool, SpawnTool, SessionSearchTool, SkillManageTool, SkillViewTool, SkillsListTool, TerminalTool, TodoTool, UserFilesCopyToWorkspaceTool, UserFilesListTool, UserFilesMkdirTool, UserFilesPublishOutputTool, UserFilesReadTool, UserFilesWriteTool, WebFetchTool, WebSearchTool, WriteFileTool, ) @dataclass(slots=True) class EngineLoadResult: """描述当前 agent runtime 已经装好的依赖。 这里同时保留两类字段： 1. `tools/skills/memory_stores/permissions` - 便于做状态展示、调试、轻量测试 2. `session_manager/tool_registry/...` - 供真正的运行时主链直接使用 """ workspace: Path config: BeaverConfig = field(default_factory=BeaverConfig) tools: list[str] = field(default_factory=list) skills: list[str] = field(default_factory=list) memory_stores: list[str] = field(default_factory=list) permissions: list[str] = field(default_factory=list) session_manager: SessionManager | None = None curated_memory_store: MemoryStore | None = None memory_service: MemoryService | None = None run_memory_store: RunMemoryStore | None = None skill_learning_store: SkillLearningStore | None = None tool_registry: ToolRegistry | None = None tool_assembler: ToolAssembler | None = None tool_executor: ToolExecutor | None = None context_builder: ContextBuilder | None = None skills_loader: SkillsLoader | None = None skill_assembler: SkillAssembler | None = None skill_spec_store: SkillSpecStore | None = None draft_service: DraftService | None = None review_service: ReviewService | None = None skill_publisher: SkillPublisher | None = None skill_learning_service: SkillLearningService | None = None skill_learning_pipeline: SkillLearningPipelineService | None = None agent_registry: AgentRegistry | None = None task_skill_resolver: TaskSkillResolver | None = None task_service: TaskService | None = None task_execution_planner: TaskExecutionPlanner | None = None mcp_manager: MCPConnectionManager | None = None mcp_report: dict[str, dict] = field(default_factory=dict) closeables: list[tuple[str, Callable[[], None]]] = field(default_factory=list, repr=False) closed: bool = False def register_closeable(self, name: str, close_fn: Callable[[], None]) -> None: """登记一个由 runtime 统一关闭的资源。""" self.closeables.append((name, close_fn)) def close(self) -> None: """按后进先出顺序关闭 runtime 资源。 这一步先保持同步、最小、可组合： 1. 只管理已经明确需要关闭的资源 2. 暂不引入 async shutdown 协议 3. 为后续 Web/Gateway lifespan 留统一入口 """ if self.closed: return errors: list[tuple[str, BaseException]] = [] for name, close_fn in reversed(self.closeables): try: close_fn() except BaseException as exc: # pragma: no cover - defensive cleanup path errors.append((name, exc)) self.closed = True if errors: parts = ", ".join(f"{name}: {exc}" for name, exc in errors) raise RuntimeError(f"Runtime shutdown failed for {parts}") class EngineLoader: """为任意 Beaver agent 装载共享 runtime 能力。 当前先做“最小可运行主链”需要的装配： - session manager - curated memory store - context builder - built-in tools - tool executor 等主链跑稳后，再把 skills、权限、MCP、delegation 逐步加进来。 """ def __init__( self, *, workspace: str | Path | None = None, config_path: str | Path | None = None, config: BeaverConfig | None = None, session_manager: SessionManager | None = None, curated_memory_store: MemoryStore | None = None, memory_service: MemoryService | None = None, run_memory_store: RunMemoryStore | None = None, skill_learning_store: SkillLearningStore | None = None, tool_registry: ToolRegistry | None = None, tool_assembler: ToolAssembler | None = None, context_builder: ContextBuilder | None = None, skills_loader: SkillsLoader | None = None, skill_assembler: SkillAssembler | None = None, skill_spec_store: SkillSpecStore | None = None, draft_service: DraftService | None = None, review_service: ReviewService | None = None, skill_publisher: SkillPublisher | None = None, skill_learning_service: SkillLearningService | None = None, skill_learning_pipeline: SkillLearningPipelineService | None = None, agent_registry: AgentRegistry | None = None, task_skill_resolver: TaskSkillResolver | None = None, task_service: TaskService | None = None, task_execution_planner: TaskExecutionPlanner | None = None, ) -> None: self.config = config or load_config(workspace=workspace, config_path=config_path) configured_workspace = self.config.agents_defaults.workspace env_workspace = os.getenv("BEAVER_WORKSPACE") self.workspace = Path(workspace or configured_workspace or env_workspace or Path.cwd()) self._session_manager = session_manager self._curated_memory_store = curated_memory_store self._memory_service = memory_service self._run_memory_store = run_memory_store self._skill_learning_store = skill_learning_store self._tool_registry = tool_registry self._tool_assembler = tool_assembler self._context_builder = context_builder self._skills_loader = skills_loader self._skill_assembler = skill_assembler self._skill_spec_store = skill_spec_store self._draft_service = draft_service self._review_service = review_service self._skill_publisher = skill_publisher self._skill_learning_service = skill_learning_service self._skill_learning_pipeline = skill_learning_pipeline self._agent_registry = agent_registry self._task_skill_resolver = task_skill_resolver self._task_service = task_service self._task_execution_planner = task_execution_planner def load(self) -> EngineLoadResult: """装配当前主链需要的最小 runtime 对象。""" workspace = self.workspace session_manager = self._session_manager or SessionManager(workspace) curated_root = workspace / "memory" / "curated" curated_memory_store = self._curated_memory_store or MemoryStore(curated_root) memory_service = self._memory_service or MemoryService(curated_root, store=curated_memory_store) memory_service.initialize() run_memory_store = self._run_memory_store or RunMemoryStore(workspace / "memory" / "runs") skill_learning_store = self._skill_learning_store or SkillLearningStore(workspace / "memory" / "skills") tool_registry = self._tool_registry or ToolRegistry() skill_spec_store = self._skill_spec_store or SkillSpecStore(workspace) skills_loader = self._skills_loader or SkillsLoader(workspace, skill_store=skill_spec_store) if self._tool_registry is None: # 这里先注册最小工具集，满足主链的 tool loop。 tool_registry.register_many( [ ObjectBackedTool(EchoTool()), ObjectBackedTool(MemoryTool(store=memory_service.get_store())), ObjectBackedTool(SessionSearchTool(db=session_manager)), ObjectBackedTool(ListDirectoryTool()), ObjectBackedTool(ReadFileTool()), ObjectBackedTool(SearchFilesTool()), ObjectBackedTool(WriteFileTool()), ObjectBackedTool(PatchFileTool()), ObjectBackedTool(UserFilesListTool()), ObjectBackedTool(UserFilesReadTool()), ObjectBackedTool(UserFilesWriteTool()), ObjectBackedTool(UserFilesMkdirTool()), ObjectBackedTool(UserFilesCopyToWorkspaceTool()), ObjectBackedTool(UserFilesPublishOutputTool()), ObjectBackedTool(WebFetchTool()), ObjectBackedTool(WebSearchTool()), ObjectBackedTool(TerminalTool()), ObjectBackedTool(ProcessTool()), ObjectBackedTool(ExecuteCodeTool()), ObjectBackedTool(TodoTool()), ObjectBackedTool(ClarifyTool()), ObjectBackedTool(SendMessageTool()), ObjectBackedTool(DelegateTool()), ObjectBackedTool(SpawnTool()), SkillsListTool(), ObjectBackedTool(SkillViewTool(loader=skills_loader)), SkillManageTool(), CronTool(), ] ) context_builder = self._context_builder or ContextBuilder() tool_assembler = self._tool_assembler or ToolAssembler() tool_executor = ToolExecutor(tool_registry) skill_assembler = self._skill_assembler or SkillAssembler(skills_loader) draft_service = self._draft_service or DraftService(skill_spec_store) review_service = self._review_service or ReviewService(skill_spec_store) skill_publisher = self._skill_publisher or SkillPublisher(skill_spec_store) evidence_selector = EvidenceSelector(run_memory_store, session_manager=session_manager) skill_learning_service = self._skill_learning_service or SkillLearningService( run_store=run_memory_store, learning_store=skill_learning_store, draft_service=draft_service, evidence_selector=evidence_selector, synthesizer=SkillDraftSynthesizer(), ) skill_learning_pipeline = self._skill_learning_pipeline or SkillLearningPipelineService( learning_store=skill_learning_store, learning_service=skill_learning_service, draft_service=draft_service, review_service=review_service, publisher=skill_publisher, safety_checker=SkillDraftSafetyChecker( allowed_tool_names={spec.name for spec in tool_registry.list_specs()}, allowed_tool_prefixes={ f"mcp_{server_id}_" for server_id in self.config.tools.mcp_servers if str(server_id).strip() }, ), evaluator=SkillDraftEvaluator(run_memory_store), ) agent_registry = self._agent_registry or AgentRegistry(workspace) task_skill_resolver = self._task_skill_resolver or TaskSkillResolver( skills_loader=skills_loader, draft_service=draft_service, ) task_service = self._task_service or TaskService(workspace / "tasks") task_execution_planner = self._task_execution_planner or TaskExecutionPlanner(task_skill_resolver=task_skill_resolver) mcp_manager = MCPConnectionManager( self.config.tools.mcp_servers, authz_config=self.config.authz, backend_identity=self.config.backend_identity, ) result = EngineLoadResult( workspace=workspace, config=self.config, tools=[spec.name for spec in tool_registry.list_specs()], skills=[record.name for record in skills_loader.list_skills(filter_unavailable=False)], memory_stores=["curated"], permissions=[], session_manager=session_manager, curated_memory_store=memory_service.get_store(), memory_service=memory_service, run_memory_store=run_memory_store, skill_learning_store=skill_learning_store, tool_registry=tool_registry, tool_assembler=tool_assembler, tool_executor=tool_executor, context_builder=context_builder, skills_loader=skills_loader, skill_assembler=skill_assembler, skill_spec_store=skill_spec_store, draft_service=draft_service, review_service=review_service, skill_publisher=skill_publisher, skill_learning_service=skill_learning_service, skill_learning_pipeline=skill_learning_pipeline, agent_registry=agent_registry, task_skill_resolver=task_skill_resolver, task_service=task_service, task_execution_planner=task_execution_planner, mcp_manager=mcp_manager, ) if self._session_manager is None: result.register_closeable("session_manager", session_manager.close) result.register_closeable("mcp_manager", lambda: _close_mcp_manager(mcp_manager)) return result def _close_mcp_manager(manager: MCPConnectionManager) -> None: try: loop = asyncio.get_running_loop() except RuntimeError: asyncio.run(manager.close()) return loop.create_task(manager.close())