From 92ea65281543367733eafca0fd39a49902ef260b Mon Sep 17 00:00:00 2001 From: Ajay Thorve Date: Fri, 17 Jul 2026 13:46:25 -0700 Subject: [PATCH 1/2] fix(auth): propagate REST submit token Signed-off-by: Ajay Thorve --- .../tests/test_job_submit_data_sources.py | 69 +++++++++++++++++++ src/aiq_agent/auth/utils.py | 34 +++++++-- tests/aiq_agent/auth/test_auth_utils.py | 19 +++++ 3 files changed, 115 insertions(+), 7 deletions(-) diff --git a/frontends/aiq_api/tests/test_job_submit_data_sources.py b/frontends/aiq_api/tests/test_job_submit_data_sources.py index 42f683e15..b9bbcd450 100644 --- a/frontends/aiq_api/tests/test_job_submit_data_sources.py +++ b/frontends/aiq_api/tests/test_job_submit_data_sources.py @@ -165,6 +165,75 @@ async def test_submit_job_forwards_selected_data_sources(submit_app): assert submitted_job.await_args.kwargs["data_sources"] == ["web_search"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("credential_source", "expected_token"), + [ + pytest.param( + "bearer-header", + "bearer-token", + id="bearer-header", + ), + pytest.param( + "id-token-cookie", + "cookie-token", + id="id-token-cookie", + ), + ], +) +async def test_submit_job_forwards_authenticated_request_token( + submit_app, + monkeypatch, + credential_source, + expected_token, +): + """REST submit captures the validated request token without requiring NAT Context.""" + import aiq_agent.auth + from aiq_agent.auth.utils import get_auth_token as real_get_auth_token + from aiq_api.auth.middleware import AuthMiddleware + + class AcceptTokenValidator: + def can_handle(self, token: str) -> bool: # noqa: ARG002 - accepts both test credentials + return True + + async def validate(self, token: str): + return ( + { + "type": "jwt", + "sub": "user-1", + "email": "user@example.com", + "name": "Test User", + "token": token, + "skip_clarifier": False, + }, + None, + ) + + app, submitted_job, _builder = submit_app + monkeypatch.setattr(aiq_agent.auth, "get_auth_token", real_get_auth_token) + app.add_middleware( + AuthMiddleware, + validators=[AcceptTokenValidator()], + require_auth=True, + external_hostnames={"testserver"}, + ) + + with TestClient(app) as client: + headers = None + if credential_source == "bearer-header": + headers = {"Authorization": f"Bearer {expected_token}"} + else: + client.cookies.set("idToken", expected_token) + response = client.post( + "/v1/jobs/async/submit", + json={"agent_type": "deep_researcher", "input": "query"}, + headers=headers, + ) + + assert response.status_code == 200 + assert submitted_job.await_args.kwargs["auth_token"] == expected_token + + @pytest.mark.asyncio async def test_submit_job_rejects_internal_agent(submit_app, monkeypatch): app, submitted_job, _builder = submit_app diff --git a/src/aiq_agent/auth/utils.py b/src/aiq_agent/auth/utils.py index 39525ea8b..480e90feb 100644 --- a/src/aiq_agent/auth/utils.py +++ b/src/aiq_agent/auth/utils.py @@ -21,9 +21,13 @@ Token sources checked in priority order by ``get_auth_token()``: -1. **NAT/AIQ Context cookies** — ``idToken`` cookie set by the frontend auth layer +1. **Registered token fetchers** — job-scoped sources such as the Dask worker + auth context. +2. **AIQ auth middleware context** — the request token already validated by the + configured ``TokenValidator``. +3. **NAT/AIQ Context cookies** — ``idToken`` cookie set by the frontend auth layer (server / web-UI mode). -2. **NAT/AIQ Context Authorization header** — ``Authorization: Bearer `` sent +4. **NAT/AIQ Context Authorization header** — ``Authorization: Bearer `` sent by API callers who authenticate with a JWT directly. """ @@ -152,10 +156,10 @@ def get_auth_token() -> str | None: Sources checked in order: - Tries registered token fetchers in priority order (highest first), - then falls back to the idToken cookie set by the frontend auth layer. - 1. NAT ``Context`` cookies — ``idToken`` key (server / web-UI mode). - 2. NAT ``Context`` Authorization header — ``Bearer `` (API callers with JWT). + 1. Registered token fetchers in priority order (highest first). + 2. AIQ auth middleware's current request token. + 3. NAT ``Context`` cookies — ``idToken`` key (server / web-UI mode). + 4. NAT ``Context`` Authorization header — ``Bearer `` (API callers with JWT). Returns: ID token string, or ``None`` if no valid token is available. @@ -172,7 +176,23 @@ def get_auth_token() -> str | None: except Exception as e: logger.debug("Registered token fetcher failed: %s", e) - # Default: Context cookies + # Custom AIQ FastAPI routes run under AuthMiddleware but outside a NAT + # SessionManager session, so their headers/cookies are not represented in + # NAT Context. TokenValidator's contract preserves the validated raw token + # in the middleware user context; use that request-scoped value before + # falling back to NAT Context. + try: + from aiq_api.auth.middleware import get_current_user + + current_user = get_current_user() + middleware_token = current_user.get("token") if isinstance(current_user, dict) else None + if isinstance(middleware_token, str) and middleware_token.strip(): + logger.debug("Using token from AIQ auth middleware context") + return middleware_token.strip() + except Exception as e: + logger.debug("Failed to retrieve token from AIQ auth middleware context: %s", e) + + # Fall back to NAT Context cookies and headers. try: from nat.builder.context import Context diff --git a/tests/aiq_agent/auth/test_auth_utils.py b/tests/aiq_agent/auth/test_auth_utils.py index 009fbcd35..cc03eed96 100644 --- a/tests/aiq_agent/auth/test_auth_utils.py +++ b/tests/aiq_agent/auth/test_auth_utils.py @@ -53,6 +53,25 @@ def test_get_auth_token_default_returns_none(): assert get_auth_token() is None +def test_get_auth_token_uses_middleware_request_token(): + """Custom FastAPI routes can read the token already validated by AuthMiddleware.""" + with patch( + "aiq_api.auth.middleware.get_current_user", + return_value={"type": "jwt", "sub": "user-1", "token": "request-token"}, + ): + assert get_auth_token() == "request-token" + + +def test_registered_token_fetcher_precedes_middleware_request_token(): + """A worker's job-scoped token remains authoritative over ambient request context.""" + register_token_fetcher(lambda: "worker-token", priority=5) + with patch( + "aiq_api.auth.middleware.get_current_user", + return_value={"type": "jwt", "sub": "user-1", "token": "request-token"}, + ): + assert get_auth_token() == "worker-token" + + def test_register_token_fetcher_basic(): """A registered fetcher that returns a token is used by get_auth_token.""" register_token_fetcher(lambda: "my-token") From fdab3ff3848560cbf502bf5fc6efb63d3147094b Mon Sep 17 00:00:00 2001 From: Ajay Thorve Date: Fri, 17 Jul 2026 14:00:11 -0700 Subject: [PATCH 2/2] sanitize auth context failure logging Signed-off-by: Ajay Thorve --- src/aiq_agent/auth/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aiq_agent/auth/utils.py b/src/aiq_agent/auth/utils.py index 480e90feb..0da0b99a2 100644 --- a/src/aiq_agent/auth/utils.py +++ b/src/aiq_agent/auth/utils.py @@ -189,8 +189,8 @@ def get_auth_token() -> str | None: if isinstance(middleware_token, str) and middleware_token.strip(): logger.debug("Using token from AIQ auth middleware context") return middleware_token.strip() - except Exception as e: - logger.debug("Failed to retrieve token from AIQ auth middleware context: %s", e) + except Exception: + logger.debug("Failed to retrieve token from AIQ auth middleware context") # Fall back to NAT Context cookies and headers. try: