Make GPU tools run to completion

CHANGELOG.md

@@ -6,6 +6,12 @@ All notable changes to lcode are documented here. The format follows
 
 ## [Unreleased]
 
+### Fixed
+
+- Tools marked `free_gpu` (such as ComfyUI's `run_workflow`) now always run to completion
+  (`wait: true`): a job left running in the background competed with lcode's reloading model for
+  the GPU.
+
 ## [0.6.0] - 2026-10-01
 
 ### Added

src/lcode/tools.py

@@ -365,6 +365,11 @@ def _mcp(self, name: str, args: dict) -> str:
             if not ok:
                 return feedback
         if mcp.needs_gpu(state, tool):
+            # The tool must finish before lcode's model comes back: a job left running in the
+            # background would compete with it for the GPU.
+            properties = (tool.get("inputSchema") or {}).get("properties") or {}
+            if (properties.get("wait") or {}).get("type") == "boolean" and arguments.get("wait") is not True:
+                arguments = {**arguments, "wait": True}
             freed = self.agent.free_gpu()
             if freed:
                 self.console.print(

tests/test_mcp.py

@@ -540,3 +540,19 @@ def test_gpu_hungry_tools_get_the_gpu(make_agent, tmp_path, monkeypatch):
     assert "freed the GPU for img (lcode-qwen3.6-35b reloads afterwards)" in output(agent)
     assert mcp_config.parse("x", {"command": "y", "free_gpu": True}).free_gpu is True
     agent.mcp.close()
+
+
+def test_gpu_tools_run_to_completion(make_agent, tmp_path, monkeypatch):
+    """A free_gpu tool with a `wait` option always waits, so it doesn't share the GPU with lcode's model."""
+    agent = make_agent([reply(tool_calls=[call("mcp__img__add", a=1, b=2, wait=False)]), reply("done")])
+    agent.mcp = ready_manager(tmp_path, server_config("img", free_gpu=["add"]))
+    state = agent.mcp.servers["img"]
+    add = next(t for t in state.tools if t["name"] == "add")
+    add["inputSchema"]["properties"]["wait"] = {"type": "boolean", "default": True}
+    sent = []
+    monkeypatch.setattr(
+        agent.mcp, "call", lambda state, tool, arguments, image_text=None: sent.append(arguments) or "ok"
+    )
+    agent.run_turn("make an image")
+    assert sent == [{"a": 1, "b": 2, "wait": True}]
+    agent.mcp.close()