fix(deep-research): source-routing guard recognizes sandbox

[KyleZheng1284]

Overview

Sandboxed deep-research runs could get stuck: the orchestrator was rejected on every tool call with "Source routing is required before any other tool call," so it never made progress.

Root cause seems to be a key-name mismatch:

• The source-routing guard decides it's safe to proceed by checking state["files"]
  for /shared/source_routing.json.
• But when a sandbox provider is enabled, CompositeBackend drops the /shared/
  prefix before storing the file, so it actually lands under /source_routing.json.
• The guard never found its expected key, so it kept blocking. (Non-sandbox runs
  keep the /shared/ prefix, so they were fine — this only hit sandbox runs.)

Fix: the guard now accepts both key forms (/shared/source_routing.json and the
sandbox /source_routing.json). Small, no behavior change for non-sandbox runs.

Validation

• pytest test_custom_middleware.py test_factory.py → 52 passed, including a new
  test that seeds the sandbox key /source_routing.json and confirms the guard opens.

• ruff check on changed files → clean.

• I ran the relevant local checks or explained why they are not applicable.

• I added or updated tests for behavior changes.

• I updated documentation for user-facing or contributor-facing changes.

• I confirmed this PR does not include secrets, credentials, or internal-only data.

• I certify this contribution under the Developer Certificate of Origin (DCO) and signed my commits with git commit -s or an equivalent sign-off.

Where should reviewers start?

'custom_middleware.py→SourceRoutingGuardMiddleware._routing_complete. Context for why the key gets stripped is in deepagents_runtime.py::_normalize_state_files`.

Related Issues

• Relates to #

Summary by CodeRabbit

• Bug Fixes
  • Improved source-routing detection so the app recognizes routing completion from either the shared or sandbox-local routing file.
  • Normal tool actions now proceed correctly in sandbox environments once routing data is present.
• Tests
  • Added coverage for sandbox-based routing to confirm the expected behavior and response flow.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: 61ec4035-4ea3-438e-958c-53b91f509043

📥 Commits

Reviewing files that changed from the base of the PR and between db86f8f and 4ee8593.

📒 Files selected for processing (2)

• src/aiq_agent/agents/deep_researcher/custom_middleware.py
• tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py

📜 Recent 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 (5)

**/*.py

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Run ruff check and ruff format validation for Python code changes

**/*.py: Python code must be linted and formatted with Ruff using line length 120, target Python 3.11, rule sets E,F,W,I,PL,UP, and isort force-single-line configuration
Never commit secrets, tokens, or environment-specific hostnames; use environment variables and SecretStr instead, resolving API keys at runtime
Never print or log secret values, including in tool output or error messages
Missing-secret paths must degrade gracefully (stub/skip), not crash or leak
Do not hand-reformat unrelated code when making changes; match the existing import and formatting style

Files:

• tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py
• src/aiq_agent/agents/deep_researcher/custom_middleware.py

**/*test*.py

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Run pytest for all behavior changes in Python code

Files:

• tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py

**

⚙️ CodeRabbit configuration file

**:

AI-Q Agent Guidance

Repository-global instructions for coding agents and for humans reviewing
agent-authored changes. These rules apply to every task in this repository.
Task-specific runbooks live in .agents/skills/ — load the
relevant skill before starting a workflow it covers.

Project overview

AI-Q is an NVIDIA AI Blueprint: an enterprise research agent built on the
NeMo Agent Toolkit (NAT). The deployed product is a research blueprint, not
a general skill runtime. New retrieval sources and tools are NAT functions;
agent behavior is driven by workflow YAML, Jinja2 prompts, and a data-source
registry — not by hard-coded logic.

Primary boundaries:

• Backend Python package: src/aiq_agent/.
• Data-source and tool packages: sources/ (each is its own package).
• Frontends and tooling: frontends/ (web UI in frontends/ui/, eval harnesses
  in frontends/benchmarks/).
• Configs, deployment, docs: configs/, deploy/, docs/.

Stay inside this repository. If your workspace also contains adjacent repos
(for example a sibling NeMo-Relay checkout), do not edit them as part of an AI-Q
change. Treat sources/* as independent packages: prefer the smallest change
scoped to the package you are touching.

Repository structure

Path	Purpose
src/aiq_agent/	Backend agent, FastAPI extensions, auth, observability, knowledge
sources/	Data-source / tool packages (e.g. tavily_web_search, google_scholar_paper_search)
configs/	Workflow YAML configs (e.g. config_cli_default.yml)
frontends/ui/	Next.js / React / TypeScript / Tailwind / KUI web UI
frontends/benchmarks/	Eval harnesses: freshqa, deepsearch_qa, deepresearch_bench
deploy/	Docker Compose and Helm/Kubernetes assets; deploy/.env for secrets
docs/source/ ...

Files:

• tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py
• src/aiq_agent/agents/deep_researcher/custom_middleware.py

src/aiq_agent/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/aiq_agent/**/*.py: Respect authenticated data sources by honoring requires_auth, per-user token pass-through, and backend token validators; apply owner guardrails before loading protected report or artifact context into an agent
Do not weaken or bypass AuthMiddleware, validators, or auth gating without a prior design discussion

Files:

• src/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/custom_middleware.py

🔇 Additional comments (2)

src/aiq_agent/agents/deep_researcher/custom_middleware.py (1)

41-45: LGTM!

Also applies to: 58-58

tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py (1)

92-104: LGTM!

---

Walkthrough

The SourceRoutingGuardMiddleware's routing-completion check is extended to accept a sandbox-local source-routing file path (/source_routing.json) in addition to the existing shared path (/shared/source_routing.json), via a new keys tuple. A corresponding unit test validates the sandbox key path.

Changes

Sandbox source-routing key support

Layer / File(s)	Summary
Multi-key routing completion check src/aiq_agent/agents/deep_researcher/custom_middleware.py, tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py	Adds _SOURCE_ROUTING_STATE_KEYS tuple containing both shared and sandbox-local file paths, updates _routing_complete() to check for any accepted key, and adds a test verifying the sandbox key opens the tool gate.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Handler
  participant SourceRoutingGuardMiddleware
  participant AgentState

  Handler->>SourceRoutingGuardMiddleware: tool call request
  SourceRoutingGuardMiddleware->>AgentState: check files for _SOURCE_ROUTING_STATE_KEYS
  AgentState-->>SourceRoutingGuardMiddleware: /shared/source_routing.json or /source_routing.json present
  SourceRoutingGuardMiddleware->>Handler: allow normal tool call
  Handler-->>SourceRoutingGuardMiddleware: ToolMessage

Loading

Related issues: None provided.

Related PRs: None provided.

Suggested labels: None provided.

Suggested reviewers: None provided.

Poem: A rabbit hopped through paths anew, one shared, one sandboxed, both checked through. Two keys now open up the gate, no longer must the shared file wait. Hop, hop, test passed — celebrate! 🐇

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows Conventional Commits and accurately summarizes the sandbox source-routing guard fix.
Description check	✅ Passed	The description includes the required Overview, Validation, reviewer guidance, and issue sections, with only minor placeholder content.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[soumilinandi]

tested PR with this, LGTM!