feat: add SSE streaming support for managed agent messages

Add `stream: bool` parameter to `POST /v1/managed-agents/{id}/messages`.
When `stream=true`, the agent processes the message synchronously and
returns an SSE stream (OpenAI-compatible format) with token-by-token
response, tool result events, and usage metadata.

This enables real-time voice assistants and chat UIs to receive agent
responses as they are generated, rather than polling for completion.

- Extend SendMessageRequest with `stream` field (default: false)
- Add _stream_managed_agent() helper using asyncio.to_thread()
- Build AgentContext from conversation history for multi-turn support
- Emit tool_results as named SSE events
- Persist agent response in DB after streaming completes
- Add 6 new tests covering streaming behavior
- Update agents.md documentation with streaming examples
This commit is contained in:
manuel.richarz
2026-03-24 14:56:58 +01:00
parent 16f601c7d7
commit b7ad2d2c12
3 changed files with 410 additions and 2 deletions
+46
View File
@@ -621,3 +621,49 @@ All agents publish events on the `EventBus` when a bus is provided:
`INFERENCE_START` / `INFERENCE_END` events are published by the `InstrumentedEngine` wrapper, not by agents directly. This keeps telemetry opt-in and transparent to agent code.
These events enable the telemetry and trace systems to record detailed interaction data automatically.
---
## Managed Agent Streaming
The Managed Agent API (`/v1/managed-agents/{id}/messages`) supports real-time SSE streaming. Send a message with `stream: true` to receive the agent's response as a Server-Sent Events stream instead of the default asynchronous queue mode.
### Streaming Messages
```bash
curl -N -X POST http://localhost:8000/v1/managed-agents/{id}/messages \
-H "Content-Type: application/json" \
-d '{"content": "What is 2+2?", "stream": true}'
```
The response follows the OpenAI SSE format:
1. **Content chunks** -- `data: {"choices": [{"delta": {"content": "token"}}]}`
2. **Tool results** (if the agent used tools) -- `event: tool_results\ndata: {"results": [...]}`
3. **Final chunk** -- `data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}`
4. **Done sentinel** -- `data: [DONE]`
When `stream: false` (the default), the endpoint behaves exactly as before -- the message is queued and the agent must be triggered separately via `/run`.
### Behavior Details
- The user message is always stored in the database before the agent runs.
- After streaming completes, the full agent response is persisted as an `agent_to_user` message.
- The agent is instantiated from the managed agent's stored `agent_type` and `config`.
- Conversation history from prior messages is automatically loaded as context.
- If the engine is not available on the server, a `503` error is returned.
### Python Example
```python
import httpx
with httpx.stream(
"POST",
"http://localhost:8000/v1/managed-agents/{id}/messages",
json={"content": "Summarize today's news", "stream": True},
) as response:
for line in response.iter_lines():
if line.startswith("data:") and "[DONE]" not in line:
print(line[5:].strip())
```
+197 -2
View File
@@ -2,16 +2,20 @@
from __future__ import annotations
import logging
from typing import Any, Dict, List, Optional, Tuple
from openjarvis.agents.manager import AgentManager
try:
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
except ImportError:
raise ImportError("fastapi and pydantic are required for server routes")
logger = logging.getLogger("openjarvis.server.agent_manager")
class CreateAgentRequest(BaseModel):
name: str
@@ -46,6 +50,7 @@ class BindChannelRequest(BaseModel):
class SendMessageRequest(BaseModel):
content: str
mode: str = "queued"
stream: bool = False # SSE streaming mode
class FeedbackRequest(BaseModel):
@@ -202,6 +207,171 @@ def build_tools_list() -> List[Dict[str, Any]]:
return items
async def _stream_managed_agent(
*,
manager: AgentManager,
agent_record: Dict[str, Any],
user_content: str,
message_id: str,
engine: Any,
bus: Any,
) -> StreamingResponse:
"""Run a managed agent and stream the response as SSE.
Instantiates the agent from its stored config, builds conversation
context from message history, executes the agent in a background
thread, and yields SSE-formatted chunks. After completion the
full response is persisted via ``manager.store_agent_response()``.
"""
import asyncio
import json
import uuid
from openjarvis.agents._stubs import AgentContext
from openjarvis.core.registry import AgentRegistry
from openjarvis.core.types import Message, Role
agent_id = agent_record["id"]
config = agent_record.get("config", {})
agent_type = agent_record.get("agent_type", "orchestrator")
model = config.get("model", getattr(engine, "_model", ""))
# Resolve the agent class from registry
agent_cls = AgentRegistry.get(agent_type)
if agent_cls is None:
# Fallback to orchestrator if the type is not registered
agent_cls = AgentRegistry.get("orchestrator")
if agent_cls is None:
raise HTTPException(
status_code=500, detail=f"Agent type '{agent_type}' not found in registry",
)
# Build agent constructor kwargs from config
agent_kwargs: Dict[str, Any] = {
"engine": engine,
"model": model,
}
if bus is not None:
agent_kwargs["bus"] = bus
if config.get("system_prompt"):
agent_kwargs["system_prompt"] = config["system_prompt"]
if config.get("temperature") is not None:
agent_kwargs["temperature"] = config["temperature"]
if config.get("max_tokens") is not None:
agent_kwargs["max_tokens"] = config["max_tokens"]
if config.get("max_turns") is not None:
agent_kwargs["max_turns"] = config["max_turns"]
try:
agent = agent_cls(**agent_kwargs)
except TypeError as exc:
logger.warning("Agent instantiation failed with all kwargs, retrying minimal: %s", exc)
agent = agent_cls(engine=engine, model=model)
# Build conversation context from existing messages
ctx = AgentContext()
messages = manager.list_messages(agent_id, limit=50)
# Messages come in DESC order, reverse for chronological
for m in reversed(messages):
# Skip the message we just stored (it will be the input)
if m["id"] == message_id:
continue
if m["direction"] == "user_to_agent":
ctx.conversation.add(Message(role=Role.USER, content=m["content"]))
elif m["direction"] == "agent_to_user":
ctx.conversation.add(Message(role=Role.ASSISTANT, content=m["content"]))
# Mark the user message as delivered
manager.mark_message_delivered(message_id)
chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
async def generate():
"""Async generator yielding SSE-formatted chunks."""
collected_content = ""
# Run agent.run() in a background thread
try:
result = await asyncio.to_thread(agent.run, user_content, context=ctx)
except Exception as exc:
logger.error("Managed agent stream error: %s", exc, exc_info=True)
error_data = {
"id": chunk_id,
"object": "chat.completion.chunk",
"model": model,
"choices": [{
"index": 0,
"delta": {"content": f"Error: {exc}"},
"finish_reason": "stop",
}],
}
yield f"data: {json.dumps(error_data)}\n\n"
yield "data: [DONE]\n\n"
return
content = result.content or ""
collected_content = content
# Emit tool results metadata if any
if result.tool_results:
tool_data = []
for tr in result.tool_results:
tool_data.append({
"tool_name": tr.tool_name,
"success": tr.success,
"output": tr.content,
"latency_ms": tr.latency_seconds * 1000,
})
yield f"event: tool_results\ndata: {json.dumps({'results': tool_data})}\n\n"
# Stream content word-by-word for real-time feel
if content:
words = content.split(" ")
for i, word in enumerate(words):
token = word if i == 0 else " " + word
chunk_data = {
"id": chunk_id,
"object": "chat.completion.chunk",
"model": model,
"choices": [{
"index": 0,
"delta": {"content": token},
"finish_reason": None,
}],
}
yield f"data: {json.dumps(chunk_data)}\n\n"
await asyncio.sleep(0.012)
# Final chunk with finish_reason
final_data = {
"id": chunk_id,
"object": "chat.completion.chunk",
"model": model,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "stop",
}],
}
yield f"data: {json.dumps(final_data)}\n\n"
yield "data: [DONE]\n\n"
# Persist agent response in DB after streaming completes
if collected_content:
try:
manager.store_agent_response(agent_id, collected_content)
except Exception as store_exc:
logger.error(
"Failed to store agent response: %s", store_exc, exc_info=True,
)
return StreamingResponse(
generate(),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
)
def create_agent_manager_router(
manager: AgentManager,
) -> Tuple[APIRouter, APIRouter, APIRouter, APIRouter]:
@@ -399,9 +569,34 @@ def create_agent_manager_router(
return {"messages": manager.list_messages(agent_id)}
@agents_router.post("/{agent_id}/messages")
def send_message(agent_id: str, req: SendMessageRequest):
async def send_message(agent_id: str, req: SendMessageRequest, request: Request):
agent_record = manager.get_agent(agent_id)
if not agent_record:
raise HTTPException(status_code=404, detail="Agent not found")
# Store user message in DB (always, regardless of stream mode)
msg = manager.send_message(agent_id, req.content, mode=req.mode)
return msg
if not req.stream:
return msg
# --- Streaming mode: run agent and return SSE response ---
engine = getattr(request.app.state, "engine", None)
bus = getattr(request.app.state, "bus", None)
if engine is None:
raise HTTPException(
status_code=503,
detail="Engine not available for streaming",
)
return await _stream_managed_agent(
manager=manager,
agent_record=agent_record,
user_content=req.content,
message_id=msg["id"],
engine=engine,
bus=bus,
)
# ── State inspection ─────────────────────────────────────
+167
View File
@@ -2,8 +2,10 @@
from __future__ import annotations
import json
import tempfile
from pathlib import Path
from unittest.mock import MagicMock
import pytest
@@ -200,3 +202,168 @@ class TestAgentManagerRoutes:
assert "channels" in state
assert "messages" in state
assert "checkpoint" in state
def test_send_message_non_stream_unchanged(self, manager, client):
"""stream=False (default) returns a normal JSON message, not SSE."""
agent = manager.create_agent(name="basic", agent_type="simple")
res = client.post(
f"/v1/managed-agents/{agent['id']}/messages",
json={"content": "hello", "stream": False},
)
assert res.status_code == 200
data = res.json()
assert data["content"] == "hello"
assert data["direction"] == "user_to_agent"
def test_send_message_stream_not_found(self, manager, client):
"""Streaming to a non-existent agent returns 404."""
res = client.post(
"/v1/managed-agents/nonexistent/messages",
json={"content": "hello", "stream": True},
)
assert res.status_code == 404
@pytest.mark.skipif(not HAS_FASTAPI, reason="fastapi not installed")
class TestAgentManagerStreaming:
"""Tests for the SSE streaming mode of the managed-agent messages endpoint."""
@pytest.fixture
def _mock_engine(self):
engine = MagicMock()
engine.engine_id = "mock"
engine._model = "test-model"
engine.health.return_value = True
return engine
@pytest.fixture
def _mock_agent_cls(self):
"""Register a mock agent class in the AgentRegistry for testing."""
from openjarvis.agents._stubs import AgentResult
from openjarvis.core.registry import AgentRegistry
class _MockStreamAgent:
agent_id = "mock_stream"
def __init__(self, engine, model, **kwargs):
self._engine = engine
self._model = model
def run(self, input_text, context=None, **kwargs):
return AgentResult(content=f"Echo: {input_text}", turns=1)
# Register under a unique key for test isolation
AgentRegistry._entries()["_test_stream"] = _MockStreamAgent
yield _MockStreamAgent
AgentRegistry._entries().pop("_test_stream", None)
@pytest.fixture
def stream_client(self, manager, _mock_engine, _mock_agent_cls):
from fastapi import FastAPI
from openjarvis.server.agent_manager_routes import create_agent_manager_router
app = FastAPI()
app.state.engine = _mock_engine
app.state.bus = None
routers = create_agent_manager_router(manager)
agents_router, templates_router, global_router, tools_router = routers
app.include_router(agents_router)
app.include_router(templates_router)
app.include_router(global_router)
app.include_router(tools_router)
return TestClient(app)
def test_send_message_stream(self, manager, stream_client, _mock_agent_cls):
"""Test streaming mode returns SSE response with [DONE] sentinel."""
agent = manager.create_agent(
name="streamer", agent_type="_test_stream",
)
resp = stream_client.post(
f"/v1/managed-agents/{agent['id']}/messages",
json={"content": "What is 2+2?", "stream": True},
)
assert resp.status_code == 200
assert "text/event-stream" in resp.headers.get("content-type", "")
# Parse SSE events
lines = resp.text.strip().split("\n")
data_lines = [ln for ln in lines if ln.startswith("data:")]
assert len(data_lines) > 0
# Last data line must be [DONE]
assert data_lines[-1].strip() == "data: [DONE]"
def test_send_message_stream_content(self, manager, stream_client, _mock_agent_cls):
"""Test streaming returns the correct agent response content."""
agent = manager.create_agent(
name="streamer2", agent_type="_test_stream",
)
resp = stream_client.post(
f"/v1/managed-agents/{agent['id']}/messages",
json={"content": "Hello world", "stream": True},
)
assert resp.status_code == 200
# Collect content tokens from stream
content = ""
for line in resp.text.strip().split("\n"):
if line.startswith("data:") and "[DONE]" not in line:
raw = line[5:].strip()
try:
data = json.loads(raw)
except json.JSONDecodeError:
continue
choices = data.get("choices", [{}])
delta_content = choices[0].get("delta", {}).get("content")
if delta_content:
content += delta_content
assert content == "Echo: Hello world"
def test_send_message_stream_stores_response(
self, manager, stream_client, _mock_agent_cls,
):
"""After streaming, agent response is persisted in the DB."""
agent = manager.create_agent(
name="streamer3", agent_type="_test_stream",
)
resp = stream_client.post(
f"/v1/managed-agents/{agent['id']}/messages",
json={"content": "persist me", "stream": True},
)
assert resp.status_code == 200
# Check messages in DB
messages = manager.list_messages(agent["id"])
# Should have both the user message and the agent response
assert len(messages) == 2
directions = {m["direction"] for m in messages}
assert "user_to_agent" in directions
assert "agent_to_user" in directions
agent_msg = next(m for m in messages if m["direction"] == "agent_to_user")
assert "persist me" in agent_msg["content"]
def test_send_message_stream_finish_reason(
self, manager, stream_client, _mock_agent_cls,
):
"""The final chunk before [DONE] has finish_reason='stop'."""
agent = manager.create_agent(
name="streamer4", agent_type="_test_stream",
)
resp = stream_client.post(
f"/v1/managed-agents/{agent['id']}/messages",
json={"content": "check finish", "stream": True},
)
# Collect all data chunks (excluding [DONE])
chunks = []
for line in resp.text.strip().split("\n"):
if line.startswith("data:") and "[DONE]" not in line:
raw = line[5:].strip()
try:
chunks.append(json.loads(raw))
except json.JSONDecodeError:
continue
# Last chunk should have finish_reason="stop"
assert chunks[-1]["choices"][0]["finish_reason"] == "stop"