fix: Run inference inside @spaces.GPU, where the weights actually exist

The Space was up, responsive, and emitting fluent multilingual noise at ~15
minutes a turn. Both symptoms were the same cause. Loading in the import window
got the model created and reported as cuda:0, but that is ZeroGPU's CUDA
*emulation*: it packs the tensors (17.6G, visible in the logs) and only
materialises them on real hardware for the duration of a @spaces.GPU call. Our
forward passes ran from a FastAPI route, so they read unmaterialised tensors --
which does not raise, it just returns garbage slowly.

app.py grows a stream_hook, None locally. app_space.py sets it to _gpu_stream, a
@spaces.GPU(duration=300) generator wrapping one whole agent turn.

The turn is the right unit. A turn is up to MAX_TOOL_STEPS generations plus the
tool calls between them, and they all mutate one KV cache; wrapping each
engine.stream() separately would put that shared state across a process boundary
on every step. One allocation per turn keeps it coherent, at the cost of not
reusing the cache between turns.

_collect() exists because ControlAgent.run() consumes self.stream directly and
would have bypassed the hook, leaving /api/chat running outside the GPU window
while /api/chat/stream ran inside it. It reads traces off the `done` event
rather than accumulating tool_end events, which also keeps the arguments that
run() dropped.

_gpu_probe stays, reduced to returning "ready": a @spaces.GPU function is only
detected when wired to a Gradio event handler, and _gpu_stream is driven by
FastAPI, so it cannot satisfy that itself.

Verified locally with a stand-in hook of the same shape that both /api/chat and
/api/chat/stream route through it, and that the MLX path is untouched when no
hook is installed.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

CLAUDE.md

@@ -210,6 +210,18 @@ routing transformers through `caching_allocator_warmup` and its direct
 unavailable**, which is what forces a model small enough to carry in bf16: `Qwen/Qwen3-8B`, ~16GB
 against ~28GB for the 14B run locally.
 
+**All inference must run inside a `@spaces.GPU` call.** Loading the model in the import window is
+necessary but not sufficient: ZeroGPU creates it under CUDA *emulation*, reports `cuda:0`, and packs
+its tensors — only a `@spaces.GPU` call materialises them on real hardware. Forward passes anywhere
+else do not fail. They read unmaterialised tensors and return fluent-looking multilingual noise at
+roughly fifteen minutes a turn, which is the worst possible failure mode: a demo that is up,
+responsive, and confidently wrong. `app.py` exposes a `stream_hook`, `None` locally; `app_space.py`
+sets it to `_gpu_stream`, a `@spaces.GPU(duration=300)` generator wrapping one whole agent turn.
+Wrapping the *turn* rather than each `engine.stream()` call is deliberate — a turn is several
+generations sharing one KV cache, and splitting them across separate calls would put that shared
+state on the far side of a process boundary each time. `_collect()` exists so `/api/chat` honours the
+hook too; `ControlAgent.run()` consumes `self.stream` directly and would bypass it.
+
 **ZeroGPU platform gotchas, each learned by having the Space fail:**
 - A `@spaces.GPU` function is only detected if wired to a real Gradio event handler. One called
   solely from a FastAPI route fails startup with "No @spaces.GPU function detected". Hence the

app.py

@@ -129,6 +129,41 @@ class ChatResponse(BaseModel):
     elapsed_seconds: float
 
 
+# Set by app_space.py to a @spaces.GPU-decorated generator. ZeroGPU attaches
+# real hardware only for the duration of such a call, so on the Space the whole
+# turn -- every tool step, every KV-cache mutation -- has to happen inside one.
+# Left None locally, where MLX needs no such thing.
+stream_hook = None
+
+
+def _stream_events(message: str, history: list[dict[str, str]]):
+    """The agent's raw events, through the GPU hook when one is installed."""
+    if stream_hook is not None:
+        yield from stream_hook(message, history)
+    else:
+        yield from get_agent().stream(message, history)
+
+
+def _collect(message: str, history: list[dict[str, str]]):
+    """`ControlAgent.run()` over `_stream_events`, so /api/chat honours the hook.
+
+    ControlAgent.run consumes self.stream directly, which would bypass the hook
+    and run inference outside the GPU window.
+    """
+    answer, plots, sources, stats = "", [], [], {}
+    traces: list[Any] = []
+    for event in _stream_events(message, history):
+        # The `done` event already carries the full traces, arguments included;
+        # accumulating tool_end events separately would only lose the arguments.
+        if event["type"] == "done":
+            answer = event["answer"]
+            traces = event["traces"]
+            plots = event["plots"]
+            sources = event["sources"]
+            stats = event["stats"]
+    return answer, traces, plots, sources, stats
+
+
 def _to_wire_events(message: str, history: list[dict[str, str]]):
     """Translate agent events into the shape the browser client consumes.
 
@@ -137,7 +172,7 @@ def _to_wire_events(message: str, history: list[dict[str, str]]):
     stable for the existing UI while the agent's own vocabulary stays clean.
     """
     thoughts: list[str] = []
