From 07e7a3f12e20e2bdc0e55c8e3bf2fcec18cfe952 Mon Sep 17 00:00:00 2001 From: alisson-anjos Date: Mon, 19 Oct 2026 10:59:47 +0000 Subject: [PATCH] Reserve an action finalization pass when agent thinking returns no answer --- bfs_agent_diagnostics.py | 3 + bfs_agent_runtime.py | 65 ++++++++++++++------ bfs_planner_agent.py | 3 +- docs/PLANNER_AGENT.md | 12 ++++ tests/test_planner_agent.py | 111 +++++++++++++++++++++++++++++++++- web/js/shotloop/view_agent.js | 1 + 6 files changed, 176 insertions(+), 19 deletions(-) diff --git a/bfs_agent_diagnostics.py b/bfs_agent_diagnostics.py index 6a0e233..c3239ba 100644 --- a/bfs_agent_diagnostics.py +++ b/bfs_agent_diagnostics.py @@ -16,6 +16,9 @@ def inspect_vlm(clip, settings, checkpoint=None): "agent_output_cap": 16384, "requested_output_tokens": settings["max_tokens"], "thinking_enabled": settings["thinking"], + "finalization_reserve_tokens": min(2048, settings["max_tokens"] // 4) + if settings["thinking"] + else 0, "warnings": [], "gpu": None, "kv_estimate": None, diff --git a/bfs_agent_runtime.py b/bfs_agent_runtime.py index cfc6bff..8bf9499 100644 --- a/bfs_agent_runtime.py +++ b/bfs_agent_runtime.py @@ -310,34 +310,65 @@ class Runtime: } ) prompt += f"\nBUDGET: renders {s.renders}/{s.settings['max_renders']}. Images are in the order specified by the observation. Return one tool action JSON." - with GPU_LOCK: - output = p._textgenerate_execute( + budget = s.settings["max_tokens"] + thinking = s.settings["thinking"] + reserve = min(2048, budget // 4) if thinking else 0 + images = visual_batch(s.images) + + def generate(text, limit, use_thinking): + return p._textgenerate_execute( TextGenerate, clip=clip, - prompt=prompt, - max_length=s.settings["max_tokens"], + prompt=text, + max_length=limit, sampling_mode={ "sampling_mode": "on", - "temperature": 0.6, + "temperature": 0.6 if use_thinking else 0.2, "top_k": 20, "top_p": 0.95, "min_p": 0.0, "seed": len(s.events), }, - thinking=s.settings["thinking"], - image=visual_batch(s.images), + thinking=use_thinking, + image=images, ) - text = str(output.args[0] or "").strip() - if not text: - reasoning_only = len(output.args) > 1 and bool(output.args[1]) - raise ValueError( - "VLM returned no final action" - + ( - " (thinking output only; the token budget may have been exhausted)" - if reasoning_only - else " (empty response)" + + with GPU_LOCK: + output = generate(prompt, budget - reserve, thinking) + text = str(output.args[0] or "").strip() + if not text and thinking: + if s.cancel.is_set() or s.status != "running": + return "" # tick checks cancellation before parsing + s.event( + "decision", + summary=( + "Thinking returned no final action. Requesting action JSON with Thinking off " + f"using the reserved {reserve} tokens; Thinking remains enabled for future decisions." + ), + phase="finalize_action", + token_budget=reserve, ) - + ". Increase Max tokens or disable Thinking, then Resume." + # Do not persist, display or feed truncated private reasoning back as instructions. + # Reuse the original task, observed state and images, with an explicit final-action request. + output = generate( + prompt + + "\nFINAL ACTION REQUIRED: The previous attempt returned no final action. " + "Return exactly one registered tool action JSON now, for example " + '{"tool":"inspect_source","arguments":{"shot":0},"summary":"Inspect source and reference"}. ' + "Choose the next useful action from the actual state above. Do not repeat the example blindly. " + "No prose, reasoning or Markdown. Do not claim a render happened unless the state confirms it.", + reserve, + False, + ) + text = str(output.args[0] or "").strip() + if not text: + raise ValueError( + "VLM returned no final action even with Thinking off during finalization. " + "Check the checkpoint/tokenizer compatibility or disable Thinking and Resume." + ) + if not text: + raise ValueError( + "VLM returned no final action with Thinking off. Check the checkpoint/tokenizer compatibility." ) return text diff --git a/bfs_planner_agent.py b/bfs_planner_agent.py index 8f5d3dd..90d1c66 100644 --- a/bfs_planner_agent.py +++ b/bfs_planner_agent.py @@ -296,9 +296,10 @@ class Session: ) return self.public() try: - action = parse_action(decide(self)) + answer = decide(self) if self.cancel.is_set() or self.status != "running": return self.public() + action = parse_action(answer) self.event("action", **action, tab=TABS.get(action["tool"], "agent")) if action["tool"] == "render": if self.renders >= self.settings["max_renders"]: diff --git a/docs/PLANNER_AGENT.md b/docs/PLANNER_AGENT.md index 12fc58e..917bf4b 100644 --- a/docs/PLANNER_AGENT.md +++ b/docs/PLANNER_AGENT.md @@ -141,3 +141,15 @@ The diagnostic does not load/offload models, allocate probe tensors or change se Numeric settings update during typing so periodic session polls do not overwrite an edit before the field loses focus. + +## Thinking that never produces a final action + +With Thinking enabled, each decision reserves up to one quarter of its output-token budget +(capped at 2,048 tokens) for finalization. If the first call returns no final action, the Agent +logs a `finalize_action` event and requests one action with Thinking off, using the original +task, observations and images. It does not feed back or persist truncated private reasoning. +The combined generation limits of both calls stay inside the configured decision budget: +8,192 tokens means at most 6,144 for the first call and 2,048 for finalization. +A valid first response needs no second call. The Thinking setting stays enabled for future +decisions. Pause/Stop between calls skips finalization; tool validation and render budgets +still apply. If finalization also fails, the Agent reports an error rather than inventing a tool action. diff --git a/tests/test_planner_agent.py b/tests/test_planner_agent.py index 55b9f83..ddb33be 100644 --- a/tests/test_planner_agent.py +++ b/tests/test_planner_agent.py @@ -76,10 +76,119 @@ class SessionTests(unittest.TestCase): ) session = self.session() with patch.dict(sys.modules, {"comfy_extras.nodes_textgen": fake}): - with self.assertRaisesRegex(ValueError, "thinking output only") as error: + with self.assertRaisesRegex( + ValueError, "Thinking off during finalization" + ) as error: runtime.Runtime(planner).decide(session) self.assertNotIn("private reasoning", str(error.exception)) + def test_thinking_finalization_respects_total_budget_and_hides_reasoning(self): + fake = types.ModuleType("comfy_extras.nodes_textgen") + fake.TextGenerate = object() + calls = [] + final = action("inspect_source", shot=0) + + def generate(*args, **kwargs): + calls.append(kwargs) + return ( + types.SimpleNamespace(args=("", "private reasoning")) + if len(calls) == 1 + else types.SimpleNamespace(args=(final, "")) + ) + + planner = types.SimpleNamespace( + _VLM={"clip": object()}, _textgenerate_execute=generate + ) + session = self.session(max_tokens=8192) + with patch.dict(sys.modules, {"comfy_extras.nodes_textgen": fake}): + result = runtime.Runtime(planner).decide(session) + self.assertEqual(core.parse_action(result)["tool"], "inspect_source") + self.assertEqual([c["thinking"] for c in calls], [True, False]) + self.assertEqual([c["max_length"] for c in calls], [6144, 2048]) + self.assertTrue(session.settings["thinking"]) + self.assertNotIn("private reasoning", calls[1]["prompt"]) + self.assertNotIn("private reasoning", json.dumps(session.public())) + self.assertEqual(session.events[-1]["phase"], "finalize_action") + + def test_valid_thinking_action_needs_no_finalization(self): + fake = types.ModuleType("comfy_extras.nodes_textgen") + fake.TextGenerate = object() + from unittest.mock import Mock + + generate = Mock( + return_value=types.SimpleNamespace(args=(action("inspect_plan"), "private")) + ) + planner = types.SimpleNamespace( + _VLM={"clip": object()}, _textgenerate_execute=generate + ) + session = self.session() + with patch.dict(sys.modules, {"comfy_extras.nodes_textgen": fake}): + runtime.Runtime(planner).decide(session) + self.assertEqual(generate.call_count, 1) + self.assertEqual(session.events, []) + + def test_no_second_call_when_thinking_is_disabled(self): + fake = types.ModuleType("comfy_extras.nodes_textgen") + fake.TextGenerate = object() + from unittest.mock import Mock + + generate = Mock(return_value=types.SimpleNamespace(args=("", ""))) + planner = types.SimpleNamespace( + _VLM={"clip": object()}, _textgenerate_execute=generate + ) + with patch.dict(sys.modules, {"comfy_extras.nodes_textgen": fake}): + with self.assertRaisesRegex(ValueError, "with Thinking off"): + runtime.Runtime(planner).decide(self.session(thinking=False)) + self.assertEqual(generate.call_count, 1) + self.assertEqual(generate.call_args.kwargs["max_length"], 4096) + + def test_pause_after_thinking_skips_finalization_and_parsing(self): + fake = types.ModuleType("comfy_extras.nodes_textgen") + fake.TextGenerate = object() + session = self.session() + calls = [] + + def generate(*args, **kwargs): + calls.append(kwargs) + session.cancel.set() + session.status = "paused" + return types.SimpleNamespace(args=("", "private")) + + planner = types.SimpleNamespace( + _VLM={"clip": object()}, _textgenerate_execute=generate + ) + with patch.dict(sys.modules, {"comfy_extras.nodes_textgen": fake}): + session.tick( + runtime.Runtime(planner).decide, + lambda *_: self.fail("must not execute"), + ) + self.assertEqual(len(calls), 1) + self.assertEqual(session.error_count, 0) + self.assertEqual(session.events, []) + + def test_finalization_does_not_bypass_registered_tool_validation(self): + fake = types.ModuleType("comfy_extras.nodes_textgen") + fake.TextGenerate = object() + from unittest.mock import Mock + + generate = Mock( + side_effect=[ + types.SimpleNamespace(args=("", "private")), + types.SimpleNamespace(args=(action("shell"), "")), + ] + ) + planner = types.SimpleNamespace( + _VLM={"clip": object()}, _textgenerate_execute=generate + ) + session = self.session() + with patch.dict(sys.modules, {"comfy_extras.nodes_textgen": fake}): + session.tick( + runtime.Runtime(planner).decide, + lambda *_: self.fail("must not execute"), + ) + self.assertIn("registered tool", session.observation["error"]) + self.assertEqual(session.renders, 0) + def test_plan_is_isolated(self): p = {"segs": [{"prompt": "old"}]} s = core.Session("task", p) diff --git a/web/js/shotloop/view_agent.js b/web/js/shotloop/view_agent.js index 7790e7d..6d8d267 100644 --- a/web/js/shotloop/view_agent.js +++ b/web/js/shotloop/view_agent.js @@ -105,6 +105,7 @@ export function agentTab(c) { h('div',`Checkpoint Thinking capability: ${a.vlmInfo.thinking_checkpoint||'unknown'}`), h('div',`Backend context: ${a.vlmInfo.context_tokens??'unknown'} tokens · Agent output cap: ${a.vlmInfo.agent_output_cap??'unknown'} · Requested output: ${a.vlmInfo.requested_output_tokens??'unknown'}`), a.vlmInfo.context_source?h('div',{class:'hint'},a.vlmInfo.context_source):null, + a.vlmInfo.finalization_reserve_tokens?h('div',{class:'hint'},`Thinking finalization reserve: ${a.vlmInfo.finalization_reserve_tokens} tokens within the decision budget; used only if the first call returns no final action.`):null, a.vlmInfo.gpu?h('div',`${a.vlmInfo.gpu.name}: ${a.vlmInfo.gpu.free_gib}/${a.vlmInfo.gpu.total_gib} GiB free · VLM weights not resident: ${a.vlmInfo.gpu.vlm_weights_not_resident_gib} GiB · Reserve: ${a.vlmInfo.gpu.reserve_gib} GiB`):null, a.vlmInfo.kv_estimate?h('div',[ h('div',`Estimated total KV capacity: ${a.vlmInfo.kv_estimate.total_tokens} tokens · Requested output KV: ${a.vlmInfo.kv_estimate.requested_output_cache_mib} MiB`),