[AjayThorve]

Overview

Fix authenticated token propagation for jobs submitted through POST /v1/jobs/async/submit.

AuthMiddleware validated Bearer and idToken credentials, but the custom async-submit route called get_auth_token() outside a NAT SessionManager request context. The submitter therefore passed None to Dask even though the request had an authenticated principal. Authenticated tools in the worker then ran without the submitting user's identity.

This change makes get_auth_token() consult the request token already validated and stored by AuthMiddleware before falling back to NAT request metadata. Registered token fetchers remain first so each Dask task's job-scoped token stays authoritative. Regression tests exercise the real async-submit route with both Bearer and idToken cookie credentials.

DCO sign-off for the squash commit

Signed-off-by: Ajay Thorve athorve@nvidia.com
Signed-off-by: Ajay Thorve AjayThorve@users.noreply.github.com

Validation

uv run --frozen pytest -q tests/aiq_agent/auth/test_auth_utils.py frontends/aiq_api/tests/test_job_submit_data_sources.py
43 passed, 1 warning in 1.40s

uv run --frozen ruff check .
All checks passed!

uv run --frozen ruff format --check .
371 files already formatted

uv run --frozen pytest
1803 passed, 13 skipped, 24 warnings in 60.54s

• 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. Existing documentation already specifies the intended async-worker token propagation contract.
• 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.
• I replaced the DCO sign-off placeholder with my GitHub commit identity and kept the required angle brackets around the email address.

Where should reviewers start?

Start with src/aiq_agent/auth/utils.py, then review test_submit_job_forwards_authenticated_request_token for the Bearer and cookie regression coverage at the actual REST submission boundary.

Related Issues

• N/A — QA-reported release issue.

Summary by CodeRabbit

Summary by CodeRabbit

• Bug Fixes

  • Improved authentication token handling for asynchronous job submissions.
  • Tokens validated by request authentication are now correctly used when submitting jobs.
  • Established clearer token-precedence behavior across registered token sources, validated request tokens, and fallback credentials.