-    for event in get_agent().stream(message, history):
+    for event in _stream_events(message, history):
         kind = event["type"]
         if kind == "text":
             yield {"type": "token", "content": event["text"]}
@@ -298,10 +333,10 @@ async def chat(req: ChatRequest) -> ChatResponse:
 
     def _run():
         with inference_lock:
-            return get_agent().run(message, req.history)
+            return _collect(message, req.history)
 
     try:
-        result = await _on_inference_thread(_run)
+        answer, traces, plots, _sources, _stats = await _on_inference_thread(_run)
     except Exception as exc:
         print(f"[chat] {type(exc).__name__}: {exc}")
         return ChatResponse(
@@ -310,9 +345,9 @@ async def chat(req: ChatRequest) -> ChatResponse:
         )
 
     return ChatResponse(
-        response=result.answer,
-        tool_traces=[{"tool": t.name, "status": t.status} for t in result.traces],
-        plots=result.plots,
+        response=answer,
+        tool_traces=[{"tool": t.get("tool"), "status": t.get("status")} for t in traces],
+        plots=plots,
         elapsed_seconds=round(time.time() - started, 2),
     )
 

app_space.py

@@ -28,6 +28,7 @@ it the retriever has no corpus and every answer falls back to model knowledge.
 from __future__ import annotations
 
 import os
+from typing import Iterator
 
 # Assigned, not setdefault: this module is the CUDA entry point by definition,
 # and a stale Space variable must not be able to select something else. One did
@@ -87,21 +88,52 @@ def _build_agent_at_import() -> None:
     # "[agent] retrieval failed". Force it onto the GPU here too. Embedding one
     # throwaway string is what actually triggers the load.
     get_embedder().encode_query("warmup")
-    print("[space] agent and embedder ready")
-
-
-_fetch_index()          # the agent prewarms retrieval, so the index must precede it
-_build_agent_at_import()
+    # Route every turn through the GPU-decorated generator above.
+    app_module.stream_hook = _gpu_stream
+    print("[space] agent and embedder ready, GPU stream hook installed")
+
+
+# One allocation per user turn, not per engine call. A turn is up to
+# MAX_TOOL_STEPS generations plus the tool executions between them, and they all
+# mutate the same KV cache; splitting them across separate @spaces.GPU calls
+# would put that shared state on the far side of a process boundary each time.
+# 300s is the ceiling for a turn that actually calls two tools.
+@spaces.GPU(duration=300)
+def _gpu_stream(message: str, history: list) -> Iterator[dict]:
+    """Run one whole agent turn with real hardware attached.
+
+    Everything about ControlAI's inference has to happen inside a call like this
+    one. ZeroGPU creates the model under CUDA *emulation* at import and packs its
+    tensors; only a @spaces.GPU call materialises them on a real device. Running
+    the forward passes anywhere else does not fail loudly -- it reads
+    unmaterialised tensors and returns fluent-looking multilingual noise at a few
+    minutes per turn, which is exactly what the Space did before this existed.
+
+    `message` and `history` are plain data. The agent is reached through the
+    module global and never passed in: ZeroGPU marshals arguments across a
+    process boundary and would try to share the model's CUDA tensors, hanging
+    with no output at all.
+    """
+    yield from app_module.get_agent().stream(message, history)
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=60)
 def _gpu_probe(text: str) -> str:
-    """Satisfies ZeroGPU's startup validation, and warms the model on first use.
+    """Satisfies ZeroGPU's startup validation.
 
-    Takes only a plain string. The agent is reached through the module global --
-    see note 4 above; passing it in is what hangs the whole Space.
+    A @spaces.GPU function is only detected if it is wired to a real Gradio event
+    handler, and _gpu_stream is driven by FastAPI rather than by Gradio, so it
+    cannot serve that purpose itself. This one exists to be wired to the hidden
+    button below.
     """
-    return app_module.get_agent().run(text).answer if text else "ready"
+    return "ready"
+
+
+# Both run at import, in this order: the index first (the agent touches retrieval
+# as it starts), then the agent -- which installs _gpu_stream, so it has to come
+# after that function exists.
+_fetch_index()
+_build_agent_at_import()
 
 
 with gr.Blocks() as _gpu_demo: