fix(auth): propagate REST submit token to async workers

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

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 <jwt>`` sent
+4. **NAT/AIQ Context Authorization header** — ``Authorization: Bearer <jwt>`` 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 <jwt>`` (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 <jwt>`` (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:
+        logger.debug("Failed to retrieve token from AIQ auth middleware context")
+
+    # Fall back to NAT Context cookies and headers.
     try:
         from nat.builder.context import Context
 

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")