diff --git a/CHANGELOG.md b/CHANGELOG.md index dd24bfa..8fe3543 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/lcode/tools.py b/src/lcode/tools.py index 1a6684d..f343f3a 100644 --- a/src/lcode/tools.py +++ b/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( diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 4196088..fae9ca3 100644 --- a/tests/test_mcp.py +++ b/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()