Compare commits

..
Author SHA1 Message Date
Elliot Slusky 063dd8ea75 fix: unify managed-agent tool resolution (#705)
* fix: unify managed-agent tool resolution

* fix: harden managed-agent tool lifecycle
2026-08-10 15:09:41 -07:00
Elliot Slusky 6af9317556 fix: route configured LiteLLM models by engine ownership (#714) 2026-08-10 12:51:16 -07:00
33 changed files with 3598 additions and 561 deletions
+4 -1
View File
@@ -466,7 +466,10 @@ export function InputArea() {
}
const totalMs = Date.now() - startTime;
const _CLOUD_PREFIXES = ['gpt-', 'o1-', 'o3-', 'o4-', 'claude-', 'gemini-', 'openrouter/', 'MiniMax-', 'chatgpt-'];
const engineLabel = _CLOUD_PREFIXES.some(p => selectedModel.startsWith(p)) ? 'cloud' : 'ollama';
const selectedOwner = useAppStore.getState().models.find((m) => m.id === selectedModel)?.owned_by;
const engineLabel = selectedOwner === 'litellm'
? 'litellm'
: _CLOUD_PREFIXES.some(p => selectedModel.startsWith(p)) ? 'cloud' : 'ollama';
const telemetry: MessageTelemetry = {
engine: engineLabel,
model_id: selectedModel,
+23 -16
View File
@@ -143,7 +143,7 @@ export function CommandPalette() {
}
}, [pullSuccess]);
const handleSelect = async (modelId: string) => {
const handleSelect = async (modelId: string, owner?: string) => {
const previousModel = selectedModel;
setSelectedModel(modelId);
setCommandPaletteOpen(false);
@@ -153,7 +153,7 @@ export function CommandPalette() {
setModelLoading(true);
addLogEntry({ timestamp: Date.now(), level: 'info', category: 'model', message: `Switching to ${modelId}...` });
try {
await preloadModel(modelId);
await preloadModel(modelId, owner);
addLogEntry({ timestamp: Date.now(), level: 'info', category: 'model', message: `${modelId} loaded` });
} catch (e: any) {
addLogEntry({ timestamp: Date.now(), level: 'error', category: 'model', message: `Failed to load ${modelId}: ${e.message}` });
@@ -255,7 +255,8 @@ export function CommandPalette() {
setSelectedIdx((i) => Math.max(i - 1, 0));
} else if (e.key === 'Enter' && tab === 'installed' && filtered.length > 0) {
e.preventDefault();
handleSelect((filtered[selectedIdx] as any).id);
const model = filtered[selectedIdx] as (typeof models)[number];
handleSelect(model.id, model.owned_by);
}
};
@@ -365,11 +366,15 @@ export function CommandPalette() {
onMouseEnter={() => setSelectedIdx(idx)}
>
<button
onClick={() => handleSelect(model.id)}
onClick={() => handleSelect(model.id, model.owned_by)}
className="flex items-center gap-3 flex-1 min-w-0 text-left cursor-pointer"
style={{ background: 'none', border: 'none', padding: 0 }}
>
<Cpu size={16} style={{ color: isActive ? 'var(--color-accent)' : 'var(--color-text-tertiary)' }} />
{model.owned_by === 'litellm' ? (
<Cloud size={16} style={{ color: isActive ? 'var(--color-accent)' : 'var(--color-text-tertiary)' }} />
) : (
<Cpu size={16} style={{ color: isActive ? 'var(--color-accent)' : 'var(--color-text-tertiary)' }} />
)}
<div className="flex-1 min-w-0">
<div className="text-sm truncate" style={{ color: isActive ? 'var(--color-accent)' : 'var(--color-text)', fontWeight: isActive ? 500 : 400 }}>
{model.id}
@@ -381,17 +386,19 @@ export function CommandPalette() {
</span>
)}
</button>
<button
onClick={() => handleDelete(model.id)}
disabled={isDeleting}
className="p-1 rounded transition-colors cursor-pointer"
style={{ color: 'var(--color-text-tertiary)', opacity: 0 }}
title="Delete model"
onMouseEnter={(e) => { e.currentTarget.style.opacity = '1'; e.currentTarget.style.color = 'var(--color-error)'; }}
onMouseLeave={(e) => { e.currentTarget.style.opacity = '0'; e.currentTarget.style.color = 'var(--color-text-tertiary)'; }}
>
{isDeleting ? <Loader2 size={14} className="animate-spin" /> : <Trash2 size={14} />}
</button>
{model.owned_by !== 'litellm' && (
<button
onClick={() => handleDelete(model.id)}
disabled={isDeleting}
className="p-1 rounded transition-colors cursor-pointer"
style={{ color: 'var(--color-text-tertiary)', opacity: 0 }}
title="Delete model"
onMouseEnter={(e) => { e.currentTarget.style.opacity = '1'; e.currentTarget.style.color = 'var(--color-error)'; }}
onMouseLeave={(e) => { e.currentTarget.style.opacity = '0'; e.currentTarget.style.color = 'var(--color-text-tertiary)'; }}
>
{isDeleting ? <Loader2 size={14} className="animate-spin" /> : <Trash2 size={14} />}
</button>
)}
</div>
);
})
+2 -2
View File
@@ -218,9 +218,9 @@ export async function deleteModel(modelName: string): Promise<void> {
const _CLOUD_PREFIXES = ['gpt-', 'o1-', 'o3-', 'o4-', 'claude-', 'gemini-', 'openrouter/'];
export async function preloadModel(modelName: string): Promise<void> {
export async function preloadModel(modelName: string, owner?: string): Promise<void> {
// Cloud models don't need Ollama preloading
if (_CLOUD_PREFIXES.some(p => modelName.startsWith(p))) {
if (owner === 'litellm' || _CLOUD_PREFIXES.some(p => modelName.startsWith(p))) {
return;
}
// Trigger Ollama to load the model into memory (empty prompt, no generation).
+4
View File
@@ -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,
+189 -96
View File
@@ -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)
+7 -1
View File
@@ -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 []
+27 -5
View File
@@ -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",
+1
View File
@@ -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,
+502
View File
@@ -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 "<unnamed>",
)
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",
]
+48 -47
View File
@@ -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")
+6
View File
@@ -35,6 +35,12 @@ def _make_engine(key: str, config: JarvisConfig) -> InferenceEngine:
"""Instantiate a registered engine with the appropriate config host."""
cls = EngineRegistry.get(key)
# LiteLLM cannot enumerate every model supported by every provider. Its
# list_models() contract therefore advertises the configured default
# model, which must be supplied when discovery constructs the engine.
if key == "litellm":
return cls(default_model=config.intelligence.default_model or None)
# gemma_cpp: pass config fields instead of host
if key == "gemma_cpp":
cfg = config.engine.gemma_cpp
+13 -2
View File
@@ -26,16 +26,19 @@ class MultiEngine(InferenceEngine):
def __init__(self, engines: list[tuple[str, InferenceEngine]]) -> None:
self._engines = engines
self._model_map: Dict[str, InferenceEngine] = {}
self._model_key_map: Dict[str, str] = {}
self._refresh_map()
def _refresh_map(self) -> None:
self._model_map.clear()
for _key, engine in self._engines:
self._model_key_map.clear()
for key, engine in self._engines:
try:
for model_id in engine.list_models():
self._model_map[model_id] = engine
self._model_key_map[model_id] = key
except Exception as exc:
logger.debug("Failed to list models for %s: %s", _key, exc)
logger.debug("Failed to list models for %s: %s", key, exc)
_CLOUD_PREFIXES = ("gpt-", "o1-", "o3-", "o4-", "claude-", "gemini-", "openrouter/")
@@ -117,6 +120,14 @@ class MultiEngine(InferenceEngine):
self._refresh_map()
return list(self._model_map.keys())
def engine_key_for(self, model: str) -> str | None:
"""Return the registry key of the engine advertising *model*."""
key = self._model_key_map.get(model)
if key is not None:
return key
self._refresh_map()
return self._model_key_map.get(model)
def health(self) -> bool:
return any(engine.health() for _key, engine in self._engines)
+36 -8
View File
@@ -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
File diff suppressed because it is too large Load Diff
+119
View File
@@ -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
+62 -18
View File
@@ -336,6 +336,34 @@ def _remember_exchange(
)
def _engine_key_for_model(engine: Any, model: str) -> str | None:
"""Resolve the engine that advertised *model* through wrapper layers."""
from openjarvis.engine.multi import MultiEngine
from openjarvis.security.guardrails import GuardrailsEngine
from openjarvis.telemetry.instrumented_engine import InstrumentedEngine
current = engine
while current is not None:
if isinstance(current, MultiEngine):
return current.engine_key_for(model)
if isinstance(current, InstrumentedEngine):
current = current._inner
continue
if isinstance(current, GuardrailsEngine):
current = current._engine
continue
engine_id = getattr(current, "engine_id", None)
return engine_id if isinstance(engine_id, str) else None
return None
def _uses_direct_cloud_router(engine: Any, model: str) -> bool:
"""Whether *model* should bypass the configured engine for direct cloud."""
from openjarvis.server.cloud_router import is_cloud_model
return is_cloud_model(model) and _engine_key_for_model(engine, model) != "litellm"
def _handle_direct(
engine,
model: str,
@@ -541,12 +569,13 @@ async def _handle_stream_tools(
tool_calls) identical to the prior plain-stream behaviour, so this never
regresses non-tool-capable engines.
"""
from openjarvis.server.cloud_router import is_cloud_model
messages = _to_messages(req.messages)
messages = _ensure_identity_prompt(messages, app_config)
chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
use_cloud = is_cloud_model(model)
use_cloud = _uses_direct_cloud_router(engine, model)
telemetry_engine = (
"cloud" if use_cloud else (_engine_key_for_model(engine, model) or "ollama")
)
query_text = ""
for _m in reversed(req.messages):
if _m.role == "user" and _m.content:
@@ -626,7 +655,7 @@ async def _handle_stream_tools(
# Tag the finish chunk with the engine label, matching _handle_stream
# so UI/telemetry consumers see the same field on the tools path.
finish_dict.setdefault("telemetry", {})
finish_dict["telemetry"]["engine"] = "cloud" if use_cloud else "ollama"
finish_dict["telemetry"]["engine"] = telemetry_engine
if complexity_info is not None:
finish_dict["complexity"] = complexity_info.model_dump()
yield f"data: {_json.dumps(finish_dict)}\n\n"
@@ -668,11 +697,7 @@ async def _handle_stream(
"""
import time
from openjarvis.server.cloud_router import (
is_cloud_model,
stream_cloud,
stream_local,
)
from openjarvis.server.cloud_router import stream_cloud, stream_local
messages = _to_messages(req.messages)
messages = _ensure_identity_prompt(messages, app_config)
@@ -687,7 +712,10 @@ async def _handle_stream(
# Route directly to the right backend — bypasses engine routing entirely
# so broken MultiEngine state can never misdirect requests.
use_cloud = is_cloud_model(model)
use_cloud = _uses_direct_cloud_router(engine, model)
telemetry_engine = (
"cloud" if use_cloud else (_engine_key_for_model(engine, model) or "ollama")
)
async def generate():
started_at = time.time()
@@ -792,7 +820,7 @@ async def _handle_stream(
query=query_text,
result=full_content,
model=model,
engine="cloud" if use_cloud else "ollama",
engine=telemetry_engine,
started_at=started_at,
ended_at=time.time(),
)
@@ -825,7 +853,7 @@ async def _handle_stream(
# We use the routing decision (use_cloud) directly rather than
# unwrapping the engine chain, which can be in a broken state.
finish_dict.setdefault("telemetry", {})
finish_dict["telemetry"]["engine"] = "cloud" if use_cloud else "ollama"
finish_dict["telemetry"]["engine"] = telemetry_engine
if complexity_info is not None:
finish_dict["complexity"] = complexity_info.model_dump()
@@ -842,24 +870,40 @@ async def _handle_stream(
@router.get("/v1/models")
async def list_models(request: Request) -> ModelListResponse:
"""List locally installed models (Ollama).
"""List selectable engine models for the installed-model picker.
Cloud models are not included here they live in the Cloud Models tab
of the UI and are selected there, not from this endpoint.
Direct cloud models live in the Cloud Models tab. Models advertised by a
configured LiteLLM engine remain here because LiteLLM owns their routing
and may use provider-qualified IDs that resemble OpenRouter IDs.
"""
from openjarvis.server.cloud_router import is_cloud_model, list_local_models
# Prefer engine.list_models() so mock engines work in tests.
# Filter out any cloud model IDs that may appear via MultiEngine.
# Filter out direct-cloud model IDs that may appear via MultiEngine, but
# retain provider-qualified IDs owned by the configured LiteLLM engine.
# Fall back to direct Ollama query only when the engine returns nothing.
engine = request.app.state.engine
all_ids = await asyncio.to_thread(engine.list_models)
model_ids = [m for m in all_ids if not is_cloud_model(m)]
model_ids = [
m
for m in all_ids
if not is_cloud_model(m) or _engine_key_for_model(engine, m) == "litellm"
]
if not model_ids:
model_ids = await list_local_models()
return ModelListResponse(
data=[ModelObject(id=mid) for mid in model_ids],
data=[
ModelObject(
id=mid,
owned_by=(
"litellm"
if _engine_key_for_model(engine, mid) == "litellm"
else "openjarvis"
),
)
for mid in model_ids
],
)
+31 -1
View File
@@ -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
+3
View File
@@ -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:
@@ -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()
+471
View File
@@ -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()
+43
View File
@@ -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
+284
View File
@@ -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 == []
+14
View File
@@ -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))
+20
View File
@@ -8,10 +8,12 @@ from openjarvis.core.config import JarvisConfig
from openjarvis.core.registry import EngineRegistry
from openjarvis.engine._base import InferenceEngine
from openjarvis.engine._discovery import (
_make_engine,
discover_engines,
discover_models,
get_engine,
)
from openjarvis.engine.litellm import LiteLLMEngine
class _FakeEngine(InferenceEngine):
@@ -131,6 +133,24 @@ class TestDiscoverModels:
assert result == {"ollama": ["m1", "m2"], "vllm": ["m3"]}
class TestLiteLLMDiscovery:
def test_configured_default_model_is_advertised(self) -> None:
"""Regression for #713: discovery must configure LiteLLM's model.
LiteLLM cannot enumerate every model supported by every provider, so
``LiteLLMEngine.list_models()`` advertises the configured default
model. Dropping that value while constructing the engine leaves the
API and Web UI with an empty model list.
"""
cfg = JarvisConfig()
cfg.intelligence.default_model = "groq/llama-3.3-70b-versatile"
EngineRegistry.register_value("litellm", LiteLLMEngine)
engine = _make_engine("litellm", cfg)
assert engine.list_models() == ["groq/llama-3.3-70b-versatile"]
class TestGetEngine:
def test_fallback_when_default_unhealthy(self) -> None:
_reg("bad", "bad")
+3
View File
@@ -120,6 +120,9 @@ async def test_multi_routes_stream_full_by_model():
engine_b.list_models = lambda: ["model-b"]
multi = MultiEngine([("a", engine_a), ("b", engine_b)])
assert multi.engine_key_for("model-a") == "a"
assert multi.engine_key_for("model-b") == "b"
assert multi.engine_key_for("missing") is None
# Route to engine A
result_a = []
+104 -1
View File
@@ -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
+127
View File
@@ -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)
+28
View File
@@ -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 = {
+80
View File
@@ -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
@@ -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
@@ -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"
+299 -1
View File
@@ -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
+59 -1
View File
@@ -3,7 +3,7 @@
from __future__ import annotations
import json
from unittest.mock import MagicMock
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -883,6 +883,64 @@ class TestModelsEndpoint:
data = resp.json()
assert len(data["data"]) == 3
def test_configured_litellm_model_is_listed(self):
"""Regression for #713: LiteLLM models must reach the Web UI."""
model = "groq/llama-3.3-70b-versatile"
engine = _make_engine(models=[model])
engine.engine_id = "litellm"
app = create_app(
engine,
model,
engine_name="litellm",
config=_test_config(),
)
with patch(
"openjarvis.server.cloud_router.list_local_models",
new_callable=AsyncMock,
) as list_local_models:
list_local_models.return_value = []
client = TestClient(app)
resp = client.get("/v1/models")
assert resp.status_code == 200
assert [item["id"] for item in resp.json()["data"]] == [model]
assert resp.json()["data"][0]["owned_by"] == "litellm"
def test_litellm_provider_model_streams_through_active_engine(self):
"""A LiteLLM ``provider/model`` ID must not bypass its engine."""
model = "groq/llama-3.3-70b-versatile"
engine = _make_engine(models=[model])
engine.engine_id = "litellm"
app = create_app(
engine,
model,
engine_name="litellm",
config=_test_config(),
)
async def direct_cloud_tokens():
yield "wrong backend"
with patch(
"openjarvis.server.cloud_router.stream_cloud",
return_value=direct_cloud_tokens(),
) as stream_cloud:
client = TestClient(app)
resp = client.post(
"/v1/chat/completions",
json={
"model": model,
"messages": [{"role": "user", "content": "hello"}],
"stream": True,
},
)
assert resp.status_code == 200
stream_cloud.assert_not_called()
assert "Hello" in resp.text
assert '"engine": "litellm"' in resp.text
# ---------------------------------------------------------------------------
# Health endpoint tests