From 543bfff0eb78a89e47900722c6a89dfeef900f90 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Wed, 13 May 2026 23:14:41 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20Replace=20synchronous=20sleep=20wit?= =?UTF-8?q?h=20asyncio=20in=20task=20worker?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refactored `_run_tool_call_task_worker` to an async function and execute it in a dedicated background `asyncio` event loop thread. This replaces `time.sleep(delay)` with `await asyncio.sleep(delay)`, resolving the performance overhead of consuming an entire OS thread just to block and wait during task initial delays. Fixes an issue where synchronous sleeps were artificially capping concurrency or tying up Python threads needlessly in the MCP handlers layer. Co-authored-by: wjohns989 <56205870+wjohns989@users.noreply.github.com> --- muninn/mcp/handlers.py | 141 ++++++++++++++++++++++++----------------- plan.md | 23 +++++++ 2 files changed, 105 insertions(+), 59 deletions(-) create mode 100644 plan.md diff --git a/muninn/mcp/handlers.py b/muninn/mcp/handlers.py index c377667..ff4864d 100644 --- a/muninn/mcp/handlers.py +++ b/muninn/mcp/handlers.py @@ -1,41 +1,62 @@ import os DEFAULT_HTTP_TIMEOUT = float(os.getenv("MUNINN_MCP_HTTP_TIMEOUT_SEC", "40")) -import time +import asyncio import logging import threading -import json -from datetime import datetime, timezone -from typing import Optional, Dict, Any, List, Union +import time +from typing import Any, Dict, List, Optional from urllib.parse import quote from muninn.version import __version__ as _MUNINN_VERSION -from .state import _SESSION_STATE from .definitions import ( - SUPPORTED_PROTOCOL_VERSIONS, TOOLS_SCHEMAS, JSON_SCHEMA_2020_12, - READ_ONLY_TOOLS, DESTRUCTIVE_TOOLS, IDEMPOTENT_TOOLS, - SUPPORTED_MODEL_PROFILES, MIMIR_TOOLS, + DESTRUCTIVE_TOOLS, + IDEMPOTENT_TOOLS, + JSON_SCHEMA_2020_12, + READ_ONLY_TOOLS, + TOOLS_SCHEMAS, ) +from .lifecycle import SERVER_URL, check_and_start_ollama, ensure_server_running, is_circuit_open +from .requests import make_request_with_retry +from .state import _SESSION_STATE from .tasks import ( - create_task, lookup_task_locked, purge_and_retain_tasks_locked, - get_registry, public_task, encode_task_cursor, decode_task_cursor, - TASK_TERMINAL_STATUSES, TASKS_LIST_PAGE_SIZE, TASK_CURSOR_PREFIX, - set_task_state_locked, emit_task_status_notification + TASK_TERMINAL_STATUSES, + TASKS_LIST_PAGE_SIZE, + create_task, + decode_task_cursor, + emit_task_status_notification, + encode_task_cursor, + get_registry, + lookup_task_locked, + public_task, + purge_and_retain_tasks_locked, + set_task_state_locked, ) from .utils import ( - get_git_info, inject_operator_profile_metadata, format_tool_result_text, - truncate_tool_text, env_flag, negotiated_protocol_version, build_initialize_instructions, - _read_operator_model_profile, get_tool_call_deadline_epoch, remaining_deadline_seconds, - startup_recovery_allowed + _read_operator_model_profile, + build_initialize_instructions, + env_flag, + format_tool_result_text, + get_git_info, + get_tool_call_deadline_epoch, + inject_operator_profile_metadata, + negotiated_protocol_version, + startup_recovery_allowed, + truncate_tool_text, ) -from .lifecycle import ensure_server_running, is_circuit_open, SERVER_URL, check_and_start_ollama -from .requests import make_request_with_retry logger = logging.getLogger("Muninn.mcp.handlers") _thread_local = threading.local() +_task_loop = asyncio.new_event_loop() +def _start_task_loop(): + asyncio.set_event_loop(_task_loop) + _task_loop.run_forever() + +threading.Thread(target=_start_task_loop, daemon=True, name="MuninnTaskLoop").start() + def _parse_positive_float_env(name: str, default: float, *, lower: float, upper: float) -> float: raw = os.environ.get(name) @@ -62,7 +83,7 @@ def handle_initialize(msg_id: Any, params: Dict[str, Any], send_error_fn, send_r requested_version = params.get("protocolVersion") negotiated_version = negotiated_protocol_version(requested_version) - + if not negotiated_version: send_error_fn(msg_id, -32602, f"Unsupported protocol version {requested_version}") return @@ -71,7 +92,7 @@ def handle_initialize(msg_id: Any, params: Dict[str, Any], send_error_fn, send_r _SESSION_STATE["protocol_version"] = negotiated_version _SESSION_STATE["client_capabilities"] = params.get("capabilities", {}) _SESSION_STATE["client_info"] = params.get("clientInfo", {}) - + # Elicitation modes client_meta = params.get("_meta", {}) if isinstance(client_meta, dict): @@ -100,11 +121,11 @@ def handle_initialize(msg_id: Any, params: Dict[str, Any], send_error_fn, send_r if isinstance(elicitation.get("form"), dict): modes.append("form") if isinstance(elicitation.get("url"), dict): modes.append("url") _SESSION_STATE["client_elicitation_modes"] = tuple(modes) - + # Operator model profile session_profile = _read_operator_model_profile("MUNINN_OPERATOR_MODEL_PROFILE") _SESSION_STATE["operator_model_profile"] = session_profile - + # Send result result = { "protocolVersion": negotiated_version, @@ -131,9 +152,8 @@ def handle_initialize(msg_id: Any, params: Dict[str, Any], send_error_fn, send_r def handle_list_tools(msg_id: Any, send_result_fn): """List available tools with schemas and hints.""" - from muninn.core.security import get_token # In Phase 10, Listing tools is allowed, but execution requires token parity. - + tools_list = [] for schema_def in TOOLS_SCHEMAS: name = schema_def["name"] @@ -145,7 +165,7 @@ def handle_list_tools(msg_id: Any, send_result_fn): # Add JSON schema and hints for SOTA clients if isinstance(tool_def["inputSchema"], dict) and "$schema" not in tool_def["inputSchema"]: tool_def["inputSchema"]["$schema"] = JSON_SCHEMA_2020_12 - + # Mapping hints to legacy annotations read_only = name in READ_ONLY_TOOLS annotations = { @@ -156,26 +176,26 @@ def handle_list_tools(msg_id: Any, send_result_fn): } tool_def["annotations"] = annotations tool_def["execution"] = {"taskSupport": "optional"} - + tools_list.append(tool_def) - + send_result_fn(msg_id, {"tools": tools_list}) def handle_call_tool_with_task(msg_id: Any, name: str, arguments: Dict[str, Any], task_request: Dict[str, Any], send_result_fn, send_notification_fn=None, worker_fn=None): """Create a task for a tool call and return immediately.""" if worker_fn is None: worker_fn = _run_tool_call_task_worker - + if send_notification_fn is None: # Fallback for environments where notification sender isn't explicitly provided send_notification_fn = lambda msg: send_result_fn(None, msg) - + ttl = task_request.get("ttl") task = create_task(name, ttl_ms=ttl) - + # Capture snapshot before starting thread to ensure 'working' status in immediate response task_snapshot = public_task(task) - + worker_args = (task["taskId"], name, arguments, send_notification_fn) # Support legacy 3-arg workers (from tests) by dropping notification sender try: @@ -186,11 +206,14 @@ def handle_call_tool_with_task(msg_id: Any, name: str, arguments: Dict[str, Any] except Exception: pass - threading.Thread( - target=worker_fn, - args=worker_args, - daemon=True - ).start() + if asyncio.iscoroutinefunction(worker_fn): + asyncio.run_coroutine_threadsafe(worker_fn(*worker_args), _task_loop) + else: + threading.Thread( + target=worker_fn, + args=worker_args, + daemon=True + ).start() logger.debug("Task accepted: tool=%s task_id=%s", name, task["taskId"]) @@ -205,13 +228,13 @@ def handle_call_tool_with_task(msg_id: Any, name: str, arguments: Dict[str, Any] } send_result_fn(msg_id, payload) -def _run_tool_call_task_worker(task_id: str, name: str, arguments: Dict[str, Any], send_notification_fn): +async def _run_tool_call_task_worker(task_id: str, name: str, arguments: Dict[str, Any], send_notification_fn): """Background worker for task-backed tool calls.""" # Delay if requested (SOTA pattern) delay = float(os.environ.get("MUNINN_MCP_TASK_WORKER_START_DELAY_MS", "0")) / 1000.0 if delay > 0: - time.sleep(delay) - + await asyncio.sleep(delay) + task = lookup_task_locked(task_id) if not task: logger.error("Task worker started for unknown task: %s", task_id) @@ -219,7 +242,7 @@ def _run_tool_call_task_worker(task_id: str, name: str, arguments: Dict[str, Any try: deadline = time.time() + (task["ttl"] / 1000.0) - res = _do_call_tool_logic(name, arguments, deadline) + res = await asyncio.to_thread(_do_call_tool_logic, name, arguments, deadline) if res is None: set_task_state_locked(task, status="failed", error={"code": -32601, "message": f"Method not found: {name}"}) else: @@ -264,7 +287,7 @@ def handle_get_task_result(msg_id: Any, params: Dict[str, Any], send_error_fn, s wait_started_at = time.monotonic() from .state import _TASKS_CONDITION - from .tasks import lookup_task_locked, TASK_TERMINAL_STATUSES, related_task_meta + from .tasks import TASK_TERMINAL_STATUSES, lookup_task_locked, related_task_meta with _TASKS_CONDITION: while True: @@ -272,11 +295,11 @@ def handle_get_task_result(msg_id: Any, params: Dict[str, Any], send_error_fn, s if not task: send_error_fn(msg_id, -32602, "Invalid params: unknown taskId") return - + status = task.get("status") if status in TASK_TERMINAL_STATUSES or status == "input_required": break - + if not should_block: send_error_fn(msg_id, -32002, f"Task result is not ready (status: {status}). Poll or use blocking mode.") return @@ -289,7 +312,7 @@ def handle_get_task_result(msg_id: Any, params: Dict[str, Any], send_error_fn, s send_error_fn(msg_id, -32002, "Task result not ready within host-safe budget.") return wait_timeout = min(wait_timeout, remaining) - + _TASKS_CONDITION.wait(timeout=wait_timeout) # Terminal state reached @@ -316,7 +339,7 @@ def handle_get_task_result(msg_id: Any, params: Dict[str, Any], send_error_fn, s result_with_meta["_meta"] = meta else: result_with_meta["_meta"] = related_task_meta(task_id) - + send_result_fn(msg_id, result_with_meta) def handle_call_tool(msg_id: Any, params: Dict[str, Any], send_error_fn, send_result_fn): @@ -334,7 +357,7 @@ def handle_call_tool(msg_id: Any, params: Dict[str, Any], send_error_fn, send_re if not isinstance(name, str) or not name: send_error_fn(msg_id, -32602, "Invalid params: tools/call requires non-empty string name") return - + task_request = params.get("task") if task_request is not None and not isinstance(task_request, dict): send_error_fn(msg_id, -32602, "Invalid params: task must be an object") @@ -354,7 +377,7 @@ def handle_call_tool(msg_id: Any, params: Dict[str, Any], send_error_fn, send_re arguments = params.get("arguments", {}) deadline = get_tool_call_deadline_epoch() - + # Validation for specific tools that need early exit if name == "delete_all_memories" and not arguments.get("confirm", False): send_error_fn(msg_id, -32602, "Must set 'confirm: true' to delete all memories") @@ -389,7 +412,7 @@ def handle_call_tool(msg_id: Any, params: Dict[str, Any], send_error_fn, send_re send_result_fn(msg_id, { "content": [{"type": "text", "text": truncated_text}] }) - + # Update metrics for success tool_metrics["response_count"] += 1 # Simplified size tracking @@ -406,12 +429,12 @@ def handle_call_tool(msg_id: Any, params: Dict[str, Any], send_error_fn, send_re # Telemetry logging matching the original wrapper elapsed_ms = (time.monotonic() - tool_call_started_monotonic) * 1000.0 outcome = "error" if tool_metrics["saw_error"] else ("success" if tool_metrics["response_count"] > 0 else "no_response") - + logger.info( "Tool call telemetry: name=%s id=%r outcome=%s elapsed_ms=%.1f responses=%d " "response_bytes_total=%d response_bytes_max=%d", - name, msg_id, outcome, elapsed_ms, - tool_metrics["response_count"], tool_metrics["response_bytes_total"], + name, msg_id, outcome, elapsed_ms, + tool_metrics["response_count"], tool_metrics["response_bytes_total"], tool_metrics["response_bytes_max"] ) setattr(_thread_local, "tool_call_metrics", None) @@ -450,7 +473,7 @@ def _do_call_tool_logic(name: str, arguments: Dict[str, Any], deadline: Optional "apply_federation_bundle": _do_apply_federation_bundle, "mimir_relay": _do_mimir_relay, } - + handler = dispatch.get(name) if not handler: return None @@ -890,7 +913,7 @@ def handle_list_tasks(msg_id: Any, params: Dict[str, Any], send_error_fn, send_r cursor = params.get("cursor") limit = min(params.get("limit", TASKS_LIST_PAGE_SIZE), TASKS_LIST_PAGE_SIZE) - + try: offset = decode_task_cursor(cursor) if cursor else 0 except ValueError as e: @@ -899,17 +922,17 @@ def handle_list_tasks(msg_id: Any, params: Dict[str, Any], send_error_fn, send_r registry = get_registry() purge_and_retain_tasks_locked(time.time()) - + all_tasks = sorted(registry.values(), key=lambda t: t.get("createdAt", ""), reverse=True) page = all_tasks[offset : offset + limit] - + next_cursor = encode_task_cursor(offset + limit) if offset + limit < len(all_tasks) else None send_result = { "tasks": [public_task(t) for t in page], } if next_cursor: send_result["nextCursor"] = next_cursor - + send_result_fn(msg_id, send_result) def handle_get_task(msg_id: Any, params: Dict[str, Any], send_error_fn, send_result_fn): @@ -936,7 +959,7 @@ def handle_cancel_task(msg_id: Any, params: Dict[str, Any], send_error_fn, send_ if not task: send_error_fn(msg_id, -32002, f"Task not found: {task_id}") return - + if task["status"] in TASK_TERMINAL_STATUSES: send_error_fn(msg_id, -32602, f"Task {task_id} is already in terminal state {task['status']}") return @@ -946,5 +969,5 @@ def handle_cancel_task(msg_id: Any, params: Dict[str, Any], send_error_fn, send_ logger.info("Cancelled task %s", task_id) if send_notification_fn: emit_task_status_notification(task, send_notification_fn) - - send_result_fn(msg_id, {"task": public_task(task)}) \ No newline at end of file + + send_result_fn(msg_id, {"task": public_task(task)}) diff --git a/plan.md b/plan.md new file mode 100644 index 0000000..db0f241 --- /dev/null +++ b/plan.md @@ -0,0 +1,23 @@ +Let's review the code to ensure `asyncio.to_thread` works. `asyncio.to_thread` is available in Python 3.9+. The repository uses Python 3.12 (as seen from `pyenv` and `pyproject.toml`). + +Let's check `muninn/mcp/handlers.py` and modify it. + +We need to instantiate the global loop: +```python +_task_loop = asyncio.new_event_loop() +def _start_task_loop(): + asyncio.set_event_loop(_task_loop) + _task_loop.run_forever() + +threading.Thread(target=_start_task_loop, daemon=True, name="MuninnTaskLoop").start() +``` + +Where to put this? `muninn/mcp/handlers.py` has a global area: +```python +logger = logging.getLogger("Muninn.mcp.handlers") + +_thread_local = threading.local() +``` +I can place it right under `_thread_local = threading.local()`. + +Let's do this modification.