mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 05:12:26 +00:00
b39dbedcc4
Rebased and cleaned-up version of PR #113 by @mricharz, resolved against current main (including Codex engine, Gemini thought_signature, and agent manager fixes merged since the original PR). MCP Transport & Client: - StreamableHTTPTransport with session tracking, SSE parsing, timeouts - MCPClient.initialize() sends proper MCP handshake (protocolVersion, capabilities, clientInfo) + notifications/initialized - Fix StdioTransport constructor: command=[command] + args - MCPRequest.to_dict() with notification support (id=None) External MCP Discovery: - _discover_external_mcp supports both url (HTTP) and command (stdio) - Per-server include_tools / exclude_tools filtering - MCP clients persisted on JarvisSystem for runtime lifetime Streaming Tool-Call Support (stream_full): - StreamChunk dataclass in _stubs.py (content, tool_calls, finish_reason, usage) - Default stream_full() on InferenceEngine ABC wraps stream() for backward compat - _OpenAICompatibleEngine.stream_full() with SSE parsing - CloudEngine: _stream_full_openai (OpenAI/OpenRouter/MiniMax/Codex routing) _stream_full_anthropic (event-based → OpenAI delta format) - InstrumentedEngine, MultiEngine: stream_full delegation - GuardrailsEngine: stream_full with post-hoc security scanning (FIXED: original PR bypassed output scanning — now accumulates and scans like stream()) Other improvements: - _prepare_anthropic_messages() extracted to eliminate duplication - Default tool_choice=auto when tools are provided (OpenAI compat engines) - MCP tool injection into managed agent streaming path - Documentation: docs/user-guide/mcp-external-servers.md Tests: ~59 new tests across 8 test files, all passing. Closes PR #113 Co-Authored-By: mricharz <mricharz@users.noreply.github.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
143 lines
4.4 KiB
Python
143 lines
4.4 KiB
Python
"""Tests for InstrumentedEngine, GuardrailsEngine, and MultiEngine stream_full
|
|
delegation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any, Dict, List
|
|
|
|
import pytest
|
|
|
|
from openjarvis.core.events import EventBus
|
|
from openjarvis.core.types import Message, Role
|
|
from openjarvis.engine._stubs import InferenceEngine, StreamChunk
|
|
from openjarvis.engine.multi import MultiEngine
|
|
from openjarvis.security.guardrails import GuardrailsEngine
|
|
from openjarvis.telemetry.instrumented_engine import InstrumentedEngine
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fake engine that yields predetermined StreamChunks via stream_full
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeStreamFullEngine(InferenceEngine):
|
|
engine_id = "fake-sf"
|
|
|
|
def __init__(self, chunks: list[StreamChunk]) -> None:
|
|
self._chunks = chunks
|
|
|
|
def generate(self, messages, *, model, **kwargs) -> Dict[str, Any]:
|
|
return {"content": "ok", "usage": {}}
|
|
|
|
async def stream(self, messages, *, model, **kwargs) -> AsyncIterator[str]:
|
|
yield "ok"
|
|
|
|
async def stream_full(
|
|
self, messages, *, model, **kwargs
|
|
) -> AsyncIterator[StreamChunk]:
|
|
for c in self._chunks:
|
|
yield c
|
|
|
|
def list_models(self) -> List[str]:
|
|
return ["fake-model"]
|
|
|
|
def health(self) -> bool:
|
|
return True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# InstrumentedEngine.stream_full delegation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_instrumented_delegates_stream_full():
|
|
"""InstrumentedEngine.stream_full delegates to inner engine."""
|
|
expected = [
|
|
StreamChunk(content="Hello"),
|
|
StreamChunk(content=" world"),
|
|
StreamChunk(finish_reason="stop"),
|
|
]
|
|
inner = _FakeStreamFullEngine(expected)
|
|
bus = EventBus(record_history=True)
|
|
engine = InstrumentedEngine(inner, bus)
|
|
|
|
result = []
|
|
async for chunk in engine.stream_full(
|
|
[Message(role=Role.USER, content="test")],
|
|
model="fake-model",
|
|
):
|
|
result.append(chunk)
|
|
|
|
assert len(result) == 3
|
|
assert result[0].content == "Hello"
|
|
assert result[1].content == " world"
|
|
assert result[2].finish_reason == "stop"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GuardrailsEngine.stream_full delegation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_guardrails_delegates_stream_full():
|
|
"""GuardrailsEngine.stream_full delegates to wrapped engine."""
|
|
expected = [
|
|
StreamChunk(content="safe output"),
|
|
StreamChunk(finish_reason="stop"),
|
|
]
|
|
inner = _FakeStreamFullEngine(expected)
|
|
engine = GuardrailsEngine(inner, scanners=[])
|
|
|
|
result = []
|
|
async for chunk in engine.stream_full(
|
|
[Message(role=Role.USER, content="test")],
|
|
model="fake-model",
|
|
):
|
|
result.append(chunk)
|
|
|
|
assert len(result) == 2
|
|
assert result[0].content == "safe output"
|
|
assert result[1].finish_reason == "stop"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MultiEngine.stream_full routing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multi_routes_stream_full_by_model():
|
|
"""MultiEngine routes stream_full to the correct engine by model name."""
|
|
chunks_a = [StreamChunk(content="from A"), StreamChunk(finish_reason="stop")]
|
|
chunks_b = [StreamChunk(content="from B"), StreamChunk(finish_reason="stop")]
|
|
|
|
engine_a = _FakeStreamFullEngine(chunks_a)
|
|
engine_a.list_models = lambda: ["model-a"]
|
|
|
|
engine_b = _FakeStreamFullEngine(chunks_b)
|
|
engine_b.list_models = lambda: ["model-b"]
|
|
|
|
multi = MultiEngine([("a", engine_a), ("b", engine_b)])
|
|
|
|
# Route to engine A
|
|
result_a = []
|
|
async for chunk in multi.stream_full(
|
|
[Message(role=Role.USER, content="test")],
|
|
model="model-a",
|
|
):
|
|
result_a.append(chunk)
|
|
|
|
assert result_a[0].content == "from A"
|
|
|
|
# Route to engine B
|
|
result_b = []
|
|
async for chunk in multi.stream_full(
|
|
[Message(role=Role.USER, content="test")],
|
|
model="model-b",
|
|
):
|
|
result_b.append(chunk)
|
|
|
|
assert result_b[0].content == "from B"
|