feat(deep-research): loop guard middleware to bound researcher calls - #436
feat(deep-research): loop guard middleware to bound researcher calls#436smasurekar wants to merge 1 commit into
Conversation
WalkthroughThe deep-research agent now supports a configurable per-worker loop guard. It limits source calls, repeated identical requests, consecutive ChangesDeep researcher loop guard
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant DeepResearchAgent
participant ResearcherRunnable
participant ResearcherLoopGuardMiddleware
participant SourceTools
DeepResearchAgent->>ResearcherRunnable: start configured research
ResearcherRunnable->>ResearcherLoopGuardMiddleware: request source or think call
ResearcherLoopGuardMiddleware->>ResearcherLoopGuardMiddleware: update invocation budgets
ResearcherLoopGuardMiddleware->>SourceTools: execute admitted source call
SourceTools-->>ResearcherLoopGuardMiddleware: return source result
ResearcherLoopGuardMiddleware-->>ResearcherRunnable: return result or guard warning
ResearcherLoopGuardMiddleware-->>ResearcherRunnable: withdraw tools after exhaustion
ResearcherRunnable-->>DeepResearchAgent: return bounded ResearchNotes
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 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 `@docs/source/customization/configuration-reference.md`:
- Around line 550-561: Update the exhaustion paragraph to describe
max_source_calls_per_query separately from max_identical_source_calls: only the
query-wide limit preserves and annotates the last search result, while an
identical repeat is not executed and returns an error tool result. Add the
max_consecutive_thinks outcome: the threshold think call still executes with a
WARNING: suffix, then only think is withdrawn from later model calls while
source tools remain available. Review nearby documentation for stale examples or
inaccurate command details.
In `@src/aiq_agent/agents/deep_researcher/custom_middleware.py`:
- Around line 1598-1636: Update _guard_source_call so reaching
max_identical_source_calls blocks only the repeated signature while allowing
other source signatures to continue; do not call _mark_exhausted on this path.
Add a _blocked_result variant or parameter that omits the “source budget is
exhausted” message for duplicate-call blocks, while preserving exhaustion for
the total source-call budget path.
In `@src/aiq_agent/agents/deep_researcher/factory.py`:
- Around line 307-316: Update build_researcher_middleware’s ToolRetryMiddleware
lookup to handle common_middleware stacks that lack that middleware without
leaking StopIteration. Either raise an explicit construction error describing
the missing ToolRetryMiddleware or append ResearcherLoopGuardMiddleware at the
end, while preserving insertion before ToolRetryMiddleware when it is present.
- Around line 448-458: Update the researcher execution path around create_agent
and _run_research_query to catch GraphRecursionError and return the partially
collected evidence as valid ResearchNotes instead of propagating RuntimeError.
Record explicit ResearchGap entries for truncated work and lower the
evidence_judgment, while preserving normal completion behavior; add a regression
test covering filesystem-tool exhaustion and verify RESEARCHER_NON_SEARCH_TURNS
= 10 covers the listed non-search turns.
In `@tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py`:
- Around line 563-628: Reduce brittle literal-text assertions in
TestPromptRendering.test_enabled_mode_adds_the_ceiling_without_removing_the_guidance
and test_enabled_mode_explains_the_graceful_exit. Retain StrictUndefined
coverage, enabled/disabled branch checks, rendered numeric limits, and the
negative ResearchNotes assertion; remove or loosen wording-only checks such as
“a backstop,” withdrawal instructions, tool names, ResearchGap, and
evidence_judgment text.
🪄 Autofix
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: ASSERTIVE
Plan: Enterprise
Run ID: bd3b752b-66f7-4a14-a4f8-d6a141fdc427
📒 Files selected for processing (12)
docs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/agent.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.pysrc/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/prompts/researcher.j2src/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pytests/aiq_agent/agents/deep_researcher/test_factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📜 Review details
⏰ Context from checks skipped due to timeout. (3)
- GitHub Check: Pytest and Coverage
- GitHub Check: Script Validation
- GitHub Check: Lint and Hooks
🧰 Additional context used
📓 Path-based instructions (9)
**/*
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*: Do not include secrets, credentials, private hostnames, internal-only logs, customer data, or generated local artifacts.
Add or update tests for behavior changes.
**/*: For substantial behavior, authentication, UI, or architecture changes, open a design discussion before coding rather than landing a large unreviewed change.
Keep changes scoped to this repository and avoid editing adjacent repositories; treat eachsources/*package independently and prefer the smallest package-scoped change.
Keep pull requests scoped, avoid unrelated files and generated artifacts, provide validation evidence, and ensure every commit has DCO sign-off.
Files:
src/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/prompts/researcher.j2src/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/agent.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_factory.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pydocs/source/customization/configuration-reference.mdtests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*.py: Runuv run ruff check .anduv run ruff format --check .for root Python changes.
Runuv run pytestfor root project Python changes.
Files:
src/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/agent.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_factory.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
**/*.{py,pyi}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{py,pyi}: Format and lint Python code with Ruff using line length 120, Python 3.11 targeting, rules E, F, W, I, PL, and UP, with single-line imports; do not reformat unrelated code.
Missing-secret paths must degrade gracefully by stubbing or skipping rather than crashing or leaking secrets.
Files:
src/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/agent.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_factory.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
**/*.{py,pyi,js,jsx,ts,tsx,yml,yaml,json,env,md}
📄 CodeRabbit inference engine (AGENTS.md)
Never commit secrets, tokens, or environment-specific hostnames; use environment variables and
SecretStr, and resolve API keys at runtime.
Files:
src/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/agent.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_factory.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pydocs/source/customization/configuration-reference.mdtests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
**/*.{py,pyi,js,jsx,ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Never print or log secret values, including in tool output or error messages.
Files:
src/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/agent.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_factory.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Respect authenticated data sources by honoringrequires_auth, passing through per-user tokens, and using backend token validators; apply owner guardrails before loading protected report or artifact context into an agent.
Do not weaken or bypassAuthMiddleware, authentication validators, or authentication gating without prior design discussion.
Files:
src/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/agent.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
src/aiq_agent/agents/**/*
⚙️ CodeRabbit configuration file
src/aiq_agent/agents/**/*: Review agent changes for research workflow correctness, graph state transitions, prompt/tool contracts,
HITL behavior, and failure handling. Flag changes that weaken source attribution, report generation,
async cancellation, checkpointing, or data-source selection without focused tests and docs.
Files:
src/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/prompts/researcher.j2src/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/agent.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
docs/source/**/*
📄 CodeRabbit inference engine (AGENTS.md)
Update canonical documentation under
docs/source/when behavior, configuration, or workflows change; do not duplicate full documentation pages in skills.
Files:
docs/source/customization/configuration-reference.md
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}
⚙️ CodeRabbit configuration file
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}: Review documentation for command accuracy, branch-name consistency, current CI and copy-pr-bot behavior, public
vs internal boundary clarity, stale examples, and links that no longer match the repository layout.
Files:
docs/source/customization/configuration-reference.md
🧠 Learnings (14)
📚 Learning: 2026-06-11T21:21:11.314Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: docs/source/contributing/code-organization.md:0-0
Timestamp: 2026-06-11T21:21:11.314Z
Learning: Applies to docs/source/contributing/src/aiq_agent/agents/deep_researcher/** : Deep researcher agent should be organized in `src/aiq_agent/agents/deep_researcher/` with implementation details documented in README.md
Applied to files:
src/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/prompts/researcher.j2src/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/agent.pytests/aiq_agent/agents/deep_researcher/test_factory.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pydocs/source/customization/configuration-reference.mdtests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.pysrc/aiq_agent/agents/deep_researcher/factory.py
📚 Learning: 2026-08-11T06:34:44.687Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/deep_researcher/register.py:302-307
Timestamp: 2026-08-11T06:34:44.687Z
Learning: In Python logging code that handles potentially sensitive exceptions, do not add `exc_info=True` solely to restore stack traces, because standard traceback formatting includes `str(exception)` and may expose provider, customer, or credential-bearing content. When sensitive-content redaction is required, log the exception type together with `log_content_metadata(exception)` instead.
Applied to files:
src/aiq_agent/agents/deep_researcher/models/__init__.pysrc/aiq_agent/agents/deep_researcher/register.pysrc/aiq_agent/agents/deep_researcher/agent.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_factory.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
📚 Learning: 2026-07-23T22:20:09.213Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: skills/aiq-deploy/SKILL.md:0-0
Timestamp: 2026-07-23T22:20:09.213Z
Learning: Applies to skills/aiq-deploy/**/* : Do not continue into deep research or deep-research completion validation unless the user asks for it or confirms the post-deployment validation prompt.
Applied to files:
src/aiq_agent/agents/deep_researcher/prompts/researcher.j2
📚 Learning: 2026-07-06T23:55:46.952Z
Learnt from: cdgamarose-nv
Repo: NVIDIA-AI-Blueprints/aiq PR: 311
File: src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2:74-81
Timestamp: 2026-07-06T23:55:46.952Z
Learning: In `tests/aiq_agent/agents/deep_researcher/test_agent.py`, avoid asserting exact substrings from the orchestrator/writer prompt templates (e.g., `src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2`) since prompt wording changes frequently. Prefer testing structural/behavioral properties instead.
Applied to files:
src/aiq_agent/agents/deep_researcher/prompts/researcher.j2tests/aiq_agent/agents/deep_researcher/test_factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-07-23T22:20:18.400Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: skills/aiq-research/SKILL.md:0-0
Timestamp: 2026-07-23T22:20:18.400Z
Learning: Applies to skills/aiq-research/scripts/aiq.py : When polling requires an execution method with escalated permissions, request explicit user approval first and explain why; tell the user when deep research is running in the background.
Applied to files:
src/aiq_agent/agents/deep_researcher/prompts/researcher.j2
📚 Learning: 2026-07-23T22:20:18.400Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: skills/aiq-research/SKILL.md:0-0
Timestamp: 2026-07-23T22:20:18.400Z
Learning: Applies to skills/aiq-research/**/SKILL.md : Treat report editing as cosmetic only and use `report_edit <JOB_ID> <EDIT_INSTRUCTIONS>`; use `research` for a new or refined investigation.
Applied to files:
src/aiq_agent/agents/deep_researcher/prompts/researcher.j2
📚 Learning: 2026-08-04T22:37:11.331Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 414
File: tests/aiq_agent/test_default_model_profiles.py:72-75
Timestamp: 2026-08-04T22:37:11.331Z
Learning: For the AIQ default-model migration, PR `#414` intentionally retains Nemotron 3 Ultra expectations for intent classification and shallow research because the Nemotron Nano 3.5 Preview public endpoint and GA slug are unavailable. The Nano intent alias and corresponding profile tests are deferred to stacked draft PR `#419` after the availability gate.
Applied to files:
src/aiq_agent/agents/deep_researcher/prompts/researcher.j2
📚 Learning: 2026-07-23T22:20:18.400Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: skills/aiq-research/SKILL.md:0-0
Timestamp: 2026-07-23T22:20:18.400Z
Learning: Applies to skills/aiq-research/**/SKILL.md : If the backend is reachable but `/chat` or async-agent routes fail, report that it is incompatible with the public research flow and offer `aiq-deploy` validation rather than fabricating an answer.
Applied to files:
src/aiq_agent/agents/deep_researcher/prompts/researcher.j2
📚 Learning: 2026-07-23T22:20:18.400Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: skills/aiq-research/SKILL.md:0-0
Timestamp: 2026-07-23T22:20:18.400Z
Learning: Applies to skills/aiq-research/**/SKILL.md : Use this skill for research-shaped requests, including deep research, AI-Q research, or requests to ask AI-Q; do not use it for installation, deployment, lifecycle, UI, CLI, Docker, Helm, or troubleshooting requests.
Applied to files:
src/aiq_agent/agents/deep_researcher/prompts/researcher.j2
📚 Learning: 2026-07-23T22:20:18.400Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: skills/aiq-research/SKILL.md:0-0
Timestamp: 2026-07-23T22:20:18.400Z
Learning: Applies to skills/aiq-research/**/SKILL.md : Present returned reports with citations and source URLs intact; do not truncate citations or source URLs.
Applied to files:
src/aiq_agent/agents/deep_researcher/prompts/researcher.j2
📚 Learning: 2026-06-11T21:21:11.314Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: docs/source/contributing/code-organization.md:0-0
Timestamp: 2026-06-11T21:21:11.314Z
Learning: Applies to docs/source/contributing/src/aiq_agent/agents/shallow_researcher/** : Shallow researcher agent should be organized in `src/aiq_agent/agents/shallow_researcher/`
Applied to files:
src/aiq_agent/agents/deep_researcher/register.pytests/aiq_agent/agents/deep_researcher/test_factory.pysrc/aiq_agent/agents/deep_researcher/factory.py
📚 Learning: 2026-07-01T23:47:04.217Z
Learnt from: KyleZheng1284
Repo: NVIDIA-AI-Blueprints/aiq PR: 298
File: tests/aiq_agent/agents/deep_researcher/sandbox/test_openshell_provider.py:163-496
Timestamp: 2026-07-01T23:47:04.217Z
Learning: In `src/aiq_agent/agents/deep_researcher/sandbox/base.py` (`SandboxProvider`), `cleanup_succeeded` is intentionally cumulative/sticky across the provider's lifetime rather than resettable: once `_cleanup_failed` is set to True (e.g., a stale session's close() fails during `_reset_session()`), it stays True even if a later terminal `close()`/`terminate()` on a replacement session succeeds. This is fail-closed by design — a stale-session close failure means the previously owned physical OpenShell sandbox's deletion was never confirmed, so a later successful cleanup of a different session must not mask that earlier unconfirmed deletion. This is covered by `test_retry_cleanup_failure_remains_terminal_failure` in `tests/aiq_agent/agents/deep_researcher/sandbox/test_openshell_provider.py`, which asserts `DeepAgentsRuntime.finalize()` remains `False` in this scenario with only `started`/`failed` cleanup events emitted.
Applied to files:
tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-08-11T06:35:46.104Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/report_rewriter/agent.py:197-200
Timestamp: 2026-08-11T06:35:46.104Z
Learning: In `src/aiq_agent/agents/report_rewriter/agent.py`, `rewrite_report` is shared by async and inline rewrite paths. It must derive its citation source allowlist only from the canonical `original_report` and durable `parent_context`. Do not widen its API to accept caller-supplied sources or change its string-only return contract solely to remove bounded duplicate parsing.
Applied to files:
src/aiq_agent/agents/deep_researcher/custom_middleware.py
📚 Learning: 2026-08-11T06:34:42.948Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/chat_researcher/agent.py:183-183
Timestamp: 2026-08-11T06:34:42.948Z
Learning: In `src/aiq_agent/agents/chat_researcher/agent.py`, `ChatResearcherAgent.validate_deep_research_tools_fn` is an injected validator contract whose returned `error_msg` can contain arbitrary sensitive text. Return `error_msg` to the caller when required, but use `log_content_metadata(error_msg)` for production logging so centralized logs do not contain the diagnostic content.
Applied to files:
src/aiq_agent/agents/deep_researcher/custom_middleware.py
🪛 ast-grep (0.45.1)
src/aiq_agent/agents/deep_researcher/custom_middleware.py
[info] 1453-1453: use jsonify instead of json.dumps for JSON output
Context: json.dumps(args, sort_keys=True, separators=(",", ":"), default=repr)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🔇 Additional comments (16)
tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py (1)
45-49: LGTM!Also applies to: 52-84, 87-192, 194-239, 242-295, 298-351, 354-396, 399-439, 442-492, 495-560, 631-690
tests/aiq_agent/agents/deep_researcher/test_factory.py (1)
523-524: LGTM!docs/source/customization/configuration-reference.md (1)
490-494: LGTM!Also applies to: 516-516
src/aiq_agent/agents/deep_researcher/models/loop_guard.py (1)
25-80: LGTM!src/aiq_agent/agents/deep_researcher/models/__init__.py (1)
16-16: LGTM!Also applies to: 44-44
src/aiq_agent/agents/deep_researcher/register.py (1)
122-129: LGTM!Also applies to: 262-262, 301-301
src/aiq_agent/agents/deep_researcher/researcher_context.py (1)
25-54: LGTM!src/aiq_agent/agents/deep_researcher/custom_middleware.py (3)
1434-1458: LGTM!
1565-1596: LGTM!
1524-1545: 🩺 Stability & AvailabilityNo guard change is required.
_request_tool_nameresolves bound tools through.name.run_research_batchinvokes the researcher withainvoke, so tool calls useawrap_tool_call. No synchronous researcher tool-call path exists in this repository.> Likely an incorrect or invalid review comment.src/aiq_agent/agents/deep_researcher/agent.py (1)
92-92: LGTM!Also applies to: 113-114, 124-124, 153-153, 216-216
src/aiq_agent/agents/deep_researcher/factory.py (2)
90-96: LGTM!
290-306: LGTM!Also applies to: 335-335, 349-353, 618-621, 669-673
src/aiq_agent/agents/deep_researcher/tools/research.py (2)
94-119: LGTM!
75-82: 🩺 Stability & AvailabilityKeep the inherited
recursion_limitremoval.langchain-core==1.4.8preserves the non-default bound limit when the invoke-time config omitsrecursion_limit.src/aiq_agent/agents/deep_researcher/prompts/researcher.j2 (1)
57-58: LGTM!Also applies to: 69-74
| async def _guard_source_call( | ||
| self, | ||
| request, | ||
| handler, | ||
| state: ResearcherRunGuardState, | ||
| tool_call: dict, | ||
| name: str, | ||
| ): | ||
| """Enforce the source-call budget and the identical-request limit for one invocation.""" | ||
| budget = self._config.max_source_calls_per_query | ||
| if state.exhausted or state.source_call_count >= budget: | ||
| self._mark_exhausted(state, "total source-call budget") | ||
| logger.warning( | ||
| "Researcher loop guard blocked source call | invocation=%s tool=%s calls=%d/%d reason=total_budget", | ||
| state.invocation_id, | ||
| name, | ||
| state.source_call_count, | ||
| budget, | ||
| ) | ||
| return self._blocked_result(tool_call, "the total source-call budget") | ||
|
|
||
| signature = _canonical_source_signature(name, tool_call.get("args", {})) | ||
| identical_count = state.source_signature_counts.get(signature, 0) | ||
| if identical_count >= self._config.max_identical_source_calls: | ||
| self._mark_exhausted(state, "repeated source-call signature") | ||
| logger.warning( | ||
| "Researcher loop guard blocked repeated source call | " | ||
| "invocation=%s tool=%s repeats=%d/%d reason=repeated_signature", | ||
| state.invocation_id, | ||
| name, | ||
| identical_count, | ||
| self._config.max_identical_source_calls, | ||
| ) | ||
| return self._blocked_result(tool_call, "the repeated source-call limit") | ||
|
|
||
| # Count before awaiting so source calls dispatched together in one assistant turn share | ||
| # a single hard ceiling instead of all passing the check and then all executing. | ||
| state.source_call_count += 1 | ||
| state.source_signature_counts[signature] = identical_count + 1 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Hitting max_identical_source_calls ends the entire source budget; the configuration description does not say so.
_guard_source_call calls _mark_exhausted(state, "repeated source-call signature") on the repeat path. _filter_tools then withdraws every source tool and think for the rest of the invocation. With the defaults, a model that issues the same query a third time loses the remaining 7 or more source calls, even though the total budget is untouched.
The class docstring states this coupling, but ResearcherLoopGuardConfig.max_identical_source_calls in src/aiq_agent/agents/deep_researcher/models/loop_guard.py (Lines 61-71) describes only a per-signature cap. An operator reading the configuration reference will not expect a duplicate query to terminate research. _blocked_result also reports "The source budget is exhausted" for this path, which reinforces the wrong mental model.
Pick one of two resolutions:
- Keep the behavior, and state it in the field description and in
docs/source/customization/configuration-reference.md: reaching this limit ends the invocation's research. - Or block only the repeated signature, let other signatures continue, and reserve exhaustion for the total budget.
♻️ Option 2: block the repeat without exhausting the invocation
signature = _canonical_source_signature(name, tool_call.get("args", {}))
identical_count = state.source_signature_counts.get(signature, 0)
if identical_count >= self._config.max_identical_source_calls:
- self._mark_exhausted(state, "repeated source-call signature")
logger.warning(
"Researcher loop guard blocked repeated source call | "
"invocation=%s tool=%s repeats=%d/%d reason=repeated_signature",
state.invocation_id,
name,
identical_count,
self._config.max_identical_source_calls,
)
return self._blocked_result(tool_call, "the repeated source-call limit")_blocked_result then needs a variant that omits the "budget is exhausted" sentence for this path.
🤖 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/aiq_agent/agents/deep_researcher/custom_middleware.py` around lines 1598
- 1636, Update _guard_source_call so reaching max_identical_source_calls blocks
only the repeated signature while allowing other source signatures to continue;
do not call _mark_exhausted on this path. Add a _blocked_result variant or
parameter that omits the “source budget is exhausted” message for duplicate-call
blocks, while preserving exhaustion for the total source-call budget path.
| middleware = list(common_middleware) | ||
| tool_retry_index = next(i for i, item in enumerate(middleware) if isinstance(item, ToolRetryMiddleware)) | ||
| middleware.insert( | ||
| tool_retry_index, | ||
| ResearcherLoopGuardMiddleware( | ||
| source_tool_names=tool_set.source_tool_names, | ||
| config=researcher_loop_guard, | ||
| ), | ||
| ) | ||
| return middleware |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
next() without a default raises an opaque error if ToolRetryMiddleware is absent.
build_common_middleware always appends ToolRetryMiddleware today, so the current call path is safe. But build_researcher_middleware accepts common_middleware from any caller. If a caller supplies a stack without ToolRetryMiddleware, this line raises a bare StopIteration during agent construction, with no indication of the cause. Inside a generator frame that surfaces as RuntimeError: generator raised StopIteration.
Fail with an explicit message, or append the guard at the end.
♻️ Proposed fix
middleware = list(common_middleware)
- tool_retry_index = next(i for i, item in enumerate(middleware) if isinstance(item, ToolRetryMiddleware))
+ tool_retry_index = next(
+ (i for i, item in enumerate(middleware) if isinstance(item, ToolRetryMiddleware)),
+ None,
+ )
+ if tool_retry_index is None:
+ raise ValueError(
+ "build_researcher_middleware requires ToolRetryMiddleware in the shared stack; "
+ "the loop guard must sit outside tool retries so a retried source request is counted once"
+ )
middleware.insert(📝 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.
| middleware = list(common_middleware) | |
| tool_retry_index = next(i for i, item in enumerate(middleware) if isinstance(item, ToolRetryMiddleware)) | |
| middleware.insert( | |
| tool_retry_index, | |
| ResearcherLoopGuardMiddleware( | |
| source_tool_names=tool_set.source_tool_names, | |
| config=researcher_loop_guard, | |
| ), | |
| ) | |
| return middleware | |
| middleware = list(common_middleware) | |
| tool_retry_index = next( | |
| (i for i, item in enumerate(middleware) if isinstance(item, ToolRetryMiddleware)), | |
| None, | |
| ) | |
| if tool_retry_index is None: | |
| raise ValueError( | |
| "build_researcher_middleware requires ToolRetryMiddleware in the shared stack; " | |
| "the loop guard must sit outside tool retries so a retried source request is counted once" | |
| ) | |
| middleware.insert( | |
| tool_retry_index, | |
| ResearcherLoopGuardMiddleware( | |
| source_tool_names=tool_set.source_tool_names, | |
| config=researcher_loop_guard, | |
| ), | |
| ) | |
| return middleware |
🤖 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/aiq_agent/agents/deep_researcher/factory.py` around lines 307 - 316,
Update build_researcher_middleware’s ToolRetryMiddleware lookup to handle
common_middleware stacks that lack that middleware without leaking
StopIteration. Either raise an explicit construction error describing the
missing ToolRetryMiddleware or append ResearcherLoopGuardMiddleware at the end,
while preserving insertion before ToolRetryMiddleware when it is present.
| class TestPromptRendering: | ||
| """The researcher prompt must stay StrictUndefined-safe and keep today's guidance.""" | ||
|
|
||
| _SOFT_GUIDANCE = "Default source budget per ResearchQuery" | ||
| _HARD_CEILING = "Hard limit (runtime-enforced)" | ||
| _WITHDRAWAL = "When source tools are withdrawn, research is over" | ||
|
|
||
| @staticmethod | ||
| def _render(**values) -> str: | ||
| return render_prompt_template( | ||
| _RESEARCHER_PROMPT, | ||
| current_datetime="2026-08-10", | ||
| user_info=None, | ||
| available_documents=[], | ||
| execution_enabled=False, | ||
| tools=[{"name": "web_search_tool", "description": "search"}], | ||
| **values, | ||
| ) | ||
|
|
||
| def test_it_renders_when_the_new_variables_are_absent(self): | ||
| """`| default(false)` is mandatory: StrictUndefined would raise on a bare `{% if %}`.""" | ||
| rendered = self._render() | ||
|
|
||
| assert self._SOFT_GUIDANCE in rendered | ||
| assert self._HARD_CEILING not in rendered | ||
|
|
||
| def test_disabled_mode_retains_todays_guidance_only(self): | ||
| """Turning the guard off must restore exactly today's prompt.""" | ||
| rendered = self._render( | ||
| researcher_loop_guard_enabled=False, | ||
| researcher_max_source_calls=10, | ||
| researcher_max_identical_source_calls=2, | ||
| ) | ||
|
|
||
| assert self._SOFT_GUIDANCE in rendered | ||
| assert "Do NOT get stuck retrying" in rendered | ||
| assert self._HARD_CEILING not in rendered | ||
| assert self._WITHDRAWAL not in rendered | ||
|
|
||
| def test_enabled_mode_adds_the_ceiling_without_removing_the_guidance(self): | ||
| """The behavioural budget and the backstop are separate instructions.""" | ||
| rendered = self._render( | ||
| researcher_loop_guard_enabled=True, | ||
| researcher_max_source_calls=10, | ||
| researcher_max_identical_source_calls=2, | ||
| ) | ||
|
|
||
| assert self._SOFT_GUIDANCE in rendered | ||
| assert "at most 10 source-tool calls" in rendered | ||
| assert "at most 2 call(s) with identical tool arguments" in rendered | ||
| assert "a backstop, not a target" in rendered | ||
|
|
||
| def test_enabled_mode_explains_the_graceful_exit(self): | ||
| """The ceiling states the limit; these bullets state what to do when it is reached.""" | ||
| rendered = self._render( | ||
| researcher_loop_guard_enabled=True, | ||
| researcher_max_source_calls=10, | ||
| researcher_max_identical_source_calls=2, | ||
| ) | ||
|
|
||
| assert self._WITHDRAWAL in rendered | ||
| assert "do not substitute `ls`, `read_file`, `glob`, `grep`" in rendered | ||
| assert "`ResearchGap`" in rendered | ||
| assert "Set `evidence_judgment` to reflect the truncation" in rendered | ||
| # Model-agnostic: never "call the ResearchNotes tool" - provider-strategy models have none. | ||
| assert "call the ResearchNotes tool" not in rendered |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Reduce exact prompt-text assertions in TestPromptRendering.
The class asserts many literal fragments from researcher.j2 ("at most 10 source-tool calls", "a backstop, not a target", "do not substitute \ls`, `read_file`, `glob`, `grep`", "Set `evidence_judgment` to reflect the truncation"). Prompt wording changes often, so these assertions break without a behavior regression. Keep the structural checks that carry real contract value — StrictUndefined rendering without the new variables, the enabled/disabled branch toggling, the rendered numeric limits, and the negative assertion on "call the ResearchNotes tool"` — and drop or loosen the pure wording assertions.
Based on learnings, avoid asserting exact substrings from deep-researcher prompt templates since prompt wording changes frequently; prefer testing structural/behavioral properties instead.
🤖 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 `@tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py` around
lines 563 - 628, Reduce brittle literal-text assertions in
TestPromptRendering.test_enabled_mode_adds_the_ceiling_without_removing_the_guidance
and test_enabled_mode_explains_the_graceful_exit. Retain StrictUndefined
coverage, enabled/disabled branch checks, rendered numeric limits, and the
negative ResearchNotes assertion; remove or loosen wording-only checks such as
“a backstop,” withdrawal instructions, tool names, ResearchGap, and
evidence_judgment text.
Source: Learnings
7d65c7c to
2ba2365
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@docs/source/customization/configuration-reference.md`:
- Around line 593-596: Update the breaker troubleshooting guidance near the
`Researcher source-call budget reached` and `Researcher loop guard blocked`
messages: instruct operators to inspect the exhaustion reason before changing
limits, avoid recommending that every triggered limit be raised, and clarify
that `max_identical_source_calls` should not be increased when repeated requests
indicate model looping; raise a limit only when the expected workload
legitimately reaches it.
🪄 Autofix
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: ASSERTIVE
Plan: Enterprise
Run ID: 4331f365-1acc-4a6d-b73f-fd6f26a176b8
📒 Files selected for processing (3)
docs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📜 Review details
⏰ Context from checks skipped due to timeout. (3)
- GitHub Check: Script Validation
- GitHub Check: Pytest and Coverage
- GitHub Check: Lint and Hooks
🧰 Additional context used
📓 Path-based instructions (9)
**/*
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*: Do not include secrets, credentials, private hostnames, internal-only logs, customer data, or generated local artifacts.
Add or update tests for behavior changes.
**/*: For substantial behavior, authentication, UI, or architecture changes, open a design discussion before coding rather than landing a large unreviewed change.
Keep changes scoped to this repository and avoid editing adjacent repositories; treat eachsources/*package independently and prefer the smallest package-scoped change.
Keep pull requests scoped, avoid unrelated files and generated artifacts, provide validation evidence, and ensure every commit has DCO sign-off.
Files:
docs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.{py,pyi,js,jsx,ts,tsx,yml,yaml,json,env,md}
📄 CodeRabbit inference engine (AGENTS.md)
Never commit secrets, tokens, or environment-specific hostnames; use environment variables and
SecretStr, and resolve API keys at runtime.
Files:
docs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
docs/source/**/*
📄 CodeRabbit inference engine (AGENTS.md)
Update canonical documentation under
docs/source/when behavior, configuration, or workflows change; do not duplicate full documentation pages in skills.
Files:
docs/source/customization/configuration-reference.md
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}
⚙️ CodeRabbit configuration file
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}: Review documentation for command accuracy, branch-name consistency, current CI and copy-pr-bot behavior, public
vs internal boundary clarity, stale examples, and links that no longer match the repository layout.
Files:
docs/source/customization/configuration-reference.md
**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*.py: Runuv run ruff check .anduv run ruff format --check .for root Python changes.
Runuv run pytestfor root project Python changes.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.{py,pyi}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{py,pyi}: Format and lint Python code with Ruff using line length 120, Python 3.11 targeting, rules E, F, W, I, PL, and UP, with single-line imports; do not reformat unrelated code.
Missing-secret paths must degrade gracefully by stubbing or skipping rather than crashing or leaking secrets.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.{py,pyi,js,jsx,ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Never print or log secret values, including in tool output or error messages.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Respect authenticated data sources by honoringrequires_auth, passing through per-user tokens, and using backend token validators; apply owner guardrails before loading protected report or artifact context into an agent.
Do not weaken or bypassAuthMiddleware, authentication validators, or authentication gating without prior design discussion.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.py
src/aiq_agent/agents/**/*
⚙️ CodeRabbit configuration file
src/aiq_agent/agents/**/*: Review agent changes for research workflow correctness, graph state transitions, prompt/tool contracts,
HITL behavior, and failure handling. Flag changes that weaken source attribution, report generation,
async cancellation, checkpointing, or data-source selection without focused tests and docs.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.py
🧠 Learnings (5)
📚 Learning: 2026-06-11T21:21:11.314Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: docs/source/contributing/code-organization.md:0-0
Timestamp: 2026-06-11T21:21:11.314Z
Learning: Applies to docs/source/contributing/src/aiq_agent/agents/deep_researcher/** : Deep researcher agent should be organized in `src/aiq_agent/agents/deep_researcher/` with implementation details documented in README.md
Applied to files:
docs/source/customization/configuration-reference.mdtests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-08-11T06:35:46.104Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/report_rewriter/agent.py:197-200
Timestamp: 2026-08-11T06:35:46.104Z
Learning: In `src/aiq_agent/agents/report_rewriter/agent.py`, `rewrite_report` is shared by async and inline rewrite paths. It must derive its citation source allowlist only from the canonical `original_report` and durable `parent_context`. Do not widen its API to accept caller-supplied sources or change its string-only return contract solely to remove bounded duplicate parsing.
Applied to files:
docs/source/customization/configuration-reference.mdtests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-08-11T06:34:44.687Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/deep_researcher/register.py:302-307
Timestamp: 2026-08-11T06:34:44.687Z
Learning: In Python logging code that handles potentially sensitive exceptions, do not add `exc_info=True` solely to restore stack traces, because standard traceback formatting includes `str(exception)` and may expose provider, customer, or credential-bearing content. When sensitive-content redaction is required, log the exception type together with `log_content_metadata(exception)` instead.
Applied to files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-07-06T23:55:46.952Z
Learnt from: cdgamarose-nv
Repo: NVIDIA-AI-Blueprints/aiq PR: 311
File: src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2:74-81
Timestamp: 2026-07-06T23:55:46.952Z
Learning: In `tests/aiq_agent/agents/deep_researcher/test_agent.py`, avoid asserting exact substrings from the orchestrator/writer prompt templates (e.g., `src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2`) since prompt wording changes frequently. Prefer testing structural/behavioral properties instead.
Applied to files:
tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-07-01T23:47:04.217Z
Learnt from: KyleZheng1284
Repo: NVIDIA-AI-Blueprints/aiq PR: 298
File: tests/aiq_agent/agents/deep_researcher/sandbox/test_openshell_provider.py:163-496
Timestamp: 2026-07-01T23:47:04.217Z
Learning: In `src/aiq_agent/agents/deep_researcher/sandbox/base.py` (`SandboxProvider`), `cleanup_succeeded` is intentionally cumulative/sticky across the provider's lifetime rather than resettable: once `_cleanup_failed` is set to True (e.g., a stale session's close() fails during `_reset_session()`), it stays True even if a later terminal `close()`/`terminate()` on a replacement session succeeds. This is fail-closed by design — a stale-session close failure means the previously owned physical OpenShell sandbox's deletion was never confirmed, so a later successful cleanup of a different session must not mask that earlier unconfirmed deletion. This is covered by `test_retry_cleanup_failure_remains_terminal_failure` in `tests/aiq_agent/agents/deep_researcher/sandbox/test_openshell_provider.py`, which asserts `DeepAgentsRuntime.finalize()` remains `False` in this scenario with only `started`/`failed` cleanup events emitted.
Applied to files:
tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
🔇 Additional comments (6)
tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py (2)
566-628: Reduce wording-specific prompt assertions.The tests still assert literal prompt text that can change without a behavior regression. Keep structural checks, branch checks, and rendered limit checks.
Based on learnings, avoid asserting exact substrings from deep-researcher prompt templates since prompt wording changes frequently; prefer testing structural or behavioral properties.
Source: Learnings
18-561: LGTM!Also applies to: 631-693
docs/source/customization/configuration-reference.md (2)
550-562: Document themax_consecutive_thinksthreshold outcome.The table states what the limit counts but not what occurs at the threshold. Document that the threshold
thinkcall executes with a warning, then onlythinkis withdrawn while source tools remain available.Source: Path instructions
490-516: LGTM!Also applies to: 564-591
src/aiq_agent/agents/deep_researcher/models/loop_guard.py (2)
20-22: 📐 Maintainability & Code QualityRun the required root Python validation.
Run
uv run ruff check .,uv run ruff format --check ., anduv run pytest. Attach the results before merge.As per coding guidelines, “Run
uv run ruff check .anduv run ruff format --check .” and “Runuv run pytest” for root Python changes.Source: Coding guidelines
25-80: LGTM!
smasurekar
changed the title
2ba2365 to
8d840da
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/aiq_agent/agents/deep_researcher/custom_middleware.py`:
- Around line 1649-1658: Update the guard around _mark_exhausted in the research
loop so an already-exhausted state preserves its existing exhaustion_reason,
including repeated-signature cases, instead of replacing it with "total
source-call budget". Log the preserved reason while retaining the total-budget
reason only when this branch newly exhausts the state.
In `@src/aiq_agent/agents/deep_researcher/models/loop_guard.py`:
- Around line 29-32: Update SOURCE_CALLS_PER_QUERY_CEILING to reference the
canonical resource_limits.DEFAULT_MAX_SOURCE_TOOL_CALLS constant instead of
duplicating the literal 100, adding the necessary import while preserving the
existing ceiling behavior.
In `@src/aiq_agent/agents/deep_researcher/tools/research.py`:
- Around line 117-128: The exception wrappers in _run_research_queries currently
expose raw exc text that run_research_batch returns to callers. Replace
user-facing messages with stable error codes or exception-type identifiers,
while preserving the original exception only through a redacted internal
diagnostic path; keep the per-item failure distinction for worker execution and
invalid ResearchNotes validation.
🪄 Autofix
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: ASSERTIVE
Plan: Enterprise
Run ID: 463afab7-2823-425f-9fb0-1e1515d31b30
📒 Files selected for processing (7)
docs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📜 Review details
🧰 Additional context used
📓 Path-based instructions (9)
**/*
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*: Do not include secrets, credentials, private hostnames, internal-only logs, customer data, or generated local artifacts.
Add or update tests for behavior changes.
**/*: For substantial behavior, authentication, UI, or architecture changes, open a design discussion before coding rather than landing a large unreviewed change.
Keep changes scoped to this repository and avoid editing adjacent repositories; treat eachsources/*package independently and prefer the smallest package-scoped change.
Keep pull requests scoped, avoid unrelated files and generated artifacts, provide validation evidence, and ensure every commit has DCO sign-off.
Files:
src/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pydocs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*.py: Runuv run ruff check .anduv run ruff format --check .for root Python changes.
Runuv run pytestfor root project Python changes.
Files:
src/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.{py,pyi}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{py,pyi}: Format and lint Python code with Ruff using line length 120, Python 3.11 targeting, rules E, F, W, I, PL, and UP, with single-line imports; do not reformat unrelated code.
Missing-secret paths must degrade gracefully by stubbing or skipping rather than crashing or leaking secrets.
Files:
src/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.{py,pyi,js,jsx,ts,tsx,yml,yaml,json,env,md}
📄 CodeRabbit inference engine (AGENTS.md)
Never commit secrets, tokens, or environment-specific hostnames; use environment variables and
SecretStr, and resolve API keys at runtime.
Files:
src/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pydocs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.{py,pyi,js,jsx,ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Never print or log secret values, including in tool output or error messages.
Files:
src/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Respect authenticated data sources by honoringrequires_auth, passing through per-user tokens, and using backend token validators; apply owner guardrails before loading protected report or artifact context into an agent.
Do not weaken or bypassAuthMiddleware, authentication validators, or authentication gating without prior design discussion.
Files:
src/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
src/aiq_agent/agents/**/*
⚙️ CodeRabbit configuration file
src/aiq_agent/agents/**/*: Review agent changes for research workflow correctness, graph state transitions, prompt/tool contracts,
HITL behavior, and failure handling. Flag changes that weaken source attribution, report generation,
async cancellation, checkpointing, or data-source selection without focused tests and docs.
Files:
src/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.py
docs/source/**/*
📄 CodeRabbit inference engine (AGENTS.md)
Update canonical documentation under
docs/source/when behavior, configuration, or workflows change; do not duplicate full documentation pages in skills.
Files:
docs/source/customization/configuration-reference.md
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}
⚙️ CodeRabbit configuration file
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}: Review documentation for command accuracy, branch-name consistency, current CI and copy-pr-bot behavior, public
vs internal boundary clarity, stale examples, and links that no longer match the repository layout.
Files:
docs/source/customization/configuration-reference.md
🧠 Learnings (7)
📚 Learning: 2026-06-11T21:21:11.314Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: docs/source/contributing/code-organization.md:0-0
Timestamp: 2026-06-11T21:21:11.314Z
Learning: Applies to docs/source/contributing/src/aiq_agent/agents/deep_researcher/** : Deep researcher agent should be organized in `src/aiq_agent/agents/deep_researcher/` with implementation details documented in README.md
Applied to files:
src/aiq_agent/agents/deep_researcher/researcher_context.pydocs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-06-11T21:21:11.314Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: docs/source/contributing/code-organization.md:0-0
Timestamp: 2026-06-11T21:21:11.314Z
Learning: Applies to docs/source/contributing/src/aiq_agent/agents/shallow_researcher/** : Shallow researcher agent should be organized in `src/aiq_agent/agents/shallow_researcher/`
Applied to files:
src/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/factory.py
📚 Learning: 2026-08-11T06:34:44.687Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/deep_researcher/register.py:302-307
Timestamp: 2026-08-11T06:34:44.687Z
Learning: In Python logging code that handles potentially sensitive exceptions, do not add `exc_info=True` solely to restore stack traces, because standard traceback formatting includes `str(exception)` and may expose provider, customer, or credential-bearing content. When sensitive-content redaction is required, log the exception type together with `log_content_metadata(exception)` instead.
Applied to files:
src/aiq_agent/agents/deep_researcher/researcher_context.pysrc/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-08-11T06:34:42.948Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/chat_researcher/agent.py:183-183
Timestamp: 2026-08-11T06:34:42.948Z
Learning: In `src/aiq_agent/agents/chat_researcher/agent.py`, `ChatResearcherAgent.validate_deep_research_tools_fn` is an injected validator contract whose returned `error_msg` can contain arbitrary sensitive text. Return `error_msg` to the caller when required, but use `log_content_metadata(error_msg)` for production logging so centralized logs do not contain the diagnostic content.
Applied to files:
src/aiq_agent/agents/deep_researcher/tools/research.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.py
📚 Learning: 2026-07-23T22:20:18.400Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: skills/aiq-research/SKILL.md:0-0
Timestamp: 2026-07-23T22:20:18.400Z
Learning: Applies to skills/aiq-research/scripts/aiq.py : When polling requires an execution method with escalated permissions, request explicit user approval first and explain why; tell the user when deep research is running in the background.
Applied to files:
src/aiq_agent/agents/deep_researcher/tools/research.py
📚 Learning: 2026-08-11T06:35:46.104Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/report_rewriter/agent.py:197-200
Timestamp: 2026-08-11T06:35:46.104Z
Learning: In `src/aiq_agent/agents/report_rewriter/agent.py`, `rewrite_report` is shared by async and inline rewrite paths. It must derive its citation source allowlist only from the canonical `original_report` and durable `parent_context`. Do not widen its API to accept caller-supplied sources or change its string-only return contract solely to remove bounded duplicate parsing.
Applied to files:
src/aiq_agent/agents/deep_researcher/tools/research.pydocs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/custom_middleware.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-07-06T23:55:46.952Z
Learnt from: cdgamarose-nv
Repo: NVIDIA-AI-Blueprints/aiq PR: 311
File: src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2:74-81
Timestamp: 2026-07-06T23:55:46.952Z
Learning: In `tests/aiq_agent/agents/deep_researcher/test_agent.py`, avoid asserting exact substrings from the orchestrator/writer prompt templates (e.g., `src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2`) since prompt wording changes frequently. Prefer testing structural/behavioral properties instead.
Applied to files:
src/aiq_agent/agents/deep_researcher/factory.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
🪛 ast-grep (0.45.1)
src/aiq_agent/agents/deep_researcher/custom_middleware.py
[info] 1456-1456: use jsonify instead of json.dumps for JSON output
Context: json.dumps(args, sort_keys=True, separators=(",", ":"), default=repr)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🔇 Additional comments (15)
src/aiq_agent/agents/deep_researcher/factory.py (4)
309-310: Handle a missingToolRetryMiddleware.Duplicate of the previous review finding.
next(...)raisesStopIterationwhen a caller supplies a middleware stack withoutToolRetryMiddleware.
54-746: 📐 Maintainability & Code QualityProvide required Python validation evidence.
Run the required repository checks for these Python changes before merge.
src/aiq_agent/agents/deep_researcher/factory.py#L54-L746: validate lint, format, and affected wiring.src/aiq_agent/agents/deep_researcher/tools/research.py#L27-L334: validate error handling and ContextVar scoping.tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py#L1-L862: run the full test suite.Run
uv run ruff check .,uv run ruff format --check ., anduv run pytest.Source: Coding guidelines
306-307: 🎯 Functional CorrectnessNo change is needed for adapted source-tool names.
Both wrappers preserve
original_tool.name, sotool_set.source_tool_namesmatching remains valid.> Likely an incorrect or invalid review comment.
464-467: 🩺 Stability & AvailabilityKeep the disabled-guard opt-out behavior.
When
enabled=false, the researcher intentionally inherits the orchestrator’srecursion_limit. The implementation, tests, and documentation are consistent. No changes are required.> Likely an incorrect or invalid review comment.tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py (1)
736-762: Reduce wording-specific prompt assertions.Duplicate of the previous review finding. These checks still couple the test to exact
researcher.j2prose instead of rendered limits and enabled or disabled behavior.Based on learnings, avoid asserting exact substrings from deep-researcher prompt templates because prompt wording changes frequently; prefer structural or behavioral properties.
Source: Learnings
docs/source/customization/configuration-reference.md (1)
614-618: Inspect the exhaustion reason before raising a limit.Duplicate of the previous review finding. Do not recommend raising every triggered breaker. Repeated identical requests can indicate a model loop rather than an undersized limit.
src/aiq_agent/agents/deep_researcher/models/loop_guard.py (2)
102-113: The description omits that this limit exhausts the whole invocation.
_guard_source_callinsrc/aiq_agent/agents/deep_researcher/custom_middleware.py(Line 1663) calls_mark_exhaustedon the repeated-signature path._filter_toolsthen withdraws every source tool. This description presents the limit as a per-signature cap only.
43-101: LGTM!Also applies to: 114-123
src/aiq_agent/agents/deep_researcher/custom_middleware.py (6)
1521-1534:_blocked_resultreports budget exhaustion for the repeated-signature path too.The message states "The source budget is exhausted." even when
reasonis "the repeated source-call limit".
1660-1672: Reachingmax_identical_source_callsexhausts the entire invocation.
_mark_exhaustedon the repeat path withdraws every source tool for the remainder of the invocation, even when the total budget is untouched.
47-49: LGTM!Also applies to: 1434-1461, 1497-1519, 1536-1549, 1606-1637
1674-1690: LGTM!
1588-1593: 🩺 Stability & AvailabilityKeep the asynchronous tool-call hook. All repository researcher execution paths use
ainvoke; no synchronous researcher entry point exists.> Likely an incorrect or invalid review comment.
1569-1578: 🩺 Stability & AvailabilityNo change required for the final-turn override.
ModelRequest.override(tool_choice=None)explicitly clearstool_choice, andcreate_agentappends structured-output tools afterrequest.tools, including when that list is empty.src/aiq_agent/agents/deep_researcher/researcher_context.py (1)
25-55: LGTM!
| if state.exhausted or state.source_call_count >= budget: | ||
| self._mark_exhausted(state, "total source-call budget") | ||
| logger.warning( | ||
| "Researcher loop guard blocked source call | invocation=%s tool=%s calls=%d/%d reason=total_budget", | ||
| state.invocation_id, | ||
| name, | ||
| state.source_call_count, | ||
| budget, | ||
| ) | ||
| return self._blocked_result(tool_call, "the total source-call budget") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
The already-exhausted branch overwrites exhaustion_reason with the wrong reason.
When state.exhausted is already True because of a repeated signature, this branch calls _mark_exhausted(state, "total source-call budget") and logs reason=total_budget. The original cause is lost, and the log misattributes the truncation. Preserve the first recorded reason.
🐛 Proposed fix to preserve the first exhaustion reason
`@staticmethod`
def _mark_exhausted(state: ResearcherRunGuardState, reason: str) -> None:
"""Record why this invocation's research budget ended."""
state.exhausted = True
- state.exhaustion_reason = reason
+ if state.exhaustion_reason is None:
+ state.exhaustion_reason = reason budget = self._config.max_source_calls_per_query
if state.exhausted or state.source_call_count >= budget:
self._mark_exhausted(state, "total source-call budget")
logger.warning(
- "Researcher loop guard blocked source call | invocation=%s tool=%s calls=%d/%d reason=total_budget",
+ "Researcher loop guard blocked source call | invocation=%s tool=%s calls=%d/%d reason=%s",
state.invocation_id,
name,
state.source_call_count,
budget,
+ state.exhaustion_reason,
)📝 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.
| if state.exhausted or state.source_call_count >= budget: | |
| self._mark_exhausted(state, "total source-call budget") | |
| logger.warning( | |
| "Researcher loop guard blocked source call | invocation=%s tool=%s calls=%d/%d reason=total_budget", | |
| state.invocation_id, | |
| name, | |
| state.source_call_count, | |
| budget, | |
| ) | |
| return self._blocked_result(tool_call, "the total source-call budget") | |
| if state.exhausted or state.source_call_count >= budget: | |
| self._mark_exhausted(state, "total source-call budget") | |
| logger.warning( | |
| "Researcher loop guard blocked source call | invocation=%s tool=%s calls=%d/%d reason=%s", | |
| state.invocation_id, | |
| name, | |
| state.source_call_count, | |
| budget, | |
| state.exhaustion_reason, | |
| ) | |
| return self._blocked_result(tool_call, "the total source-call budget") |
🤖 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/aiq_agent/agents/deep_researcher/custom_middleware.py` around lines 1649
- 1658, Update the guard around _mark_exhausted in the research loop so an
already-exhausted state preserves its existing exhaustion_reason, including
repeated-signature cases, instead of replacing it with "total source-call
budget". Log the preserved reason while retaining the total-budget reason only
when this branch newly exhausts the state.
| # Matches ``resource_limits.DEFAULT_MAX_SOURCE_TOOL_CALLS`` (the job-wide ceiling on concrete | ||
| # source calls), so one worker's logical-invocation budget can never be configured above what | ||
| # the whole job is allowed to retrieve. | ||
| SOURCE_CALLS_PER_QUERY_CEILING = 100 |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate the canonical job-wide source-call ceiling and compare it to the guard ceiling.
fd -t f 'resource_limits.py' | xargs rg -n 'DEFAULT_MAX_SOURCE_TOOL_CALLS'
rg -n 'SOURCE_CALLS_PER_QUERY_CEILING|DEFAULT_MAX_SOURCE_TOOL_CALLS' --type=py -C2Repository: NVIDIA-AI-Blueprints/aiq
Length of output: 885
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- loop_guard.py ---'
cat -n src/aiq_agent/agents/deep_researcher/models/loop_guard.py | sed -n '1,90p'
printf '%s\n' '--- resource_limits.py ---'
cat -n src/aiq_agent/agents/deep_researcher/resource_limits.py | sed -n '1,45p;250,285p'
printf '%s\n' '--- references and imports ---'
rg -n -C3 'SOURCE_CALLS_PER_QUERY_CEILING|DEFAULT_MAX_SOURCE_TOOL_CALLS' src tests --glob '*.py'Repository: NVIDIA-AI-Blueprints/aiq
Length of output: 19249
Derive SOURCE_CALLS_PER_QUERY_CEILING from DEFAULT_MAX_SOURCE_TOOL_CALLS.
Both values are currently 100, but the duplicated literal can drift. Import and reuse the canonical constant.
🤖 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/aiq_agent/agents/deep_researcher/models/loop_guard.py` around lines 29 -
32, Update SOURCE_CALLS_PER_QUERY_CEILING to reference the canonical
resource_limits.DEFAULT_MAX_SOURCE_TOOL_CALLS constant instead of duplicating
the literal 100, adding the necessary import while preserving the existing
ceiling behavior.
| except Exception as exc: # noqa: BLE001 - captured as per-item failure | ||
| raise RuntimeError(f"researcher worker failed for query {query.query!r}: {exc}") from exc | ||
|
|
||
| try: | ||
| structured = result.get("structured_response") if isinstance(result, dict) else None | ||
| if structured is None: | ||
| raise ValueError("researcher worker did not return structured ResearchNotes") | ||
| note = ResearchNotes.model_validate(structured) | ||
| except Exception as exc: # noqa: BLE001 - captured as per-item failure | ||
| raise ValueError( | ||
| f"researcher worker returned invalid ResearchNotes for query {query.query!r}: {exc}" | ||
| ) from exc |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Do not return raw exception text through the tool result.
exc can contain provider diagnostics, credentials, or source content. _run_research_queries converts this wrapped error to text, and run_research_batch returns it to the caller. Return a stable error code or exception type instead. Preserve the full exception only in a redacted internal diagnostic path.
Proposed fix
except Exception as exc: # noqa: BLE001 - captured as per-item failure
- raise RuntimeError(f"researcher worker failed for query {query.query!r}: {exc}") from exc
+ raise RuntimeError(
+ f"researcher worker failed for query {query.query!r}: {type(exc).__name__}"
+ ) from None
...
except Exception as exc: # noqa: BLE001 - captured as per-item failure
raise ValueError(
- f"researcher worker returned invalid ResearchNotes for query {query.query!r}: {exc}"
- ) from exc
+ f"researcher worker returned invalid ResearchNotes for query {query.query!r}: "
+ f"{type(exc).__name__}"
+ ) from None📝 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.
| except Exception as exc: # noqa: BLE001 - captured as per-item failure | |
| raise RuntimeError(f"researcher worker failed for query {query.query!r}: {exc}") from exc | |
| try: | |
| structured = result.get("structured_response") if isinstance(result, dict) else None | |
| if structured is None: | |
| raise ValueError("researcher worker did not return structured ResearchNotes") | |
| note = ResearchNotes.model_validate(structured) | |
| except Exception as exc: # noqa: BLE001 - captured as per-item failure | |
| raise ValueError( | |
| f"researcher worker returned invalid ResearchNotes for query {query.query!r}: {exc}" | |
| ) from exc | |
| except Exception as exc: # noqa: BLE001 - captured as per-item failure | |
| raise RuntimeError( | |
| f"researcher worker failed for query {query.query!r}: {type(exc).__name__}" | |
| ) from None | |
| try: | |
| structured = result.get("structured_response") if isinstance(result, dict) else None | |
| if structured is None: | |
| raise ValueError("researcher worker did not return structured ResearchNotes") | |
| note = ResearchNotes.model_validate(structured) | |
| except Exception as exc: # noqa: BLE001 - captured as per-item failure | |
| raise ValueError( | |
| f"researcher worker returned invalid ResearchNotes for query {query.query!r}: " | |
| f"{type(exc).__name__}" | |
| ) from None |
🤖 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/aiq_agent/agents/deep_researcher/tools/research.py` around lines 117 -
128, The exception wrappers in _run_research_queries currently expose raw exc
text that run_research_batch returns to callers. Replace user-facing messages
with stable error codes or exception-type identifiers, while preserving the
original exception only through a redacted internal diagnostic path; keep the
per-item failure distinction for worker execution and invalid ResearchNotes
validation.
Source: Coding guidelines
…vocations Signed-off-by: smasurekar <smasurekar@nvidia.com>
8d840da to
c3877c2
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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 `@docs/source/customization/configuration-reference.md`:
- Line 556: Update the max_consecutive_thinks documentation row to describe the
threshold outcome: the think call still executes with a WARNING: suffix, only
think is temporarily withdrawn while other tools remain available, and any other
tool call re-enables think for the remainder of the invocation. Review nearby
examples and boundary wording for consistency with this reversible behavior and
the permanent exhaustion withdrawal described later.
In `@src/aiq_agent/agents/deep_researcher/models/loop_guard.py`:
- Around line 102-113: Update the max_identical_source_calls description in
src/aiq_agent/agents/deep_researcher/models/loop_guard.py:102-113 to state that
reaching the limit ends the invocation’s research, not just the repeated
request, and mirror that wording in
docs/source/customization/configuration-reference.md. No code change is needed
at src/aiq_agent/agents/deep_researcher/custom_middleware.py:1664-1674; retain
the existing _mark_exhausted coupling.
In `@tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py`:
- Around line 577-583: Update
test_the_guard_sits_before_tool_retry_in_the_researcher_stack to also locate
ToolNameSanitizationMiddleware and assert its index is less than the
ResearcherLoopGuardMiddleware index. Add the middleware import alongside the
existing custom_middleware imports, preserving the existing guard-before-retry
assertion.
- Around line 374-396: Update test_no_state_installed_is_a_pass_through to
explicitly assert that the guard’s state is unset before calling _filter_tools,
using the relevant ContextVar/state accessor. Keep the existing identity
assertion, and leave test_a_healthy_invocation_keeps_every_tool focused on the
healthy-state pass-through path.
🪄 Autofix
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: ASSERTIVE
Plan: Enterprise
Run ID: 43a0994a-1796-4ebb-ba24-e15a15ac5edf
📒 Files selected for processing (4)
docs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/custom_middleware.pysrc/aiq_agent/agents/deep_researcher/models/loop_guard.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📜 Review details
⏰ Context from checks skipped due to timeout. (3)
- GitHub Check: Pytest and Coverage
- GitHub Check: Lint and Hooks
- GitHub Check: Script Validation
🧰 Additional context used
📓 Path-based instructions (9)
**/*
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*: Do not include secrets, credentials, private hostnames, internal-only logs, customer data, or generated local artifacts.
Add or update tests for behavior changes.
**/*: For substantial behavior, authentication, UI, or architecture changes, open a design discussion before coding rather than landing a large unreviewed change.
Keep changes scoped to this repository and avoid editing adjacent repositories; treat eachsources/*package independently and prefer the smallest package-scoped change.
Keep pull requests scoped, avoid unrelated files and generated artifacts, provide validation evidence, and ensure every commit has DCO sign-off.
Files:
docs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.{py,pyi,js,jsx,ts,tsx,yml,yaml,json,env,md}
📄 CodeRabbit inference engine (AGENTS.md)
Never commit secrets, tokens, or environment-specific hostnames; use environment variables and
SecretStr, and resolve API keys at runtime.
Files:
docs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
docs/source/**/*
📄 CodeRabbit inference engine (AGENTS.md)
Update canonical documentation under
docs/source/when behavior, configuration, or workflows change; do not duplicate full documentation pages in skills.
Files:
docs/source/customization/configuration-reference.md
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}
⚙️ CodeRabbit configuration file
{docs/**,README.md,CONTRIBUTING.md,SECURITY.md,CODE-OF-CONDUCT.md}: Review documentation for command accuracy, branch-name consistency, current CI and copy-pr-bot behavior, public
vs internal boundary clarity, stale examples, and links that no longer match the repository layout.
Files:
docs/source/customization/configuration-reference.md
**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
**/*.py: Runuv run ruff check .anduv run ruff format --check .for root Python changes.
Runuv run pytestfor root project Python changes.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.{py,pyi}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{py,pyi}: Format and lint Python code with Ruff using line length 120, Python 3.11 targeting, rules E, F, W, I, PL, and UP, with single-line imports; do not reformat unrelated code.
Missing-secret paths must degrade gracefully by stubbing or skipping rather than crashing or leaking secrets.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
**/*.{py,pyi,js,jsx,ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
Never print or log secret values, including in tool output or error messages.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Respect authenticated data sources by honoringrequires_auth, passing through per-user tokens, and using backend token validators; apply owner guardrails before loading protected report or artifact context into an agent.
Do not weaken or bypassAuthMiddleware, authentication validators, or authentication gating without prior design discussion.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.py
src/aiq_agent/agents/**/*
⚙️ CodeRabbit configuration file
src/aiq_agent/agents/**/*: Review agent changes for research workflow correctness, graph state transitions, prompt/tool contracts,
HITL behavior, and failure handling. Flag changes that weaken source attribution, report generation,
async cancellation, checkpointing, or data-source selection without focused tests and docs.
Files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.py
🧠 Learnings (5)
📚 Learning: 2026-06-11T21:21:11.314Z
Learnt from: CR
Repo: NVIDIA-AI-Blueprints/aiq PR: 0
File: docs/source/contributing/code-organization.md:0-0
Timestamp: 2026-06-11T21:21:11.314Z
Learning: Applies to docs/source/contributing/src/aiq_agent/agents/deep_researcher/** : Deep researcher agent should be organized in `src/aiq_agent/agents/deep_researcher/` with implementation details documented in README.md
Applied to files:
docs/source/customization/configuration-reference.mdtests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-08-11T06:35:46.104Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/report_rewriter/agent.py:197-200
Timestamp: 2026-08-11T06:35:46.104Z
Learning: In `src/aiq_agent/agents/report_rewriter/agent.py`, `rewrite_report` is shared by async and inline rewrite paths. It must derive its citation source allowlist only from the canonical `original_report` and durable `parent_context`. Do not widen its API to accept caller-supplied sources or change its string-only return contract solely to remove bounded duplicate parsing.
Applied to files:
docs/source/customization/configuration-reference.mdsrc/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-08-11T06:34:44.687Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/deep_researcher/register.py:302-307
Timestamp: 2026-08-11T06:34:44.687Z
Learning: In Python logging code that handles potentially sensitive exceptions, do not add `exc_info=True` solely to restore stack traces, because standard traceback formatting includes `str(exception)` and may expose provider, customer, or credential-bearing content. When sensitive-content redaction is required, log the exception type together with `log_content_metadata(exception)` instead.
Applied to files:
src/aiq_agent/agents/deep_researcher/models/loop_guard.pysrc/aiq_agent/agents/deep_researcher/custom_middleware.pytests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
📚 Learning: 2026-08-11T06:34:42.948Z
Learnt from: AjayThorve
Repo: NVIDIA-AI-Blueprints/aiq PR: 429
File: src/aiq_agent/agents/chat_researcher/agent.py:183-183
Timestamp: 2026-08-11T06:34:42.948Z
Learning: In `src/aiq_agent/agents/chat_researcher/agent.py`, `ChatResearcherAgent.validate_deep_research_tools_fn` is an injected validator contract whose returned `error_msg` can contain arbitrary sensitive text. Return `error_msg` to the caller when required, but use `log_content_metadata(error_msg)` for production logging so centralized logs do not contain the diagnostic content.
Applied to files:
src/aiq_agent/agents/deep_researcher/custom_middleware.py
📚 Learning: 2026-07-06T23:55:46.952Z
Learnt from: cdgamarose-nv
Repo: NVIDIA-AI-Blueprints/aiq PR: 311
File: src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2:74-81
Timestamp: 2026-07-06T23:55:46.952Z
Learning: In `tests/aiq_agent/agents/deep_researcher/test_agent.py`, avoid asserting exact substrings from the orchestrator/writer prompt templates (e.g., `src/aiq_agent/agents/deep_researcher/prompts/orchestrator.j2`) since prompt wording changes frequently. Prefer testing structural/behavioral properties instead.
Applied to files:
tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
🪛 ast-grep (0.45.1)
src/aiq_agent/agents/deep_researcher/custom_middleware.py
[info] 1456-1456: use jsonify instead of json.dumps for JSON output
Context: json.dumps(args, sort_keys=True, separators=(",", ":"), default=repr)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py
[info] 458-458: use jsonify instead of json.dumps for JSON output
Context: json.dumps(payload)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🔇 Additional comments (24)
src/aiq_agent/agents/deep_researcher/models/loop_guard.py (3)
29-32:SOURCE_CALLS_PER_QUERY_CEILINGstill duplicates the literal100.The comment documents the intended coupling to
resource_limits.DEFAULT_MAX_SOURCE_TOOL_CALLS, but the value is still an independent literal. Import the canonical constant so the two cannot drift.♻️ Proposed fix
-SOURCE_CALLS_PER_QUERY_CEILING = 100 +SOURCE_CALLS_PER_QUERY_CEILING = DEFAULT_MAX_SOURCE_TOOL_CALLSAdd the import near the other imports:
from ..resource_limits import DEFAULT_MAX_SOURCE_TOOL_CALLS
43-60: LGTM!Also applies to: 73-101, 114-123
62-72: 📐 Maintainability & Code QualityKeep the
enabled=Falserecursion-limit wording.The implementation leaves the inherited orchestrator recursion limit in place when the guard is disabled. The description is accurate and requires no change.
> Likely an incorrect or invalid review comment.src/aiq_agent/agents/deep_researcher/custom_middleware.py (3)
1508-1512:_mark_exhaustedstill overwrites the first exhaustion reason.
_guard_source_callcalls_mark_exhausted(state, "total source-call budget")at Line 1652 even whenstate.exhaustedis alreadyTruefrom the repeated-signature path. The original reason is lost and the log at Line 1654 misattributes the truncation. Keep the first recorded reason.
47-49: LGTM!Also applies to: 1434-1461, 1464-1550, 1552-1588, 1608-1639, 1676-1692
1590-1606: 🎯 Functional CorrectnessUse the asynchronous researcher path. Production code invokes the researcher runnable only through
ainvoke; synchronous tests withdraw tools before completion.tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py (11)
789-815: Prompt-wording assertions remain brittle.Lines 800, 811, 812, and 813 still assert literal wording from
researcher.j2("a backstop, not a target", the tool-name list, theevidence_judgmentsentence). Prompt wording changes without a behavior regression, so these break for no signal. Keep the structural checks that carry contract value: StrictUndefined rendering without the new variables, the enabled/disabled branch toggle, the rendered numeric limits, and the negative assertion on"call the ResearchNotes tool".Based on learnings, avoid asserting exact substrings from deep-researcher prompt templates since prompt wording changes frequently; prefer testing structural or behavioral properties instead.
Source: Learnings
55-105: LGTM!
108-218: LGTM!
221-268: LGTM!
271-340: LGTM!
399-507: LGTM!
510-556: LGTM!
585-599: LGTM!
602-679: LGTM!
682-747: LGTM!
818-915: LGTM!docs/source/customization/configuration-reference.md (7)
616-618: 📐 Maintainability & Code Quality | ⚡ Quick winThe blanket "raise it" advice is still present.
Line 616 states "A breaker that fires on healthy runs is set too low: raise it." A previous review flagged this and the thread is marked as addressed, but the sentence is unchanged. The concern stands for
max_identical_source_calls: that breaker fires because the model re-sends the same request, so raising it permits more duplicate calls instead of fixing the loop. Direct operators to read the exhaustion reason first.As per path instructions, review documentation for command accuracy and stale examples.
📝 Proposed rewording
-`Researcher loop guard forcing finalization` (WARNING, on the turn ceiling). A breaker that fires on healthy runs is -set too low: raise it. The symptom is `ResearchNotes` with spurious `ResearchGap` entries rather than an obvious -failure. +`Researcher loop guard forcing finalization` (WARNING, on the turn ceiling). Read the exhaustion reason before you +change a limit. Raise a limit only when the expected workload legitimately reaches it. Repeated identical requests +usually indicate a model loop rather than an undersized budget. The symptom is `ResearchNotes` with spurious +`ResearchGap` entries rather than an obvious failure.Source: Path instructions
490-495: LGTM!Also applies to: 517-517
559-564: LGTM!
566-577: LGTM!
579-597: LGTM!
599-607: LGTM!
609-616: 🔒 Security & PrivacyNo documentation change is required for these claims. The documented log prefixes and levels match the emitted messages.
chat_researcherresolvesdeep_research_agentby name, andresearcher_loop_guardis passed into the deep researcher configuration.> Likely an incorrect or invalid review comment.
| | `enabled` | `true` | — | Set `false` to disable every limit below at once. There is no partial-disable mode, no tool is withdrawn while it is `false`, and no recursion limit is bound. | | ||
| | `max_source_calls_per_query` | `25` | `100` | Model-issued source-tool invocations for one `ResearchQuery`. | | ||
| | `max_identical_source_calls` | `3` | `10` | Invocations of one source tool with identical arguments. The repeat is not executed. | | ||
| | `max_consecutive_thinks` | `3` | `10` | Uninterrupted `think` calls. Any other tool call resets the streak. | |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document the max_consecutive_thinks outcome, including that the withdrawal is reversible.
The row states what the limit counts and the reset rule, but not what happens at the limit. The behavior differs from the source limits in three ways that the tests pin down:
- The
thinkcall still executes; the guard appends aWARNING:suffix to its result instead of replacing it (test_at_the_threshold_the_result_is_overwritten_and_think_is_blocked). - Only
thinkis withdrawn. Source tools stay available (test_think_blocking_alone_withdraws_only_think). - The withdrawal is not permanent. Any other tool call clears
think_blockedand restoresthinkfor the rest of the invocation (test_any_other_tool_reenables_think_after_the_threshold).
The third point is the one an operator cannot infer from the current text, and it also distinguishes this path from the permanent exhaustion withdrawal described at line 568. A previous review asked for the outcome text; the split of the exhaustion paragraph landed, but the think outcome did not.
As per path instructions, review documentation for stale examples and boundary clarity.
📝 Proposed rewording
-| `max_consecutive_thinks` | `3` | `10` | Uninterrupted `think` calls. Any other tool call resets the streak. |
+| `max_consecutive_thinks` | `3` | `10` | Uninterrupted `think` calls. At the limit the call still runs, its result gains a warning, and `think` alone is withdrawn from later model calls; source tools stay available. Any other tool call resets the streak and restores `think`. |🤖 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 `@docs/source/customization/configuration-reference.md` at line 556, Update the
max_consecutive_thinks documentation row to describe the threshold outcome: the
think call still executes with a WARNING: suffix, only think is temporarily
withdrawn while other tools remain available, and any other tool call re-enables
think for the remainder of the invocation. Review nearby examples and boundary
wording for consistency with this reversible behavior and the permanent
exhaustion withdrawal described later.
Source: Path instructions
| max_identical_source_calls: int = Field( | ||
| default=3, | ||
| ge=1, | ||
| le=IDENTICAL_SOURCE_CALLS_CEILING, | ||
| description=( | ||
| "Maximum invocations of one source tool with identical tool arguments. Argument key " | ||
| "ORDER is canonicalized; case and whitespace are not, so 'AI research' and " | ||
| "'ai research' are distinct requests. Matching is per logical invocation: an " | ||
| "identical batch is caught, but two different batches sharing some queries are not " | ||
| f"deduplicated item by item. May not exceed {IDENTICAL_SOURCE_CALLS_CEILING}." | ||
| ), | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
The repeated-request limit ends all research, and the operator-facing text does not say so. _guard_source_call marks the invocation exhausted on the repeated-signature path, which withdraws every source tool and think, but the configuration description presents the limit as a per-signature cap only.
src/aiq_agent/agents/deep_researcher/models/loop_guard.py#L102-L113: state in themax_identical_source_callsdescription that reaching this limit ends the invocation's research, not only the repeated request. Mirror the same statement indocs/source/customization/configuration-reference.md.src/aiq_agent/agents/deep_researcher/custom_middleware.py#L1664-L1674: if you keep the coupling, no code change is needed here; if you prefer per-signature blocking, drop the_mark_exhaustedcall and add a_blocked_resultvariant that omits the "The source budget is exhausted" sentence.
📍 Affects 2 files
src/aiq_agent/agents/deep_researcher/models/loop_guard.py#L102-L113(this comment)src/aiq_agent/agents/deep_researcher/custom_middleware.py#L1664-L1674
🤖 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/aiq_agent/agents/deep_researcher/models/loop_guard.py` around lines 102 -
113, Update the max_identical_source_calls description in
src/aiq_agent/agents/deep_researcher/models/loop_guard.py:102-113 to state that
reaching the limit ends the invocation’s research, not just the repeated
request, and mirror that wording in
docs/source/customization/configuration-reference.md. No code change is needed
at src/aiq_agent/agents/deep_researcher/custom_middleware.py:1664-1674; retain
the existing _mark_exhausted coupling.
| def test_a_healthy_invocation_keeps_every_tool(self, state): | ||
| """Nothing is withdrawn before a limit is reached.""" | ||
| tools = self._tools() | ||
|
|
||
| assert _guard()._filter_tools(tools) is tools | ||
|
|
||
| def test_filesystem_tools_survive_exhaustion(self, state): | ||
| """Exhaustion ends searching, not reading: /shared context stays reachable. | ||
|
|
||
| This is why the source-call limits alone cannot bound the worker - | ||
| `max_model_turns_per_query` is what closes the loop. See TestTurnCeiling. | ||
| """ | ||
| state.exhausted = True | ||
|
|
||
| names = {t.name for t in _guard()._filter_tools(self._tools())} | ||
|
|
||
| assert {"read_file", "ls"} <= names | ||
|
|
||
| def test_no_state_installed_is_a_pass_through(self): | ||
| """Planner, writer, and orchestrator never install guard state, so nothing is withdrawn.""" | ||
| tools = self._tools() | ||
|
|
||
| assert _guard()._filter_tools(tools) is tools |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Assert the missing-state precondition so the two pass-through tests differ.
test_a_healthy_invocation_keeps_every_tool and test_no_state_installed_is_a_pass_through have identical bodies. Both pass through the same return tools line, because _filter_tools returns early when hidden is empty. The second test therefore does not prove that the state is None branch works; it also passes if that branch is deleted. Assert the precondition explicitly, so a leaked ContextVar from another test cannot make the test vacuous.
♻️ Proposed change
def test_no_state_installed_is_a_pass_through(self):
"""Planner, writer, and orchestrator never install guard state, so nothing is withdrawn."""
+ assert CURRENT_RESEARCHER_GUARD_STATE.get() is None
tools = self._tools()
assert _guard()._filter_tools(tools) is 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.
| def test_a_healthy_invocation_keeps_every_tool(self, state): | |
| """Nothing is withdrawn before a limit is reached.""" | |
| tools = self._tools() | |
| assert _guard()._filter_tools(tools) is tools | |
| def test_filesystem_tools_survive_exhaustion(self, state): | |
| """Exhaustion ends searching, not reading: /shared context stays reachable. | |
| This is why the source-call limits alone cannot bound the worker - | |
| `max_model_turns_per_query` is what closes the loop. See TestTurnCeiling. | |
| """ | |
| state.exhausted = True | |
| names = {t.name for t in _guard()._filter_tools(self._tools())} | |
| assert {"read_file", "ls"} <= names | |
| def test_no_state_installed_is_a_pass_through(self): | |
| """Planner, writer, and orchestrator never install guard state, so nothing is withdrawn.""" | |
| tools = self._tools() | |
| assert _guard()._filter_tools(tools) is tools | |
| def test_no_state_installed_is_a_pass_through(self): | |
| """Planner, writer, and orchestrator never install guard state, so nothing is withdrawn.""" | |
| assert CURRENT_RESEARCHER_GUARD_STATE.get() is None | |
| tools = self._tools() | |
| assert _guard()._filter_tools(tools) is 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 `@tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py` around
lines 374 - 396, Update test_no_state_installed_is_a_pass_through to explicitly
assert that the guard’s state is unset before calling _filter_tools, using the
relevant ContextVar/state accessor. Keep the existing identity assertion, and
leave test_a_healthy_invocation_keeps_every_tool focused on the healthy-state
pass-through path.
| def test_the_guard_sits_before_tool_retry_in_the_researcher_stack(self): | ||
| """Outside the retry, so a retried transient failure costs one unit and not three.""" | ||
| researcher = self._middleware_set().researcher | ||
| guard_index = next(i for i, m in enumerate(researcher) if isinstance(m, ResearcherLoopGuardMiddleware)) | ||
| retry_index = next(i for i, m in enumerate(researcher) if isinstance(m, ToolRetryMiddleware)) | ||
|
|
||
| assert guard_index < retry_index |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Also assert the guard sits after ToolNameSanitizationMiddleware.
The PR states two ordering constraints: the guard runs before ToolRetryMiddleware and after ToolNameSanitizationMiddleware. Only the first constraint is tested. The second one matters for correctness: the guard matches request.tool_call["name"] against _source_tool_names and feeds it into _canonical_source_signature. If the guard moves above sanitization, it sees raw tool names, the source budget stops matching, and the failure is silent rather than loud. Add the second assertion in this test.
As per coding guidelines, "Add or update tests for behavior changes."
♻️ Proposed addition
def test_the_guard_sits_before_tool_retry_in_the_researcher_stack(self):
"""Outside the retry, so a retried transient failure costs one unit and not three."""
researcher = self._middleware_set().researcher
guard_index = next(i for i, m in enumerate(researcher) if isinstance(m, ResearcherLoopGuardMiddleware))
retry_index = next(i for i, m in enumerate(researcher) if isinstance(m, ToolRetryMiddleware))
+ sanitize_index = next(
+ i for i, m in enumerate(researcher) if isinstance(m, ToolNameSanitizationMiddleware)
+ )
- assert guard_index < retry_index
+ assert sanitize_index < guard_index < retry_indexImport the middleware alongside the existing custom_middleware imports:
from aiq_agent.agents.deep_researcher.custom_middleware import ToolNameSanitizationMiddlewareRun the following script to confirm the symbol name and its position in the researcher stack:
#!/bin/bash
# Locate the sanitization middleware class and the researcher stack assembly.
rg -nP --type=py -C3 '\bclass\s+\w*ToolNameSanitization\w*Middleware\b'
rg -nP --type=py -C10 'researcher\s*=\s*\[' src/aiq_agent/agents/deep_researcher/factory.py🤖 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 `@tests/aiq_agent/agents/deep_researcher/test_researcher_loop_guard.py` around
lines 577 - 583, Update
test_the_guard_sits_before_tool_retry_in_the_researcher_stack to also locate
ToolNameSanitizationMiddleware and assert its index is less than the
ResearcherLoopGuardMiddleware index. Add the middleware import alongside the
existing custom_middleware imports, preserving the existing guard-before-retry
assertion.
Source: Coding guidelines
1 participant
Overview
A researcher worker could loop indefinitely — repeating the same search, thinking without acting, or burning turns until the subgraph hit the orchestrator's inherited recursion limit of 2000 and lost all its gathered evidence to a
GraphRecursionError. Prompt guidance alone did not stop it.ResearcherLoopGuardMiddlewareis a deterministic circuit breaker for one researcher worker — a singleResearchQueryinsiderun_research_batch. It runs underneath the model, so it does not depend on the model obeying the prompt. It never terminates the sub-agent: it withdraws tools and steers the model to returnResearchNoteswith explicitResearchGapentries, so a truncated worker still contributes the evidence it already gathered. State lives on aContextVar, so concurrent workers each get an isolated budget and nothing leaks across requests.Triggers (defaults; can be calibrated further):
ToolMessagereturnedthinkthinkwithdrawnTermination is enforced at three levels:
[search(A), search(B), search(C)]runs A (24) and B (25, budget hit → notice appended) and blocks C. Nothing upstream could stop C — the tool list for that turn was fixed before the model produced all three. This is also why the counter increments before awaiting the handler.request.tools. The next turn sees onlyget_verified_sources,ls,read_file,grep,glob→ emitsResearchNotes.Exposed as
researcher_loop_guardonDeepResearchAgentConfig;enabled=falsedisables all three triggers, while the derived recursion limit still applies because inheriting the parent's bound is a defect rather than guard enforcement. The guard sits beforeToolRetryMiddleware(a retried transient failure counts once) and afterToolNameSanitizationMiddleware(matched names are the sanitized ones).DCO sign-off for the squash commit
Signed-off-by: smasurekar smasurekar@nvidia.com
Validation
git commit -sor an equivalent sign-off.Where should reviewers start?
Start with
ResearcherLoopGuardMiddlewareinsrc/aiq_agent/agents/deep_researcher/custom_middleware.py— specifically_guard_source_call, where the counter increments before awaiting the handler so parallel calls in one assistant message share a single ceiling. Thenbuild_researcher_middlewareand the derivedrecursion_limitinfactory.py, and theContextVarscoping intools/research.py::_run_research_query.Related Issues
Summary by CodeRabbit
New Features
Tests