• Tests

  • Added coverage for bearer-token and cookie-based authentication scenarios for async job submission.
  • Added validation that token precedence is honored during token extraction.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[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: 8fe4b46b-4e79-426d-84ea-469efb9b1bb3

📥 Commits

Reviewing files that changed from the base of the PR and between 92ea652 and fdab3ff.

📒 Files selected for processing (1)

• src/aiq_agent/auth/utils.py

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

**/*.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:

• src/aiq_agent/auth/utils.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/auth/utils.py

{src/aiq_agent/auth/**,frontends/aiq_api/src/aiq_api/auth/**}

⚙️ CodeRabbit configuration file

{src/aiq_agent/auth/**,frontends/aiq_api/src/aiq_api/auth/**}: Review authentication changes for issuer/audience validation, token parsing, error hygiene, logging safety,
and compatibility with local and deployed modes. Do not accept changes that expose tokens, weaken validation,
or blur trusted server-side identity with client-supplied fields.

Files:

• src/aiq_agent/auth/utils.py

🔇 Additional comments (1)

src/aiq_agent/auth/utils.py (1)

192-193: LGTM!

---

Walkthrough

get_auth_token() now prioritizes registered fetchers, then reads validated middleware tokens before NAT Context fallbacks. Unit tests cover precedence, and async job submission tests verify bearer and cookie tokens reach submit_agent_job.

Changes

Authenticated token forwarding

Layer / File(s)	Summary
Middleware token retrieval and precedence src/aiq_agent/auth/utils.py, tests/aiq_agent/auth/test_auth_utils.py	get_auth_token() documents and implements registered-fetcher, middleware, cookie, and header precedence, with tests for middleware retrieval and fetcher priority.
Async submission token forwarding frontends/aiq_api/tests/test_job_submit_data_sources.py	Parametrized endpoint tests verify bearer-header and id-token-cookie credentials are validated and forwarded to submit_agent_job.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant AuthMiddleware
  participant AsyncSubmitEndpoint
  participant submit_agent_job
  Client->>AuthMiddleware: Send bearer header or idToken cookie
  AuthMiddleware-->>AsyncSubmitEndpoint: Provide validated token
  AsyncSubmitEndpoint->>submit_agent_job: Pass auth_token
  submit_agent_job-->>Client: Return HTTP 200

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows Conventional Commits and accurately summarizes the token propagation fix.
Description check	✅ Passed	The description includes all required template sections with concrete implementation, DCO, validation, reviewers, and issue context.
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.}

[coderabbitai]

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 `@src/aiq_agent/auth/utils.py`:
- Around line 192-193: Update the exception handler around token retrieval in
the auth utility to log only a static failure message. Remove the exception
object from the logger.debug call so credential-bearing details from
auth-context failures are never emitted.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: ef799947-e622-40c8-87ee-f6a65399f3d9

📥 Commits

Reviewing files that changed from the base of the PR and between b90d8a1 and 92ea652.

📒 Files selected for processing (3)

• frontends/aiq_api/tests/test_job_submit_data_sources.py
• src/aiq_agent/auth/utils.py
• tests/aiq_agent/auth/test_auth_utils.py

📜 Review details

⏰ Context from checks skipped due to timeout. (3)

• GitHub Check: Lint and Hooks
• GitHub Check: Pytest and Coverage
• GitHub Check: Script Validation

🧰 Additional context used

📓 Path-based instructions (4)

**/*.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:

• frontends/aiq_api/tests/test_job_submit_data_sources.py
• tests/aiq_agent/auth/test_auth_utils.py
• src/aiq_agent/auth/utils.py

**/*test*.py

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Run pytest for all behavior changes in Python code

Files:

• frontends/aiq_api/tests/test_job_submit_data_sources.py
• tests/aiq_agent/auth/test_auth_utils.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/auth/utils.py

{src/aiq_agent/auth/**,frontends/aiq_api/src/aiq_api/auth/**}

⚙️ CodeRabbit configuration file

{src/aiq_agent/auth/**,frontends/aiq_api/src/aiq_api/auth/**}: Review authentication changes for issuer/audience validation, token parsing, error hygiene, logging safety,
and compatibility with local and deployed modes. Do not accept changes that expose tokens, weaken validation,
or blur trusted server-side identity with client-supplied fields.

Files:

• src/aiq_agent/auth/utils.py

🧠 Learnings (1)

📚 Learning: 2026-06-14T17:49:00.640Z

Learnt from: torkian
Repo: NVIDIA-AI-Blueprints/aiq PR: 273
File: frontends/aiq_api/tests/test_sse_reconnect_cursor.py:384-401
Timestamp: 2026-06-14T17:49:00.640Z
Learning: When using `unittest.mock.patch` for code that imports dependencies inside functions/generators (e.g., inside `aiq_api.routes.jobs`), don’t patch via an attribute that doesn’t exist on the consuming module. If the generator does `from ..jobs.event_store import EventStore` inside the generator body, then `aiq_api.routes.jobs` will not have an `EventStore` attribute; patch the source class/method in its defining module instead (e.g., `aiq_api.jobs.event_store.EventStore.get_events_async`). Patching `aiq_api.routes.jobs.EventStore...` would raise `AttributeError` because that symbol is not present at module scope.

Applied to files:

• frontends/aiq_api/tests/test_job_submit_data_sources.py

🔇 Additional comments (3)

src/aiq_agent/auth/utils.py (1)

24-30: LGTM!

Also applies to: 159-162

tests/aiq_agent/auth/test_auth_utils.py (1)

56-72: LGTM!

frontends/aiq_api/tests/test_job_submit_data_sources.py (1)

168-234: LGTM!

[KyleZheng1284]

LGTM

[AjayThorve]

/merge