Open
Conversation
Contributor
WalkthroughAdds a configurable ToolUseEnforcement feature: new config types and TOML schema, a prompt fragment, PromptEngine helpers, and runtime wiring to conditionally append strict tool-use guidance to system/worker prompts based on routed model selection. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
![tembo[bot]](https://avatars.githubusercontent.com/in/1223475?s=60&v=4){kind=small}
src/agent/compactor.rs
Outdated
| ) { | ||
| Ok(prompt) => prompt, | ||
| Err(error) => { | ||
| tracing::error!(%error, "failed to append tool-use enforcement to compactor prompt"); |
Contributor
There was a problem hiding this comment.
This makes compaction bail if the enforcement fragment render fails. That feels like a non-critical failure path — I'd fall back to the base prompt and keep going.
Suggested change
| tracing::error!(%error, "failed to append tool-use enforcement to compactor prompt"); | |
| Ok(prompt) => match prompt_engine.maybe_append_tool_use_enforcement( | |
| prompt.clone(), | |
| tool_use_enforcement.as_ref(), | |
| &model_name, | |
| ) { | |
| Ok(prompt) => prompt, | |
| Err(error) => { | |
| tracing::error!(%error, "failed to append tool-use enforcement to compactor prompt"); | |
| prompt | |
| } | |
| }, |
| ) { | ||
| Ok(prompt) => prompt, | ||
| Err(error) => { | ||
| tracing::warn!(%error, "failed to append tool-use enforcement, using base cortex prompt"); |
Contributor
There was a problem hiding this comment.
Minor robustness: on append failure you can just fall back to the already-rendered base prompt instead of re-rendering the template (and potentially failing twice).
Suggested change
| tracing::warn!(%error, "failed to append tool-use enforcement, using base cortex prompt"); | |
| Ok(prompt) => match prompt_engine.maybe_append_tool_use_enforcement( | |
| prompt.clone(), | |
| tool_use_enforcement.as_ref(), | |
| &model_name, | |
| ) { | |
| Ok(prompt) => prompt, | |
| Err(error) => { | |
| tracing::warn!(%error, "failed to append tool-use enforcement, using base cortex prompt"); | |
| prompt | |
| } | |
| }, |
| let routing = rc.routing.load(); | ||
| let model_name = routing.resolve(ProcessType::Worker, None).to_string(); | ||
| let tool_use_enforcement = rc.tool_use_enforcement.load(); | ||
| let worker_system_prompt = prompt_engine |
Contributor
There was a problem hiding this comment.
Right now the enforcement fragment gets appended before the skills listing, so the final lines of the worker preamble are the skills section. If the goal is “last instruction wins”, it might be more effective to append enforcement after skills.
Suggested change
| let worker_system_prompt = prompt_engine | |
| let worker_system_prompt = prompt_engine | |
| .render_worker_prompt( | |
| &rc.instance_dir.display().to_string(), | |
| &rc.workspace_dir.display().to_string(), | |
| sandbox_enabled, | |
| sandbox_containment_active, | |
| sandbox_read_allowlist, | |
| sandbox_write_allowlist, | |
| &tool_secret_names, | |
| browser_config.persist_session, | |
| worker_status_text, | |
| ) | |
| .map_err(|e| AgentError::Other(anyhow::anyhow!("{e}")))?; | |
| let skills = rc.skills.load(); | |
| let brave_search_key = (**rc.brave_search_key.load()).clone(); | |
| // Append skills listing to worker system prompt. Suggested skills are | |
| // flagged so the worker knows the channel's intent, but it can read any | |
| // skill it decides is relevant via the read_skill tool. | |
| let system_prompt = match skills.render_worker_skills(suggested_skills, &prompt_engine) { | |
| Ok(skills_prompt) if !skills_prompt.is_empty() => { | |
| format!("{worker_system_prompt}\n\n{skills_prompt}") | |
| } | |
| Ok(_) => worker_system_prompt, | |
| Err(error) => { | |
| tracing::warn!(%error, "failed to render worker skills listing, spawning without skills context"); | |
| worker_system_prompt | |
| } | |
| }; | |
| let system_prompt = prompt_engine | |
| .maybe_append_tool_use_enforcement(system_prompt, tool_use_enforcement.as_ref(), &model_name) | |
| .map_err(|e| AgentError::Other(anyhow::anyhow!("{e}")))?; |
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/agent/channel.rs (1)
2268-2294:⚠️ Potential issue | 🟠 MajorCoalesced channel turns can bypass tool-use enforcement.
build_system_prompt()now appends enforcement, but coalesced turns usebuild_system_prompt_with_coalesce()(see Line 1433) which still returns the base prompt withoutmaybe_append_tool_use_enforcement(...). That leaves a live bypass path whenever batching is active.Suggested patch (apply same enforcement in coalesced prompt builder)
async fn build_system_prompt_with_coalesce( @@ - prompt_engine.render_channel_prompt_with_links( + let routing = rc.routing.load(); + let model_name = routing.resolve(ProcessType::Channel, None).to_string(); + let tool_use_enforcement = rc.tool_use_enforcement.load(); + + let system_prompt = prompt_engine.render_channel_prompt_with_links( empty_to_none(identity_context), empty_to_none(memory_bulletin.to_string()), empty_to_none(skills_prompt), worker_capabilities, self.conversation_context.clone(), empty_to_none(status_text), coalesce_hint, available_channels, sandbox_enabled, org_context, adapter_prompt, project_context, self.backfill_transcript.clone(), empty_to_none(working_memory), empty_to_none(channel_activity_map), - ) + )?; + + prompt_engine.maybe_append_tool_use_enforcement( + system_prompt, + tool_use_enforcement.as_ref(), + &model_name, + ) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/agent/channel.rs` around lines 2268 - 2294, build_system_prompt_with_coalesce currently returns a system prompt without invoking prompt_engine.maybe_append_tool_use_enforcement, allowing coalesced/batched channel turns to bypass tool-use enforcement; update build_system_prompt_with_coalesce (the coalesce path used for batching) to call prompt_engine.maybe_append_tool_use_enforcement with the same arguments as build_system_prompt (pass system_prompt, tool_use_enforcement.as_ref(), &model_name) after rendering the coalesced prompt so tool-use enforcement is applied consistently for both build_system_prompt() and build_system_prompt_with_coalesce().src/agent/channel_dispatch.rs (1)
1105-1117:⚠️ Potential issue | 🟠 MajorResumed builtin workers don't get tool-use enforcement.
The
resume_idle_worker_into_statepath (line 1106) callsrender_worker_promptdirectly, but unlikespawn_worker_innerand the task pickup path incortex.rs, it doesn't chain.and_then(maybe_append_tool_use_enforcement(...)). This creates an inconsistency where:
- Fresh workers → get enforcement guidance (required for outcome signaling)
- Resumed workers → skip enforcement guidance
Given that workers must signal terminal outcomes via
set_status(kind: "outcome"), add tool-use enforcement to the resume path to align with other worker spawn paths. If skipping enforcement is intentional (e.g., to preserve the exact original prompt), add an explicit comment explaining the design choice.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/agent/channel_dispatch.rs` around lines 1105 - 1117, The resume_idle_worker_into_state path currently calls prompt_engine.render_worker_prompt and returns its Result directly; update it to mirror spawn_worker_inner and the task pickup path by chaining .and_then(|p| maybe_append_tool_use_enforcement(p, &tool_definitions, &tool_secret_names)) (or the correct argument order used elsewhere) so resumed workers receive the same tool-use enforcement text; if the original omission is intentional, replace the change with an explicit comment in resume_idle_worker_into_state explaining why enforcement must be skipped and reference render_worker_prompt and maybe_append_tool_use_enforcement in the comment.
🧹 Nitpick comments (3)
src/agent/compactor.rs (2)
128-144: Tool-use enforcement may be unnecessary for the compactor.The compactor agent is built without any tools (line 247-250 shows
AgentBuilder::new(model).preamble(...).default_max_turns(1).build()with notool_server_handle). The tool-use enforcement guidance is designed to prevent models from narrating tool actions instead of executing them, but the compactor's sole job is producing a summary text.Appending tool-use enforcement guidance to a toolless agent may confuse the model or waste tokens. Consider whether this should be skipped for
ProcessType::Compactor.💡 Alternative: Skip enforcement for compactor
If enforcement should be skipped, the simplest approach would be to not call
maybe_append_tool_use_enforcementhere, or ensureToolUseEnforcement::should_inject()returnsfalsefor compactor contexts.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/agent/compactor.rs` around lines 128 - 144, The compactor prompt is being appended with tool-use enforcement even though the Compactor agent is built without tools; change the logic around prompt_engine.maybe_append_tool_use_enforcement so it is skipped for ProcessType::Compactor (or when no tools are configured) — check the routing/model context (model_name or ProcessType::Compactor) or a "has tools" flag before calling maybe_append_tool_use_enforcement and fall back to using prompt_engine.render_static("compactor") directly; keep the existing error handling and is_compacting write-unset path intact if you still call maybe_append_tool_use_enforcement.
131-151: Consider usingand_thenfor cleaner error chaining.The nested match for error handling could be flattened using
and_then, similar to how other files in this PR handle the pattern:♻️ Optional refactor using and_then
- let compactor_prompt = match prompt_engine.render_static("compactor") { - Ok(prompt) => match prompt_engine.maybe_append_tool_use_enforcement( - prompt, - tool_use_enforcement.as_ref(), - &model_name, - ) { - Ok(prompt) => prompt, - Err(error) => { - tracing::error!(%error, "failed to append tool-use enforcement to compactor prompt"); - let mut flag = is_compacting.write().await; - *flag = false; - return; - } - }, - Err(error) => { - tracing::error!(%error, "failed to render compactor prompt"); - let mut flag = is_compacting.write().await; - *flag = false; - return; - } - }; + let compactor_prompt = match prompt_engine + .render_static("compactor") + .and_then(|prompt| { + prompt_engine.maybe_append_tool_use_enforcement( + prompt, + tool_use_enforcement.as_ref(), + &model_name, + ) + }) + { + Ok(prompt) => prompt, + Err(error) => { + tracing::error!(%error, "failed to render compactor prompt"); + let mut flag = is_compacting.write().await; + *flag = false; + return; + } + };🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/agent/compactor.rs` around lines 131 - 151, Replace the nested match that builds compactor_prompt by chaining the two fallible operations using and_then: call prompt_engine.render_static("compactor").and_then(|prompt| prompt_engine.maybe_append_tool_use_enforcement(prompt, tool_use_enforcement.as_ref(), &model_name)), map or propagate the Err case to log the error with tracing::error!(%error, "...") and then set *flag = false on is_compacting.write().await before returning; ensure you still reference compactor_prompt, prompt_engine.render_static, prompt_engine.maybe_append_tool_use_enforcement, tool_use_enforcement, model_name and is_compacting so the logic and error logging remain identical while removing the nested match.src/agent/cortex.rs (1)
1499-1513: Use the already-rendered prompt instead of re-rendering on enforcement failure.When
maybe_append_tool_use_enforcementfails, the code re-renders the cortex prompt instead of reusing thepromptvariable that was already successfully rendered in the outerOk(prompt)arm. This is wasteful and the second render could theoretically fail differently.♻️ Proposed fix to reuse the already-rendered prompt
let system_prompt = match prompt_engine.render_static("cortex") { - Ok(prompt) => match prompt_engine.maybe_append_tool_use_enforcement( - prompt, + Ok(base_prompt) => match prompt_engine.maybe_append_tool_use_enforcement( + base_prompt.clone(), tool_use_enforcement.as_ref(), &model_name, ) { Ok(prompt) => prompt, Err(error) => { tracing::warn!(%error, "failed to append tool-use enforcement, using base cortex prompt"); - prompt_engine.render_static("cortex").unwrap_or_default() + base_prompt } },🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/agent/cortex.rs` around lines 1499 - 1513, The code currently re-calls prompt_engine.render_static("cortex") inside the Err branch of prompt_engine.maybe_append_tool_use_enforcement; instead, reuse the already-rendered prompt from the outer Ok(prompt) arm to avoid double-rendering and possible different failures. In the Err arm for maybe_append_tool_use_enforcement (within the match that begins with prompt_engine.render_static("cortex")), keep the original prompt variable (the one bound by the outer Ok(prompt)) and return it (after logging the warning) instead of calling render_static again; refer to prompt_engine.maybe_append_tool_use_enforcement, prompt_engine.render_static("cortex"), the prompt variable, and system_prompt to locate and update the logic.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/agent/ingestion.rs`:
- Around line 479-484: After building the ingestion_prompt with
prompt_engine.maybe_append_tool_use_enforcement and before recording the chunk
as completed, add a check for the memory_persistence_complete signal from the
model response and treat a missing signal as a chunk failure: if
memory_persistence_complete is absent, log an error with context (include
model_name, chunk id or source path), do not call the existing "mark chunk
completed" or delete the source file, and instead record the chunk as
failed/queued for retry so memory isn't lost; update the code paths that
currently always record completion after ingestion (look for
variables/operations around ingestion_prompt, memory_persistence_complete, and
the chunk completion/cleanup logic) to short-circuit on missing signal and
preserve the source for retry.
In `@src/config/types.rs`:
- Around line 245-279: The custom serde::Deserialize impl for ToolUseEnforcement
currently handles only toml::Value::String and Array, so bare TOML booleans
(toml::Value::Boolean) like tool_use_enforcement = true/false will fail; update
the deserialize function (impl Deserialize for ToolUseEnforcement) to add a
toml::Value::Boolean(b) match arm that maps true -> ToolUseEnforcement::Always
and false -> ToolUseEnforcement::Never (and keep existing String and Array
handling and error cases unchanged) so documented boolean configs are accepted.
---
Outside diff comments:
In `@src/agent/channel_dispatch.rs`:
- Around line 1105-1117: The resume_idle_worker_into_state path currently calls
prompt_engine.render_worker_prompt and returns its Result directly; update it to
mirror spawn_worker_inner and the task pickup path by chaining .and_then(|p|
maybe_append_tool_use_enforcement(p, &tool_definitions, &tool_secret_names)) (or
the correct argument order used elsewhere) so resumed workers receive the same
tool-use enforcement text; if the original omission is intentional, replace the
change with an explicit comment in resume_idle_worker_into_state explaining why
enforcement must be skipped and reference render_worker_prompt and
maybe_append_tool_use_enforcement in the comment.
In `@src/agent/channel.rs`:
- Around line 2268-2294: build_system_prompt_with_coalesce currently returns a
system prompt without invoking prompt_engine.maybe_append_tool_use_enforcement,
allowing coalesced/batched channel turns to bypass tool-use enforcement; update
build_system_prompt_with_coalesce (the coalesce path used for batching) to call
prompt_engine.maybe_append_tool_use_enforcement with the same arguments as
build_system_prompt (pass system_prompt, tool_use_enforcement.as_ref(),
&model_name) after rendering the coalesced prompt so tool-use enforcement is
applied consistently for both build_system_prompt() and
build_system_prompt_with_coalesce().
---
Nitpick comments:
In `@src/agent/compactor.rs`:
- Around line 128-144: The compactor prompt is being appended with tool-use
enforcement even though the Compactor agent is built without tools; change the
logic around prompt_engine.maybe_append_tool_use_enforcement so it is skipped
for ProcessType::Compactor (or when no tools are configured) — check the
routing/model context (model_name or ProcessType::Compactor) or a "has tools"
flag before calling maybe_append_tool_use_enforcement and fall back to using
prompt_engine.render_static("compactor") directly; keep the existing error
handling and is_compacting write-unset path intact if you still call
maybe_append_tool_use_enforcement.
- Around line 131-151: Replace the nested match that builds compactor_prompt by
chaining the two fallible operations using and_then: call
prompt_engine.render_static("compactor").and_then(|prompt|
prompt_engine.maybe_append_tool_use_enforcement(prompt,
tool_use_enforcement.as_ref(), &model_name)), map or propagate the Err case to
log the error with tracing::error!(%error, "...") and then set *flag = false on
is_compacting.write().await before returning; ensure you still reference
compactor_prompt, prompt_engine.render_static,
prompt_engine.maybe_append_tool_use_enforcement, tool_use_enforcement,
model_name and is_compacting so the logic and error logging remain identical
while removing the nested match.
In `@src/agent/cortex.rs`:
- Around line 1499-1513: The code currently re-calls
prompt_engine.render_static("cortex") inside the Err branch of
prompt_engine.maybe_append_tool_use_enforcement; instead, reuse the
already-rendered prompt from the outer Ok(prompt) arm to avoid double-rendering
and possible different failures. In the Err arm for
maybe_append_tool_use_enforcement (within the match that begins with
prompt_engine.render_static("cortex")), keep the original prompt variable (the
one bound by the outer Ok(prompt)) and return it (after logging the warning)
instead of calling render_static again; refer to
prompt_engine.maybe_append_tool_use_enforcement,
prompt_engine.render_static("cortex"), the prompt variable, and system_prompt to
locate and update the logic.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 9f68f0bd-9b5d-479f-bdd7-6edd765efca2
📒 Files selected for processing (16)
prompts/en/fragments/tool_use_enforcement.md.j2src/agent/channel.rssrc/agent/channel_dispatch.rssrc/agent/compactor.rssrc/agent/cortex.rssrc/agent/cortex_chat.rssrc/agent/ingestion.rssrc/api/agents.rssrc/config.rssrc/config/load.rssrc/config/runtime.rssrc/config/toml_schema.rssrc/config/types.rssrc/prompts/engine.rssrc/prompts/text.rssrc/tools/spawn_worker.rs
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/agent/cortex.rs (1)
3295-3314:⚠️ Potential issue | 🟡 MinorError propagation blocks all task pickups if enforcement fragment is broken.
The
.and_then()chain at lines 3307-3313 propagates anymaybe_append_tool_use_enforcementfailure, causingpickup_one_ready_taskto return an error. Per the context snippet fromsrc/prompts/engine.rs, if the "fragments/tool_use_enforcement" template is missing or broken, this will deterministically fail every pickup iteration. Since tool-use enforcement is optional guidance, fall back to the base worker prompt instead.Suggested fix
let worker_system_prompt = prompt_engine .render_worker_prompt( &deps.runtime_config.instance_dir.display().to_string(), &deps.runtime_config.workspace_dir.display().to_string(), sandbox_enabled, sandbox_containment_active, sandbox_read_allowlist, sandbox_write_allowlist, &tool_secret_names, browser_config.persist_session, worker_status_text, ) - .and_then(|prompt| { - prompt_engine.maybe_append_tool_use_enforcement( + .map(|prompt| { + match prompt_engine.maybe_append_tool_use_enforcement( prompt, tool_use_enforcement.as_ref(), &model_name, - ) + ) { + Ok(augmented) => augmented, + Err(error) => { + tracing::warn!(%error, "failed to append tool-use enforcement to worker prompt"); + prompt + } + } }) .map_err(|error| anyhow::anyhow!("failed to render worker prompt: {error}"))?;Note: This requires changing to
prompt.clone()passed tomaybe_append_tool_use_enforcementto retain ownership for the fallback:.map(|prompt| { match prompt_engine.maybe_append_tool_use_enforcement( - prompt, + prompt.clone(), tool_use_enforcement.as_ref(), &model_name, ) { Ok(augmented) => augmented, Err(error) => { tracing::warn!(%error, "failed to append tool-use enforcement to worker prompt"); prompt } } })🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/agent/cortex.rs` around lines 3295 - 3314, The current chain uses prompt_engine.render_worker_prompt(...).and_then(|prompt| prompt_engine.maybe_append_tool_use_enforcement(prompt, ...)) which propagates any enforcement-template error and blocks pickup; instead capture the base prompt from render_worker_prompt, call maybe_append_tool_use_enforcement with a cloned prompt (e.g., prompt.clone()) and on Err return Ok(prompt) so the worker_system_prompt falls back to the original prompt when enforcement rendering fails, ensuring pickup_one_ready_task continues; update the code around worker_system_prompt, using match/map_or_else or explicit match on prompt_engine.maybe_append_tool_use_enforcement to log the enforcement error but return the base prompt on failure.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@src/agent/cortex.rs`:
- Around line 3295-3314: The current chain uses
prompt_engine.render_worker_prompt(...).and_then(|prompt|
prompt_engine.maybe_append_tool_use_enforcement(prompt, ...)) which propagates
any enforcement-template error and blocks pickup; instead capture the base
prompt from render_worker_prompt, call maybe_append_tool_use_enforcement with a
cloned prompt (e.g., prompt.clone()) and on Err return Ok(prompt) so the
worker_system_prompt falls back to the original prompt when enforcement
rendering fails, ensuring pickup_one_ready_task continues; update the code
around worker_system_prompt, using match/map_or_else or explicit match on
prompt_engine.maybe_append_tool_use_enforcement to log the enforcement error but
return the base prompt on failure.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: cf4f0092-bd29-44c6-99cf-6aa9215eb9cb
📒 Files selected for processing (11)
src/agent/channel.rssrc/agent/channel_dispatch.rssrc/agent/compactor.rssrc/agent/cortex.rssrc/api/agents.rssrc/config.rssrc/config/load.rssrc/config/runtime.rssrc/config/toml_schema.rssrc/config/types.rssrc/tools/spawn_worker.rs
✅ Files skipped from review due to trivial changes (5)
- src/config/toml_schema.rs
- src/config/load.rs
- src/config.rs
- src/tools/spawn_worker.rs
- src/config/types.rs
🚧 Files skipped from review as they are similar to previous changes (4)
- src/config/runtime.rs
- src/api/agents.rs
- src/agent/channel_dispatch.rs
- src/agent/channel.rs
- Skip enforcement for toolless compactor agent - Reuse rendered prompt on enforcement failure in cortex (avoid double render) - Move enforcement after skills in worker spawn (last instruction wins) - Add enforcement to resumed builtin workers and coalesced channel turns - Return error on missing memory_persistence_complete signal in ingestion - Accept bare TOML booleans for tool_use_enforcement config Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/agent/ingestion.rs (1)
533-535: Please add a regression test for this fail-closed branch.This guard is now what keeps a chunk retryable when the model never emits
memory_persistence_complete, but the test module does not exercise that path yet. A focused test here would make future regressions much easier to catch. Based on learnings, "For changes in async/stateful paths (worker lifecycle, cancellation, retrigger, recall cache behavior), include explicit race/terminal-state reasoning in the PR summary and run targeted tests in addition tojust gate-pr."🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/agent/ingestion.rs` around lines 533 - 535, Add a regression test in the ingestion.rs test module that exercises the fail-closed branch where the model never emits the memory_persistence_complete signal: create a test (e.g., test_chunk_retry_when_no_memory_persistence_complete) that drives the chunk ingestion path to completion without emitting memory_persistence_complete, assert the function returns the Err containing the "ingestion chunk {chunk_number}/{total_chunks} for {filename} completed without memory_persistence_complete signal" message, and additionally assert the chunk is marked retryable by verifying the same retry/state indicator your code uses (e.g., the retry queue entry, chunk status field, or the function that schedules retries). Ensure the test controls the async/model behavior deterministically (mock/simulate no memory_persistence_complete) so it consistently hits the branch.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/config/types.rs`:
- Around line 266-278: The array branch that constructs Self::Custom currently
accepts empty or whitespace-only strings (e.g., tool_use_enforcement = [""]),
causing should_inject() to match always; update the mapping in the
toml::Value::Array arm so after converting each value to a String you call
.trim() and reject entries where trimmed.is_empty() by returning
D::Error::invalid_value (with a clear message like "empty or whitespace-only
string"); this ensures the patterns Vec<String> only contains non-empty trimmed
patterns before returning Ok(Self::Custom(patterns)).
---
Nitpick comments:
In `@src/agent/ingestion.rs`:
- Around line 533-535: Add a regression test in the ingestion.rs test module
that exercises the fail-closed branch where the model never emits the
memory_persistence_complete signal: create a test (e.g.,
test_chunk_retry_when_no_memory_persistence_complete) that drives the chunk
ingestion path to completion without emitting memory_persistence_complete,
assert the function returns the Err containing the "ingestion chunk
{chunk_number}/{total_chunks} for {filename} completed without
memory_persistence_complete signal" message, and additionally assert the chunk
is marked retryable by verifying the same retry/state indicator your code uses
(e.g., the retry queue entry, chunk status field, or the function that schedules
retries). Ensure the test controls the async/model behavior deterministically
(mock/simulate no memory_persistence_complete) so it consistently hits the
branch.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 3a825948-a232-4966-baa1-52d07bb57fc9
📒 Files selected for processing (6)
src/agent/channel.rssrc/agent/channel_dispatch.rssrc/agent/compactor.rssrc/agent/cortex.rssrc/agent/ingestion.rssrc/config/types.rs
✅ Files skipped from review due to trivial changes (2)
- src/agent/compactor.rs
- src/agent/channel.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- src/agent/channel_dispatch.rs
Comment on lines
+266
to
+278
| toml::Value::Array(arr) => { | ||
| let patterns: Vec<String> = arr | ||
| .into_iter() | ||
| .map(|v| { | ||
| v.as_str().map(String::from).ok_or_else(|| { | ||
| D::Error::invalid_value( | ||
| serde::de::Unexpected::Other("non-string array element"), | ||
| &"array of strings", | ||
| ) | ||
| }) | ||
| }) | ||
| .collect::<std::result::Result<_, _>>()?; | ||
| Ok(Self::Custom(patterns)) |
Contributor
There was a problem hiding this comment.
Reject blank custom match patterns.
tool_use_enforcement = [""] currently deserializes successfully, and then should_inject() turns into an unconditional match because every string contains the empty substring. Trim and reject empty entries here so malformed config does not silently behave like Always.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/config/types.rs` around lines 266 - 278, The array branch that constructs
Self::Custom currently accepts empty or whitespace-only strings (e.g.,
tool_use_enforcement = [""]), causing should_inject() to match always; update
the mapping in the toml::Value::Array arm so after converting each value to a
String you call .trim() and reject entries where trimmed.is_empty() by returning
D::Error::invalid_value (with a clear message like "empty or whitespace-only
string"); this ensures the patterns Vec<String> only contains non-empty trimmed
patterns before returning Ok(Self::Custom(patterns)).
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
tool_use_enforcementconfig support for defaults and per-agent overrides withauto,always,never, or custom model substring listsConfig
Testing
cargo test tool_use_enforcement --libjust gate-prNote
This PR adds configurable tool-use enforcement to prevent models like GPT and Gemini from narrating tool work instead of executing it. The feature supports four modes:
auto(detects based on model),always,never, or custom model substring lists. Configuration is applied at both defaults and per-agent levels, injecting enforcement guidance into all agent types (channels, branches, workers, cortex, compactor). A new reusable prompt fragment handles the injection logic. Includes test coverage for config parsing and model matching.Written by Tembo for commit 7957a4b.