Tool guard#302
Open
naamaz wants to merge 6 commits into
Open
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR introduces a comprehensive tool guard policy framework that enables per-tool compliance enforcement. It adds data models for tool-specific guards, updates policy persistence layers, implements build-time generation of policy examples and code, provides runtime enforcement via ChangesTool Guard Policy Enforcement System
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the 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 |
![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: 15
🤖 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/policy/folder_loader.py`:
- Around line 215-221: The code silently ignores non-dict frontmatter for
tool_guards; update the block that reads frontmatter.get('tool_guards') to
validate the shape and fail fast: if tool_guards_data is None do nothing, if
it's a dict continue to construct tool_guards via ToolGuard(**guard_data) as
now, but if tool_guards_data exists and is not a dict raise a clear exception
(e.g., ValueError) with a message referencing 'tool_guards' and the offending
type so malformed frontmatter is rejected rather than silently ignored.
In `@src/cuga/backend/cuga_graph/policy/storage.py`:
- Around line 185-187: The current code appends raw guard.violating_examples and
guard.compliance_examples into text_parts with no length limits, which can
exceed embedding input size; before joining each example list (used where
text_parts is built), truncate each example string to a safe max length (e.g.,
1000 chars) and/or limit the number of characters per joined block, and then
append the truncated/joined strings instead of the raw ones; implement this by
adding a small helper (e.g., truncate(text, max_chars)) or inline slicing when
building the strings that reference text_parts, guard.violating_examples, and
guard.compliance_examples so large examples are capped prior to
embedding/storage.
In
`@src/cuga/backend/cuga_graph/policy/tool_guard/tests/test_crm_finance_tool_guard_e2e.py`:
- Around line 407-424: Update the E2E to assert that the Finance tool guard
actually prevented the tool call: replace the loose asserts on
result1_initial.answer and result1_followup.answer with checks that the agent
response indicates a blocked/refused outcome (e.g., contains a specific refusal
message) AND/OR verify the CRM tool function crm_create_account_accounts_post
was not executed by instrumenting/mocking that function and asserting it was not
called; additionally add focused unit tests in the appropriate tests/ sub-folder
that mock the CRM client and the tool-guard logic to cover the blocked path, the
allowed path, and edge cases so each branch of the guard is validated.
- Around line 11-20: The test module currently sets os.environ entries at import
time (in test_crm_finance_tool_guard_e2e.py), which leaks config across tests;
instead remove those top-level os.environ writes and create a pytest fixture
(e.g., autouse or per-test) that uses monkeypatch.setenv to apply
MCP_SERVERS_FILE, CUGA_MANAGER_MODE, DYNACONF_* values before importing any CUGA
modules, and ensure all imports of cuga/backend/cuga_graph/... happen inside
that fixture or inside tests after the environment is patched so the config is
isolated and restored for other tests.
- Around line 217-285: _build_agent currently starts multiple service processes
(start_email, start_filesystem, start_crm, start_registry) and then performs
provider.initialize(), provider.get_apps(), and CugaAgent construction without
ensuring cleanup on failure; wrap the post-startup initialization (from after
starting any service through until after await
agent.policies._ensure_policy_system()) in a try/except/finally that on
exception/exit stops all started services and frees ports (e.g., call a teardown
method on AppManager such as stop_all/stop_services if available, otherwise use
the existing helpers like _kill_ports and _kill_proc to terminate registry_proc
and any started app processes), then re-raise the exception so the caller’s
finally can proceed—this ensures any partial boot will not leave ports/processes
running.
In
`@src/cuga/backend/cuga_graph/policy/tool_guard/tests/test_flight_booking_tool_guard.py`:
- Around line 14-20: The test module currently mutates process-wide environment
variables at import time (setting MCP_SERVERS_FILE, CUGA_MANAGER_MODE,
DYNACONF_* keys) which affects other tests; change this to a fixture that uses
pytest's monkeypatch to set those env vars for the test scope and restore them
afterwards. Create a fixture (e.g., def env_cleaner(monkeypatch): ...) and
inside call monkeypatch.setenv for each key used in the file instead of
assigning to os.environ at module level, then mark tests to use that fixture (or
set scope='function'/'module' as appropriate) so the environment is isolated and
restored automatically. Ensure you remove the top-level os.environ assignments
from the module so no mutation happens at import time.
- Line 35: Change the test to exercise the enforcement/runtime path by setting
USE_TOOLGUARD = True and making policy generation deterministic (replace or
inject a fixed policy_code fixture used in the runtime path instead of relying
on the current conditional), then replace print-only checks in
test_flight_booking_tool_guard with concrete assertions: assert the initial
booking attempt is blocked (inspect the response payload for the block/error
marker or assert the booking tool/mocked book_flight function was not called)
and assert the follow-up booking attempt is also blocked; update the assertions
in the block covering lines ~130-261 accordingly and/or add mocks for the tool
executor to verify no tool invocation occurred.
In
`@src/cuga/backend/cuga_graph/policy/tool_guard/tests/test_tool_guard_generation.py`:
- Around line 109-133: The test currently catches all exceptions around the call
to agent.policies.generate_tool_guard_examples and converts them into
pytest.skip, which masks real failures; update the try/except so only known
environment/precondition errors (e.g., ImportError, ModuleNotFoundError, a
specific LLM configuration error class if present) are caught and cause
pytest.skip, while letting AssertionError and any other unexpected exceptions
propagate (or be re-raised) so the test fails normally; change the except block
around the generate_tool_guard_examples call (and the analogous block at the
later location) to explicitly catch the small set of expected environment
exceptions and call pytest.skip with the error, and remove the broad except
Exception handler so other errors are not swallowed.
In
`@src/cuga/backend/cuga_graph/policy/tool_guard/tests/test_tool_guard_runtime_e2e.py`:
- Around line 417-440: Wrap the create/process/initialize/run/cleanup sequence
in a try/finally so cleanup_policies(...) and tool_guard_runtime.shutdown()
always run: create the policy_data and instantiate ToolGuardRuntime, call await
tool_guard_runtime.initialize() and await run_tests(...) inside try, and in
finally await cleanup_policies(agent, policy_system, policy_data) if policy_data
is set and await tool_guard_runtime.shutdown() if tool_guard_runtime was created
(guard against None if initialization failed). Ensure you still assign
policy_data and tool_guard_runtime before the try so they can be referenced in
the finally and await both cleanup_policies and shutdown.
- Line 30: Current test uses DELETE_ALL_POLICIES_AT_START and calls
cleanup_all_policies() which wipes every policy and the toolguard domain
directory; change the test to avoid global destructive cleanup by creating and
pointing the code under test to an isolated temporary storage/cuga directory
(use tempfile.TemporaryDirectory) or by modifying cleanup_all_policies() usage
to only delete policies created by this test (track created policy IDs and pass
them to a limited cleanup function). Concretely, update the test setup to set
the environment/config used by the code (the same place that
DELETE_ALL_POLICIES_AT_START reads) to a temp path before any cleanup runs, or
replace the global call to cleanup_all_policies() with a helper that accepts a
list of IDs gathered during test creation and deletes only those IDs.
In `@src/cuga/backend/cuga_graph/policy/tool_guard/tool_guard_buildtime.py`:
- Around line 109-112: The validation currently rejects concrete target_tool
values when the policy declares target_tools=["*"]; update the check in the
block that raises ValueError so it first allows the wildcard by treating "*" in
policy.target_tools as matching any target_tool (i.e., only raise if "*" not in
policy.target_tools and target_tool not in policy.target_tools). Modify the
logic used by add_tool_guide()/generate_examples()/generate_guard_code() that
references policy.target_tools and target_tool so wildcard policies pass
validation.
In `@src/cuga/backend/cuga_graph/policy/tool_guard/tool_guard_runtime.py`:
- Around line 111-120: The initialize() method currently loads policies and
builds tool_to_guards whenever self.policy_storage is set, ignoring the
enable_policies flag; change initialize() to check enable_policies first and
only call self.policy_storage.list_policies(...) and
self._build_tool_to_guards_mapping(...) when enable_policies is truthy (and
policy_storage is present). In short: wrap the list_policies (with
PolicyType.TOOL_GUIDE) and the
_build_tool_to_guards_mapping(tool_guide_policies) call behind the same
enable_policies feature gate so guards are not loaded when enable_policies is
False.
In `@src/cuga/sdk.py`:
- Around line 1293-1298: The policy type check is directly accessing policy.type
which raises AttributeError when the object only exposes policy_type; change
both checks to use the same accessor pattern used by list() and get(): read the
policy type via getattr(policy, "policy_type", getattr(policy, "type", None))
(or equivalent nil-safe accessor) and compare that value to
PolicyType.TOOL_GUIDE; update the check in the shown block and the similar block
around the other occurrence (the second validator near the later lines) so both
places handle objects with either policy_type or type.
- Around line 1279-1303: The code creates a ToolGuardBuildtimeManager and calls
await manager._ensure_initialized() without first ensuring the agent has been
initialized, which can race with custom tool providers; call await
self._agent._ensure_initialized() before instantiating ToolGuardBuildtimeManager
(i.e., insert await self._agent._ensure_initialized() right before manager =
ToolGuardBuildtimeManager(self._agent)), and apply the same change to the other
occurrence referenced (the block around lines 1378-1402) so the agent is always
initialized before build-time manager initialization.
- Around line 591-598: The code is currently accepting any key in tool_guards
and persisting ToolGuard objects even for tools not in
existing_policy.target_tools; update the conversion loop that builds
tool_guards_obj so it validates each tool_name against
existing_policy.target_tools and rejects invalid entries by raising a clear
error (e.g., ValueError) naming the offending tool(s) and the allowed
target_tools; ensure this check is applied before creating ToolGuard instances
and that the error message references tool_name and existing_policy.target_tools
for debugging.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6991e90b-5ab3-42b4-b026-26c2b4b46c03
📒 Files selected for processing (15)
pyproject.tomlsrc/cuga/backend/cuga_graph/policy/filesystem_sync.pysrc/cuga/backend/cuga_graph/policy/folder_loader.pysrc/cuga/backend/cuga_graph/policy/models.pysrc/cuga/backend/cuga_graph/policy/storage.pysrc/cuga/backend/cuga_graph/policy/tool_guard/tests/__init__.pysrc/cuga/backend/cuga_graph/policy/tool_guard/tests/test_crm_finance_tool_guard_e2e.pysrc/cuga/backend/cuga_graph/policy/tool_guard/tests/test_flight_booking_tool_guard.pysrc/cuga/backend/cuga_graph/policy/tool_guard/tests/test_tool_guard_generation.pysrc/cuga/backend/cuga_graph/policy/tool_guard/tests/test_tool_guard_policy.pysrc/cuga/backend/cuga_graph/policy/tool_guard/tests/test_tool_guard_runtime_e2e.pysrc/cuga/backend/cuga_graph/policy/tool_guard/tool_guard_buildtime.pysrc/cuga/backend/cuga_graph/policy/tool_guard/tool_guard_runtime.pysrc/cuga/backend/cuga_graph/policy/tool_guard/tool_invoker.pysrc/cuga/sdk.py
Comment on lines
+215
to
+221
| tool_guards_data = frontmatter.get('tool_guards') | ||
| if tool_guards_data and isinstance(tool_guards_data, dict): | ||
| from cuga.backend.cuga_graph.policy.models import ToolGuard | ||
| tool_guards = { | ||
| tool_name: ToolGuard(**guard_data) | ||
| for tool_name, guard_data in tool_guards_data.items() | ||
| } |
Contributor
There was a problem hiding this comment.
Fail fast on malformed tool_guards frontmatter.
At Line 216, a non-dict tool_guards value is silently ignored, which can accidentally disable intended guard enforcement. Reject invalid shapes explicitly instead of falling back to None.
Proposed fix
- tool_guards = None
- tool_guards_data = frontmatter.get('tool_guards')
- if tool_guards_data and isinstance(tool_guards_data, dict):
- from cuga.backend.cuga_graph.policy.models import ToolGuard
- tool_guards = {
- tool_name: ToolGuard(**guard_data)
- for tool_name, guard_data in tool_guards_data.items()
- }
+ tool_guards = None
+ if 'tool_guards' in frontmatter:
+ tool_guards_data = frontmatter.get('tool_guards')
+ if not isinstance(tool_guards_data, dict):
+ raise ValueError(f"ToolGuide {name} has invalid 'tool_guards' (must be a dictionary)")
+
+ from cuga.backend.cuga_graph.policy.models import ToolGuard
+ tool_guards = {}
+ for tool_name, guard_data in tool_guards_data.items():
+ if not isinstance(guard_data, dict):
+ raise ValueError(
+ f"ToolGuide {name} has invalid guard config for '{tool_name}' (must be a dictionary)"
+ )
+ tool_guards[tool_name] = ToolGuard(**guard_data)🤖 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/policy/folder_loader.py` around lines 215 - 221,
The code silently ignores non-dict frontmatter for tool_guards; update the block
that reads frontmatter.get('tool_guards') to validate the shape and fail fast:
if tool_guards_data is None do nothing, if it's a dict continue to construct
tool_guards via ToolGuard(**guard_data) as now, but if tool_guards_data exists
and is not a dict raise a clear exception (e.g., ValueError) with a message
referencing 'tool_guards' and the offending type so malformed frontmatter is
rejected rather than silently ignored.
Comment on lines
+185
to
+187
| text_parts.append(f"Violations: {' '.join(guard.violating_examples[:3])}") | ||
| if guard.compliance_examples: | ||
| text_parts.append(f"Compliance: {' '.join(guard.compliance_examples[:3])}") |
Contributor
There was a problem hiding this comment.
Cap tool-guard example text before embedding.
Lines 185-187 concatenate raw examples with no length bound. Large examples can blow embedding input limits and fail policy writes.
Proposed fix
if policy.tool_guards:
# Add tool guard information to search text
+ max_example_chars = 200
for tool_name, guard in policy.tool_guards.items():
text_parts.append(f"Guard for {tool_name}")
if guard.violating_examples:
- text_parts.append(f"Violations: {' '.join(guard.violating_examples[:3])}")
+ violations = " ".join(
+ ex[:max_example_chars] for ex in guard.violating_examples[:3]
+ )
+ text_parts.append(f"Violations: {violations}")
if guard.compliance_examples:
- text_parts.append(f"Compliance: {' '.join(guard.compliance_examples[:3])}")
+ compliance = " ".join(
+ ex[:max_example_chars] for ex in guard.compliance_examples[:3]
+ )
+ text_parts.append(f"Compliance: {compliance}")🤖 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/policy/storage.py` around lines 185 - 187, The
current code appends raw guard.violating_examples and guard.compliance_examples
into text_parts with no length limits, which can exceed embedding input size;
before joining each example list (used where text_parts is built), truncate each
example string to a safe max length (e.g., 1000 chars) and/or limit the number
of characters per joined block, and then append the truncated/joined strings
instead of the raw ones; implement this by adding a small helper (e.g.,
truncate(text, max_chars)) or inline slicing when building the strings that
reference text_parts, guard.violating_examples, and guard.compliance_examples so
large examples are capped prior to embedding/storage.
Comment on lines
+11
to
+20
| # ── env vars MUST be set before any cuga import ────────────────────────────── | ||
| import os | ||
|
|
||
| os.environ["MCP_SERVERS_FILE"] = "none" # registry reads from DB | ||
| os.environ["CUGA_MANAGER_MODE"] = "true" | ||
| os.environ["DYNACONF_POLICY__FILESYSTEM_SYNC"] = "false" | ||
| os.environ["DYNACONF_ADVANCED_FEATURES__ENABLE_SHELL_TOOL"] = "false" | ||
| os.environ["DYNACONF_ADVANCED_FEATURES__OPENSANDBOX_SANDBOX"] = "false" | ||
| os.environ["DYNACONF_SKILLS__ENABLED"] = "false" | ||
| os.environ["DYNACONF_SUPERVISOR__ENABLED"] = "false" |
Contributor
There was a problem hiding this comment.
Avoid setting CUGA config at module import time.
These os.environ[...] writes make the test order-dependent: if another test imported cuga first, this config is already frozen; if this file imports first, the settings leak into the rest of the pytest worker. This scenario needs process isolation, or a fixture that patches/restores the environment before importing the CUGA modules.
🤖 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/policy/tool_guard/tests/test_crm_finance_tool_guard_e2e.py`
around lines 11 - 20, The test module currently sets os.environ entries at
import time (in test_crm_finance_tool_guard_e2e.py), which leaks config across
tests; instead remove those top-level os.environ writes and create a pytest
fixture (e.g., autouse or per-test) that uses monkeypatch.setenv to apply
MCP_SERVERS_FILE, CUGA_MANAGER_MODE, DYNACONF_* values before importing any CUGA
modules, and ensure all imports of cuga/backend/cuga_graph/... happen inside
that fixture or inside tests after the environment is patched so the config is
isolated and restored for other tests.
Comment on lines
+217
to
+285
| async def _build_agent(workspace: str) -> CugaAgent: | ||
| """Start demo CRM services and return an initialized agent.""" | ||
| app_mgr = AppManager( | ||
| process_registry=_processes, | ||
| run_service=_run_service, | ||
| kill_ports=_kill_ports, | ||
| kill_process=_kill_proc, | ||
| wait_tcp=lambda p, lbl, r, i: _wait_tcp(p, lbl, r, i), | ||
| wait_http=lambda p, n: _wait_http(p, n), | ||
| ) | ||
|
|
||
| print("📁 Preparing workspace…") | ||
| app_mgr.prepare_workspace(workspace) | ||
|
|
||
| ensure_managed_mcp_file_exists(get_managed_mcp_path()) | ||
|
|
||
| ports_to_free = app_mgr.ports_for_apps(email=True, filesystem=True, crm=True) | ||
| ports_to_free += [settings.server_ports.registry] | ||
| _kill_ports(ports_to_free) | ||
|
|
||
| print("🚀 Starting tool servers…") | ||
| print(" • Email sink + MCP server") | ||
| app_mgr.start_email() | ||
|
|
||
| print(" • Filesystem MCP server") | ||
| app_mgr.start_filesystem(workspace) | ||
|
|
||
| print(" • CRM API server") | ||
| crm_db = app_mgr.prepare_crm_db(workspace) | ||
| app_mgr.start_crm(crm_db) | ||
|
|
||
| print("💾 Seeding config DB with demo_crm tool definitions…") | ||
| loop = asyncio.get_running_loop() | ||
| await loop.run_in_executor(None, setup_demo_manage_config, "demo_crm") | ||
|
|
||
| print(" • Registry server") | ||
| registry_proc = app_mgr.start_registry() | ||
| if registry_proc is None or registry_proc.poll() is not None: | ||
| raise RuntimeError("Registry failed to start") | ||
|
|
||
| from cuga.backend.cuga_graph.nodes.cuga_lite.providers.combined import ( | ||
| CombinedToolProvider, | ||
| ) | ||
|
|
||
| print("🔌 Initializing tool provider with policy enforcement enabled…") | ||
| provider = CombinedToolProvider( | ||
| app_names=["crm", "filesystem", "email"], | ||
| ) | ||
| await provider.initialize() | ||
|
|
||
| apps = await provider.get_apps() | ||
| print(f" Loaded apps: {[a.name for a in apps]}") | ||
|
|
||
| workspace_abs = os.path.abspath(workspace) | ||
| workspace_instructions = ( | ||
| "## Plan\n" | ||
| f"For the filesystem application: write or read files only from `{workspace_abs}`\n" | ||
| "For the email application: send emails only using the local SMTP sink" | ||
| ) | ||
|
|
||
| agent = CugaAgent( | ||
| tool_provider=provider, | ||
| special_instructions=workspace_instructions, | ||
| cuga_folder=os.path.join(workspace, ".cuga"), | ||
| ) | ||
|
|
||
| await agent.policies._ensure_policy_system() | ||
|
|
||
| return agent |
Contributor
There was a problem hiding this comment.
Clean up spawned services when agent bootstrap fails.
_build_agent() starts email/filesystem/CRM/registry processes before provider.initialize(), get_apps(), and CugaAgent(...) can raise. If any of those steps fails, the caller's finally never runs because _build_agent() never returned, so the ports stay occupied for later tests.
Proposed fix
async def _build_agent(workspace: str) -> CugaAgent:
"""Start demo CRM services and return an initialized agent."""
- app_mgr = AppManager(
- process_registry=_processes,
- run_service=_run_service,
- kill_ports=_kill_ports,
- kill_process=_kill_proc,
- wait_tcp=lambda p, lbl, r, i: _wait_tcp(p, lbl, r, i),
- wait_http=lambda p, n: _wait_http(p, n),
- )
-
- print("📁 Preparing workspace…")
- app_mgr.prepare_workspace(workspace)
-
- ensure_managed_mcp_file_exists(get_managed_mcp_path())
-
- ports_to_free = app_mgr.ports_for_apps(email=True, filesystem=True, crm=True)
- ports_to_free += [settings.server_ports.registry]
- _kill_ports(ports_to_free)
-
- print("🚀 Starting tool servers…")
- ...
-
- await agent.policies._ensure_policy_system()
-
- return agent
+ try:
+ app_mgr = AppManager(
+ process_registry=_processes,
+ run_service=_run_service,
+ kill_ports=_kill_ports,
+ kill_process=_kill_proc,
+ wait_tcp=lambda p, lbl, r, i: _wait_tcp(p, lbl, r, i),
+ wait_http=lambda p, n: _wait_http(p, n),
+ )
+
+ print("📁 Preparing workspace…")
+ app_mgr.prepare_workspace(workspace)
+
+ ensure_managed_mcp_file_exists(get_managed_mcp_path())
+
+ ports_to_free = app_mgr.ports_for_apps(email=True, filesystem=True, crm=True)
+ ports_to_free += [settings.server_ports.registry]
+ _kill_ports(ports_to_free)
+
+ print("🚀 Starting tool servers…")
+ ...
+
+ await agent.policies._ensure_policy_system()
+ return agent
+ except Exception:
+ _cleanup()
+ raise🤖 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/policy/tool_guard/tests/test_crm_finance_tool_guard_e2e.py`
around lines 217 - 285, _build_agent currently starts multiple service processes
(start_email, start_filesystem, start_crm, start_registry) and then performs
provider.initialize(), provider.get_apps(), and CugaAgent construction without
ensuring cleanup on failure; wrap the post-startup initialization (from after
starting any service through until after await
agent.policies._ensure_policy_system()) in a try/except/finally that on
exception/exit stops all started services and frees ports (e.g., call a teardown
method on AppManager such as stop_all/stop_services if available, otherwise use
the existing helpers like _kill_ports and _kill_proc to terminate registry_proc
and any started app processes), then re-raise the exception so the caller’s
finally can proceed—this ensures any partial boot will not leave ports/processes
running.
Comment on lines
+407
to
+424
| print("\n" + "=" * 80) | ||
| print("🧪 TEST 1: Finance Company with Finance Tool Guard (SHOULD BE BLOCKED)") | ||
| print("=" * 80) | ||
| print("Testing Finance company with Finance industry tool guard active...") | ||
| print("Expected: Account creation should be blocked (Finance guard blocks low revenue)") | ||
|
|
||
| print(f"\n📝 Query: {FINANCE_TEST_CASE['query']}\n") | ||
| result1_initial = await agent.invoke(FINANCE_TEST_CASE["query"]) | ||
| print(f"\n✅ Agent Response (Initial):\n{result1_initial.answer}\n") | ||
|
|
||
| print("-" * 80) | ||
| print(f"\n📝 Follow-up Query: {FINANCE_TEST_CASE['followup']}\n") | ||
| result1_followup = await agent.invoke(FINANCE_TEST_CASE["followup"]) | ||
| print(f"\n✅ Agent Response (Follow-up):\n{result1_followup.answer}\n") | ||
| print("=" * 80) | ||
|
|
||
| assert result1_initial.answer | ||
| assert result1_followup.answer |
Contributor
There was a problem hiding this comment.
Assert that the guard blocked the tool call.
These assertions pass whenever the agent returns any non-empty text, including the failure mode where crm_create_account_accounts_post still ran and created the account. This E2E needs an observable enforcement check for both prompts: no tool execution, no CRM record created, or a specific blocked/refused outcome.
Based on learnings: Write new unit tests in corresponding tests/ sub-folders to verify all new branches and components.
🤖 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/policy/tool_guard/tests/test_crm_finance_tool_guard_e2e.py`
around lines 407 - 424, Update the E2E to assert that the Finance tool guard
actually prevented the tool call: replace the loose asserts on
result1_initial.answer and result1_followup.answer with checks that the agent
response indicates a blocked/refused outcome (e.g., contains a specific refusal
message) AND/OR verify the CRM tool function crm_create_account_accounts_post
was not executed by instrumenting/mocking that function and asserting it was not
called; additionally add focused unit tests in the appropriate tests/ sub-folder
that mock the CRM client and the tool-guard logic to cover the blocked path, the
allowed path, and edge cases so each branch of the guard is validated.
Comment on lines
+109
to
+112
| if target_tool not in policy.target_tools: | ||
| raise ValueError( | ||
| f"Tool '{target_tool}' not in policy.target_tools: {policy.target_tools}" | ||
| ) |
Contributor
There was a problem hiding this comment.
Honor target_tools=["*"] when validating the target tool.
add_tool_guide() documents "*" as the wildcard for all tools, but this check rejects every concrete target_tool in that mode. That makes both generate_examples() and generate_guard_code() unusable for wildcard tool guides.
Suggested fix
- if target_tool not in policy.target_tools:
+ if "*" not in policy.target_tools and target_tool not in policy.target_tools:
raise ValueError(
f"Tool '{target_tool}' not in policy.target_tools: {policy.target_tools}"
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if target_tool not in policy.target_tools: | |
| raise ValueError( | |
| f"Tool '{target_tool}' not in policy.target_tools: {policy.target_tools}" | |
| ) | |
| if "*" not in policy.target_tools and target_tool not in policy.target_tools: | |
| raise ValueError( | |
| f"Tool '{target_tool}' not in policy.target_tools: {policy.target_tools}" | |
| ) |
🤖 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/policy/tool_guard/tool_guard_buildtime.py` around
lines 109 - 112, The validation currently rejects concrete target_tool values
when the policy declares target_tools=["*"]; update the check in the block that
raises ValueError so it first allows the wildcard by treating "*" in
policy.target_tools as matching any target_tool (i.e., only raise if "*" not in
policy.target_tools and target_tool not in policy.target_tools). Modify the
logic used by add_tool_guide()/generate_examples()/generate_guard_code() that
references policy.target_tools and target_tool so wildcard policies pass
validation.
Comment on lines
+111
to
+120
| # Load policies if storage is available | ||
| if self.policy_storage is not None: | ||
| policies = await self.policy_storage.list_policies( | ||
| policy_type=PolicyType.TOOL_GUIDE, enabled_only=True | ||
| ) | ||
| logger.debug(f"Found {len(policies)} ToolGuide policies") | ||
|
|
||
| # Filter to ensure we only have ToolGuide instances | ||
| tool_guide_policies = [p for p in policies if isinstance(p, ToolGuide)] | ||
| self._build_tool_to_guards_mapping(tool_guide_policies) |
Contributor
There was a problem hiding this comment.
Honor enable_policies before loading guards.
initialize() still lists policies and builds tool_to_guards whenever policy_storage is present, even if enable_policies=False. That makes the flag ineffective for callers that inject storage and can turn enforcement on unexpectedly. Keep the storage load/build path behind the same feature gate.
🤖 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/policy/tool_guard/tool_guard_runtime.py` around
lines 111 - 120, The initialize() method currently loads policies and builds
tool_to_guards whenever self.policy_storage is set, ignoring the enable_policies
flag; change initialize() to check enable_policies first and only call
self.policy_storage.list_policies(...) and
self._build_tool_to_guards_mapping(...) when enable_policies is truthy (and
policy_storage is present). In short: wrap the list_policies (with
PolicyType.TOOL_GUIDE) and the
_build_tool_to_guards_mapping(tool_guide_policies) call behind the same
enable_policies feature gate so guards are not loaded when enable_policies is
False.
Comment on lines
+591
to
+598
| # Convert tool_guards dict to ToolGuard objects | ||
| tool_guards_obj = existing_policy.tool_guards.copy() if existing_policy.tool_guards else {} | ||
| for tool_name, guard_data in tool_guards.items(): | ||
| tool_guards_obj[tool_name] = ToolGuard( | ||
| violating_examples=guard_data.get("violating_examples", []), | ||
| compliance_examples=guard_data.get("compliance_examples", []), | ||
| policy_code=guard_data.get("policy_code", ""), | ||
| ) |
Contributor
There was a problem hiding this comment.
Reject tool guards for tools the policy does not target.
This method currently accepts any key in tool_guards and persists it even when the tool is outside existing_policy.target_tools. That leaves stored policy state that the later generation APIs will refuse to use.
Suggested fix
# Convert tool_guards dict to ToolGuard objects
tool_guards_obj = existing_policy.tool_guards.copy() if existing_policy.tool_guards else {}
+ allows_all_tools = "*" in existing_policy.target_tools
for tool_name, guard_data in tool_guards.items():
+ if not allows_all_tools and tool_name not in existing_policy.target_tools:
+ raise ValueError(
+ f"Tool '{tool_name}' is not in policy.target_tools: {existing_policy.target_tools}"
+ )
tool_guards_obj[tool_name] = ToolGuard(
violating_examples=guard_data.get("violating_examples", []),
compliance_examples=guard_data.get("compliance_examples", []),
policy_code=guard_data.get("policy_code", ""),
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| # Convert tool_guards dict to ToolGuard objects | |
| tool_guards_obj = existing_policy.tool_guards.copy() if existing_policy.tool_guards else {} | |
| for tool_name, guard_data in tool_guards.items(): | |
| tool_guards_obj[tool_name] = ToolGuard( | |
| violating_examples=guard_data.get("violating_examples", []), | |
| compliance_examples=guard_data.get("compliance_examples", []), | |
| policy_code=guard_data.get("policy_code", ""), | |
| ) | |
| # Convert tool_guards dict to ToolGuard objects | |
| tool_guards_obj = existing_policy.tool_guards.copy() if existing_policy.tool_guards else {} | |
| allows_all_tools = "*" in existing_policy.target_tools | |
| for tool_name, guard_data in tool_guards.items(): | |
| if not allows_all_tools and tool_name not in existing_policy.target_tools: | |
| raise ValueError( | |
| f"Tool '{tool_name}' is not in policy.target_tools: {existing_policy.target_tools}" | |
| ) | |
| tool_guards_obj[tool_name] = ToolGuard( | |
| violating_examples=guard_data.get("violating_examples", []), | |
| compliance_examples=guard_data.get("compliance_examples", []), | |
| policy_code=guard_data.get("policy_code", ""), | |
| ) |
🤖 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/sdk.py` around lines 591 - 598, The code is currently accepting any
key in tool_guards and persisting ToolGuard objects even for tools not in
existing_policy.target_tools; update the conversion loop that builds
tool_guards_obj so it validates each tool_name against
existing_policy.target_tools and rejects invalid entries by raising a clear
error (e.g., ValueError) naming the offending tool(s) and the allowed
target_tools; ensure this check is applied before creating ToolGuard instances
and that the error message references tool_name and existing_policy.target_tools
for debugging.
Comment on lines
+1279
to
+1303
| # Ensure policy system is initialized | ||
| policy_system = await self._ensure_policy_system() | ||
| if policy_system is None: | ||
| raise RuntimeError("Policy system is disabled") | ||
|
|
||
| # Get the policy | ||
| policy_data = await self.get(policy_id) | ||
| if policy_data is None: | ||
| raise ValueError(f"Policy with ID '{policy_id}' not found") | ||
|
|
||
| policy = policy_data.get('policy') | ||
| if policy is None: | ||
| raise ValueError(f"Could not retrieve policy object for ID '{policy_id}'") | ||
|
|
||
| # Validate policy type | ||
| if policy.type != PolicyType.TOOL_GUIDE: | ||
| raise ValueError( | ||
| f"Policy must be of type 'tool_guide', got '{policy.type}'. " | ||
| f"Only tool_guide policies can generate examples." | ||
| ) | ||
|
|
||
| # Create and initialize ToolGuardBuildtimeManager | ||
| manager = ToolGuardBuildtimeManager(self._agent) | ||
| await manager._ensure_initialized() | ||
|
|
Contributor
There was a problem hiding this comment.
Initialize the agent before build-time generation.
These entry points skip await self._agent._ensure_initialized(). On a fresh agent, custom tool providers may not have run initialize() yet, so ToolGuardBuildtimeManager._ensure_initialized() can hit get_all_tools() before the provider is ready. This also bypasses the SDK’s lazy knowledge-tool injection path.
Suggested fix
- # Ensure policy system is initialized
+ # Ensure tool provider and lazy tool injection are ready
+ await self._agent._ensure_initialized()
+
+ # Ensure policy system is initialized
policy_system = await self._ensure_policy_system()
if policy_system is None:
raise RuntimeError("Policy system is disabled")Also applies to: 1378-1402
🤖 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/sdk.py` around lines 1279 - 1303, The code creates a
ToolGuardBuildtimeManager and calls await manager._ensure_initialized() without
first ensuring the agent has been initialized, which can race with custom tool
providers; call await self._agent._ensure_initialized() before instantiating
ToolGuardBuildtimeManager (i.e., insert await self._agent._ensure_initialized()
right before manager = ToolGuardBuildtimeManager(self._agent)), and apply the
same change to the other occurrence referenced (the block around lines
1378-1402) so the agent is always initialized before build-time manager
initialization.
Comment on lines
+1293
to
+1298
| # Validate policy type | ||
| if policy.type != PolicyType.TOOL_GUIDE: | ||
| raise ValueError( | ||
| f"Policy must be of type 'tool_guide', got '{policy.type}'. " | ||
| f"Only tool_guide policies can generate examples." | ||
| ) |
Contributor
There was a problem hiding this comment.
Use the same policy-type accessor this manager already uses elsewhere.
list() and get() already handle both policy_type and type, but these branches dereference policy.type directly. If the stored model only exposes policy_type, both new SDK APIs fail with AttributeError before generation starts.
Suggested fix
- if policy.type != PolicyType.TOOL_GUIDE:
+ policy_type = policy.policy_type if hasattr(policy, "policy_type") else policy.type
+ if policy_type != PolicyType.TOOL_GUIDE:
raise ValueError(
- f"Policy must be of type 'tool_guide', got '{policy.type}'. "
+ f"Policy must be of type 'tool_guide', got '{policy_type}'. "
f"Only tool_guide policies can generate examples."
)Also applies to: 1392-1397
🤖 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/sdk.py` around lines 1293 - 1298, The policy type check is directly
accessing policy.type which raises AttributeError when the object only exposes
policy_type; change both checks to use the same accessor pattern used by list()
and get(): read the policy type via getattr(policy, "policy_type",
getattr(policy, "type", None)) (or equivalent nil-safe accessor) and compare
that value to PolicyType.TOOL_GUIDE; update the check in the shown block and the
similar block around the other occurrence (the second validator near the later
lines) so both places handle objects with either policy_type or type.
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
1 participant
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.
Summary by CodeRabbit
New Features
Tests
Chores