404 lines
18 KiB
Python
404 lines
18 KiB
Python
"""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.foundation.utils.file_lock import WorkspaceWriteLock, WorkspaceWriteLockBusy
|
|
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.plugins.discovery import discover_plugins
|
|
from beaver.plugins.skills import PluginManager
|
|
from beaver.plugins.state import PluginStateStore
|
|
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
|
|
plugin_manager: PluginManager | None = None
|
|
plugins: list[dict] = field(default_factory=list)
|
|
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,
|
|
plugin_manager: PluginManager | 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._plugin_manager = plugin_manager
|
|
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")
|
|
write_lock = WorkspaceWriteLock(workspace)
|
|
skill_learning_store = self._skill_learning_store or SkillLearningStore(
|
|
workspace / "memory" / "skills",
|
|
write_lock=write_lock,
|
|
)
|
|
|
|
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(),
|
|
)
|
|
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()
|
|
},
|
|
)
|
|
discovery = discover_plugins(workspace, search_paths=self.config.plugins.search_paths)
|
|
plugin_manager = self._plugin_manager or PluginManager(
|
|
workspace=workspace,
|
|
manifests=discovery.manifests,
|
|
discovery_errors=discovery.errors,
|
|
state_store=PluginStateStore(workspace),
|
|
skill_store=skill_spec_store,
|
|
learning_store=skill_learning_store,
|
|
publisher=skill_publisher,
|
|
safety_checker=safety_checker,
|
|
write_lock=write_lock,
|
|
)
|
|
if self.config.plugins.auto_sync:
|
|
try:
|
|
plugin_manager.sync_enabled(blocking=False)
|
|
except WorkspaceWriteLockBusy:
|
|
pass
|
|
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=safety_checker,
|
|
evaluator=SkillDraftEvaluator(run_memory_store),
|
|
publish_observer=plugin_manager.on_skill_published,
|
|
)
|
|
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,
|
|
plugin_manager=plugin_manager,
|
|
plugins=_plugin_summaries(plugin_manager),
|
|
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())
|
|
|
|
|
|
def _plugin_summaries(manager: PluginManager) -> list[dict]:
|
|
summaries: list[dict] = []
|
|
for state in manager.list_plugins():
|
|
manifest = manager.manifests.get(state.plugin_id)
|
|
summaries.append(
|
|
{
|
|
"id": state.plugin_id,
|
|
"name": manifest.name if manifest is not None else state.plugin_id,
|
|
"discovered_version": manifest.version if manifest is not None else None,
|
|
"installed_version": state.installed_version,
|
|
"enabled": state.enabled,
|
|
"status": state.status,
|
|
"last_error": state.last_error,
|
|
"manifest_path": manifest.display_path if manifest is not None else state.manifest_path,
|
|
"updates_paused": state.updates_paused,
|
|
"skills": [
|
|
{
|
|
"name": name,
|
|
"status": binding.status,
|
|
"current_beaver_version": binding.current_beaver_version,
|
|
"accepted_upstream_tree_hash": binding.accepted_upstream_tree_hash,
|
|
"observed_upstream_tree_hash": binding.observed_upstream_tree_hash,
|
|
"accepted_beaver_version": binding.accepted_beaver_version,
|
|
"pending_candidate_id": binding.pending_candidate_id,
|
|
}
|
|
for name, binding in sorted(state.skills.items())
|
|
],
|
|
}
|
|
)
|
|
return summaries
|