diff --git a/Dockerfile b/Dockerfile index fa31689..b4c676d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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/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/dispatch_mcp.py /opt/voice-agent/dispatch_mcp.py diff --git a/agent/agent.py b/agent/agent.py index d084d86..b19f1c9 100644 --- a/agent/agent.py +++ b/agent/agent.py @@ -11,6 +11,7 @@ import logging import os import re import textwrap +import uuid from dotenv import load_dotenv 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 ─────────────────────────────── @@ -222,6 +231,31 @@ class _ReasoningFallbackWrapper: 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): """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 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) import time self._last_call_time: float = time.monotonic() @@ -263,6 +301,25 @@ class GemmaLLM(openai.LLM): ) 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( chat_ctx=chat_ctx, tools=tools, @@ -386,8 +443,8 @@ class GemmaLLM(openai.LLM): # ── MCP toolsets (web access + any extra configured servers) ──────────────── -def build_mcp_toolsets(user_profile: dict) -> list[mcp.MCPToolset]: - """Build isolated MCP toolsets for one authenticated voice user.""" +def build_mcp_toolsets(user_profile: dict, task_dir: str) -> list[mcp.MCPToolset]: + """Build isolated MCP toolsets for one authenticated voice user/session.""" toolsets: list[mcp.MCPToolset] = [] 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( command=python_bin, args=[dispatch_mcp_script], - env={**os.environ}, + env={**os.environ, "TASKS_DIR": task_dir}, 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", "") if extra: @@ -492,7 +549,7 @@ def build_mcp_toolsets(user_profile: dict) -> list[mcp.MCPToolset]: class VoiceAssistant(Agent): """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"] instructions = SYSTEM_PROMPT.replace("{user_display_name}", display_name) super().__init__( @@ -504,7 +561,7 @@ class VoiceAssistant(Agent): extra_body={"chat_template_kwargs": {"enable_thinking": False}}, ), 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. 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: - from agent.task_registry import registry as task_registry + from agent.task_registry import TaskRegistry 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: 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 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: while True: - await asyncio.sleep(1.0) - events = await task_registry.poll_events() - for task_id, event in events: - await _publish_task_event(task_id, event) + try: + await asyncio.sleep(0.5) + events = await task_registry.poll_events() + for task_id, event in events: + 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( - agent=VoiceAssistant(user_profile), + agent=VoiceAssistant(user_profile, task_dir), room=ctx.room, room_options=room_io.RoomOptions( # Keep the agent in the room when a participant leaves; it must diff --git a/agent/dispatch_mcp.py b/agent/dispatch_mcp.py index 187f960..4687378 100644 --- a/agent/dispatch_mcp.py +++ b/agent/dispatch_mcp.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import logging +import os from mcp.server.fastmcp import FastMCP @@ -51,8 +52,13 @@ async def _run_worker(worker: "TaskWorker", task_id: str): except ImportError: from task_registry import registry + timeout_seconds = float(os.environ.get("TASK_TIMEOUT_SECONDS", "180")) 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: logger.exception("Task %s crashed: %s", task_id, e) registry.fail(task_id, str(e)) diff --git a/agent/task_registry.py b/agent/task_registry.py index 99461c6..db8211a 100644 --- a/agent/task_registry.py +++ b/agent/task_registry.py @@ -52,12 +52,18 @@ class Task: class TaskRegistry: """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): - TASKS_DIR.mkdir(parents=True, exist_ok=True) - if not EVENTS_FILE.exists(): - EVENTS_FILE.touch() + A registry can be scoped to a task directory. Voice sessions use a unique + directory so one room can never consume another room's task results. + """ + + 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._listeners: list[Callable[[str, dict], Any]] = [] self._tail_pos = 0 @@ -109,7 +115,7 @@ class TaskRegistry: """Read new events from the file. Call periodically from the agent.""" events = [] try: - with open(EVENTS_FILE, "r") as f: + with open(self.events_file, "r", encoding="utf-8") as f: f.seek(self._tail_pos) for line in f: line = line.strip() @@ -136,12 +142,12 @@ class TaskRegistry: def _write_event(self, task_id: str, event: dict): 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") def _save_task(self, task: Task): - path = TASKS_DIR / f"{task.id}.json" - path.write_text(json.dumps(self._task_dict(task))) + path = self.tasks_dir / f"{task.id}.json" + path.write_text(json.dumps(self._task_dict(task)), encoding="utf-8") @staticmethod def _task_dict(task: Task) -> dict: diff --git a/agent/task_router.py b/agent/task_router.py new file mode 100644 index 0000000..c8cbe57 --- /dev/null +++ b/agent/task_router.py @@ -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 diff --git a/agent/task_worker.py b/agent/task_worker.py index 82da2ea..1d5bbaf 100644 --- a/agent/task_worker.py +++ b/agent/task_worker.py @@ -36,7 +36,7 @@ class TaskWorker: 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() messages = [ @@ -45,6 +45,12 @@ class TaskWorker: ] 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): if self._cancelled: reg.fail(self.task_id, "Cancelled") @@ -106,7 +112,7 @@ class TaskWorker: return ( "You are Hope's research assistant working on a background task. " "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." ) diff --git a/tests/test_background_tasks.py b/tests/test_background_tasks.py new file mode 100644 index 0000000..6a54a9e --- /dev/null +++ b/tests/test_background_tasks.py @@ -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