mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-08-14 08:52:06 +00:00
Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4efdb07dae | ||
|
|
7feda3dad3 | ||
|
|
f08ad574d3 | ||
|
|
1bfc25a860 | ||
|
|
063dd8ea75 | ||
|
|
6af9317556 |
@@ -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,
|
||||
|
||||
@@ -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>
|
||||
);
|
||||
})
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 []
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -158,7 +158,7 @@ def _show_toml_config(console: Console, config_path: Path) -> None:
|
||||
console.print(f"[dim]Loading config from: {config_path}[/dim]")
|
||||
|
||||
if config_path.exists():
|
||||
config_content = config_path.read_text()
|
||||
config_content = config_path.read_text(encoding="utf-8")
|
||||
syntax = Syntax(config_content, "toml", theme="monokai", line_numbers=True)
|
||||
console.print(Panel(syntax, title="Config File", border_style="cyan"))
|
||||
else:
|
||||
@@ -170,7 +170,7 @@ def _show_json_config(console: Console, config_path: Path) -> None:
|
||||
console.print(f"[dim]Loading config from: {config_path}[/dim]")
|
||||
|
||||
if config_path.exists():
|
||||
config_content = config_path.read_text()
|
||||
config_content = config_path.read_text(encoding="utf-8")
|
||||
|
||||
try:
|
||||
import tomllib # Python 3.11+
|
||||
@@ -375,7 +375,7 @@ def set_config(key: str, value: str) -> None:
|
||||
os.environ.get("OPENJARVIS_CONFIG", DEFAULT_CONFIG_DIR / "config.toml")
|
||||
)
|
||||
if config_path.exists():
|
||||
doc = tomlkit.parse(config_path.read_text())
|
||||
doc = tomlkit.parse(config_path.read_text(encoding="utf-8"))
|
||||
else:
|
||||
doc = tomlkit.document()
|
||||
config_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
@@ -390,7 +390,7 @@ def set_config(key: str, value: str) -> None:
|
||||
current[parts[-1]] = typed_value
|
||||
|
||||
# Write back
|
||||
config_path.write_text(tomlkit.dumps(doc))
|
||||
config_path.write_text(tomlkit.dumps(doc), encoding="utf-8")
|
||||
|
||||
console.print(f"[green]Set[/green] {key} = {value!r}")
|
||||
|
||||
|
||||
@@ -344,7 +344,9 @@ def init(
|
||||
console.print(f" Looked in: {examples_dir}")
|
||||
raise SystemExit(1)
|
||||
DEFAULT_CONFIG_DIR.mkdir(parents=True, exist_ok=True)
|
||||
DEFAULT_CONFIG_PATH.write_text(preset_path.read_text())
|
||||
DEFAULT_CONFIG_PATH.write_text(
|
||||
preset_path.read_text(encoding="utf-8"), encoding="utf-8"
|
||||
)
|
||||
console.print(
|
||||
f"[green]Preset '{preset}' installed to {DEFAULT_CONFIG_PATH}[/green]"
|
||||
)
|
||||
|
||||
+48
-47
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -9,6 +9,7 @@ import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
@@ -1305,6 +1306,160 @@ class CloudEngine(InferenceEngine):
|
||||
if chunk.text:
|
||||
yield chunk.text
|
||||
|
||||
async def _stream_full_google(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
*,
|
||||
model: str,
|
||||
temperature: float,
|
||||
max_tokens: int,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterator[StreamChunk]:
|
||||
"""Stream Google text and function-call parts as full chunks."""
|
||||
if self._google_client is None:
|
||||
raise EngineConnectionError("Google client not available")
|
||||
|
||||
system_text = ""
|
||||
contents: List[Dict[str, Any]] = []
|
||||
for message in messages:
|
||||
if message.role.value == "system":
|
||||
system_text = message.content
|
||||
elif message.role.value == "tool":
|
||||
function_response = {
|
||||
"function_response": {
|
||||
"name": message.name or "unknown",
|
||||
"response": {"result": message.content},
|
||||
}
|
||||
}
|
||||
if (
|
||||
contents
|
||||
and contents[-1]["role"] == "user"
|
||||
and contents[-1]["parts"]
|
||||
and "function_response" in contents[-1]["parts"][-1]
|
||||
):
|
||||
contents[-1]["parts"].append(function_response)
|
||||
else:
|
||||
contents.append({"role": "user", "parts": [function_response]})
|
||||
elif message.role.value == "assistant" and message.tool_calls:
|
||||
parts: List[Dict[str, Any]] = []
|
||||
if message.content:
|
||||
parts.append({"text": message.content})
|
||||
for tool_call in message.tool_calls:
|
||||
args = tool_call.arguments
|
||||
if isinstance(args, str):
|
||||
try:
|
||||
args = json.loads(args)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
args = {"input": args}
|
||||
function_call_part: Dict[str, Any] = {
|
||||
"function_call": {
|
||||
"name": tool_call.name,
|
||||
"args": args if isinstance(args, dict) else {},
|
||||
}
|
||||
}
|
||||
signature = self._thought_sigs.get(tool_call.id)
|
||||
if signature is not None:
|
||||
function_call_part["thought_signature"] = signature
|
||||
parts.append(function_call_part)
|
||||
contents.append({"role": "model", "parts": parts})
|
||||
elif message.role.value == "assistant":
|
||||
contents.append({"role": "model", "parts": [{"text": message.content}]})
|
||||
else:
|
||||
contents.append({"role": "user", "parts": [{"text": message.content}]})
|
||||
|
||||
from google.genai import types as genai_types
|
||||
|
||||
config = genai_types.GenerateContentConfig(
|
||||
temperature=temperature,
|
||||
max_output_tokens=max_tokens,
|
||||
)
|
||||
if system_text:
|
||||
config.system_instruction = system_text
|
||||
|
||||
tools = kwargs.pop("tools", None)
|
||||
if tools:
|
||||
config.tools = [{"function_declarations": _convert_tools_to_google(tools)}]
|
||||
|
||||
tool_call_count = 0
|
||||
stream_id = uuid.uuid4().hex
|
||||
final_usage: Dict[str, Any] | None = None
|
||||
for chunk in self._google_client.models.generate_content_stream(
|
||||
model=model,
|
||||
contents=contents,
|
||||
config=config,
|
||||
):
|
||||
usage_metadata = getattr(chunk, "usage_metadata", None)
|
||||
if usage_metadata is not None:
|
||||
prompt_tokens = getattr(usage_metadata, "prompt_token_count", 0) or 0
|
||||
completion_tokens = (
|
||||
getattr(usage_metadata, "candidates_token_count", 0) or 0
|
||||
)
|
||||
final_usage = {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
}
|
||||
|
||||
candidates = getattr(chunk, "candidates", None)
|
||||
parts = []
|
||||
if candidates:
|
||||
parts = getattr(candidates[0].content, "parts", []) or []
|
||||
|
||||
if parts:
|
||||
text_found = False
|
||||
calls: List[Dict[str, Any]] = []
|
||||
for part in parts:
|
||||
text = getattr(part, "text", None)
|
||||
if text:
|
||||
text_found = True
|
||||
yield StreamChunk(content=text)
|
||||
|
||||
function_call = getattr(part, "function_call", None)
|
||||
if function_call:
|
||||
name = getattr(function_call, "name", "")
|
||||
raw_args = getattr(function_call, "args", {})
|
||||
args = dict(raw_args) if hasattr(raw_args, "items") else {}
|
||||
# Gemini emits complete function-call parts, so each part is
|
||||
# a distinct invocation. The same function may legitimately
|
||||
# be called more than once in a parallel response.
|
||||
tool_index = tool_call_count
|
||||
# The engine is shared across server requests, and saved
|
||||
# thought signatures are keyed by tool-call ID. Include a
|
||||
# per-stream nonce so concurrent conversations cannot
|
||||
# overwrite each other's signatures.
|
||||
tool_id = f"google_{stream_id}_{tool_index}"
|
||||
tool_call_count += 1
|
||||
tool_call = {
|
||||
"index": tool_index,
|
||||
"id": tool_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"arguments": json.dumps(args),
|
||||
},
|
||||
}
|
||||
calls.append(tool_call)
|
||||
signature = getattr(part, "thought_signature", None)
|
||||
if signature is not None:
|
||||
tool_call["thought_signature"] = signature
|
||||
self._thought_sigs[tool_id] = signature
|
||||
if calls:
|
||||
yield StreamChunk(tool_calls=calls)
|
||||
if text_found:
|
||||
continue
|
||||
|
||||
try:
|
||||
text = chunk.text
|
||||
except (AttributeError, ValueError):
|
||||
text = None
|
||||
if text:
|
||||
yield StreamChunk(content=text)
|
||||
|
||||
yield StreamChunk(
|
||||
finish_reason="tool_calls" if tool_call_count else "stop",
|
||||
usage=final_usage,
|
||||
)
|
||||
|
||||
async def _stream_openrouter(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
@@ -1600,7 +1755,7 @@ class CloudEngine(InferenceEngine):
|
||||
async for chunk in self._stream_full_anthropic(messages, **kw):
|
||||
yield chunk
|
||||
elif _is_google_model(model):
|
||||
async for chunk in super().stream_full(messages, **kw):
|
||||
async for chunk in self._stream_full_google(messages, **kw):
|
||||
yield chunk
|
||||
else:
|
||||
async for chunk in self._stream_full_openai(messages, **kw):
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 == []
|
||||
@@ -138,6 +138,38 @@ class TestCLI:
|
||||
content = config_path.read_text()
|
||||
assert "[engine]" in content
|
||||
|
||||
def test_init_preset_uses_utf8_for_config_copy(self, tmp_path: Path) -> None:
|
||||
"""Preset installation reads and writes shipped TOML as UTF-8."""
|
||||
config_dir = tmp_path / ".openjarvis"
|
||||
config_path = config_dir / "config.toml"
|
||||
original_read_text = Path.read_text
|
||||
original_write_text = Path.write_text
|
||||
|
||||
def read_text(path: Path, *args: object, **kwargs: object) -> str:
|
||||
if path.name == "chat-simple.toml":
|
||||
assert kwargs.get("encoding") == "utf-8"
|
||||
return original_read_text(path, *args, **kwargs)
|
||||
|
||||
def write_text(path: Path, data: str, *args: object, **kwargs: object) -> int:
|
||||
if path == config_path:
|
||||
assert kwargs.get("encoding") == "utf-8"
|
||||
return original_write_text(path, data, *args, **kwargs)
|
||||
|
||||
with (
|
||||
mock.patch("openjarvis.cli.init_cmd.DEFAULT_CONFIG_DIR", config_dir),
|
||||
mock.patch("openjarvis.cli.init_cmd.DEFAULT_CONFIG_PATH", config_path),
|
||||
mock.patch.object(Path, "read_text", autospec=True, side_effect=read_text),
|
||||
mock.patch.object(
|
||||
Path, "write_text", autospec=True, side_effect=write_text
|
||||
),
|
||||
):
|
||||
result = CliRunner().invoke(cli, ["init", "--preset", "chat-simple"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "lightweight conversational AI" in config_path.read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
|
||||
|
||||
class TestStartupResilience:
|
||||
"""Importing the CLI must not force heavy/native deps (#404, #309).
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
@@ -109,18 +110,34 @@ temperature = 0.7
|
||||
except json.JSONDecodeError:
|
||||
pytest.fail(f"Output is not valid JSON: {result.output}")
|
||||
|
||||
def test_config_show_toml_displays_raw_content(self, tmp_path: Path) -> None:
|
||||
"""Test that config show toml displays the raw TOML content."""
|
||||
@pytest.mark.parametrize("output_format", ["toml", "json"])
|
||||
def test_config_show_uses_utf8_for_config_file(
|
||||
self, tmp_path: Path, output_format: str
|
||||
) -> None:
|
||||
"""Test that config show reads UTF-8 config files explicitly."""
|
||||
# Create a temporary config file
|
||||
config_file = tmp_path / "test_config.toml"
|
||||
config_file.write_text('[engine]\ndefault = "ollama"\n')
|
||||
|
||||
result = CliRunner().invoke(
|
||||
cli, ["config", "show", "toml", "--path", str(config_file)]
|
||||
config_file.write_text(
|
||||
'# Preset comment — stored as UTF-8\n[engine]\ndefault = "ollama"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
original_read_text = Path.read_text
|
||||
|
||||
def read_text(path: Path, *args: object, **kwargs: object) -> str:
|
||||
if path == config_file:
|
||||
assert kwargs.get("encoding") == "utf-8"
|
||||
return original_read_text(path, *args, **kwargs)
|
||||
|
||||
with mock.patch.object(Path, "read_text", autospec=True, side_effect=read_text):
|
||||
result = CliRunner().invoke(
|
||||
cli, ["config", "show", output_format, "--path", str(config_file)]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "[engine]" in result.output
|
||||
if output_format == "toml":
|
||||
assert "[engine]" in result.output
|
||||
else:
|
||||
assert '"engine"' in result.output
|
||||
assert "ollama" in result.output
|
||||
|
||||
def test_config_show_json_displays_parsed_content(self, tmp_path: Path) -> None:
|
||||
|
||||
@@ -60,6 +60,42 @@ class TestConfigSet:
|
||||
assert "vllm" in content
|
||||
assert "qwen2.5:3b" in content
|
||||
|
||||
def test_set_uses_utf8_for_existing_config(self, tmp_path: Path) -> None:
|
||||
"""config set preserves a UTF-8 config regardless of the system locale."""
|
||||
config_file = tmp_path / "config.toml"
|
||||
config_file.write_text(
|
||||
'# Preset comment — stored as UTF-8\n[engine]\ndefault = "ollama"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
original_read_text = Path.read_text
|
||||
original_write_text = Path.write_text
|
||||
|
||||
def read_text(path: Path, *args: object, **kwargs: object) -> str:
|
||||
if path == config_file:
|
||||
assert kwargs.get("encoding") == "utf-8"
|
||||
return original_read_text(path, *args, **kwargs)
|
||||
|
||||
def write_text(path: Path, data: str, *args: object, **kwargs: object) -> int:
|
||||
if path == config_file:
|
||||
assert kwargs.get("encoding") == "utf-8"
|
||||
return original_write_text(path, data, *args, **kwargs)
|
||||
|
||||
with (
|
||||
mock.patch.dict(os.environ, {"OPENJARVIS_CONFIG": str(config_file)}),
|
||||
mock.patch.object(Path, "read_text", autospec=True, side_effect=read_text),
|
||||
mock.patch.object(
|
||||
Path, "write_text", autospec=True, side_effect=write_text
|
||||
),
|
||||
):
|
||||
result = CliRunner().invoke(
|
||||
cli, ["config", "set", "engine.default", "vllm"]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
content = config_file.read_text(encoding="utf-8")
|
||||
assert "Preset comment — stored as UTF-8" in content
|
||||
assert "vllm" in content
|
||||
|
||||
def test_set_invalid_key_rejected(self, tmp_path: Path) -> None:
|
||||
"""config set rejects unknown keys."""
|
||||
config_file = tmp_path / "config.toml"
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -3,6 +3,8 @@ and _prepare_anthropic_messages."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from typing import Any, List
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
@@ -63,6 +65,36 @@ def _openai_tool_call_delta(
|
||||
return tc
|
||||
|
||||
|
||||
class _GoogleConfig:
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
self.__dict__.update(kwargs)
|
||||
|
||||
|
||||
def _google_stream_chunk(
|
||||
*parts: Any,
|
||||
text: str | None = None,
|
||||
usage_metadata: Any = None,
|
||||
) -> Any:
|
||||
candidates = []
|
||||
if parts:
|
||||
candidates = [SimpleNamespace(content=SimpleNamespace(parts=list(parts)))]
|
||||
return SimpleNamespace(
|
||||
text=text,
|
||||
candidates=candidates,
|
||||
usage_metadata=usage_metadata,
|
||||
)
|
||||
|
||||
|
||||
def _google_types_modules() -> dict[str, ModuleType]:
|
||||
types = ModuleType("google.genai.types")
|
||||
types.GenerateContentConfig = _GoogleConfig
|
||||
genai = ModuleType("google.genai")
|
||||
genai.types = types
|
||||
google = ModuleType("google")
|
||||
google.genai = genai
|
||||
return {"google": google, "google.genai": genai, "google.genai.types": types}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _stream_full_openai tests
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -413,6 +445,316 @@ def test_prepare_anthropic_messages_tool_calls():
|
||||
assert blocks[1]["input"] == {"city": "Berlin"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _stream_full_google tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_text_only(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Google text chunks retain their content and finish normally."""
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.return_value = iter(
|
||||
[_google_stream_chunk(text="Hello"), _google_stream_chunk(text=" world")]
|
||||
)
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {}
|
||||
messages = [Message(role=Role.USER, content="hi")]
|
||||
modules = _google_types_modules()
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in modules.items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(messages, model="gemini-2.5-flash")
|
||||
]
|
||||
|
||||
assert [chunk.content for chunk in result[:-1]] == ["Hello", " world"]
|
||||
assert result[-1].finish_reason == "stop"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_preserves_tool_calls(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Google function_call parts become OpenAI-compatible tool call chunks."""
|
||||
function_call = SimpleNamespace(name="get_weather", args={"city": "Berlin"})
|
||||
part = SimpleNamespace(
|
||||
function_call=function_call, text=None, thought_signature=b"sig"
|
||||
)
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.return_value = iter(
|
||||
[_google_stream_chunk(part)]
|
||||
)
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {}
|
||||
messages = [Message(role=Role.USER, content="weather")]
|
||||
modules = _google_types_modules()
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in modules.items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(
|
||||
messages,
|
||||
model="gemini-2.5-flash",
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
]
|
||||
|
||||
tool_call = result[0].tool_calls[0]
|
||||
assert tool_call["index"] == 0
|
||||
assert tool_call["id"].startswith("google_")
|
||||
assert tool_call["type"] == "function"
|
||||
assert tool_call["function"] == {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Berlin"}',
|
||||
}
|
||||
assert tool_call["thought_signature"] == b"sig"
|
||||
assert engine._thought_sigs[tool_call["id"]] == b"sig"
|
||||
assert result[-1].finish_reason == "tool_calls"
|
||||
config = client.models.generate_content_stream.call_args.kwargs["config"]
|
||||
assert config.tools == [
|
||||
{
|
||||
"function_declarations": [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_preserves_mixed_and_multiple_calls(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Google streams retain mixed text and multiple tool calls."""
|
||||
weather = SimpleNamespace(name="get_weather", args={"city": "Berlin"})
|
||||
calendar = SimpleNamespace(name="get_calendar", args={"day": "Monday"})
|
||||
text_part = SimpleNamespace(text="I'll check.", function_call=None)
|
||||
weather_part = SimpleNamespace(
|
||||
function_call=weather, text=None, thought_signature=None
|
||||
)
|
||||
calendar_part = SimpleNamespace(
|
||||
function_call=calendar, text=None, thought_signature=None
|
||||
)
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.return_value = iter(
|
||||
[
|
||||
_google_stream_chunk(text_part, weather_part),
|
||||
_google_stream_chunk(calendar_part),
|
||||
]
|
||||
)
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {}
|
||||
modules = _google_types_modules()
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in modules.items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(
|
||||
[Message(role=Role.USER, content="plan")], model="gemini-2.5-flash"
|
||||
)
|
||||
]
|
||||
|
||||
assert result[0].content == "I'll check."
|
||||
weather_call = result[1].tool_calls[0]
|
||||
calendar_call = result[2].tool_calls[0]
|
||||
assert weather_call["index"] == 0
|
||||
assert weather_call["function"] == {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Berlin"}',
|
||||
}
|
||||
assert calendar_call["index"] == 1
|
||||
assert calendar_call["function"] == {
|
||||
"name": "get_calendar",
|
||||
"arguments": '{"day": "Monday"}',
|
||||
}
|
||||
assert weather_call["id"] != calendar_call["id"]
|
||||
assert result[-1].finish_reason == "tool_calls"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_keeps_parallel_same_name_calls_distinct(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Parallel invocations of one function receive unique indexes and IDs."""
|
||||
paris = SimpleNamespace(name="get_weather", args={"city": "Paris"})
|
||||
london = SimpleNamespace(name="get_weather", args={"city": "London"})
|
||||
parts = [
|
||||
SimpleNamespace(function_call=paris, text=None, thought_signature=b"sig"),
|
||||
SimpleNamespace(function_call=london, text=None, thought_signature=None),
|
||||
]
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.return_value = iter(
|
||||
[_google_stream_chunk(*parts)]
|
||||
)
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {}
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in _google_types_modules().items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(
|
||||
[Message(role=Role.USER, content="Weather in Paris and London")],
|
||||
model="gemini-3-flash-preview",
|
||||
)
|
||||
]
|
||||
|
||||
calls = result[0].tool_calls
|
||||
assert [call["index"] for call in calls] == [0, 1]
|
||||
assert calls[0]["id"] != calls[1]["id"]
|
||||
assert [call["function"]["arguments"] for call in calls] == [
|
||||
'{"city": "Paris"}',
|
||||
'{"city": "London"}',
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_ids_are_unique_across_requests(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Shared engines keep signatures isolated between conversations."""
|
||||
first_part = SimpleNamespace(
|
||||
function_call=SimpleNamespace(name="get_weather", args={"city": "Paris"}),
|
||||
text=None,
|
||||
thought_signature=b"paris-sig",
|
||||
)
|
||||
second_part = SimpleNamespace(
|
||||
function_call=SimpleNamespace(name="get_weather", args={"city": "London"}),
|
||||
text=None,
|
||||
thought_signature=b"london-sig",
|
||||
)
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.side_effect = [
|
||||
iter([_google_stream_chunk(first_part)]),
|
||||
iter([_google_stream_chunk(second_part)]),
|
||||
]
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {}
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in _google_types_modules().items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
first = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(
|
||||
[Message(role=Role.USER, content="Weather in Paris")],
|
||||
model="gemini-3-flash-preview",
|
||||
)
|
||||
]
|
||||
second = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(
|
||||
[Message(role=Role.USER, content="Weather in London")],
|
||||
model="gemini-3-flash-preview",
|
||||
)
|
||||
]
|
||||
|
||||
first_id = first[0].tool_calls[0]["id"]
|
||||
second_id = second[0].tool_calls[0]["id"]
|
||||
assert first_id != second_id
|
||||
assert engine._thought_sigs[first_id] == b"paris-sig"
|
||||
assert engine._thought_sigs[second_id] == b"london-sig"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_emits_final_usage(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Google's final usage metadata is normalized onto the terminal chunk."""
|
||||
usage = SimpleNamespace(prompt_token_count=12, candidates_token_count=5)
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.return_value = iter(
|
||||
[
|
||||
_google_stream_chunk(text="Hello"),
|
||||
_google_stream_chunk(usage_metadata=usage),
|
||||
]
|
||||
)
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {}
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in _google_types_modules().items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(
|
||||
[Message(role=Role.USER, content="hi")],
|
||||
model="gemini-2.5-flash",
|
||||
)
|
||||
]
|
||||
|
||||
assert result[-1].finish_reason == "stop"
|
||||
assert result[-1].usage == {
|
||||
"prompt_tokens": 12,
|
||||
"completion_tokens": 5,
|
||||
"total_tokens": 17,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_replays_signature_on_part(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""A saved Gemini signature is replayed beside, not inside, function_call."""
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.return_value = iter([])
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {"google_get_weather_0": b"sig"}
|
||||
messages = [
|
||||
Message(role=Role.USER, content="weather"),
|
||||
Message(
|
||||
role=Role.ASSISTANT,
|
||||
content=None,
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
id="google_get_weather_0",
|
||||
name="get_weather",
|
||||
arguments='{"city": "Berlin"}',
|
||||
)
|
||||
],
|
||||
),
|
||||
Message(role=Role.TOOL, name="get_weather", content='{"temp": 20}'),
|
||||
]
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in _google_types_modules().items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(
|
||||
messages, model="gemini-3-flash-preview"
|
||||
)
|
||||
]
|
||||
|
||||
contents = client.models.generate_content_stream.call_args.kwargs["contents"]
|
||||
assert contents[1]["parts"] == [
|
||||
{
|
||||
"function_call": {
|
||||
"name": "get_weather",
|
||||
"args": {"city": "Berlin"},
|
||||
},
|
||||
"thought_signature": b"sig",
|
||||
}
|
||||
]
|
||||
assert result[-1].finish_reason == "stop"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stream_full routing tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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"
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user