fix(272): CugaLite executes code that guesses at undeclared tool output shapes#417
Open
haroldship wants to merge 16 commits into
Open
fix(272): CugaLite executes code that guesses at undeclared tool output shapes#417haroldship wants to merge 16 commits into
haroldship wants to merge 16 commits into
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds weak-schema tool detection, observed-shape session tracking, probing-aware prompt enrichment, sandbox shape capture, and code extraction truncation that stops after the first tool call needing probing. ChangesWeak Tool Output Schema Probing
Estimated code review effort: 3 (Moderate) | ~30 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
5 tasks
coderabbitai
Bot
added
the
complexity: high
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.py`:
- Around line 23-38: Apply probing isolation in extract_and_combine_codeblocks
for the recovery paths, not just the fenced-block path. The current truncation
via _truncate_after_first_probing_block only runs on well-formed ```python
blocks, so update the fallback/recovery logic that handles unclosed fences or
raw Python to apply the same tools_needing_probing filtering before returning
code. Make sure any compile/recovery branch in extract_and_combine_codeblocks
uses the same first-probing-block cutoff so probe-triggering code cannot be
combined with later dependent code in one turn.
In `@src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py`:
- Around line 158-174: The is_weak_schema_tool() helper is overmatching real
string-returning tools because it treats any success schema shaped like {"type":
"string"} as the MCP placeholder. Update prompt_utils.py so the weakness check
relies on a stronger MCP-specific signal from the tool metadata or
response_schemas source, and keep get_tool_docs() from suppressing valid schemas
for genuine string-returning tools.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 6f2d4299-2095-47d3-a65a-12aca352d5bc
📒 Files selected for processing (17)
docs/superpowers/plans/2026-06-23-weak-tool-schema-probing-stage1.mddocs/superpowers/plans/2026-06-23-weak-tool-schema-probing-stage2.mddocs/superpowers/specs/2026-06-23-weak-tool-schema-probing-design.mdsrc/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.pysrc/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/graph_nodes.pysrc/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_nodes.pysrc/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/execution/test_code_extraction.pysrc/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_graph_adapter_hooks.pysrc/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_shared_call_model.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_agent_graph_adapter.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_prepare_node_weak_schema_tools.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_prompt_utils_weak_schema.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_sandbox_node_weak_schema_shapes.py
prepare_node.py was filtering for weak-schema tools over tools_for_prompt, which collapses to just the find_tools meta-tool whenever the catalog exceeds shortlisting_tool_threshold (~prepare_node.py:230) -- the normal case once an app has more than a handful of tools. That made _weak_schema_tool_names permanently empty in any realistic deployment, silently disabling get_tools_needing_probing() (block-isolation enforcement) and _record_weak_schema_shapes (session shape-memory) for every tool actually reachable through find_tools. Only the advisory per-tool directive rendered inside find_tools' own output (which checks the real tool catalog directly, bypassing this cache) was ever active. Switch the source to tools_for_execution, the full catalog unaffected by the find_tools-shortlisting collapse. Add a regression test that exercises the shortlisting-active path (the previous test only covered shortlisting-inactive, where tools_for_prompt == tools_for_execution anyway and couldn't have caught this). Found while running the Vakra(M3) eval at broader scope (10 tasks across 9 domains) for this PR: traces showed the model writing multi-block turns that combined a probe call with code assuming its shape, with all blocks executing together -- exactly what Component B is supposed to prevent.
haroldship
force-pushed
the
feat/272-stage2-block-isolation
branch
from
June 30, 2026 13:55
bb91fbf to
70a84f9
Compare
coderabbitai
Bot
added
complexity: medium
- code_extraction: apply probing isolation to the recovery paths too. The fenced path truncates at block boundaries, but an unclosed fence (#204) or raw print() blob collapses would-be-separate blocks into one, so a probe call could still run with later dependent code. Add a line-level cutoff for those single-blob recoveries; add regression tests. - prompt_utils/mcp_manager: stop is_weak_schema_tool matching on schema shape. A genuine string-returning tool (OpenAPI text body, or an MCP tool declaring outputSchema {"type": "string"}) has a success schema identical to the synthetic no-schema placeholder, so shape-matching suppressed real schemas. The MCP manager now tags the placeholder it injects (_synthetic_placeholder) and the weak-schema check trusts that marker; update tests.
coderabbitai
Bot
added
complexity: high
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py (1)
210-214: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHide the synthetic placeholder from
find_tools()output too.This now swaps
response_docto the probing directive, butfind_tools()still copiesresponse_schemas["success"]intoTool.output_schemafor the same weak-schema tools. The result is that the markdown can show both the new directive and the fake{"type": "string"}schema, which defeats the isolation guidance this PR is adding.Suggested follow-up
- if hasattr(actual_tool, 'func') and hasattr(actual_tool.func, '_response_schemas'): + if ( + hasattr(actual_tool, 'func') + and hasattr(actual_tool.func, '_response_schemas') + and not PromptUtils.is_weak_schema_tool(actual_tool) + ): response_schemas = actual_tool.func._response_schemas if response_schemas and isinstance(response_schemas, dict) and 'success' in response_schemas: raw_output_schema = response_schemas['success']Please also add a regression around the
find_tools()markdown path so weak-schema tools do not emit the placeholder anywhere.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py` around lines 210 - 214, The weak-schema tool handling in PromptUtils still leaks the synthetic success schema through find_tools() markdown generation, so update the find_tools() path to suppress Tool.output_schema for PromptUtils.is_weak_schema_tool(tool) just like response_doc is overridden. Make the fix in the same tool-discovery flow that builds markdown from response_schemas so weak-schema tools never emit the placeholder schema anywhere, and add a regression test covering find_tools() output for a weak-schema tool to verify only the probing directive appears.src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.py (1)
39-40: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winSame-block probe + dependent code still slips through.
If a single fenced
pythonblock contains both the probing call and later shape-dependent code,_truncate_after_first_probing_block()keeps that whole block. That defeats the isolation guarantee this change is trying to add, because the probe and the dependent slicing still run in one turn. Truncate within the first matching block, not just after it.Also applies to: 63-70
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.py` around lines 39 - 40, The current truncation logic in _truncate_after_first_probing_block only drops blocks after the first probing block, so a single fenced python block can still contain both the probe and later shape-dependent code. Update the code-extraction flow in this path to truncate within the first matching block itself, splitting or cutting the block immediately after the probing call so dependent code never survives in the same turn. Use the existing tools_needing_probing check and the block-processing logic around _truncate_after_first_probing_block to locate where the in-block cutoff should happen.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.py`:
- Around line 73-84: The line-based cutoff in _truncate_after_first_probing_line
is too aggressive for multiline probing statements, so it can return invalid
Python before the probe actually completes. Update the logic to detect the first
probing call and truncate after the full statement boundary instead of the first
matching line, preserving multiline calls in recovered code while still removing
any later dependent code.
---
Outside diff comments:
In
`@src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.py`:
- Around line 39-40: The current truncation logic in
_truncate_after_first_probing_block only drops blocks after the first probing
block, so a single fenced python block can still contain both the probe and
later shape-dependent code. Update the code-extraction flow in this path to
truncate within the first matching block itself, splitting or cutting the block
immediately after the probing call so dependent code never survives in the same
turn. Use the existing tools_needing_probing check and the block-processing
logic around _truncate_after_first_probing_block to locate where the in-block
cutoff should happen.
In `@src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py`:
- Around line 210-214: The weak-schema tool handling in PromptUtils still leaks
the synthetic success schema through find_tools() markdown generation, so update
the find_tools() path to suppress Tool.output_schema for
PromptUtils.is_weak_schema_tool(tool) just like response_doc is overridden. Make
the fix in the same tool-discovery flow that builds markdown from
response_schemas so weak-schema tools never emit the placeholder schema
anywhere, and add a regression test covering find_tools() output for a
weak-schema tool to verify only the probing directive appears.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 6a838e98-dcf6-4b16-9af3-f6a96b0685d3
📒 Files selected for processing (6)
src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.pysrc/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/execution/test_code_extraction.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_prepare_node_weak_schema_tools.pysrc/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_prompt_utils_weak_schema.pysrc/cuga/backend/tools_env/registry/mcp_manager/mcp_manager.py
🚧 Files skipped from review as they are similar to previous changes (2)
- src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_prepare_node_weak_schema_tools.py
- src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_prompt_utils_weak_schema.py
| @@ -0,0 +1,637 @@ | |||
| # Weak Tool Output Schema Probing — Stage 1 (A+C) Implementation Plan | |||
| @@ -0,0 +1,222 @@ | |||
| --- | |||
| @@ -0,0 +1,222 @@ | |||
| --- | |||
…cation
The recovery paths (unclosed fence / raw Python) cut the code blob at the
first line matching a probing-tool call. A probing call spanning multiple
lines was therefore truncated mid-statement (e.g. 'res = await file_readfile('),
yielding invalid Python so the probe never executed and Stage 2 isolation was
bypassed. Grow the cut to the end of the first probing statement via an
incremental compile check, matching the recovery path's existing approach.
Also removes the docs/superpowers design spec and stage 1/2 implementation
plans per review.
coderabbitai
Bot
added
the
complexity: medium
coderabbitai
Bot
removed
the
complexity: high
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.py (1)
92-99: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider extracting the shared compile-guard logic.
The "strip
await, thencompile(..., exec)" pattern here duplicates the same technique in_recover_non_closing_python_fence(line 169). A small shared helper (e.g._compiles_as_exec(code: str) -> bool) would avoid the duplication and keep the await-stripping hack in one place.♻️ Proposed helper extraction
+def _compiles_as_exec(code: str) -> bool: + """Return True if `code` compiles as a standalone module (await-tolerant).""" + try: + compile(code.replace("await ", ""), "<string>", "exec") + return True + except SyntaxError: + return False + + def _truncate_after_first_probing_line(code: str, tools_needing_probing: frozenset[str]) -> str: ... for end in range(i + 1, len(lines) + 1): candidate = "\n".join(lines[:end]) - try: - compile(candidate.replace("await ", ""), "<string>", "exec") - except SyntaxError: - continue - return candidate + if _compiles_as_exec(candidate): + return candidate🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.py` around lines 92 - 99, The compile-guard logic is duplicated between this extraction path and `_recover_non_closing_python_fence`; factor the repeated “strip `await `, then `compile(..., "exec")`” check into a shared helper such as `_compiles_as_exec(code: str) -> bool`. Update the loop in the code extraction routine to call that helper instead of inlining the try/except, and reuse the same helper from `_recover_non_closing_python_fence` so the await-stripping behavior lives in one place.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In
`@src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.py`:
- Around line 92-99: The compile-guard logic is duplicated between this
extraction path and `_recover_non_closing_python_fence`; factor the repeated
“strip `await `, then `compile(..., "exec")`” check into a shared helper such as
`_compiles_as_exec(code: str) -> bool`. Update the loop in the code extraction
routine to call that helper instead of inlining the try/except, and reuse the
same helper from `_recover_non_closing_python_fence` so the await-stripping
behavior lives in one place.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 19d2f483-5103-4709-8d0f-e194079bc9d6
📒 Files selected for processing (2)
src/cuga/backend/cuga_graph/nodes/cuga_agent_core/execution/code_extraction.pysrc/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/execution/test_code_extraction.py
The two live-LLM tool-approval e2e tests assumed the model's first turn is the
final code that directly references digital_sales. With weak-schema probing on
(this PR), the first turn is a structure probe — the narration is natural
language ("…to inspect the returned structure") and the first approved block is
a fetch-and-print probe identical across requests (the task-specific refinement
lands in a later turn). Two brittle assertions failed as a result.
- approve_flow: assert on the *code pending approval* (what triggered the
policy) instead of the last AI chat message, which can be NL narration.
- modification_flow: drop the first-block divergence check — under probing both
requests' first approved blocks are the same structure probe by design; the
tool-reference assertion already confirms a real digital_sales call.
Verified live 3x: all 3 tests in the file pass consistently.
Collaborator
Author
|
Filed #427 for the pre-existing test-harness gap surfaced here: the tool-approval e2e helper ( |
haroldship
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The problem (#272)
CugaLite writes Python code that calls tools. To write correct code it needs to know what each tool returns — the output schema. Some tools don't declare one, and for MCP tools with no declared schema the registry was making one up: it injected a placeholder claiming the tool returns a string.
So the model confidently wrote code like this, all in one script:
By the time the wrong assumption blew up, the whole script had already run. The model was guessing the shape and executing the guess in one shot. That's the literal failure trace in #272.
The fix, in three parts
Part A — be honest in the prompt. Detect tools with no real declared schema: an empty schema, or the synthetic MCP placeholder — which the MCP manager now tags with a
_synthetic_placeholdermarker so it can't be confused with a tool that genuinely returns a bare string. Instead of rendering the fake schema, the prompt tells the model: call this tool alone in its own block andprint()the raw result; write follow-up code only after you've seen the real shape.Part B — enforce it, don't just ask. Part A is only a request; the model can still write "probe + code that uses the probe" in one response. So during code extraction, if the response contains a call to a still-unknown-schema tool, everything after that call is cut off. Only the probe runs; the model sees the real result and writes the rest next turn with actual knowledge instead of a guess. This lives behind a new adapter hook,
CoreGraphAdapter.get_tools_needing_probing(), which defaults tofrozenset()— CugaSupervisor and any adapter that hasn't opted in behave exactly as before.Part C — don't probe twice. The first time an unknown-schema tool returns a result, the sandbox records a short description of its shape (
"list of 3 items","dict with keys ['id', 'name']"). That description is injected into the system prompt on every later turn ("Observed tool output shapes — use this, no need to probe again"), and the tool drops out of Part B's enforcement set. The one-call-at-a-time restriction only applies until the shape is known.Implementation notes (the fiddly bits)
prepare_nodebuilds the unknown-schema set fromtools_for_execution, nottools_for_prompt. Whenfind_toolsshortlisting is active (the common case once an app has more than a handful of tools),tools_for_promptcollapses to just thefind_toolsmeta-tool — using it would have left the set permanently empty and silently disabled Part B and Part C for every tool actually reachable throughfind_tools.track_tool_calls— without leaking the tracked calls into the caller'stool_callsstate. Without this, Part C never recorded anything in production (the Stage 1 A+C branch was dead on arrival; there's a regression test that reproduces this)._synthetic_placeholdermarker, not schema shape. A genuine string-returning tool (OpenAPI text body, or an MCP tool that actually declaresoutputSchema: {"type": "string"}) produces asuccessschema byte-identical to the placeholder; matching on shape would have suppressed real schemas. The marker key is a plain literal on both sides to avoid a graph→registry import dependency.Staging
This PR is the union of two branches, kept separate during development so each component's contribution could be evaluated independently:
feat/272-tool-output-schema-probing— Stage 1, Parts A+C onlyfeat/272-stage2-block-isolation(this branch, off Stage 1) — adds Part BEval results (Vakra M3, 1 task, 5 runs each)
A+C alone (asking nicely) already cuts tokens/calls/time roughly in half and raises ExactMatch substantially, but doesn't move Pass@1. Adding B (forcing the split) is what converts that into actual task success — the enforcement is the part that works.
Test plan
prompt_utils,prepare_node,sandbox_node,graph_adapter) — new + extendedcode_extraction,graph_nodes,shared_nodes,graph_adapter) — new + extended, including an end-to-end reproduction of the issue's literal original multi-block tracecuga_lite+cuga_agent_core+cuga_supervisorsuites green (363 passed, 1 skipped, 3 pre-existing unrelated e2b-package-not-installed failures)ruff checkcleanSummary by CodeRabbit