From 2e8e378a8cd365dfa1e2511dc13db8d216f0cd86 Mon Sep 17 00:00:00 2001 From: CCLCK <84182037+CCLCK@users.noreply.github.com> Date: Sun, 12 Apr 2026 17:55:08 +0800 Subject: [PATCH] Add shared MCP daemon proxy runtime --- docs/global-codex-integration.md | 48 ++ openspace/cloud/embedding.py | 27 ++ openspace/cloud/search.py | 37 +- openspace/evolution_mcp_server.py | 43 +- openspace/mcp_proxy.py | 255 ++++++++++ openspace/mcp_server.py | 185 +++++++- openspace/mcp_tool_registration.py | 254 ++++++++++ openspace/shared_mcp_runtime.py | 519 +++++++++++++++++++++ openspace/skill_engine/skill_ranker.py | 37 ++ scripts/codex-openspace | 28 +- scripts/install-global-codex-openspace | 12 +- tests/conftest.py | 74 +++ tests/test_embedding_cache_optimization.py | 163 +++++++ tests/test_mcp_entrypoints.py | 104 +++++ tests/test_mcp_http_connector_transport.py | 235 ++++++++++ tests/test_mcp_proxy_runtime.py | 181 +++++++ tests/test_shared_mcp_runtime_metadata.py | 148 ++++++ 17 files changed, 2315 insertions(+), 35 deletions(-) create mode 100644 openspace/mcp_proxy.py create mode 100644 openspace/mcp_tool_registration.py create mode 100644 openspace/shared_mcp_runtime.py create mode 100644 tests/conftest.py create mode 100644 tests/test_embedding_cache_optimization.py create mode 100644 tests/test_mcp_entrypoints.py create mode 100644 tests/test_mcp_http_connector_transport.py create mode 100644 tests/test_mcp_proxy_runtime.py create mode 100644 tests/test_shared_mcp_runtime_metadata.py diff --git a/docs/global-codex-integration.md b/docs/global-codex-integration.md index 5c5e221..776b3ae 100644 --- a/docs/global-codex-integration.md +++ b/docs/global-codex-integration.md @@ -54,11 +54,59 @@ The intended machine-wide setup is: - set `OPENSPACE_WORKSPACE` - route project skills to `~/.codex/projects//skills` - include common global skills from `~/.codex/skills` + - call the shared `stdio` proxy entrypoint + - default `openspace` to `OPENSPACE_MCP_PROXY_MODE=daemon` + - default `openspace_evolution` to `OPENSPACE_MCP_PROXY_MODE=daemon` + - place per-instance daemon state under `OPENSPACE_MCP_DAEMON_STATE_DIR` unless an override is already set 3. Global `~/.codex/AGENTS.md` tells Codex: - to prefer project skill routing - to auto-run sidecar evolution for non-trivial repo work - to treat missing `git init` as a repo bootstrap issue +## Daemon / Proxy V1 + +The global and local launchers keep the same wrapper names and the same MCP config shape, but they now sit in front of a shared-daemon topology: + +- Codex still talks to stdio wrapper scripts. +- The wrapper scripts keep the existing command names but route into `openspace.mcp_proxy`. +- Both main and evolution now default to `OPENSPACE_MCP_PROXY_MODE=daemon`. +- The proxy path resolves or starts a per-instance daemon using `OPENSPACE_MCP_DAEMON_STATE_DIR`. +- The daemon owns the long-lived OpenSpace engine and serves it over localhost transport. + +This keeps the external Codex contract stable while reducing the number of overlapping OpenSpace engine processes. + +### Fallbacks + +The proxy surface supports two internal overrides: + +- `OPENSPACE_MCP_PROXY_MODE=direct` restores the old direct stdio behavior for debugging or rollback. +- `OPENSPACE_MCP_DAEMON_STATE_DIR=/custom/path` moves daemon state to a different local directory. + +The repo-local `scripts/codex-openspace` helper writes the same daemon defaults into the generated profile so local and global setups stay aligned. + +### Daemon State Metadata + +Each per-key daemon writes a JSON record under `OPENSPACE_MCP_DAEMON_STATE_DIR` named like: + +- `main-.json` +- `evolution-.json` + +For the main daemon path, the record now distinguishes two lifecycle phases: + +- `ready=true`: the daemon is reachable and `list_tools` has succeeded. +- `warmed=true`: background prewarm has completed, so the local embedding backend and candidate cache are ready. + +Useful timestamps: + +- `started_at`: child process spawn time +- `ready_at`: first confirmed MCP-ready time +- `warmed_at`: prewarm completion time + +This makes it possible to tell the difference between: + +- daemon is up but still warming +- daemon is fully warmed and ready for low-latency calls + ## Reinstalling the Global Wrappers Use: diff --git a/openspace/cloud/embedding.py b/openspace/cloud/embedding.py index 3dd424b..277e851 100644 --- a/openspace/cloud/embedding.py +++ b/openspace/cloud/embedding.py @@ -31,6 +31,7 @@ _OPENAI_BASE = "https://api.openai.com/v1" _VALID_BACKENDS = {"auto", "local", "remote"} _LOCAL_EMBEDDER = None _LOCAL_EMBEDDER_MODEL = None +_EMBEDDING_WARMUP_TEXT = "openspace skill embedding warmup" def resolve_skill_embedding_backend() -> str: @@ -56,6 +57,17 @@ def resolve_skill_embedding_model(backend: Optional[str] = None) -> str: return SKILL_REMOTE_EMBEDDING_MODEL +def using_local_skill_embeddings(backend: Optional[str] = None) -> bool: + """Return whether skill embeddings resolve to the local fastembed path.""" + backend = backend or resolve_skill_embedding_backend() + if backend == "local": + return True + if backend == "remote": + return False + remote_key, _ = _resolve_remote_embedding_api() + return not bool(remote_key) + + def _resolve_remote_embedding_api() -> Tuple[Optional[str], str]: """Resolve remote embedding credentials/base URL for skill routing.""" dedicated_key = os.environ.get("OPENSPACE_SKILL_EMBEDDING_API_KEY") @@ -176,6 +188,21 @@ def _generate_local_embedding(text: str, model_name: str) -> Optional[List[float return None +def prewarm_local_skill_embedding_backend() -> bool: + """Warm the local skill embedding backend when local routing is active. + + Returns True when the local backend is active and the embedder produced + a warmup embedding, False otherwise. + """ + backend = resolve_skill_embedding_backend() + if not using_local_skill_embeddings(backend): + return False + + model_name = resolve_skill_embedding_model(backend) + vector = _generate_local_embedding(_EMBEDDING_WARMUP_TEXT, model_name) + return vector is not None + + def generate_embedding(text: str, api_key: Optional[str] = None) -> Optional[List[float]]: """Generate skill embedding using the configured local/remote backend. diff --git a/openspace/cloud/search.py b/openspace/cloud/search.py index bd25436..046aa81 100644 --- a/openspace/cloud/search.py +++ b/openspace/cloud/search.py @@ -158,6 +158,9 @@ class SkillSearchEngine: ) -> List[Dict[str, Any]]: """Compute hybrid score = vector_score + lexical_boost.""" from openspace.cloud.embedding import cosine_similarity + from openspace.skill_engine.skill_ranker import SkillCandidate, SkillRanker + + ranker: Optional[SkillRanker] = None scored = [] for candidate in candidates: @@ -170,6 +173,32 @@ class SkillSearchEngine: ranking_signal_score = 0.0 if query_embedding: candidate_embedding = candidate.get("_embedding") + if ( + candidate_embedding is None + and candidate.get("source") == "openspace-local" + and candidate.get("_embedding_text") + ): + if ranker is None: + ranker = SkillRanker(enable_cache=True) + + cached = ranker.get_cached_embedding(candidate.get("skill_id", "")) + if cached: + candidate_embedding = cached + else: + skill_candidate = SkillCandidate( + skill_id=candidate.get("skill_id", ""), + name=candidate_name, + description=candidate.get("description", ""), + body="", + metadata=candidate, + ) + skill_candidate.embedding_text = candidate.get("_embedding_text", "") + ranker.prime_candidates([skill_candidate]) + candidate_embedding = skill_candidate.embedding + + if candidate_embedding: + candidate["_embedding"] = candidate_embedding + if candidate_embedding and isinstance(candidate_embedding, list): vector_score = cosine_similarity(query_embedding, candidate_embedding) ranking_signal_score = vector_score @@ -423,14 +452,6 @@ async def hybrid_search_skills( query_embedding: Optional[List[float]] = None try: query_embedding = await asyncio.to_thread(generate_embedding, normalized_query) - if query_embedding: - for candidate in candidates: - if not candidate.get("_embedding") and candidate.get("_embedding_text"): - candidate_embedding = await asyncio.to_thread( - generate_embedding, candidate["_embedding_text"], - ) - if candidate_embedding: - candidate["_embedding"] = candidate_embedding except Exception: pass diff --git a/openspace/evolution_mcp_server.py b/openspace/evolution_mcp_server.py index fe7a6c1..b97bcac 100644 --- a/openspace/evolution_mcp_server.py +++ b/openspace/evolution_mcp_server.py @@ -21,6 +21,7 @@ from datetime import datetime from pathlib import Path from typing import Any, Dict, Iterable, List, Optional +from openspace.mcp_tool_registration import register_evolution_tools class _MCPSafeStdout: """Stdout wrapper: binary (.buffer) -> real stdout, text (.write) -> stderr.""" @@ -504,8 +505,7 @@ async def _register_extra_skill_dirs(openspace, dirs: List[Path]) -> None: await skill_store.sync_from_registry(metas) -@mcp.tool() -async def evolve_from_context( +async def _evolve_from_context_impl( task: str, summary: str, workspace_dir: str | None = None, @@ -662,21 +662,48 @@ async def evolve_from_context( _mark_request_end() +class _DirectEvolutionToolImplementation: + async def evolve_from_context( + self, + task: str, + summary: str, + workspace_dir: str | None = None, + file_paths: list[str] | None = None, + max_skills: int = 3, + skill_dirs: list[str] | None = None, + output_dir: str | None = None, + ) -> str: + return await _evolve_from_context_impl( + task=task, + summary=summary, + workspace_dir=workspace_dir, + file_paths=file_paths, + max_skills=max_skills, + skill_dirs=skill_dirs, + output_dir=output_dir, + ) + + +register_evolution_tools(mcp, _DirectEvolutionToolImplementation()) + + def run_mcp_server() -> None: import argparse parser = argparse.ArgumentParser(description="OpenSpace Evolution MCP Server") - parser.add_argument("--transport", choices=["stdio", "sse"], default="stdio") + parser.add_argument( + "--transport", + choices=["stdio", "sse", "streamable-http"], + default="stdio", + ) parser.add_argument("--port", type=int, default=8080) args = parser.parse_args() - if args.transport == "stdio": + if args.transport == "stdio" or os.environ.get("OPENSPACE_MCP_DAEMON") == "1": _maybe_start_idle_watchdog() - if args.transport == "sse": - mcp.run(transport="sse", sse_params={"port": args.port}) - else: - mcp.run(transport="stdio") + mcp.settings.port = args.port + mcp.run(transport=args.transport) if __name__ == "__main__": diff --git a/openspace/mcp_proxy.py b/openspace/mcp_proxy.py new file mode 100644 index 0000000..e990bb3 --- /dev/null +++ b/openspace/mcp_proxy.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +import argparse +import inspect +import json +import os +from typing import Any + +from mcp.server.fastmcp import FastMCP +from mcp.types import TextContent + +from openspace.grounding.backends.mcp.client import MCPClient +from openspace.mcp_tool_registration import ( + register_evolution_tools, + register_main_tools, +) +from openspace.shared_mcp_runtime import ServerKind, ensure_daemon + + +def _proxy_mode_for(server_kind: ServerKind) -> str: + raw = os.environ.get("OPENSPACE_MCP_PROXY_MODE", "").strip().lower() + if raw in {"daemon", "direct"}: + return raw + return "daemon" + + +def _json_error(error: Any, **extra: Any) -> str: + return json.dumps({"error": str(error), **extra}, ensure_ascii=False) + + +def _extract_text_payload(result: Any) -> str: + text_parts: list[str] = [] + for item in getattr(result, "content", []): + if isinstance(item, TextContent): + text_parts.append(item.text) + continue + text = getattr(item, "text", None) + if text is not None: + text_parts.append(text) + if not text_parts: + raise RuntimeError("Remote MCP tool returned no text payload") + return "\n".join(text_parts) + + +class _RemoteProxyBase: + def __init__(self, server_kind: ServerKind): + self._server_kind = server_kind + self._client: MCPClient | None = None + self._current_url: str | None = None + + async def _get_client(self) -> MCPClient: + record = await ensure_daemon(self._server_kind) + if self._client is not None and self._current_url == record.url: + return self._client + + await self._reset_client() + + self._client = MCPClient( + config={"mcpServers": {"daemon": {"url": record.url}}}, + timeout=10.0, + sse_read_timeout=60 * 60.0, + check_dependencies=False, + ) + self._current_url = record.url + return self._client + + async def _reset_client(self) -> None: + if self._client is not None: + await self._client.close_all_sessions() + self._client = None + self._current_url = None + + async def _call_remote_tool(self, tool_name: str, args: dict[str, Any]) -> str: + for attempt in range(2): + try: + client = await self._get_client() + session = await client.create_session("daemon", auto_initialize=True) + if session is None: + raise RuntimeError("Failed to create daemon MCP session") + result = await session.connector.call_tool(tool_name, args) + return _extract_text_payload(result) + except Exception as exc: + if attempt == 0: + await self._reset_client() + continue + return _json_error(exc, status="error") + return _json_error("Unreachable proxy retry path", status="error") + + +class _MainProxyImplementation(_RemoteProxyBase): + def __init__(self): + super().__init__("main") + + async def execute_task( + self, + task: str, + workspace_dir: str | None = None, + max_iterations: int | None = None, + skill_dirs: list[str] | None = None, + search_scope: str = "all", + ) -> str: + return await self._call_remote_tool( + "execute_task", + { + "task": task, + "workspace_dir": workspace_dir, + "max_iterations": max_iterations, + "skill_dirs": skill_dirs, + "search_scope": search_scope, + }, + ) + + async def search_skills( + self, + query: str, + source: str = "all", + limit: int = 20, + auto_import: bool = True, + ) -> str: + return await self._call_remote_tool( + "search_skills", + { + "query": query, + "source": source, + "limit": limit, + "auto_import": auto_import, + }, + ) + + async def fix_skill( + self, + skill_dir: str, + direction: str, + ) -> str: + return await self._call_remote_tool( + "fix_skill", + { + "skill_dir": skill_dir, + "direction": direction, + }, + ) + + async def upload_skill( + self, + skill_dir: str, + visibility: str = "public", + origin: str | None = None, + parent_skill_ids: list[str] | None = None, + tags: list[str] | None = None, + created_by: str | None = None, + change_summary: str | None = None, + ) -> str: + return await self._call_remote_tool( + "upload_skill", + { + "skill_dir": skill_dir, + "visibility": visibility, + "origin": origin, + "parent_skill_ids": parent_skill_ids, + "tags": tags, + "created_by": created_by, + "change_summary": change_summary, + }, + ) + + +class _EvolutionProxyImplementation(_RemoteProxyBase): + def __init__(self): + super().__init__("evolution") + + async def evolve_from_context( + self, + task: str, + summary: str, + workspace_dir: str | None = None, + file_paths: list[str] | None = None, + max_skills: int = 3, + skill_dirs: list[str] | None = None, + output_dir: str | None = None, + ) -> str: + return await self._call_remote_tool( + "evolve_from_context", + { + "task": task, + "summary": summary, + "workspace_dir": workspace_dir, + "file_paths": file_paths, + "max_skills": max_skills, + "skill_dirs": skill_dirs, + "output_dir": output_dir, + }, + ) + + +def _build_fastmcp(server_kind: ServerKind) -> FastMCP: + kwargs: dict[str, Any] = {} + if "description" in inspect.signature(FastMCP.__init__).parameters: + if server_kind == "main": + kwargs["description"] = ( + "OpenSpace: Unite the Agents. Evolve the Mind. Rebuild the World." + ) + else: + kwargs["description"] = ( + "OpenSpace evolution sidecar: capture reusable skills from host-agent work." + ) + name = "OpenSpace" if server_kind == "main" else "OpenSpace Evolution" + return FastMCP(name, **kwargs) + + +def _run_proxy(server_kind: ServerKind) -> None: + if _proxy_mode_for(server_kind) == "direct": + if server_kind == "main": + from openspace.mcp_server import run_mcp_server + else: + from openspace.evolution_mcp_server import run_mcp_server + run_mcp_server() + return + + parser = argparse.ArgumentParser(description="OpenSpace MCP proxy") + parser.add_argument("--transport", choices=["stdio"], default="stdio") + parser.parse_args() + + mcp = _build_fastmcp(server_kind) + if server_kind == "main": + register_main_tools(mcp, _MainProxyImplementation()) + else: + register_evolution_tools(mcp, _EvolutionProxyImplementation()) + mcp.run(transport="stdio") + + +def run_main_mcp_proxy() -> None: + _run_proxy("main") + + +def run_evolution_mcp_proxy() -> None: + _run_proxy("evolution") + + +def main() -> None: + parser = argparse.ArgumentParser(description="OpenSpace MCP proxy") + parser.add_argument("--kind", choices=["main", "evolution"], required=True) + parser.add_argument("--transport", choices=["stdio"], default="stdio") + args = parser.parse_args() + + # Rebuild argv for the generic runner so direct fallback can reuse legacy entrypoints. + transport = args.transport + os.environ.setdefault("OPENSPACE_MCP_PROXY_MODE", _proxy_mode_for(args.kind)) + import sys + + sys.argv = [sys.argv[0], "--transport", transport] + _run_proxy(args.kind) + + +if __name__ == "__main__": + main() diff --git a/openspace/mcp_server.py b/openspace/mcp_server.py index 512672a..b1f665b 100644 --- a/openspace/mcp_server.py +++ b/openspace/mcp_server.py @@ -27,6 +27,8 @@ import time from pathlib import Path from typing import Any, Dict, List, Optional +from openspace.mcp_tool_registration import register_main_tools +from openspace.shared_mcp_runtime import update_current_daemon_status class _MCPSafeStdout: """Stdout wrapper: binary (.buffer) → real stdout, text (.write) → stderr.""" @@ -128,6 +130,8 @@ _idle_watchdog_started = False _activity_lock = threading.Lock() _active_request_count = 0 _last_activity_at = time.monotonic() +_embedding_prewarm_started = False +_embedding_prewarm_lock = threading.Lock() # Internal state: tracks bot skill directories already registered this session. _registered_skill_dirs: set = set() @@ -268,6 +272,94 @@ def _get_local_skill_registry(): return registry +def _prewarm_main_daemon_skill_embeddings() -> None: + """Warm local skill embeddings and cache local candidate vectors. + + Runs in a background thread for the main daemon path so the first user + request is less likely to pay the full fastembed/model cold-start cost. + """ + try: + from openspace.cloud.embedding import ( + prewarm_local_skill_embedding_backend, + using_local_skill_embeddings, + ) + from openspace.cloud.search import build_local_candidates + from openspace.skill_engine.skill_ranker import SkillCandidate, SkillRanker + + if not using_local_skill_embeddings(): + logger.info("Skipping main daemon embedding prewarm: remote skill embeddings active") + update_current_daemon_status("main", warmed=True) + return + + if not prewarm_local_skill_embedding_backend(): + logger.info("Main daemon embedding prewarm did not initialize a local embedder") + update_current_daemon_status( + "main", + warmed=False, + warmup_error="local embedder did not initialize", + ) + return + + registry = _get_local_skill_registry() + if not registry: + logger.info("Skipping main daemon embedding cache prewarm: no local skill registry") + update_current_daemon_status("main", warmed=True) + return + + candidates = build_local_candidates(registry.list_skills(), store=None) + if not candidates: + logger.info("Skipping main daemon embedding cache prewarm: no local candidates") + update_current_daemon_status("main", warmed=True) + return + + ranker = SkillRanker(enable_cache=True) + skill_candidates: list[SkillCandidate] = [] + for candidate in candidates: + skill_candidate = SkillCandidate( + skill_id=candidate.get("skill_id", ""), + name=candidate.get("name", ""), + description=candidate.get("description", ""), + body="", + metadata=candidate, + ) + skill_candidate.embedding_text = candidate.get("_embedding_text", "") + skill_candidates.append(skill_candidate) + + warmed = ranker.prime_candidates(skill_candidates) + logger.info( + "Main daemon skill embedding prewarm complete: %s/%s local candidates ready", + warmed, + len(skill_candidates), + ) + update_current_daemon_status("main", warmed=True, warmup_error=None) + except Exception as exc: + logger.warning("Main daemon embedding prewarm failed: %s", exc) + update_current_daemon_status("main", warmed=False, warmup_error=str(exc)) + + +def _maybe_start_main_daemon_embedding_prewarm() -> None: + global _embedding_prewarm_started + + if os.environ.get("OPENSPACE_MCP_DAEMON") != "1": + return + if os.environ.get("OPENSPACE_MCP_DISABLE_EMBEDDING_PREWARM", "").strip().lower() in { + "1", + "true", + "yes", + }: + return + + with _embedding_prewarm_lock: + if _embedding_prewarm_started: + return + threading.Thread( + target=_prewarm_main_daemon_skill_embeddings, + name="openspace-main-embedding-prewarm", + daemon=True, + ).start() + _embedding_prewarm_started = True + + def _get_cloud_client(): """Get a OpenSpaceClient instance (raises CloudError if not configured).""" from openspace.cloud.auth import get_openspace_auth @@ -594,9 +686,8 @@ def _maybe_start_idle_watchdog() -> None: _idle_watchdog_started = True -# MCP Tools (4 tools) -@mcp.tool() -async def execute_task( +# MCP tool implementations +async def _execute_task_impl( task: str, workspace_dir: str | None = None, max_iterations: int | None = None, @@ -677,8 +768,7 @@ async def execute_task( _mark_request_end() -@mcp.tool() -async def search_skills( +async def _search_skills_impl( query: str, source: str = "all", limit: int = 20, @@ -783,8 +873,7 @@ async def search_skills( _mark_request_end() -@mcp.tool() -async def fix_skill( +async def _fix_skill_impl( skill_dir: str, direction: str, ) -> str: @@ -904,8 +993,7 @@ async def fix_skill( _mark_request_end() -@mcp.tool() -async def upload_skill( +async def _upload_skill_impl( skill_dir: str, visibility: str = "public", origin: str | None = None, @@ -977,22 +1065,89 @@ async def upload_skill( finally: _mark_request_end() + +class _DirectMainToolImplementation: + async def execute_task( + self, + task: str, + workspace_dir: str | None = None, + max_iterations: int | None = None, + skill_dirs: list[str] | None = None, + search_scope: str = "all", + ) -> str: + return await _execute_task_impl( + task=task, + workspace_dir=workspace_dir, + max_iterations=max_iterations, + skill_dirs=skill_dirs, + search_scope=search_scope, + ) + + async def search_skills( + self, + query: str, + source: str = "all", + limit: int = 20, + auto_import: bool = True, + ) -> str: + return await _search_skills_impl( + query=query, + source=source, + limit=limit, + auto_import=auto_import, + ) + + async def fix_skill( + self, + skill_dir: str, + direction: str, + ) -> str: + return await _fix_skill_impl(skill_dir=skill_dir, direction=direction) + + async def upload_skill( + self, + skill_dir: str, + visibility: str = "public", + origin: str | None = None, + parent_skill_ids: list[str] | None = None, + tags: list[str] | None = None, + created_by: str | None = None, + change_summary: str | None = None, + ) -> str: + return await _upload_skill_impl( + skill_dir=skill_dir, + visibility=visibility, + origin=origin, + parent_skill_ids=parent_skill_ids, + tags=tags, + created_by=created_by, + change_summary=change_summary, + ) + + +register_main_tools(mcp, _DirectMainToolImplementation()) + + def run_mcp_server() -> None: """Console-script entry point for ``openspace-mcp``.""" import argparse parser = argparse.ArgumentParser(description="OpenSpace MCP Server") - parser.add_argument("--transport", choices=["stdio", "sse"], default="stdio") + parser.add_argument( + "--transport", + choices=["stdio", "sse", "streamable-http"], + default="stdio", + ) parser.add_argument("--port", type=int, default=8080) args = parser.parse_args() - if args.transport == "stdio": + if args.transport == "stdio" or os.environ.get("OPENSPACE_MCP_DAEMON") == "1": _maybe_start_idle_watchdog() + if args.transport == "streamable-http": + _maybe_start_main_daemon_embedding_prewarm() - if args.transport == "sse": - mcp.run(transport="sse", sse_params={"port": args.port}) - else: - mcp.run(transport="stdio") + mcp.settings.port = args.port + mcp.run(transport=args.transport) if __name__ == "__main__": diff --git a/openspace/mcp_tool_registration.py b/openspace/mcp_tool_registration.py new file mode 100644 index 0000000..be9af0b --- /dev/null +++ b/openspace/mcp_tool_registration.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +from typing import Protocol + +from mcp.server.fastmcp import FastMCP + + +class MainMCPToolImplementation(Protocol): + async def execute_task( + self, + task: str, + workspace_dir: str | None = None, + max_iterations: int | None = None, + skill_dirs: list[str] | None = None, + search_scope: str = "all", + ) -> str: ... + + async def search_skills( + self, + query: str, + source: str = "all", + limit: int = 20, + auto_import: bool = True, + ) -> str: ... + + async def fix_skill( + self, + skill_dir: str, + direction: str, + ) -> str: ... + + async def upload_skill( + self, + skill_dir: str, + visibility: str = "public", + origin: str | None = None, + parent_skill_ids: list[str] | None = None, + tags: list[str] | None = None, + created_by: str | None = None, + change_summary: str | None = None, + ) -> str: ... + + +class EvolutionMCPToolImplementation(Protocol): + async def evolve_from_context( + self, + task: str, + summary: str, + workspace_dir: str | None = None, + file_paths: list[str] | None = None, + max_skills: int = 3, + skill_dirs: list[str] | None = None, + output_dir: str | None = None, + ) -> str: ... + + +def register_main_tools(mcp: FastMCP, impl: MainMCPToolImplementation) -> None: + @mcp.tool() + async def execute_task( + task: str, + workspace_dir: str | None = None, + max_iterations: int | None = None, + skill_dirs: list[str] | None = None, + search_scope: str = "all", + ) -> str: + """Execute a task with OpenSpace's full grounding engine. + + OpenSpace will: + 1. Auto-register bot skills from skill_dirs (if provided) + 2. Search for relevant skills (scope controls local vs cloud+local) + 3. Attempt skill-guided execution → fallback to pure tools + 4. Auto-analyze → auto-evolve (FIX/DERIVED/CAPTURED) if needed + + If skills are auto-evolved, the response includes ``evolved_skills`` + with ``upload_ready: true``. Call ``upload_skill`` with just the + ``skill_dir`` + ``visibility`` to upload — metadata is pre-saved. + + Note: This call blocks until the task completes (may take minutes). + Set MCP client tool-call timeout ≥ 600 seconds. + + Args: + task: The task instruction (natural language). + workspace_dir: Working directory. Defaults to OPENSPACE_WORKSPACE env. + max_iterations: Max agent iterations (default: 20). + skill_dirs: Bot's skill directories to auto-register so OpenSpace + can select and track them. Directories are re-scanned + on every call to discover skills created since the last + invocation. + search_scope: Skill search scope before execution. + "all" (default) — local + cloud; falls back to local + if no API key is configured. + "local" — local SkillRegistry only (fast, no cloud). + """ + return await impl.execute_task( + task=task, + workspace_dir=workspace_dir, + max_iterations=max_iterations, + skill_dirs=skill_dirs, + search_scope=search_scope, + ) + + @mcp.tool() + async def search_skills( + query: str, + source: str = "all", + limit: int = 20, + auto_import: bool = True, + ) -> str: + """Search skills across local registry and cloud community. + + Standalone search for browsing / discovery. Use this when the bot + wants to find available skills, then decide whether to handle the + task locally or delegate to ``execute_task``. + + **Scope difference from execute_task**: + - ``search_skills`` returns results to the bot for decision-making. + - ``execute_task``'s internal search feeds directly into execution + (the bot never sees the search results). + + Uses hybrid ranking: BM25 → embedding re-rank → lexical boost. + Embedding requires OPENAI_API_KEY; falls back to lexical-only without it. + + Args: + query: Search query text (natural language or keywords). + source: "all" (cloud + local), "local", or "cloud". Default: "all". + limit: Maximum results to return (default: 20). + auto_import: Auto-download top public cloud skills (default: True). + """ + return await impl.search_skills( + query=query, + source=source, + limit=limit, + auto_import=auto_import, + ) + + @mcp.tool() + async def fix_skill( + skill_dir: str, + direction: str, + ) -> str: + """Manually fix a broken skill. + + This is the **only** manual evolution entry point. DERIVED and + CAPTURED evolutions are triggered automatically by ``execute_task`` + (they need a task to run). Use ``fix_skill`` when: + + - A skill's instructions are wrong or outdated + - The bot knows exactly which skill is broken and what to fix + - Auto-evolution inside ``execute_task`` didn't catch the issue + + The skill does NOT need to be pre-registered in OpenSpace — + provide the skill directory path and OpenSpace will register it + automatically before fixing. + + After fixing, the new skill is saved locally and ``.upload_meta.json`` + is pre-written. Call ``upload_skill`` with just ``skill_dir`` + + ``visibility`` to upload. + + Args: + skill_dir: Path to the broken skill directory (must contain SKILL.md). + direction: What's broken and how to fix it. Be specific: + e.g. "The API endpoint changed from v1 to v2" or + "Add retry logic for HTTP 429 rate limit errors". + """ + return await impl.fix_skill(skill_dir=skill_dir, direction=direction) + + @mcp.tool() + async def upload_skill( + skill_dir: str, + visibility: str = "public", + origin: str | None = None, + parent_skill_ids: list[str] | None = None, + tags: list[str] | None = None, + created_by: str | None = None, + change_summary: str | None = None, + ) -> str: + """Upload a local skill to the cloud. + + For evolved skills (from ``execute_task`` or ``fix_skill``), most + metadata is **pre-saved** in ``.upload_meta.json``. The bot only + needs to provide: + + - ``skill_dir`` — path to the skill directory + - ``visibility`` — "public" or "private" + + All other parameters are optional overrides. If omitted, pre-saved + values are used. If no pre-saved values exist, sensible defaults + are applied. + + **origin + parent_skill_ids constraints** (enforced by cloud): + - imported / captured → parent_skill_ids must be empty + - derived → at least 1 parent + - fixed → exactly 1 parent + + Args: + skill_dir: Path to skill directory (must contain SKILL.md). + visibility: "public" or "private". This is the one thing the + bot MUST decide. + origin: Override origin. Default: from .upload_meta.json or "imported". + parent_skill_ids: Override parents. Default: from .upload_meta.json. + tags: Override tags. Default: from .upload_meta.json. + created_by: Override creator. Default: from .upload_meta.json. + change_summary: Override summary. Default: from .upload_meta.json. + """ + return await impl.upload_skill( + skill_dir=skill_dir, + visibility=visibility, + origin=origin, + parent_skill_ids=parent_skill_ids, + tags=tags, + created_by=created_by, + change_summary=change_summary, + ) + + +def register_evolution_tools( + mcp: FastMCP, + impl: EvolutionMCPToolImplementation, +) -> None: + @mcp.tool() + async def evolve_from_context( + task: str, + summary: str, + workspace_dir: str | None = None, + file_paths: list[str] | None = None, + max_skills: int = 3, + skill_dirs: list[str] | None = None, + output_dir: str | None = None, + ) -> str: + """Capture reusable skills from a completed host-agent task. + + Use this when the main task was already handled by another agent + (for example Codex Desktop) and OpenSpace should only spend provider + tokens on post-task skill capture. + + Args: + task: Short description of the completed task. + summary: What changed, what was learned, and what seems reusable. + workspace_dir: Repository/workspace path. Defaults to OPENSPACE_WORKSPACE. + file_paths: Optional files worth emphasizing when planning captures. + max_skills: Maximum number of new skills to capture. + skill_dirs: Optional additional skill directories to register first. + output_dir: Override directory for new skills. Defaults to the first + OPENSPACE_HOST_SKILL_DIRS entry. + """ + return await impl.evolve_from_context( + task=task, + summary=summary, + workspace_dir=workspace_dir, + file_paths=file_paths, + max_skills=max_skills, + skill_dirs=skill_dirs, + output_dir=output_dir, + ) diff --git a/openspace/shared_mcp_runtime.py b/openspace/shared_mcp_runtime.py new file mode 100644 index 0000000..aa69d9d --- /dev/null +++ b/openspace/shared_mcp_runtime.py @@ -0,0 +1,519 @@ +from __future__ import annotations + +import asyncio +import contextlib +import hashlib +import json +import os +import signal +import socket +import subprocess +import sys +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Literal + +from openspace.config.loader import get_agent_config +from openspace.grounding.backends.mcp.client import MCPClient +from openspace.host_detection import ( + build_grounding_config_path, + build_llm_kwargs, + load_runtime_env, +) +from openspace.utils.logging import Logger + +logger = Logger.get_logger(__name__) + +ServerKind = Literal["main", "evolution"] + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_EXPECTED_TOOL_NAMES: dict[ServerKind, tuple[str, ...]] = { + "main": ("execute_task", "search_skills", "fix_skill", "upload_skill"), + "evolution": ("evolve_from_context",), +} +_SERVER_MODULES: dict[ServerKind, str] = { + "main": "openspace.mcp_server", + "evolution": "openspace.evolution_mcp_server", +} + + +@dataclass(frozen=True) +class MCPDaemonIdentity: + server_kind: ServerKind + workspace: str + resolved_model: str + llm_kwargs_fingerprint: str + backend_scope: tuple[str, ...] + host_skill_dirs: tuple[str, ...] + grounding_config_fingerprint: str + instance_key: str + state_dir: str + + @property + def metadata_path(self) -> Path: + return Path(self.state_dir) / f"{self.server_kind}-{self.instance_key}.json" + + @property + def lock_path(self) -> Path: + return Path(self.state_dir) / f"{self.server_kind}-{self.instance_key}.lock" + + @property + def log_path(self) -> Path: + return Path(self.state_dir) / f"{self.server_kind}-{self.instance_key}.log" + + +@dataclass(frozen=True) +class MCPDaemonRecord: + server_kind: ServerKind + instance_key: str + pid: int + port: int + workspace: str + resolved_model: str + llm_kwargs_fingerprint: str + backend_scope: list[str] + host_skill_dirs: list[str] + grounding_config_fingerprint: str + started_at: float + log_path: str + ready: bool = False + warmed: bool = False + ready_at: float | None = None + warmed_at: float | None = None + warmup_error: str | None = None + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.port}/mcp" + + +class _FileLock: + def __init__(self, path: Path): + self._path = path + self._handle = None + + def __enter__(self): + self._path.parent.mkdir(parents=True, exist_ok=True) + self._handle = self._path.open("a+", encoding="utf-8") + if os.name == "nt": + import msvcrt + + while True: + try: + msvcrt.locking(self._handle.fileno(), msvcrt.LK_LOCK, 1) + break + except OSError: + time.sleep(0.1) + else: + import fcntl + + fcntl.flock(self._handle.fileno(), fcntl.LOCK_EX) + return self + + def __exit__(self, exc_type, exc, tb): + if not self._handle: + return + try: + if os.name == "nt": + import msvcrt + + self._handle.seek(0) + msvcrt.locking(self._handle.fileno(), msvcrt.LK_UNLCK, 1) + else: + import fcntl + + fcntl.flock(self._handle.fileno(), fcntl.LOCK_UN) + finally: + self._handle.close() + self._handle = None + + +def _default_state_dir() -> Path: + override = os.environ.get("OPENSPACE_MCP_DAEMON_STATE_DIR", "").strip() + if override: + return Path(override).expanduser().resolve() + + if sys.platform == "darwin": + base = Path.home() / "Library" / "Application Support" + elif os.name == "nt": + base = Path(os.environ.get("LOCALAPPDATA", Path.home() / "AppData" / "Local")) + else: + base = Path(os.environ.get("XDG_STATE_HOME", Path.home() / ".local" / "state")) + return (base / "openspace" / "mcp-daemons").resolve() + + +def _canonical_workspace() -> Path: + workspace = Path(os.environ.get("OPENSPACE_WORKSPACE") or os.getcwd()).expanduser() + workspace = workspace.resolve() + try: + proc = subprocess.run( + ["git", "-C", str(workspace), "rev-parse", "--show-toplevel"], + check=False, + capture_output=True, + text=True, + ) + if proc.returncode == 0 and proc.stdout.strip(): + return Path(proc.stdout.strip()).resolve() + except Exception: + pass + return workspace + + +def _effective_backend_scope(server_kind: ServerKind) -> list[str]: + raw = os.environ.get("OPENSPACE_BACKEND_SCOPE", "").strip() + if raw: + parts = [part.strip().lower() for part in raw.split(",") if part.strip()] + return sorted(dict.fromkeys(parts)) + + if server_kind == "evolution": + return ["shell", "system"] + + agent_cfg = get_agent_config("GroundingAgent") or {} + parts = agent_cfg.get("backend_scope") or ["gui", "shell", "mcp", "web", "system"] + return sorted(dict.fromkeys(str(part).strip().lower() for part in parts if str(part).strip())) + + +def _effective_host_skill_dirs() -> list[str]: + raw = os.environ.get("OPENSPACE_HOST_SKILL_DIRS", "").strip() + if not raw: + return [] + + normalized: list[str] = [] + for item in raw.split(","): + item = item.strip() + if not item: + continue + resolved = str(Path(item).expanduser().resolve()) + if resolved not in normalized: + normalized.append(resolved) + return normalized + + +def _fingerprint_payload(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _grounding_config_fingerprint() -> str: + config_path = build_grounding_config_path() + if not config_path: + return "none" + + path = Path(config_path) + if path.is_file(): + return hashlib.sha256(path.read_bytes()).hexdigest() + return hashlib.sha256(str(path).encode("utf-8")).hexdigest() + + +def compute_daemon_identity(server_kind: ServerKind) -> MCPDaemonIdentity: + load_runtime_env() + + workspace = _canonical_workspace() + env_model = os.environ.get("OPENSPACE_MODEL", "") + resolved_model, llm_kwargs = build_llm_kwargs(env_model) + backend_scope = _effective_backend_scope(server_kind) + host_skill_dirs = _effective_host_skill_dirs() + grounding_config_fingerprint = _grounding_config_fingerprint() + llm_kwargs_fingerprint = _fingerprint_payload(llm_kwargs) + + key_payload = { + "server_kind": server_kind, + "workspace": str(workspace), + "resolved_model": resolved_model, + "llm_kwargs_fingerprint": llm_kwargs_fingerprint, + "backend_scope": backend_scope, + "host_skill_dirs": host_skill_dirs, + "grounding_config_fingerprint": grounding_config_fingerprint, + } + + return MCPDaemonIdentity( + server_kind=server_kind, + workspace=str(workspace), + resolved_model=resolved_model, + llm_kwargs_fingerprint=llm_kwargs_fingerprint, + backend_scope=tuple(backend_scope), + host_skill_dirs=tuple(host_skill_dirs), + grounding_config_fingerprint=grounding_config_fingerprint, + instance_key=_fingerprint_payload(key_payload)[:32], + state_dir=str(_default_state_dir()), + ) + + +def _read_record(path: Path) -> MCPDaemonRecord | None: + if not path.is_file(): + return None + try: + return MCPDaemonRecord(**json.loads(path.read_text(encoding="utf-8"))) + except Exception as exc: + logger.warning("Failed to read daemon metadata %s: %s", path, exc) + return None + + +def _write_record(path: Path, record: MCPDaemonRecord) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = path.with_suffix(path.suffix + ".tmp") + tmp_path.write_text( + json.dumps(asdict(record), ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + tmp_path.replace(path) + + +def _metadata_paths( + server_kind: ServerKind, + instance_key: str, + state_dir: str, +) -> tuple[Path, Path]: + state_path = Path(state_dir) + return ( + state_path / f"{server_kind}-{instance_key}.json", + state_path / f"{server_kind}-{instance_key}.lock", + ) + + +def update_current_daemon_status( + server_kind: ServerKind, + *, + ready: bool | None = None, + warmed: bool | None = None, + warmup_error: str | None = None, +) -> MCPDaemonRecord | None: + instance_key = os.environ.get("OPENSPACE_MCP_INSTANCE_KEY", "").strip() + state_dir = os.environ.get("OPENSPACE_MCP_DAEMON_STATE_DIR", "").strip() + if not instance_key or not state_dir: + return None + + metadata_path, lock_path = _metadata_paths(server_kind, instance_key, state_dir) + with _FileLock(lock_path): + record = _read_record(metadata_path) + if record is None: + return None + now = time.time() + + updates: dict[str, Any] = {} + if ready is not None: + updates["ready"] = ready + if ready and record.ready_at is None: + updates["ready_at"] = now + if warmed is not None: + updates["warmed"] = warmed + if warmed and record.warmed_at is None: + updates["warmed_at"] = now + if warmup_error is not None: + updates["warmup_error"] = warmup_error + + if not updates: + return record + + updated = MCPDaemonRecord( + **{ + **asdict(record), + **updates, + } + ) + _write_record(metadata_path, updated) + return updated + + +def _pid_exists(pid: int) -> bool: + if pid <= 0: + return False + try: + os.kill(pid, 0) + return True + except OSError: + return False + + +def _expected_process_marker(server_kind: ServerKind) -> str: + return _SERVER_MODULES[server_kind] + + +def _pid_matches_server(record: MCPDaemonRecord) -> bool: + if os.name == "nt": + return _pid_exists(record.pid) + try: + proc = subprocess.run( + ["ps", "-o", "command=", "-p", str(record.pid)], + check=False, + capture_output=True, + text=True, + ) + except Exception: + return False + command = proc.stdout.strip() + return bool(command) and _expected_process_marker(record.server_kind) in command + + +def _terminate_record_process(record: MCPDaemonRecord) -> None: + if not _pid_exists(record.pid) or not _pid_matches_server(record): + return + + with contextlib.suppress(Exception): + os.kill(record.pid, signal.SIGTERM) + deadline = time.monotonic() + 3.0 + while time.monotonic() < deadline: + if not _pid_exists(record.pid): + return + time.sleep(0.1) + with contextlib.suppress(Exception): + os.kill(record.pid, signal.SIGKILL) + + +def _pick_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + sock.listen(1) + return int(sock.getsockname()[1]) + + +def _spawn_daemon(identity: MCPDaemonIdentity, port: int) -> MCPDaemonRecord: + env = os.environ.copy() + env["OPENSPACE_MCP_DAEMON"] = "1" + env["OPENSPACE_MCP_INSTANCE_KEY"] = identity.instance_key + env["OPENSPACE_MCP_DAEMON_STATE_DIR"] = identity.state_dir + env["OPENSPACE_WORKSPACE"] = identity.workspace + env["OPENSPACE_MODEL"] = identity.resolved_model + env["OPENSPACE_BACKEND_SCOPE"] = ",".join(identity.backend_scope) + if identity.host_skill_dirs: + env["OPENSPACE_HOST_SKILL_DIRS"] = ",".join(identity.host_skill_dirs) + else: + env.pop("OPENSPACE_HOST_SKILL_DIRS", None) + + log_path = identity.log_path + log_path.parent.mkdir(parents=True, exist_ok=True) + log_handle = log_path.open("ab") + + popen_kwargs: dict[str, Any] = { + "cwd": identity.workspace, + "env": env, + "stdin": subprocess.DEVNULL, + "stdout": log_handle, + "stderr": subprocess.STDOUT, + } + if os.name == "nt": + popen_kwargs["creationflags"] = ( + subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.DETACHED_PROCESS + ) + else: + popen_kwargs["start_new_session"] = True + + proc = subprocess.Popen( + [ + sys.executable, + "-m", + _SERVER_MODULES[identity.server_kind], + "--transport", + "streamable-http", + "--port", + str(port), + ], + **popen_kwargs, + ) + log_handle.close() + return MCPDaemonRecord( + server_kind=identity.server_kind, + instance_key=identity.instance_key, + pid=proc.pid, + port=port, + workspace=identity.workspace, + resolved_model=identity.resolved_model, + llm_kwargs_fingerprint=identity.llm_kwargs_fingerprint, + backend_scope=list(identity.backend_scope), + host_skill_dirs=list(identity.host_skill_dirs), + grounding_config_fingerprint=identity.grounding_config_fingerprint, + started_at=time.time(), + log_path=str(log_path), + ) + + +async def _probe_record(record: MCPDaemonRecord) -> bool: + client = MCPClient( + config={"mcpServers": {"daemon": {"url": record.url}}}, + timeout=5.0, + sse_read_timeout=15.0, + max_retries=1, + retry_interval=0.1, + check_dependencies=False, + ) + try: + session = await client.create_session("daemon", auto_initialize=True) + if session is None: + return False + tools = await session.list_tools() + actual = {tool.name for tool in tools} + expected = set(_EXPECTED_TOOL_NAMES[record.server_kind]) + return actual == expected + except Exception: + return False + finally: + with contextlib.suppress(Exception): + await client.close_all_sessions() + + +async def _wait_until_ready(record: MCPDaemonRecord, timeout_seconds: float = 15.0) -> bool: + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + if _pid_exists(record.pid) and await _probe_record(record): + return True + await asyncio.sleep(0.25) + return False + + +async def ensure_daemon(server_kind: ServerKind) -> MCPDaemonRecord: + identity = compute_daemon_identity(server_kind) + identity.metadata_path.parent.mkdir(parents=True, exist_ok=True) + + with _FileLock(identity.lock_path): + existing = _read_record(identity.metadata_path) + if existing and _pid_exists(existing.pid) and await _probe_record(existing): + if not existing.ready or (server_kind != "main" and not existing.warmed): + now = time.time() + refreshed = MCPDaemonRecord( + **{ + **asdict(existing), + "ready": True, + "ready_at": existing.ready_at or now, + "warmed": (existing.warmed or server_kind != "main"), + "warmed_at": ( + existing.warmed_at + or (now if (existing.warmed or server_kind != "main") else None) + ), + } + ) + _write_record(identity.metadata_path, refreshed) + return refreshed or existing + return existing + + if existing: + _terminate_record_process(existing) + with contextlib.suppress(FileNotFoundError): + identity.metadata_path.unlink() + + last_error: Exception | None = None + for _ in range(3): + record = _spawn_daemon(identity, _pick_free_port()) + _write_record(identity.metadata_path, record) + if await _wait_until_ready(record): + now = time.time() + updated = MCPDaemonRecord( + **{ + **asdict(record), + "ready": True, + "ready_at": now, + "warmed": (server_kind != "main"), + "warmed_at": (now if server_kind != "main" else None), + } + ) + _write_record(identity.metadata_path, updated) + return updated + + last_error = RuntimeError( + f"Daemon for key={identity.instance_key} did not become ready" + ) + _terminate_record_process(record) + with contextlib.suppress(FileNotFoundError): + identity.metadata_path.unlink() + + raise last_error or RuntimeError("Failed to start daemon") diff --git a/openspace/skill_engine/skill_ranker.py b/openspace/skill_engine/skill_ranker.py index 8498f2b..25d2b61 100644 --- a/openspace/skill_engine/skill_ranker.py +++ b/openspace/skill_engine/skill_ranker.py @@ -172,6 +172,43 @@ class SkillRanker: self._save_cache() return emb + def get_cached_embedding(self, skill_id: str) -> Optional[List[float]]: + """Return a cached embedding without computing a new one.""" + return self._embedding_cache.get(skill_id) + + def prime_candidates(self, candidates: List[SkillCandidate]) -> int: + """Populate embeddings for candidates, saving cache once at the end. + + Returns the number of candidates that ended with an embedding, whether + loaded from cache or computed during this call. + """ + warmed = 0 + cache_changed = False + + for candidate in candidates: + if candidate.embedding: + warmed += 1 + continue + + cached = self._embedding_cache.get(candidate.skill_id) + if cached: + candidate.embedding = cached + warmed += 1 + continue + + text = self._build_embedding_text(candidate) + emb = self._generate_embedding(text) + if emb: + candidate.embedding = emb + self._embedding_cache[candidate.skill_id] = emb + warmed += 1 + cache_changed = True + + if cache_changed: + self._save_cache() + + return warmed + def invalidate_cache(self, skill_id: str) -> None: """Remove a skill's cached embedding (e.g. after evolution).""" self._embedding_cache.pop(skill_id, None) diff --git a/scripts/codex-openspace b/scripts/codex-openspace index bb459c2..cb7bccf 100755 --- a/scripts/codex-openspace +++ b/scripts/codex-openspace @@ -46,6 +46,7 @@ sync_profile_dir() { bootstrap_profile_home() { mkdir -p "$PROFILE_HOME" + mkdir -p "$PROFILE_HOME/state/openspace" sync_profile_dir "plugins" sync_profile_dir "skills" @@ -94,12 +95,35 @@ enabled = false url = "https://mcp.linear.app/mcp" [mcp_servers.openspace] -command = "$REPO_ROOT/.venv/bin/openspace-mcp" -args = ["--transport", "stdio"] +command = "$REPO_ROOT/.venv/bin/python" +args = ["-m", "openspace.mcp_proxy", "--kind", "main", "--transport", "stdio"] [mcp_servers.openspace.env] OPENSPACE_HOST_SKILL_DIRS = "$PROFILE_HOME/skills" OPENSPACE_WORKSPACE = "$REPO_ROOT" +OPENSPACE_MCP_PROXY_MODE = "daemon" +OPENSPACE_MCP_DAEMON_STATE_DIR = "$PROFILE_HOME/state/openspace" +OPENSPACE_MODEL = "$OPENSPACE_MODEL" +OPENSPACE_LLM_API_KEY = "$OPENSPACE_LLM_API_KEY" +OPENSPACE_LLM_API_BASE = "$OPENSPACE_LLM_API_BASE" +OPENSPACE_LLM_OPENAI_STREAM_COMPAT = "$OPENSPACE_LLM_OPENAI_STREAM_COMPAT" +OPENSPACE_SKILL_EMBEDDING_BACKEND = "$OPENSPACE_SKILL_EMBEDDING_BACKEND" +OPENSPACE_SKILL_EMBEDDING_MODEL = "$OPENSPACE_SKILL_EMBEDDING_MODEL" +OPENSPACE_SKILL_EMBEDDING_API_KEY = "${OPENSPACE_SKILL_EMBEDDING_API_KEY:-}" +OPENSPACE_SKILL_EMBEDDING_API_BASE = "${OPENSPACE_SKILL_EMBEDDING_API_BASE:-}" +EMBEDDING_API_KEY = "${EMBEDDING_API_KEY:-}" +EMBEDDING_BASE_URL = "${EMBEDDING_BASE_URL:-}" +EMBEDDING_MODEL = "${EMBEDDING_MODEL:-}" + +[mcp_servers.openspace_evolution] +command = "$REPO_ROOT/.venv/bin/python" +args = ["-m", "openspace.mcp_proxy", "--kind", "evolution", "--transport", "stdio"] + +[mcp_servers.openspace_evolution.env] +OPENSPACE_HOST_SKILL_DIRS = "$PROFILE_HOME/skills" +OPENSPACE_WORKSPACE = "$REPO_ROOT" +OPENSPACE_MCP_PROXY_MODE = "daemon" +OPENSPACE_MCP_DAEMON_STATE_DIR = "$PROFILE_HOME/state/openspace" OPENSPACE_MODEL = "$OPENSPACE_MODEL" OPENSPACE_LLM_API_KEY = "$OPENSPACE_LLM_API_KEY" OPENSPACE_LLM_API_BASE = "$OPENSPACE_LLM_API_BASE" diff --git a/scripts/install-global-codex-openspace b/scripts/install-global-codex-openspace index 1ff0b27..0155d7c 100644 --- a/scripts/install-global-codex-openspace +++ b/scripts/install-global-codex-openspace @@ -19,6 +19,7 @@ cat > "$BIN_DIR/openspace-global-mcp" < "$BIN_DIR/openspace-evolution-global-mcp" < "$BIN_DIR/openspace-evolution-global-mcp" < None: + module = ModuleType("fastembed") + module.TextEmbedding = _DummyTextEmbedding + monkeypatch.setitem(sys.modules, "fastembed", module) + + +def _reset_embedding_state(monkeypatch) -> None: + monkeypatch.setattr(embedding, "_LOCAL_EMBEDDER", None, raising=False) + monkeypatch.setattr(embedding, "_LOCAL_EMBEDDER_MODEL", None, raising=False) + _DummyTextEmbedding.instances.clear() + + +def test_load_local_embedder_reuses_same_model_instance(monkeypatch) -> None: + _install_fastembed_stub(monkeypatch) + _reset_embedding_state(monkeypatch) + + first = embedding._load_local_embedder("unit-model") + second = embedding._load_local_embedder("unit-model") + third = embedding._load_local_embedder("other-model") + + assert first is second + assert third is not first + assert [instance.model_name for instance in _DummyTextEmbedding.instances] == [ + "unit-model", + "other-model", + ] + + +def test_generate_embedding_reuses_prewarmed_local_embedder(monkeypatch) -> None: + _install_fastembed_stub(monkeypatch) + _reset_embedding_state(monkeypatch) + monkeypatch.setenv("OPENSPACE_SKILL_EMBEDDING_BACKEND", "local") + monkeypatch.setenv("OPENSPACE_SKILL_EMBEDDING_MODEL", "unit-model") + + first = embedding.generate_embedding("alpha") + second = embedding.generate_embedding("beta") + + assert first == [5.0, 10.0] + assert second == [4.0, 10.0] + assert len(_DummyTextEmbedding.instances) == 1 + assert _DummyTextEmbedding.instances[0].embed_inputs == [["alpha"], ["beta"]] + + +def test_skill_ranker_reuses_persisted_embedding_cache_between_instances( + monkeypatch, + tmp_path, +) -> None: + calls: list[str] = [] + + monkeypatch.setattr( + "openspace.cloud.embedding.resolve_skill_embedding_model", + lambda backend=None: "unit-model", + ) + + def fake_generate_embedding(text: str, api_key=None): + calls.append(text) + return [float(len(text)), 1.0] + + monkeypatch.setattr( + SkillRanker, + "_generate_embedding", + staticmethod(fake_generate_embedding), + ) + + first_ranker = SkillRanker(cache_dir=tmp_path, enable_cache=True) + candidate = SkillCandidate( + skill_id="skill-1", + name="alpha", + description="beta", + body="gamma", + ) + first_ranker.hybrid_rank("query text", [candidate], top_k=1) + + cache_file = tmp_path / "skill_embeddings_unit-model_v2.pkl" + assert cache_file.exists() + assert calls == [ + "query text", + embedding.build_skill_embedding_text("alpha", "beta", "gamma"), + ] + + calls.clear() + + second_ranker = SkillRanker(cache_dir=tmp_path, enable_cache=True) + assert "skill-1" in second_ranker._embedding_cache + + second_candidate = SkillCandidate( + skill_id="skill-1", + name="alpha", + description="beta", + body="gamma", + ) + second_ranker.hybrid_rank("query text", [second_candidate], top_k=1) + + assert calls == ["query text"] + + +def test_skill_search_engine_uses_ranker_cache_for_local_candidates(monkeypatch) -> None: + events: list[tuple[str, str]] = [] + + class _DummyRanker: + def __init__(self, enable_cache: bool = True): + self.enable_cache = enable_cache + + def get_cached_embedding(self, skill_id: str): + events.append(("cached", skill_id)) + return [0.5, 0.5] + + def prime_candidates(self, candidates): + events.append(("prime", candidates[0].skill_id)) + return 1 + + monkeypatch.setattr( + "openspace.skill_engine.skill_ranker.SkillRanker", + _DummyRanker, + ) + monkeypatch.setattr( + "openspace.cloud.embedding.cosine_similarity", + lambda a, b: 0.75, + ) + + engine = SkillSearchEngine() + scored = engine._score_phase( + candidates=[ + { + "skill_id": "skill-local", + "name": "Local Skill", + "description": "demo", + "source": "openspace-local", + "_embedding_text": "Local Skill\ndemo", + } + ], + query_tokens=["local"], + query_embedding=[1.0, 1.0], + ) + + assert events == [("cached", "skill-local")] + assert scored[0]["vector_score"] == 0.75 diff --git a/tests/test_mcp_entrypoints.py b/tests/test_mcp_entrypoints.py new file mode 100644 index 0000000..4fe6f63 --- /dev/null +++ b/tests/test_mcp_entrypoints.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import argparse +import importlib +from types import SimpleNamespace + +import pytest + + +ENTRYPOINT_MODULES = [ + "openspace.mcp_server", + "openspace.evolution_mcp_server", +] + + +@pytest.mark.parametrize("module_name", ENTRYPOINT_MODULES) +def test_stdio_entrypoint_uses_stdio_transport(module_name, monkeypatch) -> None: + module = importlib.import_module(module_name) + calls: list[tuple[tuple[object, ...], dict[str, object]]] = [] + watchdog_calls: list[bool] = [] + + monkeypatch.setattr( + argparse.ArgumentParser, + "parse_args", + lambda self: SimpleNamespace(transport="stdio", port=9123), + ) + monkeypatch.setattr( + module.mcp, + "run", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + monkeypatch.setattr( + module, + "_maybe_start_idle_watchdog", + lambda: watchdog_calls.append(True), + ) + + module.run_mcp_server() + + assert watchdog_calls == [True] + assert calls == [((), {"transport": "stdio"})] + assert module.mcp.settings.port == 9123 + + +@pytest.mark.parametrize("module_name", ENTRYPOINT_MODULES) +def test_sse_entrypoint_does_not_forward_sse_params(module_name, monkeypatch) -> None: + module = importlib.import_module(module_name) + calls: list[tuple[tuple[object, ...], dict[str, object]]] = [] + watchdog_calls: list[bool] = [] + + monkeypatch.setattr( + argparse.ArgumentParser, + "parse_args", + lambda self: SimpleNamespace(transport="sse", port=9123), + ) + monkeypatch.setattr( + module.mcp, + "run", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + monkeypatch.setattr( + module, + "_maybe_start_idle_watchdog", + lambda: watchdog_calls.append(True), + ) + + module.run_mcp_server() + + assert watchdog_calls == [] + assert calls == [((), {"transport": "sse"})] + assert module.mcp.settings.port == 9123 + + +@pytest.mark.parametrize("module_name", ENTRYPOINT_MODULES) +def test_streamable_http_entrypoint_uses_watchdog_for_daemon( + module_name, + monkeypatch, +) -> None: + module = importlib.import_module(module_name) + calls: list[tuple[tuple[object, ...], dict[str, object]]] = [] + watchdog_calls: list[bool] = [] + + monkeypatch.setattr( + argparse.ArgumentParser, + "parse_args", + lambda self: SimpleNamespace(transport="streamable-http", port=9234), + ) + monkeypatch.setattr( + module.mcp, + "run", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + monkeypatch.setattr( + module, + "_maybe_start_idle_watchdog", + lambda: watchdog_calls.append(True), + ) + monkeypatch.setenv("OPENSPACE_MCP_DAEMON", "1") + + module.run_mcp_server() + + assert watchdog_calls == [True] + assert calls == [((), {"transport": "streamable-http"})] + assert module.mcp.settings.port == 9234 diff --git a/tests/test_mcp_http_connector_transport.py b/tests/test_mcp_http_connector_transport.py new file mode 100644 index 0000000..04c6735 --- /dev/null +++ b/tests/test_mcp_http_connector_transport.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +import importlib.util +import logging +import sys +from pathlib import Path +from types import ModuleType + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +BASE_FILE = REPO_ROOT / "openspace/grounding/backends/mcp/transport/connectors/base.py" +CORE_TM_BASE_FILE = REPO_ROOT / "openspace/grounding/core/transport/task_managers/base.py" +HTTP_FILE = REPO_ROOT / "openspace/grounding/backends/mcp/transport/connectors/http.py" + + +class _DummyStreamableHttpConnectionManager: + instances: list["_DummyStreamableHttpConnectionManager"] = [] + + def __init__(self, url, headers, timeout, read_timeout): + self.url = url + self.headers = headers + self.timeout = timeout + self.read_timeout = read_timeout + self.started = False + self.stopped = False + _DummyStreamableHttpConnectionManager.instances.append(self) + + async def start(self, timeout=None): + self.started = True + self.timeout_used = timeout + return "read-stream", "write-stream" + + def get_streams(self): + return ("read-stream", "write-stream") + + async def stop(self): + self.stopped = True + + +class _ForbiddenSseConnectionManager: + def __init__(self, *args, **kwargs): + raise AssertionError( + "SSE fallback should not be constructed when streamable HTTP succeeds" + ) + + +class _DummyClientSession: + def __init__(self, read_stream, write_stream, sampling_callback=None): + self.read_stream = read_stream + self.write_stream = write_stream + self.sampling_callback = sampling_callback + self.entered = False + self.initialized = False + self.tools_listed = False + self.exited = False + + async def __aenter__(self): + self.entered = True + return self + + async def initialize(self): + self.initialized = True + + async def list_tools(self): + self.tools_listed = True + return [] + + async def __aexit__(self, exc_type, exc, tb): + self.exited = True + + +def _install_package_stub( + monkeypatch, + module_name: str, + **attributes, +) -> ModuleType: + module = ModuleType(module_name) + module.__path__ = [] # mark as package + for key, value in attributes.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, module_name, module) + return module + + +class _BaseConnectorStub: + @classmethod + def __class_getitem__(cls, item): + return cls + + def __init__(self, connection_manager): + self._connection_manager = connection_manager + self._connection = None + self._connected = False + + async def _cleanup_on_connect_failure(self): + if self._connection_manager and hasattr(self._connection_manager, "stop"): + maybe_awaitable = self._connection_manager.stop() + if hasattr(maybe_awaitable, "__await__"): + await maybe_awaitable + self._connection = None + + async def _after_disconnect(self): + return None + + +class _BaseConnectionManagerStub: + @classmethod + def __class_getitem__(cls, item): + return cls + + +class _AsyncContextConnectionManagerStub(_BaseConnectionManagerStub): + pass + + +class _PlaceholderConnectionManagerStub: + def __init__(self, *args, **kwargs): + self._connection = None + + async def start(self, timeout=None): + return self._connection + + async def stop(self, timeout=5.0): + return None + + def get_streams(self): + return self._connection + + +def _load_module(module_name: str, file_path: Path) -> ModuleType: + spec = importlib.util.spec_from_file_location(module_name, file_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def _load_http_module(monkeypatch) -> ModuleType: + # Stub the package layers so we can load the target file without + # importing the broader MCP package tree and its optional deps. + _install_package_stub( + monkeypatch, + "openspace.grounding.core.transport.connectors", + BaseConnector=_BaseConnectorStub, + ) + _install_package_stub( + monkeypatch, + "openspace.grounding.core.transport.task_managers", + BaseConnectionManager=_BaseConnectionManagerStub, + AsyncContextConnectionManager=_AsyncContextConnectionManagerStub, + PlaceholderConnectionManager=_PlaceholderConnectionManagerStub, + ) + _install_package_stub( + monkeypatch, + "openspace.utils.logging", + Logger=type( + "Logger", + (), + {"get_logger": staticmethod(logging.getLogger)}, + ), + ) + _install_package_stub( + monkeypatch, + "openspace.grounding.backends.mcp.transport.task_managers", + SseConnectionManager=type("SseConnectionManager", (), {}), + StreamableHttpConnectionManager=type( + "StreamableHttpConnectionManager", (), {} + ), + ) + _install_package_stub( + monkeypatch, + "openspace.grounding.backends.mcp.transport.connectors", + ) + _install_package_stub( + monkeypatch, + "openspace.grounding.backends.mcp.transport", + ) + _install_package_stub( + monkeypatch, + "openspace.grounding.backends.mcp", + ) + + _load_module( + "openspace.grounding.backends.mcp.transport.connectors.base", + BASE_FILE, + ) + _load_module( + "openspace.grounding.core.transport.task_managers.base", + CORE_TM_BASE_FILE, + ) + return _load_module( + "openspace.grounding.backends.mcp.transport.connectors.http", + HTTP_FILE, + ) + + +@pytest.mark.asyncio +async def test_http_connector_prefers_streamable_http(monkeypatch) -> None: + http_module = _load_http_module(monkeypatch) + + monkeypatch.setattr( + http_module, + "StreamableHttpConnectionManager", + _DummyStreamableHttpConnectionManager, + ) + monkeypatch.setattr( + http_module, + "SseConnectionManager", + _ForbiddenSseConnectionManager, + ) + monkeypatch.setattr(http_module, "ClientSession", _DummyClientSession) + + connector = http_module.HttpConnector("http://127.0.0.1:8123/mcp") + + await connector.connect() + + assert connector.transport_type == "streamable HTTP" + assert isinstance( + connector._connection_manager, _DummyStreamableHttpConnectionManager + ) + assert connector._connection == ("read-stream", "write-stream") + assert connector.client_session.entered is True + assert connector.client_session.initialized is True + assert connector.client_session.tools_listed is True + + client_session = connector.client_session + await connector.disconnect() + + assert client_session.exited is True + assert connector._connected is False + assert connector._connection is None + assert connector._connection_manager.stopped is True diff --git a/tests/test_mcp_proxy_runtime.py b/tests/test_mcp_proxy_runtime.py new file mode 100644 index 0000000..210622e --- /dev/null +++ b/tests/test_mcp_proxy_runtime.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path + +from openspace import mcp_proxy +from openspace import shared_mcp_runtime +from openspace.mcp_tool_registration import register_main_tools + + +def test_proxy_mode_defaults_follow_split_rollout(monkeypatch) -> None: + monkeypatch.delenv("OPENSPACE_MCP_PROXY_MODE", raising=False) + + assert mcp_proxy._proxy_mode_for("main") == "daemon" + assert mcp_proxy._proxy_mode_for("evolution") == "daemon" + + monkeypatch.setenv("OPENSPACE_MCP_PROXY_MODE", "daemon") + assert mcp_proxy._proxy_mode_for("main") == "daemon" + assert mcp_proxy._proxy_mode_for("evolution") == "daemon" + + +def test_proxy_registration_is_lazy(monkeypatch) -> None: + async def _fail_if_called(server_kind): + raise AssertionError(f"ensure_daemon should not run during tool registration ({server_kind})") + + monkeypatch.setattr(mcp_proxy, "ensure_daemon", _fail_if_called) + + mcp = mcp_proxy._build_fastmcp("main") + register_main_tools(mcp, mcp_proxy._MainProxyImplementation()) + + assert {tool.name for tool in mcp._tool_manager.list_tools()} == { + "execute_task", + "search_skills", + "fix_skill", + "upload_skill", + } + + +def test_compute_daemon_identity_normalizes_repo_workspace(monkeypatch) -> None: + repo_root = Path(__file__).resolve().parents[1] + nested_path = repo_root / "openspace" + + monkeypatch.setattr( + shared_mcp_runtime, + "build_llm_kwargs", + lambda model: ("resolved-model", {"api_base": "http://unit.test/v1"}), + ) + monkeypatch.setattr(shared_mcp_runtime, "build_grounding_config_path", lambda: None) + monkeypatch.setattr( + shared_mcp_runtime, + "get_agent_config", + lambda name: {"backend_scope": ["mcp", "shell"]}, + ) + monkeypatch.delenv("OPENSPACE_BACKEND_SCOPE", raising=False) + monkeypatch.delenv("OPENSPACE_HOST_SKILL_DIRS", raising=False) + + monkeypatch.setenv("OPENSPACE_WORKSPACE", str(repo_root)) + root_identity = shared_mcp_runtime.compute_daemon_identity("main") + + monkeypatch.setenv("OPENSPACE_WORKSPACE", str(nested_path)) + nested_identity = shared_mcp_runtime.compute_daemon_identity("main") + + assert root_identity.workspace == str(repo_root.resolve()) + assert nested_identity.workspace == str(repo_root.resolve()) + assert root_identity.instance_key == nested_identity.instance_key + + +def test_compute_daemon_identity_changes_when_skill_dirs_change(monkeypatch, tmp_path) -> None: + monkeypatch.setattr( + shared_mcp_runtime, + "build_llm_kwargs", + lambda model: ("resolved-model", {"api_base": "http://unit.test/v1"}), + ) + monkeypatch.setattr(shared_mcp_runtime, "build_grounding_config_path", lambda: None) + monkeypatch.setattr( + shared_mcp_runtime, + "get_agent_config", + lambda name: {"backend_scope": ["shell", "mcp"]}, + ) + monkeypatch.setenv("OPENSPACE_WORKSPACE", str(tmp_path)) + monkeypatch.delenv("OPENSPACE_BACKEND_SCOPE", raising=False) + + first = tmp_path / "skills-a" + second = tmp_path / "skills-b" + first.mkdir() + second.mkdir() + + monkeypatch.setenv("OPENSPACE_HOST_SKILL_DIRS", str(first)) + first_identity = shared_mcp_runtime.compute_daemon_identity("main") + + monkeypatch.setenv("OPENSPACE_HOST_SKILL_DIRS", f"{first},{second}") + second_identity = shared_mcp_runtime.compute_daemon_identity("main") + + assert first_identity.host_skill_dirs == (str(first.resolve()),) + assert second_identity.host_skill_dirs == ( + str(first.resolve()), + str(second.resolve()), + ) + assert first_identity.instance_key != second_identity.instance_key + + +async def _ready_probe(record): + return True + + +def test_ensure_daemon_marks_main_ready_but_not_warmed(monkeypatch, tmp_path) -> None: + identity = shared_mcp_runtime.MCPDaemonIdentity( + server_kind="main", + workspace=str(tmp_path), + resolved_model="model", + llm_kwargs_fingerprint="llm", + backend_scope=("shell",), + host_skill_dirs=(str(tmp_path),), + grounding_config_fingerprint="cfg", + instance_key="main-key", + state_dir=str(tmp_path), + ) + + monkeypatch.setattr(shared_mcp_runtime, "compute_daemon_identity", lambda kind: identity) + monkeypatch.setattr(shared_mcp_runtime, "_pick_free_port", lambda: 12345) + monkeypatch.setattr(shared_mcp_runtime, "_spawn_daemon", lambda ident, port: shared_mcp_runtime.MCPDaemonRecord( + server_kind="main", + instance_key=ident.instance_key, + pid=4321, + port=port, + workspace=ident.workspace, + resolved_model=ident.resolved_model, + llm_kwargs_fingerprint=ident.llm_kwargs_fingerprint, + backend_scope=list(ident.backend_scope), + host_skill_dirs=list(ident.host_skill_dirs), + grounding_config_fingerprint=ident.grounding_config_fingerprint, + started_at=1.0, + log_path=str(identity.log_path), + )) + monkeypatch.setattr(shared_mcp_runtime, "_wait_until_ready", _ready_probe) + monkeypatch.setattr(shared_mcp_runtime, "_pid_exists", lambda pid: True) + + record = asyncio.run(shared_mcp_runtime.ensure_daemon("main")) + + assert record.ready is True + assert record.warmed is False + assert record.ready_at is not None + assert record.warmed_at is None + + +def test_update_current_daemon_status_marks_warmed(monkeypatch, tmp_path) -> None: + metadata_path = tmp_path / "main-key.json" + lock_path = tmp_path / "main-key.lock" + record = shared_mcp_runtime.MCPDaemonRecord( + server_kind="main", + instance_key="key", + pid=4321, + port=12345, + workspace=str(tmp_path), + resolved_model="model", + llm_kwargs_fingerprint="llm", + backend_scope=["shell"], + host_skill_dirs=[str(tmp_path)], + grounding_config_fingerprint="cfg", + started_at=1.0, + log_path=str(tmp_path / "main-key.log"), + ready=True, + warmed=False, + ready_at=2.0, + ) + shared_mcp_runtime._write_record(metadata_path, record) + lock_path.touch() + + monkeypatch.setenv("OPENSPACE_MCP_INSTANCE_KEY", "key") + monkeypatch.setenv("OPENSPACE_MCP_DAEMON_STATE_DIR", str(tmp_path)) + + updated = shared_mcp_runtime.update_current_daemon_status( + "main", + warmed=True, + warmup_error=None, + ) + + assert updated is not None + assert updated.ready is True + assert updated.warmed is True + assert updated.warmed_at is not None diff --git a/tests/test_shared_mcp_runtime_metadata.py b/tests/test_shared_mcp_runtime_metadata.py new file mode 100644 index 0000000..bb2f4fc --- /dev/null +++ b/tests/test_shared_mcp_runtime_metadata.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import threading +import time +from pathlib import Path + +from openspace import shared_mcp_runtime + + +def _build_identity(state_dir: Path) -> shared_mcp_runtime.MCPDaemonIdentity: + return shared_mcp_runtime.MCPDaemonIdentity( + server_kind="main", + workspace="/Users/admin/PycharmProjects/openspace", + resolved_model="unit-model", + llm_kwargs_fingerprint="llm-fingerprint", + backend_scope=("shell", "mcp"), + host_skill_dirs=("/tmp/unit-skills",), + grounding_config_fingerprint="grounding-fingerprint", + instance_key="unit-instance-key", + state_dir=str(state_dir), + ) + + +def _build_record(identity: shared_mcp_runtime.MCPDaemonIdentity) -> shared_mcp_runtime.MCPDaemonRecord: + return shared_mcp_runtime.MCPDaemonRecord( + server_kind=identity.server_kind, + instance_key=identity.instance_key, + pid=4242, + port=56789, + workspace=identity.workspace, + resolved_model=identity.resolved_model, + llm_kwargs_fingerprint=identity.llm_kwargs_fingerprint, + backend_scope=list(identity.backend_scope), + host_skill_dirs=list(identity.host_skill_dirs), + grounding_config_fingerprint=identity.grounding_config_fingerprint, + started_at=100.0, + log_path=str(Path(identity.state_dir) / "main-unit-instance-key.log"), + ready=False, + warmed=False, + ) + + +def test_daemon_metadata_round_trip_includes_ready_and_warmed(tmp_path) -> None: + identity = _build_identity(tmp_path) + record = _build_record(identity) + metadata_path = identity.metadata_path + metadata_path.parent.mkdir(parents=True, exist_ok=True) + shared_mcp_runtime._write_record(metadata_path, record) + + assert metadata_path.is_file() + + initial = shared_mcp_runtime._read_record(metadata_path) + assert initial is not None + assert initial.ready is False + assert initial.warmed is False + + assert initial.server_kind == "main" + assert initial.instance_key == identity.instance_key + + +def test_update_current_daemon_status_marks_ready_then_warmed_for_main_daemon(monkeypatch, tmp_path) -> None: + identity = _build_identity(tmp_path) + metadata_path = identity.metadata_path + metadata_path.parent.mkdir(parents=True, exist_ok=True) + shared_mcp_runtime._write_record(metadata_path, _build_record(identity)) + + monkeypatch.setenv("OPENSPACE_MCP_INSTANCE_KEY", identity.instance_key) + monkeypatch.setenv("OPENSPACE_MCP_DAEMON_STATE_DIR", identity.state_dir) + + monkeypatch.setattr(shared_mcp_runtime.time, "time", lambda: 101.0) + ready_record = shared_mcp_runtime.update_current_daemon_status("main", ready=True) + assert ready_record is not None + assert ready_record.ready is True + assert ready_record.warmed is False + assert ready_record.ready_at == 101.0 + assert ready_record.warmed_at is None + + monkeypatch.setattr(shared_mcp_runtime.time, "time", lambda: 107.5) + warmed_record = shared_mcp_runtime.update_current_daemon_status("main", warmed=True) + assert warmed_record is not None + assert warmed_record.ready is True + assert warmed_record.warmed is True + assert warmed_record.ready_at == 101.0 + assert warmed_record.warmed_at == 107.5 + + reloaded = shared_mcp_runtime._read_record(metadata_path) + assert reloaded is not None + assert reloaded.ready is True + assert reloaded.warmed is True + assert reloaded.ready_at == 101.0 + assert reloaded.warmed_at == 107.5 + + +def test_spawn_daemon_exports_metadata_env_for_background_updates(monkeypatch, tmp_path) -> None: + identity = _build_identity(tmp_path) + captured: dict[str, object] = {} + + class _FakeProcess: + def __init__(self, argv, **kwargs): + captured["argv"] = argv + captured["env"] = kwargs["env"] + self.pid = 9898 + + monkeypatch.setattr(shared_mcp_runtime.subprocess, "Popen", _FakeProcess) + + record = shared_mcp_runtime._spawn_daemon(identity, 45678) + + env = captured["env"] + assert isinstance(env, dict) + assert env["OPENSPACE_MCP_DAEMON"] == "1" + assert env["OPENSPACE_MCP_INSTANCE_KEY"] == identity.instance_key + assert env["OPENSPACE_MCP_DAEMON_STATE_DIR"] == identity.state_dir + assert record.instance_key == identity.instance_key + assert record.port == 45678 + + +def test_update_current_daemon_status_timestamps_after_lock_wait(monkeypatch, tmp_path) -> None: + identity = _build_identity(tmp_path) + metadata_path = identity.metadata_path + metadata_path.parent.mkdir(parents=True, exist_ok=True) + shared_mcp_runtime._write_record(metadata_path, _build_record(identity)) + + monkeypatch.setenv("OPENSPACE_MCP_INSTANCE_KEY", identity.instance_key) + monkeypatch.setenv("OPENSPACE_MCP_DAEMON_STATE_DIR", identity.state_dir) + + current_time = {"value": 101.0} + monkeypatch.setattr(shared_mcp_runtime.time, "time", lambda: current_time["value"]) + + result_holder: dict[str, shared_mcp_runtime.MCPDaemonRecord | None] = {} + + with shared_mcp_runtime._FileLock(identity.lock_path): + worker = threading.Thread( + target=lambda: result_holder.setdefault( + "record", + shared_mcp_runtime.update_current_daemon_status("main", warmed=True), + ), + daemon=True, + ) + worker.start() + time.sleep(0.1) + current_time["value"] = 107.5 + + worker.join(timeout=2.0) + + updated = result_holder.get("record") + assert updated is not None + assert updated.warmed is True + assert updated.warmed_at == 107.5