diff --git a/src/openjarvis/agents/_stubs.py b/src/openjarvis/agents/_stubs.py index 3d800b40..3cbf8dbd 100644 --- a/src/openjarvis/agents/_stubs.py +++ b/src/openjarvis/agents/_stubs.py @@ -57,6 +57,10 @@ class BaseAgent(ABC): agent_id: str accepts_tools: bool = False + # Plain conversational agents may opt into the managed runtime's generic + # function-calling loop. Specialized agents keep their own execution + # class even when process-wide MCP tools are available. + supports_managed_tool_fallback: bool = False def __init__( self, diff --git a/src/openjarvis/agents/executor.py b/src/openjarvis/agents/executor.py index ed7126c5..14b70601 100644 --- a/src/openjarvis/agents/executor.py +++ b/src/openjarvis/agents/executor.py @@ -2,7 +2,9 @@ from __future__ import annotations +import json import logging +import threading import time from typing import TYPE_CHECKING, Any @@ -14,6 +16,7 @@ from openjarvis.agents.errors import ( classify_error, retry_delay, ) +from openjarvis.agents.tool_resolver import resolve_agent_tools from openjarvis.core.events import EventBus, EventType if TYPE_CHECKING: @@ -33,6 +36,32 @@ _MAX_RETRIES = 3 _AGENT_TICK_DEFAULT_MODEL = "gemma4:31b" +def _tool_calls_for_storage(result: AgentResult) -> list[dict[str, Any]] | None: + """Convert executor tool results to the managed-message storage contract.""" + + calls: list[dict[str, Any]] = [] + for tool_result in result.tool_results: + metadata = getattr(tool_result, "metadata", {}) or {} + arguments = metadata.get("arguments", "") + if not isinstance(arguments, str): + try: + arguments = json.dumps(arguments, sort_keys=True) + except (TypeError, ValueError): + arguments = json.dumps(str(arguments)) + calls.append( + { + "tool": getattr(tool_result, "tool_name", ""), + "arguments": arguments, + "result": getattr(tool_result, "content", "") or "", + "success": bool(getattr(tool_result, "success", False)), + # SSE and the frontend persist/display latency in milliseconds. + "latency": float(getattr(tool_result, "latency_seconds", 0.0) or 0.0) + * 1000.0, + } + ) + return calls or None + + class AgentExecutor: """Executes a single tick for a managed agent. @@ -51,6 +80,7 @@ class AgentExecutor: self._manager = manager self._bus = event_bus self._trace_store = trace_store + self._toolkit_local = threading.local() def set_system(self, system: Any) -> None: """Deferred system injection — called after JarvisSystem is constructed.""" @@ -63,27 +93,6 @@ class AgentExecutor: except Exception: pass # Non-critical - def _inject_tool_deps(self, tool: Any) -> None: - """Inject runtime dependencies into a tool instance. - - Mirrors SystemBuilder._inject_tool_deps (system.py:920-945) - but uses the lightweight system's references. - """ - if self._system is None: - return - name = getattr(getattr(tool, "spec", None), "name", "") - if name == "llm": - if hasattr(tool, "_engine"): - tool._engine = self._system.engine - if hasattr(tool, "_model"): - tool._model = self._system.model - elif name == "retrieval" or name.startswith("memory_"): - if hasattr(tool, "_backend"): - tool._backend = getattr(self._system, "memory_backend", None) - elif name.startswith("channel_"): - if hasattr(tool, "_channel"): - tool._channel = getattr(self._system, "channel_backend", None) - def run_ephemeral( self, agent_type: str, @@ -248,7 +257,20 @@ class AgentExecutor: raise last_error or FatalError("max retries exhausted") def _invoke_agent(self, agent: dict) -> AgentResult: - """Invoke the actual agent run. Tests mock this method.""" + """Invoke one agent while owning every resource its resolver opens.""" + + previous = getattr(self._toolkit_local, "current", None) + self._toolkit_local.current = None + try: + return self._invoke_agent_impl(agent) + finally: + current = getattr(self._toolkit_local, "current", None) + if current is not None: + current.close() + self._toolkit_local.current = previous + + def _invoke_agent_impl(self, agent: dict) -> AgentResult: + """Implementation split out so the wrapper owns resolver lifetime.""" from openjarvis.agents import AgentRegistry agent_type = agent.get("agent_type", "monitor_operative") @@ -257,6 +279,10 @@ class AgentExecutor: raise FatalError(f"Unknown agent type: {agent_type}") config = agent.get("config", {}) + agent_accepts_tools = bool(getattr(agent_cls, "accepts_tools", False)) + supports_tool_fallback = bool( + getattr(agent_cls, "supports_managed_tool_fallback", False) + ) # Resolve engine + model from JarvisSystem engine = self._system.engine if self._system else None @@ -300,64 +326,88 @@ class AgentExecutor: except Exception: pass # Fall back to configured model - # Resolve tools from config via ToolRegistry - tool_names = config.get("tools", []) - if isinstance(tool_names, str): - tool_names = [t.strip() for t in tool_names.split(",") if t.strip()] + mcp_tools: list[Any] = [] + mcp_clients: list[Any] = [] + if ( + config.get("mcp_tools", True) is not False + and self._system is not None + and (agent_accepts_tools or supports_tool_fallback) + ): + provider = getattr( + self._system, + "get_managed_agent_mcp_tools", + None, + ) + if callable(provider): + try: + mcp_tools, mcp_clients = provider() + except Exception as exc: + logger.warning("Managed-agent MCP discovery failed: %s", exc) + else: + mcp_tools = list(getattr(self._system, "mcp_tools", []) or []) + mcp_clients = list(getattr(self._system, "_mcp_clients", []) or []) - tool_instances: list[Any] = [] - if tool_names: - try: - from openjarvis.server.agent_manager_routes import ( - _ensure_registries_populated, - ) + if not mcp_tools: + try: + from openjarvis.tools.mcp_adapter import MCPToolAdapter - _ensure_registries_populated() - except ImportError: - pass - from openjarvis.core.registry import ToolRegistry + pool = ( + getattr( + getattr(self._system, "tool_executor", None), + "_tools", + {}, + ) + or {} + ) + mcp_tools = [ + tool + for tool in pool.values() + if isinstance(tool, MCPToolAdapter) + ] + except Exception: + mcp_tools = [] - for tname in tool_names: - if ToolRegistry.contains(tname): - try: - tool_cls = ToolRegistry.get(tname) - tool = tool_cls() - self._inject_tool_deps(tool) - tool_instances.append(tool) - except Exception: - logger.warning("Failed to instantiate tool %s", tname) + resolved_toolkit = resolve_agent_tools( + agent, + engine=engine, + model=model, + memory_backend=getattr(self._system, "memory_backend", None), + channel_backend=getattr(self._system, "channel_backend", None), + mcp_tools=mcp_tools, + mcp_clients=mcp_clients, + knowledge_db_path=getattr(self._system, "knowledge_db_path", None), + ) + self._toolkit_local.current = resolved_toolkit + tool_instances = resolved_toolkit.instances + logger.info( + "Agent %s: resolved %d tools (%s)", + agent["name"], + len(tool_instances), + ", ".join(resolved_toolkit.by_name) or "none", + ) - # Pull tools already discovered by SystemBuilder (e.g. external MCP - # adapters) that aren't in the static ToolRegistry. Without this, - # agents declaring MCP-discovered tools in their template would - # silently fall back to natives only. - if ( - self._system is not None - and getattr(self._system, "tool_executor", None) is not None - ): - mcp_pool = getattr(self._system.tool_executor, "_tools", {}) or {} - existing = {t.spec.name for t in tool_instances} - for tname in tool_names: - if tname in existing: - continue - pooled = mcp_pool.get(tname) - if pooled is not None: - tool_instances.append(pooled) + execution_agent_cls = agent_cls + if tool_instances and not agent_accepts_tools and supports_tool_fallback: + # Managed SSE already runs configured tools through a native + # function-calling loop regardless of the selected class. Use the + # same capability for immediate/scheduled ticks instead of + # silently discarding the resolved toolkit for SimpleAgent and + # other explicitly compatible non-tool classes. + from openjarvis.agents.orchestrator import OrchestratorAgent - if tool_instances: - logger.info( - "Agent %s: resolved %d/%d tools", - agent["name"], - len(tool_instances), - len(tool_names), - ) + execution_agent_cls = OrchestratorAgent + logger.info( + "Agent %s: %s does not accept tools; using %s for this " + "tool-enabled tick", + agent["name"], + agent_cls.__name__, + execution_agent_cls.__name__, + ) # Construct agent instance agent_kwargs: dict[str, Any] = {} sys_prompt = config.get("system_prompt") - if sys_prompt is not None: - agent_kwargs["system_prompt"] = sys_prompt - if getattr(agent_cls, "accepts_tools", False) and tool_instances: + if getattr(execution_agent_cls, "accepts_tools", False) and tool_instances: agent_kwargs["tools"] = tool_instances # Hand the agent our EventBus so its ToolExecutor can publish # TOOL_CALL_START/END — without this, ToolExecutor's ``self._bus`` @@ -379,7 +429,7 @@ class AgentExecutor: # recall / persistence paths. import inspect - init_sig = inspect.signature(agent_cls.__init__) + init_sig = inspect.signature(execution_agent_cls.__init__) accepts_var_kw = any( p.kind == inspect.Parameter.VAR_KEYWORD for p in init_sig.parameters.values() @@ -388,6 +438,16 @@ class AgentExecutor: def _accepts(name: str) -> bool: return accepts_var_kw or name in init_sig.parameters + # Unsupported kwargs used to trigger the broad TypeError fallback + # below, which retried with a bare constructor and silently discarded + # valid prompt/state wiring. Filter by the selected class's signature + # before construction instead. + if sys_prompt is not None and _accepts("system_prompt"): + agent_kwargs["system_prompt"] = sys_prompt + agent_kwargs = { + name: value for name, value in agent_kwargs.items() if _accepts(name) + } + state_kwargs: dict[str, Any] = {} if _accepts("operator_id"): state_kwargs["operator_id"] = agent["id"] @@ -404,23 +464,49 @@ class AgentExecutor: # agents, mirroring the one-shot `jarvis ask` path so they no # longer apply to CLI calls only (#376). cfg = getattr(self._system, "config", None) - if cfg is not None and _accepts("prompt_builder"): + if _accepts("prompt_builder") and ( + cfg is not None or sys_prompt is not None + ): from openjarvis.prompt.builder import SystemPromptBuilder state_kwargs["prompt_builder"] = SystemPromptBuilder( - agent_template=getattr(cfg.agent, "default_system_prompt", "") - or "", - memory_files_config=cfg.memory_files, - system_prompt_config=cfg.system_prompt, + agent_template=( + sys_prompt + if sys_prompt is not None + else getattr( + getattr(cfg, "agent", None), + "default_system_prompt", + "", + ) + or "" + ), + memory_files_config=getattr(cfg, "memory_files", None), + system_prompt_config=getattr(cfg, "system_prompt", None), ) try: - agent_instance = agent_cls(engine, model, **agent_kwargs, **state_kwargs) - except TypeError: try: - agent_instance = agent_cls(engine, model, **agent_kwargs) + agent_instance = execution_agent_cls( + engine, + model, + **agent_kwargs, + **state_kwargs, + ) except TypeError: - agent_instance = agent_cls(engine, model) + try: + agent_instance = execution_agent_cls( + engine, + model, + **agent_kwargs, + ) + except TypeError: + agent_instance = execution_agent_cls(engine, model) + except Exception: + resolved_toolkit.close() + raise + + if resolved_toolkit.mcp_clients: + agent_instance._mcp_clients = resolved_toolkit.mcp_clients # Inject the managed-agent UUID into the agent's ToolExecutor so # emitted TOOL_CALL_START/END events carry it; the trace subscriber @@ -436,7 +522,7 @@ class AgentExecutor: agent["name"], len(tool_instances), ", ".join(t.spec.name for t in tool_instances) or "none", - agent_cls.__name__, + execution_agent_cls.__name__, ) # Build input from instruction + summary_memory + pending messages. @@ -551,21 +637,24 @@ class AgentExecutor: len(input_text), ) _t0 = time.time() - result = agent_instance.run(input_text, context=agent_ctx) - - # Retry once if the model returned empty content (common with - # Qwen3.5 thinking mode consuming all tokens). - if not (result.content or "").strip(): - self._set_activity( - agent["id"], - "Retrying (empty response)...", - ) - logger.warning( - "Agent %s: empty content, retrying once", - agent["name"], - ) + try: result = agent_instance.run(input_text, context=agent_ctx) + # Retry once if the model returned empty content (common with + # Qwen3.5 thinking mode consuming all tokens). + if not (result.content or "").strip(): + self._set_activity( + agent["id"], + "Retrying (empty response)...", + ) + logger.warning( + "Agent %s: empty content, retrying once", + agent["name"], + ) + result = agent_instance.run(input_text, context=agent_ctx) + finally: + resolved_toolkit.close() + _elapsed = time.time() - _t0 logger.info( "Agent %s: agent.run() completed in %.1fs, " @@ -655,7 +744,11 @@ class AgentExecutor: # message keeps the complete report. The old [:2000] slices # double-truncated and cut findings off mid-sentence. self._manager.update_summary_memory(agent_id, result.content) - self._manager.store_agent_response(agent_id, result.content) + self._manager.store_agent_response( + agent_id, + result.content, + tool_calls=_tool_calls_for_storage(result), + ) # Budget enforcement (post-tick check) agent_data = self._manager.get_agent(agent_id) diff --git a/src/openjarvis/agents/orchestrator.py b/src/openjarvis/agents/orchestrator.py index cca03d2c..7241cb38 100644 --- a/src/openjarvis/agents/orchestrator.py +++ b/src/openjarvis/agents/orchestrator.py @@ -57,6 +57,7 @@ class OrchestratorAgent(ToolUsingAgent): max_tokens: Optional[int] = None, mode: str = "function_calling", system_prompt: Optional[str] = None, + prompt_builder: Optional[Any] = None, parallel_tools: bool = True, interactive: bool = False, confirm_callback=None, @@ -71,6 +72,7 @@ class OrchestratorAgent(ToolUsingAgent): max_tokens=max_tokens, interactive=interactive, confirm_callback=confirm_callback, + prompt_builder=prompt_builder, ) self._mode = mode self._system_prompt = system_prompt @@ -214,7 +216,11 @@ class OrchestratorAgent(ToolUsingAgent): self._emit_turn_start(input) # Build initial messages - messages = self._build_messages(input, context) + messages = self._build_messages( + input, + context, + system_prompt=self._system_prompt, + ) # Get OpenAI-format tool definitions openai_tools = self._executor.get_openai_tools() if self._tools else [] diff --git a/src/openjarvis/agents/scheduler.py b/src/openjarvis/agents/scheduler.py index 157bcf40..5ac6f72e 100644 --- a/src/openjarvis/agents/scheduler.py +++ b/src/openjarvis/agents/scheduler.py @@ -123,15 +123,35 @@ class AgentScheduler: self._thread.start() logger.info("Agent scheduler started") - def stop(self) -> None: - """Stop the scheduler background thread.""" + def request_stop(self) -> None: + """Prevent new scheduled ticks without waiting for the worker.""" + self._stop_event.set() if self._bus: self._bus.unsubscribe(EventType.AGENT_TICK_END, self._on_tick_event) - if self._thread is not None: - self._thread.join(timeout=10) + + def wait_stopped(self, timeout: float = 10.0) -> bool: + """Wait for an active tick to finish, retaining live thread state.""" + + thread = self._thread + if thread is None: + return True + if thread is threading.current_thread(): + return False + thread.join(timeout=timeout) + if thread.is_alive(): + logger.warning("Agent scheduler did not stop within %.1fs", timeout) + return False + if self._thread is thread: self._thread = None - logger.info("Agent scheduler stopped") + return True + + def stop(self, timeout: float = 10.0) -> None: + """Stop dispatching and wait for the scheduler worker.""" + + self.request_stop() + if self.wait_stopped(timeout=timeout): + logger.info("Agent scheduler stopped") def _loop(self) -> None: """Main scheduler loop.""" @@ -160,6 +180,8 @@ class AgentScheduler: ] for agent_id, info in due: + if self._stop_event.is_set(): + break agent = self._manager.get_agent(agent_id) if agent is None or agent["status"] in ( "paused", diff --git a/src/openjarvis/agents/simple.py b/src/openjarvis/agents/simple.py index 9af56b59..748cb5f8 100644 --- a/src/openjarvis/agents/simple.py +++ b/src/openjarvis/agents/simple.py @@ -13,6 +13,7 @@ class SimpleAgent(BaseAgent): """Single-turn agent: query -> model -> response. No tool calling.""" agent_id = "simple" + supports_managed_tool_fallback = True def run( self, diff --git a/src/openjarvis/agents/tool_resolver.py b/src/openjarvis/agents/tool_resolver.py new file mode 100644 index 00000000..404b92b1 --- /dev/null +++ b/src/openjarvis/agents/tool_resolver.py @@ -0,0 +1,502 @@ +"""Canonical managed-agent tool resolution. + +Managed agents can run through streaming HTTP, immediate/scheduled ticks, or +the persistent-agent CLI. Those paths must bind the same live tool instances: +agent-type grants first, then configured native tools, then MCP adapters. +""" + +from __future__ import annotations + +import importlib +import logging +import sys +import weakref +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Iterable, Mapping + +logger = logging.getLogger(__name__) + +BROWSER_SUB_TOOLS = ( + "browser_navigate", + "browser_click", + "browser_type", + "browser_screenshot", + "browser_extract", + "browser_axtree", +) + +_MEMORY_TOOLS = frozenset( + {"retrieval", "memory_store", "memory_search", "memory_index", "memory_retrieve"} +) +_CHANNEL_TOOLS = frozenset({"channel_send", "channel_list", "channel_status"}) + + +class _SpecOverrideTool: + """Delegate execution while exposing an agent-configured OpenAI schema.""" + + def __init__(self, wrapped: Any, advertised_spec: dict[str, Any]) -> None: + self._wrapped = wrapped + self._advertised_spec = advertised_spec + + @property + def spec(self) -> Any: + base = self._wrapped.spec + function = self._advertised_spec.get("function", {}) + return replace( + base, + name=function.get("name", base.name), + description=function.get("description", base.description), + parameters=function.get("parameters", base.parameters), + ) + + def execute(self, **params: Any) -> Any: + return self._wrapped.execute(**params) + + def to_openai_function(self) -> dict[str, Any]: + return self._advertised_spec + + def __getattr__(self, name: str) -> Any: + return getattr(self._wrapped, name) + + +def _tool_name(tool: Any) -> str: + try: + return str(tool.spec.name) + except Exception: + return "" + + +def _spec_name(spec: Mapping[str, Any]) -> str: + function = spec.get("function") + if not isinstance(function, Mapping): + return "" + name = function.get("name") + return str(name) if name else "" + + +def _openai_spec(tool: Any) -> dict[str, Any]: + to_openai_function = getattr(tool, "to_openai_function", None) + if callable(to_openai_function): + try: + advertised = to_openai_function() + except Exception: + logger.debug( + "Failed to build advertised schema for tool %r; falling back " + "to its ToolSpec", + _tool_name(tool), + exc_info=True, + ) + else: + if isinstance(advertised, Mapping) and _spec_name(advertised): + return dict(advertised) + logger.debug( + "Tool %r returned an invalid advertised schema; falling back " + "to its ToolSpec", + _tool_name(tool), + ) + + spec = tool.spec + return { + "type": "function", + "function": { + "name": spec.name, + "description": spec.description, + "parameters": spec.parameters, + }, + } + + +def _close_resources(resources: tuple[Any, ...]) -> None: + for resource in reversed(resources): + close = getattr(resource, "close", None) + if callable(close): + try: + close() + except Exception: + logger.debug("Failed to close resolved tool resource", exc_info=True) + + +@dataclass +class ResolvedAgentTools: + """One resolved toolkit, with views for agent loops and raw streaming.""" + + instances: list[Any] = field(default_factory=list) + extra_specs: list[dict[str, Any]] = field(default_factory=list) + advertised_specs: list[dict[str, Any]] = field(default_factory=list) + mcp_clients: list[Any] = field(default_factory=list) + owned_resources: list[Any] = field(default_factory=list, repr=False) + _closed: bool = field(default=False, init=False, repr=False) + _finalizer: weakref.finalize = field(init=False, repr=False) + + def __post_init__(self) -> None: + # This fallback covers exceptions anywhere after resolution, including + # before an executor/response installs its normal explicit cleanup. + self._finalizer = weakref.finalize( + self, + _close_resources, + tuple(self.owned_resources), + ) + + @property + def by_name(self) -> dict[str, Any]: + return {name: tool for tool in self.instances if (name := _tool_name(tool))} + + @property + def openai_specs(self) -> list[dict[str, Any]]: + specs: list[dict[str, Any]] = [] + seen: set[str] = set() + advertised = self.advertised_specs + if not advertised: + advertised = [*map(_openai_spec, self.instances), *self.extra_specs] + for spec in advertised: + name = _spec_name(spec) + if name and name in seen: + continue + specs.append(spec) + if name: + seen.add(name) + return specs + + def close(self) -> None: + """Close request-local resources without touching shared MCP clients.""" + + if self._closed: + return + self._closed = True + self._finalizer() + + def __enter__(self) -> ResolvedAgentTools: + return self + + def __exit__(self, *exc_info: object) -> None: + self.close() + + +def ensure_registries_populated() -> None: + """Populate tool/channel registries, including after tests clear them.""" + + from openjarvis.core.registry import ChannelRegistry, ToolRegistry + + try: + import openjarvis.channels # noqa: F401 + except Exception: + pass + + try: + import openjarvis.tools # noqa: F401 + except Exception: + pass + + browser_modules = ("openjarvis.tools.browser", "openjarvis.tools.browser_axtree") + for module_name in browser_modules: + try: + importlib.import_module(module_name) + except Exception: + pass + + if not ChannelRegistry.keys(): + for module_name in list(sys.modules): + if module_name.startswith( + "openjarvis.channels." + ) and not module_name.endswith("_stubs"): + try: + importlib.reload(sys.modules[module_name]) + except Exception: + pass + + if not ToolRegistry.keys(): + for module_name in list(sys.modules): + if ( + module_name.startswith("openjarvis.tools.") + and not module_name.endswith("_stubs") + and not module_name.endswith("agent_tools") + ): + try: + importlib.reload(sys.modules[module_name]) + except Exception: + pass + + if not any(ToolRegistry.contains(name) for name in BROWSER_SUB_TOOLS): + for module_name in browser_modules: + module = sys.modules.get(module_name) + if module is not None: + try: + importlib.reload(module) + except Exception: + pass + + +def instantiate_registered_tool( + tool_cls: Any, + name: str, + *, + engine: Any, + model: str, + memory_backend: Any = None, + channel_backend: Any = None, +) -> Any: + """Instantiate a registry tool with its runtime dependencies.""" + + if name in _MEMORY_TOOLS: + if memory_backend is None: + logger.warning( + "Memory tool %r instantiated without a backend — calls will " + "return no results.", + name, + ) + return tool_cls(backend=memory_backend) + if name in _CHANNEL_TOOLS: + if channel_backend is None: + logger.warning( + "Channel tool %r instantiated without a channel — calls will " + "fail with 'No channel backend configured'.", + name, + ) + return tool_cls(channel=channel_backend) + if name == "llm": + return tool_cls(engine=engine, model=model) + return tool_cls() + + +def build_deep_research_tools( + engine: Any, + model: str, + knowledge_db_path: str | Path | None = None, +) -> list[Any]: + """Construct the live knowledge tools granted to ``deep_research``.""" + + if not knowledge_db_path: + from openjarvis.core.config import DEFAULT_CONFIG_DIR + + knowledge_db_path = DEFAULT_CONFIG_DIR / "knowledge.db" + + path = Path(knowledge_db_path) + if not path.exists(): + return [] + + from openjarvis.connectors.retriever import TwoStageRetriever + from openjarvis.connectors.store import KnowledgeStore + from openjarvis.tools.knowledge_search import KnowledgeSearchTool + from openjarvis.tools.knowledge_sql import KnowledgeSQLTool + from openjarvis.tools.scan_chunks import ScanChunksTool + from openjarvis.tools.think import ThinkTool + + store = KnowledgeStore(str(path)) + try: + retriever = TwoStageRetriever(store) + return [ + KnowledgeSearchTool(retriever=retriever), + KnowledgeSQLTool(store=store), + ScanChunksTool(store=store, engine=engine, model=model), + ThinkTool(), + ] + except Exception: + store.close() + raise + + +def _normalized_tool_config(tool_config: Any) -> list[Any]: + if not tool_config: + return [] + if isinstance(tool_config, str): + return [part.strip() for part in tool_config.split(",") if part.strip()] + if isinstance(tool_config, Mapping): + return [dict(tool_config)] + try: + return list(tool_config) + except TypeError: + return [] + + +def resolve_agent_tools( + agent_record: Mapping[str, Any], + *, + engine: Any, + model: str, + memory_backend: Any = None, + channel_backend: Any = None, + mcp_tools: Iterable[Any] = (), + mcp_clients: Iterable[Any] = (), + knowledge_db_path: str | Path | None = None, +) -> ResolvedAgentTools: + """Resolve the effective live toolkit for a managed agent. + + Resolution is stable and first-wins: agent-type grants take precedence + over configured registry tools, which take precedence over MCP adapters. + ``config["mcp_tools"] = false`` excludes MCP adapters from this agent; + process-wide runtimes may still own connections used by other agents. + """ + + ensure_registries_populated() + from openjarvis.core.registry import ChannelRegistry, ToolRegistry + + config = agent_record.get("config") or {} + if not isinstance(config, Mapping): + config = {} + + instances: list[Any] = [] + extra_specs: list[dict[str, Any]] = [] + advertised_specs: list[dict[str, Any]] = [] + owned_resources: list[Any] = [] + seen: set[str] = set() + + def add_instance( + tool: Any, + *, + advertised_spec: dict[str, Any] | None = None, + ) -> None: + name = _tool_name(tool) + if not name or name in seen: + return + instances.append(tool) + advertised_specs.append(advertised_spec or _openai_spec(tool)) + seen.add(name) + + use_mcp = config.get("mcp_tools", True) is not False + mcp_tool_list = list(mcp_tools) if use_mcp else [] + mcp_by_name: dict[str, Any] = {} + for tool in mcp_tool_list: + name = _tool_name(tool) + if name and name not in mcp_by_name: + mcp_by_name[name] = tool + + if agent_record.get("agent_type") == "deep_research": + granted_tools = build_deep_research_tools( + engine=engine, + model=model, + knowledge_db_path=knowledge_db_path, + ) + owned_ids: set[int] = set() + for tool in granted_tools: + resource = getattr(tool, "_store", None) + if ( + resource is not None + and callable(getattr(resource, "close", None)) + and id(resource) not in owned_ids + ): + owned_resources.append(resource) + owned_ids.add(id(resource)) + add_instance(tool) + + for entry in _normalized_tool_config(config.get("tools")): + if isinstance(entry, Mapping): + raw_spec = entry if isinstance(entry, dict) else dict(entry) + name = _spec_name(raw_spec) + if name and name in seen: + continue + + backing_tool = None + if name and not ChannelRegistry.contains(name): + if ToolRegistry.contains(name): + try: + backing_tool = instantiate_registered_tool( + ToolRegistry.get(name), + name, + engine=engine, + model=model, + memory_backend=memory_backend, + channel_backend=channel_backend, + ) + except Exception as exc: + logger.warning( + "Could not instantiate tool '%s' (%s) — " + "advertising its custom spec without execution", + name, + exc, + ) + elif name in mcp_by_name: + backing_tool = mcp_by_name[name] + + if backing_tool is not None: + add_instance( + _SpecOverrideTool(backing_tool, raw_spec), + advertised_spec=raw_spec, + ) + else: + logger.warning( + "Custom tool spec '%s' has no registered or MCP execution " + "backend — dropping", + name or "", + ) + continue + if not isinstance(entry, str): + continue + + names = BROWSER_SUB_TOOLS if entry == "browser" else (entry,) + for name in names: + if name in seen: + continue + if ChannelRegistry.contains(name): + continue + if not ToolRegistry.contains(name): + logger.warning( + "Tool '%s' referenced in agent config but not in ToolRegistry", + name, + ) + continue + try: + add_instance( + instantiate_registered_tool( + ToolRegistry.get(name), + name, + engine=engine, + model=model, + memory_backend=memory_backend, + channel_backend=channel_backend, + ) + ) + except Exception as exc: + logger.warning( + "Could not instantiate tool '%s' (%s) — dropping", name, exc + ) + + if use_mcp: + for tool in mcp_tool_list: + add_instance(tool) + + return ResolvedAgentTools( + instances=instances, + extra_specs=extra_specs, + advertised_specs=advertised_specs, + mcp_clients=list(mcp_clients) if use_mcp else [], + owned_resources=owned_resources, + ) + + +def resolve_tool_specs(tool_config: Any) -> list[dict[str, Any]]: + """Compatibility view for callers that only need configured specs.""" + + specs: list[dict[str, Any]] = [] + seen: set[str] = set() + for entry in _normalized_tool_config(tool_config): + if isinstance(entry, dict): + specs.append(entry) + name = _spec_name(entry) + if name: + seen.add(name) + continue + resolved = resolve_agent_tools( + {"config": {"tools": [entry]}}, + engine=None, + model="", + ) + for spec in resolved.openai_specs: + name = _spec_name(spec) + if name and name in seen: + continue + specs.append(spec) + if name: + seen.add(name) + return specs + + +__all__ = [ + "BROWSER_SUB_TOOLS", + "ResolvedAgentTools", + "build_deep_research_tools", + "ensure_registries_populated", + "instantiate_registered_tool", + "resolve_agent_tools", + "resolve_tool_specs", +] diff --git a/src/openjarvis/cli/serve.py b/src/openjarvis/cli/serve.py index 6c9e415a..a9e346a3 100644 --- a/src/openjarvis/cli/serve.py +++ b/src/openjarvis/cli/serve.py @@ -279,6 +279,15 @@ def serve( # (which would re-discover the engine, re-resolve tools, re-open the channel, # etc.). See the scheduler block near the bottom of this function (#263). resolved_tools: list = [] + managed_mcp_tools: list = [] + mcp_clients: list = [] + try: + from openjarvis.mcp.loader import load_mcp_tools_from_config + + managed_mcp_tools, mcp_clients = load_mcp_tools_from_config(config.tools.mcp) + except Exception as exc: + logger.warning("Managed-agent MCP tools failed to load: %s", exc) + if agent_key: try: import openjarvis.agents # noqa: F401 @@ -290,11 +299,6 @@ def serve( if sec.capability_policy is not None: agent_kwargs["capability_policy"] = sec.capability_policy - # MCP transports persisted on the agent at the bottom of - # this block — initialise here so the reference is valid - # even when accepts_tools is False (#461). - mcp_clients: list = [] - # Load tools for agents that support them if getattr(agent_cls, "accepts_tools", False): import openjarvis.tools # noqa: F401 # trigger registration @@ -331,12 +335,13 @@ def serve( # MCP server tools from config.tools.mcp.servers # (#461 — these were silently dropped). - from openjarvis.mcp.loader import load_mcp_tools_from_config - - mcp_tools, mcp_clients = load_mcp_tools_from_config( - config.tools.mcp, - allowed_names=allowed if configured else None, - ) + mcp_tools = managed_mcp_tools + if configured: + mcp_tools = [ + tool + for tool in managed_mcp_tools + if tool.spec.name in allowed + ] if mcp_tools: existing = {t.spec.name for t in tools} for t in mcp_tools: @@ -389,10 +394,6 @@ def serve( channel_agent = config.channel.default_agent or agent_key or "simple" _channel_tools: list = [] - # MCP transports persisted at function scope (= server-process - # lifetime); see the comment near the channel-MCP-load block - # below. Initialise here so it's always bound. #461. - _channel_mcp_clients: list = [] if channel_agent: try: import openjarvis.agents @@ -432,29 +433,23 @@ def serve( elif isinstance(_tcls, BaseTool): _channel_tools.append(_tcls) - # MCP tools for the channel agent too (#461). - from openjarvis.mcp.loader import ( - load_mcp_tools_from_config, - ) - - _ch_mcp_tools, _ch_mcp_clients = load_mcp_tools_from_config( - config.tools.mcp, - allowed_names=_allowed if configured else None, - ) + # Reuse the process-owned MCP pool so channels do not + # open a second transport to every configured server. + _ch_mcp_tools = managed_mcp_tools + if configured: + _ch_mcp_tools = [ + tool + for tool in managed_mcp_tools + if tool.spec.name in _allowed + ] if _ch_mcp_tools: _existing = {t.spec.name for t in _channel_tools} for t in _ch_mcp_tools: if t.spec.name not in _existing: _channel_tools.append(t) _existing.add(t.spec.name) - # Hold a reference at module / function scope — - # the channel agent is constructed inside - # JarvisSystem below; we extend its lifetime by - # keeping the list bound here. - _channel_mcp_clients = _ch_mcp_clients except Exception as exc: logger.warning("Channel tools failed to load: %s", exc) - _channel_mcp_clients = [] _wire_system = JarvisSystem( config=config, @@ -464,6 +459,8 @@ def serve( model=model_name, agent_name=channel_agent, tools=_channel_tools, + mcp_tools=managed_mcp_tools, + _mcp_clients=mcp_clients, ) _wire_system.wire_channel(channel_bridge) @@ -481,23 +478,24 @@ def serve( # Create app from openjarvis.server.app import create_app - # Set up memory backend for context injection. Built before the scheduler - # block so the executor's JarvisSystem can reference it (#263). + # Set up the memory backend for storage tools, API routes, and optional + # prompt-context injection. ``context_from_memory`` controls only the last + # of those, so disabling it must not leave explicit memory_* tools with a + # null backend. Built before the scheduler so AgentExecutor can reuse it. memory_backend = None - if config.agent.context_from_memory: - try: - import openjarvis.tools.storage # noqa: F401 - from openjarvis.core.registry import MemoryRegistry + try: + import openjarvis.tools.storage # noqa: F401 + from openjarvis.core.registry import MemoryRegistry - mem_key = config.memory.default_backend - if MemoryRegistry.contains(mem_key): - memory_backend = MemoryRegistry.create( - mem_key, - db_path=config.memory.db_path, - ) - console.print(" Memory: [cyan]active[/cyan]") - except Exception as exc: - logger.debug("Memory backend init failed: %s", exc) + mem_key = config.memory.default_backend + if MemoryRegistry.contains(mem_key): + memory_backend = MemoryRegistry.create( + mem_key, + db_path=config.memory.db_path, + ) + console.print(" Memory: [cyan]active[/cyan]") + except Exception as exc: + logger.debug("Memory backend init failed: %s", exc) # Automatic long-term memory service (background fact extraction). memory_service = None @@ -592,6 +590,7 @@ def serve( agent=agent, agent_name=agent_key or "", tools=resolved_tools, + mcp_tools=managed_mcp_tools, tool_executor=_sched_tool_executor, memory_backend=memory_backend, telemetry_store=telem_store, @@ -600,6 +599,7 @@ def serve( capability_policy=sec.capability_policy, agent_manager=agent_manager, agent_executor=executor, + _mcp_clients=mcp_clients, ) executor.set_system(system) @@ -691,10 +691,13 @@ def serve( channel_bridge=channel_bridge, config=config, memory_backend=memory_backend, + own_memory_backend=memory_backend is not None, memory_service=memory_service, speech_backend=speech_backend, agent_manager=agent_manager, agent_scheduler=agent_scheduler, + mcp_tools=managed_mcp_tools, + mcp_clients=mcp_clients, api_key=api_key, webhook_config=webhook_config, cors_origins=config.server.cors_origins, @@ -723,6 +726,4 @@ def serve( "authenticated requests to your instance." ) - import uvicorn - uvicorn.run(app, host=bind_host, port=bind_port, log_level="info") diff --git a/src/openjarvis/mcp/client.py b/src/openjarvis/mcp/client.py index 5b495922..a18b2032 100644 --- a/src/openjarvis/mcp/client.py +++ b/src/openjarvis/mcp/client.py @@ -3,6 +3,7 @@ from __future__ import annotations import itertools +import threading from typing import Any, Dict, List from openjarvis.mcp.protocol import MCPError, MCPRequest, MCPResponse @@ -24,18 +25,31 @@ class MCPClient: self._initialized = False self._capabilities: Dict[str, Any] = {} self._id_counter = itertools.count(1) + # A client may be shared by server, scheduled, and channel agents. + # Keep each transport request/response exchange atomic so stdio + # readers cannot consume another thread's JSON-RPC response. + self._request_lock = threading.RLock() + # Closing must not wait for ``_request_lock``: transport.close() is + # what interrupts a request that is blocked in a transport read. + # An event lets queued requests fail before touching that transport, + # while this separate lock keeps close itself idempotent. + self._closed = threading.Event() + self._transport_closed = threading.Event() + self._close_lock = threading.Lock() def _next_id(self) -> int: return next(self._id_counter) def _send(self, method: str, params: Dict[str, Any] | None = None) -> MCPResponse: """Send a request and check for errors.""" - request = MCPRequest( - method=method, - params=params or {}, - id=self._next_id(), - ) - response = self._transport.send(request) + with self._request_lock: + self._raise_if_closed() + request = MCPRequest( + method=method, + params=params or {}, + id=self._next_id(), + ) + response = self._transport.send(request) if response.error is not None: raise MCPError( code=response.error.get("code", -1), @@ -44,6 +58,10 @@ class MCPClient: ) return response + def _raise_if_closed(self) -> None: + if self._closed.is_set(): + raise RuntimeError("MCP client is closed") + def initialize(self) -> Dict[str, Any]: """Perform the MCP initialize handshake. @@ -75,7 +93,9 @@ class MCPClient: params=params or {}, id=None, # None → no id field in JSON (notification) ) - self._transport.send_notification(request) + with self._request_lock: + self._raise_if_closed() + self._transport.send_notification(request) def list_tools(self) -> List[ToolSpec]: """Discover available tools from the server. @@ -114,7 +134,15 @@ class MCPClient: def close(self) -> None: """Close the transport connection.""" - self._transport.close() + # Do not acquire _request_lock here. A transport request can be stuck + # waiting for a server response, and closing the underlying transport + # is the mechanism that unblocks it. + with self._close_lock: + if self._transport_closed.is_set(): + return + self._closed.set() + self._transport.close() + self._transport_closed.set() def __enter__(self) -> MCPClient: return self diff --git a/src/openjarvis/server/agent_manager_routes.py b/src/openjarvis/server/agent_manager_routes.py index 5b33491c..afb4a591 100644 --- a/src/openjarvis/server/agent_manager_routes.py +++ b/src/openjarvis/server/agent_manager_routes.py @@ -4,9 +4,22 @@ from __future__ import annotations import logging import re as _re +import threading from typing import Any, Dict, List, Optional, Tuple from openjarvis.agents.manager import AgentManager +from openjarvis.agents.tool_resolver import ( + BROWSER_SUB_TOOLS as _BROWSER_SUB_TOOLS, +) +from openjarvis.agents.tool_resolver import ( + build_deep_research_tools, + instantiate_registered_tool, + resolve_agent_tools, + resolve_tool_specs, +) +from openjarvis.agents.tool_resolver import ( + ensure_registries_populated as _ensure_registries_populated, +) try: from fastapi import APIRouter, HTTPException, Request @@ -16,6 +29,44 @@ except ImportError: raise ImportError("fastapi and pydantic are required for server routes") logger = logging.getLogger("openjarvis.server.agent_manager") +_MEMORY_BACKEND_LOCK_SETUP = threading.Lock() +_MCP_LOCK_SETUP = threading.Lock() + + +def _start_managed_worker(app_state: Any, target: Any, *, name: str) -> Any: + """Start and track a managed-agent worker for orderly app shutdown.""" + + lock = getattr(app_state, "_managed_worker_lock", None) + if lock is None: + lock = threading.Lock() + app_state._managed_worker_lock = lock + workers = getattr(app_state, "_managed_workers", None) + if workers is None: + workers = set() + app_state._managed_workers = workers + + thread: threading.Thread + + def _run() -> None: + try: + target() + finally: + with lock: + workers.discard(thread) + + with lock: + if getattr(app_state, "_managed_runtime_stopping", False): + raise RuntimeError("managed-agent runtime is shutting down") + thread = threading.Thread(target=_run, daemon=True, name=name) + workers.add(thread) + try: + # Shutdown must never snapshot an added-but-not-started thread, + # because joining such a thread raises instead of draining it. + thread.start() + except Exception: + workers.discard(thread) + raise + return thread class CreateAgentRequest(BaseModel): @@ -59,24 +110,15 @@ class FeedbackRequest(BaseModel): reason: Optional[str] = None -_BROWSER_SUB_TOOLS = { - "browser_navigate", - "browser_click", - "browser_type", - "browser_screenshot", - "browser_extract", - "browser_axtree", -} - - def _resolve_memory_backend(config: Any) -> Any: """Instantiate the configured memory backend, or None if unavailable. - Mirrors serve.py's setup so an agent tick that runs through the server - can actually use its memory_* tools. Returns None (so the tools degrade - gracefully) when memory context is disabled or the backend can't load. + Memory storage is a tool dependency, independent of whether retrieved + context is injected into prompts. ``context_from_memory`` controls only + that prompt enrichment; explicitly configured ``memory_*`` tools still + need a backend when it is false. """ - if config is None or not getattr(config.agent, "context_from_memory", False): + if config is None: return None try: import openjarvis.tools.storage # noqa: F401 @@ -90,26 +132,141 @@ def _resolve_memory_backend(config: Any) -> Any: return None +def _memory_backend_lock(app_state: Any) -> threading.Lock: + """Return the per-runtime lock, creating it safely for lightweight tests.""" + + lock = getattr(app_state, "_memory_backend_lock", None) + if lock is not None: + return lock + with _MEMORY_BACKEND_LOCK_SETUP: + lock = getattr(app_state, "_memory_backend_lock", None) + if lock is None: + lock = threading.Lock() + app_state._memory_backend_lock = lock + return lock + + +def _get_or_create_memory_backend(app_state: Any, config: Any = None) -> Any: + """Return one runtime-owned memory backend, even under concurrent ticks.""" + + if app_state is None: + return None + backend = getattr(app_state, "memory_backend", None) + if backend is not None: + return backend + + lock = _memory_backend_lock(app_state) + with lock: + backend = getattr(app_state, "memory_backend", None) + if backend is not None: + return backend + if getattr(app_state, "_managed_runtime_stopping", False): + return None + backend = _resolve_memory_backend( + config if config is not None else getattr(app_state, "config", None) + ) + if backend is not None: + app_state.memory_backend = backend + app_state._owns_memory_backend = True + return backend + + +def _mcp_state_lock(app_state: Any, attr: str) -> threading.Lock: + """Return a named per-runtime MCP lock, creating it atomically.""" + + lock = getattr(app_state, attr, None) + if lock is not None: + return lock + with _MCP_LOCK_SETUP: + lock = getattr(app_state, attr, None) + if lock is None: + lock = threading.Lock() + setattr(app_state, attr, lock) + return lock + + +def _register_mcp_client(app_state: Any, client: Any) -> bool: + """Publish a client before blocking I/O so shutdown can interrupt it.""" + + lock = _mcp_state_lock(app_state, "_mcp_clients_lock") + with lock: + if getattr(app_state, "_managed_runtime_stopping", False): + return False + clients = getattr(app_state, "_mcp_clients", None) + if not isinstance(clients, list): + clients = list(clients or []) + app_state._mcp_clients = clients + clients.append(client) + return True + + +def _unregister_mcp_client(app_state: Any, client: Any) -> None: + """Remove a client that failed discovery or exposed no usable tools.""" + + lock = _mcp_state_lock(app_state, "_mcp_clients_lock") + with lock: + clients = getattr(app_state, "_mcp_clients", None) + if isinstance(clients, list): + try: + clients.remove(client) + except ValueError: + pass + + class _LightweightSystem: """Minimal system facade for the executor — avoids rebuilding the full JarvisSystem (which picks a random model from Ollama).""" - def __init__(self, engine: Any, model: str, config: Any = None): + def __init__( + self, + engine: Any, + model: str, + config: Any = None, + runtime: Any = None, + ): self.engine = engine self.model = model self.config = config + self._runtime = runtime # Wire the configured memory backend so an agent's memory_store / # memory_retrieve tools work when the tick runs through the server. # The executor injects system.memory_backend into those tools; this # facade previously left it None, so they reported "No memory backend # configured" even though the backend was configured and active. - self.memory_backend = _resolve_memory_backend(config) + self.memory_backend = _get_or_create_memory_backend(runtime, config) + self.channel_backend = None + self.mcp_tools: list[Any] = [] + self._mcp_clients: list[Any] = [] + self.knowledge_db_path = None + if runtime is not None: + self.channel_backend = getattr(runtime, "channel_backend", None) or getattr( + runtime, "channel_bridge", None + ) + self.knowledge_db_path = getattr(runtime, "knowledge_db_path", None) + + def get_managed_agent_mcp_tools(self) -> tuple[list[Any], list[Any]]: + """Lazily discover MCP tools only after the agent allows them.""" + + if self.mcp_tools: + return self.mcp_tools, self._mcp_clients + if self._runtime is None: + return [], [] + + runtime_tools = list(getattr(self._runtime, "mcp_tools", []) or []) + if runtime_tools: + self.mcp_tools = runtime_tools + else: + _, adapters = _get_mcp_tools(self._runtime) + self.mcp_tools = list(adapters.values()) + self._mcp_clients = list(getattr(self._runtime, "_mcp_clients", []) or []) + return self.mcp_tools, self._mcp_clients def _make_lightweight_system( engine: Any, model: str, config: Any = None, + runtime: Any = None, ) -> _LightweightSystem: """Build a minimal system with a fresh inference engine. @@ -154,10 +311,10 @@ def _make_lightweight_system( ) except Exception: pass # telemetry is optional - return _LightweightSystem(plain_engine, model, cfg) + return _LightweightSystem(plain_engine, model, cfg, runtime) except Exception: pass - return _LightweightSystem(engine, model, config) + return _LightweightSystem(engine, model, config, runtime) def _parse_param_count(model_name: str) -> float: @@ -193,73 +350,6 @@ def _pick_recommended_model( } -def _ensure_registries_populated() -> None: - """Ensure ToolRegistry and ChannelRegistry are populated. - - If the registries are empty (e.g. cleared by test fixtures) but the - modules are already cached in sys.modules, reload the individual - submodules to re-execute their @register decorators. - """ - import importlib - import sys - - from openjarvis.core.registry import ChannelRegistry, ToolRegistry - - # First, try a normal import (works if modules haven't been imported yet) - try: - import openjarvis.channels # noqa: F401 - except Exception: - pass - - try: - import openjarvis.tools # noqa: F401 - except Exception: - pass - - # Also try to import browser tools (not included in openjarvis.tools.__init__) - for _browser_mod in ("openjarvis.tools.browser", "openjarvis.tools.browser_axtree"): - try: - importlib.import_module(_browser_mod) - except Exception: - pass - - # If registries are still empty, reload individual submodules from sys.modules - if not ChannelRegistry.keys(): - for mod_name in list(sys.modules): - if mod_name.startswith("openjarvis.channels.") and not mod_name.endswith( - "_stubs" - ): - try: - importlib.reload(sys.modules[mod_name]) - except Exception: - pass - - if not ToolRegistry.keys(): - for mod_name in list(sys.modules): - if ( - mod_name.startswith("openjarvis.tools.") - and not mod_name.endswith("_stubs") - and not mod_name.endswith("agent_tools") - ): - try: - importlib.reload(sys.modules[mod_name]) - except Exception: - pass - - # After reloading tools, also try browser tools if still not registered - if not any(ToolRegistry.contains(n) for n in _BROWSER_SUB_TOOLS): - for _browser_mod in ( - "openjarvis.tools.browser", - "openjarvis.tools.browser_axtree", - ): - mod = sys.modules.get(_browser_mod) - if mod is not None: - try: - importlib.reload(mod) - except Exception: - pass - - def build_tools_list() -> List[Dict[str, Any]]: """Build unified tools list from ToolRegistry + ChannelRegistry.""" import os @@ -347,93 +437,10 @@ def build_tools_list() -> List[Dict[str, Any]]: return items -def _resolve_tool_specs( - tool_config: Any, -) -> List[Dict[str, Any]]: - """Convert a template's ``tools`` config into OpenAI-format function specs. +def _resolve_tool_specs(tool_config: Any) -> List[Dict[str, Any]]: + """Compatibility wrapper for callers of the former route-local helper.""" - The template TOML stores tools as a list of string names (e.g. - ``["file_read", "shell_exec"]``). Engines expect OpenAI-shaped dicts: - ``{"type": "function", "function": {"name, description, parameters"}}``. - - Special handling: - * Dict entries pass through as-is (allows advanced configs to - supply fully-formed specs). - * ``browser`` is a synthetic display-only meta-tool that expands - to the 6 real browser sub-tools (browser_navigate, click, …). - * Channel names (``slack``, ``gmail``, …) come from the - ``ChannelRegistry`` and are not directly callable by the LLM — - they're destinations for ``channel_send``. Silently skip them. - * Unknown tool names are dropped with a warning. - """ - if not tool_config: - return [] - - from openjarvis.core.registry import ChannelRegistry, ToolRegistry - - _ensure_registries_populated() - - def _spec_dict_for(name: str) -> Optional[Dict[str, Any]]: - try: - spec = ToolRegistry.get(name)().spec - except Exception as exc: - logger.warning( - "Could not build spec for tool '%s' (%s) — dropping", - name, - exc, - ) - return None - return { - "type": "function", - "function": { - "name": spec.name, - "description": spec.description, - "parameters": spec.parameters, - }, - } - - resolved: List[Dict[str, Any]] = [] - seen: set = set() - - for entry in tool_config: - if isinstance(entry, dict): - resolved.append(entry) - continue - if not isinstance(entry, str): - continue - - # Expand the synthetic "browser" meta-tool into its sub-tools. - if entry == "browser": - for sub in _BROWSER_SUB_TOOLS: - if sub in seen or not ToolRegistry.contains(sub): - continue - spec_dict = _spec_dict_for(sub) - if spec_dict: - resolved.append(spec_dict) - seen.add(sub) - continue - - # Channels (slack, gmail, …) live in ChannelRegistry and aren't - # callable by the LLM. Skip silently — the agent talks to them - # through the `channel_send` tool with a `channel` argument. - if ChannelRegistry.contains(entry): - continue - - if not ToolRegistry.contains(entry): - logger.warning( - "Tool '%s' referenced in agent config but not in ToolRegistry", - entry, - ) - continue - - if entry in seen: - continue - spec_dict = _spec_dict_for(entry) - if spec_dict: - resolved.append(spec_dict) - seen.add(entry) - - return resolved + return resolve_tool_specs(tool_config) # Per-agent sampler params forwarded to the engine when present in config. @@ -547,82 +554,35 @@ def _instantiate_managed_tool( model: str, app_state: Any, ) -> Any: - """Instantiate a tool with the same dependency injection as the canonical - ``cli/ask.py`` path, so ``memory_*`` / ``channel_*`` / ``llm`` tools work - instead of silently failing with "No backend configured" (#395). - """ - try: - from openjarvis.cli.ask import ( - _CHANNEL_TOOLS, - _MEMORY_TOOLS, - _get_memory_backend, - ) - except Exception: # pragma: no cover - cli import should always succeed - _MEMORY_TOOLS = frozenset() - _CHANNEL_TOOLS = frozenset() - _get_memory_backend = None + """Compatibility adapter for tests and non-managed route callers.""" - if name in _MEMORY_TOOLS: - backend = getattr(app_state, "memory_backend", None) if app_state else None - if backend is None and app_state is not None and _get_memory_backend: - cfg = getattr(app_state, "config", None) - if cfg is not None: - backend = _get_memory_backend(cfg) - if backend is None: - logger.warning( - "Memory tool %r instantiated without a backend — calls will " - "return no results.", - name, - ) - return tool_cls(backend=backend) - if name in _CHANNEL_TOOLS: - channel = getattr(app_state, "channel_bridge", None) if app_state else None - if channel is None: - logger.warning( - "Channel tool %r instantiated without a channel — calls will " - "fail with 'No channel backend configured'.", - name, - ) - return tool_cls(channel=channel) - if name == "llm": - return tool_cls(engine=engine, model=model) - return tool_cls() + memory_backend = _get_or_create_memory_backend( + app_state, + getattr(app_state, "config", None) if app_state is not None else None, + ) + channel_backend = None + if app_state is not None: + channel_backend = getattr(app_state, "channel_backend", None) or getattr( + app_state, "channel_bridge", None + ) + return instantiate_registered_tool( + tool_cls, + name, + engine=engine, + model=model, + memory_backend=memory_backend, + channel_backend=channel_backend, + ) def _build_deep_research_tools( engine: Any, model: str, knowledge_db_path: str = "", -) -> list: - """Build the 4 DeepResearch tools from a KnowledgeStore. +) -> list[Any]: + """Compatibility wrapper for existing channel and server integrations.""" - Returns an empty list if the knowledge DB does not exist. - """ - from pathlib import Path - - if not knowledge_db_path: - from openjarvis.core.config import DEFAULT_CONFIG_DIR - - knowledge_db_path = str(DEFAULT_CONFIG_DIR / "knowledge.db") - - if not Path(knowledge_db_path).exists(): - return [] - - from openjarvis.connectors.retriever import TwoStageRetriever - from openjarvis.connectors.store import KnowledgeStore - from openjarvis.tools.knowledge_search import KnowledgeSearchTool - from openjarvis.tools.knowledge_sql import KnowledgeSQLTool - from openjarvis.tools.scan_chunks import ScanChunksTool - from openjarvis.tools.think import ThinkTool - - store = KnowledgeStore(knowledge_db_path) - retriever = TwoStageRetriever(store) - return [ - KnowledgeSearchTool(retriever=retriever), - KnowledgeSQLTool(store=store), - ScanChunksTool(store=store, engine=engine, model=model), - ThinkTool(), - ] + return build_deep_research_tools(engine, model, knowledge_db_path) def _merge_tool_call_fragments( @@ -659,10 +619,44 @@ def _get_mcp_tools(app_state: Any) -> Tuple[List[Dict[str, Any]], Dict[str, Any] Lazily discovers MCP tools from config and caches them on ``app_state`` so that subsequent requests reuse the same connections. """ + lock = _mcp_state_lock(app_state, "_mcp_discovery_lock") + with lock: + return _get_mcp_tools_locked(app_state) + + +def _get_mcp_tools_locked( + app_state: Any, +) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Discover MCP tools while the runtime's discovery lock is held.""" + + if getattr(app_state, "_managed_runtime_stopping", False): + return [], {} + cached = getattr(app_state, "_mcp_tools_cache", None) if cached is not None: return cached + preloaded = list(getattr(app_state, "mcp_tools", []) or []) + if preloaded: + adapters_by_name: Dict[str, Any] = {} + for tool in preloaded: + spec = getattr(tool, "spec", None) + if spec is not None and spec.name not in adapters_by_name: + adapters_by_name[spec.name] = tool + openai_tools = [ + { + "type": "function", + "function": { + "name": tool.spec.name, + "description": tool.spec.description, + "parameters": tool.spec.parameters, + }, + } + for tool in adapters_by_name.values() + ] + app_state._mcp_tools_cache = (openai_tools, adapters_by_name) + return app_state._mcp_tools_cache + import json as _json from openjarvis.core.config import load_config @@ -683,9 +677,6 @@ def _get_mcp_tools(app_state: Any) -> Tuple[List[Dict[str, Any]], Dict[str, Any] from openjarvis.mcp.transport import StdioTransport, StreamableHTTPTransport from openjarvis.tools.mcp_adapter import MCPToolProvider - # Keep clients alive so transports persist for tool calls at runtime - mcp_clients: list = getattr(app_state, "_mcp_clients", []) - try: server_list = _json.loads(app_config.tools.mcp.servers) except (_json.JSONDecodeError, TypeError) as exc: @@ -704,6 +695,7 @@ def _get_mcp_tools(app_state: Any) -> Tuple[List[Dict[str, Any]], Dict[str, Any] command = cfg.get("command", "") args = cfg.get("args", []) + client = None try: if url: transport = StreamableHTTPTransport(url=url, token=token) @@ -717,8 +709,11 @@ def _get_mcp_tools(app_state: Any) -> Tuple[List[Dict[str, Any]], Dict[str, Any] continue client = MCPClient(transport) + if not _register_mcp_client(app_state, client): + client.close() + client = None + break client.initialize() - mcp_clients.append(client) provider = MCPToolProvider(client) discovered = provider.discover() @@ -731,8 +726,16 @@ def _get_mcp_tools(app_state: Any) -> Tuple[List[Dict[str, Any]], Dict[str, Any] if exclude_tools: discovered = [t for t in discovered if t.spec.name not in exclude_tools] + staged: list[tuple[Any, Any]] = [] + staged_names = set(adapters_by_name) for adapter in discovered: spec = adapter.spec + if spec.name in staged_names: + continue + staged.append((adapter, spec)) + staged_names.add(spec.name) + + for adapter, spec in staged: openai_tools.append( { "type": "function", @@ -745,22 +748,37 @@ def _get_mcp_tools(app_state: Any) -> Tuple[List[Dict[str, Any]], Dict[str, Any] ) adapters_by_name[spec.name] = adapter + if not staged: + client.close() + _unregister_mcp_client(app_state, client) + client = None + logger.info( "Discovered %d MCP tools from server '%s'", - len(discovered), + len(staged), name, ) except Exception as exc: + if client is not None: + try: + client.close() + except Exception: + logger.debug("Failed to close unusable MCP client", exc_info=True) + else: + _unregister_mcp_client(app_state, client) logger.warning( "Failed to discover MCP tools from '%s': %s", name, exc, ) - app_state._mcp_clients = mcp_clients if openai_tools: - app_state._mcp_tools_cache = (openai_tools, adapters_by_name) - return openai_tools, adapters_by_name + clients_lock = _mcp_state_lock(app_state, "_mcp_clients_lock") + with clients_lock: + if not getattr(app_state, "_managed_runtime_stopping", False): + app_state._mcp_tools_cache = (openai_tools, adapters_by_name) + return openai_tools, adapters_by_name + return [], {} def _sse_chunk(chunk_id: str, model: str, content: str) -> str: @@ -824,6 +842,8 @@ async def _stream_managed_agent( import json import uuid + from starlette.background import BackgroundTask + from openjarvis.core.types import Message, Role agent_id = agent_record["id"] @@ -861,16 +881,36 @@ async def _stream_managed_agent( Message(role=Role.SYSTEM, content=final_system_prompt.strip()) ) - # Resolve agent type and class for DeepResearch tool wiring + # Resolve one live toolkit for every managed-agent path. The same + # instances are advertised to the model and used for execution below. agent_type = agent_record.get("agent_type", "") - if agent_type == "deep_research": - dr_tools = _build_deep_research_tools( - engine=engine, - model=model, - ) - # Store on app_state so streaming loop can access them - if app_state is not None and dr_tools: - app_state._dr_tools = dr_tools + mcp_adapters: Dict[str, Any] = {} + mcp_clients: list[Any] = [] + if app_state is not None and config.get("mcp_tools", True) is not False: + try: + _, mcp_adapters = _get_mcp_tools(app_state) + mcp_clients = list(getattr(app_state, "_mcp_clients", [])) + except Exception as exc: + logger.warning( + "Failed to get MCP tools for managed agent: %s", + exc, + exc_info=True, + ) + + memory_backend = _get_or_create_memory_backend(app_state, app_config) + channel_backend = getattr(app_state, "channel_backend", None) or getattr( + app_state, "channel_bridge", None + ) + resolved_toolkit = resolve_agent_tools( + agent_record, + engine=engine, + model=model, + memory_backend=memory_backend, + channel_backend=channel_backend, + mcp_tools=mcp_adapters.values(), + mcp_clients=mcp_clients, + knowledge_db_path=getattr(app_state, "knowledge_db_path", None), + ) # Load prior conversation context (DESC order, reverse for chronological). # Replaying recorded tool_calls (assistant tool-use + tool results) keeps @@ -887,15 +927,14 @@ async def _stream_managed_agent( chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" # For deep_research agents: run the full agent loop, not raw streaming - if agent_type == "deep_research" and app_state is not None: - dr_tools = getattr(app_state, "_dr_tools", None) + if agent_type == "deep_research": + dr_tools = resolved_toolkit.instances if dr_tools: async def generate_deep_research(): """Run DeepResearchAgent in thread, stream progress + result.""" import asyncio import queue - import threading import time as _dr_time from openjarvis.agents.deep_research import DeepResearchAgent @@ -927,6 +966,8 @@ async def _stream_managed_agent( interactive=True, confirm_callback=lambda _prompt: True, ) + if resolved_toolkit.mcp_clients: + dr_agent._mcp_clients = resolved_toolkit.mcp_clients # Wrap the executor to capture tool calls original_execute = dr_agent._executor.execute @@ -997,6 +1038,8 @@ async def _stream_managed_agent( agent_metadata = result.metadata or {} except Exception as exc: content = f"Error: {exc}" + finally: + resolved_toolkit.close() elapsed = _dr_time.time() - _dr_start @@ -1028,8 +1071,22 @@ async def _stream_managed_agent( } ) - thread = threading.Thread(target=_run_agent, daemon=True) - thread.start() + try: + _start_managed_worker( + app_state, + _run_agent, + name=f"managed-agent-deep-research-{agent_id}", + ) + except Exception as exc: + resolved_toolkit.close() + progress_q.put( + { + "type": "error", + "content": f"Error: {exc}", + "metadata": {}, + "elapsed": _dr_time.time() - _dr_start, + } + ) # Collect tool calls from deep-research so we can persist them # alongside the final response (and the UI can re-render them @@ -1163,35 +1220,26 @@ async def _stream_managed_agent( }, ) - # Build extra kwargs for stream_full (e.g. tools from config). - # Template stores tool names as strings; convert to OpenAI function specs - # so the engine can actually bind them to the model. + # The canonical resolver exposes the same live instances as OpenAI specs + # for engines that run the generic streaming loop. stream_kwargs: Dict[str, Any] = {} - resolved_tools = _resolve_tool_specs(config.get("tools")) - if resolved_tools: - stream_kwargs["tools"] = resolved_tools + if resolved_toolkit.openai_specs: + stream_kwargs["tools"] = resolved_toolkit.openai_specs + + from openjarvis.tools._stubs import ToolExecutor + + resolved_by_name = resolved_toolkit.by_name + stream_tool_executor = ToolExecutor( + tools=resolved_toolkit.instances, + bus=bus, + interactive=True, + confirm_callback=lambda _prompt: True, + ) # Forward any per-agent sampler params (repetition_penalty, top_p, …) so # locally-hosted models can be tuned per agent (#386). stream_kwargs.update(_sampler_kwargs(config)) - # Discover MCP tools and merge into stream_kwargs - mcp_adapters: Dict[str, Any] = {} - if app_state is not None: - try: - mcp_openai_tools, mcp_adapters = _get_mcp_tools(app_state) - if mcp_openai_tools: - existing_tools = stream_kwargs.get("tools", []) - stream_kwargs["tools"] = existing_tools + mcp_openai_tools - logger.info( - "Added %d MCP tools to streaming request", - len(mcp_openai_tools), - ) - except Exception as exc: - logger.warning( - "Failed to get MCP tools for streaming: %s", exc, exc_info=True - ) - # Shared state between the generator and the BackgroundTask that # runs after the SSE response completes (or the client disconnects # mid-stream). Starlette guarantees the BackgroundTask runs in both @@ -1234,6 +1282,12 @@ async def _stream_managed_agent( except Exception as _qc_exc: logger.warning("Log query_complete failed: %s", _qc_exc) + def _finalize_stream() -> None: + try: + _persist_final() + finally: + resolved_toolkit.close() + async def generate(): """Async generator yielding SSE-formatted chunks with real token streaming.""" @@ -1369,66 +1423,21 @@ async def _stream_managed_agent( tool_start_ms = _time.monotonic() * 1000 try: - # Try MCP adapter first (external tools) - mcp_adapter = mcp_adapters.get(tool_name) - if mcp_adapter is not None: - try: - parsed_args = json.loads(tool_args) if tool_args else {} - except (json.JSONDecodeError, TypeError): - parsed_args = {} - result = mcp_adapter.execute(**parsed_args) + if tool_name in resolved_by_name: + result = stream_tool_executor.execute( + MsgToolCall( + id=tc["id"], + name=tool_name, + arguments=tool_args, + ) + ) tool_result_content = result.content + tool_succeeded = bool(result.success) else: - # Try to use ToolExecutor if tools are configured - from openjarvis.core.registry import ToolRegistry - from openjarvis.tools._stubs import ( - ToolCall as StubToolCall, + logger.warning( + "Tool '%s' was not included in the resolved toolkit", + tool_name, ) - from openjarvis.tools._stubs import ( - ToolExecutor, - ) - - tool_cls = ToolRegistry.get(tool_name) - if tool_cls is not None: - # Inject backend / channel / engine the same - # way cli/ask.py does, else memory_* / channel_* - # / llm tools fail with "No backend configured" - # on every call (#395). - tool_instance = _instantiate_managed_tool( - tool_cls, - tool_name, - engine=engine, - model=model, - app_state=app_state, - ) - # Tools the user explicitly added to this - # agent's toolkit are considered pre-approved — - # selecting them in the wizard is the - # confirmation. Without this, tools that have - # `requires_confirmation=True` (shell_exec, - # apply_patch) would fail with "requires - # confirmation but no callback available" on - # every call. - executor = ToolExecutor( - tools=[tool_instance], - bus=bus, - interactive=True, - confirm_callback=lambda _prompt: True, - ) - result = executor.execute( - StubToolCall( - id=tc["id"], - name=tool_name, - arguments=tool_args, - ), - ) - tool_result_content = result.content - else: - logger.warning( - "Tool '%s' not found in registry or MCP adapters", - tool_name, - ) - tool_succeeded = True except Exception as tool_exc: logger.error( "Tool execution error for %s: %s", @@ -1516,13 +1525,11 @@ async def _stream_managed_agent( yield f"data: {json.dumps(final_data)}\n\n" yield "data: [DONE]\n\n" - from starlette.background import BackgroundTask - return StreamingResponse( generate(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, - background=BackgroundTask(_persist_final), + background=BackgroundTask(_finalize_stream), ) @@ -1605,8 +1612,6 @@ def create_agent_manager_router( @agents_router.post("/{agent_id}/run") async def run_agent(agent_id: str, request: Request): - import threading - agent = manager.get_agent(agent_id) if not agent: raise HTTPException(status_code=404, detail="Agent not found") @@ -1644,6 +1649,7 @@ def create_agent_manager_router( server_engine, server_model, server_config, + request.app.state, ) executor.set_system(system) # The route handler above already called start_tick() to @@ -1668,7 +1674,18 @@ def create_agent_manager_router( f"ERROR: {exc}", ) - threading.Thread(target=_run_tick, daemon=True).start() + try: + _start_managed_worker( + request.app.state, + _run_tick, + name=f"managed-agent-run-{agent_id}", + ) + except RuntimeError as exc: + try: + manager.end_tick(agent_id) + except Exception: + pass + raise HTTPException(status_code=503, detail=str(exc)) return {"status": "running", "agent_id": agent_id} # ── Recover ────────────────────────────────────────────── @@ -1753,10 +1770,6 @@ def create_agent_manager_router( engine = getattr(request.app.state, "engine", None) if engine: - from openjarvis.server.agent_manager_routes import ( - _build_deep_research_tools, - ) - tools = _build_deep_research_tools(engine=engine, model="") if tools: from openjarvis.agents.deep_research import ( @@ -1825,11 +1838,7 @@ def create_agent_manager_router( engine = getattr(request.app.state, "engine", None) dr_agent = None if engine: - from openjarvis.server.agent_manager_routes import ( - _build_deep_research_tools as _bdr, - ) - - tools = _bdr(engine=engine, model="") + tools = _build_deep_research_tools(engine=engine, model="") if tools: from openjarvis.agents.deep_research import ( DeepResearchAgent, @@ -1971,7 +1980,6 @@ def create_agent_manager_router( # Non-streaming immediate: trigger a background tick so the # agent processes the message, then return the stored msg. # Re-use the server's existing system (correct model/engine). - import threading import time as _time from openjarvis.agents.executor import AgentExecutor @@ -1999,6 +2007,7 @@ def create_agent_manager_router( _srv_engine, _srv_model, _srv_config, + request.app.state, ) executor.set_system(system) logger.info( @@ -2030,10 +2039,14 @@ def create_agent_manager_router( f"ERROR: {exc}", ) - threading.Thread( - target=_immediate_tick, - daemon=True, - ).start() + try: + _start_managed_worker( + request.app.state, + _immediate_tick, + name=f"managed-agent-immediate-{agent_id}", + ) + except RuntimeError as exc: + raise HTTPException(status_code=503, detail=str(exc)) return msg # --- Streaming mode: run agent and return SSE response --- diff --git a/src/openjarvis/server/app.py b/src/openjarvis/server/app.py index 35d42844..55c2dbd5 100644 --- a/src/openjarvis/server/app.py +++ b/src/openjarvis/server/app.py @@ -4,6 +4,7 @@ from __future__ import annotations import logging import pathlib +import threading import time from fastapi import FastAPI @@ -21,6 +22,8 @@ from openjarvis.server.routes import router from openjarvis.server.upload_router import router as upload_router logger = logging.getLogger(__name__) +_MANAGED_SHUTDOWN_GRACE_SECONDS = 0.25 +_MANAGED_SHUTDOWN_DRAIN_SECONDS = 10.0 def _restore_sendblue_bindings(app: FastAPI) -> None: @@ -151,10 +154,13 @@ def create_app( channel_bridge=None, config=None, memory_backend=None, + own_memory_backend: bool = False, memory_service=None, speech_backend=None, agent_manager=None, agent_scheduler=None, + mcp_tools=None, + mcp_clients=None, api_key: str = "", webhook_config: dict | None = None, cors_origins: list[str] | None = None, @@ -221,16 +227,129 @@ def create_app( ) app.state.channel_bridge = channel_bridge app.state.config = config + app.state._memory_backend_lock = threading.Lock() app.state.memory_backend = memory_backend + app.state._owns_memory_backend = bool(own_memory_backend) app.state.memory_service = memory_service app.state.speech_backend = speech_backend app.state.agent_manager = agent_manager app.state.agent_scheduler = agent_scheduler + app.state.mcp_tools = list(mcp_tools or []) + app.state._mcp_discovery_lock = threading.Lock() + app.state._mcp_clients_lock = threading.Lock() + app.state._mcp_clients = list(mcp_clients or []) + app.state._managed_worker_lock = threading.Lock() + app.state._managed_workers: set[threading.Thread] = set() + app.state._managed_runtime_stopping = False app.state.session_start = time.time() # Exposed so WebSocket handlers can authenticate the handshake (the HTTP # AuthMiddleware never sees WS upgrade requests). Empty = auth disabled. app.state.api_key = api_key + @app.on_event("shutdown") + async def _shutdown_managed_runtime() -> None: + # Quiesce every producer before touching the shared MCP pool. Route + # workers are registered under this lock, so none can slip in after + # the snapshot. The scheduler has a two-phase stop because closing an + # MCP transport may be what releases an in-flight tick. + with app.state._managed_worker_lock: + app.state._managed_runtime_stopping = True + managed_workers = list(app.state._managed_workers) + + # Stop external listener threads before draining ticks or closing the + # shared MCP pool. Channel callbacks are wired to that same pool by + # ``serve`` and otherwise could race teardown or survive app restart. + channel_bridge = getattr(app.state, "channel_bridge", None) + disconnect_channels = getattr(channel_bridge, "disconnect", None) + if callable(disconnect_channels): + try: + disconnect_channels() + except Exception: + logger.debug("Channel bridge shutdown failed", exc_info=True) + + def _join_workers(timeout: float) -> None: + deadline = time.monotonic() + timeout + for thread in managed_workers: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + thread.join(timeout=remaining) + + scheduler = getattr(app.state, "agent_scheduler", None) + scheduler_wait = None + scheduler_drained = True + if scheduler is not None: + try: + request_stop = getattr(scheduler, "request_stop", None) + wait_stopped = getattr(scheduler, "wait_stopped", None) + if callable(request_stop) and callable(wait_stopped): + request_stop() + scheduler_wait = wait_stopped + scheduler_drained = bool( + wait_stopped(timeout=_MANAGED_SHUTDOWN_GRACE_SECONDS) + ) + else: + scheduler.stop() + scheduler_drained = not bool( + getattr(scheduler, "is_running", False) + ) + except Exception: + scheduler_drained = False + logger.debug("Agent scheduler shutdown failed", exc_info=True) + + # Give normal work a brief chance to finish before cancellation. + _join_workers(timeout=_MANAGED_SHUTDOWN_GRACE_SECONDS) + with app.state._mcp_clients_lock: + mcp_clients_to_close = list(app.state._mcp_clients) + for client in mcp_clients_to_close: + try: + client.close() + except Exception: + logger.debug("MCP client shutdown failed", exc_info=True) + + # Transport closure interrupts blocked MCP reads. Drain the workers a + # second time so shutdown does not return while they still own runtime + # state. Any stragglers can no longer issue transport requests because + # MCPClient marks itself closed before closing its transport. + if scheduler_wait is not None: + try: + scheduler_drained = bool( + scheduler_wait(timeout=_MANAGED_SHUTDOWN_DRAIN_SECONDS) + ) + except Exception: + scheduler_drained = False + logger.debug("Agent scheduler drain failed", exc_info=True) + _join_workers(timeout=_MANAGED_SHUTDOWN_DRAIN_SECONDS) + alive = [thread.name for thread in managed_workers if thread.is_alive()] + if alive: + logger.warning("Managed workers did not stop during shutdown: %s", alive) + + # A backend created by ``serve`` or lazily by a managed route belongs + # to this app process. Close it only after every tracked consumer has + # been drained; injected/borrowed backends remain the caller's concern. + owned_memory_backend = None + runtime_drained = scheduler_drained and not alive + if runtime_drained: + with app.state._memory_backend_lock: + if app.state._owns_memory_backend: + owned_memory_backend = app.state.memory_backend + app.state.memory_backend = None + app.state._owns_memory_backend = False + else: + # A live worker may itself hold _memory_backend_lock while opening + # the backend. Respect the bounded shutdown deadline: do not wait + # on that lock or mutate ownership until every consumer is gone. + logger.warning( + "Skipping memory backend cleanup because managed runtime " + "consumers did not stop" + ) + close_memory = getattr(owned_memory_backend, "close", None) + if callable(close_memory): + try: + close_memory() + except Exception: + logger.debug("Memory backend shutdown failed", exc_info=True) + # Wire up trace store if traces are enabled. # # We deliberately do NOT subscribe the trace store to the bus. The chat diff --git a/src/openjarvis/system/builder.py b/src/openjarvis/system/builder.py index 743bb7b1..e0fde9c6 100644 --- a/src/openjarvis/system/builder.py +++ b/src/openjarvis/system/builder.py @@ -48,6 +48,7 @@ class SystemBuilder: self._sessions: Optional[bool] = None self._speech: Optional[bool] = None self._mcp_clients: List = [] + self._mcp_tools: List[BaseTool] = [] def engine(self, key: str) -> SystemBuilder: self._engine_key = key @@ -113,6 +114,33 @@ class SystemBuilder: def build(self) -> JarvisSystem: """Construct a fully wired JarvisSystem.""" + # Discovery state belongs to one build only. Once a system is + # returned, that system owns the clients and adapters captured below; + # retaining them here would make a reused builder hand closed clients + # from an earlier system to the next one. + self._clear_mcp_discovery_state(close_clients=True) + try: + system = self._build() + except BaseException: + # No system took ownership, so release any clients opened before + # the build failed. + self._clear_mcp_discovery_state(close_clients=True) + raise + self._clear_mcp_discovery_state(close_clients=False) + return system + + def _clear_mcp_discovery_state(self, *, close_clients: bool) -> None: + if close_clients: + for client in getattr(self, "_mcp_clients", []): + try: + client.close() + except Exception: + logger.debug("Error closing unowned MCP client", exc_info=True) + self._mcp_clients = [] + self._mcp_tools = [] + + def _build(self) -> JarvisSystem: + """Build one system using fresh, build-local MCP discovery state.""" config = self._config bus = self._bus or get_event_bus() @@ -291,6 +319,7 @@ class SystemBuilder: model=model, agent_name=agent_name, tools=tool_list, + mcp_tools=list(self._mcp_tools), tool_executor=tool_executor, memory_backend=memory_backend, channel_backend=channel_backend, @@ -440,7 +469,7 @@ class SystemBuilder: else: tools = [] - if config.tools.mcp.servers: + if config.tools.mcp.enabled and config.tools.mcp.servers: try: import json @@ -449,6 +478,7 @@ class SystemBuilder: for server_cfg in server_list: try: external_tools = self._discover_external_mcp(server_cfg) + self._mcp_tools.extend(external_tools) if tool_names: external_tools = [ t diff --git a/src/openjarvis/system/core.py b/src/openjarvis/system/core.py index 95f48179..e0d803f1 100644 --- a/src/openjarvis/system/core.py +++ b/src/openjarvis/system/core.py @@ -86,6 +86,9 @@ class JarvisSystem: skill_manager: Optional[SkillManager] = None _learning_orchestrator: Optional[LearningOrchestrator] = None _mcp_clients: List[MCPClient] = field(default_factory=list) + # Keep newly added fields after every pre-existing positional field so + # older positional JarvisSystem(...) calls retain their original meaning. + mcp_tools: List[BaseTool] = field(default_factory=list) @property def security(self) -> SecurityContext: diff --git a/tests/agents/test_executor_tool_persistence.py b/tests/agents/test_executor_tool_persistence.py new file mode 100644 index 00000000..18c325c8 --- /dev/null +++ b/tests/agents/test_executor_tool_persistence.py @@ -0,0 +1,127 @@ +"""Regression tests for managed-agent tool-call persistence.""" + +from __future__ import annotations + +import json + +import pytest + +from openjarvis.agents._stubs import AgentResult +from openjarvis.agents.executor import AgentExecutor, _tool_calls_for_storage +from openjarvis.agents.manager import AgentManager +from openjarvis.core.events import EventBus +from openjarvis.core.types import ToolResult + + +def test_tool_results_are_serialized_for_managed_messages() -> None: + result = AgentResult( + content="Finished", + tool_results=[ + ToolResult( + tool_name="knowledge_search", + content="Found the requested note", + success=True, + latency_seconds=0.42, + metadata={ + "arguments": { + "query": "financial independence", + "limit": 3, + } + }, + ), + ToolResult( + tool_name="shell_exec", + content="Permission denied", + success=False, + latency_seconds=1.25, + metadata={"arguments": '{"command":"whoami"}'}, + ), + ], + ) + + calls = _tool_calls_for_storage(result) + + assert calls is not None + assert len(calls) == 2 + + knowledge_call = calls[0] + assert knowledge_call["tool"] == "knowledge_search" + assert isinstance(knowledge_call["arguments"], str) + assert json.loads(knowledge_call["arguments"]) == { + "query": "financial independence", + "limit": 3, + } + assert knowledge_call["result"] == "Found the requested note" + assert knowledge_call["success"] is True + assert knowledge_call["latency"] == pytest.approx(420.0) + + failed_call = calls[1] + assert failed_call["arguments"] == '{"command":"whoami"}' + assert failed_call["result"] == "Permission denied" + assert failed_call["success"] is False + assert failed_call["latency"] == pytest.approx(1250.0) + + +def test_no_tool_results_serialize_as_none() -> None: + assert _tool_calls_for_storage(AgentResult(content="Plain response")) is None + + +def test_finalize_tick_persists_tool_calls_round_trip(tmp_path) -> None: + manager = AgentManager(str(tmp_path / "agents.db")) + try: + agent = manager.create_agent("researcher") + manager.start_tick(agent["id"]) + result = AgentResult( + content="Answer grounded in the knowledge base", + tool_results=[ + ToolResult( + tool_name="knowledge_search", + content="Matching source text", + success=True, + latency_seconds=0.007, + metadata={"arguments": {"query": "grounded answer"}}, + ) + ], + ) + + executor = AgentExecutor(manager, EventBus()) + executor._finalize_tick( + agent["id"], + result, + error=None, + duration=0.01, + ) + + messages = manager.list_messages(agent["id"]) + assert len(messages) == 1 + stored = messages[0] + assert stored["content"] == result.content + assert stored["direction"] == "agent_to_user" + assert stored["tool_calls"] == _tool_calls_for_storage(result) + assert isinstance(stored["tool_calls"][0]["arguments"], str) + assert json.loads(stored["tool_calls"][0]["arguments"]) == { + "query": "grounded answer" + } + assert stored["tool_calls"][0]["latency"] == pytest.approx(7.0) + finally: + manager.close() + + +def test_finalize_tick_without_tools_stores_null_tool_calls(tmp_path) -> None: + manager = AgentManager(str(tmp_path / "agents.db")) + try: + agent = manager.create_agent("plain-agent") + manager.start_tick(agent["id"]) + + executor = AgentExecutor(manager, EventBus()) + executor._finalize_tick( + agent["id"], + AgentResult(content="No tools needed"), + error=None, + duration=0.01, + ) + + stored = manager.list_messages(agent["id"])[0] + assert stored["tool_calls"] is None + finally: + manager.close() diff --git a/tests/agents/test_executor_tools.py b/tests/agents/test_executor_tools.py index 70e26132..ad94ddcb 100644 --- a/tests/agents/test_executor_tools.py +++ b/tests/agents/test_executor_tools.py @@ -2,13 +2,91 @@ from __future__ import annotations +import gc +import sqlite3 +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from openjarvis.agents._stubs import AgentResult from openjarvis.agents.executor import AgentExecutor from openjarvis.agents.manager import AgentManager +from openjarvis.agents.tool_resolver import ResolvedAgentTools +from openjarvis.connectors.store import KnowledgeStore +from openjarvis.core.config import MemoryFilesConfig, SystemPromptConfig from openjarvis.core.events import EventBus +from openjarvis.core.registry import AgentRegistry, ToolRegistry +from openjarvis.core.types import Role, ToolResult +from openjarvis.tools._stubs import BaseTool, ToolSpec from tests.agents.fake_engine import FakeEngine from tests.agents.scenario_harness import FakeSystem +class _CapturingToolAgent: + """Minimal agent that exposes the toolkit received by AgentExecutor.""" + + accepts_tools = True + captured_tools = [] + captured_search_result = None + + def __init__(self, engine, model, *, tools=None, **kwargs): + self.engine = engine + self.model = model + type(self).captured_tools = list(tools or []) + + def run(self, input_text, context=None): + tools_by_name = {tool.spec.name: tool for tool in self.captured_tools} + search = tools_by_name.get("knowledge_search") + if search is not None: + type(self).captured_search_result = search.execute( + query="EXECUTOR_RESOLVER_SENTINEL" + ) + return AgentResult(content="captured") + + +class _NonToolAgent: + """Agent class whose run method must not swallow a configured toolkit.""" + + accepts_tools = False + supports_managed_tool_fallback = True + runs = 0 + + def __init__(self, engine, model, **kwargs): + pass + + def run(self, input_text, context=None): + type(self).runs += 1 + raise AssertionError("non-tool agent should use the managed tool loop") + + +class _SpecializedNonToolAgent: + """Non-tool agent that must retain its specialized execution path.""" + + accepts_tools = False + runs = 0 + + def __init__(self, engine, model, **kwargs): + pass + + def run(self, input_text, context=None): + type(self).runs += 1 + return AgentResult(content="specialized response") + + +class _ExecutorProbeTool(BaseTool): + tool_id = "executor_probe" + calls = 0 + + @property + def spec(self) -> ToolSpec: + return ToolSpec(name=self.tool_id, description="Executor parity probe") + + def execute(self, **params) -> ToolResult: + type(self).calls += 1 + return ToolResult(tool_name=self.tool_id, content="probe-result") + + def _register_agent(): """Re-register MonitorOperativeAgent (cleared by autouse fixture).""" from openjarvis.agents.monitor_operative import MonitorOperativeAgent @@ -102,3 +180,396 @@ def test_executor_handles_string_tools(tmp_path): result_agent = mgr.get_agent(agent["id"]) assert result_agent["status"] == "idle" mgr.close() + + +def test_executor_uses_tool_loop_for_non_tool_agent_with_configured_tools(tmp_path): + """Immediate/scheduled ticks match SSE instead of discarding tools.""" + + AgentRegistry.register_value("non_tool_probe", _NonToolAgent) + ToolRegistry.register_value(_ExecutorProbeTool.tool_id, _ExecutorProbeTool) + _NonToolAgent.runs = 0 + _ExecutorProbeTool.calls = 0 + + engine = FakeEngine( + [ + { + "tool_calls": [ + { + "id": "call-executor-probe", + "name": _ExecutorProbeTool.tool_id, + "arguments": "{}", + } + ] + }, + {"content": "tool-backed final response"}, + ] + ) + system = FakeSystem(engine=engine) + system.config = SimpleNamespace( + agent=SimpleNamespace(default_system_prompt="GLOBAL_DEFAULT"), + memory_files=MemoryFilesConfig(persona_name="none"), + system_prompt=SystemPromptConfig(), + ) + manager = AgentManager(db_path=str(tmp_path / "agents.db")) + agent = manager.create_agent( + "non-tool with tools", + agent_type="non_tool_probe", + config={ + "model": "test-model", + "tools": [_ExecutorProbeTool.tool_id], + "instruction": "Use the probe.", + "system_prompt": "NON_TOOL_SYSTEM_SENTINEL", + }, + ) + + try: + AgentExecutor(manager, EventBus(), system=system).execute_tick(agent["id"]) + + assert _NonToolAgent.runs == 0 + assert _ExecutorProbeTool.calls == 1 + assert engine.call_count == 2 + assert any( + message.role is Role.SYSTEM + and message.content == "NON_TOOL_SYSTEM_SENTINEL" + for message in engine.last_messages or [] + ) + refreshed = manager.get_agent(agent["id"]) + assert refreshed["status"] == "idle" + assert refreshed["total_runs"] == 1 + responses = [ + message + for message in manager.list_messages(agent["id"]) + if message["direction"] == "agent_to_user" + ] + assert responses[-1]["content"] == "tool-backed final response" + assert responses[-1]["tool_calls"][0]["tool"] == "executor_probe" + finally: + manager.close() + + +def test_simple_agent_uses_global_mcp_tools_without_native_tool_config(tmp_path): + """Fallback-compatible simple agents preserve SSE/global-MCP parity.""" + + from openjarvis.agents.simple import SimpleAgent + + AgentRegistry.register_value("simple", SimpleAgent) + _ExecutorProbeTool.calls = 0 + provider = MagicMock(return_value=([_ExecutorProbeTool()], [])) + engine = FakeEngine( + [ + { + "tool_calls": [ + { + "id": "call-global-mcp-probe", + "name": _ExecutorProbeTool.tool_id, + "arguments": "{}", + } + ] + }, + {"content": "global MCP response"}, + ] + ) + system = SimpleNamespace( + engine=engine, + model="test-model", + config=None, + memory_backend=None, + channel_backend=None, + session_store=None, + knowledge_db_path=None, + get_managed_agent_mcp_tools=provider, + ) + manager = AgentManager(db_path=str(tmp_path / "agents.db")) + agent = manager.create_agent( + "simple global MCP", + agent_type="simple", + config={"model": "test-model", "instruction": "Use MCP."}, + ) + + try: + AgentExecutor(manager, EventBus(), system=system).execute_tick(agent["id"]) + + provider.assert_called_once_with() + assert _ExecutorProbeTool.calls == 1 + assert engine.call_count == 2 + responses = [ + message + for message in manager.list_messages(agent["id"]) + if message["direction"] == "agent_to_user" + ] + assert responses[-1]["content"] == "global MCP response" + finally: + manager.close() + + +def test_simple_agent_without_tools_keeps_its_custom_system_prompt(tmp_path): + """Signature filtering must not discard prompt-builder state on retry.""" + + from openjarvis.agents.simple import SimpleAgent + + AgentRegistry.register_value("simple", SimpleAgent) + engine = FakeEngine([{"content": "custom prompt response"}]) + system = FakeSystem(engine=engine) + system.config = SimpleNamespace( + agent=SimpleNamespace(default_system_prompt="GLOBAL_DEFAULT"), + memory_files=MemoryFilesConfig(persona_name="none"), + system_prompt=SystemPromptConfig(), + ) + manager = AgentManager(db_path=str(tmp_path / "agents.db")) + agent = manager.create_agent( + "simple custom prompt", + agent_type="simple", + config={ + "model": "test-model", + "instruction": "Answer directly.", + "system_prompt": "SIMPLE_CUSTOM_SYSTEM_SENTINEL", + "mcp_tools": False, + }, + ) + + try: + AgentExecutor(manager, EventBus(), system=system).execute_tick(agent["id"]) + + assert engine.call_count == 1 + assert any( + message.role is Role.SYSTEM + and message.content == "SIMPLE_CUSTOM_SYSTEM_SENTINEL" + for message in engine.last_messages or [] + ) + finally: + manager.close() + + +def test_specialized_non_tool_agent_is_not_replaced_by_generic_tool_loop(tmp_path): + """Configured/global tools never replace a non-opted-in agent class.""" + + AgentRegistry.register_value("specialized_non_tool", _SpecializedNonToolAgent) + ToolRegistry.register_value(_ExecutorProbeTool.tool_id, _ExecutorProbeTool) + _SpecializedNonToolAgent.runs = 0 + _ExecutorProbeTool.calls = 0 + provider = MagicMock(return_value=([_ExecutorProbeTool()], [])) + system = SimpleNamespace( + engine=FakeEngine([{"content": "unused"}]), + model="test-model", + config=None, + memory_backend=None, + channel_backend=None, + session_store=None, + knowledge_db_path=None, + get_managed_agent_mcp_tools=provider, + ) + manager = AgentManager(db_path=str(tmp_path / "agents.db")) + agent = manager.create_agent( + "specialized with configured tool", + agent_type="specialized_non_tool", + config={ + "model": "test-model", + "instruction": "Keep the specialized path.", + "tools": [_ExecutorProbeTool.tool_id], + }, + ) + + try: + AgentExecutor(manager, EventBus(), system=system).execute_tick(agent["id"]) + + provider.assert_not_called() + assert _SpecializedNonToolAgent.runs == 1 + assert _ExecutorProbeTool.calls == 0 + responses = [ + message + for message in manager.list_messages(agent["id"]) + if message["direction"] == "agent_to_user" + ] + assert responses[-1]["content"] == "specialized response" + finally: + manager.close() + + +def test_executor_grants_deep_research_live_knowledge_tools(tmp_path): + """Immediate ticks receive the same live Deep Research grant as SSE.""" + + AgentRegistry.register_value("deep_research", _CapturingToolAgent) + _CapturingToolAgent.captured_tools = [] + _CapturingToolAgent.captured_search_result = None + + knowledge_db_path = tmp_path / "knowledge.db" + with KnowledgeStore(db_path=knowledge_db_path) as store: + store.store( + "The EXECUTOR_RESOLVER_SENTINEL decision was approved.", + source="test", + doc_type="note", + ) + + manager = AgentManager(db_path=str(tmp_path / "agents.db")) + agent = manager.create_agent( + "researcher", + agent_type="deep_research", + config={ + "model": "agent-selected-model", + # These duplicate two agent-type grants and must not replace them. + "tools": ["knowledge_search", "think"], + "instruction": "Find the sentinel.", + }, + ) + manager.send_message(agent["id"], "Search the knowledge base.", mode="immediate") + + system = SimpleNamespace( + engine=FakeEngine([{"content": "unused"}]), + model="system-model", + memory_backend=None, + channel_backend=None, + tool_executor=None, + _mcp_clients=[], + knowledge_db_path=knowledge_db_path, + config=None, + session_store=None, + ) + executor = AgentExecutor(manager=manager, event_bus=EventBus(), system=system) + + try: + executor.execute_tick(agent["id"]) + + tools_by_name = { + tool.spec.name: tool for tool in _CapturingToolAgent.captured_tools + } + assert set(tools_by_name) == { + "knowledge_search", + "knowledge_sql", + "scan_chunks", + "think", + } + result = _CapturingToolAgent.captured_search_result + assert result is not None + assert result.success is True + assert "EXECUTOR_RESOLVER_SENTINEL" in result.content + assert tools_by_name["scan_chunks"]._model == "agent-selected-model" + assert manager.get_agent(agent["id"])["status"] == "idle" + with pytest.raises(sqlite3.ProgrammingError): + tools_by_name["knowledge_sql"]._store._conn.execute("SELECT 1") + finally: + manager.close() + + +def test_executor_mcp_opt_out_does_not_call_lazy_provider(tmp_path): + """An opted-out tick must not trigger request-local MCP discovery.""" + + AgentRegistry.register_value("capturing", _CapturingToolAgent) + provider = MagicMock(side_effect=AssertionError("MCP discovery must stay lazy")) + system = SimpleNamespace( + engine=FakeEngine([{"content": "unused"}]), + model="system-model", + memory_backend=None, + channel_backend=None, + tool_executor=None, + _mcp_clients=[], + config=None, + session_store=None, + get_managed_agent_mcp_tools=provider, + ) + manager = AgentManager(db_path=str(tmp_path / "agents.db")) + agent = manager.create_agent( + "no-mcp", + agent_type="capturing", + config={"model": "test-model", "mcp_tools": False}, + ) + + try: + AgentExecutor(manager, EventBus(), system=system).execute_tick(agent["id"]) + provider.assert_not_called() + assert manager.get_agent(agent["id"])["status"] == "idle" + finally: + manager.close() + + +def test_executor_preserves_custom_dict_tool_schema(tmp_path): + """Executor-based agents see the same custom schema advertised by SSE.""" + + from openjarvis.tools.think import ThinkTool + + AgentRegistry.register_value("capturing", _CapturingToolAgent) + ToolRegistry.register_value("think", ThinkTool) + _CapturingToolAgent.captured_tools = [] + custom_spec = { + "type": "function", + "function": { + "name": "think", + "description": "Agent-specific thinking schema", + "parameters": { + "type": "object", + "properties": {"thought": {"type": "string"}}, + "required": ["thought"], + }, + }, + } + system = SimpleNamespace( + engine=FakeEngine([{"content": "unused"}]), + model="test-model", + memory_backend=None, + channel_backend=None, + tool_executor=None, + mcp_tools=[], + _mcp_clients=[], + config=None, + session_store=None, + ) + manager = AgentManager(db_path=str(tmp_path / "agents.db")) + agent = manager.create_agent( + "custom-schema", + agent_type="capturing", + config={"model": "test-model", "tools": [custom_spec]}, + ) + + try: + AgentExecutor(manager, EventBus(), system=system).execute_tick(agent["id"]) + assert len(_CapturingToolAgent.captured_tools) == 1 + configured_tool = _CapturingToolAgent.captured_tools[0] + assert configured_tool.to_openai_function() == custom_spec + assert configured_tool.spec.description == "Agent-specific thinking schema" + assert configured_tool.execute(thought="same instance").success is True + finally: + manager.close() + + +def test_executor_closes_resolver_resources_when_pre_run_setup_fails( + tmp_path, + monkeypatch, +): + """The resolver finalizer covers failures before agent.run is reached.""" + + AgentRegistry.register_value("capturing", _CapturingToolAgent) + resource = MagicMock() + + def _resolve(*args, **kwargs): + return ResolvedAgentTools(owned_resources=[resource]) + + monkeypatch.setattr("openjarvis.agents.executor.resolve_agent_tools", _resolve) + system = SimpleNamespace( + engine=FakeEngine([{"content": "unused"}]), + model="test-model", + memory_backend=None, + channel_backend=None, + tool_executor=None, + mcp_tools=[], + _mcp_clients=[], + config=None, + session_store=None, + ) + manager = AgentManager(db_path=str(tmp_path / "agents.db")) + agent = manager.create_agent( + "cleanup", + agent_type="capturing", + config={"model": "test-model"}, + ) + monkeypatch.setattr( + manager, + "get_pending_messages", + MagicMock(side_effect=RuntimeError("pre-run setup failed")), + ) + + try: + with pytest.raises(RuntimeError, match="pre-run setup failed"): + AgentExecutor(manager, EventBus(), system=system)._invoke_agent(agent) + gc.collect() + resource.close.assert_called_once_with() + finally: + manager.close() diff --git a/tests/agents/test_scheduler.py b/tests/agents/test_scheduler.py index b4ac38ec..7a7b41da 100644 --- a/tests/agents/test_scheduler.py +++ b/tests/agents/test_scheduler.py @@ -3,6 +3,7 @@ from __future__ import annotations import tempfile +import threading import time from pathlib import Path from unittest.mock import MagicMock @@ -117,6 +118,48 @@ class TestSchedulerBasic: assert executor.execute_tick.call_count >= 1 executor.execute_tick.assert_called_with(agent["id"]) + def test_two_phase_stop_retains_and_drains_active_worker(self, manager): + """Shutdown quiesces later ticks and can wait again after cancellation.""" + + from openjarvis.agents.scheduler import AgentScheduler + + started = threading.Event() + release = threading.Event() + calls: list[str] = [] + + class _BlockingExecutor: + def execute_tick(self, agent_id): + calls.append(agent_id) + started.set() + release.wait(timeout=2) + + scheduler = AgentScheduler( + manager=manager, + executor=_BlockingExecutor(), + tick_interval=0.01, + ) + agents = [ + manager.create_agent( + name=f"test-{index}", + agent_type="monitor_operative", + config={"schedule_type": "interval", "schedule_value": 0}, + ) + for index in range(2) + ] + for agent in agents: + scheduler.register_agent(agent["id"]) + + scheduler.start() + assert started.wait(timeout=1) + scheduler.request_stop() + assert scheduler.wait_stopped(timeout=0.01) is False + assert scheduler._thread is not None + + release.set() + assert scheduler.wait_stopped(timeout=1) is True + assert scheduler._thread is None + assert calls == [agents[0]["id"]] + def test_skips_paused_agents(self, manager): from openjarvis.agents.scheduler import AgentScheduler diff --git a/tests/agents/test_tool_resolver.py b/tests/agents/test_tool_resolver.py new file mode 100644 index 00000000..0d640400 --- /dev/null +++ b/tests/agents/test_tool_resolver.py @@ -0,0 +1,284 @@ +"""Focused tests for canonical managed-agent tool resolution (#688).""" + +from __future__ import annotations + +from collections import Counter + +import pytest + +from openjarvis.agents import tool_resolver +from openjarvis.connectors.store import KnowledgeStore +from openjarvis.core.registry import ToolRegistry +from openjarvis.core.types import ToolResult +from openjarvis.tools import description_loader +from openjarvis.tools._stubs import BaseTool, ToolSpec + + +class _AlphaTool(BaseTool): + tool_id = "alpha" + + @property + def spec(self) -> ToolSpec: + return ToolSpec(name="alpha", description="Alpha test tool") + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name="alpha", content="alpha", success=True) + + +class _BetaTool(BaseTool): + tool_id = "beta" + + @property + def spec(self) -> ToolSpec: + return ToolSpec(name="beta", description="Beta test tool") + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name="beta", content="beta", success=True) + + +class _NativeSharedTool(BaseTool): + tool_id = "shared" + + @property + def spec(self) -> ToolSpec: + return ToolSpec(name="shared", description="Native shared tool") + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name="shared", content="native", success=True) + + +class _MCPSharedTool(BaseTool): + tool_id = "shared" + + @property + def spec(self) -> ToolSpec: + return ToolSpec(name="shared", description="MCP name collision") + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name="shared", content="mcp", success=True) + + +class _MCPOnlyTool(BaseTool): + tool_id = "mcp_only" + + @property + def spec(self) -> ToolSpec: + return ToolSpec(name="mcp_only", description="MCP-only test tool") + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name="mcp_only", content="mcp-only", success=True) + + +@pytest.fixture(autouse=True) +def _use_explicit_test_registrations(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep these unit tests independent of import-time registry population.""" + + monkeypatch.setattr(tool_resolver, "ensure_registries_populated", lambda: None) + + +def test_deep_research_grants_are_live_deduplicated_and_use_selected_model( + tmp_path, +) -> None: + """Agent-type grants must beat duplicate bare configured tools.""" + + db_path = tmp_path / "knowledge.db" + with KnowledgeStore(db_path=db_path) as store: + store.store( + "The RESOLVER_SENTINEL decision was approved.", + source="test", + doc_type="note", + ) + + engine = object() + resolved = tool_resolver.resolve_agent_tools( + { + "agent_type": "deep_research", + "config": { + # Both names are already supplied by the agent-type grant. + "tools": ["knowledge_search", "think", "think"], + }, + }, + engine=engine, + model="agent-selected-model", + knowledge_db_path=db_path, + ) + + try: + names = [tool.spec.name for tool in resolved.instances] + assert set(names) == { + "knowledge_search", + "knowledge_sql", + "scan_chunks", + "think", + } + assert all(count == 1 for count in Counter(names).values()) + + search = resolved.by_name["knowledge_search"] + result = search.execute(query="RESOLVER_SENTINEL") + assert result.success is True + assert "RESOLVER_SENTINEL" in result.content + + scan = resolved.by_name["scan_chunks"] + assert scan._engine is engine + assert scan._model == "agent-selected-model" + finally: + # All three knowledge tools share this store connection. + resolved.by_name["knowledge_sql"]._store.close() + + +@pytest.mark.parametrize( + "tool_config", + [ + ["alpha", "beta", "alpha"], + " alpha, beta, alpha ", + ], +) +def test_configured_tools_normalize_lists_and_comma_separated_strings( + tool_config, +) -> None: + ToolRegistry.register_value("alpha", _AlphaTool) + ToolRegistry.register_value("beta", _BetaTool) + + resolved = tool_resolver.resolve_agent_tools( + {"agent_type": "simple", "config": {"tools": tool_config}}, + engine=object(), + model="test-model", + ) + + assert [tool.spec.name for tool in resolved.instances] == ["alpha", "beta"] + assert [spec["function"]["name"] for spec in resolved.openai_specs] == [ + "alpha", + "beta", + ] + + +def test_registered_tool_advertisement_matches_to_openai_function( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Runtime description overrides must reach canonical advertisements.""" + + ToolRegistry.register_value("alpha", _AlphaTool) + monkeypatch.setattr( + description_loader, + "get_tool_description_override", + lambda name: "Runtime alpha description" if name == "alpha" else None, + ) + + resolved = tool_resolver.resolve_agent_tools( + {"agent_type": "simple", "config": {"tools": ["alpha"]}}, + engine=object(), + model="test-model", + ) + + tool = resolved.by_name["alpha"] + assert resolved.openai_specs == [tool.to_openai_function()] + assert ( + resolved.openai_specs[0]["function"]["description"] + == "Runtime alpha description" + ) + + +def test_explicit_config_schema_takes_priority_over_tool_advertisement( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ToolRegistry.register_value("alpha", _AlphaTool) + monkeypatch.setattr( + description_loader, + "get_tool_description_override", + lambda name: "Runtime alpha description" if name == "alpha" else None, + ) + custom_spec = { + "type": "function", + "function": { + "name": "alpha", + "description": "Agent-specific alpha description", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, + } + + resolved = tool_resolver.resolve_agent_tools( + {"agent_type": "simple", "config": {"tools": [custom_spec]}}, + engine=object(), + model="test-model", + ) + + assert resolved.openai_specs == [custom_spec] + assert resolved.by_name["alpha"].to_openai_function() == custom_spec + + +def test_invalid_tool_advertisement_falls_back_to_tool_spec() -> None: + class _InvalidAdvertisementTool(_AlphaTool): + def to_openai_function(self) -> dict[str, object]: + raise RuntimeError("broken advertisement") + + ToolRegistry.register_value("alpha", _InvalidAdvertisementTool) + + resolved = tool_resolver.resolve_agent_tools( + {"agent_type": "simple", "config": {"tools": ["alpha"]}}, + engine=object(), + model="test-model", + ) + + assert resolved.openai_specs == [ + { + "type": "function", + "function": { + "name": "alpha", + "description": "Alpha test tool", + "parameters": {}, + }, + } + ] + + +def test_mcp_tools_merge_after_native_tools_without_name_collisions() -> None: + ToolRegistry.register_value("shared", _NativeSharedTool) + mcp_shared = _MCPSharedTool() + mcp_only = _MCPOnlyTool() + client = object() + + resolved = tool_resolver.resolve_agent_tools( + { + "agent_type": "simple", + "config": {"tools": ["shared", "shared"]}, + }, + engine=object(), + model="test-model", + mcp_tools=[mcp_shared, mcp_only, mcp_only], + mcp_clients=[client], + ) + + assert [tool.spec.name for tool in resolved.instances] == ["shared", "mcp_only"] + assert isinstance(resolved.by_name["shared"], _NativeSharedTool) + assert resolved.by_name["mcp_only"] is mcp_only + assert resolved.mcp_clients == [client] + assert [spec["function"]["name"] for spec in resolved.openai_specs] == [ + "shared", + "mcp_only", + ] + + +def test_mcp_tools_can_be_disabled_per_agent() -> None: + ToolRegistry.register_value("shared", _NativeSharedTool) + + class _MustNotIterate: + def __iter__(self): + raise AssertionError("MCP tools must not be inspected after opt-out") + + resolved = tool_resolver.resolve_agent_tools( + { + "agent_type": "simple", + "config": {"tools": ["shared"], "mcp_tools": False}, + }, + engine=object(), + model="test-model", + mcp_tools=_MustNotIterate(), + mcp_clients=_MustNotIterate(), + ) + + assert [tool.spec.name for tool in resolved.instances] == ["shared"] + assert resolved.mcp_clients == [] diff --git a/tests/cli/test_serve_single_build.py b/tests/cli/test_serve_single_build.py index 1a0fddf8..0453d49c 100644 --- a/tests/cli/test_serve_single_build.py +++ b/tests/cli/test_serve_single_build.py @@ -101,6 +101,7 @@ def _run_serve(tmp_path, monkeypatch, *, build_spy, set_system_spy): (``uvicorn.run`` is a no-op) and no real engine is contacted. """ from openjarvis.core.config import JarvisConfig + from openjarvis.core.registry import MemoryRegistry _repopulate_registries() @@ -112,6 +113,9 @@ def _run_serve(tmp_path, monkeypatch, *, build_spy, set_system_spy): config.sessions.enabled = True config.sessions.db_path = str(tmp_path / "sessions.db") config.memory.db_path = str(tmp_path / "memory.db") + # Disabling prompt-context injection must not disable the backend needed + # by explicitly configured memory tools in managed-agent ticks. + config.agent.context_from_memory = False config.telemetry.enabled = False config.traces.enabled = False config.channel.enabled = False @@ -122,6 +126,16 @@ def _run_serve(tmp_path, monkeypatch, *, build_spy, set_system_spy): config.intelligence.default_model = "test-model" engine = _fake_engine() + # Keep this wiring test independent of the optional native memory runtime. + # The assertion is that serve resolves and passes a backend even when + # prompt-context injection is disabled, not that SQLite itself works. + memory_backend = MagicMock(name="memory_backend") + monkeypatch.setattr(MemoryRegistry, "contains", MagicMock(return_value=True)) + monkeypatch.setattr( + MemoryRegistry, + "create", + MagicMock(return_value=memory_backend), + ) monkeypatch.setattr(serve_mod, "load_config", lambda *a, **k: config) monkeypatch.setattr(serve_mod, "get_engine", lambda *a, **k: ("mock", engine)) diff --git a/tests/mcp/test_client.py b/tests/mcp/test_client.py index 0ee692fb..c23bd868 100644 --- a/tests/mcp/test_client.py +++ b/tests/mcp/test_client.py @@ -2,10 +2,15 @@ from __future__ import annotations +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import MagicMock + import pytest from openjarvis.mcp.client import MCPClient -from openjarvis.mcp.protocol import MCPError +from openjarvis.mcp.protocol import MCPError, MCPResponse from openjarvis.mcp.server import MCPServer from openjarvis.mcp.transport import InProcessTransport from openjarvis.tools._stubs import ToolSpec @@ -101,3 +106,101 @@ class TestMCPClient: result = client.call_tool("think") # Think tool echoes empty thought assert result["isError"] is False + + def test_shared_client_serializes_transport_round_trips(self): + """Concurrent agents cannot consume one another's MCP responses.""" + + class _ConcurrencyProbeTransport: + def __init__(self): + self.active = 0 + self.max_active = 0 + self.lock = threading.Lock() + + def send(self, request): + with self.lock: + self.active += 1 + self.max_active = max(self.max_active, self.active) + time.sleep(0.01) + with self.lock: + self.active -= 1 + return MCPResponse(result={"tools": []}, id=request.id) + + def send_notification(self, request): + return None + + def close(self): + return None + + transport = _ConcurrencyProbeTransport() + shared_client = MCPClient(transport) + + with ThreadPoolExecutor(max_workers=8) as pool: + list(pool.map(lambda _: shared_client.list_tools(), range(24))) + + assert transport.max_active == 1 + + def test_close_interrupts_blocked_request_and_rejects_queued_request(self): + """Shutdown reaches the transport without waiting on an in-flight call.""" + + class _BlockingTransport: + def __init__(self): + self.send_started = threading.Event() + self.send_released = threading.Event() + self.close_called = threading.Event() + self.send_count = 0 + + def send(self, request): + self.send_count += 1 + self.send_started.set() + self.send_released.wait() + raise RuntimeError("transport closed") + + def send_notification(self, request): + return None + + def close(self): + self.close_called.set() + self.send_released.set() + + transport = _BlockingTransport() + shared_client = MCPClient(transport) + + with ThreadPoolExecutor(max_workers=3) as pool: + blocked_request = pool.submit(shared_client.list_tools) + assert transport.send_started.wait(timeout=1) + + queued_request = pool.submit(shared_client.list_tools) + close_call = pool.submit(shared_client.close) + + try: + close_reached_transport = transport.close_called.wait(timeout=1) + finally: + # Keep the test failure-safe against a regression that makes + # close wait behind the blocked request. + transport.send_released.set() + + close_call.result(timeout=1) + assert close_reached_transport + with pytest.raises(RuntimeError, match="transport closed"): + blocked_request.result(timeout=1) + with pytest.raises(RuntimeError, match="MCP client is closed"): + queued_request.result(timeout=1) + + assert transport.send_count == 1 + + def test_close_retries_transport_cleanup_after_failure(self): + """A failed close keeps requests blocked but permits cleanup retry.""" + + transport = MagicMock() + transport.close.side_effect = [RuntimeError("terminate timed out"), None] + client = MCPClient(transport) + + with pytest.raises(RuntimeError, match="terminate timed out"): + client.close() + with pytest.raises(RuntimeError, match="MCP client is closed"): + client.list_tools() + + client.close() + client.close() + + assert transport.close.call_count == 2 diff --git a/tests/mcp/test_discovery.py b/tests/mcp/test_discovery.py index 1bfde7cf..9dc15097 100644 --- a/tests/mcp/test_discovery.py +++ b/tests/mcp/test_discovery.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from unittest.mock import MagicMock, patch import pytest @@ -181,6 +182,132 @@ class TestClientPersistence: assert len(builder._mcp_clients) == 3 +def test_builder_retains_full_mcp_pool_for_managed_agents() -> None: + """Global primary-agent filters must not trim managed-agent MCP tools.""" + + from openjarvis.core.config import JarvisConfig + from openjarvis.system import SystemBuilder + + config = JarvisConfig() + config.tools.mcp.servers = json.dumps( + [{"name": "test", "url": "http://localhost:8080/mcp"}] + ) + external = _make_mock_tool("mcp_only") + builder = SystemBuilder(config).tools(["native_only"]) + + with ( + patch("openjarvis.mcp.server.MCPServer") as mcp_server_cls, + patch.object( + builder, + "_discover_external_mcp", + return_value=[external], + ), + ): + mcp_server_cls.return_value.get_tools.return_value = [] + primary_tools = builder._resolve_tools( + config, + engine=MagicMock(), + model="test-model", + memory_backend=None, + ) + + assert primary_tools == [] + assert builder._mcp_tools == [external] + + +def test_builder_global_mcp_disable_prevents_discovery() -> None: + """A global MCP disable is honored by every managed-agent entry path.""" + + from openjarvis.core.config import JarvisConfig + from openjarvis.system import SystemBuilder + + config = JarvisConfig() + config.tools.mcp.enabled = False + config.tools.mcp.servers = json.dumps( + [{"name": "disabled", "url": "http://localhost:8080/mcp"}] + ) + builder = SystemBuilder(config) + + with ( + patch("openjarvis.mcp.server.MCPServer") as mcp_server_cls, + patch.object(builder, "_discover_external_mcp") as discover, + ): + mcp_server_cls.return_value.get_tools.return_value = [] + builder._resolve_tools( + config, + engine=MagicMock(), + model="test-model", + memory_backend=None, + ) + + discover.assert_not_called() + assert builder._mcp_tools == [] + + +def test_reused_builder_transfers_only_current_build_mcp_state() -> None: + """Each built system exclusively owns its own MCP clients and tools.""" + + from openjarvis.core.config import JarvisConfig + from openjarvis.system import SystemBuilder + + config = JarvisConfig() + config.telemetry.enabled = False + config.traces.enabled = False + config.skills.enabled = False + config.agent_manager.enabled = False + config.tools.mcp.servers = json.dumps( + [{"name": "test", "url": "http://localhost:8080/mcp"}] + ) + + engine = MagicMock(spec=["health", "can_serve", "generate", "list_models", "close"]) + engine.health.return_value = True + first_tool = _make_mock_tool("first_mcp_tool") + second_tool = _make_mock_tool("second_mcp_tool") + first_client = MagicMock() + second_client = MagicMock() + discoveries = iter([(first_tool, first_client), (second_tool, second_client)]) + + builder = ( + SystemBuilder(config) + .engine_instance(engine) + .model("test-model") + .tools([]) + .telemetry(False) + .traces(False) + .speech(False) + ) + + def _discover(_server_cfg): + tool, client = next(discoveries) + builder._mcp_clients.append(client) + return [tool] + + with ( + patch.object(builder, "_discover_external_mcp", side_effect=_discover), + patch.object(builder, "_resolve_memory", return_value=None), + ): + first_system = builder.build() + assert first_system.mcp_tools == [first_tool] + assert first_system._mcp_clients == [first_client] + assert builder._mcp_tools == [] + assert builder._mcp_clients == [] + first_system.close() + + second_system = builder.build() + + try: + assert second_system.mcp_tools == [second_tool] + assert second_system._mcp_clients == [second_client] + assert first_client not in second_system._mcp_clients + assert builder._mcp_tools == [] + assert builder._mcp_clients == [] + finally: + second_system.close() + + first_client.close.assert_called_once() + second_client.close.assert_called_once() + + class TestStringConfig: @patch(_PATCH_PROVIDER) @patch(_PATCH_CLIENT) diff --git a/tests/sdk/test_system.py b/tests/sdk/test_system.py index 4502ea47..79f0c138 100644 --- a/tests/sdk/test_system.py +++ b/tests/sdk/test_system.py @@ -12,6 +12,34 @@ from openjarvis.system import JarvisSystem, SystemBuilder class TestJarvisSystem: + def test_new_fields_do_not_shift_existing_positional_arguments(self): + """Adding mcp_tools must not reinterpret legacy positional calls.""" + config = JarvisConfig() + bus = EventBus() + engine = MagicMock() + agent = MagicMock() + tools = [MagicMock()] + tool_executor = MagicMock() + memory_backend = MagicMock() + + system = JarvisSystem( + config, + bus, + engine, + "mock", + "test-model", + agent, + "simple", + tools, + tool_executor, + memory_backend, + ) + + assert system.tools is tools + assert system.tool_executor is tool_executor + assert system.memory_backend is memory_backend + assert system.mcp_tools == [] + def test_ask_direct_mode(self): engine = MagicMock() engine.generate.return_value = { diff --git a/tests/server/test_agent_manager_routes.py b/tests/server/test_agent_manager_routes.py index 6dda4fbc..aded7532 100644 --- a/tests/server/test_agent_manager_routes.py +++ b/tests/server/test_agent_manager_routes.py @@ -4,6 +4,8 @@ from __future__ import annotations import json import tempfile +import threading +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -602,3 +604,81 @@ class TestLightweightSystemEngineResolution: engine=MagicMock(), model="m", config=self._cfg(None, "llamacpp") ) assert captured["key"] == "llamacpp" + + def test_caches_tool_memory_backend_when_prompt_context_is_disabled( + self, + monkeypatch, + ): + pytest.importorskip("fastapi") + from openjarvis.server import agent_manager_routes as amr + + backend = object() + resolver = MagicMock(return_value=backend) + monkeypatch.setattr(amr, "_resolve_memory_backend", resolver) + config = SimpleNamespace( + agent=SimpleNamespace(context_from_memory=False), + memory=SimpleNamespace(default_backend="sqlite", db_path="memory.db"), + ) + runtime = SimpleNamespace( + memory_backend=None, + _owns_memory_backend=False, + channel_backend=None, + channel_bridge=None, + knowledge_db_path=None, + ) + + system = amr._LightweightSystem( + engine=MagicMock(), + model="m", + config=config, + runtime=runtime, + ) + + resolver.assert_called_once_with(config) + assert system.memory_backend is backend + assert runtime.memory_backend is backend + assert runtime._owns_memory_backend is True + + def test_memory_backend_lazy_init_is_synchronized(self, monkeypatch): + pytest.importorskip("fastapi") + from openjarvis.server import agent_manager_routes as amr + + backend = object() + resolver_calls = 0 + calls_lock = threading.Lock() + duplicate_entered = threading.Event() + start = threading.Barrier(8) + + def _resolve(config): + nonlocal resolver_calls + with calls_lock: + resolver_calls += 1 + call_number = resolver_calls + if call_number > 1: + duplicate_entered.set() + # A check-then-create race lets another worker enter while the + # first resolver is blocked here. The locked implementation times + # out once, publishes the backend, and all other workers reuse it. + if call_number == 1: + duplicate_entered.wait(timeout=0.2) + return backend + + monkeypatch.setattr(amr, "_resolve_memory_backend", _resolve) + config = SimpleNamespace() + runtime = SimpleNamespace( + memory_backend=None, + _owns_memory_backend=False, + _managed_runtime_stopping=False, + ) + + def _get_backend(): + start.wait(timeout=2) + return amr._get_or_create_memory_backend(runtime, config) + + with ThreadPoolExecutor(max_workers=8) as pool: + results = list(pool.map(lambda _: _get_backend(), range(8))) + + assert resolver_calls == 1 + assert results == [backend] * 8 + assert runtime.memory_backend is backend + assert runtime._owns_memory_backend is True diff --git a/tests/server/test_deep_research_tools_wiring.py b/tests/server/test_deep_research_tools_wiring.py index f568471b..5de649af 100644 --- a/tests/server/test_deep_research_tools_wiring.py +++ b/tests/server/test_deep_research_tools_wiring.py @@ -2,7 +2,9 @@ from __future__ import annotations +import json from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -15,6 +17,82 @@ except ImportError: HAS_FASTAPI = False from openjarvis.connectors.store import KnowledgeStore +from openjarvis.core.registry import ToolRegistry +from openjarvis.core.types import Role, ToolResult +from openjarvis.tools._stubs import BaseTool, ToolSpec + + +class _ConfiguredResearchProbe(BaseTool): + """Configured native tool used to exercise the Deep Research SSE path.""" + + tool_id = "configured_research_probe_682" + calls = 0 + + @property + def spec(self) -> ToolSpec: + return ToolSpec( + name=self.tool_id, + description="Configured Deep Research probe", + parameters={ + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + }, + ) + + def execute(self, **params) -> ToolResult: + type(self).calls += 1 + return ToolResult( + tool_name=self.tool_id, + content=f"configured:{params['value']}", + ) + + +class _MCPResearchProbe(BaseTool): + """MCP-shaped adapter that must be merged into the same toolkit.""" + + tool_id = "mcp_research_probe_682" + + @property + def spec(self) -> ToolSpec: + return ToolSpec(name=self.tool_id, description="MCP Deep Research probe") + + def execute(self, **params) -> ToolResult: + return ToolResult(tool_name=self.tool_id, content="mcp") + + +class _ScriptedDeepResearchEngine: + """Call the configured probe once, then return a final answer.""" + + def __init__(self) -> None: + self.turns = 0 + self.advertised_names: list[str] = [] + self.observed_tool_result = "" + + def generate(self, messages, *, model, **kwargs): + self.turns += 1 + self.advertised_names = [ + spec["function"]["name"] for spec in kwargs.get("tools", []) + ] + if self.turns == 1: + return { + "content": "", + "tool_calls": [ + { + "id": "call-configured-research-probe", + "type": "function", + "function": { + "name": _ConfiguredResearchProbe.tool_id, + "arguments": json.dumps({"value": "sentinel"}), + }, + } + ], + "usage": {}, + } + + tool_messages = [message for message in messages if message.role is Role.TOOL] + self.observed_tool_result = tool_messages[-1].content + return {"content": "complete", "tool_calls": [], "usage": {}} @pytest.mark.skipif(not HAS_FASTAPI, reason="fastapi not installed") @@ -53,3 +131,111 @@ def test_deep_research_tools_returns_empty_when_no_db() -> None: ) assert tools == [] + + +@pytest.mark.skipif(not HAS_FASTAPI, reason="fastapi not installed") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "with_knowledge_db", + [False, True], + ids=["without-knowledge-db", "with-knowledge-db"], +) +async def test_server_deep_research_merges_and_executes_all_tool_sources( + tmp_path: Path, + with_knowledge_db: bool, + monkeypatch, +) -> None: + """Configured and MCP tools reach Deep Research with or without its DB.""" + + from openjarvis.server import agent_manager_routes as routes + + start_worker = MagicMock(wraps=routes._start_managed_worker) + monkeypatch.setattr(routes, "_start_managed_worker", start_worker) + + db_path = tmp_path / "knowledge.db" + if with_knowledge_db: + store = KnowledgeStore(str(db_path)) + store.store("test content", source="test", doc_type="note") + store.close() + + if not ToolRegistry.contains(_ConfiguredResearchProbe.tool_id): + ToolRegistry.register_value( + _ConfiguredResearchProbe.tool_id, + _ConfiguredResearchProbe, + ) + _ConfiguredResearchProbe.calls = 0 + + mcp_tool = _MCPResearchProbe() + app_state = SimpleNamespace( + config=SimpleNamespace(memory_files=None, system_prompt=None), + memory_backend=None, + channel_backend=None, + channel_bridge=None, + knowledge_db_path=str(db_path), + _mcp_clients=[object()], + _mcp_tools_cache=( + [mcp_tool.to_openai_function()], + {mcp_tool.spec.name: mcp_tool}, + ), + ) + manager = MagicMock() + manager.list_messages.return_value = [] + engine = _ScriptedDeepResearchEngine() + + response = await routes._stream_managed_agent( + manager=manager, + agent_record={ + "id": "agent-deep-research-682", + "name": "Deep Research Agent", + "agent_type": "deep_research", + "config": { + "model": "test-model", + "max_turns": 3, + "tools": [_ConfiguredResearchProbe.tool_id], + }, + }, + user_content="Use the configured research probe", + message_id="message-deep-research-682", + engine=engine, + bus=None, + app_state=app_state, + ) + + body_parts: list[str] = [] + async for part in response.body_iterator: + body_parts.append(part.decode() if isinstance(part, bytes) else part) + + expected_names = { + _ConfiguredResearchProbe.tool_id, + _MCPResearchProbe.tool_id, + } + knowledge_names = { + "knowledge_search", + "knowledge_sql", + "scan_chunks", + "think", + } + if with_knowledge_db: + expected_names.update(knowledge_names) + + assert set(engine.advertised_names) == expected_names + assert len(engine.advertised_names) == len(expected_names) + assert not with_knowledge_db or knowledge_names.issubset(engine.advertised_names) + assert with_knowledge_db or knowledge_names.isdisjoint(engine.advertised_names) + assert engine.turns == 2 + assert _ConfiguredResearchProbe.calls == 1 + assert engine.observed_tool_result == "configured:sentinel" + assert "data: [DONE]" in "".join(body_parts) + start_worker.assert_called_once() + assert start_worker.call_args.kwargs["name"].startswith( + "managed-agent-deep-research-" + ) + assert app_state._managed_workers == set() + + manager.store_agent_response.assert_called_once() + stored = manager.store_agent_response.call_args + assert stored.args[:2] == ("agent-deep-research-682", "complete") + persisted_calls = stored.kwargs["tool_calls"] + assert persisted_calls[0]["tool"] == _ConfiguredResearchProbe.tool_id + assert persisted_calls[0]["result"] == "configured:sentinel" + assert persisted_calls[0]["success"] is True diff --git a/tests/server/test_managed_agent_resolved_tools.py b/tests/server/test_managed_agent_resolved_tools.py new file mode 100644 index 00000000..c1ca4f66 --- /dev/null +++ b/tests/server/test_managed_agent_resolved_tools.py @@ -0,0 +1,302 @@ +"""SSE regression coverage for canonical managed-agent tool resolution.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +pytest.importorskip("fastapi") + +from openjarvis.core.registry import ToolRegistry # noqa: E402 +from openjarvis.core.types import Role, ToolResult # noqa: E402 +from openjarvis.engine._stubs import StreamChunk # noqa: E402 +from openjarvis.tools._stubs import BaseTool, ToolSpec # noqa: E402 + + +class _StatefulConfiguredTool(BaseTool): + """A tool whose result identifies the exact instance that executed.""" + + tool_id = "stateful_probe" + instances: list["_StatefulConfiguredTool"] = [] + + def __init__(self) -> None: + self.instance_id = len(self.instances) + 1 + self.calls = 0 + self.instances.append(self) + + @property + def spec(self) -> ToolSpec: + return ToolSpec( + name="stateful_probe", + description=f"configured-instance-{self.instance_id}", + parameters={ + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + }, + ) + + def execute(self, **params) -> ToolResult: + self.calls += 1 + return ToolResult( + tool_name=self.spec.name, + content=( + f"instance={self.instance_id};calls={self.calls};" + f"value={params['value']}" + ), + ) + + +class _CollidingMCPTool(BaseTool): + """An MCP-shaped collision that must lose to the configured native tool.""" + + tool_id = "mcp_stateful_probe" + + def __init__(self) -> None: + self.calls = 0 + + @property + def spec(self) -> ToolSpec: + return ToolSpec( + name="stateful_probe", + description="mcp-collision", + parameters={"type": "object", "properties": {}}, + ) + + def execute(self, **params) -> ToolResult: + self.calls += 1 + return ToolResult(tool_name=self.spec.name, content="wrong MCP instance") + + +class _ToolCallingEngine: + """Advertise the toolkit, request one call, then observe its result.""" + + def __init__( + self, + tool_name: str = "stateful_probe", + arguments: dict | None = None, + ) -> None: + self.tool_name = tool_name + self.arguments = arguments or {"value": "sentinel"} + self.turns = 0 + self.advertised_specs: list[dict] = [] + self.observed_tool_result = "" + + async def stream_full(self, messages, *, model, **kwargs): + self.turns += 1 + self.advertised_specs = list(kwargs.get("tools", [])) + + if self.turns == 1: + yield StreamChunk( + tool_calls=[ + { + "index": 0, + "id": f"call-{self.tool_name}", + "type": "function", + "function": { + "name": self.tool_name, + "arguments": json.dumps(self.arguments), + }, + } + ], + finish_reason="tool_calls", + ) + return + + tool_messages = [message for message in messages if message.role is Role.TOOL] + self.observed_tool_result = tool_messages[-1].content + yield StreamChunk(content="complete") + yield StreamChunk(finish_reason="stop") + + +class _FinalOnlyEngine: + async def stream_full(self, messages, *, model, **kwargs): + yield StreamChunk(content="complete") + yield StreamChunk(finish_reason="stop") + + +@pytest.mark.asyncio +async def test_sse_advertises_and_executes_the_same_resolved_tool_instance() -> None: + """The schema and dispatch map must come from one first-wins toolkit.""" + + from openjarvis.server.agent_manager_routes import _stream_managed_agent + + _StatefulConfiguredTool.instances.clear() + ToolRegistry.register_value("stateful_probe", _StatefulConfiguredTool) + + colliding_mcp = _CollidingMCPTool() + mcp_spec = colliding_mcp.to_openai_function() + app_state = SimpleNamespace( + config=SimpleNamespace(memory_files=None, system_prompt=None), + memory_backend=None, + channel_backend=None, + channel_bridge=None, + _mcp_clients=[object()], + _mcp_tools_cache=( + [mcp_spec], + {"stateful_probe": colliding_mcp}, + ), + ) + manager = MagicMock() + manager.list_messages.return_value = [] + engine = _ToolCallingEngine() + custom_spec = { + "type": "function", + "function": { + "name": "stateful_probe", + "description": "custom configured schema", + "parameters": { + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + }, + }, + } + + response = await _stream_managed_agent( + manager=manager, + agent_record={ + "id": "agent-stateful", + "name": "Stateful Agent", + "agent_type": "simple", + "config": { + "model": "test-model", + "max_turns": 3, + "tools": [custom_spec], + }, + }, + user_content="Use the stateful probe", + message_id="message-stateful", + engine=engine, + bus=None, + app_state=app_state, + ) + + body_parts: list[str] = [] + async for part in response.body_iterator: + body_parts.append(part.decode() if isinstance(part, bytes) else part) + + assert engine.turns == 2 + assert len(_StatefulConfiguredTool.instances) == 1 + configured_instance = _StatefulConfiguredTool.instances[0] + assert configured_instance.calls == 1 + assert colliding_mcp.calls == 0 + + advertised = [ + spec + for spec in engine.advertised_specs + if spec.get("function", {}).get("name") == "stateful_probe" + ] + assert len(advertised) == 1 + assert advertised[0] is custom_spec + assert advertised[0]["function"]["description"] == "custom configured schema" + assert engine.observed_tool_result == "instance=1;calls=1;value=sentinel" + assert "data: [DONE]" in "".join(body_parts) + + +@pytest.mark.asyncio +async def test_sse_mcp_opt_out_skips_discovery(monkeypatch) -> None: + """Opting out skips request-local discovery and hides MCP specs.""" + + from openjarvis.server import agent_manager_routes as routes + + discovery = MagicMock(side_effect=AssertionError("MCP discovery must not run")) + monkeypatch.setattr(routes, "_get_mcp_tools", discovery) + manager = MagicMock() + manager.list_messages.return_value = [] + app_state = SimpleNamespace( + config=SimpleNamespace(memory_files=None, system_prompt=None), + memory_backend=None, + channel_backend=None, + channel_bridge=None, + ) + + response = await routes._stream_managed_agent( + manager=manager, + agent_record={ + "id": "agent-no-mcp", + "name": "No MCP", + "agent_type": "simple", + "config": {"model": "test-model", "mcp_tools": False}, + }, + user_content="Answer directly", + message_id="message-no-mcp", + engine=_FinalOnlyEngine(), + bus=None, + app_state=app_state, + ) + + async for _ in response.body_iterator: + pass + + discovery.assert_not_called() + + +@pytest.mark.asyncio +async def test_sse_memory_tools_resolve_backend_when_context_injection_is_off( + monkeypatch, +) -> None: + """Prompt context opt-out must not disable explicit memory tools.""" + + from openjarvis.server import agent_manager_routes as routes + from openjarvis.tools.storage_tools import MemoryStoreTool + + if not ToolRegistry.contains("memory_store"): + ToolRegistry.register_value("memory_store", MemoryStoreTool) + + backend = MagicMock() + backend.store.return_value = "doc-1" + resolver = MagicMock(return_value=backend) + monkeypatch.setattr(routes, "_resolve_memory_backend", resolver) + + manager = MagicMock() + manager.list_messages.return_value = [] + app_config = SimpleNamespace( + memory_files=None, + system_prompt=None, + agent=SimpleNamespace(context_from_memory=False), + memory=SimpleNamespace(default_backend="sqlite", db_path="memory.db"), + ) + app_state = SimpleNamespace( + config=app_config, + memory_backend=None, + channel_backend=None, + channel_bridge=None, + _mcp_clients=[], + _mcp_tools_cache=([], {}), + ) + engine = _ToolCallingEngine( + tool_name="memory_store", + arguments={"content": "remember me"}, + ) + + response = await routes._stream_managed_agent( + manager=manager, + agent_record={ + "id": "agent-memory-tool", + "name": "Memory Tool Agent", + "agent_type": "simple", + "config": { + "model": "test-model", + "max_turns": 3, + "tools": ["memory_store"], + }, + }, + user_content="Remember this", + message_id="message-memory-tool", + engine=engine, + bus=None, + app_state=app_state, + ) + + async for _ in response.body_iterator: + pass + + resolver.assert_called_once_with(app_config) + assert app_state.memory_backend is backend + assert app_state._owns_memory_backend is True + backend.store.assert_called_once_with("remember me", source="") + assert engine.observed_tool_result == "Stored as doc-1" diff --git a/tests/server/test_mcp_tools_cache.py b/tests/server/test_mcp_tools_cache.py index d3bcdfa4..ccc6b26a 100644 --- a/tests/server/test_mcp_tools_cache.py +++ b/tests/server/test_mcp_tools_cache.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import threading from unittest.mock import MagicMock, patch import pytest @@ -118,7 +119,7 @@ def test_does_not_cache_empty_results(mock_load_config: MagicMock): with ( patch("openjarvis.mcp.transport.StreamableHTTPTransport"), - patch("openjarvis.mcp.client.MCPClient"), + patch("openjarvis.mcp.client.MCPClient") as MockClient, patch("openjarvis.tools.mcp_adapter.MCPToolProvider") as MockProvider, ): # First call: discovery returns empty @@ -127,6 +128,8 @@ def test_does_not_cache_empty_results(mock_load_config: MagicMock): tools1, _ = _get_mcp_tools(app_state) assert len(tools1) == 0 + MockClient.return_value.close.assert_called_once_with() + assert getattr(app_state, "_mcp_clients", []) == [] # Verify no cache was set (empty result) assert getattr(app_state, "_mcp_tools_cache", None) is None @@ -152,3 +155,298 @@ def test_handles_config_load_failure(mock_load_config: MagicMock): assert tools == [] assert adapters == {} + + +@patch("openjarvis.core.config.load_config") +def test_uses_preloaded_full_system_pool(mock_load_config: MagicMock): + """Server and scheduled paths reuse one unfiltered MCP discovery.""" + + from openjarvis.server.agent_manager_routes import _get_mcp_tools + + adapter = _make_adapter("preloaded_tool") + app_state = _FakeAppState() + app_state.mcp_tools = [adapter] + + tools, adapters = _get_mcp_tools(app_state) + + mock_load_config.assert_not_called() + assert tools[0]["function"]["name"] == "preloaded_tool" + assert adapters == {"preloaded_tool": adapter} + + +@patch("openjarvis.core.config.load_config") +def test_preloaded_duplicate_names_are_first_wins(mock_load_config: MagicMock): + """SSE and executor paths choose the same adapter on name collisions.""" + + from openjarvis.server.agent_manager_routes import _get_mcp_tools + + first = _make_adapter("duplicate") + second = _make_adapter("duplicate") + app_state = _FakeAppState() + app_state.mcp_tools = [first, second] + + tools, adapters = _get_mcp_tools(app_state) + + mock_load_config.assert_not_called() + assert len(tools) == 1 + assert adapters == {"duplicate": first} + + +def test_app_shutdown_stops_scheduler_before_closing_shared_mcp_clients() -> None: + """Shutdown quiesces and drains every user of the shared MCP pool.""" + + from fastapi.testclient import TestClient + + from openjarvis.core.config import JarvisConfig + from openjarvis.server.agent_manager_routes import _start_managed_worker + from openjarvis.server.app import create_app + + events: list[str] = [] + release_worker = threading.Event() + worker_finished = threading.Event() + + class _Scheduler: + def request_stop(self): + events.append("scheduler-stop") + + def wait_stopped(self, timeout=10): + events.append("scheduler-wait") + return True + + scheduler = _Scheduler() + mcp_client = MagicMock() + memory_backend = MagicMock() + channel_bridge = MagicMock() + + def _close_mcp(): + events.append("mcp") + release_worker.set() + + mcp_client.close.side_effect = _close_mcp + memory_backend.close.side_effect = lambda: events.append("memory") + channel_bridge.disconnect.side_effect = lambda: events.append("channel") + config = JarvisConfig() + config.analytics.enabled = False + config.traces.enabled = False + + app = create_app( + MagicMock(), + "test-model", + config=config, + channel_bridge=channel_bridge, + agent_scheduler=scheduler, + mcp_clients=[mcp_client], + memory_backend=memory_backend, + own_memory_backend=True, + ) + + def _worker(): + release_worker.wait(timeout=2) + events.append("worker-finished") + worker_finished.set() + + _start_managed_worker(app.state, _worker, name="test-managed-worker") + + with TestClient(app): + pass + + mcp_client.close.assert_called_once_with() + memory_backend.close.assert_called_once_with() + channel_bridge.disconnect.assert_called_once_with() + assert worker_finished.is_set() + assert app.state.memory_backend is None + assert app.state._owns_memory_backend is False + assert app.state._managed_runtime_stopping is True + assert app.state._managed_workers == set() + assert events.index("channel") < events.index("mcp") + assert events.index("scheduler-stop") < events.index("mcp") + assert events.index("mcp") < events.index("worker-finished") + assert events.index("worker-finished") < events.index("memory") + with pytest.raises(RuntimeError, match="shutting down"): + _start_managed_worker(app.state, lambda: None, name="too-late-worker") + + +def test_shutdown_interrupts_mcp_client_during_lazy_initialization() -> None: + """A client is registered before initialize() can block on transport I/O.""" + + from fastapi.testclient import TestClient + + from openjarvis.core.config import JarvisConfig + from openjarvis.server.agent_manager_routes import ( + _get_mcp_tools, + _start_managed_worker, + ) + from openjarvis.server.app import create_app + + initialize_started = threading.Event() + initialize_released = threading.Event() + discovery_finished = threading.Event() + + class _BlockingClient: + def __init__(self): + self.closed = False + self.close_calls = 0 + + def initialize(self): + initialize_started.set() + initialize_released.wait(timeout=2) + if self.closed: + raise RuntimeError("transport closed during initialize") + + def close(self): + self.close_calls += 1 + self.closed = True + initialize_released.set() + + client = _BlockingClient() + config = JarvisConfig() + config.analytics.enabled = False + config.traces.enabled = False + app = create_app(MagicMock(), "test-model", config=config) + mcp_config = _make_config( + servers_json=json.dumps([{"name": "blocking", "url": "http://localhost:9999"}]) + ) + + def _discover(): + try: + _get_mcp_tools(app.state) + finally: + discovery_finished.set() + + with ( + patch("openjarvis.core.config.load_config", return_value=mcp_config), + patch("openjarvis.mcp.transport.StreamableHTTPTransport"), + patch("openjarvis.mcp.client.MCPClient", return_value=client), + ): + _start_managed_worker( + app.state, + _discover, + name="blocking-mcp-discovery", + ) + assert initialize_started.wait(timeout=2) + with app.state._mcp_clients_lock: + assert client in app.state._mcp_clients + + with TestClient(app): + pass + + assert client.close_calls >= 1 + assert discovery_finished.is_set() + assert app.state._managed_workers == set() + assert getattr(app.state, "_mcp_tools_cache", None) is None + + +def test_app_shutdown_closes_lazily_created_memory_backend(monkeypatch) -> None: + """A backend opened by a managed route is owned and closed by the app.""" + + from fastapi.testclient import TestClient + + from openjarvis.core.config import JarvisConfig + from openjarvis.server import agent_manager_routes as routes + from openjarvis.server.app import create_app + + backend = MagicMock() + monkeypatch.setattr(routes, "_resolve_memory_backend", lambda config: backend) + config = JarvisConfig() + config.analytics.enabled = False + config.traces.enabled = False + app = create_app(MagicMock(), "test-model", config=config) + + assert routes._get_or_create_memory_backend(app.state, config) is backend + assert app.state._owns_memory_backend is True + + with TestClient(app): + pass + + backend.close.assert_called_once_with() + assert app.state.memory_backend is None + + +def test_app_shutdown_keeps_owned_memory_open_for_live_worker(monkeypatch) -> None: + """A timed-out worker must never resume against a closed backend.""" + + from fastapi.testclient import TestClient + + from openjarvis.core.config import JarvisConfig + from openjarvis.server import app as app_module + from openjarvis.server.agent_manager_routes import _start_managed_worker + + monkeypatch.setattr(app_module, "_MANAGED_SHUTDOWN_GRACE_SECONDS", 0.01) + monkeypatch.setattr(app_module, "_MANAGED_SHUTDOWN_DRAIN_SECONDS", 0.01) + + release_worker = threading.Event() + worker_holds_memory_lock = threading.Event() + shutdown_finished = threading.Event() + shutdown_errors: list[BaseException] = [] + backend = MagicMock() + config = JarvisConfig() + config.analytics.enabled = False + config.traces.enabled = False + app = app_module.create_app( + MagicMock(), + "test-model", + config=config, + memory_backend=backend, + own_memory_backend=True, + ) + + def _hold_memory_lock(): + with app.state._memory_backend_lock: + worker_holds_memory_lock.set() + release_worker.wait(timeout=2) + + worker = _start_managed_worker( + app.state, + _hold_memory_lock, + name="memory-using-straggler", + ) + assert worker_holds_memory_lock.wait(timeout=2) + + def _shutdown_app(): + try: + with TestClient(app): + pass + except BaseException as exc: + shutdown_errors.append(exc) + finally: + shutdown_finished.set() + + shutdown_thread = threading.Thread(target=_shutdown_app, daemon=True) + shutdown_thread.start() + + try: + assert shutdown_finished.wait(timeout=1) + assert shutdown_errors == [] + backend.close.assert_not_called() + assert app.state.memory_backend is backend + assert app.state._owns_memory_backend is True + finally: + release_worker.set() + worker.join(timeout=2) + shutdown_thread.join(timeout=2) + + +def test_app_shutdown_leaves_borrowed_memory_backend_open() -> None: + """An injected backend remains owned by its caller unless opted in.""" + + from fastapi.testclient import TestClient + + from openjarvis.core.config import JarvisConfig + from openjarvis.server.app import create_app + + backend = MagicMock() + config = JarvisConfig() + config.analytics.enabled = False + config.traces.enabled = False + app = create_app( + MagicMock(), + "test-model", + config=config, + memory_backend=backend, + ) + + with TestClient(app): + pass + + backend.close.assert_not_called() + assert app.state.memory_backend is backend