[BUG] AgentServer runs sync handlers on the event loop, blocking all concurrent requests (incl. /health)

[!WARNING] Before submitting a PR, please make sure that:

• A maintainer has triaged this issue and applied the ready label
• This issue has no assignee
• No duplicate PR exists

PRs not meeting these requirements may be automatically closed.

Issues Policy acknowledgement

• I have read and agree to submit bug reports in accordance with the issues policy

Where did you encounter this bug?

Local serving via mlflow.genai.agent_server.AgentServer

MLflow version

• mlflow: 3.12.0 (bug also present on master — see code links below)

System information

• macOS 15 / Python 3.11.14, also reproducible on Linux
• uvicorn (default AgentServer.run path)

Describe the problem

AgentServer calls sync registered handlers directly on the asyncio event loop, so a single slow request blocks every concurrent request on the server — including /health.

In _handle_invoke_request:

if inspect.iscoroutinefunction(func):
    result = await func(request)
else:
    result = func(request)   # ← sync handler runs ON the event loop

And in _generate for streaming:

else:
    for chunk in func(request):   # ← sync generator iterated ON the event loop
        ...
        yield f"data: {json.dumps(chunk)}\n\n"

A sync (non-coroutine) handler — the natural shape for most agent frameworks, whose tool/LLM loops are synchronous — therefore freezes the loop for its entire duration. For streaming, the loop is blocked for the full time between two chunks (e.g. while a long-running tool executes).

Measured impact: with a handler containing a 60s tool call, a second concurrent request waited 66 s instead of ~5 s; even trivial endpoints on the same app (health checks, static routes mounted on server.app) were blocked.

Note the streaming case is doubly unfortunate: if the sync generator were handed directly to Starlette's StreamingResponse, Starlette would iterate it in its threadpool automatically — the async _generate wrapper defeats that built-in protection.

Suggested fix

• Non-streaming: result = await asyncio.to_thread(func, request) for non-coroutine handlers.
• Streaming: iterate sync generators off the loop. Two options:
  • starlette.concurrency.iterate_in_threadpool(func(request)), or
  • run the generator on a single dedicated worker thread and bridge chunks through an asyncio.Queue. The single-thread variant is safer for MLflow's own tracing: span context attach/detach via contextvars happens inside the generator, and resuming it on a different pool thread per __next__ can produce "token was created in a different Context" errors.

Happy to send a PR if maintainers agree on the approach.

Code to reproduce issue

"""repro.py — run, then execute the curl commands below."""
import time

from mlflow.genai.agent_server import AgentServer, invoke


@invoke()
def non_streaming(request):
    time.sleep(30)  # stands in for a slow tool / LLM call
    return {
        "output": [{
            "type": "message", "role": "assistant", "status": "completed",
            "id": "m-1",
            "content": [{"type": "output_text", "text": "done", "annotations": []}],
        }],
    }


server = AgentServer(agent_type="ResponsesAgent")
app = server.app

# uvicorn repro:app --port 8000

# Terminal 1 — slow request:
curl -s -X POST localhost:8000/invocations -H 'Content-Type: application/json' \
  -d '{"input":[{"role":"user","content":"hi"}]}' -w 'slow: %{time_total}s\n' &

# Terminal 2 (immediately after) — even /health is blocked:
curl -s localhost:8000/health -w 'health: %{time_total}s\n'

Expected: health returns in milliseconds. Actual: health waits ~30 s for the slow invocation to finish.

What component(s) does this bug affect?

• area/genai: LLM/GenAI serving (mlflow.genai.agent_server)

Workaround

Register async handlers that move the sync work off the loop themselves (asyncio.to_thread for invoke; a dedicated-thread + queue bridge for streaming generators). This works because the iscoroutinefunction / isasyncgenfunction branches avoid the blocking paths — but it shouldn't be required for the default sync-handler shape the decorators advertise.