fix: reliably dispatch and return background tasks
Rebuild & Deploy / rebuild (push) Successful in 47s
Rebuild & Deploy / rebuild (push) Successful in 47s
This commit is contained in:
1 parent
746ca8dca9
commit
321421e37f
7 files changed
+307
-24
No files matched your search
@@ -45,6 +45,7 @@ COPY agent/weather_mcp.py /opt/voice-agent/weather_mcp.py
|
|||||||
|
|
||||||
COPY agent/skills_mcp.py /opt/voice-agent/skills_mcp.py
|
COPY agent/skills_mcp.py /opt/voice-agent/skills_mcp.py
|
||||||
COPY agent/task_registry.py /opt/voice-agent/task_registry.py
|
COPY agent/task_registry.py /opt/voice-agent/task_registry.py
|
||||||
|
COPY agent/task_router.py /opt/voice-agent/task_router.py
|
||||||
COPY agent/task_worker.py /opt/voice-agent/task_worker.py
|
COPY agent/task_worker.py /opt/voice-agent/task_worker.py
|
||||||
COPY agent/dispatch_mcp.py /opt/voice-agent/dispatch_mcp.py
|
COPY agent/dispatch_mcp.py /opt/voice-agent/dispatch_mcp.py
|
||||||
|
|
||||||
|
|||||||
+102
-12
@@ -11,6 +11,7 @@ import logging
|
|||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
import textwrap
|
import textwrap
|
||||||
|
import uuid
|
||||||
|
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
from livekit.agents import (
|
from livekit.agents import (
|
||||||
@@ -157,6 +158,14 @@ SYSTEM_PROMPT = textwrap.dedent("""\
|
|||||||
""")
|
""")
|
||||||
|
|
||||||
|
|
||||||
|
# ── Background-task routing ─────────────────────────────────────────────────
|
||||||
|
|
||||||
|
try:
|
||||||
|
from agent.task_router import should_force_background_task
|
||||||
|
except ImportError:
|
||||||
|
from task_router import should_force_background_task
|
||||||
|
|
||||||
|
|
||||||
# ── Gemma LLM with reasoning_content fallback ───────────────────────────────
|
# ── Gemma LLM with reasoning_content fallback ───────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
@@ -222,6 +231,31 @@ class _ReasoningFallbackWrapper:
|
|||||||
return await self._inner.collect()
|
return await self._inner.collect()
|
||||||
|
|
||||||
|
|
||||||
|
class _ForcedToolCallStream(lk_llm.LLMStream):
|
||||||
|
"""Emit a dispatch_task call without asking the small LLM to route it."""
|
||||||
|
|
||||||
|
def __init__(self, llm, *, chat_ctx, tools, conn_options, description: str):
|
||||||
|
self._description = description
|
||||||
|
super().__init__(llm, chat_ctx=chat_ctx, tools=tools, conn_options=conn_options)
|
||||||
|
|
||||||
|
async def _run(self) -> None:
|
||||||
|
self._event_ch.send_nowait(
|
||||||
|
lk_llm.ChatChunk(
|
||||||
|
id=f"forced-dispatch-{uuid.uuid4().hex}",
|
||||||
|
delta=lk_llm.ChoiceDelta(
|
||||||
|
role="assistant",
|
||||||
|
tool_calls=[
|
||||||
|
lk_llm.FunctionToolCall(
|
||||||
|
name="dispatch_task",
|
||||||
|
arguments=json.dumps({"description": self._description}),
|
||||||
|
call_id=f"dispatch-{uuid.uuid4().hex}",
|
||||||
|
)
|
||||||
|
],
|
||||||
|
),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class GemmaLLM(openai.LLM):
|
class GemmaLLM(openai.LLM):
|
||||||
"""OpenAI-compatible LLM (llama.cpp Gemma 4) with thinking disabled."""
|
"""OpenAI-compatible LLM (llama.cpp Gemma 4) with thinking disabled."""
|
||||||
|
|
||||||
@@ -233,6 +267,10 @@ class GemmaLLM(openai.LLM):
|
|||||||
IDLE_COMPACTION_SECONDS = 300 # also compact after 5 min of silence
|
IDLE_COMPACTION_SECONDS = 300 # also compact after 5 min of silence
|
||||||
|
|
||||||
def __init__(self, *args, **kwargs):
|
def __init__(self, *args, **kwargs):
|
||||||
|
# llama.cpp's Gemma tool template ignores named tool_choice when OpenAI
|
||||||
|
# strict schemas are enabled. Non-strict schemas make forced dispatch
|
||||||
|
# (and ordinary MCP calls) work reliably.
|
||||||
|
kwargs.setdefault("_strict_tool_schema", False)
|
||||||
super().__init__(*args, **kwargs)
|
super().__init__(*args, **kwargs)
|
||||||
import time
|
import time
|
||||||
self._last_call_time: float = time.monotonic()
|
self._last_call_time: float = time.monotonic()
|
||||||
@@ -263,6 +301,25 @@ class GemmaLLM(openai.LLM):
|
|||||||
)
|
)
|
||||||
chat_ctx = self._compact_context(chat_ctx, conn_options)
|
chat_ctx = self._compact_context(chat_ctx, conn_options)
|
||||||
|
|
||||||
|
# Force dispatch for clear background-work requests. Without this,
|
||||||
|
# Gemma often says "I'll get on that" but emits no tool call at all.
|
||||||
|
items = list(chat_ctx.items)
|
||||||
|
latest = items[-1] if items else None
|
||||||
|
latest_role = getattr(latest, "role", None)
|
||||||
|
latest_text = self._msg_content_text(latest) if latest_role == "user" else ""
|
||||||
|
dispatch_available = any(
|
||||||
|
getattr(tool, "id", None) == "dispatch_task" for tool in (tools or [])
|
||||||
|
)
|
||||||
|
if dispatch_available and should_force_background_task(latest_text):
|
||||||
|
logger.info("Routing background request to dispatch_task: %.120s", latest_text)
|
||||||
|
return _ForcedToolCallStream(
|
||||||
|
self,
|
||||||
|
chat_ctx=chat_ctx,
|
||||||
|
tools=tools or [],
|
||||||
|
conn_options=conn_options,
|
||||||
|
description=latest_text,
|
||||||
|
)
|
||||||
|
|
||||||
stream = super().chat(
|
stream = super().chat(
|
||||||
chat_ctx=chat_ctx,
|
chat_ctx=chat_ctx,
|
||||||
tools=tools,
|
tools=tools,
|
||||||
@@ -386,8 +443,8 @@ class GemmaLLM(openai.LLM):
|
|||||||
|
|
||||||
# ── MCP toolsets (web access + any extra configured servers) ────────────────
|
# ── MCP toolsets (web access + any extra configured servers) ────────────────
|
||||||
|
|
||||||
def build_mcp_toolsets(user_profile: dict) -> list[mcp.MCPToolset]:
|
def build_mcp_toolsets(user_profile: dict, task_dir: str) -> list[mcp.MCPToolset]:
|
||||||
"""Build isolated MCP toolsets for one authenticated voice user."""
|
"""Build isolated MCP toolsets for one authenticated voice user/session."""
|
||||||
toolsets: list[mcp.MCPToolset] = []
|
toolsets: list[mcp.MCPToolset] = []
|
||||||
|
|
||||||
python_bin = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".venv", "bin", "python")
|
python_bin = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".venv", "bin", "python")
|
||||||
@@ -459,12 +516,12 @@ def build_mcp_toolsets(user_profile: dict) -> list[mcp.MCPToolset]:
|
|||||||
mcp_server=mcp.MCPServerStdio(
|
mcp_server=mcp.MCPServerStdio(
|
||||||
command=python_bin,
|
command=python_bin,
|
||||||
args=[dispatch_mcp_script],
|
args=[dispatch_mcp_script],
|
||||||
env={**os.environ},
|
env={**os.environ, "TASKS_DIR": task_dir},
|
||||||
client_session_timeout_seconds=60,
|
client_session_timeout_seconds=60,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
logger.info("Dispatch MCP toolset enabled")
|
logger.info("Dispatch MCP toolset enabled (%s)", task_dir)
|
||||||
|
|
||||||
extra = os.environ.get("EXTRA_MCP_SERVERS", "")
|
extra = os.environ.get("EXTRA_MCP_SERVERS", "")
|
||||||
if extra:
|
if extra:
|
||||||
@@ -492,7 +549,7 @@ def build_mcp_toolsets(user_profile: dict) -> list[mcp.MCPToolset]:
|
|||||||
class VoiceAssistant(Agent):
|
class VoiceAssistant(Agent):
|
||||||
"""The conversational agent. LLM is the brain; STT/TTS are senses."""
|
"""The conversational agent. LLM is the brain; STT/TTS are senses."""
|
||||||
|
|
||||||
def __init__(self, user_profile: dict) -> None:
|
def __init__(self, user_profile: dict, task_dir: str) -> None:
|
||||||
display_name = user_profile.get("display_name") or user_profile["user_id"]
|
display_name = user_profile.get("display_name") or user_profile["user_id"]
|
||||||
instructions = SYSTEM_PROMPT.replace("{user_display_name}", display_name)
|
instructions = SYSTEM_PROMPT.replace("{user_display_name}", display_name)
|
||||||
super().__init__(
|
super().__init__(
|
||||||
@@ -504,7 +561,7 @@ class VoiceAssistant(Agent):
|
|||||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||||
),
|
),
|
||||||
instructions=instructions,
|
instructions=instructions,
|
||||||
tools=build_mcp_toolsets(user_profile),
|
tools=build_mcp_toolsets(user_profile, task_dir),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -614,11 +671,15 @@ async def handle_job(ctx: JobContext) -> None:
|
|||||||
# audio is arriving from the participant.
|
# audio is arriving from the participant.
|
||||||
logger.info("user state -> %s", ev.new_state)
|
logger.info("user state -> %s", ev.new_state)
|
||||||
|
|
||||||
# ── Background task events → room data channel ─────────────────────────
|
# ── Background task events → UI and spoken result ──────────────────────
|
||||||
|
# Every room gets a private event directory. This prevents concurrent or
|
||||||
|
# reconnecting sessions from consuming and announcing each other's tasks.
|
||||||
|
task_dir = f"/tmp/tasks/{ctx.room.name}"
|
||||||
try:
|
try:
|
||||||
from agent.task_registry import registry as task_registry
|
from agent.task_registry import TaskRegistry
|
||||||
except ImportError:
|
except ImportError:
|
||||||
from task_registry import registry as task_registry
|
from task_registry import TaskRegistry
|
||||||
|
task_registry = TaskRegistry(task_dir)
|
||||||
|
|
||||||
async def _publish_task_event(task_id: str, event: dict) -> None:
|
async def _publish_task_event(task_id: str, event: dict) -> None:
|
||||||
payload = json.dumps({"type": "task_event", "task_id": task_id, **event})
|
payload = json.dumps({"type": "task_event", "task_id": task_id, **event})
|
||||||
@@ -629,17 +690,46 @@ async def handle_job(ctx: JobContext) -> None:
|
|||||||
except Exception as e: # noqa: BLE001
|
except Exception as e: # noqa: BLE001
|
||||||
logger.warning("failed to publish task event: %s", e)
|
logger.warning("failed to publish task event: %s", e)
|
||||||
|
|
||||||
|
def _announce_task_result(task_id: str, event: dict) -> None:
|
||||||
|
event_type = event.get("event")
|
||||||
|
if event_type == "completed":
|
||||||
|
result = str(event.get("result") or "").strip()
|
||||||
|
if result:
|
||||||
|
logger.info("Announcing completed task %s: %.160s", task_id, result)
|
||||||
|
session.say(result, allow_interruptions=True, add_to_chat_ctx=True)
|
||||||
|
elif event_type == "failed":
|
||||||
|
logger.warning("Background task %s failed: %s", task_id, event.get("error"))
|
||||||
|
session.say(
|
||||||
|
"I couldn't finish that background task. You can ask me to try again.",
|
||||||
|
allow_interruptions=True,
|
||||||
|
add_to_chat_ctx=True,
|
||||||
|
)
|
||||||
|
|
||||||
async def _poll_tasks() -> None:
|
async def _poll_tasks() -> None:
|
||||||
while True:
|
while True:
|
||||||
await asyncio.sleep(1.0)
|
try:
|
||||||
|
await asyncio.sleep(0.5)
|
||||||
events = await task_registry.poll_events()
|
events = await task_registry.poll_events()
|
||||||
for task_id, event in events:
|
for task_id, event in events:
|
||||||
await _publish_task_event(task_id, event)
|
await _publish_task_event(task_id, event)
|
||||||
|
_announce_task_result(task_id, event)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
# A transient room or speech error must not permanently kill
|
||||||
|
# result delivery for every later task.
|
||||||
|
logger.exception("Background task event poller failed")
|
||||||
|
|
||||||
asyncio.create_task(_poll_tasks())
|
poll_task = asyncio.create_task(_poll_tasks(), name=f"task-poller-{ctx.room.name}")
|
||||||
|
|
||||||
|
async def _stop_task_poller(*_args) -> None:
|
||||||
|
poll_task.cancel()
|
||||||
|
await asyncio.gather(poll_task, return_exceptions=True)
|
||||||
|
|
||||||
|
ctx.add_shutdown_callback(_stop_task_poller)
|
||||||
|
|
||||||
await session.start(
|
await session.start(
|
||||||
agent=VoiceAssistant(user_profile),
|
agent=VoiceAssistant(user_profile, task_dir),
|
||||||
room=ctx.room,
|
room=ctx.room,
|
||||||
room_options=room_io.RoomOptions(
|
room_options=room_io.RoomOptions(
|
||||||
# Keep the agent in the room when a participant leaves; it must
|
# Keep the agent in the room when a participant leaves; it must
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
|
import os
|
||||||
|
|
||||||
from mcp.server.fastmcp import FastMCP
|
from mcp.server.fastmcp import FastMCP
|
||||||
|
|
||||||
@@ -51,8 +52,13 @@ async def _run_worker(worker: "TaskWorker", task_id: str):
|
|||||||
except ImportError:
|
except ImportError:
|
||||||
from task_registry import registry
|
from task_registry import registry
|
||||||
|
|
||||||
|
timeout_seconds = float(os.environ.get("TASK_TIMEOUT_SECONDS", "180"))
|
||||||
try:
|
try:
|
||||||
await worker.run()
|
await asyncio.wait_for(worker.run(), timeout=timeout_seconds)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
message = f"Timed out after {timeout_seconds:g} seconds"
|
||||||
|
logger.error("Task %s %s", task_id, message.lower())
|
||||||
|
registry.fail(task_id, message)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.exception("Task %s crashed: %s", task_id, e)
|
logger.exception("Task %s crashed: %s", task_id, e)
|
||||||
registry.fail(task_id, str(e))
|
registry.fail(task_id, str(e))
|
||||||
|
|||||||
+15
-9
@@ -52,12 +52,18 @@ class Task:
|
|||||||
|
|
||||||
class TaskRegistry:
|
class TaskRegistry:
|
||||||
"""File-backed registry. Writers (MCP process) append to events.jsonl.
|
"""File-backed registry. Writers (MCP process) append to events.jsonl.
|
||||||
Readers (main agent) tail the file and dispatch to subscribers."""
|
Readers (main agent) tail the file and dispatch to subscribers.
|
||||||
|
|
||||||
def __init__(self):
|
A registry can be scoped to a task directory. Voice sessions use a unique
|
||||||
TASKS_DIR.mkdir(parents=True, exist_ok=True)
|
directory so one room can never consume another room's task results.
|
||||||
if not EVENTS_FILE.exists():
|
"""
|
||||||
EVENTS_FILE.touch()
|
|
||||||
|
def __init__(self, tasks_dir: str | Path | None = None):
|
||||||
|
self.tasks_dir = Path(tasks_dir) if tasks_dir is not None else TASKS_DIR
|
||||||
|
self.events_file = self.tasks_dir / "events.jsonl"
|
||||||
|
self.tasks_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
if not self.events_file.exists():
|
||||||
|
self.events_file.touch()
|
||||||
self._tasks: dict[str, Task] = {}
|
self._tasks: dict[str, Task] = {}
|
||||||
self._listeners: list[Callable[[str, dict], Any]] = []
|
self._listeners: list[Callable[[str, dict], Any]] = []
|
||||||
self._tail_pos = 0
|
self._tail_pos = 0
|
||||||
@@ -109,7 +115,7 @@ class TaskRegistry:
|
|||||||
"""Read new events from the file. Call periodically from the agent."""
|
"""Read new events from the file. Call periodically from the agent."""
|
||||||
events = []
|
events = []
|
||||||
try:
|
try:
|
||||||
with open(EVENTS_FILE, "r") as f:
|
with open(self.events_file, "r", encoding="utf-8") as f:
|
||||||
f.seek(self._tail_pos)
|
f.seek(self._tail_pos)
|
||||||
for line in f:
|
for line in f:
|
||||||
line = line.strip()
|
line = line.strip()
|
||||||
@@ -136,12 +142,12 @@ class TaskRegistry:
|
|||||||
|
|
||||||
def _write_event(self, task_id: str, event: dict):
|
def _write_event(self, task_id: str, event: dict):
|
||||||
payload = json.dumps({"task_id": task_id, **event})
|
payload = json.dumps({"task_id": task_id, **event})
|
||||||
with open(EVENTS_FILE, "a") as f:
|
with open(self.events_file, "a", encoding="utf-8") as f:
|
||||||
f.write(payload + "\n")
|
f.write(payload + "\n")
|
||||||
|
|
||||||
def _save_task(self, task: Task):
|
def _save_task(self, task: Task):
|
||||||
path = TASKS_DIR / f"{task.id}.json"
|
path = self.tasks_dir / f"{task.id}.json"
|
||||||
path.write_text(json.dumps(self._task_dict(task)))
|
path.write_text(json.dumps(self._task_dict(task)), encoding="utf-8")
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _task_dict(task: Task) -> dict:
|
def _task_dict(task: Task) -> dict:
|
||||||
|
|||||||
@@ -0,0 +1,58 @@
|
|||||||
|
"""Deterministic routing hints for background voice tasks."""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import re
|
||||||
|
|
||||||
|
# Gemma is a small model and will sometimes promise to do work without actually
|
||||||
|
# calling dispatch_task. Force that tool for high-confidence background intents;
|
||||||
|
# the main prompt remains responsible for less obvious cases.
|
||||||
|
_BACKGROUND_TASK_RE = re.compile(
|
||||||
|
r"\b(?:"
|
||||||
|
r"latest|news|current events?|headlines?|prices?|scores?|"
|
||||||
|
r"research|look\s+up|search(?:\s+for)?|find\s+out|investigate|"
|
||||||
|
r"compare|comparison|pros\s+and\s+cons|"
|
||||||
|
r"plan|itinerary|schedule|"
|
||||||
|
r"write|draft|compose|rewrite|"
|
||||||
|
r"summari[sz]e|summary|analy[sz]e|analysis|review|report|"
|
||||||
|
r"recommend|recommendation|options?|developments?"
|
||||||
|
r")\b",
|
||||||
|
re.IGNORECASE,
|
||||||
|
)
|
||||||
|
_MEMORY_INTENT_RE = re.compile(
|
||||||
|
r"\b(?:remember|memory|memories|forget|write down|my name|"
|
||||||
|
r"what do you know about me|preference|preferences)\b",
|
||||||
|
re.IGNORECASE,
|
||||||
|
)
|
||||||
|
_DIRECT_WEATHER_TIME_RE = re.compile(
|
||||||
|
r"\b(?:weather|temperature|forecast|what time|current time|what day|today's date)\b",
|
||||||
|
re.IGNORECASE,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def should_force_background_task(text: str) -> bool:
|
||||||
|
"""Return whether a fresh user turn should be forced through dispatch_task.
|
||||||
|
|
||||||
|
This deliberately favors precise intent signals over trying to classify every
|
||||||
|
possible request. Memory stays on its authenticated MCP and straightforward
|
||||||
|
weather/time requests stay on their purpose-built tools.
|
||||||
|
"""
|
||||||
|
normalized = " ".join((text or "").split())
|
||||||
|
if not normalized or _MEMORY_INTENT_RE.search(normalized):
|
||||||
|
return False
|
||||||
|
if re.search(r"\b(?:my|saved)\s+notes?\b", normalized, re.IGNORECASE):
|
||||||
|
return False
|
||||||
|
|
||||||
|
match = _BACKGROUND_TASK_RE.search(normalized)
|
||||||
|
if not match:
|
||||||
|
return bool(re.search(r"\b(?:tell me about|give me an overview of)\b", normalized, re.I))
|
||||||
|
|
||||||
|
# Straightforward weather/time requests stay on the fast dedicated MCP.
|
||||||
|
# Comparisons, research, plans, and written reports still go background.
|
||||||
|
if _DIRECT_WEATHER_TIME_RE.search(normalized) and not re.search(
|
||||||
|
r"\b(?:compare|comparison|research|investigate|plan|write|draft|"
|
||||||
|
r"summari[sz]e|analy[sz]e|review|report)\b",
|
||||||
|
normalized,
|
||||||
|
re.IGNORECASE,
|
||||||
|
):
|
||||||
|
return False
|
||||||
|
return True
|
||||||
@@ -36,7 +36,7 @@ class TaskWorker:
|
|||||||
|
|
||||||
reg.add_step(self.task_id, "thinking", f"Starting: {self.description}")
|
reg.add_step(self.task_id, "thinking", f"Starting: {self.description}")
|
||||||
|
|
||||||
client = AsyncOpenAI(base_url=base_url, api_key=api_key)
|
client = AsyncOpenAI(base_url=base_url, api_key=api_key, timeout=60.0)
|
||||||
tools = self._load_tools()
|
tools = self._load_tools()
|
||||||
|
|
||||||
messages = [
|
messages = [
|
||||||
@@ -45,6 +45,12 @@ class TaskWorker:
|
|||||||
]
|
]
|
||||||
|
|
||||||
max_iterations = 10
|
max_iterations = 10
|
||||||
|
try:
|
||||||
|
await self._run_loop(client, model, tools, messages, max_iterations, reg)
|
||||||
|
finally:
|
||||||
|
await client.close()
|
||||||
|
|
||||||
|
async def _run_loop(self, client, model, tools, messages, max_iterations, reg):
|
||||||
for iteration in range(max_iterations):
|
for iteration in range(max_iterations):
|
||||||
if self._cancelled:
|
if self._cancelled:
|
||||||
reg.fail(self.task_id, "Cancelled")
|
reg.fail(self.task_id, "Cancelled")
|
||||||
@@ -106,7 +112,7 @@ class TaskWorker:
|
|||||||
return (
|
return (
|
||||||
"You are Hope's research assistant working on a background task. "
|
"You are Hope's research assistant working on a background task. "
|
||||||
"Use your tools to gather information thoroughly. When you have "
|
"Use your tools to gather information thoroughly. When you have "
|
||||||
"enough, provide a final answer in 2-4 sentences suitable for "
|
"enough, provide a final answer in 1-3 sentences suitable for "
|
||||||
"speaking aloud. Plain text only, no markdown or formatting."
|
"speaking aloud. Plain text only, no markdown or formatting."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,116 @@
|
|||||||
|
"""Regression tests for background-task routing and event isolation."""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
import uuid
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).parents[1] / "agent"))
|
||||||
|
from task_registry import TaskRegistry # noqa: E402
|
||||||
|
from task_router import should_force_background_task # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"utterance",
|
||||||
|
[
|
||||||
|
"What are the top news stories today?",
|
||||||
|
"Compare the PS5 Pro and Xbox Series X for me.",
|
||||||
|
"Plan a weekend trip to Denver.",
|
||||||
|
"Write a short birthday message for my sister.",
|
||||||
|
"Write a birthday note for Sam.",
|
||||||
|
"Look up the current price of a PlayStation 5.",
|
||||||
|
"Research the latest solid state battery developments.",
|
||||||
|
"Tell me about quantum computing.",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_nontrivial_requests_force_dispatch(utterance):
|
||||||
|
assert should_force_background_task(utterance)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"utterance",
|
||||||
|
[
|
||||||
|
"What's my name?",
|
||||||
|
"Remember that I prefer tea.",
|
||||||
|
"Read my saved notes.",
|
||||||
|
"What's the weather in Raleigh?",
|
||||||
|
"What time is it?",
|
||||||
|
"Hello there.",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_direct_requests_do_not_force_dispatch(utterance):
|
||||||
|
assert not should_force_background_task(utterance)
|
||||||
|
|
||||||
|
|
||||||
|
def test_task_event_directories_are_isolated(tmp_path):
|
||||||
|
first = TaskRegistry(tmp_path / "room-one")
|
||||||
|
second = TaskRegistry(tmp_path / "room-two")
|
||||||
|
|
||||||
|
task = first.create("private task")
|
||||||
|
first.complete(task.id, "private result")
|
||||||
|
|
||||||
|
first_events = asyncio.run(first.poll_events())
|
||||||
|
second_events = asyncio.run(second.poll_events())
|
||||||
|
|
||||||
|
assert [event[1]["event"] for event in first_events] == ["created", "completed"]
|
||||||
|
assert second_events == []
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.slow
|
||||||
|
def test_dispatch_mcp_task_reaches_terminal_event():
|
||||||
|
"""The MCP request returns immediately while its child task stays alive."""
|
||||||
|
from conftest import AGENT_DIR_IN_CONTAINER, CONTAINER_PYTHON, docker_exec_python
|
||||||
|
|
||||||
|
task_dir = f"/tmp/task-test-{uuid.uuid4().hex}"
|
||||||
|
code = f"""
|
||||||
|
import asyncio, json, os
|
||||||
|
from mcp import ClientSession, StdioServerParameters
|
||||||
|
from mcp.client.stdio import stdio_client
|
||||||
|
|
||||||
|
async def main():
|
||||||
|
task_dir = {task_dir!r}
|
||||||
|
os.makedirs(task_dir, exist_ok=True)
|
||||||
|
env = {{**os.environ, "TASKS_DIR": task_dir, "TASK_TIMEOUT_SECONDS": "60"}}
|
||||||
|
params = StdioServerParameters(
|
||||||
|
command={CONTAINER_PYTHON!r},
|
||||||
|
args=[{AGENT_DIR_IN_CONTAINER + '/dispatch_mcp.py'!r}],
|
||||||
|
env=env,
|
||||||
|
)
|
||||||
|
async with stdio_client(params) as (read, write):
|
||||||
|
async with ClientSession(read, write) as session:
|
||||||
|
await session.initialize()
|
||||||
|
result = await session.call_tool(
|
||||||
|
"dispatch_task",
|
||||||
|
{{"description": "Write a warm two-sentence birthday greeting for Sam."}},
|
||||||
|
)
|
||||||
|
assert not result.isError
|
||||||
|
for _ in range(60):
|
||||||
|
await asyncio.sleep(0.5)
|
||||||
|
with open(os.path.join(task_dir, "events.jsonl"), encoding="utf-8") as f:
|
||||||
|
events = [json.loads(line) for line in f if line.strip()]
|
||||||
|
terminal = [e for e in events if e.get("event") in ("completed", "failed")]
|
||||||
|
if terminal:
|
||||||
|
print(json.dumps(terminal[-1]))
|
||||||
|
return
|
||||||
|
raise RuntimeError("background task produced no terminal event")
|
||||||
|
|
||||||
|
asyncio.run(main())
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
proc = docker_exec_python(code, timeout=90, check=False)
|
||||||
|
assert proc.returncode == 0, proc.stderr[-1200:]
|
||||||
|
terminal = json.loads(proc.stdout.strip().splitlines()[-1])
|
||||||
|
assert terminal["event"] == "completed", terminal
|
||||||
|
assert terminal.get("result")
|
||||||
|
finally:
|
||||||
|
# Best-effort cleanup; a failed container connection should not mask the assertion.
|
||||||
|
try:
|
||||||
|
from conftest import docker_exec
|
||||||
|
|
||||||
|
docker_exec("rm", "-rf", task_dir, check=False)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
Reference in new issue
Block a user