fix(agent): free the LLM's VRAM on the normal path after a workflow runs

api/routers/agent.py

@@ -420,6 +420,25 @@ async def list_ollama_models(ollama_url: str = "http://localhost:11434"):
             return {"models": []}
 
 
+async def _unload_llm_after_workflow(
+    client: httpx.AsyncClient,
+    request: AgentChatRequest,
+    actions_done: list[ActionDone],
+) -> None:
+    """Free the Ollama model's VRAM once a workflow has been dispatched, so the
+    workflow gets the full GPU. Best-effort — never fail the chat over it."""
+    if not any(a.tool == "run_workflow" for a in actions_done):
+        return
+    try:
+        await client.post(
+            f"{request.ollama_url}/api/generate",
+            json={"model": request.model, "keep_alive": 0},
+            timeout=5.0,
+        )
+    except Exception:
+        pass
+
+
 @router.post("/chat", response_model=AgentChatResponse)
 async def agent_chat(request: AgentChatRequest):
     messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
@@ -488,6 +507,7 @@ async def agent_chat(request: AgentChatRequest):
 
             tool_calls = msg.get("tool_calls") or []
             if not tool_calls:
+                await _unload_llm_after_workflow(client, request, actions_done)
                 combined_thinking = "\n\n---\n\n".join(all_thinking) if all_thinking else None
                 return AgentChatResponse(
                     message=clean_content,
@@ -501,17 +521,8 @@ async def agent_chat(request: AgentChatRequest):
                 actions_done.append(ActionDone(tool=fn["name"], result=result_text, payload=payload))
                 messages.append({"role": "tool", "content": result_text})
 
-        has_workflow = any(a.tool == "run_workflow" for a in actions_done)
-        if has_workflow:
-            # Unload LLM from VRAM immediately so the workflow has full GPU memory
-            try:
-                await client.post(
-                    f"{request.ollama_url}/api/generate",
-                    json={"model": request.model, "keep_alive": 0},
-                    timeout=5.0,
-                )
-            except Exception:
-                pass
+        # Loop exhausted without a final answer: still free VRAM if a workflow ran.
+        await _unload_llm_after_workflow(client, request, actions_done)
 
     combined_thinking = "\n\n---\n\n".join(all_thinking) if all_thinking else None
     return AgentChatResponse(message="Reached maximum tool iterations.", actions=actions_done, thinking=combined_thinking)

api/tests/test_agent_workflow_unload.py

@@ -0,0 +1,77 @@
+import asyncio
+import json
+import unittest
+from unittest import mock
+
+import httpx
+
+import routers.agent as agent
+
+
+class _ScriptedOllama:
+    """MockTransport wiring: serves a scripted /api/chat sequence and records
+    every /api/generate (VRAM keep_alive) call agent_chat makes."""
+
+    def __init__(self, chat_bodies) -> None:
+        self._chat = list(chat_bodies)
+        self.generate_calls: list[dict] = []
+        self._real = httpx.AsyncClient
+
+    def _handler(self, request: httpx.Request) -> httpx.Response:
+        path = request.url.path
+        if path == "/api/chat":
+            return httpx.Response(200, json=self._chat.pop(0))
+        if path == "/api/generate":
+            self.generate_calls.append(json.loads(request.content))
+            return httpx.Response(200, json={})
+        return httpx.Response(404, json={})
+
+    def __call__(self, *args, **kwargs):
+        kwargs["transport"] = httpx.MockTransport(self._handler)
+        return self._real(*args, **kwargs)
+
+
+def _assistant(content: str = "", tool_calls=None) -> dict:
+    msg: dict = {"role": "assistant", "content": content}
+    if tool_calls:
+        msg["tool_calls"] = tool_calls
+    return {"message": msg}
+
+
+def _run(chat_bodies, context) -> _ScriptedOllama:
+    scripted = _ScriptedOllama(chat_bodies)
+    request = agent.AgentChatRequest(
+        messages=[agent.ChatMessage(role="user", content="make me a thing")],
+        ollama_url="http://ollama.test",
+        model="llama-test",
+        context=context,
+    )
+    with mock.patch.object(agent.httpx, "AsyncClient", scripted):
+        asyncio.run(agent.agent_chat(request))
+    return scripted
+
+
+class WorkflowVramUnloadTests(unittest.TestCase):
+    def test_llm_unloaded_on_the_normal_return_after_a_workflow(self) -> None:
+        # Round 1 dispatches a workflow; round 2 is the final answer (no tools) and
+        # returns early. The VRAM unload must still fire — before the fix it lived
+        # after the loop and this common path skipped it entirely.
+        chat = [
+            _assistant(tool_calls=[
+                {"function": {"name": "run_workflow", "arguments": {"workflow_id": "wf1"}}},
+            ]),
+            _assistant(content="Running your workflow now."),
+        ]
+        scripted = _run(chat, {"workflows": [{"id": "wf1", "name": "My Workflow"}]})
+        self.assertTrue(
+            any(call.get("keep_alive") == 0 for call in scripted.generate_calls),
+            f"expected a keep_alive:0 unload, got {scripted.generate_calls}",
+        )
+
+    def test_no_unload_when_no_workflow_was_dispatched(self) -> None:
+        scripted = _run([_assistant(content="Here is some info.")], {})
+        self.assertEqual(scripted.generate_calls, [])
+
+
+if __name__ == "__main__":
+    unittest.main()