diff --git a/.secrets.baseline b/.secrets.baseline index 0e6f9882d..b8985fdb5 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -290,5 +290,5 @@ } ] }, - "generated_at": "2026-05-22T20:01:44Z" + "generated_at": "2026-07-02T05:57:01Z" } diff --git a/README.md b/README.md index c77490406..78707e5de 100644 --- a/README.md +++ b/README.md @@ -380,7 +380,7 @@ For development, contribution, and documentation, refer to: - [ ] **[NeMo Guardrails](https://github.com/NVIDIA-NeMo/Guardrails) Integration:** Enhance safety and security guardrails. - [ ] **[NVIDIA Dynamo](https://github.com/ai-dynamo/dynamo) Integration:** Reduce latency via priority scheduling at scale. -- [ ] **MCP Authentication:** Implement secure login/auth for MCP connections. +- [x] **Per-user MCP OAuth:** Connect each signed-in user to protected MCP data sources through the UI. - [x] **Skills & Sandboxing:** Support built-in deep research skills with job-scoped sandbox execution. - [ ] **Custom Skill Management:** Add UI and lifecycle controls for user-provided skill bundles. - [ ] **Dynamic Model Routing:** Allow sub-agents to automatically select the optimal model per task. diff --git a/configs/config_web_frag_mcp_auth.yml b/configs/config_web_frag_mcp_auth.yml new file mode 100644 index 000000000..4e5b1c9ea --- /dev/null +++ b/configs/config_web_frag_mcp_auth.yml @@ -0,0 +1,297 @@ +# Web mode with Foundational RAG PLUS a protected per-user OAuth MCP data source +# (example: Google Drive). This is config_web_frag.yml with the per-user MCP auth +# wiring added, kept as a separate config so the base web config stays minimal. +# +# Features: +# - Web search enabled by default +# - Knowledge retrieval using Foundational RAG. +# Requires a RAG server and ingest server to be running (not deployed by this +# blueprint). Example RAG deployment: https://github.com/NVIDIA-AI-Blueprints/rag/tree/main +# - Protected per-user MCP source (`gdrive`): AIQ surfaces connection state on the +# source card, gates job submission until connected (409 mcp_auth_required), and +# owns the OAuth connect/callback. See the `authentication` / `object_stores` +# sections below for setup (MCP_GDRIVE_URL, AIQ_PUBLIC_URL, MCP_TOKEN_STORE_TYPE). + +general: + use_uvloop: true + telemetry: + logging: + console: + _type: console + level: INFO + # tracing: + # langsmith: # Optional: LangSmith tracing - requires langsmith API key. Set using `export LANGSMITH_API_KEY=` + # _type: langsmith + # project: nvidia-aiq + + front_end: + _type: aiq_api + runner_class: aiq_api.plugin.AIQAPIWorker + # ========================================================================= + # Knowledge API is automatically enabled when knowledge_retrieval function + # is configured + # ========================================================================= + # Async Job API Settings + # ========================================================================= + # Async job infrastructure database (NAT JobStore + EventStore) + # Used by: /v1/jobs/async routes, SSE streaming, job status persistence + # Requires async driver for SQLite (aiosqlite) or PostgreSQL (asyncpg) + # Environment overrides: + # - NAT_JOB_STORE_DB_URL (direct override) + # - NAT_JOB_STORE_DB_URL_DEV / NAT_JOB_STORE_DB_URL_PROD (via NAT_ENV) + db_url: ${NAT_JOB_STORE_DB_URL:-sqlite+aiosqlite:///./jobs.db} + # Job expiry - how long completed jobs stay in database before cleanup + expiry_seconds: 86400 # 24 hours (min: 600, max: 604800/7 days) + cors: + allow_origin_regex: 'http://localhost(:\d+)?|http://127.0.0.1(:\d+)?' + allow_methods: + - GET + - POST + - DELETE + - OPTIONS + allow_headers: + - "*" + allow_credentials: true + expose_headers: + - "*" + +llms: + nemotron_llm_intent: + _type: nim + model_name: nvidia/nemotron-3-super-120b-a12b + base_url: "https://integrate.api.nvidia.com/v1" + temperature: 0.5 + top_p: 0.9 + max_tokens: 4096 + num_retries: 5 + chat_template_kwargs: + enable_thinking: true + + nemotron_super_llm: + _type: nim + model_name: nvidia/nemotron-3-super-120b-a12b + base_url: "https://integrate.api.nvidia.com/v1" + temperature: 0.7 + top_p: 0.7 + max_tokens: 65536 + num_retries: 5 + chat_template_kwargs: + enable_thinking: true + +functions: + # ========================================================================= + # Data Source Registry + # ========================================================================= + # Central registry that controls: + # 1. UI toggles — each source appears as an on/off switch in the frontend + # 2. Per-message filtering — users can select active sources per request + # 3. Tool auto-inheritance — agents with no explicit `tools` list receive + # every tool listed here (use `exclude_tools` on agents to specialize) + # + # Source entry fields: + # id, name, description, tools, requires_auth (default: false), + # default_enabled (default: true) + # + # See docs/source/customization/tools-and-sources.md for full details. + # ========================================================================= + data_sources: + _type: data_source_registry + sources: + - id: web_search + name: "Web Search" + description: "Search the web for real-time information." + tools: + - web_search_tool + - advanced_web_search_tool + - id: knowledge_layer + name: "Knowledge Base" + description: "Search uploaded documents and files." + tools: + - knowledge_search + # Protected per-user MCP source (example: a Google Drive MCP server). `per_user_auth` + # makes AIQ surface connection state on the source card, gate job submission + # until connected (409 mcp_auth_required), and own the OAuth connect/callback. + - id: gdrive + name: "Google Drive" + description: "Search and read your authorized Google Drive files." + default_enabled: false + # requires_auth (the old AIQ-login gate) is left false so the source card + # is visible/Connectable in no-auth local runs; per_user_auth is the gate. + requires_auth: false + per_user_auth: + required: true + provider: google + mcp_server_id: gdrive + auth_provider: mcp_oauth2_gdrive # -> the `authentication` entry below + # Clean names + descriptions so the research agent recognizes these as the + # way to access the user's Drive (the raw MCP tools are terse/blank, which + # made the small model fall back to web search / "no access"). + tool_overrides: + gdrive_search: + alias: google_drive_search + description: >- + Search the USER'S connected Google Drive for files by name or content. + Use this whenever the user asks to find or look up a document in their + Google Drive. Returns matching files with their URLs/ids. + gdrive_get_file: + alias: google_drive_read_file + description: >- + Read the full text contents of a specific Google Drive file given its + URL or id (from google_drive_search). Use this to read or summarize a + document the user referenced in their Google Drive. + gdrive_get_metadata: + alias: google_drive_file_metadata + description: >- + Get metadata (name, type, owner, last modified) for a Google Drive file + by URL or id. + # NOTE: no `tools:` here, and intentionally NO config-declared function group. + # A per_user_mcp_client declared in config gets built by NAT's per-user + # interactive (WebSocket) session builder, which crashes for a user with no + # token. Instead the async-job worker builds the MCP client IN CODE per-job + # (from this source's auth_provider) with the owner's token — see + # aiq_api.mcp_auth.runtime_tools. This keeps the interactive WS path untouched. + + web_search_tool: + _type: tavily_web_search + max_results: 5 + max_content_length: 1000 + + advanced_web_search_tool: + _type: tavily_web_search + max_results: 2 + advanced_search: true + + # Knowledge Retrieval (see sources/knowledge_layer/KNOWLEDGE-LAYER-SETUP.md) + knowledge_search: + _type: knowledge_retrieval + backend: foundational_rag + collection_name: ${COLLECTION_NAME:-test_collection} + top_k: 5 + rag_url: ${RAG_SERVER_URL:-http://localhost:8081} + ingest_url: ${RAG_INGEST_URL:-http://localhost:8082} + timeout: 300 + + # Paper Search (optional - requires SERPER_API_KEY) + # Uncomment the block below and set SERPER_API_KEY to enable academic paper search. + # paper_search_tool: + # _type: paper_search + # max_results: 5 + # serper_api_key: ${SERPER_API_KEY} + + # ========================================================================= + # Agents — inherit all registry tools; use exclude_tools to specialize + # ========================================================================= + intent_classifier: + _type: intent_classifier + llm: nemotron_llm_intent + # tools: omitted -> inherits all from data_source_registry + # exclude_tools: [] + + clarifier_agent: + _type: clarifier_agent + llm: nemotron_super_llm + planner_llm: nemotron_super_llm + # tools: omitted -> inherits all from data_source_registry + # exclude_tools: [] + max_turns: 3 + enable_plan_approval: true + log_response_max_chars: 2000 + verbose: true + + shallow_research_agent: + _type: shallow_research_agent + llm: nemotron_super_llm + # tools: omitted -> inherits all from data_source_registry + exclude_tools: + - advanced_web_search_tool + max_llm_turns: 10 + max_tool_iterations: 5 + + deep_research_agent: + _type: deep_research_agent + enable_citation_verification: true + orchestrator_llm: nemotron_super_llm + source_router_llm: nemotron_super_llm + researcher_llm: nemotron_super_llm + planner_llm: nemotron_super_llm + writer_llm: nemotron_super_llm + # tools: omitted -> inherits all from data_source_registry + exclude_tools: + - web_search_tool + +# ========================================================================= +# Protected per-user MCP source wiring (example: Google Drive) — OPTIONAL +# ========================================================================= +# Uncomment together with the `gdrive` data_sources entry above to enable a +# per-user OAuth-protected MCP source. +# +# How it fits together: +# - `mcp_oauth2_gdrive` (authentication): standard NAT MCP OAuth provider. AIQ +# treats this as the single source of truth — it derives the shared token +# storage, redirect_uri, scopes and client_id from here, and reuses NAT's +# well-known discovery for the authorize/token endpoints. +# - `mcp_gdrive` (function_groups): the job-time per_user_mcp_client. It reads +# the same token AIQ wrote at connect time from the same object store. +# +# Requirements (see frontends/aiq_api/src/aiq_api/mcp_auth/factory.py): +# * token_storage_object_store MUST be a shared, persistent object store — the +# API process writes the token and a separate worker process reads it. +# * redirect_uri MUST point at AIQ's callback: /v1/auth/mcp/gdrive/callback +# and be reachable by the user's browser (i.e. AIQ's public URL). +# * client_id is OPTIONAL. Omitted -> NAT dynamic client registration (DCR): +# connect works, but silent token refresh won't survive the API->worker +# process boundary, so users Reconnect when a source shows "expired" (the UI +# supports this). Set a fixed (e.g. ECI public) client_id for silent refresh. +# +# Example below wires a per-user OAuth-protected MCP server (e.g. a Google Drive +# MCP service). Point MCP_GDRIVE_URL at your provider's MCP endpoint. +# Set: MCP_GDRIVE_URL=https://your-mcp-server.example.com/mcp +# AIQ_PUBLIC_URL= +# +# NOTE: deliberately NO `function_groups: mcp_gdrive` here. A per_user_mcp_client +# declared in config is built by NAT's per-user interactive (WebSocket) session +# builder, which fails for a user with no token and breaks interactive chat. The +# async-job worker instead builds the per-user MCP client IN CODE, per job, from +# the `mcp_oauth2_gdrive` auth provider below (its server_url) + the job owner's +# token. See aiq_api.mcp_auth.runtime_tools.open_per_user_mcp_tools. + +authentication: + mcp_oauth2_gdrive: + _type: mcp_oauth2 + server_url: ${MCP_GDRIVE_URL:-https://your-mcp-server.example.com/mcp} + redirect_uri: ${AIQ_PUBLIC_URL:-http://localhost:8000}/v1/auth/mcp/gdrive/callback + token_storage_object_store: mcp_token_store + # scopes are discovered from MaaS protected-resource metadata; pin only to override. + # client_id: ${MCP_GDRIVE_CLIENT_ID} # optional — see note above + # client_secret: ${MCP_GDRIVE_CLIENT_SECRET} + +# A shared object store is required so the API process (connect) and the worker +# process (job) see the same token. Two modes, selected by MCP_TOKEN_STORE_TYPE: +# * aiq_sqlite (default) — a SQLite file, serviceless and shared across the API +# and worker processes on the same host. Works locally with nothing deployed. +# * redis — a networked store for multi-replica / multi-host deployments +# (requires nvidia-nat-redis; deploy sets MCP_TOKEN_STORE_TYPE=redis). +# Fields not used by the active _type are ignored, so both sets can coexist here. +# +# aiq_sqlite requirement: the API and worker processes must resolve db_path to the +# SAME file, so they must share a filesystem (single host / shared volume). Prefer +# an ABSOLUTE path — a relative path only works if both processes share a working +# directory. Separate hosts/pods do NOT share a filesystem; use redis there. +object_stores: + mcp_token_store: + _type: ${MCP_TOKEN_STORE_TYPE:-aiq_sqlite} + bucket_name: mcp-tokens + # aiq_sqlite (use an absolute path in any multi-process deployment): + db_path: ${MCP_TOKEN_DB:-./mcp_tokens.db} + # redis (S3/MySQL also available via their NAT packages): + host: ${REDIS_HOST:-localhost} + port: ${REDIS_PORT:-6379} + password: ${REDIS_PASSWORD} + # ttl: 2592000 # optional, seconds (both modes) + +workflow: + _type: chat_deepresearcher_agent + enable_escalation: true + enable_clarifier: true + use_async_deep_research: true + checkpoint_db: ${AIQ_CHECKPOINT_DB:-./checkpoints.db} diff --git a/deploy/.env.example b/deploy/.env.example index 03b5d7130..472f0f49f 100644 --- a/deploy/.env.example +++ b/deploy/.env.example @@ -47,6 +47,47 @@ DASK_DISTRIBUTED__LOGGING__DISTRIBUTED=warning # ----------------------------------------------------------------------------- # AIQ_SUMMARY_DB= +# ----------------------------------------------------------------------------- +# Per-user MCP auth token store (config_web_frag_mcp_auth.yml) +# Shared, persistent store the API process (connect) and the Dask worker (job) +# both reach to read/write per-user MCP tokens. +# +# MCP_TOKEN_STORE_TYPE selects the backend: +# * aiq_sqlite (config default) — a SQLite file, no service required. Good for +# local/single-host runs; only spans processes that share the filesystem. +# Override the path with MCP_TOKEN_DB (default ./mcp_tokens.db). +# * redis — networked, safe across replicas/hosts. Compose and Helm set this. +# +# Compose defaults to the bundled `redis` service, which is unauthenticated and +# reachable only on the internal Docker network. For production, point these at a +# managed Redis and set REDIS_PASSWORD — which also requires uncommenting +# `password: ${REDIS_PASSWORD}` in config_web_frag_mcp_auth.yml. +# ----------------------------------------------------------------------------- +# MCP_TOKEN_STORE_TYPE=redis +# MCP_TOKEN_DB=./mcp_tokens.db +# REDIS_HOST=redis +# REDIS_PORT=6379 +# REDIS_PASSWORD= + +# ----------------------------------------------------------------------------- +# Per-user MCP source connection (config_web_frag_mcp_auth.yml gdrive source) +# Required to actually CONNECT a protected source — separate from the token store +# above. On connect, AIQ discovers the OAuth endpoints by reaching the MCP server, +# so MCP_GDRIVE_URL must be set to a reachable server (the placeholder default is +# not). Without these the gdrive card shows but connect fails with +# "Source is not configured for MCP OAuth". Compose passes these in via env_file. +# +# MCP_GDRIVE_URL - the MCP server endpoint (must be reachable from the backend) +# AIQ_PUBLIC_URL - public base URL of AIQ's API; used for the OAuth redirect_uri +# (/v1/auth/mcp/gdrive/callback) +# Optionally pin a pre-registered OAuth client instead of dynamic registration: +# MCP_GDRIVE_CLIENT_ID / MCP_GDRIVE_CLIENT_SECRET +# ----------------------------------------------------------------------------- +# MCP_GDRIVE_URL=https://your-mcp-server.example.com/mcp +# AIQ_PUBLIC_URL=http://localhost:8000 +# MCP_GDRIVE_CLIENT_ID= +# MCP_GDRIVE_CLIENT_SECRET= + # ----------------------------------------------------------------------------- # Environment Variables for Modal Sandbox for Skills Execution (optional) # ----------------------------------------------------------------------------- diff --git a/deploy/Dockerfile b/deploy/Dockerfile index 0c5eb3916..3c13f928c 100644 --- a/deploy/Dockerfile +++ b/deploy/Dockerfile @@ -69,8 +69,14 @@ COPY configs/ ./configs/ # to avoid leaking .env, Helm charts, compose files, or other dev artifacts. COPY deploy/entrypoint.py deploy/start_web.py ./deploy/ -# Install dependencies using uv sync -RUN uv sync --frozen --no-dev --no-install-workspace +# Install dependencies using uv sync. We do NOT pass --no-dev: the workspace +# source packages (tavily-web-search, exa-web-search, ...) and their runtime +# deps (e.g. langchain-tavily) are declared in the [dependency-groups] dev group, +# so --no-dev would omit them and the editable source installs below would +# import-fail at runtime. Syncing with the dev group installs those deps from the +# frozen lock (still pinned, no re-resolve). TODO(deploy): move the source +# packages to a dedicated non-dev group to keep test-only tooling out of the image. +RUN uv sync --frozen # Install workspace packages (without CLI for base) RUN uv pip install --no-deps -e . \ diff --git a/deploy/compose/README.md b/deploy/compose/README.md index 5d8704161..68a96293c 100644 --- a/deploy/compose/README.md +++ b/deploy/compose/README.md @@ -120,6 +120,29 @@ Services started: - `aiq-blueprint-ui` (port 3000) - `postgres` (port 5432) +### Per-user MCP authentication + +Per-user MCP authentication is optional. To try it locally, set the MCP source and OAuth variables documented in `deploy/.env.example`, then apply the dedicated override with the default stack: + +```bash +cd deploy/compose +docker compose --env-file ../.env \ + -f docker-compose.yaml \ + -f docker-compose.per-user-auth.yaml \ + up -d --build +``` + +The override selects `config_web_frag_mcp_auth.yml` and starts a persistent Redis service on the private Compose network. Redis is not published on a host port and has no password, so this stack is for local development only. Production deployments should use a managed Redis service and set `BACKEND_CONFIG`, `MCP_TOKEN_STORE_TYPE=redis`, `REDIS_HOST`, `REDIS_PORT`, and, when required, `REDIS_PASSWORD` in `deploy/.env` without applying the local override. + +Stop the local per-user-auth stack with the same file set: + +```bash +docker compose --env-file ../.env \ + -f docker-compose.yaml \ + -f docker-compose.per-user-auth.yaml \ + down +``` + ## Foundational RAG (FRAG) prerequisites If you switch the backend to `configs/config_web_frag.yml`, you must run a compatible RAG server and ingest server separately and set: diff --git a/deploy/compose/docker-compose.per-user-auth.yaml b/deploy/compose/docker-compose.per-user-auth.yaml new file mode 100644 index 000000000..5e9a60518 --- /dev/null +++ b/deploy/compose/docker-compose.per-user-auth.yaml @@ -0,0 +1,31 @@ +# Optional local stack for configs/config_web_frag_mcp_auth.yml. +# Apply this file together with docker-compose.yaml; it is not standalone. + +services: + aiq-agent: + environment: + CONFIG_FILE: /app/configs/config_web_frag_mcp_auth.yml + MCP_TOKEN_STORE_TYPE: redis + REDIS_HOST: redis + REDIS_PORT: "6379" + depends_on: + redis: + condition: service_healthy + + redis: + image: redis:7-alpine + command: ["redis-server", "--appendonly", "yes"] + volumes: + - redis-data:/data + networks: + - aiq-network + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 5s + timeout: 5s + retries: 5 + restart: unless-stopped + +volumes: + redis-data: + driver: local diff --git a/deploy/compose/docker-compose.yaml b/deploy/compose/docker-compose.yaml index 81fd0c2eb..77335d041 100644 --- a/deploy/compose/docker-compose.yaml +++ b/deploy/compose/docker-compose.yaml @@ -140,7 +140,6 @@ volumes: driver: local postgres-data: driver: local - networks: aiq-network: driver: bridge diff --git a/deploy/helm/README.md b/deploy/helm/README.md index adf442391..9d5004781 100644 --- a/deploy/helm/README.md +++ b/deploy/helm/README.md @@ -196,6 +196,7 @@ The backend loads a workflow config at startup. Switch configs with `--set`: |-------------|-------------| | `configs/config_web_default_llamaindex.yml` | Default — LlamaIndex backend (no external RAG required) | | `configs/config_web_frag.yml` | Foundational RAG mode (requires a running RAG service) | +| `configs/config_web_frag_mcp_auth.yml` | Foundational RAG with optional per-user MCP authentication | ```bash helm upgrade --install aiq aiq2-web-2.0.0.tgz -n ns-aiq \ @@ -206,6 +207,36 @@ helm upgrade --install aiq aiq2-web-2.0.0.tgz -n ns-aiq \ --set aiq.apps.backend.env.CONFIG_FILE=configs/config_web_frag.yml ``` +### Per-user MCP authentication with external Redis + +The chart does not install Redis. When selecting the per-user authentication config, provide a Redis service that the backend and its workers can both reach. The deployer owns its availability, persistence, networking, and backup. + +Add `REDIS_PASSWORD` to the existing `aiq-credentials` Secret when the Redis service requires authentication, then create a values file: + +```yaml +# aiq-per-user-auth-values.yaml +aiq: + apps: + backend: + env: + CONFIG_FILE: configs/config_web_frag_mcp_auth.yml + MCP_TOKEN_STORE_TYPE: redis + REDIS_HOST: redis.example.com + REDIS_PORT: "6379" + secretEnv: + REDIS_PASSWORD: REDIS_PASSWORD +``` + +Apply it to either the downloaded chart or a source-chart installation: + +```bash +helm upgrade --install aiq aiq2-web-2.0.0.tgz -n ns-aiq \ + --wait --timeout 10m \ + -f aiq-per-user-auth-values.yaml +``` + +Configure `MCP_GDRIVE_URL`, `AIQ_PUBLIC_URL`, and any OAuth client credentials required by the protected MCP source through the same `env` and `secretEnv` maps. NAT 1.8's Redis object store in this image supports host, port, database, and optional password; this example does not support TLS, ACL usernames, Sentinel, or Redis Cluster. + ## FRAG Integration To use the Foundational RAG (FRAG) config, you need a running NVIDIA RAG Blueprint deployment. See the [RAG Blueprint Helm deployment guide](https://github.com/NVIDIA-AI-Blueprints/rag/blob/develop/docs/deploy-helm.md) for setup instructions. diff --git a/docs/source/customization/mcp-tools.md b/docs/source/customization/mcp-tools.md index 1c4d60c2d..2dfe2b520 100644 --- a/docs/source/customization/mcp-tools.md +++ b/docs/source/customization/mcp-tools.md @@ -8,22 +8,20 @@ Model Context Protocol (MCP) is an open protocol that standardizes how applicati context to LLM applications. The AIQ Blueprint is built on the NVIDIA NeMo Agent toolkit (NAT), so AIQ can use MCP servers as data sources through NAT function groups. -This guide is written for AIQ 2.1 deployments running NAT `1.6.0` or later. Verify your installed -version with `uv pip show nvidia-nat`. +This guide targets AIQ deployments pinned to NAT `1.8.0`. Verify your installed version with +`uv pip show nvidia-nat`. ## What this guide covers -**Supported in AIQ 2.1:** +**Supported:** - Connect AIQ to an unauthenticated MCP server. - Connect AIQ to an MCP server with backend service-account credentials. +- Connect each signed-in user to a protected MCP server through the AIQ UI using MCP OAuth. - Forward the signed-in AIQ user's identity to a downstream service from a custom AIQ tool. -**Planned for AIQ 2.2 / 2.3:** +**Not yet first-party:** -- Native per-user MCP OAuth driven by the AIQ UI. NAT 1.6 ships the protocol-level support - (`mcp_oauth2`, `per_user_mcp_client`); the AIQ UI cannot yet drive per-MCP consent. See the - short [planning note](#per-user-mcp-oauth-planned) below. - A first-party AIQ-token pass-through MCP auth provider. Today, if your MCP server trusts the AIQ user's bearer token, you must implement and register a custom NAT auth provider in your deployment package. @@ -41,14 +39,14 @@ For the full NAT MCP reference: | MCP server has no per-user auth | `mcp_client` function group | [Connect AIQ to an MCP Server](#connect-aiq-to-an-mcp-server) | | MCP server uses backend / app credentials | `mcp_client` + `mcp_service_account` | [Service-Account MCP Servers](#service-account-mcp-servers) | | Downstream API trusts the AIQ user's bearer token | Custom AIQ tool using `get_auth_token()` | [Forwarding AIQ User Identity](#forwarding-aiq-user-identity-from-a-tool) | -| MCP server requires per-user OAuth consent | Planned for AIQ 2.2 / 2.3 | [Per-User MCP OAuth (planned)](#per-user-mcp-oauth-planned) | +| MCP server requires per-user OAuth consent | `per_user_mcp_client` + `mcp_oauth2` | [Per-User MCP OAuth](#per-user-mcp-oauth) | ## Prerequisites Install NAT and the MCP package on the same release line as your AIQ deployment: ```bash -uv pip install "nvidia-nat[mcp]==1.6.0" nvidia-nat-mcp==1.6.0 +uv pip install "nvidia-nat[mcp]==1.8.0" nvidia-nat-mcp==1.8.0 nvidia-nat-redis==1.8.0 ``` Keep `nvidia-nat`, `nvidia-nat-core`, `nvidia-nat-eval`, and `nvidia-nat-mcp` on the same release @@ -276,34 +274,59 @@ functions: endpoint: ${INTERNAL_SEARCH_URL} ``` -This is the AIQ-user-identity MCP pattern fully supported in 2.1. The two alternatives — +This is the supported AIQ-user-identity MCP pattern. The two alternatives — protocol-level pass-through via a custom NAT auth provider, and an auth-forwarding MCP proxy — are -viable in NAT but are not first-class in AIQ 2.1; treat them as deployment-side extensions. +viable in NAT but are not first-class in AIQ; treat them as deployment-side extensions. For the broader auth context (UI sign-in flow, validator registration, headless API callers), see [Authentication](../deployment/authentication.md). -## Per-User MCP OAuth (planned) - -NAT 1.6 ships the protocol-level building blocks for per-user MCP OAuth — `mcp_oauth2` (auth -provider for MCP OAuth flows) and `per_user_mcp_client` (function group with per-user token -storage). What AIQ 2.1 **does not yet** ship is the UI integration that drives this flow: the -data-source API does not return per-MCP auth status, connect / disconnect URLs, scopes, or token -expiry, and the UI has no per-source "Connect" / "Reconnect" controls. - -Until that lands, the recommended patterns for AIQ deployments remain: - -- **Service-account MCP** when the access can be shared at the application level. -- **AIQ user-identity tools** (the section above) when the downstream service trusts the AIQ - bearer token. - -Beyond the UI, two further gaps exist in 2.1: AIQ's `/v1/data_sources` does not surface per-MCP -auth status (connect URL, scopes, token expiry, error state), and async deep research jobs cannot -yet resolve per-user MCP tokens inside Dask workers. The full per-user MCP OAuth integration — -backend status APIs, UI controls, and worker-side token resolution — is tracked on the AIQ 2.2 / -2.3 roadmap. Refer to the +## Per-User MCP OAuth + +Use per-user MCP OAuth when each AIQ user must authorize the upstream MCP server with their own +identity. The reference configuration is +[`configs/config_web_frag_mcp_auth.yml`](../../../configs/config_web_frag_mcp_auth.yml). It combines +an OAuth-protected data source, NAT's `per_user_mcp_client`, an `mcp_oauth2` provider, and a shared +token object store. + +Set these values before starting AIQ: + +- `MCP_GDRIVE_URL`: protected streamable-HTTP MCP endpoint. The example calls the source `gdrive`, + but the mechanism is not Google Drive-specific. +- `AIQ_PUBLIC_URL`: externally reachable AIQ origin used to construct the OAuth callback URL. +- `MCP_GDRIVE_CLIENT_ID` and `MCP_GDRIVE_CLIENT_SECRET`: only when the MCP authorization server + requires a pre-registered OAuth client rather than dynamic client registration. +- `MCP_TOKEN_STORE_TYPE`: `aiq_sqlite` for a single-host example or `redis` for multi-process and + multi-host deployments. + +The UI reads connection status from `/v1/data_sources` and presents Connect, Reconnect, and +Disconnect actions for the protected source. AIQ owns the OAuth callback and stores the resulting +token under the current AIQ user identity. Submitting a job with a disconnected protected source +fails with `409 mcp_auth_required`; AIQ does not silently run the job without that source. + +Both interactive WebSocket sessions and REST-submitted async jobs resolve the user's MCP tools. +Async workers open their own MCP client for the job and read the same token store as the API, so +the API and workers must share that store: + +- `aiq_sqlite` requires the API and worker processes to share the same absolute database path on + one host. +- `redis` is the supported example for multiple processes, hosts, or Kubernetes pods. Each + protected source must reference an object store with a distinct bucket or namespace; + configuration fails closed when two sources reference the same object-store configuration. + +This example targets the AIQ web/API deployment, which owns the connect and callback routes and +supplies user identity to jobs. Raw NAT CLI runs do not provide that browser OAuth lifecycle; use +an unauthenticated or service-account MCP configuration for standalone CLI execution. + +For a local Redis-backed stack, use the +[per-user-auth Compose override](https://github.com/NVIDIA-AI-Blueprints/aiq/blob/develop/deploy/compose/README.md#per-user-mcp-authentication). +For a released chart, provide an external Redis service as described in the +[Helm deployment guide](https://github.com/NVIDIA-AI-Blueprints/aiq/blob/develop/deploy/helm/README.md#per-user-mcp-authentication-with-external-redis). +The default Compose and Helm deployments remain Redis-free when this example is not selected. + +See the [NAT MCP authentication guide](https://docs.nvidia.com/nemo/agent-toolkit/latest/components/auth/mcp-auth/index.html) -if you want to follow NAT's MCP OAuth surface directly. +for protocol details. ## Security Guidance @@ -312,9 +335,11 @@ if you want to follow NAT's MCP OAuth surface directly. - Store secrets in environment variables or a secret manager, not in YAML checked into source control. - Use service-account MCP auth only when shared app-level access is acceptable. +- Use `per_user_auth: true` for upstream MCP OAuth; `requires_auth: true` only gates a source on + AIQ sign-in and does not authorize the upstream MCP server. - Keep token forwarding scoped to trusted internal services and HTTPS endpoints. -- Mark user-authenticated data sources with `requires_auth: true` so the UI can prevent - unauthenticated use. +- Use `requires_auth: true` for sources that depend on AIQ sign-in but do not have a separate + upstream OAuth connection. ## Troubleshooting diff --git a/frontends/aiq_api/src/aiq_api/auth/middleware.py b/frontends/aiq_api/src/aiq_api/auth/middleware.py index cf87821ac..f635f4e2d 100644 --- a/frontends/aiq_api/src/aiq_api/auth/middleware.py +++ b/frontends/aiq_api/src/aiq_api/auth/middleware.py @@ -193,12 +193,24 @@ def build_request_trace_tags( "/v1/jobs/async/agents", "/v1/jobs/async/submit", "/v1/jobs/async/job/", # prefix — matches /v1/jobs/async/job/{id}/* + "/v1/auth/mcp/", # prefix — per-user MCP auth: {id}/status, {id}/connect, {id}/callback ] # External paths that require no token (monitoring, etc.) AUTH_EXEMPT_PATHS: set[str] = {"/health", "/docs", "/redoc", "/openapi.json"} +def _is_oauth_callback_path(path: str) -> bool: + """The MCP OAuth redirect callback (``/v1/auth/mcp/{source_id}/callback``). + + It must be auth-exempt: the provider redirects the user's browser here with no + AIQ token. It is secured by the unguessable OAuth ``state`` (bound to the + principal + source when the flow was started via /connect), not by a request + token. ``/status`` and ``/connect`` are NOT exempt — they need the principal. + """ + return path.startswith("/v1/auth/mcp/") and path.endswith("/callback") + + def _load_external_hostnames() -> set[str]: """Read ``AIQ_EXTERNAL_HOSTNAMES`` env var; fall back to staging hostname.""" env = os.getenv("AIQ_EXTERNAL_HOSTNAMES", "") @@ -388,7 +400,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: await self._send_json(send, 404, {"detail": "Not found"}) return - if is_external and path in AUTH_EXEMPT_PATHS: + if is_external and (path in AUTH_EXEMPT_PATHS or _is_oauth_callback_path(path)): user = {"type": "anonymous", "skip_clarifier": True} await self._call_app(scope, receive, send, headers, user) return diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index dbdb531cc..40c00eb3d 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -449,6 +449,7 @@ async def run_agent_job( auth_token: str | None = None, initial_files: dict[str, Any] | None = None, output_metadata: dict[str, Any] | None = None, + owner_user_id: str | None = None, ): """ Dask task to run any registered agent with cancellation support and telemetry. @@ -482,6 +483,9 @@ async def run_agent_job( data sources that require authentication (requires_auth: true). initial_files: Optional DeepAgents virtual filesystem files to seed into state. output_metadata: Optional metadata to persist alongside the final report. + owner_user_id: Canonical per-user key (``principal_user_id``), set on the NAT + Context so per_user_mcp_client retrieves the token the owner connected + via /v1/auth/mcp/{id}/connect. """ # Propagate auth token into the current async task's context so tools @@ -559,6 +563,14 @@ async def run_agent_job( provider, llm = await _create_llm_provider(builder, fn_config) + # Bind the job owner's identity on the NAT context before tools are built, + # so per_user_mcp_client resolves the token this user connected via + # /v1/auth/mcp/{id}/connect (keyed by principal_user_id). + if owner_user_id: + from nat.builder.context import ContextState + + ContextState.get().user_id.set(owner_user_id) + # Resolve tools: use explicit list or auto-inherit from data_source_registry tool_refs = fn_config.tools if not tool_refs: @@ -680,41 +692,57 @@ async def run_agent_job( callbacks.append(AgentEventCallback(event_store)) callbacks.append(nat_profiler_callback) - # Instantiate agent with callbacks - agent = _create_agent_instance( - agent_cls=agent_cls, - llm_provider=provider, - llm=llm, - tools=tools, - fn_config=fn_config, - verbose=verbose, - callbacks=callbacks, - job_id=job_id, - # Artifact harvesting rides 284's job store + event stream: the same db_url - # backs the SqlArtifactStore, and event_store.store carries artifact SSE - # events. Inert unless sandbox.artifact_capture is enabled in config. - artifact_db_url=db_url, - artifact_emit=event_store.store, - ) + # Resolve per-user MCP source tools for the job owner (Context.user_id + # set above); connections stay open via mcp_stack for the agent run. + # Best-effort: the helper never raises, so this can't break a job. + from contextlib import AsyncExitStack - # Capture the runtime so the terminal path can release the sandbox. None for - # agents without a sandbox runtime; close()/terminate() are then no-ops. - sandbox_runtime = getattr(agent, "deepagents_runtime", None) - - # Run agent - LLM/tool events will be nested under workflow span - result = await _run_agent( - agent=agent, - input_text=input_text, - builder=builder, - config=config, - function_name=agent_config_name, - function_config=fn_config, - monitor=cancellation_monitor, - available_documents=available_documents, - data_sources=data_sources, - event_store=event_store, - initial_files=initial_files, - ) + from ..mcp_auth.runtime_tools import open_per_user_mcp_tools + + async with AsyncExitStack() as mcp_stack: + mcp_tools = await open_per_user_mcp_tools( + builder=builder, + data_sources=data_sources, + exit_stack=mcp_stack, + wrapper_type=LLMFrameworkEnum.LANGCHAIN, + ) + agent_tools = [*tools, *mcp_tools] if mcp_tools else tools + + # Instantiate agent with callbacks + agent = _create_agent_instance( + agent_cls=agent_cls, + llm_provider=provider, + llm=llm, + tools=agent_tools, + fn_config=fn_config, + verbose=verbose, + callbacks=callbacks, + job_id=job_id, + # Artifact harvesting rides 284's job store + event stream: the same db_url + # backs the SqlArtifactStore, and event_store.store carries artifact SSE + # events. Inert unless sandbox.artifact_capture is enabled in config. + artifact_db_url=db_url, + artifact_emit=event_store.store, + ) + + # Capture the runtime so the terminal path can release the sandbox. None for + # agents without a sandbox runtime; close()/terminate() are then no-ops. + sandbox_runtime = getattr(agent, "deepagents_runtime", None) + + # Run agent - LLM/tool events will be nested under workflow span + result = await _run_agent( + agent=agent, + input_text=input_text, + builder=builder, + config=config, + function_name=agent_config_name, + function_config=fn_config, + monitor=cancellation_monitor, + available_documents=available_documents, + data_sources=data_sources, + event_store=event_store, + initial_files=initial_files, + ) # Emit WORKFLOW_END event for Phoenix context.intermediate_step_manager.push_intermediate_step( diff --git a/frontends/aiq_api/src/aiq_api/jobs/submit.py b/frontends/aiq_api/src/aiq_api/jobs/submit.py index d79e2597f..7f93e5bda 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/submit.py +++ b/frontends/aiq_api/src/aiq_api/jobs/submit.py @@ -31,6 +31,7 @@ from aiq_agent.auth import Principal from aiq_agent.auth import get_current_principal from aiq_api.auth import get_current_trace_tags +from aiq_api.mcp_auth.provider import principal_user_id from ..registry import get_agent_config from .access import _make_no_auth_principal @@ -252,6 +253,21 @@ async def submit_agent_job( if principal is None: raise RuntimeError("Verified current principal required for async job submission") + # Preflight protected MCP sources before enqueue. The REST submit route also + # does this (returning 409), but programmatic submitters — notably the chat + # researcher's async deep-research path — call this directly and would otherwise + # bypass the check, so this is the single chokepoint both paths share. Skipped + # when no MCP auth provider is active in this process (nothing to enforce). + from aiq_api.mcp_auth.active import get_active_mcp_auth_provider + from aiq_api.mcp_auth.preflight import McpAuthRequiredError + from aiq_api.mcp_auth.preflight import evaluate_mcp_auth + + mcp_provider = get_active_mcp_auth_provider() + if mcp_provider is not None: + block = await evaluate_mcp_auth(mcp_provider, principal, data_sources) + if block is not None: + raise McpAuthRequiredError(block) + job_store = JobStore(scheduler_address=scheduler_address, db_url=db_url) resolved_job_id = job_store.ensure_job_id(job_id) loop = asyncio.get_running_loop() @@ -293,6 +309,7 @@ async def _rollback_partial_submission() -> None: auth_token, initial_files, output_metadata, + principal_user_id(principal), ], ) except IntegrityError as e: diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/__init__.py b/frontends/aiq_api/src/aiq_api/mcp_auth/__init__.py new file mode 100644 index 000000000..998cb76e3 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/__init__.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-user MCP auth control plane for AIQ. + +AIQ owns the product control plane (status surfacing, connect actions, +structured auth-required responses, submit-time preflight). The actual OAuth +mechanics and per-user token storage are NAT's. AIQ mints the provider +authorization URL and completes the callback using NAT's public OAuth/token +primitives, then writes the resulting token into the *same* NAT token storage +that the headless job-time ``per_user_mcp_client`` reads from. See +``aiq-nat-oauth-execution-gap`` design notes. +""" + +from .models import McpAuthRequiredResponse +from .models import McpAuthRequiredSource +from .models import PerUserAuthInfo +from .models import SourceAuthStatusResponse +from .models import SourceConnectResponse +from .provider import AuthStatus +from .provider import ProtectedSourceAuthProvider +from .provider import SourceAuthChallenge +from .provider import SourceAuthState +from .provider import principal_user_id + +__all__ = [ + "AuthStatus", + "McpAuthRequiredResponse", + "McpAuthRequiredSource", + "PerUserAuthInfo", + "ProtectedSourceAuthProvider", + "SourceAuthChallenge", + "SourceAuthState", + "SourceAuthStatusResponse", + "SourceConnectResponse", + "principal_user_id", +] diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/active.py b/frontends/aiq_api/src/aiq_api/mcp_auth/active.py new file mode 100644 index 000000000..f25c46b96 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/active.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Process-global handle to the per-user MCP auth provider. + +The provider is built once at route registration in the API process. Both the +REST submit route and programmatic submitters (e.g. the chat researcher's async +deep-research submit) run in that same process, so they can reach the live +provider through this module to run the connect-state preflight before enqueue. + +It is intentionally a simple module-global: there is one provider per API +process and no per-request state. Returns ``None`` when MCP auth was never +registered (e.g. unit tests that call ``submit_agent_job`` directly), in which +case callers skip the preflight — there is nothing to enforce. +""" + +from __future__ import annotations + +from .provider import ProtectedSourceAuthProvider + +_active_provider: ProtectedSourceAuthProvider | None = None + + +def set_active_mcp_auth_provider(provider: ProtectedSourceAuthProvider | None) -> None: + """Register the provider built at route setup as the process-wide instance.""" + global _active_provider + _active_provider = provider + + +def get_active_mcp_auth_provider() -> ProtectedSourceAuthProvider | None: + """Return the registered provider, or ``None`` if MCP auth is not configured.""" + return _active_provider diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/factory.py b/frontends/aiq_api/src/aiq_api/mcp_auth/factory.py new file mode 100644 index 000000000..274095221 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/factory.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build a :class:`NatMcpAuthProvider` from the registry + NAT's mcp_oauth2 config. + +The NAT ``authentication: mcp_oauth2`` provider is the single source of truth. +A protected source's registry entry points at it via ``per_user_auth.auth_provider`` +(falling back to ``mcp_server_id``); AIQ resolves that provider through the +builder and derives: + + * the shared token storage (from ``token_storage_object_store``) — so the token + AIQ writes at connect time is the same one the job-time ``per_user_mcp_client`` + reads; + * ``redirect_uri`` / ``scopes`` / ``client_id`` / ``use_pkce``; + * the authorize/token endpoints, via NAT's own discovery (RFC 8414 / RFC 9728 + well-known), reusing NAT rather than a parallel config surface. + +Deployment notes for the cross-process design: + 1. ``token_storage_object_store`` must name a shared, persistent object store + (the API process writes the token; a separate worker process reads it). + This is REQUIRED. + 2. ``client_id`` (manual registration) is OPTIONAL. Without it, NAT performs + dynamic client registration (DCR) — the connect flow works, but the refresh + token is bound to the connect-process's registered client, so the worker + cannot silently refresh it. That degrades gracefully to the "expired -> + Reconnect" UX. Set a fixed (e.g. ECI public) ``client_id`` only if you want + silent cross-process refresh. + +Endpoint discovery runs once at startup per source and is guarded: if it fails +(server unreachable, no well-known), the source is left unconfigured and its +status surfaces as ``error`` rather than crashing route registration. +""" + +from __future__ import annotations + +import logging + +import httpx + +from aiq_agent.common.data_source_registry import get_all_sources +from nat.authentication.token_storage import InMemoryTokenStorage +from nat.authentication.token_storage import ObjectStoreTokenStorage +from nat.authentication.token_storage import TokenStorageBase + +from .nat_provider import NatMcpAuthProvider +from .nat_provider import OAuthSourceSettings + +logger = logging.getLogger(__name__) + + +async def _resolve_token_storage(builder, cfg, source_id: str) -> TokenStorageBase: + """Resolve the token storage for a source from the mcp_oauth2 config. + + Uses the configured object store (shared across processes) when present; + otherwise falls back to a process-local in-memory store with a loud warning + (dev only — tokens will not be visible to job workers). + """ + object_store_name = getattr(cfg, "token_storage_object_store", None) + if object_store_name: + object_store = await builder.get_object_store_client(object_store_name) + return ObjectStoreTokenStorage(object_store) + logger.warning( + "Source '%s': mcp_oauth2 provider has no token_storage_object_store; using a process-local " + "in-memory token store. Tokens connected via the API will NOT be visible to job workers. " + "Set token_storage_object_store to a shared object store for real deployments.", + source_id, + ) + return InMemoryTokenStorage() + + +async def _probe_for_oauth_challenge(server_url: str) -> httpx.Response | None: + """Send an unauthenticated MCP request to elicit the 401 + WWW-Authenticate. + + MCP servers point at their authorization server via the 401's RFC 9728 + ``resource_metadata`` hint. Returns the 401 response (for NAT discovery) or + None if the server didn't challenge (NAT then falls back to well-known). + """ + if not server_url: + return None + try: + async with httpx.AsyncClient(timeout=15) as client: + response = await client.post( + server_url, + json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}, + headers={"Accept": "application/json, text/event-stream"}, + ) + if response.status_code == 401: + return response + logger.debug("MCP probe of %s returned %s (no 401 challenge)", server_url, response.status_code) + except Exception as exc: + logger.warning("MCP auth probe failed for %s: %s", server_url, exc) + return None + + +async def _resolve_oauth_settings(source_id: str, pua, nat_provider, cfg) -> OAuthSourceSettings | None: + """Derive AIQ-side OAuth settings from the NAT provider, discovering endpoints. + + Returns None (source left unconfigured -> status 'error') if endpoints or a + client_id cannot be resolved. + + NAT private-API surface (review on every nvidia-nat upgrade; pinned to 1.8.0). + All accesses below are defensive (``getattr`` defaults + the surrounding + ``try/except``), so a renamed/removed member degrades to status 'error' with a + warning rather than crashing: + + * ``nat_provider._discover_and_register(response=...)`` -- coroutine; runs + well-known discovery + RFC 7591 dynamic client registration, populating the + cached members below. + * ``nat_provider._cached_endpoints`` -- object | None; ``.authorization_url`` + and ``.token_url`` (str-able). + * ``nat_provider._cached_credentials`` -- object | None; ``.client_id`` and + ``.client_secret`` (str). + * ``nat_provider._effective_scopes`` -- Iterable[str] | None; scopes resolved + from protected-resource metadata. + * ``nat_provider._discoverer`` -- object | None; ``._resource_from_metadata`` + (str | None) is the RFC 9728 resource identifier. + """ + redirect_uri = str(cfg.redirect_uri) if getattr(cfg, "redirect_uri", None) else "" + scopes = list(getattr(cfg, "scopes", None) or []) + client_id = getattr(cfg, "client_id", None) + client_secret = getattr(cfg, "client_secret", None) + server_url = str(getattr(cfg, "server_url", "") or "") + # OAuth resource indicator (RFC 8707/9728); default to the server URL, matching + # NAT's `resource = _resource_from_metadata or server_url`. + resource = server_url or None + + authorization_url = token_url = None + # Reuse NAT's discovery (well-known metadata + RFC 7591 DCR when no client_id + # is set). This module is the explicit NAT integration seam, so reaching + # discovery internals is acceptable (version-pinned to the resolved nat release). + # MCP servers (e.g. NVIDIA MaaS) advertise their AS via the 401 WWW-Authenticate + # header (RFC 9728 resource_metadata), not the root well-known, so we probe for + # that 401 first; without it unauthenticated discovery cannot locate the AS. + try: + challenge = await _probe_for_oauth_challenge(str(getattr(cfg, "server_url", "") or "")) + await nat_provider._discover_and_register(response=challenge) # noqa: SLF001 — NAT seam + endpoints = getattr(nat_provider, "_cached_endpoints", None) + credentials = getattr(nat_provider, "_cached_credentials", None) + if endpoints is not None: + authorization_url = str(endpoints.authorization_url) + token_url = str(endpoints.token_url) + if credentials is not None and not client_id: + client_id = credentials.client_id + client_secret = credentials.client_secret + # Prefer scopes resolved by discovery (protected-resource metadata) when + # the config didn't pin any. + discovered_scopes = getattr(nat_provider, "_effective_scopes", None) + if discovered_scopes: + scopes = list(discovered_scopes) + # RFC 9728 resource identifier from protected-resource metadata, mirroring + # NAT: `_discoverer._resource_from_metadata or server_url` (auth_provider.py). + discoverer = getattr(nat_provider, "_discoverer", None) + resource = getattr(discoverer, "_resource_from_metadata", None) or resource + except Exception as exc: + logger.warning("Source '%s': OAuth endpoint discovery failed: %s", source_id, exc) + + if not (authorization_url and token_url and client_id): + logger.warning( + "Source '%s': could not resolve OAuth endpoints/client_id from provider '%s'; " + "connect will be unavailable. Ensure the mcp_oauth2 server is reachable and client_id is set.", + source_id, + pua.auth_provider or pua.mcp_server_id or source_id, + ) + return None + + return OAuthSourceSettings( + source_id=source_id, + mcp_server_id=pua.mcp_server_id or source_id, + provider=pua.provider, + authorization_url=authorization_url, + token_url=token_url, + client_id=client_id, + client_secret=client_secret, + scopes=scopes, + redirect_uri=redirect_uri, + use_pkce=bool(getattr(cfg, "use_pkce", True)), + token_endpoint_auth_method=(getattr(cfg, "token_endpoint_auth_method", None) or "client_secret_post"), + resource=resource, + ) + + +async def build_mcp_auth_provider(builder) -> NatMcpAuthProvider: + """Construct the provider from registry metadata + NAT mcp_oauth2 config.""" + settings_by_source: dict[str, OAuthSourceSettings] = {} + storages: dict[str, TokenStorageBase] = {} + # Guard against credential cross-contamination: NAT's ObjectStoreTokenStorage + # keys tokens by user only (``tokens/{sha256(user_id)}``), so two protected + # sources sharing one token-storage object store would overwrite each other's + # credentials for the same user. AIQ can't fix this with a per-source key + # prefix, because NAT's job-time per_user_mcp_client reads the token itself via + # the source's mcp_oauth2 provider (unprefixed) — a prefix here would desync + # that read. So each protected source needs its OWN token-storage bucket; fail + # closed (skip the later source) when one is reused rather than silently + # clobbering tokens. + claimed_stores: dict[str, str] = {} # object_store name -> first source_id to claim it + + for source in get_all_sources(): + pua = source.per_user_auth + if pua is None or not pua.required: + continue + ref = pua.auth_provider or pua.mcp_server_id or source.id + try: + nat_provider = await builder.get_auth_provider(ref) + except Exception as exc: + logger.warning( + "Source '%s' declares per_user_auth but NAT auth provider '%s' is not configured: %s", + source.id, + ref, + exc, + ) + continue + + cfg = getattr(nat_provider, "config", None) + if cfg is None: + logger.warning("Source '%s': auth provider '%s' has no config; skipping", source.id, ref) + continue + + object_store_name = getattr(cfg, "token_storage_object_store", None) + if object_store_name: + prior = claimed_stores.get(object_store_name) + if prior is not None: + logger.error( + "Source '%s' shares token_storage_object_store '%s' with source '%s'. NAT keys tokens " + "per user only, so sharing one store lets these sources overwrite each other's " + "credentials. Give each protected source its own token-storage object store (distinct " + "bucket). Skipping '%s' — it will surface as unconfigured until this is resolved.", + source.id, + object_store_name, + prior, + source.id, + ) + continue + claimed_stores[object_store_name] = source.id + + try: + storage = await _resolve_token_storage(builder, cfg, source.id) + except Exception as exc: + logger.error("Source '%s': could not resolve token storage: %s", source.id, exc) + continue + + settings = await _resolve_oauth_settings(source.id, pua, nat_provider, cfg) + if settings is None: + continue + + settings_by_source[source.id] = settings + storages[source.id] = storage + + if settings_by_source: + logger.info("MCP auth configured for sources: %s", ", ".join(sorted(settings_by_source))) + return NatMcpAuthProvider( + settings_by_source=settings_by_source, + token_storage_resolver=lambda s: storages[s.source_id], + ) diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/models.py b/frontends/aiq_api/src/aiq_api/mcp_auth/models.py new file mode 100644 index 000000000..c31ed0e54 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/models.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""API response models for per-user MCP auth.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from pydantic import BaseModel +from pydantic import Field + +AuthStatusLiteral = Literal["connected", "not_connected", "expired", "error"] + + +class PerUserAuthInfo(BaseModel): + """Per-user MCP auth block attached to a data source in API responses. + + Combines the static declaration (``required``/``type``/``provider``/ + ``mcp_server_id`` from the registry) with the current user's dynamic state + (``status``/``expires_at``/``last_error``) and the action URLs. + + ``connect_url`` is a stable AIQ API surface. ``auth_url`` is a short-lived + provider/NAT login URL and is only populated when AIQ has intentionally + started or reused an auth challenge — never from a read-only listing. + """ + + required: bool = False + type: Literal["mcp_oauth2"] = "mcp_oauth2" + provider: str | None = None + mcp_server_id: str | None = None + status: AuthStatusLiteral | None = None + connect_url: str | None = None + auth_url: str | None = None + expires_at: datetime | None = None + last_error: str | None = None + + +class SourceAuthStatusResponse(BaseModel): + """Response for ``GET /v1/auth/mcp/{source_id}/status``.""" + + source_id: str + status: AuthStatusLiteral + expires_at: datetime | None = None + connect_url: str | None = None + last_error: str | None = None + + +class SourceConnectResponse(BaseModel): + """Response for ``POST /v1/auth/mcp/{source_id}/connect``.""" + + source_id: str + status: Literal["auth_required", "connected"] = "auth_required" + auth_url: str | None = None + expires_at: datetime | None = Field( + default=None, + description="Expiry of the auth challenge (auth_url), not of the eventual token", + ) + + +class McpAuthRequiredSource(BaseModel): + """A single blocked source in a 409 mcp_auth_required response.""" + + source_id: str + status: AuthStatusLiteral + connect_url: str + auth_url: str | None = None + + +class McpAuthRequiredResponse(BaseModel): + """Body of the 409 returned by submit preflight when sources need auth.""" + + error: Literal["mcp_auth_required"] = "mcp_auth_required" + message: str = "One or more selected data sources require connection before this job can start." + sources: list[McpAuthRequiredSource] = Field(default_factory=list) diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/nat_provider.py b/frontends/aiq_api/src/aiq_api/mcp_auth/nat_provider.py new file mode 100644 index 000000000..8a4447e46 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/nat_provider.py @@ -0,0 +1,304 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""NAT-backed implementation of :class:`ProtectedSourceAuthProvider`. + +This implementation never invokes NAT's blocking, browser-opening MCP auth +flow. Instead it mints the provider authorization URL with ``authlib`` (the +same library NAT uses), completes the ``code -> token`` exchange in AIQ's own +callback, and writes the resulting token into NAT's *public* token storage in +the exact shape NAT writes (``AuthResult`` with a ``BearerTokenCred`` and the +raw token dict incl. ``refresh_token``). The headless job-time +``per_user_mcp_client`` then finds that token via ``token_storage.retrieve`` +and never has to authenticate interactively. + +The per-source token storage MUST be a shared, persistent backend (an +``ObjectStore``) — the connect endpoint runs in the API process while the job +runs in a separate worker process, so in-memory storage would not be visible +across the boundary. +""" + +from __future__ import annotations + +import asyncio +import logging +import secrets +from collections.abc import Callable +from dataclasses import dataclass +from dataclasses import field +from datetime import UTC +from datetime import datetime +from datetime import timedelta + +import pkce +from authlib.integrations.httpx_client import AsyncOAuth2Client +from pydantic import BaseModel +from pydantic import Field +from pydantic import SecretStr + +from aiq_agent.auth import Principal +from nat.authentication.token_storage import TokenStorageBase +from nat.data_models.authentication import AuthResult +from nat.data_models.authentication import BearerTokenCred + +from .provider import SourceAuthChallenge +from .provider import SourceAuthState +from .provider import principal_user_id + +logger = logging.getLogger(__name__) + +_DEFAULT_CHALLENGE_TTL = timedelta(seconds=300) + + +class OAuthSourceSettings(BaseModel): + """OAuth2 settings for one protected MCP source. + + These mirror the fields NAT's ``mcp_oauth2`` provider uses. ``mcp_server_id`` + is the NAT auth-provider/server key — the token written here must be + retrievable by the same key the job-time provider uses (which is the per-user + key within that server's token storage). + """ + + source_id: str + mcp_server_id: str + provider: str | None = None + authorization_url: str + token_url: str + client_id: str + client_secret: SecretStr | None = None + scopes: list[str] = Field(default_factory=list) + redirect_uri: str + use_pkce: bool = True + token_endpoint_auth_method: str = "client_secret_post" + resource: str | None = Field( + default=None, + description=( + "OAuth resource indicator (RFC 8707 / RFC 9728) added to the authorization request, matching NAT's " + "MCPOAuth2Provider. Derived from the protected-resource metadata's `resource` (falling back to the " + "server URL). Authorization-request only — NAT does not send it at token exchange." + ), + ) + + +@dataclass +class _PendingFlow: + source_id: str + user_id: str + client: AsyncOAuth2Client + verifier: str | None + settings: OAuthSourceSettings + expires_at: datetime + + +@dataclass +class NatMcpAuthProvider: + """Concrete :class:`ProtectedSourceAuthProvider` backed by NAT primitives. + + Args: + settings_by_source: OAuth settings per protected source id. + token_storage_resolver: Maps settings -> the NAT ``TokenStorageBase`` for + that source. In production this resolves a shared ``ObjectStore``; in + tests it can return an ``InMemoryTokenStorage`` per source. + challenge_ttl: How long a minted ``auth_url`` / pending flow stays valid. + now: Clock injection for tests. + """ + + settings_by_source: dict[str, OAuthSourceSettings] + token_storage_resolver: Callable[[OAuthSourceSettings], TokenStorageBase] + challenge_ttl: timedelta = _DEFAULT_CHALLENGE_TTL + now: Callable[[], datetime] = field(default_factory=lambda: lambda: datetime.now(UTC)) + + # DEPLOYMENT CONSTRAINT — single API replica (or sticky sessions): + # ``_pending`` holds in-PROCESS OAuth flow state (PKCE verifier + authlib + # client) keyed by the OAuth ``state``. The browser hits /connect on one + # replica and the provider redirects /callback back to the API; the callback + # MUST land on the SAME process that minted the state, or complete_callback + # raises "Unknown or expired auth state". The Helm chart pins the backend to + # replicas: 1, so this holds today. Scaling the API beyond one replica + # requires either session affinity on the ingress/service (route a user's + # /connect and /callback to the same pod) or moving pending-flow state to a + # shared store (e.g. Redis). The token store is already shared/cross-process; + # only this short-lived (challenge_ttl) pending state is process-local. + _pending: dict[str, _PendingFlow] = field(default_factory=dict) + _storage_cache: dict[str, TokenStorageBase] = field(default_factory=dict) + _lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + # ── storage helpers ── + def _storage(self, settings: OAuthSourceSettings) -> TokenStorageBase: + cached = self._storage_cache.get(settings.source_id) + if cached is None: + cached = self.token_storage_resolver(settings) + self._storage_cache[settings.source_id] = cached + return cached + + def is_protected(self, source_id: str) -> bool: + """Whether this source is configured for per-user MCP OAuth.""" + return source_id in self.settings_by_source + + # ── ProtectedSourceAuthProvider ── + async def get_status(self, principal: Principal, source_id: str) -> SourceAuthState: + settings = self.settings_by_source.get(source_id) + if settings is None: + # Not configured for OAuth — treat as an error so the UI can surface it + # rather than silently claiming "connected". + return SourceAuthState(status="error", last_error="Source is not configured for MCP OAuth") + + user_id = principal_user_id(principal) + try: + auth_result = await self._storage(settings).retrieve(user_id) + except Exception as exc: # storage backend failure — report, don't crash the listing + logger.warning("Token storage read failed for source=%s: %s", source_id, exc) + return SourceAuthState(status="error", last_error="Could not read auth state") + + if auth_result is None or not auth_result.credentials: + return SourceAuthState(status="not_connected") + if auth_result.is_expired(): + return SourceAuthState(status="expired", expires_at=auth_result.token_expires_at) + return SourceAuthState(status="connected", expires_at=auth_result.token_expires_at) + + async def start_auth(self, principal: Principal, source_id: str) -> SourceAuthChallenge: + settings = self.settings_by_source.get(source_id) + if settings is None: + raise ValueError(f"Source '{source_id}' is not configured for MCP OAuth") + + user_id = principal_user_id(principal) + state = secrets.token_urlsafe(24) + + client = AsyncOAuth2Client( + client_id=settings.client_id, + client_secret=(settings.client_secret.get_secret_value() if settings.client_secret else None), + redirect_uri=settings.redirect_uri, + scope=" ".join(settings.scopes) if settings.scopes else None, + code_challenge_method="S256" if settings.use_pkce else None, + token_endpoint_auth_method=settings.token_endpoint_auth_method, + ) + + verifier = challenge = None + if settings.use_pkce: + verifier, challenge = pkce.generate_pkce_pair() + + # RFC 8707 resource indicator on the authorize request, matching NAT's + # MCPOAuth2Provider (authorization_kwargs={"resource": ...}). Authorize-only: + # NAT does not send it at token exchange, so complete_callback omits it too. + extra = {"resource": settings.resource} if settings.resource else {} + auth_url, _ = client.create_authorization_url( + settings.authorization_url, + state=state, + code_verifier=verifier if settings.use_pkce else None, + code_challenge=challenge if settings.use_pkce else None, + **extra, + ) + + expires_at = self.now() + self.challenge_ttl + async with self._lock: + stale = self._prune_locked() + self._pending[state] = _PendingFlow( + source_id=source_id, + user_id=user_id, + client=client, + verifier=verifier, + settings=settings, + expires_at=expires_at, + ) + await self._aclose_flows(stale) + logger.info("Started MCP auth challenge for source=%s user=%s state=%s", source_id, user_id, state[:8]) + return SourceAuthChallenge(source_id=source_id, auth_url=auth_url, state=state, expires_at=expires_at) + + async def require_connected( + self, + principal: Principal, + source_ids: list[str], + ) -> list[SourceAuthChallenge]: + blocked: list[SourceAuthChallenge] = [] + for source_id in source_ids: + if not self.is_protected(source_id): + continue # unprotected / unknown sources never block submission + state = await self.get_status(principal, source_id) + if state.status == "connected": + continue + # Best-effort: mint an auth_url so the client can act immediately. If + # minting fails, still report the source as blocked (connect_url only). + try: + challenge = await self.start_auth(principal, source_id) + except Exception as exc: + logger.warning("Could not mint auth_url during preflight for source=%s: %s", source_id, exc) + challenge = SourceAuthChallenge(source_id=source_id, auth_url="", state="") + blocked.append(challenge) + return blocked + + # ── callback completion (AIQ owns the redirect route) ── + async def complete_callback(self, state: str, authorization_response_url: str) -> str: + """Exchange the callback's code for a token and persist it. Returns source_id. + + Raises ``KeyError`` for an unknown/expired state and propagates token + exchange errors to the caller (the route maps them to an HTML error). + """ + async with self._lock: + stale = self._prune_locked() + flow = self._pending.pop(state, None) + await self._aclose_flows(stale) + if flow is None: + raise KeyError("Unknown or expired auth state") + + try: + token = await flow.client.fetch_token( + url=flow.settings.token_url, + authorization_response=authorization_response_url, + code_verifier=flow.verifier, + state=state, + ) + finally: + await flow.client.aclose() + + auth_result = _auth_result_from_token(token) + await self._storage(flow.settings).store(flow.user_id, auth_result) + logger.info("Completed MCP auth for source=%s user=%s", flow.source_id, flow.user_id) + return flow.source_id + + def _prune_locked(self) -> list[_PendingFlow]: + """Pop expired flows and return them so the caller can close their clients. + + Returns the popped flows rather than discarding them: each holds an + ``AsyncOAuth2Client`` (an httpx client pool) that must be ``aclose()``d to + avoid leaking connection state. Pruning runs under ``self._lock`` and is + synchronous, so the async close happens outside the lock via + :meth:`_aclose_flows`. + """ + now = self.now() + expired = [s for s, f in self._pending.items() if f.expires_at <= now] + return [self._pending.pop(s) for s in expired] + + @staticmethod + async def _aclose_flows(flows: list[_PendingFlow]) -> None: + for flow in flows: + try: + await flow.client.aclose() + except Exception as exc: # best-effort cleanup; never fail the caller + logger.debug("Error closing expired auth flow client: %s", exc) + + +def _auth_result_from_token(token: dict) -> AuthResult: + """Build a NAT ``AuthResult`` in the same shape NAT's OAuth provider writes.""" + expires_at: datetime | None = None + if token.get("expires_at"): + expires_at = datetime.fromtimestamp(float(token["expires_at"]), tz=UTC) + access_token = token.get("access_token") + if not access_token: + raise ValueError("Token response missing access_token") + return AuthResult( + credentials=[BearerTokenCred(token=SecretStr(access_token))], + token_expires_at=expires_at, + raw=dict(token), + ) diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/preflight.py b/frontends/aiq_api/src/aiq_api/mcp_auth/preflight.py new file mode 100644 index 000000000..0abe7edba --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/preflight.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Transport-agnostic preflight for per-user MCP auth. + +This is the single source of truth for "may this job be enqueued given which +protected sources the caller selected?". It is reused by two enforcement points +so they cannot drift: + + * the REST ``/v1/jobs/async/submit`` route, which wraps a block in a 409 + JSON response (``_preflight_mcp_auth``); and + * ``submit_agent_job`` itself, which raises :class:`McpAuthRequiredError` so + that programmatic callers (e.g. the chat researcher's async deep-research + submit) cannot bypass the route-level check. +""" + +from __future__ import annotations + +import logging + +from aiq_agent.auth import Principal +from aiq_api.auth.errors import AuthError + +from .models import McpAuthRequiredResponse +from .models import McpAuthRequiredSource +from .provider import ProtectedSourceAuthProvider +from .serialize import connect_url_for + +logger = logging.getLogger(__name__) + + +async def evaluate_mcp_auth( + provider: ProtectedSourceAuthProvider, + principal: Principal, + data_sources: list[str] | None, +) -> McpAuthRequiredResponse | None: + """Return a block descriptor if a selected protected source needs connecting. + + ``data_sources is None`` means the job may use any tool, so every protected + source must be connected; an explicit list restricts the check to those ids. + Returns ``None`` when nothing is blocked (no protected sources selected, or + all connected). + """ + from aiq_agent.common.data_source_registry import get_all_sources + from aiq_agent.common.data_source_registry import get_source + + if data_sources is None: + protected_ids = [s.id for s in get_all_sources() if s.per_user_auth and s.per_user_auth.required] + else: + protected_ids = [ + sid for sid in data_sources if (src := get_source(sid)) and src.per_user_auth and src.per_user_auth.required + ] + if not protected_ids: + return None + + challenges = await provider.require_connected(principal, protected_ids) + if not challenges: + return None + + blocked: list[McpAuthRequiredSource] = [] + for challenge in challenges: + state = await provider.get_status(principal, challenge.source_id) + blocked.append( + McpAuthRequiredSource( + source_id=challenge.source_id, + status=state.status if state.status != "connected" else "not_connected", + connect_url=connect_url_for(challenge.source_id), + auth_url=challenge.auth_url or None, + ) + ) + return McpAuthRequiredResponse(sources=blocked) + + +class McpAuthRequiredError(AuthError): + """Raised by ``submit_agent_job`` when a selected protected source is not connected. + + Subclasses :class:`AuthError` so the chat researcher's deep-research node + surfaces ``str(self)`` to the user (instead of a generic failure) and the + structured ``response`` is available for clients that want connect URLs. + """ + + error_code = "mcp_auth_required" + + def __init__(self, response: McpAuthRequiredResponse) -> None: + self.response = response + names = ", ".join(s.source_id for s in response.sources) or "the selected data source" + super().__init__( + f"Connect the following data source(s) before starting deep research: {names}. " + "Open the data sources panel, click Connect to sign in, then try again." + ) diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/provider.py b/frontends/aiq_api/src/aiq_api/mcp_auth/provider.py new file mode 100644 index 000000000..bf9001700 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/provider.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Adapter boundary for protected MCP source auth. + +This Protocol is the single seam between AIQ's product control plane and NAT's +MCP OAuth mechanics. All NAT calls live behind it (see ``nat_provider``) so the +route handlers and submit preflight never import NAT auth internals directly. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime +from typing import Literal +from typing import Protocol +from typing import runtime_checkable + +from aiq_agent.auth import Principal + +AuthStatus = Literal["connected", "not_connected", "expired", "error"] + + +def principal_user_id(principal: Principal) -> str: + """Canonical per-user key for NAT token storage. + + This MUST match the ``user_id`` that the headless worker sets on the NAT + ``Context`` at job time, otherwise a token connected here will not be found + when ``per_user_mcp_client`` looks it up during execution. Keying by + ``type:sub`` keeps anonymous/no-auth principals distinct from verified ones. + """ + return f"{principal.type}:{principal.sub}" + + +@dataclass(frozen=True) +class SourceAuthState: + """Current per-user auth state for a single source.""" + + status: AuthStatus + expires_at: datetime | None = None + last_error: str | None = None + + +@dataclass(frozen=True) +class SourceAuthChallenge: + """Result of starting (or reusing) an auth challenge for a source. + + ``auth_url`` is the provider login URL to hand to the client. ``state`` is + the opaque OAuth state bound to (principal, source) used by the callback to + complete the flow. ``expires_at`` is the challenge's expiry, not the token's. + """ + + source_id: str + auth_url: str + state: str + expires_at: datetime | None = None + + +@runtime_checkable +class ProtectedSourceAuthProvider(Protocol): + """Product-facing interface AIQ depends on; NAT lives behind the impl.""" + + async def get_status(self, principal: Principal, source_id: str) -> SourceAuthState: + """Return the current per-user auth state for ``source_id`` (read-only).""" + ... + + async def start_auth(self, principal: Principal, source_id: str) -> SourceAuthChallenge: + """Start or resume the OAuth flow and return a provider login URL.""" + ... + + async def require_connected( + self, + principal: Principal, + source_ids: list[str], + ) -> list[SourceAuthChallenge]: + """Return a challenge per *blocked* source; empty list means all connected. + + Used by submit preflight. Implementations should attempt to include an + ``auth_url`` when one can be safely minted, but may return a challenge + with an empty ``auth_url`` if only ``connect_url`` can be offered. + """ + ... diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/runtime_tools.py b/frontends/aiq_api/src/aiq_api/mcp_auth/runtime_tools.py new file mode 100644 index 000000000..18d30096b --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/runtime_tools.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-job runtime resolution of per-user MCP source tools (built in code). + +AIQ agents inherit data-source tools at *build* time, but a per-user MCP source's +tools are per-user/dynamic — the MCP client connects and enumerates tools using +the *user's* token. Two consequences shaped this design: + + * The tools can't be inherited statically by agents (no user at build time). + * A ``per_user_mcp_client`` declared in the *config* is built by NAT's per-user + *interactive* (WebSocket) session builder, which fails for a user with no + token — breaking interactive chat. So we do NOT declare it in config. + +Instead the headless async-job worker builds the per-user MCP client **in code**, +per job, after it has set ``Context.user_id`` to the job owner: it reads the MCP +endpoint from the source's ``mcp_oauth2`` auth provider, connects with the owner's +stored token (no interactive flow), enumerates the tools, and wraps them for the +agent. The client stays open via the caller's ``AsyncExitStack`` for the run. +""" + +from __future__ import annotations + +import logging +from contextlib import AsyncExitStack + +from aiq_agent.common.data_source_registry import get_all_sources +from nat.builder.framework_enum import LLMFrameworkEnum + +logger = logging.getLogger(__name__) + + +class PerUserMcpSourceUnavailableError(RuntimeError): + """An explicitly-selected protected MCP source could not be resolved at run time. + + Raised (rather than silently continuing) when the caller selected specific data + sources and one of them is a configured per-user MCP source whose tools cannot + be built — typically because the owner's token is missing or expired. Surfacing + this lets the client prompt a reconnect instead of returning a web-only answer + that misrepresents which sources were actually used. + """ + + def __init__(self, source_ids: list[str]) -> None: + self.source_ids = source_ids + names = ", ".join(source_ids) + super().__init__( + f"The following selected data source(s) are not connected (or the connection expired): {names}. " + "Reconnect them in the data sources panel and try again." + ) + + +def _resolve_type_registry(builder): + """Resolve NAT's type registry through dependency-tracking child builders.""" + current = builder + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + registry = getattr(current, "_registry", None) + if registry is not None: + return registry + current = getattr(current, "_workflow_builder", None) + raise TypeError(f"Could not resolve a NAT type registry from builder {type(builder).__name__}") + + +async def _token_usable(builder, cfg, source_id: str) -> bool: + """Return whether the job owner has a usable (present, non-expired) token. + + If a token exists but is expired or has no credentials, it is deleted so the + next ``get_status`` reports the source as disconnected (the UI then shows + Reconnect) instead of falsely "connected". Best-effort: any error resolving + the store is treated as "not usable" so we skip rather than crash the job. + """ + from nat.builder.context import Context + + from .factory import _resolve_token_storage + + user_id = Context.get().user_id + if not user_id: + return False + try: + storage = await _resolve_token_storage(builder, cfg, source_id) + auth = await storage.retrieve(user_id) + except Exception as exc: + logger.warning("Source '%s': could not read token state: %s", source_id, exc) + return False + + if auth is None or not auth.credentials: + return False # nothing stored -> status is already not_connected + if auth.is_expired(): + # Invalidate so the card stops showing "connected" and prompts Reconnect. + try: + await storage.delete(user_id) + except Exception as exc: + logger.debug("Source '%s': failed to delete expired token: %s", source_id, exc) + logger.warning("Source '%s': stored token is expired; invalidated. User must reconnect.", source_id) + return False + return True + + +async def open_per_user_mcp_tools( + *, + builder, + data_sources: list[str] | None, + exit_stack: AsyncExitStack, + wrapper_type: LLMFrameworkEnum | str = LLMFrameworkEnum.LANGCHAIN, +) -> list: + """Build per-user MCP clients for selected protected sources; return their tools. + + Args: + builder: The per-job ``WorkflowBuilder`` (resolves the auth provider + the + framework tool wrapper). + data_sources: Selected source ids, or ``None`` meaning "all" (so every + connected protected source's tools are made available). + exit_stack: An ``AsyncExitStack`` whose lifetime spans the agent run; the + MCP client contexts are entered here and torn down when it closes. + wrapper_type: Agent framework to wrap tools for. + + Returns: + Framework-wrapped tools (possibly empty). Best-effort for the "all" case + (``data_sources is None``): a failure resolving one source is logged and + skipped so it never breaks the job. + + Raises: + PerUserMcpSourceUnavailableError: when ``data_sources`` is an explicit list + and one of the selected, configured per-user MCP sources cannot be + resolved (missing/expired token, unreachable server). The caller asked + for those sources specifically, so failing is preferable to silently + answering without them. + + Precondition: ``Context.user_id`` must already be set to the job owner. + """ + from nat.plugins.mcp.client.client_config import MCPServerConfig + from nat.plugins.mcp.client.client_config import MCPToolOverrideConfig + from nat.plugins.mcp.client.client_config import PerUserMCPClientConfig + from nat.plugins.mcp.client.client_impl import per_user_mcp_client_function_group + + selected = None if data_sources is None else {s.lower() for s in data_sources} + tools: list = [] + # Explicitly-selected per-user sources we couldn't resolve -> fail closed below. + unavailable: list[str] = [] + + for source in get_all_sources(): + pua = source.per_user_auth + if pua is None or not pua.required or not pua.auth_provider: + continue + explicitly_selected = selected is not None and source.id.lower() in selected + if selected is not None and not explicitly_selected: + continue + + try: + # The mcp_oauth2 provider's server_url is the MCP endpoint; reuse it so + # the connect flow and the job-time client target the same server, and + # the client authenticates via the same provider (stored token lookup). + provider = await builder.get_auth_provider(pua.auth_provider) + server_url = str(getattr(getattr(provider, "config", None), "server_url", "") or "") + if not server_url: + logger.warning( + "Source '%s': auth provider '%s' has no server_url; cannot build MCP client.", + source.id, + pua.auth_provider, + ) + if explicitly_selected: + unavailable.append(source.id) + continue + + # Reconcile UI status with reality: the data-source card reports + # "connected" from an offline token read, but the token can be expired + # while the card still says connected. The use-site is authoritative — + # if the owner's stored token is missing/expired here, invalidate it so + # the next get_status returns not_connected/expired (UI -> Reconnect) + # and skip, rather than failing the live MCP call and silently dropping + # the tool while the UI keeps claiming connected. + if not await _token_usable(builder, getattr(provider, "config", None), source.id): + if explicitly_selected: + unavailable.append(source.id) + continue + + # Give terse/blank MCP tools clear names + descriptions so the agent + # reliably selects them over web search (declared on the source). + tool_overrides = { + name: MCPToolOverrideConfig(alias=ov.get("alias"), description=ov.get("description")) + for name, ov in (pua.tool_overrides or {}).items() + } + client_cfg = PerUserMCPClientConfig( + server=MCPServerConfig(transport="streamable-http", url=server_url, auth_provider=pua.auth_provider), + tool_overrides=tool_overrides, + ) + group = await exit_stack.enter_async_context(per_user_mcp_client_function_group(client_cfg, builder)) + fns = await group.get_accessible_functions() + # Resolve the tool wrapper via the builder's type registry rather than + # `builder._registry` directly: in server mode `builder` is a ChildBuilder, + # which has no `_registry` (it delegates to its parent), so the attribute + # access raised AttributeError and this whole block was silently swallowed + # — dropping the selected source's tools and falling back to web search. + wrapper = _resolve_type_registry(builder).get_tool_wrapper(llm_framework=wrapper_type) + wrapped = [wrapper.build_fn(name, fn, builder) for name, fn in fns.items()] + tools.extend(wrapped) + + # Map these runtime-resolved tools to their data source so the agents' + # citation/source capture treats their results as sources. Without this, + # get_source_id_for_tool returns None for them and shallow research raises + # EmptySourceRegistryError ("no sources captured") even on a successful read. + from aiq_agent.common.data_source_registry import register_tool_sources + + register_tool_sources({getattr(t, "name", ""): source.id for t in wrapped if getattr(t, "name", "")}) + logger.info("Resolved %d per-user MCP tool(s) for source '%s'.", len(wrapped), source.id) + except Exception: + logger.exception( + "Failed to resolve per-user MCP tools for source '%s'; continuing without them.", + source.id, + ) + if explicitly_selected: + unavailable.append(source.id) + + # Fail closed for sources the caller singled out but we couldn't resolve (e.g. + # a token that expired between submit-time preflight and job execution). For the + # "all" case (data_sources is None) we stay best-effort — the user didn't ask + # for these specifically, so a missing one shouldn't sink the whole run. + if unavailable: + raise PerUserMcpSourceUnavailableError(unavailable) + + return tools diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/serialize.py b/frontends/aiq_api/src/aiq_api/mcp_auth/serialize.py new file mode 100644 index 000000000..e3181a302 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/serialize.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Helpers to project registry + auth state into API models.""" + +from __future__ import annotations + +from aiq_agent.auth import Principal +from aiq_agent.common.data_source_registry import DataSourceMeta + +from .models import PerUserAuthInfo +from .provider import ProtectedSourceAuthProvider + +# Statuses for which the client should be offered a way to (re)connect. +_ACTIONABLE = {"not_connected", "expired", "error"} + + +def connect_url_for(source_id: str) -> str: + return f"/v1/auth/mcp/{source_id}/connect" + + +def status_url_for(source_id: str) -> str: + return f"/v1/auth/mcp/{source_id}/status" + + +async def build_listing_auth_info( + provider: ProtectedSourceAuthProvider, + principal: Principal, + source: DataSourceMeta, +) -> PerUserAuthInfo | None: + """Build the ``per_user_auth`` block for ``GET /v1/data_sources``. + + Read-only: never mints ``auth_url`` (no OAuth state is created here). Returns + ``None`` for sources with no per-user auth declaration so the field is + omitted entirely. + """ + pua = source.per_user_auth + if pua is None: + return None + + info = PerUserAuthInfo( + required=pua.required, + type=pua.type, + provider=pua.provider, + mcp_server_id=pua.mcp_server_id, + ) + if not pua.required: + return info + + state = await provider.get_status(principal, source.id) + info.status = state.status + info.expires_at = state.expires_at + info.last_error = state.last_error + if state.status in _ACTIONABLE: + info.connect_url = connect_url_for(source.id) + return info diff --git a/frontends/aiq_api/src/aiq_api/mcp_auth/sqlite_object_store.py b/frontends/aiq_api/src/aiq_api/mcp_auth/sqlite_object_store.py new file mode 100644 index 000000000..756111e91 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/mcp_auth/sqlite_object_store.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A SQLite-backed NAT ``ObjectStore`` for serviceless, cross-process storage. + +NAT ships only ``in_memory`` (process-local) in core, plus ``redis``/``s3``/ +``mysql`` in separate packages — all of which are network services. The per-user +MCP token store, however, is written by the API process (OAuth ``/connect`` and +``/callback``) and read by a *separate* Dask worker process at job time, so an +in-memory store can't bridge that gap and a network service can't run "with no +deployment". + +A file on local disk is the one option that is both serviceless **and** visible +across processes: both processes open the same SQLite file (WAL mode), exactly +how this project's job/checkpoint/summary stores already span the same two +processes locally. The same class points at a path locally and is swapped for +``redis`` in deployment via the config's env-selectable ``_type``. + +This implements NAT's ``ObjectStore`` interface only; the token serialization, +refresh, and expiry logic stay in NAT's ``ObjectStoreTokenStorage`` / +``mcp_oauth2`` on top — nothing here is token-specific. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import time + +import aiosqlite +from pydantic import Field + +from nat.builder.builder import Builder +from nat.cli.register_workflow import register_object_store +from nat.data_models.object_store import KeyAlreadyExistsError +from nat.data_models.object_store import NoSuchKeyError +from nat.data_models.object_store import ObjectStoreBaseConfig +from nat.object_store.interfaces import ObjectStore +from nat.object_store.models import ObjectStoreItem +from nat.utils.type_utils import override + +logger = logging.getLogger(__name__) + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS objects ( + key TEXT PRIMARY KEY, + data BLOB NOT NULL, + content_type TEXT, + metadata TEXT, + expires_at REAL +) +""" + + +def _restrict_file(path: str) -> None: + """chmod ``path`` to owner-only (0600) if it exists; best-effort.""" + try: + if os.path.exists(path): + os.chmod(path, 0o600) + except OSError as exc: # e.g. unsupported on the platform — don't crash startup + logger.warning("Could not restrict permissions on %s: %s", path, exc) + + +def _ensure_private_file(path: str) -> None: + """Ensure the token DB file exists with owner-only (0600) permissions. + + Creating it ourselves with ``O_CREAT`` + ``0o600`` closes the window where + sqlite would otherwise create it at the umask default (0644), and we chmod + unconditionally so an existing 0644 file is tightened on next open. + """ + if not os.path.exists(path): + try: + fd = os.open(path, os.O_CREAT | os.O_WRONLY, 0o600) + os.close(fd) + except OSError as exc: + logger.warning("Could not pre-create token DB %s with private perms: %s", path, exc) + return + _restrict_file(path) + + +class AiqSqliteObjectStoreConfig(ObjectStoreBaseConfig, name="aiq_sqlite"): + """SQLite-file object store — serviceless and shared across processes. + + Suited to local/single-host runs where standing up Redis/S3/MySQL is + undesirable. For multi-replica deployments use a networked object store + (e.g. ``redis``) instead, since a SQLite file is only shared by processes + that can see the same path. + """ + + db_path: str = Field( + default="./mcp_tokens.db", + description="Path to the SQLite database file (created on first use).", + ) + bucket_name: str | None = Field( + default=None, + description="Optional key prefix so multiple logical buckets can share one file.", + ) + ttl: int | None = Field( + default=None, + description="TTL in seconds for stored objects (None = no expiration).", + ) + + +class AiqSqliteObjectStore(ObjectStore): + """ObjectStore backed by a single SQLite file. + + Cross-process visibility relies on WAL mode + a busy timeout: each process + opens its own connection to the same file, and committed writes from one + process are readable by the others. + """ + + def __init__(self, db_path: str, bucket_name: str | None = None, ttl: int | None = None) -> None: + self._db_path = db_path + self._prefix = f"{bucket_name}/" if bucket_name else "" + self._ttl = ttl + self._db: aiosqlite.Connection | None = None + self._lock = asyncio.Lock() + + # ── connection / helpers ── + async def _conn(self) -> aiosqlite.Connection: + if self._db is None: + # The serialized objects hold plaintext access/refresh tokens, so the + # file must never be group/world-readable. Create it 0600 BEFORE sqlite + # opens it (default umask 022 would otherwise yield 0644), and tighten + # an existing file too. + _ensure_private_file(self._db_path) + db = await aiosqlite.connect(self._db_path) + # WAL allows concurrent readers across processes alongside one writer; + # busy_timeout waits out a peer's write lock instead of failing fast. + await db.execute("PRAGMA journal_mode=WAL") + await db.execute("PRAGMA busy_timeout=5000") + await db.execute(_SCHEMA) + await db.commit() + # WAL/SHM sidecars are created by sqlite on first write and also carry + # token bytes; restrict them once they exist. + for suffix in ("-wal", "-shm"): + _restrict_file(self._db_path + suffix) + self._db = db + return self._db + + def _k(self, key: str) -> str: + return f"{self._prefix}{key}" + + def _expiry(self) -> float | None: + return time.time() + self._ttl if self._ttl else None + + @staticmethod + def _expired(expires_at: float | None) -> bool: + return expires_at is not None and expires_at <= time.time() + + def _row_to_item(self, row: tuple) -> ObjectStoreItem: + data, content_type, metadata = row[0], row[1], row[2] + return ObjectStoreItem( + data=data, + content_type=content_type, + metadata=json.loads(metadata) if metadata else None, + ) + + # ── ObjectStore interface ── + @override + async def put_object(self, key: str, item: ObjectStoreItem) -> None: + k = self._k(key) + async with self._lock: + db = await self._conn() + # An expired row is logically absent, so let a fresh put replace it. + async with db.execute("SELECT expires_at FROM objects WHERE key = ?", (k,)) as cur: + existing = await cur.fetchone() + if existing is not None and not self._expired(existing[0]): + raise KeyAlreadyExistsError(key) + await db.execute( + "INSERT OR REPLACE INTO objects (key, data, content_type, metadata, expires_at) VALUES (?, ?, ?, ?, ?)", + (k, item.data, item.content_type, json.dumps(item.metadata) if item.metadata else None, self._expiry()), + ) + await db.commit() + + @override + async def upsert_object(self, key: str, item: ObjectStoreItem) -> None: + k = self._k(key) + async with self._lock: + db = await self._conn() + await db.execute( + "INSERT OR REPLACE INTO objects (key, data, content_type, metadata, expires_at) VALUES (?, ?, ?, ?, ?)", + (k, item.data, item.content_type, json.dumps(item.metadata) if item.metadata else None, self._expiry()), + ) + await db.commit() + + @override + async def get_object(self, key: str) -> ObjectStoreItem: + k = self._k(key) + async with self._lock: + db = await self._conn() + async with db.execute( + "SELECT data, content_type, metadata, expires_at FROM objects WHERE key = ?", (k,) + ) as cur: + row = await cur.fetchone() + if row is None: + raise NoSuchKeyError(key) + if self._expired(row[3]): + await db.execute("DELETE FROM objects WHERE key = ?", (k,)) + await db.commit() + raise NoSuchKeyError(key) + return self._row_to_item(row) + + @override + async def delete_object(self, key: str) -> None: + k = self._k(key) + async with self._lock: + db = await self._conn() + cur = await db.execute("DELETE FROM objects WHERE key = ?", (k,)) + await db.commit() + if cur.rowcount == 0: + raise NoSuchKeyError(key) + + async def aclose(self) -> None: + if self._db is not None: + await self._db.close() + self._db = None + + +@register_object_store(config_type=AiqSqliteObjectStoreConfig) +async def aiq_sqlite_object_store(config: AiqSqliteObjectStoreConfig, builder: Builder): + store = AiqSqliteObjectStore(db_path=config.db_path, bucket_name=config.bucket_name, ttl=config.ttl) + # Log the ABSOLUTE path: the API and worker must resolve db_path to the same + # file. With a relative path that only holds if they share a working dir — a + # cwd divergence otherwise creates two files and silently loses tokens. The + # absolute path here makes such a mismatch visible in each process's logs. + logger.info("SQLite object store initialized at %s", os.path.abspath(config.db_path)) + try: + yield store + finally: + await store.aclose() diff --git a/frontends/aiq_api/src/aiq_api/register.py b/frontends/aiq_api/src/aiq_api/register.py index 44f6f4a04..3ca50d4cf 100644 --- a/frontends/aiq_api/src/aiq_api/register.py +++ b/frontends/aiq_api/src/aiq_api/register.py @@ -21,6 +21,9 @@ """ # Import the plugin registration - this makes it discoverable by NAT +# Importing the module runs its @register_object_store decorator, adding the +# serviceless SQLite token store (_type: aiq_sqlite) to NAT's object-store registry. +from .mcp_auth import sqlite_object_store as _sqlite_object_store # noqa: F401 from .plugin import register_aiq_api __all__ = ["register_aiq_api"] diff --git a/frontends/aiq_api/src/aiq_api/routes/auth.py b/frontends/aiq_api/src/aiq_api/routes/auth.py new file mode 100644 index 000000000..f702540e4 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/routes/auth.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-user MCP auth routes: status, connect, and OAuth callback. + +AIQ owns these routes (the control plane). The actual OAuth mechanics live +behind :class:`ProtectedSourceAuthProvider`. The callback completes NAT's token +exchange and persists the token, then closes the popup. +""" + +from __future__ import annotations + +import html +import json +import logging + +from fastapi import FastAPI +from fastapi import HTTPException +from fastapi import Request +from fastapi.responses import HTMLResponse + +from aiq_agent.common.data_source_registry import get_source + +from ..jobs.access import require_verified_principal +from ..mcp_auth.models import SourceAuthStatusResponse +from ..mcp_auth.models import SourceConnectResponse +from ..mcp_auth.nat_provider import NatMcpAuthProvider +from ..mcp_auth.serialize import connect_url_for + +logger = logging.getLogger(__name__) + +_CALLBACK_HTML = """ +{title} + +

{message}

+ +""" + + +def _callback_page(source_id: str, *, ok: bool, message: str) -> HTMLResponse: + # `message` can carry a provider-controlled value (e.g. the OAuth `error` + # query param reflected by the callback), so HTML-escape it before it lands in + # the page body — otherwise a crafted `error` executes script in AIQ's origin. + # `source_id` is a validated registry id and json.dumps escapes it for the JS + # string context. + page = _CALLBACK_HTML.format( + title="Connected" if ok else "Authentication failed", + message=html.escape(message), + source_id_js=json.dumps(source_id), + ok_js="true" if ok else "false", + ) + return HTMLResponse(content=page, status_code=200 if ok else 400, headers={"Cache-Control": "no-cache"}) + + +def _require_protected_source(source_id: str): + """Return the registry source or raise 404 if it isn't a protected MCP source.""" + source = get_source(source_id) + if source is None or source.per_user_auth is None or not source.per_user_auth.required: + raise HTTPException(404, f"Unknown protected data source: {source_id}") + return source + + +def register_mcp_auth_routes(app: FastAPI, provider: NatMcpAuthProvider) -> None: + """Register the per-user MCP auth routes against ``provider``.""" + + @app.get( + "/v1/auth/mcp/{source_id}/status", + response_model=SourceAuthStatusResponse, + tags=["mcp auth"], + summary="Get per-user auth status for a protected source", + ) + async def mcp_auth_status(source_id: str) -> SourceAuthStatusResponse: + _require_protected_source(source_id) + principal = require_verified_principal() + state = await provider.get_status(principal, source_id) + connect_url = connect_url_for(source_id) if state.status != "connected" else None + return SourceAuthStatusResponse( + source_id=source_id, + status=state.status, + expires_at=state.expires_at, + connect_url=connect_url, + last_error=state.last_error, + ) + + @app.post( + "/v1/auth/mcp/{source_id}/connect", + response_model=SourceConnectResponse, + tags=["mcp auth"], + summary="Start (or resume) the OAuth flow for a protected source", + ) + async def mcp_auth_connect(source_id: str) -> SourceConnectResponse: + _require_protected_source(source_id) + principal = require_verified_principal() + + # If already connected, don't mint a new challenge. + state = await provider.get_status(principal, source_id) + if state.status == "connected": + return SourceConnectResponse(source_id=source_id, status="connected") + + try: + challenge = await provider.start_auth(principal, source_id) + except ValueError as exc: + # Source declared but not configured for OAuth in this deployment. + raise HTTPException(503, str(exc)) from exc + except Exception as exc: + logger.exception("Failed to start MCP auth for source=%s", source_id) + raise HTTPException(502, "Could not start authentication flow") from exc + + return SourceConnectResponse( + source_id=source_id, + status="auth_required", + auth_url=challenge.auth_url, + expires_at=challenge.expires_at, + ) + + @app.get( + "/v1/auth/mcp/{source_id}/callback", + tags=["mcp auth"], + summary="OAuth redirect callback for a protected source", + include_in_schema=False, + ) + async def mcp_auth_callback(source_id: str, request: Request) -> HTMLResponse: + _require_protected_source(source_id) + state = request.query_params.get("state") + if not state: + return _callback_page(source_id, ok=False, message="Missing state. Please restart the connection.") + + error = request.query_params.get("error") + if error: + return _callback_page(source_id, ok=False, message=f"Authorization was denied: {error}") + + try: + completed_source = await provider.complete_callback(state, str(request.url)) + except KeyError: + return _callback_page(source_id, ok=False, message="This connection link has expired. Please try again.") + except Exception: + logger.exception("MCP auth callback failed for source=%s", source_id) + return _callback_page(source_id, ok=False, message="Authentication failed. Please try again.") + + if completed_source != source_id: + # State was bound to a different source — refuse to cross the streams. + logger.warning("Callback source mismatch: path=%s flow=%s", source_id, completed_source) + return _callback_page(source_id, ok=False, message="Connection mismatch. Please try again.") + + return _callback_page(source_id, ok=True, message="Connected. You can close this window.") + + logger.info("Registered /v1/auth/mcp/{source_id} status, connect, and callback routes") diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index f7d46e584..96121cb49 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -53,6 +53,7 @@ from nat.builder.framework_enum import LLMFrameworkEnum from ..jobs.access import require_verified_principal +from ..mcp_auth.models import PerUserAuthInfo from ..registry import AGENT_REGISTRY from ..registry import get_agent_config @@ -226,6 +227,18 @@ async def _get_agent_available_source_ids(builder: WorkflowBuilder, agent_config sid = get_source_id_for_tool(name) if sid is not None: source_ids.add(sid) + + # Per-user MCP sources (e.g. Google Drive) contribute NO static tools — their + # tools are resolved per-user at run time by open_per_user_mcp_tools, so they + # never appear in the loop above. Treat a configured protected source as an + # available runtime candidate so submit validation doesn't 422 it; connectivity + # is enforced separately by the MCP auth preflight (409 mcp_auth_required). + from aiq_agent.common.data_source_registry import get_all_sources + + for source in get_all_sources(): + pua = source.per_user_auth + if pua is not None and pua.required: + source_ids.add(source.id) return sorted(source_ids) @@ -307,6 +320,24 @@ async def _validate_data_sources_for_agent( ) +async def _preflight_mcp_auth(provider, principal, data_sources: list[str] | None): + """Return a 409 JSONResponse if any selected protected source is not connected, else None. + + Thin HTTP wrapper over the shared :func:`evaluate_mcp_auth`; the same check + runs inside ``submit_agent_job`` (raising instead) so programmatic submitters + cannot bypass it. Source existence is validated earlier by + ``_validate_data_sources_for_agent``. + """ + from fastapi.responses import JSONResponse + + from ..mcp_auth.preflight import evaluate_mcp_auth + + body = await evaluate_mcp_auth(provider, principal, data_sources) + if body is None: + return None + return JSONResponse(status_code=409, content=body.model_dump(mode="json")) + + class JobStatusResponse(BaseModel): """Job status response.""" @@ -411,7 +442,15 @@ class DataSource(BaseModel): id: str = Field(..., description="Unique identifier for the data source") name: str = Field(..., description="Display name") description: str | None = Field(default=None, description="Human-readable description") + default_enabled: bool = Field( + default=True, + description="Whether the source is toggled on by default in the UI (from registry metadata)", + ) requires_auth: bool = Field(default=False, description="Whether user authentication is required") + per_user_auth: PerUserAuthInfo | None = Field( + default=None, + description="Per-user MCP OAuth state for a protected source (omitted for unprotected sources)", + ) async def register_job_routes(app: FastAPI, builder: WorkflowBuilder, worker: FastApiFrontEndPluginWorker) -> None: @@ -436,6 +475,21 @@ async def register_job_routes(app: FastAPI, builder: WorkflowBuilder, worker: Fa from ..jobs.report_context import to_initial_files from ..jobs.submit import JobIdConflictError from ..jobs.submit import submit_agent_job as submit_authorized_job + from ..mcp_auth.factory import build_mcp_auth_provider + from ..mcp_auth.preflight import McpAuthRequiredError + from ..mcp_auth.serialize import build_listing_auth_info + from .auth import register_mcp_auth_routes + + # Per-user MCP auth control plane. The provider is shared by the data-source + # listing, the status/connect/callback routes, and submit preflight so a flow + # started via /connect can be completed by /callback in the same process. + mcp_auth_provider = await build_mcp_auth_provider(builder) + # Publish the provider process-wide so submit_agent_job() can run the same + # connect-state preflight for programmatic submitters, not just this REST route. + from ..mcp_auth.active import set_active_mcp_auth_provider + + set_active_mcp_auth_provider(mcp_auth_provider) + register_mcp_auth_routes(app, mcp_auth_provider) if not get_all_sources(): logger.warning( @@ -467,16 +521,22 @@ async def list_agents() -> AgentListResponse: summary="List data sources", ) async def list_data_sources() -> list[DataSource]: - """List available data sources dynamically from the registry.""" - return [ - DataSource( - id=source.id, - name=source.name, - description=source.description, - requires_auth=source.requires_auth, + """List available data sources, including the current user's per-source auth state.""" + principal = require_verified_principal() + sources = [] + for source in get_all_sources(): + per_user_auth = await build_listing_auth_info(mcp_auth_provider, principal, source) + sources.append( + DataSource( + id=source.id, + name=source.name, + description=source.description, + default_enabled=source.default_enabled, + requires_auth=source.requires_auth, + per_user_auth=per_user_auth, + ) ) - for source in get_all_sources() - ] + return sources logger.info("Registered /v1/data_sources and /v1/jobs/async/agents routes") @@ -549,7 +609,12 @@ async def health_check(): ), responses={ 400: {"description": "Unknown agent type or invalid request"}, - 409: {"description": "A custom job_id was supplied that collides with an existing job"}, + 409: { + "description": ( + "A custom job_id was supplied that collides with an existing job, or a selected " + "protected data source requires per-user OAuth connection" + ) + }, 422: {"description": "One or more unknown or agent-unavailable data source IDs"}, 500: {"description": "Failed to persist async job authorization metadata"}, 503: {"description": "Dask scheduler not available"}, @@ -591,6 +656,13 @@ async def submit_job( if _sandbox_caps_configured() and _agent_uses_sandbox(builder, agent_config.config_name): await _enforce_sandbox_concurrency(db_url, principal) + # Preflight protected MCP sources: block before enqueue if a selected + # protected source is not connected. When data_sources is None the job + # may use any tool, so every protected source must be connected. + mcp_block = await _preflight_mcp_auth(mcp_auth_provider, principal, req.data_sources) + if mcp_block is not None: + return mcp_block + # Propagate auth token to Dask worker for requires_auth data sources from aiq_agent.auth import get_auth_token @@ -608,6 +680,14 @@ async def submit_job( ) except JobIdConflictError: raise HTTPException(409, f"Job already exists: {req.job_id}") + except McpAuthRequiredError as e: + # submit_agent_job runs the same MCP preflight and raises if a selected + # protected source became disconnected between the route preflight above + # and enqueue. Surface the SAME 409 mcp_auth_required contract instead of + # letting it fall through to the generic 500 handler. + from fastapi.responses import JSONResponse + + return JSONResponse(status_code=409, content=e.response.model_dump(mode="json")) except RuntimeError as e: # The principal is resolved above, so a RuntimeError here is an # availability/config failure (e.g. scheduler not configured), not an @@ -1436,6 +1516,7 @@ async def _sse_generator_postgres(job_store, job_id: str, db_url: str, start_eve Achieves sub-10ms latency compared to 500ms polling interval. """ import asyncio + import time import asyncpg @@ -1450,6 +1531,12 @@ async def _sse_generator_postgres(job_store, job_id: str, db_url: str, start_eve sequence_id = start_event_id terminal_statuses = {JobStatus.SUCCESS.value, JobStatus.FAILURE.value, JobStatus.INTERRUPTED.value} is_reconnect = start_event_id > 0 + # Emit an SSE keepalive comment after this many seconds of silence so an + # upstream idle timeout (OpenShift router / edge / proxy) never closes the + # connection. job.heartbeat only starts once the worker runs, so it does not + # cover worker cold-start on the first request — this keepalive does. + SSE_KEEPALIVE_INTERVAL = 15.0 + last_keepalive = time.monotonic() def format_sse(event_type: str, data: dict, event_id: int | None = None) -> str: """Format an SSE frame and advance (or set) the monotonic event sequence id.""" @@ -1560,6 +1647,13 @@ def notification_handler(connection, pid, channel_name, payload): last_event_id = db_event_id event_type = event.pop("type", "event") yield format_sse(event_type, event, db_event_id) + # Keepalive during silent periods (e.g. worker cold-start) so an + # upstream idle timeout never closes the connection. + if fallback_events: + last_keepalive = time.monotonic() + elif (time.monotonic() - last_keepalive) >= SSE_KEEPALIVE_INTERVAL: + last_keepalive = time.monotonic() + yield ": keepalive\n\n" job = await job_store.get_job(job_id) if not job: @@ -1624,6 +1718,7 @@ async def _sse_generator_polling(job_store, job_id: str, db_url: str, start_even Supports graceful shutdown via the SSE connection manager. """ import asyncio + import time from nat.front_ends.fastapi.async_jobs.job_store import JobStatus @@ -1638,6 +1733,11 @@ async def _sse_generator_polling(job_store, job_id: str, db_url: str, start_even is_reconnect = start_event_id > 0 in_replay_mode = True replay_mode_announced = False + # Emit an SSE keepalive comment after this many seconds of silent live polling + # so an upstream idle timeout never closes the connection (e.g. during worker + # cold-start, before the first event or 30s heartbeat arrives). + SSE_KEEPALIVE_INTERVAL = 15.0 + last_keepalive = time.monotonic() def format_sse(event_type: str, data: dict, event_id: int | None = None) -> str: """Format an SSE frame and advance (or set) the monotonic event sequence id.""" @@ -1677,6 +1777,7 @@ def format_sse(event_type: str, data: dict, event_id: int | None = None) -> str: events = await EventStore.get_events_async(db_url, job_id, last_event_id, limit) if events: + last_keepalive = time.monotonic() logger.info(f"SSE: Fetched {len(events)} events for job {job_id} (after_id={last_event_id})") elif job.status in terminal_statuses: logger.warning(f"SSE: No events found for completed job {job_id} (after_id={last_event_id})") @@ -1732,6 +1833,12 @@ def format_sse(event_type: str, data: dict, event_id: int | None = None) -> str: replay_mode_announced = True yield format_sse("stream.mode", {"mode": "live"}) + # Keepalive during silent live periods (e.g. worker cold-start) so the + # idle-looking connection isn't closed by an upstream idle timeout. + if (time.monotonic() - last_keepalive) >= SSE_KEEPALIVE_INTERVAL: + last_keepalive = time.monotonic() + yield ": keepalive\n\n" + shutdown_signaled = await connection_manager.wait_or_shutdown(0.5) if shutdown_signaled: logger.info("SSE stream closing for job %s due to server shutdown (during wait)", job_id) diff --git a/frontends/aiq_api/tests/test_auth.py b/frontends/aiq_api/tests/test_auth.py index d3a56acd8..7c70e3a1f 100644 --- a/frontends/aiq_api/tests/test_auth.py +++ b/frontends/aiq_api/tests/test_auth.py @@ -897,6 +897,20 @@ def test_path_allowed_exact_and_prefix(self) -> None: assert mw._path_allowed("/v1/jobs/async/job/abc/result") is True assert mw._path_allowed("/v1/jobs/async/job") is True assert mw._path_allowed("/nope") is False + # Per-user MCP auth routes must be reachable externally. + assert mw._path_allowed("/v1/auth/mcp/gdrive/status") is True + assert mw._path_allowed("/v1/auth/mcp/gdrive/connect") is True + assert mw._path_allowed("/v1/auth/mcp/gdrive/callback") is True + + def test_mcp_oauth_callback_is_auth_exempt(self) -> None: + from aiq_api.auth.middleware import _is_oauth_callback_path + + # Only the OAuth redirect callback is exempt (no AIQ token; secured by state). + assert _is_oauth_callback_path("/v1/auth/mcp/gdrive/callback") is True + # status/connect must still require auth (they need the principal). + assert _is_oauth_callback_path("/v1/auth/mcp/gdrive/status") is False + assert _is_oauth_callback_path("/v1/auth/mcp/gdrive/connect") is False + assert _is_oauth_callback_path("/v1/data_sources") is False @pytest.mark.asyncio async def test_non_http_passthrough(self) -> None: 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 a6e17ec90..1d0045ee4 100644 --- a/frontends/aiq_api/tests/test_job_submit_data_sources.py +++ b/frontends/aiq_api/tests/test_job_submit_data_sources.py @@ -643,3 +643,25 @@ async def test_validation_does_not_call_get_all_tool_refs_when_fn_config_tools_i _, kwargs = builder.get_tools.await_args assert kwargs["tool_names"] == ["knowledge_search_tool"] submitted_job.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_list_data_sources_exposes_default_enabled(submit_app): + """GET /v1/data_sources must surface the registry's default_enabled (not hardcode True).""" + app, _submitted_job, _builder = submit_app + # Re-populate at request time: list_data_sources() reads the registry per request. + reset_registry() + populate_from_config( + [ + {"id": "web_search", "name": "Web Search", "description": "x"}, # default_enabled -> True + {"id": "off_by_default", "name": "Off", "description": "y", "default_enabled": False}, + ] + ) + + with TestClient(app) as client: + response = client.get("/v1/data_sources") + + assert response.status_code == 200 + by_id = {s["id"]: s for s in response.json()} + assert by_id["web_search"]["default_enabled"] is True + assert by_id["off_by_default"]["default_enabled"] is False diff --git a/frontends/aiq_api/tests/test_mcp_auth_factory.py b/frontends/aiq_api/tests/test_mcp_auth_factory.py new file mode 100644 index 000000000..515a8ed3d --- /dev/null +++ b/frontends/aiq_api/tests/test_mcp_auth_factory.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from aiq_agent.common.data_source_registry import populate_from_config +from aiq_agent.common.data_source_registry import reset_registry +from aiq_api.mcp_auth.factory import build_mcp_auth_provider + + +class _FakeProvider: + """Stand-in for NAT's MCPOAuth2Provider with discovery already resolved.""" + + def __init__(self, *, discover_ok: bool = True, client_id: str | None = "client-xyz"): + self.config = SimpleNamespace( + server_url="", # empty -> factory probe is skipped (offline test) + redirect_uri="https://aiq.example/v1/auth/mcp/gdrive/callback", + scopes=["https://www.googleapis.com/auth/drive.readonly"], + client_id=client_id, + client_secret="shh", # pragma: allowlist secret + use_pkce=True, + token_endpoint_auth_method="client_secret_post", + token_storage_object_store="mcp_token_store", + ) + self._discover_ok = discover_ok + self._cached_endpoints = None + self._cached_credentials = None + + async def _discover_and_register(self, response=None): + if not self._discover_ok: + raise RuntimeError("server unreachable") + self._cached_endpoints = SimpleNamespace( + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + ) + self._cached_credentials = SimpleNamespace(client_id="client-xyz", client_secret="shh") + + +class _FakeBuilder: + def __init__(self, provider): + self._provider = provider + + async def get_auth_provider(self, name): + if self._provider is None: + raise KeyError(name) + return self._provider + + async def get_object_store_client(self, name): + return SimpleNamespace(name=name) # ObjectStoreTokenStorage only needs an object + + +@pytest.fixture(autouse=True) +def _registry(): + reset_registry() + yield + reset_registry() + + +def _register_gdrive(): + populate_from_config( + [ + { + "id": "gdrive", + "name": "Google Drive", + "requires_auth": True, + "per_user_auth": { + "required": True, + "provider": "google", + "mcp_server_id": "gdrive", + "auth_provider": "mcp_oauth2_gdrive", + }, + }, + ] + ) + + +def test_factory_resolves_settings_from_nat_provider(): + _register_gdrive() + provider = build_from(_FakeBuilder(_FakeProvider())) + assert provider.is_protected("gdrive") + settings = provider.settings_by_source["gdrive"] + assert settings.authorization_url == "https://accounts.google.com/o/oauth2/v2/auth" + assert settings.token_url == "https://oauth2.googleapis.com/token" + assert settings.client_id == "client-xyz" + assert settings.redirect_uri.endswith("/v1/auth/mcp/gdrive/callback") + assert settings.scopes == ["https://www.googleapis.com/auth/drive.readonly"] + + +def test_factory_skips_source_when_provider_missing(): + _register_gdrive() + provider = build_from(_FakeBuilder(None)) # get_auth_provider raises + assert not provider.is_protected("gdrive") # left unconfigured -> status will be 'error' + + +def test_factory_skips_source_when_discovery_fails(): + _register_gdrive() + provider = build_from(_FakeBuilder(_FakeProvider(discover_ok=False, client_id=None))) + assert not provider.is_protected("gdrive") + + +def test_factory_skips_second_source_sharing_token_store(): + """Two protected sources must not share one token-storage object store. + + NAT keys tokens per user only, so a shared store would let the sources + overwrite each other's credentials. The factory fails closed: the first + source claims the store, the second is skipped (left unconfigured). + """ + populate_from_config( + [ + { + "id": "gdrive", + "name": "Google Drive", + "requires_auth": True, + "per_user_auth": {"required": True, "mcp_server_id": "gdrive", "auth_provider": "mcp_oauth2_gdrive"}, + }, + { + "id": "notion", + "name": "Notion", + "requires_auth": True, + "per_user_auth": {"required": True, "mcp_server_id": "notion", "auth_provider": "mcp_oauth2_notion"}, + }, + ] + ) + # Both auth providers resolve to a store named "mcp_token_store" (the fake's + # default), so the second source collides with the first. + provider = build_from(_FakeBuilder(_FakeProvider())) + assert provider.is_protected("gdrive") + assert not provider.is_protected("notion") + + +def test_factory_ignores_unprotected_sources(): + populate_from_config([{"id": "web_search", "name": "Web", "description": "x"}]) + provider = build_from(_FakeBuilder(_FakeProvider())) + assert provider.settings_by_source == {} + + +def build_from(builder): + return asyncio.run(build_mcp_auth_provider(builder)) diff --git a/frontends/aiq_api/tests/test_mcp_auth_provider.py b/frontends/aiq_api/tests/test_mcp_auth_provider.py new file mode 100644 index 000000000..98176c5da --- /dev/null +++ b/frontends/aiq_api/tests/test_mcp_auth_provider.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from datetime import UTC +from datetime import datetime +from datetime import timedelta +from unittest.mock import AsyncMock + +import pytest + +from aiq_agent.auth import Principal +from aiq_api.mcp_auth.nat_provider import NatMcpAuthProvider +from aiq_api.mcp_auth.nat_provider import OAuthSourceSettings +from aiq_api.mcp_auth.nat_provider import _auth_result_from_token +from aiq_api.mcp_auth.provider import ProtectedSourceAuthProvider +from aiq_api.mcp_auth.provider import principal_user_id +from nat.authentication.token_storage import InMemoryTokenStorage + +PRINCIPAL = Principal(type="jwt", sub="user-1") + + +def _settings() -> OAuthSourceSettings: + return OAuthSourceSettings( + source_id="gdrive", + mcp_server_id="gdrive", + provider="google", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + client_id="client-123", + client_secret="secret", # pragma: allowlist secret + scopes=["https://www.googleapis.com/auth/drive.readonly"], + redirect_uri="https://aiq.example/v1/auth/mcp/gdrive/callback", + ) + + +@pytest.fixture +def provider() -> tuple[NatMcpAuthProvider, InMemoryTokenStorage]: + store = InMemoryTokenStorage() + prov = NatMcpAuthProvider( + settings_by_source={"gdrive": _settings()}, + token_storage_resolver=lambda _s: store, + ) + return prov, store + + +def test_satisfies_protocol(provider): + prov, _ = provider + assert isinstance(prov, ProtectedSourceAuthProvider) + + +def test_status_not_connected_initially(provider): + prov, _ = provider + state = asyncio.run(prov.get_status(PRINCIPAL, "gdrive")) + assert state.status == "not_connected" + + +def test_status_error_for_unconfigured_source(provider): + prov, _ = provider + state = asyncio.run(prov.get_status(PRINCIPAL, "not-configured")) + assert state.status == "error" and state.last_error + + +def test_start_auth_mints_provider_url_with_pkce_and_state(provider): + prov, _ = provider + challenge = asyncio.run(prov.start_auth(PRINCIPAL, "gdrive")) + assert challenge.auth_url.startswith("https://accounts.google.com/o/oauth2/v2/auth?") + assert f"state={challenge.state}" in challenge.auth_url + assert "code_challenge=" in challenge.auth_url + assert "client_id=client-123" in challenge.auth_url + assert challenge.state in prov._pending # pending flow registered for the callback + # Back-compat: with no resource configured, no resource indicator is added. + assert "resource=" not in challenge.auth_url + + +def test_start_auth_includes_resource_when_set(): + # Mirrors NAT's MCPOAuth2Provider, which adds the RFC 8707 resource indicator to + # the authorize request (authorization_kwargs={"resource": ...}). + from urllib.parse import parse_qs + from urllib.parse import urlparse + + settings = _settings().model_copy(update={"resource": "https://maas.example/maas/gdrive/mcp"}) + prov = NatMcpAuthProvider( + settings_by_source={"gdrive": settings}, + token_storage_resolver=lambda _s: InMemoryTokenStorage(), + ) + challenge = asyncio.run(prov.start_auth(PRINCIPAL, "gdrive")) + params = parse_qs(urlparse(challenge.auth_url).query) + assert params.get("resource") == ["https://maas.example/maas/gdrive/mcp"] + + +def test_connected_after_token_stored(provider): + prov, store = provider + expires = datetime.now(UTC) + timedelta(hours=1) + asyncio.run( + store.store( + principal_user_id(PRINCIPAL), + _auth_result_from_token({"access_token": "tok", "refresh_token": "r", "expires_at": expires.timestamp()}), + ) + ) + state = asyncio.run(prov.get_status(PRINCIPAL, "gdrive")) + assert state.status == "connected" + assert abs((state.expires_at - expires).total_seconds()) < 2 + + +def test_expired_token_reports_expired(provider): + prov, store = provider + past = datetime.now(UTC) - timedelta(minutes=5) + asyncio.run( + store.store( + principal_user_id(PRINCIPAL), + _auth_result_from_token({"access_token": "old", "expires_at": past.timestamp()}), + ) + ) + state = asyncio.run(prov.get_status(PRINCIPAL, "gdrive")) + assert state.status == "expired" + + +def test_require_connected_blocks_only_disconnected_protected(provider): + prov, _ = provider + blocked = asyncio.run(prov.require_connected(PRINCIPAL, ["gdrive", "web_search"])) + # web_search is not configured/protected -> ignored; gdrive blocked with an auth_url + assert [c.source_id for c in blocked] == ["gdrive"] + assert blocked[0].auth_url.startswith("https://accounts.google.com") + + +def test_complete_callback_exchanges_and_stores(provider, monkeypatch): + prov, store = provider + challenge = asyncio.run(prov.start_auth(PRINCIPAL, "gdrive")) + + # Stub the network token exchange on the pending flow's authlib client. + flow = prov._pending[challenge.state] + expires = datetime.now(UTC) + timedelta(hours=1) + flow.client.fetch_token = AsyncMock( + return_value={"access_token": "fresh", "refresh_token": "r", "expires_at": expires.timestamp()} + ) + flow.client.aclose = AsyncMock() + + completed = asyncio.run( + prov.complete_callback(challenge.state, f"https://aiq.example/cb?code=abc&state={challenge.state}") + ) + assert completed == "gdrive" + # Token now retrievable under the principal's key -> status connected + state = asyncio.run(prov.get_status(PRINCIPAL, "gdrive")) + assert state.status == "connected" + assert challenge.state not in prov._pending + + +def test_complete_callback_unknown_state_raises(provider): + prov, _ = provider + with pytest.raises(KeyError): + asyncio.run(prov.complete_callback("nope", "https://aiq.example/cb?code=x&state=nope")) + + +def test_prune_closes_expired_flow_client(): + # An abandoned challenge (started, never completed) must have its + # AsyncOAuth2Client closed when pruned, otherwise it leaks httpx connections. + clock = [datetime(2026, 1, 1, tzinfo=UTC)] + prov = NatMcpAuthProvider( + settings_by_source={"gdrive": _settings()}, + token_storage_resolver=lambda _s: InMemoryTokenStorage(), + challenge_ttl=timedelta(minutes=5), + now=lambda: clock[0], + ) + challenge = asyncio.run(prov.start_auth(PRINCIPAL, "gdrive")) + closer = AsyncMock() + prov._pending[challenge.state].client.aclose = closer + + # Advance past the TTL so the flow is expired, then trigger a prune. + clock[0] = clock[0] + timedelta(minutes=10) + asyncio.run(prov.start_auth(PRINCIPAL, "gdrive")) + + closer.assert_awaited_once() + assert challenge.state not in prov._pending diff --git a/frontends/aiq_api/tests/test_mcp_auth_routes.py b/frontends/aiq_api/tests/test_mcp_auth_routes.py new file mode 100644 index 000000000..8e770c60d --- /dev/null +++ b/frontends/aiq_api/tests/test_mcp_auth_routes.py @@ -0,0 +1,187 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from datetime import UTC +from datetime import datetime +from datetime import timedelta +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from aiq_agent.auth import Principal +from aiq_agent.common.data_source_registry import populate_from_config +from aiq_agent.common.data_source_registry import reset_registry +from aiq_api.mcp_auth.nat_provider import NatMcpAuthProvider +from aiq_api.mcp_auth.nat_provider import OAuthSourceSettings +from aiq_api.mcp_auth.nat_provider import _auth_result_from_token +from aiq_api.mcp_auth.provider import principal_user_id +from aiq_api.routes import auth as auth_routes +from nat.authentication.token_storage import InMemoryTokenStorage + +PRINCIPAL = Principal(type="jwt", sub="user-1") + + +def _settings() -> OAuthSourceSettings: + return OAuthSourceSettings( + source_id="gdrive", + mcp_server_id="gdrive", + provider="google", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + client_id="client-123", + client_secret="secret", # pragma: allowlist secret + scopes=["drive.readonly"], + redirect_uri="https://aiq.example/v1/auth/mcp/gdrive/callback", + ) + + +@pytest.fixture(autouse=True) +def registry(): + reset_registry() + populate_from_config( + [ + {"id": "web_search", "name": "Web Search", "description": "x"}, + { + "id": "gdrive", + "name": "Google Drive", + "description": "Drive", + "requires_auth": True, + "per_user_auth": {"required": True, "provider": "google", "mcp_server_id": "gdrive"}, + }, + ] + ) + yield + reset_registry() + + +@pytest.fixture +def client(monkeypatch): + store = InMemoryTokenStorage() + provider = NatMcpAuthProvider(settings_by_source={"gdrive": _settings()}, token_storage_resolver=lambda _s: store) + monkeypatch.setattr(auth_routes, "require_verified_principal", lambda: PRINCIPAL) + app = FastAPI() + auth_routes.register_mcp_auth_routes(app, provider) + return TestClient(app), provider, store + + +def test_status_not_connected_offers_connect_url(client): + tc, _, _ = client + resp = tc.get("/v1/auth/mcp/gdrive/status") + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "not_connected" + assert body["connect_url"] == "/v1/auth/mcp/gdrive/connect" + + +def test_status_unknown_source_404(client): + tc, _, _ = client + assert tc.get("/v1/auth/mcp/web_search/status").status_code == 404 # not a protected source + assert tc.get("/v1/auth/mcp/nope/status").status_code == 404 + + +def test_connect_returns_auth_url_structured(client): + tc, _, _ = client + resp = tc.post("/v1/auth/mcp/gdrive/connect") + assert resp.status_code == 200 + body = resp.json() + # The auth_url is returned in a structured response (not UI-only) for all clients. + assert body["status"] == "auth_required" + assert body["auth_url"].startswith("https://accounts.google.com/o/oauth2/v2/auth?") + + +def test_connect_when_already_connected(client): + tc, provider, store = client + import asyncio + + expires = datetime.now(UTC) + timedelta(hours=1) + asyncio.run( + store.store( + principal_user_id(PRINCIPAL), + _auth_result_from_token({"access_token": "t", "expires_at": expires.timestamp()}), + ) + ) + body = tc.post("/v1/auth/mcp/gdrive/connect").json() + assert body["status"] == "connected" and body.get("auth_url") is None + + +def test_connect_then_callback_completes(client): + tc, provider, store = client + auth_url = tc.post("/v1/auth/mcp/gdrive/connect").json()["auth_url"] + state = auth_url.split("state=")[1].split("&")[0] + + # Stub the token exchange on the pending flow. + flow = provider._pending[state] + expires = datetime.now(UTC) + timedelta(hours=1) + flow.client.fetch_token = AsyncMock( + return_value={"access_token": "fresh", "refresh_token": "r", "expires_at": expires.timestamp()} + ) + flow.client.aclose = AsyncMock() + + cb = tc.get(f"/v1/auth/mcp/gdrive/callback?code=abc&state={state}") + assert cb.status_code == 200 + assert "Connected" in cb.text + assert tc.get("/v1/auth/mcp/gdrive/status").json()["status"] == "connected" + + +def test_callback_unknown_state_renders_error(client): + tc, _, _ = client + cb = tc.get("/v1/auth/mcp/gdrive/callback?code=abc&state=bogus") + assert cb.status_code == 400 + assert "expired" in cb.text.lower() + + +def test_callback_provider_denied(client): + tc, _, _ = client + cb = tc.get("/v1/auth/mcp/gdrive/callback?error=access_denied&state=whatever") + assert cb.status_code == 400 + assert "denied" in cb.text.lower() + + +def test_callback_error_is_html_escaped(client): + """A provider-controlled `error` must not inject markup into the AIQ origin.""" + tc, _, _ = client + payload = "

" + cb = tc.get("/v1/auth/mcp/gdrive/callback", params={"error": payload, "state": "whatever"}) + assert cb.status_code == 400 + # The raw script tag must not appear; the escaped form must. + assert "" not in cb.text + assert "<script>" in cb.text + + +def test_preflight_blocks_disconnected_with_409(client, monkeypatch): + import asyncio + + from aiq_api.routes import jobs as jobs_routes + + _, provider, _ = client + result = asyncio.run(jobs_routes._preflight_mcp_auth(provider, PRINCIPAL, ["gdrive", "web_search"])) + assert result is not None and result.status_code == 409 + body = json.loads(result.body) + assert body["error"] == "mcp_auth_required" + assert len(body["sources"]) == 1 + src = body["sources"][0] + assert src["source_id"] == "gdrive" + assert src["connect_url"] == "/v1/auth/mcp/gdrive/connect" + assert src["auth_url"].startswith("https://accounts.google.com") + + +def test_preflight_allows_when_connected(client): + import asyncio + + from aiq_api.routes import jobs as jobs_routes + + _, provider, store = client + expires = datetime.now(UTC) + timedelta(hours=1) + asyncio.run( + store.store( + principal_user_id(PRINCIPAL), + _auth_result_from_token({"access_token": "t", "expires_at": expires.timestamp()}), + ) + ) + result = asyncio.run(jobs_routes._preflight_mcp_auth(provider, PRINCIPAL, ["gdrive"])) + assert result is None # connected -> no block diff --git a/frontends/aiq_api/tests/test_runtime_tools_token.py b/frontends/aiq_api/tests/test_runtime_tools_token.py new file mode 100644 index 000000000..adaef0e0a --- /dev/null +++ b/frontends/aiq_api/tests/test_runtime_tools_token.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from contextlib import AsyncExitStack +from datetime import UTC +from datetime import datetime +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from pydantic import SecretStr + +from aiq_api.mcp_auth import runtime_tools +from aiq_api.mcp_auth.runtime_tools import PerUserMcpSourceUnavailableError +from nat.authentication.token_storage import InMemoryTokenStorage +from nat.builder.context import ContextState +from nat.data_models.authentication import AuthResult +from nat.data_models.authentication import BearerTokenCred + +USER = "verified:alice" + + +def _auth(*, expired: bool) -> AuthResult: + delta = timedelta(hours=-1) if expired else timedelta(hours=1) + return AuthResult( + credentials=[BearerTokenCred(token=SecretStr("tok"))], + token_expires_at=datetime.now(UTC) + delta, + ) + + +async def _check(stored: AuthResult | None) -> tuple[bool, bool]: + """Returns (usable, token_still_present_after).""" + store = InMemoryTokenStorage() + ContextState.get().user_id.set(USER) + if stored is not None: + await store.store(USER, stored) + + async def _fake_resolve(builder, cfg, source_id): + return store + + # _token_usable does `from .factory import _resolve_token_storage`, so patch it there. + with patch("aiq_api.mcp_auth.factory._resolve_token_storage", _fake_resolve): + usable = await runtime_tools._token_usable(builder=None, cfg=None, source_id="gdrive") + still_there = (await store.retrieve(USER)) is not None + return usable, still_there + + +def test_valid_token_is_usable_and_kept(): + usable, still_there = asyncio.run(_check(_auth(expired=False))) + assert usable is True + assert still_there is True + + +def test_expired_token_is_not_usable_and_invalidated(): + # The core fix: an expired token must be reported unusable AND deleted, so the + # next get_status flips the card to Reconnect instead of false "connected". + usable, still_there = asyncio.run(_check(_auth(expired=True))) + assert usable is False + assert still_there is False + + +def test_missing_token_is_not_usable(): + usable, still_there = asyncio.run(_check(None)) + assert usable is False + assert still_there is False + + +def _gdrive_source(): + return SimpleNamespace( + id="gdrive", + per_user_auth=SimpleNamespace(required=True, auth_provider="mcp_oauth2_gdrive", tool_overrides={}), + ) + + +class _Builder: + async def get_auth_provider(self, name): + return SimpleNamespace(config=SimpleNamespace(server_url="https://mcp.example/mcp")) + + +async def _open(data_sources): + async with AsyncExitStack() as stack: + return await runtime_tools.open_per_user_mcp_tools( + builder=_Builder(), data_sources=data_sources, exit_stack=stack + ) + + +async def _not_usable(*_args, **_kwargs): + return False + + +def test_explicitly_selected_unusable_source_fails_closed(): + # A source the caller named specifically that we can't resolve (expired/missing + # token) must raise rather than silently answering without it. + with ( + patch("aiq_api.mcp_auth.runtime_tools.get_all_sources", return_value=[_gdrive_source()]), + patch("aiq_api.mcp_auth.runtime_tools._token_usable", _not_usable), + ): + with pytest.raises(PerUserMcpSourceUnavailableError) as exc: + asyncio.run(_open(["gdrive"])) + assert "gdrive" in str(exc.value) + + +def test_all_sources_mode_skips_unusable_without_raising(): + # When data_sources is None ("all"), stay best-effort: an unusable source is + # skipped, not fatal. + with ( + patch("aiq_api.mcp_auth.runtime_tools.get_all_sources", return_value=[_gdrive_source()]), + patch("aiq_api.mcp_auth.runtime_tools._token_usable", _not_usable), + ): + tools = asyncio.run(_open(None)) + assert tools == [] diff --git a/frontends/aiq_api/tests/test_sqlite_object_store.py b/frontends/aiq_api/tests/test_sqlite_object_store.py new file mode 100644 index 000000000..148f3ac62 --- /dev/null +++ b/frontends/aiq_api/tests/test_sqlite_object_store.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import stat +import sys + +import pytest + +from aiq_api.mcp_auth.sqlite_object_store import AiqSqliteObjectStore +from nat.data_models.object_store import KeyAlreadyExistsError +from nat.data_models.object_store import NoSuchKeyError +from nat.object_store.models import ObjectStoreItem + + +@pytest.fixture +def db_path(tmp_path): + return str(tmp_path / "tokens.db") + + +def _item(data: bytes = b'{"tok": 1}', **kw) -> ObjectStoreItem: + return ObjectStoreItem(data=data, content_type=kw.get("content_type"), metadata=kw.get("metadata")) + + +async def test_cross_connection_round_trip(db_path): + """A second connection (proxy for the worker process) reads what the first wrote. + + This is the property Redis exists for: the API process writes the token and a + *separate* Dask worker process reads it at job time. Two independent + connections to the same file stand in for those two processes. + """ + writer = AiqSqliteObjectStore(db_path, bucket_name="mcp-tokens") + await writer.put_object("alice", _item(b"secret", metadata={"u": "alice"})) + + reader = AiqSqliteObjectStore(db_path, bucket_name="mcp-tokens") + got = await reader.get_object("alice") + assert got.data == b"secret" + assert got.metadata == {"u": "alice"} + + await writer.aclose() + await reader.aclose() + + +async def test_put_rejects_existing_key(db_path): + store = AiqSqliteObjectStore(db_path) + await store.put_object("k", _item()) + with pytest.raises(KeyAlreadyExistsError): + await store.put_object("k", _item()) + await store.aclose() + + +async def test_upsert_overwrites(db_path): + store = AiqSqliteObjectStore(db_path) + await store.put_object("k", _item(b"v1")) + await store.upsert_object("k", _item(b"v2")) + assert (await store.get_object("k")).data == b"v2" + await store.aclose() + + +async def test_get_and_delete_missing_raise(db_path): + store = AiqSqliteObjectStore(db_path) + with pytest.raises(NoSuchKeyError): + await store.get_object("nope") + with pytest.raises(NoSuchKeyError): + await store.delete_object("nope") + await store.aclose() + + +async def test_delete_removes(db_path): + store = AiqSqliteObjectStore(db_path) + await store.put_object("k", _item()) + await store.delete_object("k") + with pytest.raises(NoSuchKeyError): + await store.get_object("k") + await store.aclose() + + +async def test_expired_object_is_absent(db_path): + # Negative TTL => already expired on write; reads must treat it as missing. + store = AiqSqliteObjectStore(db_path, ttl=-1) + await store.upsert_object("k", _item()) + with pytest.raises(NoSuchKeyError): + await store.get_object("k") + # And an expired key does not block a fresh put. + fresh = AiqSqliteObjectStore(db_path) + await fresh.put_object("k", _item(b"new")) + assert (await fresh.get_object("k")).data == b"new" + await store.aclose() + await fresh.aclose() + + +async def test_bucket_prefix_isolates_keys(db_path): + a = AiqSqliteObjectStore(db_path, bucket_name="a") + b = AiqSqliteObjectStore(db_path, bucket_name="b") + await a.put_object("same", _item(b"in-a")) + await b.put_object("same", _item(b"in-b")) + assert (await a.get_object("same")).data == b"in-a" + assert (await b.get_object("same")).data == b"in-b" + await a.aclose() + await b.aclose() + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX file modes not meaningful on Windows") +async def test_db_file_is_owner_only(db_path): + """The token DB (and its WAL/SHM sidecars) must not be group/world-readable. + + The stored objects contain plaintext access/refresh tokens, so a 0644 file + would let other local users read the credential store. + """ + store = AiqSqliteObjectStore(db_path, bucket_name="mcp-tokens") + # Force the connection (and thus file creation + WAL sidecars) to happen. + await store.put_object("alice", _item(b"secret")) + + for path in (db_path, db_path + "-wal", db_path + "-shm"): + if os.path.exists(path): + mode = stat.S_IMODE(os.stat(path).st_mode) + assert mode & 0o077 == 0, f"{path} is accessible to group/other: {oct(mode)}" + + await store.aclose() + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX file modes not meaningful on Windows") +async def test_existing_loose_db_is_tightened(db_path): + """An already-present 0644 DB file is chmod'd to 0600 on next open.""" + # Simulate a pre-existing world-readable file left by an older build. + with open(db_path, "wb"): + pass + os.chmod(db_path, 0o644) + + store = AiqSqliteObjectStore(db_path) + await store.put_object("k", _item()) + assert stat.S_IMODE(os.stat(db_path).st_mode) & 0o077 == 0 + await store.aclose() + + +async def test_registered_with_nat(): + """The store is discoverable by NAT under _type: aiq_sqlite.""" + from nat.runtime.loader import PluginTypes + from nat.runtime.loader import discover_and_register_plugins + + discover_and_register_plugins(PluginTypes.ALL) + from nat.cli.type_registry import GlobalTypeRegistry + + names = [getattr(i, "local_name", "") for i in GlobalTypeRegistry.get().get_registered_object_stores()] + assert "aiq_sqlite" in names diff --git a/frontends/aiq_api/tests/test_submit_mcp_auth_guard.py b/frontends/aiq_api/tests/test_submit_mcp_auth_guard.py new file mode 100644 index 000000000..bd87c2d5e --- /dev/null +++ b/frontends/aiq_api/tests/test_submit_mcp_auth_guard.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""submit_agent_job() must run the per-user MCP auth preflight before enqueue. + +This guards the programmatic submit path (e.g. the chat researcher's async +deep-research submit), which bypasses the REST route's 409 preflight. The check +lives in submit_agent_job so both paths share one chokepoint. +""" + +from __future__ import annotations + +import asyncio +from datetime import UTC +from datetime import datetime +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from aiq_agent.auth import Principal +from aiq_agent.common.data_source_registry import populate_from_config +from aiq_agent.common.data_source_registry import reset_registry +from aiq_api.jobs import submit as submit_mod +from aiq_api.mcp_auth import active as active_mod +from aiq_api.mcp_auth.nat_provider import NatMcpAuthProvider +from aiq_api.mcp_auth.nat_provider import OAuthSourceSettings +from aiq_api.mcp_auth.nat_provider import _auth_result_from_token +from aiq_api.mcp_auth.preflight import McpAuthRequiredError +from aiq_api.mcp_auth.provider import principal_user_id +from nat.authentication.token_storage import InMemoryTokenStorage + +PRINCIPAL = Principal(type="jwt", sub="user-1", email="u@example.com") + + +def _settings() -> OAuthSourceSettings: + return OAuthSourceSettings( + source_id="gdrive", + mcp_server_id="gdrive", + provider="google", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + client_id="client-123", + client_secret="secret", # pragma: allowlist secret + scopes=["drive.readonly"], + redirect_uri="https://aiq.example/v1/auth/mcp/gdrive/callback", + ) + + +class _FakeJobStore: + submitted = False + + def __init__(self, **_kwargs): + pass + + def ensure_job_id(self, job_id): + return job_id or "job-1" + + async def submit_job(self, *, job_id, expiry_seconds, job_fn, job_args): + _FakeJobStore.submitted = True + + +@pytest.fixture(autouse=True) +def registry(): + reset_registry() + populate_from_config( + [ + {"id": "web_search", "name": "Web Search", "description": "x"}, + { + "id": "gdrive", + "name": "Google Drive", + "description": "Drive", + "requires_auth": True, + "per_user_auth": {"required": True, "provider": "google", "mcp_server_id": "gdrive"}, + }, + ] + ) + yield + reset_registry() + + +@pytest.fixture +def patched(monkeypatch): + import nat.front_ends.fastapi.async_jobs.job_store as js_mod + + monkeypatch.setenv("NAT_DASK_SCHEDULER_ADDRESS", "tcp://localhost:8786") + monkeypatch.setattr(js_mod, "JobStore", _FakeJobStore) + monkeypatch.setattr( + submit_mod, + "get_agent_config", + lambda _t: SimpleNamespace(class_path="pkg.mod.Agent", config_name="deep_research_agent", public=True), + ) + monkeypatch.setattr(submit_mod, "create_job_access", MagicMock()) + _FakeJobStore.submitted = False + + store = InMemoryTokenStorage() + provider = NatMcpAuthProvider(settings_by_source={"gdrive": _settings()}, token_storage_resolver=lambda _s: store) + active_mod.set_active_mcp_auth_provider(provider) + yield store + active_mod.set_active_mcp_auth_provider(None) + + +def _submit(data_sources): + return asyncio.run( + submit_mod.submit_agent_job( + agent_type="deep_researcher", + input_text="query", + owner="u@example.com", + principal=PRINCIPAL, + data_sources=data_sources, + ) + ) + + +def test_blocks_when_protected_source_not_connected(patched): + with pytest.raises(McpAuthRequiredError) as exc: + _submit(["gdrive"]) + assert _FakeJobStore.submitted is False # never enqueued + assert [s.source_id for s in exc.value.response.sources] == ["gdrive"] + assert "gdrive" in str(exc.value) # user-facing message names the source + + +def test_blocks_when_data_sources_none_and_protected_disconnected(patched): + # data_sources=None means "any tool" -> every protected source must be connected. + with pytest.raises(McpAuthRequiredError): + _submit(None) + assert _FakeJobStore.submitted is False + + +def test_allows_unprotected_only_selection(patched): + job_id = _submit(["web_search"]) # no protected source selected + assert job_id == "job-1" + assert _FakeJobStore.submitted is True + + +def test_allows_when_protected_source_connected(patched): + store = patched + expires = datetime.now(UTC) + timedelta(hours=1) + asyncio.run( + store.store( + principal_user_id(PRINCIPAL), + _auth_result_from_token({"access_token": "t", "expires_at": expires.timestamp()}), + ) + ) + job_id = _submit(["gdrive"]) + assert job_id == "job-1" + assert _FakeJobStore.submitted is True + + +def test_no_active_provider_skips_guard(patched): + # When MCP auth is not configured in this process, there is nothing to enforce. + active_mod.set_active_mcp_auth_provider(None) + job_id = _submit(["gdrive"]) + assert job_id == "job-1" + assert _FakeJobStore.submitted is True diff --git a/frontends/aiq_api/tests/test_submit_owner_user_id.py b/frontends/aiq_api/tests/test_submit_owner_user_id.py new file mode 100644 index 000000000..1890892c9 --- /dev/null +++ b/frontends/aiq_api/tests/test_submit_owner_user_id.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""submit_agent_job must forward the owner's canonical user_id as the last job arg +so the worker can bind it on the NAT Context and per_user_mcp_client finds the token.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from aiq_agent.auth import Principal +from aiq_api.jobs import submit as submit_mod +from aiq_api.mcp_auth.provider import principal_user_id + + +class _FakeJobStore: + last_job_args = None + + def __init__(self, **_kwargs): + pass + + def ensure_job_id(self, job_id): + return job_id or "job-1" + + async def submit_job(self, *, job_id, expiry_seconds, job_fn, job_args): + _FakeJobStore.last_job_args = job_args + + +@pytest.fixture +def patched(monkeypatch): + import nat.front_ends.fastapi.async_jobs.job_store as js_mod + + monkeypatch.setenv("NAT_DASK_SCHEDULER_ADDRESS", "tcp://localhost:8786") + monkeypatch.setattr(js_mod, "JobStore", _FakeJobStore) + monkeypatch.setattr( + submit_mod, + "get_agent_config", + lambda _t: SimpleNamespace(class_path="pkg.mod.Agent", config_name="deep_research_agent", public=True), + ) + monkeypatch.setattr(submit_mod, "create_job_access", MagicMock()) + _FakeJobStore.last_job_args = None + + +def test_submit_forwards_owner_user_id_as_last_arg(patched): + principal = Principal(type="jwt", sub="user-1", email="u@example.com") + asyncio.run( + submit_mod.submit_agent_job( + agent_type="deep_researcher", + input_text="query", + owner="u@example.com", + principal=principal, + auth_token="token-1", + ) + ) + job_args = _FakeJobStore.last_job_args + assert job_args is not None + # Owner user_id is appended last. Trailing worker args are: + # data_sources, auth_token, initial_files, output_metadata, owner_user_id. + assert job_args[-1] == principal_user_id(principal) == "jwt:user-1" + assert job_args[-4] == "token-1" + + +def test_context_user_id_binding_mechanism(): + """The worker binds owner_user_id via ContextState.user_id; Context.user_id reads it. + + Guards the contract runner.run_agent_job relies on (and that per_user_mcp_client + reads via Context.get().user_id) against NAT-side changes. + """ + from nat.builder.context import Context + from nat.builder.context import ContextState + + token = ContextState.get().user_id.set("jwt:user-9") + try: + assert Context.get().user_id == "jwt:user-9" + finally: + ContextState.get().user_id.reset(token) diff --git a/frontends/ui/src/adapters/api/data-sources-client.ts b/frontends/ui/src/adapters/api/data-sources-client.ts index dae976482..1a221e1b5 100644 --- a/frontends/ui/src/adapters/api/data-sources-client.ts +++ b/frontends/ui/src/adapters/api/data-sources-client.ts @@ -19,6 +19,29 @@ const getBaseUrl = (): string => { // Types // ============================================================================ +/** Per-user MCP auth status for a protected source. */ +export type PerUserAuthStatus = 'connected' | 'not_connected' | 'expired' | 'error' + +/** Per-user MCP OAuth block attached to a protected data source (mirrors the API). */ +export interface PerUserAuthInfoFromAPI { + required: boolean + /** Auth mechanism (only mcp_oauth2 today) */ + type?: 'mcp_oauth2' + /** Provider identifier, e.g. 'google' */ + provider?: string | null + /** MCP server/auth-provider key */ + mcp_server_id?: string | null + status?: PerUserAuthStatus | null + /** Stable URL to (re)start the connect flow */ + connect_url?: string | null + /** Short-lived provider login URL (only present when an auth challenge was started) */ + auth_url?: string | null + /** Token expiry (ISO timestamp) */ + expires_at?: string | null + /** Last error detail, if status is 'error' */ + last_error?: string | null +} + export interface DataSourceFromAPI { /** Unique identifier for the data source */ id: string @@ -32,6 +55,8 @@ export interface DataSourceFromAPI { default_enabled?: boolean /** Whether the source requires user authentication */ requires_auth?: boolean + /** Per-user MCP OAuth state (present only for protected MCP sources) */ + per_user_auth?: PerUserAuthInfoFromAPI | null } export interface DataSourcesResponse { diff --git a/frontends/ui/src/adapters/api/index.ts b/frontends/ui/src/adapters/api/index.ts index 8a803c9d4..0300b8d3b 100644 --- a/frontends/ui/src/adapters/api/index.ts +++ b/frontends/ui/src/adapters/api/index.ts @@ -72,8 +72,20 @@ export type { DataSourcesClientOptions, DataSourceFromAPI, DataSourcesResponse, + PerUserAuthInfoFromAPI, + PerUserAuthStatus, } from './data-sources-client' +// MCP Auth Client (per-user OAuth control plane) +export { createMcpAuthClient, openAuthPopupAndWait } from './mcp-auth-client' +export type { + McpAuthClient, + McpAuthClientOptions, + SourceAuthStatusResponse, + SourceConnectResponse, + AuthPopupResult, +} from './mcp-auth-client' + // Documents Schemas export { DocumentFileStatusSchema, diff --git a/frontends/ui/src/adapters/api/mcp-auth-client.spec.ts b/frontends/ui/src/adapters/api/mcp-auth-client.spec.ts new file mode 100644 index 000000000..430c84b13 --- /dev/null +++ b/frontends/ui/src/adapters/api/mcp-auth-client.spec.ts @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, test, expect, vi, beforeEach, afterEach } from 'vitest' +import { openAuthPopupAndWait } from './mcp-auth-client' +import type { PerUserAuthStatus } from './data-sources-client' + +describe('openAuthPopupAndWait', () => { + let fakePopup: { closed: boolean; close: ReturnType } + let openSpy: ReturnType + + beforeEach(() => { + vi.useFakeTimers() + fakePopup = { closed: false, close: vi.fn(() => { fakePopup.closed = true }) } + openSpy = vi.fn(() => fakePopup) + vi.stubGlobal('open', openSpy) + }) + + afterEach(() => { + vi.useRealTimers() + vi.unstubAllGlobals() + vi.clearAllMocks() + }) + + // The core regression: clicking "Reconnect" starts from an `expired` status, so + // the first status probe sees `expired`. The popup must stay open through that + // baseline and only resolve once auth actually completes (`connected`). + test('does not close on the pre-existing expired status (Reconnect)', async () => { + const statuses: PerUserAuthStatus[] = ['expired', 'expired', 'connected'] + let i = 0 + const pollStatus = vi.fn(async () => statuses[Math.min(i++, statuses.length - 1)]) + + const result = openAuthPopupAndWait('https://provider/auth', 'gdrive', { + pollStatus, + pollIntervalMs: 1000, + }) + + // First two probes return the baseline `expired` — popup must remain open. + await vi.advanceTimersByTimeAsync(1000) + expect(fakePopup.close).not.toHaveBeenCalled() + await vi.advanceTimersByTimeAsync(1000) + expect(fakePopup.close).not.toHaveBeenCalled() + + // Third probe flips to `connected` → resolve as success and close the popup. + await vi.advanceTimersByTimeAsync(1000) + await expect(result).resolves.toEqual({ ok: true, sourceId: 'gdrive' }) + expect(fakePopup.close).toHaveBeenCalled() + }) + + test('resolves failure only on a NEW error, not the pre-existing one', async () => { + // Baseline already `error`; staying `error` must not resolve (it's the state + // we're trying to fix). A transition expired -> error would, but unchanged + // error should keep waiting until the popup closes / times out. + const pollStatus = vi.fn(async (): Promise => 'error') + const result = openAuthPopupAndWait('https://provider/auth', 'gdrive', { + pollStatus, + pollIntervalMs: 1000, + }) + + await vi.advanceTimersByTimeAsync(3000) + expect(fakePopup.close).not.toHaveBeenCalled() + + // User gives up and closes the popup → resolve empty (caller re-checks status). + fakePopup.closed = true + await vi.advanceTimersByTimeAsync(700) + await expect(result).resolves.toEqual({}) + }) + + test('resolves success via callback postMessage from the popup', async () => { + const result = openAuthPopupAndWait('https://provider/auth', 'gdrive', {}) + // A genuine completion is posted by the callback page via window.opener, so + // event.source is the popup we opened. + const evt = new MessageEvent('message', { data: { type: 'mcp-auth', source_id: 'gdrive', ok: true } }) + Object.defineProperty(evt, 'source', { value: fakePopup }) + window.dispatchEvent(evt) + await expect(result).resolves.toEqual({ ok: true, sourceId: 'gdrive' }) + expect(fakePopup.close).toHaveBeenCalled() + }) + + test('ignores an mcp-auth message from an untrusted source', async () => { + const result = openAuthPopupAndWait('https://provider/auth', 'gdrive', {}) + // No event.source (a spoofed / cross-window message) must NOT be accepted. + window.dispatchEvent(new MessageEvent('message', { data: { type: 'mcp-auth', source_id: 'gdrive', ok: true } })) + // The spoofed message is ignored; the promise only settles when the popup + // actually closes (empty result), and the popup is never closed by the spoof. + fakePopup.closed = true + await vi.advanceTimersByTimeAsync(700) + await expect(result).resolves.toEqual({}) + expect(fakePopup.close).not.toHaveBeenCalled() + }) + + test('resolves empty when the popup is blocked', async () => { + openSpy.mockReturnValueOnce(null) + await expect(openAuthPopupAndWait('https://provider/auth', 'gdrive', {})).resolves.toEqual({}) + }) +}) diff --git a/frontends/ui/src/adapters/api/mcp-auth-client.ts b/frontends/ui/src/adapters/api/mcp-auth-client.ts new file mode 100644 index 000000000..6e5f81c57 --- /dev/null +++ b/frontends/ui/src/adapters/api/mcp-auth-client.ts @@ -0,0 +1,218 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +/** + * MCP Auth API Client + * + * Per-user MCP OAuth control plane: read a protected source's connection + * status, start the connect flow (returns a provider login URL), and a helper + * to open that URL in a popup and resolve once the popup closes or the callback + * page posts back. + */ + +import { apiConfig } from './config' +import type { PerUserAuthStatus } from './data-sources-client' + +const getBaseUrl = (): string => { + const isBrowser = typeof window !== 'undefined' + return isBrowser ? '' : apiConfig.baseUrl +} + +const apiUrl = (path: string): string => { + const baseUrl = getBaseUrl() + return baseUrl ? `${baseUrl}${path}` : `/api${path}` +} + +// ============================================================================ +// Types +// ============================================================================ + +export interface SourceAuthStatusResponse { + source_id: string + status: PerUserAuthStatus + expires_at?: string | null + connect_url?: string | null + last_error?: string | null +} + +export interface SourceConnectResponse { + source_id: string + status: 'auth_required' | 'connected' + auth_url?: string | null + expires_at?: string | null +} + +export interface McpAuthClientOptions { + authToken?: string +} + +// ============================================================================ +// Client Factory +// ============================================================================ + +export const createMcpAuthClient = (options: McpAuthClientOptions = {}) => { + const { authToken } = options + + const getHeaders = (): Record => { + const headers: Record = { 'Content-Type': 'application/json' } + if (authToken) { + headers['Authorization'] = `Bearer ${authToken}` + } + return headers + } + + return { + /** Read the current per-user auth status for a protected source. */ + async getStatus(sourceId: string, signal?: AbortSignal): Promise { + const response = await fetch(apiUrl(`/v1/auth/mcp/${encodeURIComponent(sourceId)}/status`), { + method: 'GET', + headers: getHeaders(), + signal, + }) + if (!response.ok) { + throw new Error(`Failed to fetch auth status: ${response.statusText}`) + } + return response.json() + }, + + /** Start (or resume) the OAuth flow; returns a provider login URL to open. */ + async connect(sourceId: string): Promise { + const response = await fetch(apiUrl(`/v1/auth/mcp/${encodeURIComponent(sourceId)}/connect`), { + method: 'POST', + headers: getHeaders(), + }) + if (!response.ok) { + throw new Error(`Failed to start connection: ${response.statusText}`) + } + return response.json() + }, + } +} + +export type McpAuthClient = ReturnType + +// ============================================================================ +// Popup helper +// ============================================================================ + +export interface AuthPopupResult { + /** True if the callback page reported success via postMessage. Undefined if we + * only observed the window closing (caller should re-check status). */ + ok?: boolean + /** The source id reported by the callback page, if any. */ + sourceId?: string +} + +export interface AuthPopupOptions { + /** + * Optional backend status probe. When provided, the popup also resolves as + * soon as the source reports a terminal status. This is the reliable signal: + * the provider's pages typically send `Cross-Origin-Opener-Policy`, which + * severs `window.opener` so the callback's `postMessage` and the parent's + * `popup.closed` check both go silent — leaving the card stuck on + * "Connecting…". The callback persists the token before it messages the + * opener, so a status probe sees the result regardless. + */ + pollStatus?: () => Promise + /** How often to probe backend status, in ms. Default 1500. */ + pollIntervalMs?: number + /** Stop waiting after this long, in ms, so a never-finished login can't poll + * forever. Default 180000 (3 min). */ + timeoutMs?: number +} + +/** + * Open the provider login URL in a popup and resolve when it closes, the + * callback page posts back, or (when `pollStatus` is supplied) the backend + * reports a terminal status. The caller should still re-fetch the source status + * after this resolves to confirm the connection. + */ +export function openAuthPopupAndWait( + authUrl: string, + sourceId: string, + options: AuthPopupOptions = {} +): Promise { + const { pollStatus, pollIntervalMs = 1500, timeoutMs = 180_000 } = options + return new Promise((resolve) => { + const popup = window.open(authUrl, `mcp-auth-${sourceId}`, 'popup,width=520,height=680') + + // Popup blocked — fall back to a same-tab redirect is too disruptive, so + // resolve immediately and let the caller surface the URL / re-check status. + if (!popup) { + resolve({}) + return + } + + let settled = false + const finish = (result: AuthPopupResult) => { + if (settled) return + settled = true + window.removeEventListener('message', onMessage) + clearInterval(poll) + if (statusPoll !== undefined) clearInterval(statusPoll) + clearTimeout(timeout) + resolve(result) + } + + const onMessage = (event: MessageEvent) => { + // Only trust messages from the popup we opened: the callback page posts via + // window.opener.postMessage, so a genuine completion has event.source === popup. + // Reject anything else (a synthetic or cross-window message) so an unrelated + // page can't spoof an auth-complete. If COOP severs the opener relationship + // event.source won't match — the status poll below is the fallback for that. + if (event.source !== popup) return + const data = event.data + if (data && data.type === 'mcp-auth' && data.source_id === sourceId) { + try { + popup.close() + } catch { + /* ignore */ + } + finish({ ok: !!data.ok, sourceId }) + } + } + window.addEventListener('message', onMessage) + + const poll = setInterval(() => { + if (popup.closed) { + finish({}) + } + }, 700) + + // Authoritative resolve path: poll the backend until the source reaches a + // terminal status. Survives the COOP opener-severing described above. + // + // `expired`/`error` are the PRE-EXISTING states that make the card show + // "Reconnect", so they're also what the first poll sees before the user has + // authenticated. Treating them as terminal would close the popup ~1.5s after + // it opens (UI: "Session expired") before login can complete. So we snapshot + // the baseline on the first probe and only resolve on a real transition: + // `connected` (success), or a NEW `error` that differs from the baseline. + let baselineStatus: PerUserAuthStatus | undefined + const statusPoll: ReturnType | undefined = pollStatus + ? setInterval(() => { + void (async () => { + let status: PerUserAuthStatus | undefined + try { + status = await pollStatus() + } catch { + return // transient probe failure — keep polling + } + if (baselineStatus === undefined) baselineStatus = status + const succeeded = status === 'connected' + const newlyErrored = status === 'error' && baselineStatus !== 'error' + if (succeeded || newlyErrored) { + try { + popup.close() + } catch { + /* ignore */ + } + finish({ ok: succeeded, sourceId }) + } + })() + }, pollIntervalMs) + : undefined + + const timeout = setTimeout(() => finish({}), timeoutMs) + }) +} diff --git a/frontends/ui/src/app/api/jobs/async/[...path]/route.ts b/frontends/ui/src/app/api/jobs/async/[...path]/route.ts index 50c4eecec..3d0b34802 100644 --- a/frontends/ui/src/app/api/jobs/async/[...path]/route.ts +++ b/frontends/ui/src/app/api/jobs/async/[...path]/route.ts @@ -42,6 +42,27 @@ const buildBackendUrl = (path: string[]): string => { return `${backendBase}/v1/jobs/async/${pathString}` } +/** + * Forward a non-2xx backend response to the browser. + * + * If the backend returned a JSON body (e.g. the 409 `mcp_auth_required` payload + * with its `sources`/`auth_url` details), pass it through verbatim with the + * original status so the client can act on the structure. Only fall back to the + * BACKEND_ERROR envelope for non-JSON error bodies, where there is nothing + * structured to preserve. + */ +const forwardBackendError = (status: number, errorText: string): NextResponse => { + try { + const parsed = JSON.parse(errorText) + return NextResponse.json(parsed, { status }) + } catch { + return NextResponse.json( + { error: { code: 'BACKEND_ERROR', message: `Backend returned ${status}: ${errorText}` } }, + { status } + ) + } +} + /** * Get auth headers from request, including idToken cookie. * Returns empty object when REQUIRE_AUTH=false to prevent user identification. @@ -117,19 +138,7 @@ export async function GET( if (!response.ok) { const errorText = await response.text() console.error('[Deep Research API] Backend error:', response.status, errorText) - - return new NextResponse( - JSON.stringify({ - error: { - code: 'BACKEND_ERROR', - message: `Backend returned ${response.status}: ${errorText}`, - }, - }), - { - status: response.status, - headers: { 'Content-Type': 'application/json' }, - } - ) + return forwardBackendError(response.status, errorText) } // For SSE streams, pass through the response body @@ -245,19 +254,7 @@ export async function POST( if (!response.ok) { const errorText = await response.text() console.error('[Deep Research API] Backend error:', response.status, errorText) - - return new NextResponse( - JSON.stringify({ - error: { - code: 'BACKEND_ERROR', - message: `Backend returned ${response.status}: ${errorText}`, - }, - }), - { - status: response.status, - headers: { 'Content-Type': 'application/json' }, - } - ) + return forwardBackendError(response.status, errorText) } // Return JSON response diff --git a/frontends/ui/src/features/chat/store.ts b/frontends/ui/src/features/chat/store.ts index 2e8ef16c4..9a617d6b2 100644 --- a/frontends/ui/src/features/chat/store.ts +++ b/frontends/ui/src/features/chat/store.ts @@ -378,17 +378,32 @@ const patchConversationMessageById = ( return didPatch ? { ...conversation, messages, updatedAt: new Date() } : conversation } +/** + * A protected per-user source (e.g. Google Drive) must be connected before it + * can be part of the active selection — otherwise it appears in "Selected Data + * Sources" and is submitted while unusable. Mirrors the card toggle, "Enable + * All", and the initial-fetch gate in the layout store. + */ +const isSelectableDataSource = (source: { + per_user_auth?: { required?: boolean; status?: string | null } | null +}): boolean => !(source.per_user_auth?.required && source.per_user_auth.status !== 'connected') + const getDefaultEnabledDataSourceIds = (): string[] => { const layoutStore = useLayoutStore.getState() - return layoutStore.availableDataSources?.map((source) => source.id) ?? [] + return (layoutStore.availableDataSources ?? []).filter(isSelectableDataSource).map((source) => source.id) } const restoreConversationDataSources = (conversation: Conversation): void => { const layoutStore = useLayoutStore.getState() if (conversation.enabledDataSourceIds) { - const availableIds = new Set(layoutStore.availableDataSources?.map((source) => source.id) ?? []) - const validIds = conversation.enabledDataSourceIds.filter((id) => availableIds.has(id)) + // Only restore sources that are still available AND currently selectable — + // a protected source saved as enabled must not come back while it's not + // connected (e.g. an old session that had Google Drive on). + const selectableIds = new Set( + (layoutStore.availableDataSources ?? []).filter(isSelectableDataSource).map((source) => source.id) + ) + const validIds = conversation.enabledDataSourceIds.filter((id) => selectableIds.has(id)) layoutStore.setEnabledDataSources(validIds) return } diff --git a/frontends/ui/src/features/layout/components/DataConnectionCard.tsx b/frontends/ui/src/features/layout/components/DataConnectionCard.tsx index 254cfa53f..076a0fca2 100644 --- a/frontends/ui/src/features/layout/components/DataConnectionCard.tsx +++ b/frontends/ui/src/features/layout/components/DataConnectionCard.tsx @@ -10,8 +10,8 @@ 'use client' -import { type FC } from 'react' -import { Flex, Text, Switch } from '@/adapters/ui' +import { type FC, useCallback, useState } from 'react' +import { Flex, Text, Switch, Button } from '@/adapters/ui' import { Globe } from '@/adapters/ui/icons' import type { DataSource } from '../data-sources' @@ -28,6 +28,16 @@ interface DataConnectionCardProps { unavailableReason?: string /** Callback when toggle state changes */ onToggle: (id: string, enabled: boolean) => void + /** Start the per-user OAuth connect flow for a protected source */ + onConnect?: (id: string) => void | Promise +} + +/** Human-readable status line for a protected MCP source. */ +const STATUS_LABELS: Record = { + connected: 'Connected', + not_connected: 'Not connected', + expired: 'Session expired', + error: 'Connection error', } /** @@ -41,16 +51,38 @@ export const DataConnectionCard: FC = ({ isBusy = false, unavailableReason, onToggle, + onConnect, }) => { - // Combine availability and busy state + const [connecting, setConnecting] = useState(false) + + // Per-user MCP auth state (present only for protected sources). + const perUserAuth = source.perUserAuth + const isProtected = !!perUserAuth?.required + const authStatus = perUserAuth?.status ?? undefined + // A protected source must be connected before it can be enabled. + const needsConnect = isProtected && authStatus !== 'connected' + + // Combine availability and busy state. A protected-but-unconnected source + // cannot be toggled — the Connect action replaces the switch. const isDisabled = !isAvailable || isBusy + const canToggle = !isDisabled && !needsConnect const handleToggle = () => { - if (!isDisabled) { + if (canToggle) { onToggle(source.id, !isEnabled) } } + const handleConnect = useCallback(async () => { + if (connecting || !onConnect) return + setConnecting(true) + try { + await onConnect(source.id) + } finally { + setConnecting(false) + } + }, [connecting, onConnect, source.id]) + const handleCardClick = () => { handleToggle() } @@ -72,15 +104,19 @@ export const DataConnectionCard: FC = ({ align="center" justify="between" role="button" - tabIndex={isDisabled ? -1 : 0} + tabIndex={canToggle ? 0 : -1} onClick={handleCardClick} onKeyDown={handleCardKeyDown} className={`border-base rounded-lg border p-3 transition-colors ${ - isDisabled ? 'cursor-not-allowed opacity-50' : 'cursor-pointer hover:bg-surface-raised-50' + isDisabled + ? 'cursor-not-allowed opacity-50' + : canToggle + ? 'cursor-pointer hover:bg-surface-raised-50' + : '' }`} aria-pressed={isEnabled} - aria-disabled={isDisabled} - aria-label={`${source.name}: ${isEnabled ? 'enabled' : 'disabled'}${isDisabled ? ' (disabled)' : ''}`} + aria-disabled={!canToggle} + aria-label={`${source.name}: ${isEnabled ? 'enabled' : 'disabled'}${!canToggle ? ' (disabled)' : ''}`} title={ isBusy ? 'Data source changes disabled during active operations' @@ -106,17 +142,43 @@ export const DataConnectionCard: FC = ({ {source.description} + {isProtected && ( + + {perUserAuth?.lastError && authStatus === 'error' + ? perUserAuth.lastError + : STATUS_LABELS[authStatus ?? 'not_connected']} + + )} {/* eslint-disable-next-line jsx-a11y/click-events-have-key-events, jsx-a11y/no-static-element-interactions */}

- + {needsConnect ? ( + + ) : ( + + )}
) diff --git a/frontends/ui/src/features/layout/components/DataSourcesPanel.spec.tsx b/frontends/ui/src/features/layout/components/DataSourcesPanel.spec.tsx index 40b14a8e7..3c24538c2 100644 --- a/frontends/ui/src/features/layout/components/DataSourcesPanel.spec.tsx +++ b/frontends/ui/src/features/layout/components/DataSourcesPanel.spec.tsx @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -import { render, screen } from '@/test-utils' +import { render, screen, waitFor } from '@/test-utils' import userEvent from '@testing-library/user-event' import { vi, describe, test, expect, beforeEach } from 'vitest' import { DataSourcesPanel } from './DataSourcesPanel' @@ -13,6 +13,7 @@ const mockSetDataSourcesPanelTab = vi.fn() const mockToggleDataSource = vi.fn() const mockSetEnabledDataSources = vi.fn() const mockFetchDataSources = vi.fn() +const mockRefreshDataSourceStatus = vi.fn() const mockDataSources = [ { id: 'web_search', name: 'Web Search', description: 'Search the web', requires_auth: false }, @@ -35,6 +36,7 @@ vi.mock('../store', () => ({ dataSourcesLoading: false, dataSourcesError: null, fetchDataSources: mockFetchDataSources, + refreshDataSourceStatus: mockRefreshDataSourceStatus, } return selector ? selector(state) : state }), @@ -49,19 +51,34 @@ vi.mock('@/adapters/auth', () => ({ })), })) -// Mock child components +// Mock the MCP auth client + popup so we can drive the connect flow. +const mockConnect = vi.fn() +const mockGetStatus = vi.fn() +const mockOpenAuthPopupAndWait = vi.fn() +vi.mock('@/adapters/api', async (importOriginal) => ({ + ...(await importOriginal()), + createMcpAuthClient: vi.fn(() => ({ connect: mockConnect, getStatus: mockGetStatus })), + openAuthPopupAndWait: (...args: unknown[]) => mockOpenAuthPopupAndWait(...args), +})) + +// Mock child components. The connect button lets tests trigger handleConnect. vi.mock('./DataConnectionCard', () => ({ DataConnectionCard: ({ source, isEnabled, isAvailable, + onConnect, }: { source: { id: string; name: string } isEnabled: boolean isAvailable: boolean + onConnect?: (id: string) => void }) => (
{source.name} - {isEnabled ? 'enabled' : 'disabled'} - {isAvailable ? 'available' : 'unavailable'} +
), })) @@ -95,6 +112,7 @@ describe('DataSourcesPanel', () => { dataSourcesLoading: false, dataSourcesError: null, fetchDataSources: mockFetchDataSources, + refreshDataSourceStatus: mockRefreshDataSourceStatus, } return selector ? selector(state) : state }) @@ -151,6 +169,7 @@ describe('DataSourcesPanel', () => { dataSourcesLoading: false, dataSourcesError: null, fetchDataSources: mockFetchDataSources, + refreshDataSourceStatus: mockRefreshDataSourceStatus, } return selector ? selector(state) : state }) @@ -181,6 +200,7 @@ describe('DataSourcesPanel', () => { dataSourcesLoading: false, dataSourcesError: null, fetchDataSources: mockFetchDataSources, + refreshDataSourceStatus: mockRefreshDataSourceStatus, } return selector ? selector(state) : state }) @@ -205,6 +225,7 @@ describe('DataSourcesPanel', () => { dataSourcesLoading: false, dataSourcesError: null, fetchDataSources: mockFetchDataSources, + refreshDataSourceStatus: mockRefreshDataSourceStatus, } return selector ? selector(state) : state }) @@ -232,6 +253,43 @@ describe('DataSourcesPanel', () => { expect(mockSetEnabledDataSources).toHaveBeenCalled() }) + test('enable all excludes protected sources that are not connected', async () => { + const user = userEvent.setup() + vi.mocked(useLayoutStore).mockImplementation((selector?: (s: any) => any) => { + const state = { + rightPanel: 'data-sources', + closeRightPanel: mockCloseRightPanel, + openRightPanel: mockOpenRightPanel, + dataSourcesPanelTab: 'connections', + setDataSourcesPanelTab: mockSetDataSourcesPanelTab, + enabledDataSourceIds: [], // nothing enabled -> click enables all + toggleDataSource: mockToggleDataSource, + setEnabledDataSources: mockSetEnabledDataSources, + availableDataSources: [ + { id: 'web_search', name: 'Web Search', description: 'Search', requires_auth: false }, + { + id: 'gdrive', + name: 'Google Drive', + description: 'Drive', + requires_auth: true, + per_user_auth: { required: true, status: 'not_connected' }, + }, + ], + dataSourcesLoading: false, + dataSourcesError: null, + fetchDataSources: mockFetchDataSources, + refreshDataSourceStatus: mockRefreshDataSourceStatus, + } + return selector ? selector(state) : state + }) + + render() + await user.click(screen.getByRole('button', { name: /all available connections/i })) + + // gdrive is protected + not connected, so it must be excluded from bulk enable. + expect(mockSetEnabledDataSources).toHaveBeenCalledWith(['web_search']) + }) + test('shows correct enabled state for data connection cards', () => { render() @@ -258,6 +316,7 @@ describe('DataSourcesPanel', () => { dataSourcesLoading: false, dataSourcesError: 'Failed to connect to server', fetchDataSources: mockFetchDataSources, + refreshDataSourceStatus: mockRefreshDataSourceStatus, } return selector ? selector(state) : state }) @@ -283,6 +342,7 @@ describe('DataSourcesPanel', () => { dataSourcesLoading: false, dataSourcesError: 'Network error', fetchDataSources: mockFetchDataSources, + refreshDataSourceStatus: mockRefreshDataSourceStatus, } return selector ? selector(state) : state }) @@ -331,6 +391,7 @@ describe('DataSourcesPanel', () => { dataSourcesLoading: false, dataSourcesError: null, fetchDataSources: mockFetchDataSources, + refreshDataSourceStatus: mockRefreshDataSourceStatus, } return selector ? selector(state) : state }) @@ -385,4 +446,49 @@ describe('DataSourcesPanel', () => { expect(screen.getByTestId('connection-card-bug_tracker')).toHaveTextContent('available') }) }) + + describe('connect flow (handleConnect)', () => { + test('opens the auth popup when the source requires auth', async () => { + const user = userEvent.setup() + mockConnect.mockResolvedValue({ status: 'auth_required', auth_url: 'https://provider/oauth' }) + mockOpenAuthPopupAndWait.mockResolvedValue({ ok: true }) + + render() + await user.click(screen.getByTestId('connect-knowledge_base')) + + await waitFor(() => expect(mockOpenAuthPopupAndWait).toHaveBeenCalled()) + expect(mockOpenAuthPopupAndWait).toHaveBeenCalledWith( + 'https://provider/oauth', + 'knowledge_base', + expect.objectContaining({ pollStatus: expect.any(Function) }) + ) + // finally: always refresh statuses after the attempt + await waitFor(() => expect(mockFetchDataSources).toHaveBeenCalledWith('valid-token')) + }) + + test('does not open the popup when no auth is required, but still refreshes', async () => { + const user = userEvent.setup() + mockConnect.mockResolvedValue({ status: 'connected', auth_url: null }) + + render() + await user.click(screen.getByTestId('connect-knowledge_base')) + + await waitFor(() => expect(mockFetchDataSources).toHaveBeenCalledWith('valid-token')) + expect(mockOpenAuthPopupAndWait).not.toHaveBeenCalled() + }) + + test('shows a failure banner and still refreshes when connect throws', async () => { + const user = userEvent.setup() + mockConnect.mockRejectedValue(new Error('network down')) + + render() + await user.click(screen.getByTestId('connect-knowledge_base')) + + // connectError banner surfaces the source name + error detail + expect(await screen.findByText(/Couldn't connect Knowledge Base\. network down/)).toBeInTheDocument() + // finally still runs despite the failure + await waitFor(() => expect(mockFetchDataSources).toHaveBeenCalledWith('valid-token')) + expect(mockOpenAuthPopupAndWait).not.toHaveBeenCalled() + }) + }) }) diff --git a/frontends/ui/src/features/layout/components/DataSourcesPanel.tsx b/frontends/ui/src/features/layout/components/DataSourcesPanel.tsx index 5eb37787c..da7131de7 100644 --- a/frontends/ui/src/features/layout/components/DataSourcesPanel.tsx +++ b/frontends/ui/src/features/layout/components/DataSourcesPanel.tsx @@ -10,8 +10,9 @@ 'use client' -import { type FC, memo, useCallback, useMemo } from 'react' +import { type FC, memo, useCallback, useEffect, useMemo, useState } from 'react' import { Flex, Text, SidePanel, SegmentedControl, Switch, Button, Banner } from '@/adapters/ui' +import { createMcpAuthClient, openAuthPopupAndWait } from '@/adapters/api' import { useShallow } from 'zustand/react/shallow' import { Globe, LoadingSpinner } from '@/adapters/ui/icons' import { useAuth } from '@/adapters/auth' @@ -61,10 +62,24 @@ export const DataSourcesPanel: FC = memo(function DataSou const toggleDataSource = useLayoutStore((s) => s.toggleDataSource) const setEnabledDataSources = useLayoutStore((s) => s.setEnabledDataSources) const fetchDataSources = useLayoutStore((s) => s.fetchDataSources) + const refreshDataSourceStatus = useLayoutStore((s) => s.refreshDataSourceStatus) + + // Refresh per-source auth status each time the panel opens so a token that was + // invalidated server-side (e.g. expired and dropped at job time) shows as + // Reconnect instead of a stale "connected". Selection-preserving, so it won't + // reset the user's enabled sources. + useEffect(() => { + if (isOpen) { + void refreshDataSourceStatus(idToken) + } + }, [isOpen, idToken, refreshDataSourceStatus]) // Check if current session is busy with operations const isBusy = useIsCurrentSessionBusy() + // Error surfaced when a protected-source connect attempt fails. + const [connectError, setConnectError] = useState(null) + // Check if user has valid auth token const hasValidToken = !!idToken @@ -81,8 +96,19 @@ export const DataSourcesPanel: FC = memo(function DataSou name: source.name, description: source.description ?? '', category: source.category ?? 'enterprise', - defaultEnabled: true, + defaultEnabled: source.default_enabled ?? true, requiresAuth: source.requires_auth ?? false, + perUserAuth: source.per_user_auth + ? { + required: source.per_user_auth.required, + provider: source.per_user_auth.provider, + mcpServerId: source.per_user_auth.mcp_server_id, + status: source.per_user_auth.status, + connectUrl: source.per_user_auth.connect_url, + expiresAt: source.per_user_auth.expires_at, + lastError: source.per_user_auth.last_error, + } + : undefined, })) }, [availableDataSources]) @@ -114,6 +140,40 @@ export const DataSourcesPanel: FC = memo(function DataSou [toggleDataSource, enabledDataSourceIds, saveDataSourcesToConversation, onSourceToggle] ) + // Start (or resume) the per-user OAuth flow for a protected source: get a + // provider login URL, open it in a popup, then refresh statuses. On failure + // (network/popup errors) surface a banner and still resync in `finally`, so + // the card never sticks in "Connecting…". + const handleConnect = useCallback( + async (sourceId: string) => { + setConnectError(null) + const sourceName = displaySources.find((s) => s.id === sourceId)?.name ?? sourceId + try { + const client = createMcpAuthClient({ authToken: idToken }) + const { status, auth_url } = await client.connect(sourceId) + if (status === 'auth_required' && auth_url) { + // Pass a status probe so the popup resolves once the backend records + // the connection, even when the provider's COOP headers sever the + // opener and the postMessage / popup-close signals never arrive. + await openAuthPopupAndWait(auth_url, sourceId, { + pollStatus: () => client.getStatus(sourceId).then((s) => s.status), + }) + } + } catch (err) { + console.error('[DataSourcesPanel] Failed to connect data source', sourceId, err) + const detail = err instanceof Error ? err.message : 'Please try again.' + setConnectError(`Couldn't connect ${sourceName}. ${detail}`) + } finally { + try { + await fetchDataSources(idToken) + } catch (err) { + console.error('[DataSourcesPanel] Failed to refresh data sources', err) + } + } + }, + [idToken, fetchDataSources, displaySources] + ) + const handleTabChange = useCallback( (value: string) => { setDataSourcesPanelTab(value as DataSourcesPanelTab) @@ -142,13 +202,22 @@ export const DataSourcesPanel: FC = memo(function DataSou availableSources.some((s) => s.id === id) ).length const availableCount = availableSources.length - const allAvailableEnabled = enabledAvailableCount === availableCount && availableCount > 0 + // Treat this as a master on/off switch: it stays on while any available source + // is enabled, and turns everything off when clicked. + const anyAvailableEnabled = enabledAvailableCount > 0 const handleToggleAll = useCallback(() => { - const updatedIds = allAvailableEnabled ? [] : availableSources.map((s) => s.id) + // Mirror DataConnectionCard's per-card gate: a protected source must be + // connected before it can be enabled. "Enable All" must skip protected + // sources that aren't connected, otherwise it bypasses that gate. + const updatedIds = anyAvailableEnabled + ? [] + : availableSources + .filter((s) => !(s.perUserAuth?.required && s.perUserAuth.status !== 'connected')) + .map((s) => s.id) setEnabledDataSources(updatedIds) saveDataSourcesToConversation(updatedIds) - }, [allAvailableEnabled, setEnabledDataSources, availableSources, saveDataSourcesToConversation]) + }, [anyAvailableEnabled, setEnabledDataSources, availableSources, saveDataSourcesToConversation]) return ( = memo(function DataSou )} + {/* Connect failure feedback so a failed attempt isn't silent */} + {connectError && ( + + {connectError} + + )} + {/* All Connections Toggle */} All Connections @@ -232,12 +308,12 @@ export const DataSourcesPanel: FC = memo(function DataSou className={`border-base mb-4 rounded-lg border p-3 transition-colors ${ isBusy ? 'cursor-not-allowed opacity-50' : 'cursor-pointer hover:bg-surface-raised-50' }`} - aria-pressed={allAvailableEnabled} + aria-pressed={anyAvailableEnabled} aria-disabled={isBusy} aria-label={ isBusy ? 'All available connections (disabled during operations)' - : `All available connections: ${allAvailableEnabled ? 'enabled' : 'disabled'}` + : `All available connections: ${anyAvailableEnabled ? 'enabled' : 'disabled'}` } title={isBusy ? 'Data source changes disabled during active operations' : undefined} > @@ -248,13 +324,13 @@ export const DataSourcesPanel: FC = memo(function DataSou
e.stopPropagation()}> = memo(function DataSou !isSourceAvailable ? 'Sign in required to access this data source' : undefined } onToggle={handleToggle} + onConnect={handleConnect} /> ) })} diff --git a/frontends/ui/src/features/layout/data-sources.ts b/frontends/ui/src/features/layout/data-sources.ts index 353ad3d08..26f21d1a3 100644 --- a/frontends/ui/src/features/layout/data-sources.ts +++ b/frontends/ui/src/features/layout/data-sources.ts @@ -11,6 +11,20 @@ /** Category types for organizing data sources */ export type DataSourceCategory = 'web' | 'enterprise' | 'storage' | 'collaboration' +/** Per-user MCP auth status for a protected source. */ +export type PerUserAuthStatus = 'connected' | 'not_connected' | 'expired' | 'error' + +/** Per-user MCP OAuth state attached to a protected data source (UI shape). */ +export interface PerUserAuth { + required: boolean + provider?: string | null + mcpServerId?: string | null + status?: PerUserAuthStatus | null + connectUrl?: string | null + expiresAt?: string | null + lastError?: string | null +} + /** Data source configuration interface */ export interface DataSource { /** Unique identifier matching backend source IDs */ @@ -25,4 +39,6 @@ export interface DataSource { defaultEnabled: boolean /** Whether the source requires user authentication */ requiresAuth: boolean + /** Per-user MCP OAuth state (present only for protected MCP sources) */ + perUserAuth?: PerUserAuth } diff --git a/frontends/ui/src/features/layout/store.spec.ts b/frontends/ui/src/features/layout/store.spec.ts index 6a994a677..2681d7ad3 100644 --- a/frontends/ui/src/features/layout/store.spec.ts +++ b/frontends/ui/src/features/layout/store.spec.ts @@ -204,5 +204,86 @@ describe('useLayoutStore', () => { ]) expect(useLayoutStore.getState().knowledgeLayerAvailable).toBe(true) }) + + test('does not auto-enable a protected source that is not connected', async () => { + mockGetDataSources.mockResolvedValueOnce({ + data_sources: [ + { id: 'web_search', name: 'Web Search', requires_auth: false }, + { + id: 'gdrive', + name: 'Google Drive', + requires_auth: false, + per_user_auth: { required: true, status: 'not_connected' }, + }, + ], + knowledge_layer: false, + }) + + await useLayoutStore.getState().fetchDataSources('token-1') + + // gdrive is not connected -> excluded from the initial selection. + expect(useLayoutStore.getState().enabledDataSourceIds).toEqual(['web_search']) + }) + + test('auto-enables a protected source once it is connected', async () => { + mockGetDataSources.mockResolvedValueOnce({ + data_sources: [ + { + id: 'gdrive', + name: 'Google Drive', + requires_auth: false, + per_user_auth: { required: true, status: 'connected' }, + }, + ], + knowledge_layer: false, + }) + + await useLayoutStore.getState().fetchDataSources('token-1') + + expect(useLayoutStore.getState().enabledDataSourceIds).toEqual(['gdrive']) + }) + }) + + describe('refreshDataSourceStatus', () => { + test('drops a protected source that is no longer connected from the selection', async () => { + useLayoutStore.setState({ enabledDataSourceIds: ['web_search', 'gdrive'] }) + mockGetDataSources.mockResolvedValueOnce({ + data_sources: [ + { id: 'web_search', name: 'Web Search', requires_auth: false }, + { + id: 'gdrive', + name: 'Google Drive', + requires_auth: false, + per_user_auth: { required: true, status: 'expired' }, + }, + ], + knowledge_layer: false, + }) + + await useLayoutStore.getState().refreshDataSourceStatus('token-1') + + // gdrive's token expired -> reconciled out; web_search preserved. + expect(useLayoutStore.getState().enabledDataSourceIds).toEqual(['web_search']) + }) + + test('keeps a still-connected protected source and preserves other selections', async () => { + useLayoutStore.setState({ enabledDataSourceIds: ['gdrive', 'confluence'] }) + mockGetDataSources.mockResolvedValueOnce({ + data_sources: [ + { + id: 'gdrive', + name: 'Google Drive', + requires_auth: false, + per_user_auth: { required: true, status: 'connected' }, + }, + // 'confluence' is absent from this response -> must be preserved, not dropped. + ], + knowledge_layer: false, + }) + + await useLayoutStore.getState().refreshDataSourceStatus('token-1') + + expect(useLayoutStore.getState().enabledDataSourceIds).toEqual(['gdrive', 'confluence']) + }) }) }) diff --git a/frontends/ui/src/features/layout/store.ts b/frontends/ui/src/features/layout/store.ts index cdefd7acd..42f34f85b 100644 --- a/frontends/ui/src/features/layout/store.ts +++ b/frontends/ui/src/features/layout/store.ts @@ -88,9 +88,18 @@ export const useLayoutStore = create()( const client = createDataSourcesClient({ authToken }) const response = await client.getDataSources() - // Start with every returned source enabled. Auth-aware cleanup still - // runs through disableAuthRequiredSources when a user loses access. - const enabledIds = response.data_sources.map((source) => source.id) + // Start with every returned source enabled, EXCEPT protected per-user + // sources that aren't connected yet: enabling those would put an + // unusable source into the selection (shown in "Selected Data Sources" + // and submitted), which the card toggle and "Enable All" already refuse + // to do. The user connects such a source and then enables it. Auth-aware + // cleanup still runs through disableAuthRequiredSources on access loss. + const enabledIds = response.data_sources + .filter( + (source) => + !(source.per_user_auth?.required && source.per_user_auth.status !== 'connected') + ) + .map((source) => source.id) set( { @@ -116,6 +125,42 @@ export const useLayoutStore = create()( } }, + refreshDataSourceStatus: async (authToken?: string) => { + // Selection-preserving refresh: update the source list/auth status without + // touching loading/error flags or resetting non-protected selections. The + // one exception is a protected source whose status is no longer + // 'connected' (e.g. the token expired since it was enabled) — it can no + // longer be used, so drop it from the selection here rather than leaving it + // shown in "Selected Data Sources" and submitted while unusable. + try { + const client = createDataSourcesClient({ authToken }) + const response = await client.getDataSources() + set( + (state) => { + const stillUsable = new Set( + response.data_sources + .filter((s) => !(s.per_user_auth?.required && s.per_user_auth.status !== 'connected')) + .map((s) => s.id) + ) + return { + availableDataSources: response.data_sources, + knowledgeLayerAvailable: response.knowledge_layer, + // Only ever removes now-unusable protected ids; other selections + // (incl. sources absent from this response) are preserved. + enabledDataSourceIds: state.enabledDataSourceIds.filter((id) => { + const src = response.data_sources.find((s) => s.id === id) + return !(src?.per_user_auth?.required) || stillUsable.has(id) + }), + } + }, + false, + 'refreshDataSourceStatus' + ) + } catch { + // Best-effort: keep the previously loaded state on failure. + } + }, + disableAuthRequiredSources: () => set( (state) => ({ diff --git a/frontends/ui/src/features/layout/types.ts b/frontends/ui/src/features/layout/types.ts index 71421daca..4660cd93a 100644 --- a/frontends/ui/src/features/layout/types.ts +++ b/frontends/ui/src/features/layout/types.ts @@ -75,6 +75,12 @@ export interface LayoutActions { setTheme: (theme: ThemeMode) => void /** Fetch data sources from API. Only web_search is enabled by default */ fetchDataSources: (authToken?: string) => Promise + /** + * Refresh per-source metadata/auth status WITHOUT changing the user's enabled + * selection. Used to pick up server-side status changes (e.g. a token that was + * invalidated at job time) so a stale "connected" card flips to Reconnect. + */ + refreshDataSourceStatus: (authToken?: string) => Promise /** Disable sources that require authentication */ disableAuthRequiredSources: () => void /** Set available data sources (from API) */ diff --git a/pyproject.toml b/pyproject.toml index 488339882..fe0175630 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "nvidia-nat[langchain,async_endpoints,phoenix,mcp]==1.8.0", "nvidia-nat-eval==1.8.0", "nvidia-nat-profiler==1.8.0", + "nvidia-nat-redis==1.8.0", "nvidia-nat-security[guardrails]==1.8.0", "deepagents>=0.6.5", "langgraph-checkpoint-postgres>=3.0.0", @@ -202,7 +203,7 @@ dev = [ ] [tool.uv] -prerelease = "if-necessary-or-explicit" +prerelease = "allow" override-dependencies = [ "authlib>=1.6.11,<2", "cryptography>=46.0.6,<47", diff --git a/skills/aiq-research/BENCHMARK.md b/skills/aiq-research/BENCHMARK.md index 7e54c2ac4..10df39b03 100644 --- a/skills/aiq-research/BENCHMARK.md +++ b/skills/aiq-research/BENCHMARK.md @@ -7,7 +7,7 @@ This benchmark summarizes 3-Tier Evaluation from NVSkills-Eval results for the s ## Evaluation Summary - Skill: `aiq-research` -- Evaluation date: 2026-06-29 +- Evaluation date: 2026-07-02 - NVSkills-Eval profile: `external` - Environment: `astra-sandbox` - Dataset: 3 evaluation tasks @@ -54,11 +54,11 @@ Task composition is derived from the evaluation dataset when possible. Entries w | Dimension | Num | `claude-code` | `codex` | |---|---:|---:|---:| -| Security | 3 | 100% (+0%) | 100% (+33%) | -| Correctness | 3 | 47% (+28%) | 71% (+28%) | -| Discoverability | 3 | 33% (+25%) | 37% (+23%) | -| Effectiveness | 3 | 42% (+24%) | 69% (+26%) | -| Efficiency | 3 | 47% (+21%) | 43% (+14%) | +| Security | 3 | 67% (-33%) | 100% (+17%) | +| Correctness | 3 | 93% (+6%) | 80% (+17%) | +| Discoverability | 3 | 100% (-0%) | 54% (+16%) | +| Effectiveness | 3 | 65% (-1%) | 69% (+14%) | +| Efficiency | 3 | 85% (-3%) | 56% (+14%) | Score values show skill-assisted performance. Values in parentheses show uplift versus the no-skill baseline when baseline data is available. diff --git a/skills/aiq-research/scripts/aiq.py b/skills/aiq-research/scripts/aiq.py index 8a7d75cb2..b6891fab6 100644 --- a/skills/aiq-research/scripts/aiq.py +++ b/skills/aiq-research/scripts/aiq.py @@ -88,6 +88,20 @@ def _int_const(value: str) -> int: _CHAT_JOB_ID_RE = re.compile(rf"Job ID:\s*([0-9a-f-]{{{JOB_ID_HEX_DASH_LENGTH}}})", re.IGNORECASE) +class McpAuthRequiredError(RuntimeError): + """The API returned 409 ``mcp_auth_required``. + + One or more selected protected data sources must be connected (per-user MCP + OAuth) before the job can run. Carries the structured ``sources`` list so the + caller can surface each source's status and connect/auth URL. + """ + + def __init__(self, payload: dict[str, Any]) -> None: + self.payload = payload + self.sources: list[dict[str, Any]] = payload.get("sources") or [] + super().__init__(payload.get("message") or "Data source connection required") + + def _validate_base_url(url: str) -> str: """Validate and normalize the configured AI-Q server base URL.""" raw = (url or "").strip() @@ -137,6 +151,13 @@ def _validate_agent_type(agent_type: str) -> str: return value +def _validate_source_id(source_id: str) -> str: + value = source_id.strip() + if not _AGENT_TYPE_RE.fullmatch(value): # same charset: [A-Za-z0-9_.-] + raise RuntimeError("Invalid source_id") + return value + + def _api_request( method: str, path: str, @@ -162,6 +183,15 @@ def _api_request( payload = resp.read().decode("utf-8") except urllib.error.HTTPError as exc: error_body = exc.read().decode("utf-8", errors="replace") + # A 409 mcp_auth_required carries a structured body listing the protected + # sources to connect; surface it rather than collapsing to "HTTP 409". + if exc.code == 409: + try: + parsed = json.loads(error_body) + except json.JSONDecodeError: + parsed = None + if isinstance(parsed, dict) and parsed.get("error") == "mcp_auth_required": + raise McpAuthRequiredError(parsed) from exc print(f"HTTP {exc.code}: {error_body[:ERROR_BODY_PREVIEW_CHARS]}", file=sys.stderr) raise RuntimeError(f"HTTP {exc.code}") from exc except urllib.error.URLError as exc: @@ -274,6 +304,39 @@ def cancel_job(job_id: str) -> dict[str, Any]: return _api_request("POST", f"/v1/jobs/async/job/{_validate_job_id(job_id)}/cancel") +def list_data_sources() -> Any: + """GET /v1/data_sources — available sources + this user's per-source auth state.""" + return _api_request("GET", "/v1/data_sources") + + +def source_auth_status(source_id: str) -> dict[str, Any]: + """GET /v1/auth/mcp/{id}/status — per-user MCP auth status for one source.""" + return _api_request("GET", f"/v1/auth/mcp/{_validate_source_id(source_id)}/status") + + +def connect_source(source_id: str) -> dict[str, Any]: + """POST /v1/auth/mcp/{id}/connect — start the OAuth flow; returns an auth_url.""" + return _api_request("POST", f"/v1/auth/mcp/{_validate_source_id(source_id)}/connect") + + +def _print_mcp_auth_required(err: McpAuthRequiredError) -> None: + """Render a 409 mcp_auth_required as actionable connect guidance on stderr.""" + base = _validate_base_url(AIQ_SERVER_URL) + print(f"\n{err}\n", file=sys.stderr) + for src in err.sources: + sid = src.get("source_id", "?") + status = src.get("status", "?") + print(f" - {sid} ({status})", file=sys.stderr) + auth_url = src.get("auth_url") + connect_url = src.get("connect_url") + if auth_url: + print(f" Sign in: open this URL in a browser -> {auth_url}", file=sys.stderr) + elif connect_url: + print(f" Start the connect flow: aiq.py connect {sid}", file=sys.stderr) + print(f" (POST {base}{connect_url} returns the sign-in URL)", file=sys.stderr) + print("\nConnect the source(s) above, then re-run the command.", file=sys.stderr) + + def stream_job(job_id: str) -> None: """Print server-sent event payloads for an async AI-Q job.""" for line in _stream_request(f"/v1/jobs/async/job/{_validate_job_id(job_id)}/stream"): @@ -362,6 +425,9 @@ def _print_usage() -> None: print(" research [agent_type] Submit async job, poll, and return report") print(" research_poll Resume polling an existing async job") print(" cancel Cancel a running async job") + print(" data_sources List data sources + per-user auth state") + print(" source_status Per-user MCP auth status for one source") + print(" connect Start MCP OAuth; prints the sign-in URL") print() print(f"Environment: AIQ_SERVER_URL defaults to {DEFAULT_SERVER_URL}") @@ -586,6 +652,27 @@ def _command_cancel(args: list[str]) -> None: print(json.dumps(cancel_job(job_id), indent=JSON_INDENT_SPACES)) +def _command_data_sources(_args: list[str]) -> None: + print(json.dumps(list_data_sources(), indent=JSON_INDENT_SPACES)) + + +def _command_source_status(args: list[str]) -> None: + source_id = _require_arg(args, "Usage: aiq.py source_status ") + print(json.dumps(source_auth_status(source_id), indent=JSON_INDENT_SPACES)) + + +def _command_connect(args: list[str]) -> None: + source_id = _require_arg(args, "Usage: aiq.py connect ") + result = connect_source(source_id) + auth_url = result.get("auth_url") + if result.get("status") == "connected": + print(f"{source_id} is already connected.", file=sys.stderr) + elif auth_url: + print("Open this URL in a browser to sign in, then re-run your command:", file=sys.stderr) + print(auth_url) + print(json.dumps(result, indent=JSON_INDENT_SPACES), file=sys.stderr) + + def main() -> None: """Dispatch the command-line interface.""" if len(sys.argv) < MIN_COMMAND_ARG_COUNT: @@ -606,6 +693,9 @@ def main() -> None: "research": _command_research, "research_poll": _command_research_poll, "cancel": _command_cancel, + "data_sources": _command_data_sources, + "source_status": _command_source_status, + "connect": _command_connect, } handler = commands.get(cmd) if handler is None: @@ -620,4 +710,8 @@ def main() -> None: if __name__ == "__main__": - main() + try: + main() + except McpAuthRequiredError as err: + _print_mcp_auth_required(err) + sys.exit(2) diff --git a/skills/aiq-research/skill-card.md b/skills/aiq-research/skill-card.md index 419443e88..e74981604 100644 --- a/skills/aiq-research/skill-card.md +++ b/skills/aiq-research/skill-card.md @@ -9,14 +9,14 @@ NVIDIA
### License/Terms of Use:
Apache 2.0
## Use Case:
-Developers and engineers use this skill to submit deep research queries to a local or self-hosted NVIDIA AI-Q Blueprint backend and retrieve structured research reports with citations.
+Developers and engineers use this skill to submit deep research queries to a locally or self-hosted NVIDIA AI-Q Blueprint backend and retrieve comprehensive research reports with citations.
### Deployment Geography for Use:
Global
## Requirements / Dependencies:
**Requires API Key or External Credential:** [No]
-**Credential Type(s):** [None]
+**Credential Type(s):** [None]
Do not include secrets in prompts/logs/output; use least-privilege credentials; rotate keys as appropriate.
@@ -26,24 +26,25 @@ Mitigation: Review and scan skill before deployment.
## Reference(s):
- [NVIDIA AI-Q Blueprint Repository](https://github.com/NVIDIA-AI-Blueprints/aiq)
-- [Helper Script](scripts/aiq.py)
-- [Deployment Skill](../aiq-deploy/SKILL.md)
+- [DeepResearch Bench Paper](https://arxiv.org/pdf/2506.11763)
+- [DeepResearch Bench Leaderboard](https://huggingface.co/spaces/muset-ai/DeepResearch-Bench-Leaderboard)
+- [Helper script](scripts/aiq.py)
## Skill Output:
-**Output Type(s):** [Shell commands, API Calls, Analysis]
-**Output Format:** [Markdown with inline bash code blocks]
+**Output Type(s):** [Analysis, API Calls]
+**Output Format:** [JSON and Markdown research reports with citations]
**Output Parameters:** [1D]
**Other Properties Related to Output:** [None]
## Evaluation Agents Used:
-- claude-code
-- codex
+- Claude Code (`claude-code`)
+- Codex (`codex`)
## Evaluation Tasks:
-Evaluated against 3 evaluation tasks using NVSkills-Eval `external` profile in `astra-sandbox` environment.
+Evaluated against 3 evaluation tasks from NVSkills-Eval external profile (3 positive skill-activation cases).
## Evaluation Metrics Used:
Reported benchmark dimensions:
@@ -67,11 +68,11 @@ Underlying evaluation signals used in this run:
## Evaluation Results:
| Dimension | Num | `claude-code` | `codex` | |---|---:|---:|---:| -| Security | 3 | 100% (+0%) | 100% (+33%) | -| Correctness | 3 | 47% (+28%) | 71% (+28%) | -| Discoverability | 3 | 33% (+25%) | 37% (+23%) | -| Effectiveness | 3 | 42% (+24%) | 69% (+26%) | -| Efficiency | 3 | 47% (+21%) | 43% (+14%) | +| Security | 3 | 67% (-33%) | 100% (+17%) | +| Correctness | 3 | 93% (+6%) | 80% (+17%) | +| Discoverability | 3 | 100% (-0%) | 54% (+16%) | +| Effectiveness | 3 | 65% (-1%) | 69% (+14%) | +| Efficiency | 3 | 85% (-3%) | 56% (+14%) | ## Skill Version(s):
2.1.0 (source: frontmatter)
diff --git a/skills/aiq-research/skill.oms.sig b/skills/aiq-research/skill.oms.sig index 5fa238f0d..fbc5a5dd2 100644 --- a/skills/aiq-research/skill.oms.sig +++ b/skills/aiq-research/skill.oms.sig @@ -1 +1 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"x509CertificateChain":{"certificates":[{"rawBytes":"MIICgzCCAgmgAwIBAgIUKIyS7SxNteQIiWzK1dWj85E6520wCgYIKoZIzj0EAwMwVTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjEpMCcGA1UEAwwgTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBJQ0EgMDEwHhcNMjYwNDAxMDAwMDAwWhcNMjgwNDIyMTUzMzA5WjBUMQswCQYDVQQGEwJVUzEbMBkGA1UECgwSTlZJRElBIENvcnBvcmF0aW9uMSgwJgYDVQQDDB9OVklESUEgQWdlbnQgU2tpbGxzIFNpZ25pbmcgMDAxMHYwEAYHKoZIzj0CAQYFK4EEACIDYgAEYoRM9bQl/dGlwSRNi6bTpIJUXH8Nv9GciP6LSflJYYMLCc296kpyuTSsk5ddbAWiDcFX3C/ydX3jwc+qCLYP6uHy9XphyLjOQ27Yb2J6rBLVtRBS1mgGco/Gr7fL6ODco4GaMIGXMB0GA1UdDgQWBBRQ/5ZW3nJ6lmo9SVk7I15o7UGmpTAfBgNVHSMEGDAWgBRPGpILxMBBleJSsBGjrMKsby1CgjAMBgNVHRMBAf8EAjAAMA4GA1UdDwEB/wQEAwIHgDA3BggrBgEFBQcBAQQrMCkwJwYIKwYBBQUHMAGGG2h0dHA6Ly9vY3NwLm5kaXMubnZpZGlhLmNvbTAKBggqhkjOPQQDAwNoADBlAjAUygu/GiOCIXrgGr4SmLgeEVDcEitfFUv7ALbvLVGVyMysB3mxmO/uInZfXzWcJZsCMQDxuoxj4ZmO30jhkPIcCxGFCOvnUsnfU3TfGcouYm4M6iRpbKvtVnHPiy4bi6pcKf0="},{"rawBytes":"MIICiDCCAg6gAwIBAgIUZsIuSv9NkpJCNqtYEfCouVv5BzowCgYIKoZIzj0EAwMwUTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjElMCMGA1UEAwwcTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBDQTAgFw0yNjA0MDEwMDAwMDBaGA85OTk5MTIzMTIzNTk1OVowVTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjEpMCcGA1UEAwwgTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBJQ0EgMDEwdjAQBgcqhkjOPQIBBgUrgQQAIgNiAASI72cR3ctKGg4VWnB3bNja6g1Z2PnOmFEopkPof+QeIcPk9rT+g9MjJnq51EQXL93a7C2GJ9J985G4o2V85VD7wJ1RaXhluHW2rf3y8bQGeAYaKMr5s/hUgn+M3/9WlWejgaAwgZ0wHQYDVR0OBBYEFE8akgvEwEGV4lKwEaOswqxvLUKCMB8GA1UdIwQYMBaAFItnoAjjfuCEUvzyvWyI2vOGvwPjMBIGA1UdEwEB/wQIMAYBAf8CAQAwDgYDVR0PAQH/BAQDAgEGMDcGCCsGAQUFBwEBBCswKTAnBggrBgEFBQcwAYYbaHR0cDovL29jc3AubmRpcy5udmlkaWEuY29tMAoGCCqGSM49BAMDA2gAMGUCMQCeIMMfAbyzPDacw2MxG+Yt1cikrJX/DVxiGfXuHmkkXn6VgSzE79+lkqDErpVO2gYCMCNEColOyvUvkzZGUEI1hQ3PfMgi3FIo9tHoBKMw4/wGBLFpu/0ubtmbBXM6/UMOEw=="},{"rawBytes":"MIICRTCCAcygAwIBAgIUeJdY3rV86EdvFmG7L8LJBsyQFYkwCgYIKoZIzj0EAwMwUTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjElMCMGA1UEAwwcTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBDQTAgFw0yNjA0MDEwMDAwMDBaGA85OTk5MTIzMTIzNTk1OVowUTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjElMCMGA1UEAwwcTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBDQTB2MBAGByqGSM49AgEGBSuBBAAiA2IABAYpiXCDjJ9NT2eSDhyHJVSw1Tbze18cGG2F/578oWvHxg23eQAhNRYdq88i1iOshZSO6C29doKui5Xpmo/7Ctw9Sx4PP2RzOmIuOLCuTdNtKcTRwi4GEsd5BAFvWj42M6NjMGEwHQYDVR0OBBYEFItnoAjjfuCEUvzyvWyI2vOGvwPjMB8GA1UdIwQYMBaAFItnoAjjfuCEUvzyvWyI2vOGvwPjMA8GA1UdEwEB/wQFMAMBAf8wDgYDVR0PAQH/BAQDAgEGMAoGCCqGSM49BAMDA2cAMGQCMCwtAjWLaNwgGWNCgdyNoTyvNhqWRECRJV2r3+7w8g0PL6NHLOsbkgE09BH95h8XlgIwTaQmbbUh2ChAJ5TA1wRiVDnCcvbzHlZl2jM2FcwQQZlk19LOAbyGMRixbu2Ww/rj"}]},"tlogEntries":[]},"dsseEnvelope":{"payload":"ewogICJfdHlwZSI6ICJodHRwczovL2luLXRvdG8uaW8vU3RhdGVtZW50L3YxIiwKICAic3ViamVjdCI6IFsKICAgIHsKICAgICAgIm5hbWUiOiAiYWlxLXJlc2VhcmNoIiwKICAgICAgImRpZ2VzdCI6IHsKICAgICAgICAic2hhMjU2IjogImMxNTJhOTk1MmZjMjg2YjNlZjY1NTkyZDc4ZTcxODk5NGJkZjUyYWNjZGI5OGRiYmQzN2Y0MTFhMTQ3MzM3MTUiCiAgICAgIH0KICAgIH0KICBdLAogICJwcmVkaWNhdGVUeXBlIjogImh0dHBzOi8vbW9kZWxfc2lnbmluZy9zaWduYXR1cmUvdjEuMCIsCiAgInByZWRpY2F0ZSI6IHsKICAgICJzZXJpYWxpemF0aW9uIjogewogICAgICAibWV0aG9kIjogImZpbGVzIiwKICAgICAgImlnbm9yZV9wYXRocyI6IFsKICAgICAgICAiLmdpdGlnbm9yZSIsCiAgICAgICAgIi5naXRhdHRyaWJ1dGVzIiwKICAgICAgICAiLmdpdGh1YiIsCiAgICAgICAgIi5naXQiCiAgICAgIF0sCiAgICAgICJoYXNoX3R5cGUiOiAic2hhMjU2IiwKICAgICAgImFsbG93X3N5bWxpbmtzIjogZmFsc2UKICAgIH0sCiAgICAicmVzb3VyY2VzIjogWwogICAgICB7CiAgICAgICAgImRpZ2VzdCI6ICJlYjYwYWRjZDhjMjFjMWRhNmZmY2FjYTc1YjNiZGU0ZjU1MzIxNjljZjNjYjEzODVkM2E4OWRkNGM2NmNkMDNjIiwKICAgICAgICAibmFtZSI6ICJCRU5DSE1BUksubWQiLAogICAgICAgICJhbGdvcml0aG0iOiAic2hhMjU2IgogICAgICB9LAogICAgICB7CiAgICAgICAgImRpZ2VzdCI6ICIyNDgxNmZjZTRjZjM3YTA3MWE5OTlmMWMzYTEyZjU3ZWI5OGNkZTA3MTA1NjZkMmRmNmI1YTlmYWIzN2Y3OTBkIiwKICAgICAgICAibmFtZSI6ICJTS0lMTC5tZCIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogIjY3ZDNjYWEyYWMzOWU3ZTcxZjc1OTNiOTk4MGUzNTIxMzU5MTI5YTA3ZWI4ZjMwOGNhMTdjZDBlODZhYjZjOTciLAogICAgICAgICJuYW1lIjogImV2YWxzL2Jhc2ljLXByb2R1Y3QuanNvbiIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogImFmZjI1ZDc4NDIxMDE5NzFjYTU5M2FiYjdkMTAxYTBkZmE4MGFiZWM3N2IzYjE1OTY3ODEyNWQxM2FlNmY4MDQiLAogICAgICAgICJuYW1lIjogImV2YWxzL2V2YWxzLmpzb24iLAogICAgICAgICJhbGdvcml0aG0iOiAic2hhMjU2IgogICAgICB9LAogICAgICB7CiAgICAgICAgImRpZ2VzdCI6ICIzMzkyMzEwNzZmZDEyYjQ4YWQzMmQ4OGZmYmUxZmU4YmEzMDdmYWZiZGI2NjY0N2UzNmMyZjdhMzEyYjJmNzlhIiwKICAgICAgICAibmFtZSI6ICJzY3JpcHRzL2FpcS5weSIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogImVlNDdlYzY0NGU3NDdmZWNmNWNjYTYwYjM2N2I1NDEyZDUxN2E1ZGFkNTFiMTFiODI3YThmZDljOTgwMTNmODEiLAogICAgICAgICJuYW1lIjogInNraWxsLWNhcmQubWQiLAogICAgICAgICJhbGdvcml0aG0iOiAic2hhMjU2IgogICAgICB9CiAgICBdCiAgfQp9","payloadType":"application/vnd.in-toto+json","signatures":[{"sig":"MGUCMQDQCTZOUIoV6wTrazkHIIBL5Vw0uohingxyL1msXvR2SQwyzkN23Bq4WzECVsi3NbICMCWLvuP1bPW3eN5AN7sN+ipFEYWJzWmEInUlDI/FDSG0wyOG3WDAfQ3FgAL0VTN3JQ==","keyid":""}]}} \ No newline at end of file +{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"x509CertificateChain":{"certificates":[{"rawBytes":"MIICgzCCAgmgAwIBAgIUKIyS7SxNteQIiWzK1dWj85E6520wCgYIKoZIzj0EAwMwVTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjEpMCcGA1UEAwwgTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBJQ0EgMDEwHhcNMjYwNDAxMDAwMDAwWhcNMjgwNDIyMTUzMzA5WjBUMQswCQYDVQQGEwJVUzEbMBkGA1UECgwSTlZJRElBIENvcnBvcmF0aW9uMSgwJgYDVQQDDB9OVklESUEgQWdlbnQgU2tpbGxzIFNpZ25pbmcgMDAxMHYwEAYHKoZIzj0CAQYFK4EEACIDYgAEYoRM9bQl/dGlwSRNi6bTpIJUXH8Nv9GciP6LSflJYYMLCc296kpyuTSsk5ddbAWiDcFX3C/ydX3jwc+qCLYP6uHy9XphyLjOQ27Yb2J6rBLVtRBS1mgGco/Gr7fL6ODco4GaMIGXMB0GA1UdDgQWBBRQ/5ZW3nJ6lmo9SVk7I15o7UGmpTAfBgNVHSMEGDAWgBRPGpILxMBBleJSsBGjrMKsby1CgjAMBgNVHRMBAf8EAjAAMA4GA1UdDwEB/wQEAwIHgDA3BggrBgEFBQcBAQQrMCkwJwYIKwYBBQUHMAGGG2h0dHA6Ly9vY3NwLm5kaXMubnZpZGlhLmNvbTAKBggqhkjOPQQDAwNoADBlAjAUygu/GiOCIXrgGr4SmLgeEVDcEitfFUv7ALbvLVGVyMysB3mxmO/uInZfXzWcJZsCMQDxuoxj4ZmO30jhkPIcCxGFCOvnUsnfU3TfGcouYm4M6iRpbKvtVnHPiy4bi6pcKf0="},{"rawBytes":"MIICiDCCAg6gAwIBAgIUZsIuSv9NkpJCNqtYEfCouVv5BzowCgYIKoZIzj0EAwMwUTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjElMCMGA1UEAwwcTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBDQTAgFw0yNjA0MDEwMDAwMDBaGA85OTk5MTIzMTIzNTk1OVowVTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjEpMCcGA1UEAwwgTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBJQ0EgMDEwdjAQBgcqhkjOPQIBBgUrgQQAIgNiAASI72cR3ctKGg4VWnB3bNja6g1Z2PnOmFEopkPof+QeIcPk9rT+g9MjJnq51EQXL93a7C2GJ9J985G4o2V85VD7wJ1RaXhluHW2rf3y8bQGeAYaKMr5s/hUgn+M3/9WlWejgaAwgZ0wHQYDVR0OBBYEFE8akgvEwEGV4lKwEaOswqxvLUKCMB8GA1UdIwQYMBaAFItnoAjjfuCEUvzyvWyI2vOGvwPjMBIGA1UdEwEB/wQIMAYBAf8CAQAwDgYDVR0PAQH/BAQDAgEGMDcGCCsGAQUFBwEBBCswKTAnBggrBgEFBQcwAYYbaHR0cDovL29jc3AubmRpcy5udmlkaWEuY29tMAoGCCqGSM49BAMDA2gAMGUCMQCeIMMfAbyzPDacw2MxG+Yt1cikrJX/DVxiGfXuHmkkXn6VgSzE79+lkqDErpVO2gYCMCNEColOyvUvkzZGUEI1hQ3PfMgi3FIo9tHoBKMw4/wGBLFpu/0ubtmbBXM6/UMOEw=="},{"rawBytes":"MIICRTCCAcygAwIBAgIUeJdY3rV86EdvFmG7L8LJBsyQFYkwCgYIKoZIzj0EAwMwUTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjElMCMGA1UEAwwcTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBDQTAgFw0yNjA0MDEwMDAwMDBaGA85OTk5MTIzMTIzNTk1OVowUTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjElMCMGA1UEAwwcTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBDQTB2MBAGByqGSM49AgEGBSuBBAAiA2IABAYpiXCDjJ9NT2eSDhyHJVSw1Tbze18cGG2F/578oWvHxg23eQAhNRYdq88i1iOshZSO6C29doKui5Xpmo/7Ctw9Sx4PP2RzOmIuOLCuTdNtKcTRwi4GEsd5BAFvWj42M6NjMGEwHQYDVR0OBBYEFItnoAjjfuCEUvzyvWyI2vOGvwPjMB8GA1UdIwQYMBaAFItnoAjjfuCEUvzyvWyI2vOGvwPjMA8GA1UdEwEB/wQFMAMBAf8wDgYDVR0PAQH/BAQDAgEGMAoGCCqGSM49BAMDA2cAMGQCMCwtAjWLaNwgGWNCgdyNoTyvNhqWRECRJV2r3+7w8g0PL6NHLOsbkgE09BH95h8XlgIwTaQmbbUh2ChAJ5TA1wRiVDnCcvbzHlZl2jM2FcwQQZlk19LOAbyGMRixbu2Ww/rj"}]},"tlogEntries":[]},"dsseEnvelope":{"payload":"ewogICJfdHlwZSI6ICJodHRwczovL2luLXRvdG8uaW8vU3RhdGVtZW50L3YxIiwKICAic3ViamVjdCI6IFsKICAgIHsKICAgICAgIm5hbWUiOiAiYWlxLXJlc2VhcmNoIiwKICAgICAgImRpZ2VzdCI6IHsKICAgICAgICAic2hhMjU2IjogImE5N2I3YWZiOWFmYzU3MmRlNWY5NDcwZjdhYzQ4OTZiZjAzNzdhYzAyYzk5MTUyNjA5NmM3NDQzNDYxNTE3MzUiCiAgICAgIH0KICAgIH0KICBdLAogICJwcmVkaWNhdGVUeXBlIjogImh0dHBzOi8vbW9kZWxfc2lnbmluZy9zaWduYXR1cmUvdjEuMCIsCiAgInByZWRpY2F0ZSI6IHsKICAgICJyZXNvdXJjZXMiOiBbCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogIjAzNDczMDI4NDEzYjE2NjJmY2U1OGFjYTM1OWI0OTkyYTEwODViN2UwODk1OTUyNGY3ZGI5Mjc0NjBiY2IwNDkiLAogICAgICAgICJuYW1lIjogIkJFTkNITUFSSy5tZCIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogIjkxNzJkOTM4M2E5NGI2MjA4YmRkZjllZjg2YTJhZTU1ZWYzYmJiZGEzNWMwYzgzZGZjNjVhMzJmMjlmNjRlMWMiLAogICAgICAgICJuYW1lIjogIlNLSUxMLm1kIiwKICAgICAgICAiYWxnb3JpdGhtIjogInNoYTI1NiIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiNjdkM2NhYTJhYzM5ZTdlNzFmNzU5M2I5OTgwZTM1MjEzNTkxMjlhMDdlYjhmMzA4Y2ExN2NkMGU4NmFiNmM5NyIsCiAgICAgICAgIm5hbWUiOiAiZXZhbHMvYmFzaWMtcHJvZHVjdC5qc29uIiwKICAgICAgICAiYWxnb3JpdGhtIjogInNoYTI1NiIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiYWZmMjVkNzg0MjEwMTk3MWNhNTkzYWJiN2QxMDFhMGRmYTgwYWJlYzc3YjNiMTU5Njc4MTI1ZDEzYWU2ZjgwNCIsCiAgICAgICAgIm5hbWUiOiAiZXZhbHMvZXZhbHMuanNvbiIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogImUyOTQ0ZmRjYmNmMjEwNjI0N2Q3NzgzYzUwOTcyNDYxNjljMDM2ZTg5MDkyMWFiMGM1OWI3ZjA4NTEwNzY1NjMiLAogICAgICAgICJuYW1lIjogInNjcmlwdHMvYWlxLnB5IiwKICAgICAgICAiYWxnb3JpdGhtIjogInNoYTI1NiIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiOTFkY2M5NGIyOTZhNDI3OTI2YzI1MWZmNDhiMzFlY2FmYTFiNjdjZTc2ZTI3ZWQ3NDFiMDg5NGFmNDZlYzNmMSIsCiAgICAgICAgIm5hbWUiOiAic2tpbGwtY2FyZC5tZCIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiCiAgICAgIH0KICAgIF0sCiAgICAic2VyaWFsaXphdGlvbiI6IHsKICAgICAgImhhc2hfdHlwZSI6ICJzaGEyNTYiLAogICAgICAiYWxsb3dfc3ltbGlua3MiOiBmYWxzZSwKICAgICAgIm1ldGhvZCI6ICJmaWxlcyIsCiAgICAgICJpZ25vcmVfcGF0aHMiOiBbCiAgICAgICAgIi5naXRhdHRyaWJ1dGVzIiwKICAgICAgICAiLmdpdGlnbm9yZSIsCiAgICAgICAgIi5naXQiLAogICAgICAgICIuZ2l0aHViIgogICAgICBdCiAgICB9CiAgfQp9","payloadType":"application/vnd.in-toto+json","signatures":[{"sig":"MGUCMGziCkQjT1fm+zUkkSgUCbXnFCE91vhAo6QT3BRZPzNRMDwFEVczCbwn1V4L17LgSAIxAIdHjROhoRWdUEYS5qM45WVyJSv4T6Mjns7jq4E8bW/SvjUPy3GXOY9km0yNZqsS5w==","keyid":""}]}} \ No newline at end of file diff --git a/src/aiq_agent/agents/chat_researcher/agent.py b/src/aiq_agent/agents/chat_researcher/agent.py index 5a1217121..62662b258 100644 --- a/src/aiq_agent/agents/chat_researcher/agent.py +++ b/src/aiq_agent/agents/chat_researcher/agent.py @@ -302,7 +302,16 @@ async def shallow_research_node(state: ChatResearcherState) -> dict[str, Any]: async def deep_research_node(state: ChatResearcherState) -> dict[str, Any]: if self.deep_research_job_submitter is not None: - job_id = await self.deep_research_job_submitter(state) + try: + job_id = await self.deep_research_job_submitter(state) + except Exception as e: + # Surface auth failures to the user verbatim instead of failing the + # turn — submit_agent_job raises McpAuthRequiredError (an AuthError) + # before enqueue when a selected source isn't connected. + if _AuthError and isinstance(e, _AuthError): + logger.warning("Auth required before deep research submit: %s", e) + return {"messages": [AIMessage(content=str(e))]} + raise escalation = _job_escalation_message(_ESCALATION_KIND_DEEP_RESEARCH, job_id) return {"messages": [AIMessage(content=escalation)]} diff --git a/src/aiq_agent/agents/chat_researcher/register.py b/src/aiq_agent/agents/chat_researcher/register.py index fac28770a..ade23cd06 100644 --- a/src/aiq_agent/agents/chat_researcher/register.py +++ b/src/aiq_agent/agents/chat_researcher/register.py @@ -319,6 +319,22 @@ def validate_deep_research_tools(data_sources: list[str] | None) -> tuple[bool, """ from aiq_agent.common import format_tool_unavailability_error from aiq_agent.common import validate_tool_availability + from aiq_agent.common.data_source_registry import get_source + + # Per-user MCP sources (e.g. Google Drive) contribute NO tools to the + # static startup list — their tools are resolved per-user at run time by + # open_per_user_mcp_tools (in the async worker). So a selected, configured + # per-user MCP source is a valid runtime tool candidate here; rejecting it + # against the static list would block the job before it ever submits. + # Connectivity is enforced separately by the submit preflight + # (submit_agent_job -> evaluate_mcp_auth, which returns mcp_auth_required + # when not connected), and the authoritative tool check runs after + # resolution. + if data_sources: + for source_id in data_sources: + source = get_source(source_id) + if source and source.per_user_auth and source.per_user_auth.required: + return True, "" selected_tools = filter_tools_by_sources(deep_research_tools, data_sources) @@ -433,12 +449,17 @@ async def _build_report_seed_files(state: ChatResearcherState) -> dict[str, str] scheduler_address = os.environ.get("NAT_DASK_SCHEDULER_ADDRESS") if scheduler_address: from aiq_agent.auth import get_auth_token - from aiq_agent.auth import get_current_principal + from aiq_api.jobs.access import require_verified_principal from aiq_api.jobs.submit import submit_agent_job async def _submit_deep_job(state: ChatResearcherState) -> str: - principal = get_current_principal() - owner = principal.email if principal and principal.email else "anonymous" + # Resolve the principal the same way the connect/status routes do, so the + # job's per-user MCP token key (principal_user_id) matches where the token + # was stored at connect time. Keying off `owner` (email) instead would + # mismatch the connect key ("{type}:{sub}") and trigger interactive re-auth + # for an already-connected source. + principal = require_verified_principal() + owner = principal.email or principal.sub query = state.original_query if not query: if not state.messages: @@ -471,7 +492,7 @@ async def _submit_deep_job(state: ChatResearcherState) -> str: agent_type="deep_researcher", input_text=input_text, owner=owner, - principal=principal, + principal=principal, # ensures worker token key == connect-time key available_documents=available_docs, data_sources=state.data_sources, auth_token=get_auth_token(), diff --git a/src/aiq_agent/agents/clarifier/agent.py b/src/aiq_agent/agents/clarifier/agent.py index 5779e5f0d..ebcaf1170 100644 --- a/src/aiq_agent/agents/clarifier/agent.py +++ b/src/aiq_agent/agents/clarifier/agent.py @@ -59,6 +59,7 @@ from aiq_agent.common import LLMProvider from aiq_agent.common import LLMRole +from aiq_agent.common import format_data_source_tools from aiq_agent.common import get_latest_user_query from aiq_agent.common import load_prompt from aiq_agent.common import render_prompt_template @@ -531,10 +532,15 @@ async def agent_node(state: ClarifierAgentState): tools_info = [ {"name": getattr(t, "name", ""), "description": getattr(t, "description", "")} for t in self.tools ] + # Selected data sources (e.g. Google Drive) become tools in the research + # phase but are NOT bound to the clarifier, so surface them explicitly — + # otherwise the LLM refuses Drive/URL requests it thinks it can't access. + connected_sources = format_data_source_tools(state.data_sources) if state.data_sources else [] rendered_system_prompt = render_prompt_template( self.system_prompt, clarifier_result=state.clarifier_log, available_documents=state.available_documents or [], + connected_sources=connected_sources, tools=tools_info, tool_names=[t["name"] for t in tools_info], ) diff --git a/src/aiq_agent/agents/clarifier/prompts/research_clarification.j2 b/src/aiq_agent/agents/clarifier/prompts/research_clarification.j2 index a27455777..8cbcccd8d 100644 --- a/src/aiq_agent/agents/clarifier/prompts/research_clarification.j2 +++ b/src/aiq_agent/agents/clarifier/prompts/research_clarification.j2 @@ -151,6 +151,14 @@ The user has uploaded the following file(s); they are ingested and will be inclu If the user's request refers to "the paper", "my file", "the document", or a filename (e.g. {{ available_documents[0].file_name }}), treat it as referring to these uploaded files. Do NOT ask which paper or which file — proceed with the understanding that the uploaded file(s) are the subject. Only ask for clarification on other missing dimensions (e.g. comparison criteria), not the document identity. {% endif %} +{% if connected_sources and connected_sources | length > 0 %} +## Connected data sources (IMPORTANT) +The research phase has tools for these connected sources; it CAN access them even though you cannot read them yourself:{% for s in connected_sources %} +- **{{ s.name }}**: {{ s.description }}{% endfor %} + +If the user references a resource served by a connected source (e.g. a Google Drive link or file when Google Drive is connected), treat it as accessible — the research tools will fetch it. Do NOT refuse, claim you cannot access it, or ask the user to paste its contents. Only ask for clarification on genuine research-scope gaps (e.g. what to focus on), not on access. +{% endif %} + {% if clarifier_result is defined and clarifier_result %} ## Clarification Context So Far {{ clarifier_result }} diff --git a/src/aiq_agent/agents/deep_researcher/agent.py b/src/aiq_agent/agents/deep_researcher/agent.py index 0b0f66457..d48968cb3 100644 --- a/src/aiq_agent/agents/deep_researcher/agent.py +++ b/src/aiq_agent/agents/deep_researcher/agent.py @@ -138,6 +138,7 @@ def __init__( self.middleware_set = build_deep_research_middleware_set( tool_set=self.tool_set, source_registry_middleware=self.source_registry_middleware, + enable_source_router=self.enable_source_router, ) self.source_tool_names = self.tool_set.source_tool_names diff --git a/src/aiq_agent/agents/deep_researcher/custom_middleware.py b/src/aiq_agent/agents/deep_researcher/custom_middleware.py index c2379fce7..09bbe659f 100644 --- a/src/aiq_agent/agents/deep_researcher/custom_middleware.py +++ b/src/aiq_agent/agents/deep_researcher/custom_middleware.py @@ -37,6 +37,40 @@ # Path to this agent's prompts directory _PROMPTS_DIR = Path(__file__).parent / "prompts" +_SOURCE_ROUTING_PATH = "/shared/source_routing.json" + + +class SourceRoutingGuardMiddleware(AgentMiddleware): + """Require the source-router handoff before other orchestrator tool calls.""" + + def __init__(self, *, enabled: bool, required_subagent: str = "source-router-agent") -> None: + self.enabled = enabled + self.required_subagent = required_subagent + + @staticmethod + def _routing_complete(state: object) -> bool: + files = state.get("files", {}) if isinstance(state, dict) else getattr(state, "files", {}) + return isinstance(files, dict) and _SOURCE_ROUTING_PATH in files + + async def awrap_tool_call(self, request, handler): + """Block out-of-order calls until the source router writes its route file.""" + if not self.enabled or self._routing_complete(request.state): + return await handler(request) + + tool_call = request.tool_call + args = tool_call.get("args") or {} + if tool_call.get("name") == "task" and args.get("subagent_type") == self.required_subagent: + return await handler(request) + + return ToolMessage( + content=( + "Source routing is required before any other tool call. " + f"Call task with subagent_type={self.required_subagent!r}." + ), + tool_call_id=tool_call.get("id", "source-routing-guard"), + name=tool_call.get("name"), + status="error", + ) class EmptyContentFixMiddleware(AgentMiddleware): diff --git a/src/aiq_agent/agents/deep_researcher/factory.py b/src/aiq_agent/agents/deep_researcher/factory.py index 88dc6ab8c..cf08c4ae3 100644 --- a/src/aiq_agent/agents/deep_researcher/factory.py +++ b/src/aiq_agent/agents/deep_researcher/factory.py @@ -44,6 +44,7 @@ from .custom_middleware import EmptyContentFixMiddleware from .custom_middleware import PlanPersistenceMiddleware from .custom_middleware import SourceRegistryMiddleware +from .custom_middleware import SourceRoutingGuardMiddleware from .custom_middleware import TodoSuppressionMiddleware from .custom_middleware import ToolNameSanitizationMiddleware from .custom_middleware import ToolResultPruningMiddleware @@ -215,6 +216,37 @@ def build_common_middleware( ] +def build_orchestrator_middleware( + *, + tool_set: DeepResearchToolSet, + source_registry_middleware: SourceRegistryMiddleware, + enable_source_router: bool, + research_batch_tool_name: str = "run_research_batch", +) -> list[Any]: + """Middleware for the orchestrator. + + Unlike the shared stack, the tool-name sanitizer allowlist here is restricted + to the tools the orchestrator is actually bound to — helper tools, + ``run_research_batch``, and the filesystem tools — deliberately excluding the + source tools. The orchestrator routes all source access through + ``run_research_batch`` (which runs the researcher, where the source tools + live), so a source-tool name emitted by the orchestrator must not be treated + as a valid direct call. + """ + valid_tool_names = {tool.name for tool in tool_set.helper_tools} + valid_tool_names.add(research_batch_tool_name) + valid_tool_names.update(FILESYSTEM_TOOL_NAMES) + return [ + EmptyContentFixMiddleware(), + SourceRoutingGuardMiddleware(enabled=enable_source_router, required_subagent=SOURCE_ROUTER_AGENT), + ToolNameSanitizationMiddleware(valid_tool_names=sorted(valid_tool_names)), + ToolRetryMiddleware(max_retries=3, backoff_factor=2.0, initial_delay=1.0), + source_registry_middleware, + ToolResultPruningMiddleware(keep_last_n=10, max_chars=2000), + ModelRetryMiddleware(max_retries=2, backoff_factor=2.0, initial_delay=1.0), + ] + + def build_source_router_middleware(*, extra_valid_tool_names: Sequence[str] = ()) -> list[Any]: """Build minimal middleware for the source-router-agent.""" return [ @@ -229,6 +261,7 @@ def build_deep_research_middleware_set( *, tool_set: DeepResearchToolSet, source_registry_middleware: SourceRegistryMiddleware, + enable_source_router: bool = True, ) -> DeepResearchMiddlewareSet: """Build researcher, writer, and orchestrator middleware stacks.""" @@ -244,7 +277,11 @@ def common(extra_valid_tool_names: Sequence[str] = ()) -> list[Any]: researcher=common(), planner=common(), writer=common(), - orchestrator=common(["run_research_batch"]), + orchestrator=build_orchestrator_middleware( + tool_set=tool_set, + source_registry_middleware=source_registry_middleware, + enable_source_router=enable_source_router, + ), ) @@ -484,13 +521,21 @@ def build_deep_research_graph( source_registry_middleware=source_registry_middleware, ) + orchestrator_tools = [*context.tool_set.helper_tools, research_batch_tool] agent = create_deep_agent( model=context.llm_provider.get(LLMRole.ORCHESTRATOR), - tools=[*context.tool_set.helper_tools, research_batch_tool], + tools=orchestrator_tools, system_prompt=context.render_prompt( "orchestrator", clarifier_result=context.state.clarifier_result, - tools=context.tool_set.tools_info, + # Advertise only the tools the orchestrator can actually call. Source + # tools (incl. per-user MCP tools like Google Drive) are NOT directly + # callable here — the orchestrator delegates all source access through + # run_research_batch to the researcher, which holds those tools. Listing + # them under "Available Tools" made the orchestrator call them directly + # (e.g. per_user_mcp_client__google_drive_read_file), which the runtime + # rejects since they aren't bound to this agent. + tools=[{"name": t.name, "description": t.description} for t in orchestrator_tools], enable_source_router=context.enable_source_router, max_research_concurrency=context.max_research_concurrency, execution_enabled=context.runtime.execution_enabled, diff --git a/src/aiq_agent/agents/shallow_researcher/register.py b/src/aiq_agent/agents/shallow_researcher/register.py index 4f5faef85..2baae940a 100644 --- a/src/aiq_agent/agents/shallow_researcher/register.py +++ b/src/aiq_agent/agents/shallow_researcher/register.py @@ -95,20 +95,58 @@ async def shallow_research_agent(config: ShallowResearchAgentConfig, builder: Bu verbose = is_verbose(config.verbose) callbacks = [VerboseTraceCallback()] if verbose else [] - agent = ShallowResearcherAgent( - llm_provider=provider, - tools=tools, - max_llm_turns=config.max_llm_turns, - max_tool_iterations=config.max_tool_iterations, - callbacks=callbacks, - ) + # No shared agent is built here: it is (re)built per request inside _run, since + # the active tool set depends on the request's data_sources and the per-user MCP + # tools resolved at run time. async def _run(state: ShallowResearchAgentState) -> ShallowResearchAgentState: + from contextlib import AsyncExitStack + try: data_sources = state.data_sources selected_tools = filter_tools_by_sources(tools, data_sources) - active_agent = agent - if data_sources is not None and selected_tools != tools: + + if all_mapped_tools_filtered_out(tools, selected_tools, data_sources): + logger.warning("Shallow research received data_sources with no matching tools") + + # Per-user MCP tools (e.g. a connected Google Drive) are resolved at RUN + # time: this runs per request with Context.user_id set, so we can build the + # user's MCP client and add its tools to this turn. Build-time inheritance + # can't (the agent would be a shared, user-less instance) — see + # aiq_api.mcp_auth.runtime_tools. The client stays open via mcp_stack for the run. + async with AsyncExitStack() as mcp_stack: + try: + from aiq_api.jobs.access import require_verified_principal + from aiq_api.mcp_auth.provider import principal_user_id + from aiq_api.mcp_auth.runtime_tools import PerUserMcpSourceUnavailableError + from aiq_api.mcp_auth.runtime_tools import open_per_user_mcp_tools + from nat.builder.context import ContextState + + # Align the MCP token-lookup key with where connect stored it. The + # interactive session sets Context.user_id to NAT's user id — a *different* + # derivation than principal_user_id (used by connect/status). Without this, + # per_user_mcp_client looks under the wrong key, finds no token, and triggers + # interactive re-auth even though the source shows connected. Set with no + # reset: the client reads Context.user_id at build and on each tool call, so + # it must stay set for the whole turn. + ContextState.get().user_id.set(principal_user_id(require_verified_principal())) + + mcp_tools = await open_per_user_mcp_tools( + builder=builder, data_sources=data_sources, exit_stack=mcp_stack + ) + if mcp_tools: + selected_tools = [*selected_tools, *mcp_tools] + except PerUserMcpSourceUnavailableError as exc: + # The user explicitly selected a protected source we can't resolve + # (e.g. token expired). Surface a reconnect message instead of + # silently answering without it. + from langchain_core.messages import AIMessage + + return ShallowResearchAgentState(messages=state.messages + [AIMessage(content=str(exc))]) + except Exception: + logger.exception("Failed to resolve per-user MCP tools for shallow research; continuing") + + # Build the agent with this turn's tool set. active_agent = ShallowResearcherAgent( llm_provider=provider, tools=selected_tools, @@ -117,32 +155,21 @@ async def _run(state: ShallowResearchAgentState) -> ShallowResearchAgentState: callbacks=callbacks, ) - if all_mapped_tools_filtered_out(tools, selected_tools, data_sources): - logger.warning("Shallow research received data_sources with no matching tools") - - # Validate tool availability before starting shallow research - # At least one tool must be available - # This prevents the agent from trying to reason about unavailable tools - # Check selected_tools directly - they already reflect data_sources filtering - from aiq_agent.common import format_user_facing_tool_error - from aiq_agent.common import validate_tool_availability + # Require at least one available tool, else the agent would reason about + # tools it can't call. selected_tools already reflects filtering + MCP. + from aiq_agent.common import format_user_facing_tool_error + from aiq_agent.common import validate_tool_availability - is_valid, _, unavailable_tools = validate_tool_availability( - selected_tools, research_type="shallow research" - ) - - # Fail if no tools are available - if not is_valid: - error_msg = format_user_facing_tool_error("shallow research", unavailable_tools) - - # Return error state with error message - this prevents the agent from running - from langchain_core.messages import AIMessage + is_valid, _, unavailable_tools = validate_tool_availability( + selected_tools, research_type="shallow research" + ) + if not is_valid: + error_msg = format_user_facing_tool_error("shallow research", unavailable_tools) + from langchain_core.messages import AIMessage - error_state = ShallowResearchAgentState(messages=state.messages + [AIMessage(content=error_msg)]) - return error_state + return ShallowResearchAgentState(messages=state.messages + [AIMessage(content=error_msg)]) - result = await active_agent.run(state) - return result + return await active_agent.run(state) except Exception: logger.exception("Error in shallow research execution.") raise diff --git a/src/aiq_agent/common/data_source_registry.py b/src/aiq_agent/common/data_source_registry.py index 311660b50..9cbd98d99 100644 --- a/src/aiq_agent/common/data_source_registry.py +++ b/src/aiq_agent/common/data_source_registry.py @@ -48,6 +48,7 @@ import logging from dataclasses import dataclass from typing import Any +from typing import Literal from pydantic import BaseModel from pydantic import Field @@ -63,6 +64,42 @@ logger = logging.getLogger(__name__) +class PerUserAuthConfig(BaseModel): + """Static per-user MCP auth declaration for a protected data source. + + This is the configuration-time metadata (declared in YAML). The dynamic, + per-request state (status / connect_url / auth_url / expires_at) is computed + by the API layer from this config plus the user's token storage state — see + ``aiq_api.mcp_auth``. + """ + + required: bool = Field(default=False, description="Whether per-user MCP auth is required for this source") + type: Literal["mcp_oauth2"] = Field(default="mcp_oauth2", description="Auth mechanism (only mcp_oauth2 today)") + provider: str | None = Field(default=None, description="Provider identifier, e.g. 'google'") + mcp_server_id: str | None = Field( + default=None, + description="MCP server / auth-provider key used to look up the user's token in NAT token storage", + ) + auth_provider: str | None = Field( + default=None, + description=( + "Name of the NAT `authentication` (mcp_oauth2) provider for this source, resolved via the builder " + "to derive OAuth settings + shared token storage (connect flow) and to build the per-user MCP " + "client in the worker at job time. Defaults to mcp_server_id when unset. NOTE: no config-declared " + "per_user_mcp_client function group is used — see aiq_api.mcp_auth.runtime_tools for why (it breaks " + "the interactive WebSocket path)." + ), + ) + tool_overrides: dict[str, dict[str, str]] | None = Field( + default=None, + description=( + "Optional per-tool overrides for this source's MCP tools, keyed by the server's tool name, each " + "mapping to {alias, description}. Use to give terse/blank MCP tools clear names and descriptions so " + "the agent reliably selects them (e.g. gdrive_get_file -> a descriptive 'read a Drive file' tool)." + ), + ) + + @dataclass(frozen=True) class DataSourceMeta: """Metadata for a registered data source.""" @@ -72,6 +109,7 @@ class DataSourceMeta: description: str default_enabled: bool = True requires_auth: bool = False + per_user_auth: PerUserAuthConfig | None = None # ── Global state ── @@ -96,6 +134,10 @@ class DataSourceEntry(BaseModel): description: str = Field(default="", description="Human-readable description") default_enabled: bool = Field(default=True, description="Whether enabled by default") requires_auth: bool = Field(default=False, description="Whether the source requires user authentication") + per_user_auth: PerUserAuthConfig | None = Field( + default=None, + description="Per-user MCP OAuth declaration for a protected source (see PerUserAuthConfig)", + ) tools: list[FunctionRef] = Field( default_factory=list, description="NAT functions or function groups that belong to this data source", @@ -126,12 +168,16 @@ def _populate(sources: list[dict], group_names: set[str] | None = None) -> None: for entry in sources: source_id = entry["id"] + per_user_auth = entry.get("per_user_auth") + if per_user_auth is not None and not isinstance(per_user_auth, PerUserAuthConfig): + per_user_auth = PerUserAuthConfig.model_validate(per_user_auth) _registry[source_id] = DataSourceMeta( id=source_id, name=entry.get("name", source_id.replace("_", " ").title()), description=entry.get("description", ""), default_enabled=entry.get("default_enabled", True), requires_auth=entry.get("requires_auth", False), + per_user_auth=per_user_auth, ) for ref_name in entry.get("tools", []): if ref_name in group_names: @@ -160,6 +206,7 @@ async def data_source_registry_fn(config: DataSourceRegistryConfig, builder: Any "description": entry.description, "default_enabled": entry.default_enabled, "requires_auth": entry.requires_auth, + "per_user_auth": entry.per_user_auth, "tools": [str(ref) for ref in entry.tools], } for entry in config.sources @@ -248,6 +295,40 @@ def set_group_source_map(mapping: dict[str, str]) -> None: _rebuild_prefix_index() +def register_tool_sources(mapping: dict[str, str]) -> None: + """Merge additional exact tool→source mappings into the registry. + + Used for tools that only become known at runtime — notably a per-user MCP + source's tools (resolved per request), which aren't in the static config. + Registering them lets ``get_source_id_for_tool`` resolve their source so + citation/source capture treats their results as sources (otherwise shallow + research raises EmptySourceRegistryError). Keys are deterministic, so merging + globally is idempotent. + + Re-registering the same tool→source pair is a no-op. A key that already maps + to a *different* source is a collision (e.g. two sources exposing a tool with + the same name); we keep the existing owner and skip the conflicting entry + rather than silently reassigning citation ownership. + """ + if not mapping: + return + global _tool_source_map + merged = dict(_tool_source_map) + for tool, source in mapping.items(): + existing = merged.get(tool) + if existing is not None and existing != source: + logger.warning( + "Tool '%s' already maps to source '%s'; skipping conflicting registration to '%s'.", + tool, + existing, + source, + ) + continue + merged[tool] = source + _tool_source_map = merged + _rebuild_prefix_index() + + # ── Test helpers ── def populate_from_config(sources: list[dict], group_names: set[str] | None = None) -> None: """Populate registry and maps directly (for testing without NAT startup). diff --git a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py index 61d0376ce..53e691024 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py +++ b/tests/aiq_agent/agents/deep_researcher/test_custom_middleware.py @@ -27,6 +27,7 @@ from aiq_agent.agents.deep_researcher.custom_middleware import PlanPersistenceMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import SourceRegistryMiddleware +from aiq_agent.agents.deep_researcher.custom_middleware import SourceRoutingGuardMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import TodoSuppressionMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import ToolNameSanitizationMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import ToolVisibilityMiddleware @@ -36,6 +37,72 @@ from aiq_agent.common.data_source_registry import reset_registry +class TestSourceRoutingGuardMiddleware: + """Tests for the orchestrator's required source-routing transition.""" + + @staticmethod + def _request(tool_name: str, *, args: dict | None = None, files: dict | None = None) -> MagicMock: + request = MagicMock() + request.tool_call = { + "name": tool_name, + "args": args or {}, + "id": "tc1", + } + request.state = {"files": files or {}} + return request + + @pytest.mark.asyncio + async def test_blocks_other_tools_before_source_routing(self): + """An orchestrator cannot infer source absence from filesystem inspection before routing.""" + middleware = SourceRoutingGuardMiddleware(enabled=True) + handler = AsyncMock(return_value=ToolMessage(content="[]", tool_call_id="tc1")) + + result = await middleware.awrap_tool_call(self._request("ls", args={"path": "/shared"}), handler) + + handler.assert_not_awaited() + assert result.status == "error" + assert "source-router-agent" in str(result.content) + + @pytest.mark.asyncio + async def test_allows_source_router_task_before_routing(self): + """The required source-router task remains executable while the gate is closed.""" + middleware = SourceRoutingGuardMiddleware(enabled=True) + expected = ToolMessage(content="Source routing complete.", tool_call_id="tc1") + handler = AsyncMock(return_value=expected) + request = self._request("task", args={"subagent_type": "source-router-agent"}) + + result = await middleware.awrap_tool_call(request, handler) + + handler.assert_awaited_once_with(request) + assert result is expected + + @pytest.mark.asyncio + async def test_allows_normal_tools_after_routing_file_exists(self): + """The gate opens once the source-router output is present in virtual state.""" + middleware = SourceRoutingGuardMiddleware(enabled=True) + expected = ToolMessage(content="[]", tool_call_id="tc1") + handler = AsyncMock(return_value=expected) + request = self._request("ls", files={"/shared/source_routing.json": {"content": "{}"}}) + + result = await middleware.awrap_tool_call(request, handler) + + handler.assert_awaited_once_with(request) + assert result is expected + + @pytest.mark.asyncio + async def test_disabled_guard_is_noop(self): + """Workflows with source routing disabled preserve their existing tool behavior.""" + middleware = SourceRoutingGuardMiddleware(enabled=False) + expected = ToolMessage(content="[]", tool_call_id="tc1") + handler = AsyncMock(return_value=expected) + request = self._request("ls") + + result = await middleware.awrap_tool_call(request, handler) + + handler.assert_awaited_once_with(request) + assert result is expected + + class TestToolNameSanitizationMiddleware: """Tests for ToolNameSanitizationMiddleware.""" diff --git a/tests/aiq_agent/agents/deep_researcher/test_factory.py b/tests/aiq_agent/agents/deep_researcher/test_factory.py index 72614b1c2..7b25c0cfe 100644 --- a/tests/aiq_agent/agents/deep_researcher/test_factory.py +++ b/tests/aiq_agent/agents/deep_researcher/test_factory.py @@ -24,6 +24,7 @@ from langchain_core.tools import tool from aiq_agent.agents.deep_researcher.custom_middleware import SourceRegistryMiddleware +from aiq_agent.agents.deep_researcher.custom_middleware import SourceRoutingGuardMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import ToolNameSanitizationMiddleware from aiq_agent.agents.deep_researcher.custom_middleware import ToolVisibilityMiddleware from aiq_agent.agents.deep_researcher.deepagents_runtime import DeepAgentsRuntime @@ -94,6 +95,10 @@ def _sanitizer(middleware: list[object]) -> ToolNameSanitizationMiddleware: return next(item for item in middleware if isinstance(item, ToolNameSanitizationMiddleware)) +def _routing_guard(middleware: list[object]) -> SourceRoutingGuardMiddleware: + return next(item for item in middleware if isinstance(item, SourceRoutingGuardMiddleware)) + + def _graph_context( *, runtime: DeepAgentsRuntime | None = None, @@ -145,11 +150,41 @@ def test_middleware_set_adds_orchestrator_batch_tool_name(): assert "write_file" in researcher_sanitizer.valid_tool_names assert "run_research_batch" not in researcher_sanitizer.valid_tool_names assert "run_research_batch" in orchestrator_sanitizer.valid_tool_names + # The orchestrator allowlist matches the tools it is actually bound to: helper + # tools, run_research_batch, and filesystem tools — but NOT source tools, which + # it must reach only via run_research_batch (not call directly). + assert "web_search_tool" not in orchestrator_sanitizer.valid_tool_names + assert "think" in orchestrator_sanitizer.valid_tool_names + assert "get_verified_sources" in orchestrator_sanitizer.valid_tool_names + assert "read_file" in orchestrator_sanitizer.valid_tool_names assert registry in middleware_set.researcher assert registry in middleware_set.writer + assert registry in middleware_set.orchestrator assert tool_set.writer_tools != tool_set.researcher_tools +def test_middleware_set_configures_orchestrator_source_routing_guard(): + """Only the orchestrator stack enforces the configured source-router transition.""" + registry, tool_set, _ = _tool_set_and_middleware() + + enabled = build_deep_research_middleware_set( + tool_set=tool_set, + source_registry_middleware=registry, + enable_source_router=True, + ) + disabled = build_deep_research_middleware_set( + tool_set=tool_set, + source_registry_middleware=registry, + enable_source_router=False, + ) + + assert _routing_guard(enabled.orchestrator).enabled is True + assert _routing_guard(disabled.orchestrator).enabled is False + assert not any(isinstance(item, SourceRoutingGuardMiddleware) for item in enabled.researcher) + assert not any(isinstance(item, SourceRoutingGuardMiddleware) for item in enabled.planner) + assert not any(isinstance(item, SourceRoutingGuardMiddleware) for item in enabled.writer) + + def test_subagents_route_tools_and_writer_skills(): """Source-router excludes source tools, planner receives source tools, and writer receives configured skills.""" runtime = DeepAgentsRuntime( diff --git a/tests/aiq_agent/common/test_data_source_registry.py b/tests/aiq_agent/common/test_data_source_registry.py index d432334f9..b54adbf45 100644 --- a/tests/aiq_agent/common/test_data_source_registry.py +++ b/tests/aiq_agent/common/test_data_source_registry.py @@ -24,6 +24,7 @@ from aiq_agent.common.data_source_registry import get_source from aiq_agent.common.data_source_registry import get_source_id_for_tool from aiq_agent.common.data_source_registry import populate_from_config +from aiq_agent.common.data_source_registry import register_tool_sources from aiq_agent.common.data_source_registry import reset_registry from aiq_agent.common.data_source_registry import set_group_source_map from aiq_agent.common.data_source_registry import set_tool_source_map @@ -166,6 +167,27 @@ def test_returns_metadata_for_known(self): assert meta.name == "Known" +class TestRegisterToolSources: + def test_adds_new_mappings(self): + register_tool_sources({"gdrive_search": "gdrive"}) + assert get_source_id_for_tool("gdrive_search") == "gdrive" + + def test_reregistering_same_pair_is_idempotent(self): + register_tool_sources({"gdrive_search": "gdrive"}) + register_tool_sources({"gdrive_search": "gdrive"}) + assert get_source_id_for_tool("gdrive_search") == "gdrive" + + def test_collision_keeps_existing_owner_and_skips(self, caplog): + register_tool_sources({"shared_tool": "source_a"}) + with caplog.at_level("WARNING"): + register_tool_sources({"shared_tool": "source_b", "new_tool": "source_b"}) + + # Existing owner preserved; conflicting remap skipped; unrelated key still added. + assert get_source_id_for_tool("shared_tool") == "source_a" + assert get_source_id_for_tool("new_tool") == "source_b" + assert "shared_tool" in caplog.text + + class TestToolSourceMap: def test_get_source_id_for_tool_returns_none_when_empty(self): assert get_source_id_for_tool("any_tool") is None diff --git a/tests/aiq_agent/fastapi_extensions/test_deep_research.py b/tests/aiq_agent/fastapi_extensions/test_deep_research.py index b8fb6c682..7c9fc7001 100644 --- a/tests/aiq_agent/fastapi_extensions/test_deep_research.py +++ b/tests/aiq_agent/fastapi_extensions/test_deep_research.py @@ -204,8 +204,18 @@ async def test_routes_not_registered_without_dask(self): await register_job_routes(mock_app, mock_builder, mock_worker) - mock_app.post.assert_not_called() - assert mock_app.get.call_count == 2 + # Async job submission/control routes require Dask + a job store and must + # NOT be registered. The always-on control-plane routes (data sources, + # agent list, and per-user MCP auth status/connect/callback) are still + # registered regardless of Dask availability. + post_paths = [c.args[0] for c in mock_app.post.call_args_list if c.args] + assert "/v1/jobs/async/submit" not in post_paths + assert "/v1/auth/mcp/{source_id}/connect" in post_paths + get_paths = [c.args[0] for c in mock_app.get.call_args_list if c.args] + assert "/v1/jobs/async/agents" in get_paths + assert "/v1/data_sources" in get_paths + assert "/v1/auth/mcp/{source_id}/status" in get_paths + assert "/v1/auth/mcp/{source_id}/callback" in get_paths @pytest.mark.asyncio async def test_routes_not_registered_without_job_store(self): @@ -221,8 +231,18 @@ async def test_routes_not_registered_without_job_store(self): await register_job_routes(mock_app, mock_builder, mock_worker) - mock_app.post.assert_not_called() - assert mock_app.get.call_count == 2 + # Async job submission/control routes require Dask + a job store and must + # NOT be registered. The always-on control-plane routes (data sources, + # agent list, and per-user MCP auth status/connect/callback) are still + # registered regardless of Dask availability. + post_paths = [c.args[0] for c in mock_app.post.call_args_list if c.args] + assert "/v1/jobs/async/submit" not in post_paths + assert "/v1/auth/mcp/{source_id}/connect" in post_paths + get_paths = [c.args[0] for c in mock_app.get.call_args_list if c.args] + assert "/v1/jobs/async/agents" in get_paths + assert "/v1/data_sources" in get_paths + assert "/v1/auth/mcp/{source_id}/status" in get_paths + assert "/v1/auth/mcp/{source_id}/callback" in get_paths @pytest.mark.asyncio async def test_routes_registered_with_dask(self): diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index 8579ab59c..e1d6da7d6 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -469,8 +469,9 @@ async def test_submit_agent_job_passes_initial_files_and_output_metadata(self): assert result == "test-job-id" job_args = mock_job_store.submit_job.call_args.kwargs["job_args"] - assert job_args[-2] == initial_files - assert job_args[-1] == output_metadata + # Trailing worker args: data_sources, auth_token, initial_files, output_metadata, principal_user_id + assert job_args[-3] == initial_files + assert job_args[-2] == output_metadata @pytest.mark.asyncio async def test_submit_with_custom_job_id(self): diff --git a/tests/deploy/test_compose_deployment.py b/tests/deploy/test_compose_deployment.py new file mode 100644 index 000000000..03034c676 --- /dev/null +++ b/tests/deploy/test_compose_deployment.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import shutil +import subprocess +from pathlib import Path +from typing import Any + +import pytest +import yaml + +from nat.utils.io.yaml_tools import yaml_load + +REPO_ROOT = Path(__file__).resolve().parents[2] +COMPOSE_PATH = REPO_ROOT / "deploy" / "compose" / "docker-compose.yaml" +PER_USER_AUTH_COMPOSE_PATH = REPO_ROOT / "deploy" / "compose" / "docker-compose.per-user-auth.yaml" +MCP_CONFIG_PATH = REPO_ROOT / "configs" / "config_web_frag_mcp_auth.yml" + + +def load_compose() -> dict[str, Any]: + with COMPOSE_PATH.open(encoding="utf-8") as compose_file: + return yaml.safe_load(compose_file) + + +def test_default_compose_does_not_provision_or_configure_redis(): + compose = load_compose() + services = compose["services"] + backend = services["aiq-agent"] + + assert "redis" not in services + assert "redis-data" not in compose.get("volumes", {}) + assert "redis" not in backend.get("depends_on", {}) + + backend_env = { + entry.split("=", maxsplit=1)[0] for entry in backend.get("environment", []) if isinstance(entry, str) + } + assert backend_env.isdisjoint({"MCP_TOKEN_STORE_TYPE", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD"}) + + +def test_per_user_auth_compose_adds_private_redis_token_store(tmp_path: Path): + if shutil.which("docker") is None: + pytest.skip("docker is required to validate Compose merge behavior") + + compose_version = subprocess.run( + ["docker", "compose", "version"], + capture_output=True, + check=False, + text=True, + ) + if compose_version.returncode != 0: + pytest.skip("docker compose is required to validate Compose merge behavior") + + # The runtime env file is intentionally untracked. Render from a temporary + # copy so this merge test works in a clean checkout without developer secrets. + base_compose = load_compose() + base_compose["services"]["aiq-agent"].pop("env_file", None) + test_compose_path = tmp_path / "docker-compose.yaml" + test_compose_path.write_text(yaml.safe_dump(base_compose), encoding="utf-8") + + result = subprocess.run( + [ + "docker", + "compose", + "--env-file", + str(REPO_ROOT / "deploy" / ".env.example"), + "-f", + str(test_compose_path), + "-f", + str(PER_USER_AUTH_COMPOSE_PATH), + "config", + "--format", + "json", + "--no-env-resolution", + ], + capture_output=True, + check=True, + text=True, + ) + compose = json.loads(result.stdout) + backend = compose["services"]["aiq-agent"] + redis = compose["services"]["redis"] + + assert "redis-data" in compose["volumes"] + assert backend["depends_on"]["redis"]["condition"] == "service_healthy" + assert backend["environment"]["CONFIG_FILE"] == "/app/configs/config_web_frag_mcp_auth.yml" + assert backend["environment"]["MCP_TOKEN_STORE_TYPE"] == "redis" + assert backend["environment"]["REDIS_HOST"] == "redis" + assert backend["environment"]["REDIS_PORT"] == "6379" + assert "ports" not in redis + + +def test_mcp_example_config_accepts_an_optional_external_redis_password(monkeypatch): + monkeypatch.delenv("MCP_TOKEN_STORE_TYPE", raising=False) + monkeypatch.delenv("REDIS_PASSWORD", raising=False) + + without_password = yaml_load(MCP_CONFIG_PATH)["object_stores"]["mcp_token_store"] + assert without_password["_type"] == "aiq_sqlite" + assert without_password["password"] is None + + monkeypatch.setenv("MCP_TOKEN_STORE_TYPE", "redis") + test_password = "managed-redis-password" # pragma: allowlist secret + monkeypatch.setenv("REDIS_PASSWORD", test_password) + with_password = yaml_load(MCP_CONFIG_PATH)["object_stores"]["mcp_token_store"] + assert with_password["_type"] == "redis" + assert with_password["password"] == test_password diff --git a/tests/deploy/test_helm_deployment_k8s.py b/tests/deploy/test_helm_deployment_k8s.py index 2401d2ea7..4a18f0a0e 100644 --- a/tests/deploy/test_helm_deployment_k8s.py +++ b/tests/deploy/test_helm_deployment_k8s.py @@ -94,3 +94,24 @@ def test_chart_renders_app_host_aliases(tmp_path: Path): assert backend_deployment["spec"]["template"]["spec"]["hostAliases"] == [ {"ip": "127.0.0.1", "hostnames": ["aiq.local"]} ] + + +def test_default_chart_does_not_provision_or_configure_redis(): + manifests = render_chart() + + redis_resources = { + (manifest.get("kind"), manifest.get("metadata", {}).get("name")) + for manifest in manifests + if manifest.get("metadata", {}).get("name", "").startswith("aiq-redis") + } + assert redis_resources == set() + + backend_deployment = next( + manifest + for manifest in manifests + if manifest.get("kind") == "Deployment" and manifest["metadata"]["name"] == "aiq-backend" + ) + backend_env = { + item["name"] for item in backend_deployment["spec"]["template"]["spec"]["containers"][0].get("env", []) + } + assert backend_env.isdisjoint({"MCP_TOKEN_STORE_TYPE", "REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD"}) diff --git a/uv.lock b/uv.lock index 563d2512b..7c874c272 100644 --- a/uv.lock +++ b/uv.lock @@ -7,6 +7,9 @@ resolution-markers = [ "python_full_version < '3.12'", ] +[options] +prerelease-mode = "allow" + [manifest] members = [ "aiq-agent", @@ -205,6 +208,7 @@ dependencies = [ { name = "nvidia-nat-core" }, { name = "nvidia-nat-eval" }, { name = "nvidia-nat-profiler" }, + { name = "nvidia-nat-redis" }, { name = "nvidia-nat-security", extra = ["guardrails"] }, { name = "pillow" }, { name = "psycopg", extra = ["binary"] }, @@ -278,6 +282,7 @@ requires-dist = [ { name = "nvidia-nat-core", specifier = "==1.8.0" }, { name = "nvidia-nat-eval", specifier = "==1.8.0" }, { name = "nvidia-nat-profiler", specifier = "==1.8.0" }, + { name = "nvidia-nat-redis", specifier = "==1.8.0" }, { name = "nvidia-nat-security", extras = ["guardrails"], specifier = "==1.8.0" }, { name = "nvidia-sphinx-theme", marker = "extra == 'docs'", specifier = ">=0.0.9.post1" }, { name = "pillow", specifier = ">=12.2.0,<13" }, @@ -519,6 +524,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/45/19/cc8bd127d28a43da249aa955cfd164cf8fd534e79e42cea96c4854d72fd0/ast_serialize-0.5.0-cp39-abi3-win_arm64.whl", hash = "sha256:92a31c9c20d25a076edaeec76b128a3535d74a24f340b9a8a7e96c9b86dc9642", size = 1081181, upload-time = "2026-05-17T17:48:28.122Z" }, ] +[[package]] +name = "async-timeout" +version = "5.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a5/ae/136395dfbfe00dfc94da3f3e136d0b13f394cba8f4841120e34226265780/async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3", size = 9274, upload-time = "2024-11-06T16:41:39.6Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/ba/e2081de779ca30d473f21f5b30e0e737c438205440784c7dfc81efc2b029/async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c", size = 6233, upload-time = "2024-11-06T16:41:37.9Z" }, +] + [[package]] name = "asyncpg" version = "0.31.0" @@ -669,38 +683,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/88/c6/92fcd42f1ba33e1184263f25bfabf3d27c383410470f169e4b8163bf9c17/beautifulsoup4-4.15.0-py3-none-any.whl", hash = "sha256:d6f88de62e1d4e38ecb1077eb9724cd0eff29d2a08ca16a401e9b9e93f117cf9", size = 109924, upload-time = "2026-06-07T16:44:21.566Z" }, ] -[[package]] -name = "blis" -version = "1.3.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d0/d0/d8cc8c9a4488a787e7fa430f6055e5bd1ddb22c340a751d9e901b82e2efe/blis-1.3.3.tar.gz", hash = "sha256:034d4560ff3cc43e8aa37e188451b0440e3261d989bb8a42ceee865607715ecd", size = 2644873, upload-time = "2025-11-17T12:28:30.511Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a1/0a/a4c8736bc497d386b0ffc76d321f478c03f1a4725e52092f93b38beb3786/blis-1.3.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:e10c8d3e892b1dbdff365b9d00e08291876fc336915bf1a5e9f188ed087e1a91", size = 6925522, upload-time = "2025-11-17T12:27:29.199Z" }, - { url = "https://files.pythonhosted.org/packages/83/5a/3437009282f23684ecd3963a8b034f9307cdd2bf4484972e5a6b096bf9ac/blis-1.3.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:66e6249564f1db22e8af1e0513ff64134041fa7e03c8dd73df74db3f4d8415a7", size = 1232787, upload-time = "2025-11-17T12:27:30.996Z" }, - { url = "https://files.pythonhosted.org/packages/d1/0e/82221910d16259ce3017c1442c468a3f206a4143a96fbba9f5b5b81d62e8/blis-1.3.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7260da065958b4e5475f62f44895ef9d673b0f47dcf61b672b22b7dae1a18505", size = 2844596, upload-time = "2025-11-17T12:27:32.601Z" }, - { url = "https://files.pythonhosted.org/packages/6c/93/ab547f1a5c23e20bca16fbcf04021c32aac3f969be737ea4980509a7ca90/blis-1.3.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9327a6ca67de8ae76fe071e8584cc7f3b2e8bfadece4961d40f2826e1cda2df", size = 11377746, upload-time = "2025-11-17T12:27:35.342Z" }, - { url = "https://files.pythonhosted.org/packages/6e/a6/7733820aa62da32526287a63cd85c103b2b323b186c8ee43b7772ff7017c/blis-1.3.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:c4ae70629cf302035d268858a10ca4eb6242a01b2dc8d64422f8e6dcb8a8ee74", size = 3041954, upload-time = "2025-11-17T12:27:37.479Z" }, - { url = "https://files.pythonhosted.org/packages/87/53/e39d67fd3296b649772780ca6aab081412838ecb54e0b0c6432d01626a50/blis-1.3.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:45866a9027d43b93e8b59980a23c5d7358b6536fc04606286e39fdcfce1101c2", size = 14251222, upload-time = "2025-11-17T12:27:39.705Z" }, - { url = "https://files.pythonhosted.org/packages/ea/44/b749f8777b020b420bceaaf60f66432fc30cc904ca5b69640ec9cbef11ed/blis-1.3.3-cp311-cp311-win_amd64.whl", hash = "sha256:27f82b8633030f8d095d2b412dffa7eb6dbc8ee43813139909a20012e54422ea", size = 6171233, upload-time = "2025-11-17T12:27:41.921Z" }, - { url = "https://files.pythonhosted.org/packages/16/d1/429cf0cf693d4c7dc2efed969bd474e315aab636e4a95f66c4ed7264912d/blis-1.3.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2a1c74e100665f8e918ebdbae2794576adf1f691680b5cdb8b29578432f623ef", size = 6929663, upload-time = "2025-11-17T12:27:44.482Z" }, - { url = "https://files.pythonhosted.org/packages/11/69/363c8df8d98b3cc97be19aad6aabb2c9c53f372490d79316bdee92d476e7/blis-1.3.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3f6c595185176ce021316263e1a1d636a3425b6c48366c1fd712d08d0b71849a", size = 1230939, upload-time = "2025-11-17T12:27:46.19Z" }, - { url = "https://files.pythonhosted.org/packages/96/2a/fbf65d906d823d839076c5150a6f8eb5ecbc5f9135e0b6510609bda1e6b7/blis-1.3.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d734b19fba0be7944f272dfa7b443b37c61f9476d9ab054a9ac53555ceadd2e0", size = 2818835, upload-time = "2025-11-17T12:27:48.167Z" }, - { url = "https://files.pythonhosted.org/packages/d5/ad/58deaa3ad856dd3cc96493e40ffd2ed043d18d4d304f85a65cde1ccbf644/blis-1.3.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ef6d6e2b599a3a2788eb6d9b443533961265aa4ec49d574ed4bb846e548dcdb", size = 11366550, upload-time = "2025-11-17T12:27:49.958Z" }, - { url = "https://files.pythonhosted.org/packages/78/82/816a7adfe1f7acc8151f01ec86ef64467a3c833932d8f19f8e06613b8a4e/blis-1.3.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8c888438ae99c500422d50698e3028b65caa8ebb44e24204d87fda2df64058f7", size = 3023686, upload-time = "2025-11-17T12:27:52.062Z" }, - { url = "https://files.pythonhosted.org/packages/1e/e2/0e93b865f648b5519360846669a35f28ee8f4e1d93d054f6850d8afbabde/blis-1.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8177879fd3590b5eecdd377f9deafb5dc8af6d684f065bd01553302fb3fcf9a7", size = 14250939, upload-time = "2025-11-17T12:27:53.847Z" }, - { url = "https://files.pythonhosted.org/packages/20/07/fb43edc2ff0a6a367e4a94fc39eb3b85aa1e55e24cc857af2db145ce9f0d/blis-1.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:f20f7ad69aaffd1ce14fe77de557b6df9b61e0c9e582f75a843715d836b5c8af", size = 6192759, upload-time = "2025-11-17T12:27:56.176Z" }, - { url = "https://files.pythonhosted.org/packages/e6/f7/d26e62d9be3d70473a63e0a5d30bae49c2fe138bebac224adddcdef8a7ce/blis-1.3.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1e647341f958421a86b028a2efe16ce19c67dba2a05f79e8f7e80b1ff45328aa", size = 6928322, upload-time = "2025-11-17T12:27:57.965Z" }, - { url = "https://files.pythonhosted.org/packages/4a/78/750d12da388f714958eb2f2fd177652323bbe7ec528365c37129edd6eb84/blis-1.3.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d563160f874abb78a57e346f07312c5323f7ad67b6370052b6b17087ef234a8e", size = 1229635, upload-time = "2025-11-17T12:28:00.118Z" }, - { url = "https://files.pythonhosted.org/packages/e8/36/eac4199c5b200a5f3e93cad197da8d26d909f218eb444c4f552647c95240/blis-1.3.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:30b8a5b90cb6cb81d1ada9ae05aa55fb8e70d9a0ae9db40d2401bb9c1c8f14c4", size = 2815650, upload-time = "2025-11-17T12:28:02.544Z" }, - { url = "https://files.pythonhosted.org/packages/bf/51/472e7b36a6bedb5242a9757e7486f702c3619eff76e256735d0c8b1679c6/blis-1.3.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9f5c53b277f6ac5b3ca30bc12ebab7ea16c8f8c36b14428abb56924213dc127", size = 11359008, upload-time = "2025-11-17T12:28:04.589Z" }, - { url = "https://files.pythonhosted.org/packages/84/da/d0dfb6d6e6321ae44df0321384c32c322bd07b15740d7422727a1a49fc5d/blis-1.3.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6297e7616c158b305c9a8a4e47ca5fc9b0785194dd96c903b1a1591a7ca21ddf", size = 3011959, upload-time = "2025-11-17T12:28:06.862Z" }, - { url = "https://files.pythonhosted.org/packages/20/c5/2b0b5e556fa0364ed671051ea078a6d6d7b979b1cfef78d64ad3ca5f0c7f/blis-1.3.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:3f966ca74f89f8a33e568b9a1d71992fc9a0d29a423e047f0a212643e21b5458", size = 14232456, upload-time = "2025-11-17T12:28:08.779Z" }, - { url = "https://files.pythonhosted.org/packages/31/07/4cdc81a47bf862c0b06d91f1bc6782064e8b69ac9b5d4ff51d97e4ff03da/blis-1.3.3-cp313-cp313-win_amd64.whl", hash = "sha256:7a0fc4b237a3a453bdc3c7ab48d91439fcd2d013b665c46948d9eaf9c3e45a97", size = 6192624, upload-time = "2025-11-17T12:28:14.197Z" }, -] - [[package]] name = "boto3" version = "1.43.25" @@ -819,15 +801,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8c/7b/1fc1c09cc0756cf25861a3be10565915953876da48bb228fb9a672b20a42/cachetools-7.1.4-py3-none-any.whl", hash = "sha256:323dc4127934744db5b54eb4924482d7edafbf9554e820d1531c2e08c0e4ef54", size = 16761, upload-time = "2026-05-21T22:40:41.845Z" }, ] -[[package]] -name = "catalogue" -version = "2.0.10" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/38/b4/244d58127e1cdf04cf2dc7d9566f0d24ef01d5ce21811bab088ecc62b5ea/catalogue-2.0.10.tar.gz", hash = "sha256:4f56daa940913d3f09d589c191c74e5a6d51762b3a9e37dd53b7437afd6cda15", size = 19561, upload-time = "2023-09-25T06:29:24.962Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9e/96/d32b941a501ab566a16358d68b6eb4e4acc373fab3c3c4d7d9e649f7b4bb/catalogue-2.0.10-py3-none-any.whl", hash = "sha256:58c2de0020aa90f4a2da7dfad161bf7b3b054c86a5f09fcedc0b2b740c109a9f", size = 17325, upload-time = "2023-09-25T06:29:23.337Z" }, -] - [[package]] name = "cbor2" version = "6.1.2" @@ -1046,15 +1019,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/0d/67e5b4109ea4a837e80daa87c2c696711955e40449a97e8926672534def2/click-8.4.1-py3-none-any.whl", hash = "sha256:482be17c6991b8c19c5429a1e995d9b0efdbb63172824c41f99965dc0ade8ec2", size = 116639, upload-time = "2026-05-22T04:08:35.26Z" }, ] -[[package]] -name = "cloudpathlib" -version = "0.24.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/06/19/58bc6b5d7d0f81c7209b05445af477e147c486552f96665a5912211839b9/cloudpathlib-0.24.0.tar.gz", hash = "sha256:c521a984e77b47e656fe78e20a7e3e260e0ab45fc69e33ac01094227c979e34a", size = 53600, upload-time = "2026-04-30T00:54:43.265Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c2/5b/ba933f896d9b0b07608d575a8501e2b4e32166b60d84c430a4a7285ebe64/cloudpathlib-0.24.0-py3-none-any.whl", hash = "sha256:b1c51e2d2ec7dc4fed6538991f4aea849d6cf11a7e6b9069f86e461aa1f9b5b4", size = 63214, upload-time = "2026-04-30T00:54:42.06Z" }, -] - [[package]] name = "cloudpickle" version = "3.1.2" @@ -1073,15 +1037,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] -[[package]] -name = "confection" -version = "1.3.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ca/65/efd0fe8a936fc8ca2978cb7b82581fb20d901c6039e746a808f746b7647b/confection-1.3.3.tar.gz", hash = "sha256:f0f6810d567ff73993fe74d218ca5e1ffb6a44fb03f391257fc5d033546cbfaa", size = 54895, upload-time = "2026-03-24T18:45:24.331Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8d/e4/d66708bdf0d92fb4d49b22cdff4b10cec38aca5dcd7e81d909bb55c65cd7/confection-1.3.3-py3-none-any.whl", hash = "sha256:b9fef9ee84b237ef4611ec3eb5797b70e13063e6310ad9f15536373f5e313c82", size = 35902, upload-time = "2026-03-24T18:45:22.664Z" }, -] - [[package]] name = "contourpy" version = "1.3.3" @@ -1322,46 +1277,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, ] -[[package]] -name = "cymem" -version = "2.0.13" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/2f0fbb32535c3731b7c2974c569fb9325e0a38ed5565a08e1139a3b71e82/cymem-2.0.13.tar.gz", hash = "sha256:1c91a92ae8c7104275ac26bd4d29b08ccd3e7faff5893d3858cb6fadf1bc1588", size = 12320, upload-time = "2025-11-14T14:58:36.902Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/10/64/1db41f7576a6b69f70367e3c15e968fd775ba7419e12059c9966ceb826f8/cymem-2.0.13-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:673183466b0ff2e060d97ec5116711d44200b8f7be524323e080d215ee2d44a5", size = 43587, upload-time = "2025-11-14T14:57:22.39Z" }, - { url = "https://files.pythonhosted.org/packages/81/13/57f936fc08551323aab3f92ff6b7f4d4b89d5b4e495c870a67cb8d279757/cymem-2.0.13-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bee2791b3f6fc034ce41268851462bf662ff87e8947e35fb6dd0115b4644a61f", size = 43139, upload-time = "2025-11-14T14:57:23.363Z" }, - { url = "https://files.pythonhosted.org/packages/32/a6/9345754be51e0479aa387b7b6cffc289d0fd3201aaeb8dade4623abd1e02/cymem-2.0.13-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f3aee3adf16272bca81c5826eed55ba3c938add6d8c9e273f01c6b829ecfde22", size = 245063, upload-time = "2025-11-14T14:57:24.839Z" }, - { url = "https://files.pythonhosted.org/packages/d6/01/6bc654101526fa86e82bf6b05d99b2cd47c30a333cfe8622c26c0592beb2/cymem-2.0.13-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:30c4e75a3a1d809e89106b0b21803eb78e839881aa1f5b9bd27b454bc73afde3", size = 244496, upload-time = "2025-11-14T14:57:26.42Z" }, - { url = "https://files.pythonhosted.org/packages/c4/fb/853b7b021e701a1f41687f3704d5f469aeb2a4f898c3fbb8076806885955/cymem-2.0.13-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:ec99efa03cf8ec11c8906aa4d4cc0c47df393bc9095c9dd64b89b9b43e220b04", size = 243287, upload-time = "2025-11-14T14:57:27.542Z" }, - { url = "https://files.pythonhosted.org/packages/d4/2b/0e4664cafc581de2896d75000651fd2ce7094d33263f466185c28ffc96e4/cymem-2.0.13-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:c90a6ecba994a15b17a3f45d7ec74d34081df2f73bd1b090e2adc0317e4e01b6", size = 248287, upload-time = "2025-11-14T14:57:29.055Z" }, - { url = "https://files.pythonhosted.org/packages/21/0f/f94c6950edbfc2aafb81194fc40b6cacc8e994e9359d3cb4328c5705b9b5/cymem-2.0.13-cp311-cp311-win_amd64.whl", hash = "sha256:ce821e6ba59148ed17c4567113b8683a6a0be9c9ac86f14e969919121efb61a5", size = 40116, upload-time = "2025-11-14T14:57:30.592Z" }, - { url = "https://files.pythonhosted.org/packages/00/df/2455eff6ac0381ff165db6883b311f7016e222e3dd62185517f8e8187ed0/cymem-2.0.13-cp311-cp311-win_arm64.whl", hash = "sha256:0dca715e708e545fd1d97693542378a00394b20a37779c1ae2c8bdbb43acef79", size = 36349, upload-time = "2025-11-14T14:57:31.573Z" }, - { url = "https://files.pythonhosted.org/packages/c9/52/478a2911ab5028cb710b4900d64aceba6f4f882fcb13fd8d40a456a1b6dc/cymem-2.0.13-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:e8afbc5162a0fe14b6463e1c4e45248a1b2fe2cbcecc8a5b9e511117080da0eb", size = 43745, upload-time = "2025-11-14T14:57:32.52Z" }, - { url = "https://files.pythonhosted.org/packages/f9/71/f0f8adee945524774b16af326bd314a14a478ed369a728a22834e6785a18/cymem-2.0.13-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c9251d889348fe79a75e9b3e4d1b5fa651fca8a64500820685d73a3acc21b6a8", size = 42927, upload-time = "2025-11-14T14:57:33.827Z" }, - { url = "https://files.pythonhosted.org/packages/62/6d/159780fe162ff715d62b809246e5fc20901cef87ca28b67d255a8d741861/cymem-2.0.13-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:742fc19764467a49ed22e56a4d2134c262d73a6c635409584ae3bf9afa092c33", size = 258346, upload-time = "2025-11-14T14:57:34.917Z" }, - { url = "https://files.pythonhosted.org/packages/eb/12/678d16f7aa1996f947bf17b8cfb917ea9c9674ef5e2bd3690c04123d5680/cymem-2.0.13-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f190a92fe46197ee64d32560eb121c2809bb843341733227f51538ce77b3410d", size = 260843, upload-time = "2025-11-14T14:57:36.503Z" }, - { url = "https://files.pythonhosted.org/packages/31/5d/0dd8c167c08cd85e70d274b7235cfe1e31b3cebc99221178eaf4bbb95c6f/cymem-2.0.13-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d670329ee8dbbbf241b7c08069fe3f1d3a1a3e2d69c7d05ea008a7010d826298", size = 254607, upload-time = "2025-11-14T14:57:38.036Z" }, - { url = "https://files.pythonhosted.org/packages/b7/c9/d6514a412a1160aa65db539836b3d47f9b59f6675f294ec34ae32f867c82/cymem-2.0.13-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a84ba3178d9128b9ffb52ce81ebab456e9fe959125b51109f5b73ebdfc6b60d6", size = 262421, upload-time = "2025-11-14T14:57:39.265Z" }, - { url = "https://files.pythonhosted.org/packages/dd/fe/3ee37d02ca4040f2fb22d34eb415198f955862b5dd47eee01df4c8f5454c/cymem-2.0.13-cp312-cp312-win_amd64.whl", hash = "sha256:2ff1c41fd59b789579fdace78aa587c5fc091991fa59458c382b116fc36e30dc", size = 40176, upload-time = "2025-11-14T14:57:40.706Z" }, - { url = "https://files.pythonhosted.org/packages/94/fb/1b681635bfd5f2274d0caa8f934b58435db6c091b97f5593738065ddb786/cymem-2.0.13-cp312-cp312-win_arm64.whl", hash = "sha256:6bbd701338df7bf408648191dff52472a9b334f71bcd31a21a41d83821050f67", size = 35959, upload-time = "2025-11-14T14:57:41.682Z" }, - { url = "https://files.pythonhosted.org/packages/ce/0f/95a4d1e3bebfdfa7829252369357cf9a764f67569328cd9221f21e2c952e/cymem-2.0.13-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:891fd9030293a8b652dc7fb9fdc79a910a6c76fc679cd775e6741b819ffea476", size = 43478, upload-time = "2025-11-14T14:57:42.682Z" }, - { url = "https://files.pythonhosted.org/packages/bf/a0/8fc929cc29ae466b7b4efc23ece99cbd3ea34992ccff319089c624d667fd/cymem-2.0.13-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:89c4889bd16513ce1644ccfe1e7c473ba7ca150f0621e66feac3a571bde09e7e", size = 42695, upload-time = "2025-11-14T14:57:43.741Z" }, - { url = "https://files.pythonhosted.org/packages/4a/b3/deeb01354ebaf384438083ffe0310209ef903db3e7ba5a8f584b06d28387/cymem-2.0.13-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:45dcaba0f48bef9cc3d8b0b92058640244a95a9f12542210b51318da97c2cf28", size = 250573, upload-time = "2025-11-14T14:57:44.81Z" }, - { url = "https://files.pythonhosted.org/packages/36/36/bc980b9a14409f3356309c45a8d88d58797d02002a9d794dd6c84e809d3a/cymem-2.0.13-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e96848faaafccc0abd631f1c5fb194eac0caee4f5a8777fdbb3e349d3a21741c", size = 254572, upload-time = "2025-11-14T14:57:46.023Z" }, - { url = "https://files.pythonhosted.org/packages/fd/dd/a12522952624685bd0f8968e26d2ed6d059c967413ce6eb52292f538f1b0/cymem-2.0.13-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e02d3e2c3bfeb21185d5a4a70790d9df40629a87d8d7617dc22b4e864f665fa3", size = 248060, upload-time = "2025-11-14T14:57:47.605Z" }, - { url = "https://files.pythonhosted.org/packages/08/11/5dc933ddfeb2dfea747a0b935cb965b9a7580b324d96fc5f5a1b5ff8df29/cymem-2.0.13-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:fece5229fd5ecdcd7a0738affb8c59890e13073ae5626544e13825f26c019d3c", size = 254601, upload-time = "2025-11-14T14:57:48.861Z" }, - { url = "https://files.pythonhosted.org/packages/70/66/d23b06166864fa94e13a98e5922986ce774832936473578febce64448d75/cymem-2.0.13-cp313-cp313-win_amd64.whl", hash = "sha256:38aefeb269597c1a0c2ddf1567dd8605489b661fa0369c6406c1acd433b4c7ba", size = 40103, upload-time = "2025-11-14T14:57:50.396Z" }, - { url = "https://files.pythonhosted.org/packages/2f/9e/c7b21271ab88a21760f3afdec84d2bc09ffa9e6c8d774ad9d4f1afab0416/cymem-2.0.13-cp313-cp313-win_arm64.whl", hash = "sha256:717270dcfd8c8096b479c42708b151002ff98e434a7b6f1f916387a6c791e2ad", size = 36016, upload-time = "2025-11-14T14:57:51.611Z" }, - { url = "https://files.pythonhosted.org/packages/7f/28/d3b03427edc04ae04910edf1c24b993881c3ba93a9729a42bcbb816a1808/cymem-2.0.13-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:7e1a863a7f144ffb345397813701509cfc74fc9ed360a4d92799805b4b865dd1", size = 46429, upload-time = "2025-11-14T14:57:52.582Z" }, - { url = "https://files.pythonhosted.org/packages/35/a9/7ed53e481f47ebfb922b0b42e980cec83e98ccb2137dc597ea156642440c/cymem-2.0.13-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:c16cb80efc017b054f78998c6b4b013cef509c7b3d802707ce1f85a1d68361bf", size = 46205, upload-time = "2025-11-14T14:57:53.64Z" }, - { url = "https://files.pythonhosted.org/packages/61/39/a3d6ad073cf7f0fbbb8bbf09698c3c8fac11be3f791d710239a4e8dd3438/cymem-2.0.13-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0d78a27c88b26c89bd1ece247d1d5939dba05a1dae6305aad8fd8056b17ddb51", size = 296083, upload-time = "2025-11-14T14:57:55.922Z" }, - { url = "https://files.pythonhosted.org/packages/36/0c/20697c8bc19f624a595833e566f37d7bcb9167b0ce69de896eba7cfc9c2d/cymem-2.0.13-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6d36710760f817194dacb09d9fc45cb6a5062ed75e85f0ef7ad7aeeb13d80cc3", size = 286159, upload-time = "2025-11-14T14:57:57.106Z" }, - { url = "https://files.pythonhosted.org/packages/82/d4/9326e3422d1c2d2b4a8fb859bdcce80138f6ab721ddafa4cba328a505c71/cymem-2.0.13-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c8f30971cadd5dcf73bcfbbc5849b1f1e1f40db8cd846c4aa7d3b5e035c7b583", size = 288186, upload-time = "2025-11-14T14:57:58.334Z" }, - { url = "https://files.pythonhosted.org/packages/ed/bc/68da7dd749b72884dc22e898562f335002d70306069d496376e5ff3b6153/cymem-2.0.13-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:9d441d0e45798ec1fd330373bf7ffa6b795f229275f64016b6a193e6e2a51522", size = 290353, upload-time = "2025-11-14T14:58:00.562Z" }, - { url = "https://files.pythonhosted.org/packages/50/23/dbf2ad6ecd19b99b3aab6203b1a06608bbd04a09c522d836b854f2f30f73/cymem-2.0.13-cp313-cp313t-win_amd64.whl", hash = "sha256:d1c950eebb9f0f15e3ef3591313482a5a611d16fc12d545e2018cd607f40f472", size = 44764, upload-time = "2025-11-14T14:58:01.793Z" }, - { url = "https://files.pythonhosted.org/packages/54/3f/35701c13e1fc7b0895198c8b20068c569a841e0daf8e0b14d1dc0816b28f/cymem-2.0.13-cp313-cp313t-win_arm64.whl", hash = "sha256:042e8611ef862c34a97b13241f5d0da86d58aca3cecc45c533496678e75c5a1f", size = 38964, upload-time = "2025-11-14T14:58:02.87Z" }, -] - [[package]] name = "dask" version = "2026.3.0" @@ -3661,46 +3576,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/08/7036c080d7117f28a4af526d794aab6a84463126db031b007717c1a6676e/multidict-6.7.1-py3-none-any.whl", hash = "sha256:55d97cc6dae627efa6a6e548885712d4864b81110ac76fa4e534c03819fa4a56", size = 12319, upload-time = "2026-01-26T02:46:44.004Z" }, ] -[[package]] -name = "murmurhash" -version = "1.0.15" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/23/2e/88c147931ea9725d634840d538622e94122bceaf346233349b7b5c62964b/murmurhash-1.0.15.tar.gz", hash = "sha256:58e2b27b7847f9e2a6edf10b47a8c8dd70a4705f45dccb7bf76aeadacf56ba01", size = 13291, upload-time = "2025-11-14T09:51:15.272Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6b/ca/77d3e69924a8eb4508bb4f0ad34e46adbeedeb93616a71080e61e53dad71/murmurhash-1.0.15-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f32307fb9347680bb4fe1cbef6362fb39bd994f1b59abd8c09ca174e44199081", size = 27397, upload-time = "2025-11-14T09:50:03.077Z" }, - { url = "https://files.pythonhosted.org/packages/e6/53/a936f577d35b245d47b310f29e5e9f09fcac776c8c992f1ab51a9fb0cee2/murmurhash-1.0.15-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:539d8405885d1d19c005f3a2313b47e8e54b0ee89915eb8dfbb430b194328e6c", size = 27692, upload-time = "2025-11-14T09:50:04.144Z" }, - { url = "https://files.pythonhosted.org/packages/4d/64/5f8cfd1fd9cbeb43fcff96672f5bd9e7e1598d1c970f808ecd915490dc20/murmurhash-1.0.15-cp311-cp311-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:c4cd739a00f5a4602201b74568ddabae46ec304719d9be752fd8f534a9464b5e", size = 128396, upload-time = "2025-11-14T09:50:05.268Z" }, - { url = "https://files.pythonhosted.org/packages/ac/10/d9ce29d559a75db0d8a3f13ea12c7f541ec9de2afca38dc70418b890eedb/murmurhash-1.0.15-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:44d211bcc3ec203c47dac06f48ee871093fcbdffa6652a6cc5ea7180306680a8", size = 128687, upload-time = "2025-11-14T09:50:06.527Z" }, - { url = "https://files.pythonhosted.org/packages/48/cd/dc97ab7e68cdfa1537a56e36dbc846c5a66701cc39ecee2d4399fe61996c/murmurhash-1.0.15-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:f9bf47101354fb1dc4b2e313192566f04ba295c28a37e2f71c692759acc1ba3c", size = 128198, upload-time = "2025-11-14T09:50:08.062Z" }, - { url = "https://files.pythonhosted.org/packages/53/73/32f2aaa22c1e4afae337106baf0c938abf36a6cc879cfee83a00461bbbf7/murmurhash-1.0.15-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:3c69b4d3bcd6233782a78907fe10b9b7a796bdc5d28060cf097d067bec280a5d", size = 127214, upload-time = "2025-11-14T09:50:09.265Z" }, - { url = "https://files.pythonhosted.org/packages/82/ed/812103a7f353eba2d83655b08205e13a38c93b4db0692f94756e1eb44516/murmurhash-1.0.15-cp311-cp311-win_amd64.whl", hash = "sha256:e43a69496342ce530bdd670264cb7c8f45490b296e4764c837ce577e3c7ebd53", size = 25241, upload-time = "2025-11-14T09:50:10.373Z" }, - { url = "https://files.pythonhosted.org/packages/eb/5f/2c511bdd28f7c24da37a00116ffd0432b65669d098f0d0260c66ac0ffdc2/murmurhash-1.0.15-cp311-cp311-win_arm64.whl", hash = "sha256:f3e99a6ee36ef5372df5f138e3d9c801420776d3641a34a49e5c2555f44edba7", size = 23216, upload-time = "2025-11-14T09:50:11.651Z" }, - { url = "https://files.pythonhosted.org/packages/b6/46/be8522d3456fdccf1b8b049c6d82e7a3c1114c4fc2cfe14b04cba4b3e701/murmurhash-1.0.15-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d37e3ae44746bca80b1a917c2ea625cf216913564ed43f69d2888e5df97db0cb", size = 27884, upload-time = "2025-11-14T09:50:13.133Z" }, - { url = "https://files.pythonhosted.org/packages/ed/cc/630449bf4f6178d7daf948ce46ad00b25d279065fc30abd8d706be3d87e0/murmurhash-1.0.15-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0861cb11039409eaf46878456b7d985ef17b6b484103a6fc367b2ecec846891d", size = 27855, upload-time = "2025-11-14T09:50:14.859Z" }, - { url = "https://files.pythonhosted.org/packages/ff/30/ea8f601a9bf44db99468696efd59eb9cff1157cd55cb586d67116697583f/murmurhash-1.0.15-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:5a301decfaccfec70fe55cb01dde2a012c3014a874542eaa7cc73477bb749616", size = 134088, upload-time = "2025-11-14T09:50:15.958Z" }, - { url = "https://files.pythonhosted.org/packages/c9/de/c40ce8c0877d406691e735b8d6e9c815f36a82b499d358313db5dbe219d7/murmurhash-1.0.15-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:32c6fde7bd7e9407003370a07b5f4addacabe1556ad3dc2cac246b7a2bba3400", size = 133978, upload-time = "2025-11-14T09:50:17.572Z" }, - { url = "https://files.pythonhosted.org/packages/47/84/bd49963ecd84ebab2fe66595e2d1ed41d5e8b5153af5dc930f0bd827007c/murmurhash-1.0.15-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:5d8b43a7011540dc3c7ce66f2134df9732e2bc3bbb4a35f6458bc755e48bde26", size = 132956, upload-time = "2025-11-14T09:50:18.742Z" }, - { url = "https://files.pythonhosted.org/packages/4f/7c/2530769c545074417c862583f05f4245644599f1e9ff619b3dfe2969aafc/murmurhash-1.0.15-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:43bf4541892ecd95963fcd307bf1c575fc0fee1682f41c93007adee71ca2bb40", size = 134184, upload-time = "2025-11-14T09:50:19.941Z" }, - { url = "https://files.pythonhosted.org/packages/84/a4/b249b042f5afe34d14ada2dc4afc777e883c15863296756179652e081c44/murmurhash-1.0.15-cp312-cp312-win_amd64.whl", hash = "sha256:f4ac15a2089dc42e6eb0966622d42d2521590a12c92480aafecf34c085302cca", size = 25647, upload-time = "2025-11-14T09:50:21.049Z" }, - { url = "https://files.pythonhosted.org/packages/13/bf/028179259aebc18fd4ba5cae2601d1d47517427a537ab44336446431a215/murmurhash-1.0.15-cp312-cp312-win_arm64.whl", hash = "sha256:4a70ca4ae19e600d9be3da64d00710e79dde388a4d162f22078d64844d0ebdda", size = 23338, upload-time = "2025-11-14T09:50:22.359Z" }, - { url = "https://files.pythonhosted.org/packages/29/2f/ba300b5f04dae0409202d6285668b8a9d3ade43a846abee3ef611cb388d5/murmurhash-1.0.15-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:fe50dc70e52786759358fd1471e309b94dddfffb9320d9dfea233c7684c894ba", size = 27861, upload-time = "2025-11-14T09:50:23.804Z" }, - { url = "https://files.pythonhosted.org/packages/34/02/29c19d268e6f4ea1ed2a462c901eed1ed35b454e2cbc57da592fad663ac6/murmurhash-1.0.15-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1349a7c23f6092e7998ddc5bd28546cc31a595afc61e9fdb3afc423feec3d7ad", size = 27840, upload-time = "2025-11-14T09:50:25.146Z" }, - { url = "https://files.pythonhosted.org/packages/e2/63/58e2de2b5232cd294c64092688c422196e74f9fa8b3958bdf02d33df24b9/murmurhash-1.0.15-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b3ba6d05de2613535b5a9227d4ad8ef40a540465f64660d4a8800634ae10e04f", size = 133080, upload-time = "2025-11-14T09:50:26.566Z" }, - { url = "https://files.pythonhosted.org/packages/aa/9a/d13e2e9f8ba1ced06840921a50f7cece0a475453284158a3018b72679761/murmurhash-1.0.15-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fa1b70b3cc2801ab44179c65827bbd12009c68b34e9d9ce7125b6a0bd35af63c", size = 132648, upload-time = "2025-11-14T09:50:27.788Z" }, - { url = "https://files.pythonhosted.org/packages/b2/e1/47994f1813fa205c84977b0ff51ae6709f8539af052c7491a5f863d82bdc/murmurhash-1.0.15-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:213d710fb6f4ef3bc11abbfad0fa94a75ffb675b7dc158c123471e5de869f9af", size = 131502, upload-time = "2025-11-14T09:50:29.339Z" }, - { url = "https://files.pythonhosted.org/packages/b9/ea/90c1fd00b4aeb704fb5e84cd666b33ffd7f245155048071ffbb51d2bb57d/murmurhash-1.0.15-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b65a5c4e7f5d71f7ccac2d2b60bdf7092d7976270878cfec59d5a66a533db823", size = 132736, upload-time = "2025-11-14T09:50:30.545Z" }, - { url = "https://files.pythonhosted.org/packages/00/db/da73462dbfa77f6433b128d2120ba7ba300f8c06dc4f4e022c38d240a5f5/murmurhash-1.0.15-cp313-cp313-win_amd64.whl", hash = "sha256:9aba94c5d841e1904cd110e94ceb7f49cfb60a874bbfb27e0373622998fb7c7c", size = 25682, upload-time = "2025-11-14T09:50:31.624Z" }, - { url = "https://files.pythonhosted.org/packages/bb/83/032729ef14971b938fbef41ee125fc8800020ee229bd35178b6ede8ee934/murmurhash-1.0.15-cp313-cp313-win_arm64.whl", hash = "sha256:263807eca40d08c7b702413e45cca75ecb5883aa337237dc5addb660f1483378", size = 23370, upload-time = "2025-11-14T09:50:33.264Z" }, - { url = "https://files.pythonhosted.org/packages/10/83/7547d9205e9bd2f8e5dfd0b682cc9277594f98909f228eb359489baec1df/murmurhash-1.0.15-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:694fd42a74b7ce257169d14c24aa616aa6cd4ccf8abe50eca0557e08da99d055", size = 29955, upload-time = "2025-11-14T09:50:34.488Z" }, - { url = "https://files.pythonhosted.org/packages/b7/c7/3afd5de7a5b3ae07fe2d3a3271b327ee1489c58ba2b2f2159bd31a25edb9/murmurhash-1.0.15-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a2ea4546ba426390beff3cd10db8f0152fdc9072c4f2583ec7d8aa9f3e4ac070", size = 30108, upload-time = "2025-11-14T09:50:35.53Z" }, - { url = "https://files.pythonhosted.org/packages/02/69/d6637ee67d78ebb2538c00411f28ea5c154886bbe1db16c49435a8a4ab16/murmurhash-1.0.15-cp313-cp313t-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:34e5a91139c40b10f98d0b297907f5d5267b4b1b2e5dd2eb74a021824f751b98", size = 164054, upload-time = "2025-11-14T09:50:36.591Z" }, - { url = "https://files.pythonhosted.org/packages/ab/4c/89e590165b4c7da6bf941441212a721a270195332d3aacfdfdf527d466ca/murmurhash-1.0.15-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:dc35606868a5961cf42e79314ca0bddf5a400ce377b14d83192057928d6252ec", size = 168153, upload-time = "2025-11-14T09:50:37.856Z" }, - { url = "https://files.pythonhosted.org/packages/07/7a/95c42df0c21d2e413b9fcd17317a7587351daeb264dc29c6aec1fdbd26f8/murmurhash-1.0.15-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:43cc6ac3b91ca0f7a5ae9c063ba4d6c26972c97fd7c25280ecc666413e4c5535", size = 164345, upload-time = "2025-11-14T09:50:39.346Z" }, - { url = "https://files.pythonhosted.org/packages/d0/22/9d02c880a88b83bb3ce7d6a38fb727373ab78d82e5f3d8d9fc5612219f90/murmurhash-1.0.15-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:847d712136cb462f0e4bd6229ee2d9eb996d8854eb8312dff3d20c8f5181fda5", size = 161990, upload-time = "2025-11-14T09:50:40.689Z" }, - { url = "https://files.pythonhosted.org/packages/9a/e3/750232524e0dc262e8dcede6536dafc766faadd9a52f1d23746b02948ad8/murmurhash-1.0.15-cp313-cp313t-win_amd64.whl", hash = "sha256:2680851af6901dbe66cc4aa7ef8e263de47e6e1b425ae324caa571bdf18f8d58", size = 28812, upload-time = "2025-11-14T09:50:41.971Z" }, - { url = "https://files.pythonhosted.org/packages/ff/89/4ad9d215ef6ade89f27a72dc4e86b98ef1a43534cc3e6a6900a362a0bf0a/murmurhash-1.0.15-cp313-cp313t-win_arm64.whl", hash = "sha256:189a8de4d657b5da9efd66601b0636330b08262b3a55431f2379097c986995d0", size = 25398, upload-time = "2025-11-14T09:50:43.023Z" }, -] - [[package]] name = "mypy" version = "2.1.0" @@ -4101,6 +3976,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b0/b4/b7f5ce2e3ddee8518cf2314a7cf1f69f46da0e056d253d8d4bcb4952ba57/nvidia_nat_profiler-1.8.0-py3-none-any.whl", hash = "sha256:32ef0c7e83cdc72029b97b3899497dec6864ac3db9378a3842565bbe269f0dce", size = 108154, upload-time = "2026-06-17T00:24:38.546Z" }, ] +[[package]] +name = "nvidia-nat-redis" +version = "1.8.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nat-core" }, + { name = "redis" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/c5/ad7e01178a8c68eaed57bee3cc686c31f9a87d27f42193622831e4722e16/nvidia_nat_redis-1.8.0-py3-none-any.whl", hash = "sha256:e617230e304ba029bf4f89da1adb7133df4651dfe90338aa9771b0a855b1d753", size = 58896, upload-time = "2026-06-17T00:40:21.281Z" }, +] + [[package]] name = "nvidia-nat-security" version = "1.8.0" @@ -4573,15 +4460,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8b/c8/cdbc975f5b634e249cfa6597e37c50f3078412474f21c015e508bfbfe3c3/pdfplumber-0.11.9-py3-none-any.whl", hash = "sha256:33ec5580959ba524e9100138746e090879504c42955df1b8a997604dd326c443", size = 60045, upload-time = "2026-01-05T08:10:27.512Z" }, ] -[[package]] -name = "phonenumbers" -version = "9.0.33" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/75/37/dfc4cf24169f1a7169ebaedaf896c818f0add8603409d1e748e3085ccdc0/phonenumbers-9.0.33.tar.gz", hash = "sha256:9ab8a02b940b90c64f3866c0b25a30e567ddf7bb9836a3e11efdb0478f65fc1c", size = 2306756, upload-time = "2026-06-22T10:23:33.428Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e6/29/f7e30e3dbd3c7e3d9c4a55006112c04ee62b4765a31f21bcc28c253ac3f1/phonenumbers-9.0.33-py2.py3-none-any.whl", hash = "sha256:ba1d0da52711d5fdda6b2b673b2621fe80774fc5d1b2e5a6ef783396b0343186", size = 2595422, upload-time = "2026-06-22T10:23:29.925Z" }, -] - [[package]] name = "pillow" version = "12.2.0" @@ -4743,42 +4621,6 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/c5/b0/e66e9f6e07a0b37aa0f5703c46f54bafbdf65dfba63994247676b19076c4/prefixspan-0.5.2.tar.gz", hash = "sha256:8fbd2f94b3a7f4399d04f9bd6aa214b830fb7828799c472cd43dda10b03f671c", size = 10404, upload-time = "2018-09-29T06:12:27.708Z" } -[[package]] -name = "preshed" -version = "3.0.13" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cymem" }, - { name = "murmurhash" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/43/75/fe6b7bbd0dea530a001b0e24c331b21a0be2786e402abf3c57f5dce43d4b/preshed-3.0.13.tar.gz", hash = "sha256:d75f718bbfd97e992f7827e0fa7faf6a91bdd9c922d5baa4b50d62731396cb89", size = 18338, upload-time = "2026-03-23T08:57:31.378Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7c/d1/7bc39738388b38ff48cecbb326a9b2bb3f422bb32097be92e010f3162395/preshed-3.0.13-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5268c0e6fa96f50cdf87f516c2d4b32563c12706ee768e75c00e8d0098acd545", size = 136718, upload-time = "2026-03-23T08:56:23.889Z" }, - { url = "https://files.pythonhosted.org/packages/f6/65/de465b6801740140c2b5d2db6c312ca7937dcfd0442f1ae7d50dee529544/preshed-3.0.13-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:df642547a1a94079978a0ea8f4593ab4b8d3bd43f767bef0ef64d9a214f8c4c9", size = 137261, upload-time = "2026-03-23T08:56:25.303Z" }, - { url = "https://files.pythonhosted.org/packages/89/83/478ee078746a4a413c841542caebd2ea74b659475b8bf5f2e3724b6fe655/preshed-3.0.13-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:09397592d333a77f88454e72b7f1f941b2afaf040b392b9e74898dbc4648cdf5", size = 821010, upload-time = "2026-03-23T08:56:26.455Z" }, - { url = "https://files.pythonhosted.org/packages/ee/2e/1ac761e973966893cd3a0ad3256360365276e2d1e779e351448981a1156a/preshed-3.0.13-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f8e6fe0620ed0f96a246d46447055c447e071cd8222731a045c235e8a758c918", size = 823096, upload-time = "2026-03-23T08:56:28.126Z" }, - { url = "https://files.pythonhosted.org/packages/3c/51/7824cfd85dd7fe547888de20228ebd87d9acd3708206d30b82211e382d23/preshed-3.0.13-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:502f93f49a22788203f02d3067d4ea077a0cca3864de6a792eae12e7ce589e14", size = 1812148, upload-time = "2026-03-23T08:56:29.755Z" }, - { url = "https://files.pythonhosted.org/packages/34/48/32160a24705d56179de6af838c10a0c735c955dae5f9e4bb344750b79bc2/preshed-3.0.13-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:acd4d89abeca3678c5d8c89b3cd351314465bc67c7fa053d2644f8513e543386", size = 1881154, upload-time = "2026-03-23T08:56:31.49Z" }, - { url = "https://files.pythonhosted.org/packages/ed/22/0344b50f8b1ad9e3aac08099c47e1aba91c81602fd117d2673f6606ecae6/preshed-3.0.13-cp311-cp311-win_amd64.whl", hash = "sha256:de87fbabb0f37c3c92d4dd9b94fc82ab73cdab4247cdfbd57ab3926caa983919", size = 122219, upload-time = "2026-03-23T08:56:32.74Z" }, - { url = "https://files.pythonhosted.org/packages/33/c4/812eeaa568510f396e27edab01100ca71418f032fd7098b107f12e572361/preshed-3.0.13-cp311-cp311-win_arm64.whl", hash = "sha256:5e2753779832e411e93eb727f3d409c0a6b7408e5ce4dd868076d8ece48c7693", size = 109308, upload-time = "2026-03-23T08:56:33.839Z" }, - { url = "https://files.pythonhosted.org/packages/39/fb/ccff23c44c04088c248539005fcda78b9014512a34d170c5360f02ad908b/preshed-3.0.13-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:5d14eea14bd01291388928991d7df7d60b9fd19ae970e55006eb4d29b0c1e8eb", size = 138497, upload-time = "2026-03-23T08:56:35.321Z" }, - { url = "https://files.pythonhosted.org/packages/8e/ce/cad5a8145881a771e6c0d002f2e585fc19b962f120860b54d32af5baa342/preshed-3.0.13-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f05b08ce92399c0655b5e0eb5a1cc1f9e295703ed3aabdfaf6538dfa8ae23d57", size = 138010, upload-time = "2026-03-23T08:56:36.399Z" }, - { url = "https://files.pythonhosted.org/packages/a7/a2/c5fed4fb3e946699259d11e4036a3cfdd8c89b3e542e3077d46781642425/preshed-3.0.13-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:62cf7f3113132891d6bba70ff547ad81c6fe50a31930bbbb8499f1d47cd122b7", size = 861498, upload-time = "2026-03-23T08:56:37.67Z" }, - { url = "https://files.pythonhosted.org/packages/51/94/8c9bc48a6ea4903f53a1a0031ce8e35687526949f25821762ef21493c007/preshed-3.0.13-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8b8de3f58043070a354477995acdd98626ce43e4193c708ebd0f694e467f5155", size = 868988, upload-time = "2026-03-23T08:56:39.324Z" }, - { url = "https://files.pythonhosted.org/packages/b6/df/ecd2f40055ff52527ca117ffbfafb888c1a3079b59fbabe03c5b8f9b7240/preshed-3.0.13-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:183b339956a9e1d7a4a00038a3b9587a734db9e8bd915939a49791bd1b372156", size = 1847382, upload-time = "2026-03-23T08:56:40.89Z" }, - { url = "https://files.pythonhosted.org/packages/e6/88/bdb244e40284ded3632a9f88c23bc80230bd7b2ae4a8b7f2cc91adead7a8/preshed-3.0.13-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2e77bed56aded7cbe5d28d6bd2178bc5b13eda0e0e464dab205fb578fa915000", size = 1919236, upload-time = "2026-03-23T08:56:42.616Z" }, - { url = "https://files.pythonhosted.org/packages/a0/c9/c91ea56342e6c364fc69b444a1ac5432327857199c44032c9cc9dc4c3a23/preshed-3.0.13-cp312-cp312-win_amd64.whl", hash = "sha256:04d8f13f2986e5d11af5ac51f55ce3106c70c41b483d20ea392e6180bdd0f870", size = 122938, upload-time = "2026-03-23T08:56:44.271Z" }, - { url = "https://files.pythonhosted.org/packages/b2/0b/6a99d99619fd83b14c696e2489caed7070647488d4d3ac0b723d35db2de0/preshed-3.0.13-cp312-cp312-win_arm64.whl", hash = "sha256:19318dc1cd8cac6663c6c830bf7e0002d2de853769fb03e056774e97c21bedfd", size = 109194, upload-time = "2026-03-23T08:56:45.346Z" }, - { url = "https://files.pythonhosted.org/packages/0e/2a/401158195d6dc7f6aef0b354d74d0e95c9da124499448c2b3dbb95b71204/preshed-3.0.13-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c0d0c14187dc0078d8a63bf190ec045a4d13e7748b6caeb557a7d575e411410b", size = 137289, upload-time = "2026-03-23T08:56:46.516Z" }, - { url = "https://files.pythonhosted.org/packages/88/8f/e20e64573988528785447a6893b2e7ab287ecfd85b3888e978b28812fd20/preshed-3.0.13-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7770987c2e57497cd26124a9be5f652b5b3ccd0def89859ab0da8bca6144a3de", size = 136847, upload-time = "2026-03-23T08:56:47.572Z" }, - { url = "https://files.pythonhosted.org/packages/b9/72/18168f881359c4482d312f8dc196371bdd61c1583a52b34390da4c88bbea/preshed-3.0.13-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4a7bc48220de579be6bdb0a8715482cf36e2a625a6fd5ad26c9f43485a4a23b5", size = 831478, upload-time = "2026-03-23T08:56:48.769Z" }, - { url = "https://files.pythonhosted.org/packages/fd/3a/3543476091087102775568cea9885dde3453569e9aeee365809108de572f/preshed-3.0.13-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e5c8462472f790c16708306aef3a102a762bd19dfe3d2f8ee08bd5e12f51b835", size = 839913, upload-time = "2026-03-23T08:56:49.937Z" }, - { url = "https://files.pythonhosted.org/packages/cf/65/b13f01329decc44ef53cfb6b4601ba85382dcb2a4ec78d9250f03a418066/preshed-3.0.13-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:c046736239cc8d72670749b79b526e4111839a2fc461a58545d212797649129c", size = 1816452, upload-time = "2026-03-23T08:56:51.233Z" }, - { url = "https://files.pythonhosted.org/packages/d1/c7/f1a996c6832234efd4d543041b582418d41ac480ee55c557ec9e65344637/preshed-3.0.13-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7c333f18e9a81c8a6de0603fd8781e17115324b117c445ca91abdf7bfb1abe49", size = 1888978, upload-time = "2026-03-23T08:56:52.591Z" }, - { url = "https://files.pythonhosted.org/packages/e3/b9/96fb71499049885ce19545903fdd38877bbc2be0da47e37c04d01f3e9f66/preshed-3.0.13-cp313-cp313-win_amd64.whl", hash = "sha256:461327f8dd36520dcf1fd55a671e0c3c2c97a2d95e22fc85faa31173f4785dda", size = 122134, upload-time = "2026-03-23T08:56:54.392Z" }, - { url = "https://files.pythonhosted.org/packages/ef/a7/32a4903019d936a2316fdd330bedddac287ac26326107d24fb76a1fbc60a/preshed-3.0.13-cp313-cp313-win_arm64.whl", hash = "sha256:35d6c5acb3ee3b12b87a551913063f0cec784055c2af16e028c19fe875f079d0", size = 108497, upload-time = "2026-03-23T08:56:55.816Z" }, -] - [[package]] name = "primp" version = "1.3.1" @@ -5587,6 +5429,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/73/e8/2bdf3ca2090f68bb3d75b44da7bbc71843b19c9f2b9cb9b0f4ab7a5a4329/pyyaml-6.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb", size = 140246, upload-time = "2025-09-25T21:32:34.663Z" }, ] +[[package]] +name = "redis" +version = "4.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "async-timeout", marker = "python_full_version <= '3.11.2'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/73/88/63d802c2b18dd9eaa5b846cbf18917c6b2882f20efda398cc16a7500b02c/redis-4.6.0.tar.gz", hash = "sha256:585dc516b9eb042a619ef0a39c3d7d55fe81bdb4df09a52c9cdde0d07bf1aa7d", size = 4561721, upload-time = "2023-06-25T13:13:57.139Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/20/2e/409703d645363352a20c944f5d119bdae3eb3034051a53724a7c5fee12b8/redis-4.6.0-py3-none-any.whl", hash = "sha256:e2b03db868160ee4591de3cb90d40ebb50a90dd302138775937f6a42b7ed183c", size = 241149, upload-time = "2023-06-25T13:13:54.563Z" }, +] + [[package]] name = "referencing" version = "0.37.0" @@ -5688,18 +5542,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, ] -[[package]] -name = "requests-file" -version = "3.0.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "requests" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/3c/f8/5dc70102e4d337063452c82e1f0d95e39abfe67aa222ed8a5ddeb9df8de8/requests_file-3.0.1.tar.gz", hash = "sha256:f14243d7796c588f3521bd423c5dea2ee4cc730e54a3cac9574d78aca1272576", size = 6967, upload-time = "2025-10-20T18:56:42.279Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e1/d5/de8f089119205a09da657ed4784c584ede8381a0ce6821212a6d4ca47054/requests_file-3.0.1-py2.py3-none-any.whl", hash = "sha256:d0f5eb94353986d998f80ac63c7f146a307728be051d4d1cd390dbdb59c10fa2", size = 4514, upload-time = "2025-10-20T18:56:41.184Z" }, -] - [[package]] name = "requests-oauthlib" version = "2.0.0" @@ -5983,18 +5825,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] -[[package]] -name = "smart-open" -version = "7.7.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "wrapt" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/db/c6/22e7a2acd5d27941e85e0d7ede398da5abe2e4677d2265c924157247c32e/smart_open-7.7.1.tar.gz", hash = "sha256:9414ba5733e28309f29b28a303b0f1054ad23fe0275f1a1b600c80a724f4bd1a", size = 54952, upload-time = "2026-06-26T07:56:35.309Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8a/c5/14732c140fa0af0e59345949eaf9016c75ad3f076586f2413e3d3566a8da/smart_open-7.7.1-py3-none-any.whl", hash = "sha256:cb62dc45f519bf39b612564d326d2a17556f3a6056e3c3a86a07215d670d45bc", size = 65061, upload-time = "2026-06-26T07:56:33.685Z" }, -] - [[package]] name = "sniffio" version = "1.3.1" @@ -6307,41 +6137,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/42/89/81b2907cda14e566b9bf215e2ad82fc9b349edf07d2010756ffdb902f328/sqlite_vec-0.1.9-py3-none-win_amd64.whl", hash = "sha256:4a28dc12fa4b53d7b1dced22da2488fade444e96b5d16fd2d698cd670675cf32", size = 292804, upload-time = "2026-03-31T08:02:36.035Z" }, ] -[[package]] -name = "srsly" -version = "2.5.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "catalogue" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/2b/db/f794f219a6c788b881252d2536a8c4a97d2bdaadc690391e1cb53d123d71/srsly-2.5.3.tar.gz", hash = "sha256:08f98dbecbff3a31466c4ae7c833131f59d3655a0ad8ac749e6e2c149e2b0680", size = 490881, upload-time = "2026-03-23T11:56:59.865Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9a/36/5d7bb412d52e9cca787f9bfe838b596367189b254e50bf90f234a97184bf/srsly-2.5.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:785a09216ac31570fb301ddb9f61ee73d1f18f8b9561f712dce0b8ac8628bc88", size = 656760, upload-time = "2026-03-23T11:55:47.155Z" }, - { url = "https://files.pythonhosted.org/packages/d6/dc/124f008cd2be3e887e972cbdeb17c5aee0f42093eca02c7cfd63bb5daf19/srsly-2.5.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0017c7d2a0cd9a4f1bdc00d946b45edcf90bb0e271e8f084c1ce542bf6708c32", size = 657503, upload-time = "2026-03-23T11:55:48.681Z" }, - { url = "https://files.pythonhosted.org/packages/35/8a/2c97244ebab125d55f1bfb7bb94e9572b3e819410dffd6a040eca1112350/srsly-2.5.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:66ebae2c70305987341519ec1a720072a3cb3e4b1d52ac0e9e841f4d02658d3d", size = 1139161, upload-time = "2026-03-23T11:55:50.179Z" }, - { url = "https://files.pythonhosted.org/packages/fc/ea/ecd396188f7591d80b89665f7af9e3ae02e42683daef57033ad7993ad3f9/srsly-2.5.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4ca4a068f6e14d84113a02fcb875c6b50a6285a12938c0e7a157eb3a63c50a86", size = 1142438, upload-time = "2026-03-23T11:55:52.607Z" }, - { url = "https://files.pythonhosted.org/packages/9c/65/143e2e143c53d498ad0956f69d0e09189aa7a6e0ee6017758c285ba1ab2d/srsly-2.5.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:e283fa2a8f7350fb9fb70ecdee28d59d39c92f4c7f1cc90a44d6b86db3b3a8b3", size = 1101783, upload-time = "2026-03-23T11:55:53.906Z" }, - { url = "https://files.pythonhosted.org/packages/6b/86/1392a5593de0cd3d08c2d6c071b877c84358a37f63172c4e9cb71706842d/srsly-2.5.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9ffc97e22730ea97b00f7c303ccc60b1305e786afadb2a4a46578dafa4d29da0", size = 1115876, upload-time = "2026-03-23T11:55:55.624Z" }, - { url = "https://files.pythonhosted.org/packages/d4/a5/6193aa4c08e488821538fcbce2282449e228fd2183ed67d118bb5ccd8b54/srsly-2.5.3-cp311-cp311-win_amd64.whl", hash = "sha256:f09b551f6c3e334652831ac68c770ee4284741ce0a3895bf1ccf2a1178d66cdd", size = 651733, upload-time = "2026-03-23T11:55:56.964Z" }, - { url = "https://files.pythonhosted.org/packages/66/a8/a73181743b6d237026615ca75c3fb3e4780736f1390550a7350d0c7f1149/srsly-2.5.3-cp311-cp311-win_arm64.whl", hash = "sha256:21cf09e417d3e4f3fbf7dd337fd6d948c97abd01896b9b4cb80e81cd9778a73a", size = 639124, upload-time = "2026-03-23T11:55:58.532Z" }, - { url = "https://files.pythonhosted.org/packages/02/cc/e9f7fcec4cc92ad8bad6316c4241638b8cf7380382d4489d94ec6c436452/srsly-2.5.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:71e51c046ccbeefb86524c6b1e17574f579c6ac4dc8ea4a09437d3e8f88342d3", size = 658379, upload-time = "2026-03-23T11:55:59.85Z" }, - { url = "https://files.pythonhosted.org/packages/21/e4/fea4512e9785f58509b2cf67d993323848e583161b5fcfdc7dd9d7c1f3df/srsly-2.5.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:2f73c0db911552e94fe2016e1759d261d2f47926f68826664cada3723c87006a", size = 658513, upload-time = "2026-03-23T11:56:01.239Z" }, - { url = "https://files.pythonhosted.org/packages/20/b1/53591681b6ff2699a4f97b2d5552ba196eaa6a979b0873605f4c04b5f7ee/srsly-2.5.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5c1ac27ae5f4bb9163c7d2c45fc8ec173aac3d92e32086d9472b326c5c6e570e", size = 1172265, upload-time = "2026-03-23T11:56:02.589Z" }, - { url = "https://files.pythonhosted.org/packages/4e/c9/741e29f534919a944a16da4184924b1d3404c4bf60716ab2b91be771d1e3/srsly-2.5.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:99026bcd9cbd3211cc36517400b04ca0fc5d3e412b14daf84ee6e65f67d9a2d8", size = 1180873, upload-time = "2026-03-23T11:56:03.944Z" }, - { url = "https://files.pythonhosted.org/packages/89/57/5554f786eccf78b2750d6ac63be126e1b67badec2cb409dd611cf6f8c52b/srsly-2.5.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:07d682679e639eb46ff7e6da4a92714f4d5ffe351d088ee66f221e9b1f8865bb", size = 1120437, upload-time = "2026-03-23T11:56:05.283Z" }, - { url = "https://files.pythonhosted.org/packages/eb/95/9b4f73b1be3692f86d72ccc131c8e50f26f824d5c8830a59390bcc5b60ef/srsly-2.5.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8e0542d85d6b55cf2934050d6ffcb1cd76c768dcf9572e7467002cf087bb366d", size = 1137376, upload-time = "2026-03-23T11:56:06.613Z" }, - { url = "https://files.pythonhosted.org/packages/5a/de/89ca640ca1953c4612279ce515d0af35658df3c06cdb324329bc91b4a7e1/srsly-2.5.3-cp312-cp312-win_amd64.whl", hash = "sha256:598f1e494c18cacb978299d77125415a586417081959f8ec3f068b32d97f8933", size = 652459, upload-time = "2026-03-23T11:56:07.994Z" }, - { url = "https://files.pythonhosted.org/packages/6d/4f/7ab6d49e36d9cc72ee15746cabd116eb6f338be8a06c1882968ee9d6c7d7/srsly-2.5.3-cp312-cp312-win_arm64.whl", hash = "sha256:4b1b721cd3ad1a9b2343519aadc786a4d09d5c0666962d49852eb12d6ec3fe26", size = 638411, upload-time = "2026-03-23T11:56:09.31Z" }, - { url = "https://files.pythonhosted.org/packages/9d/5c/12901e3794f4158abc6da750725aad6c2afddb1e4227b300fe7c71f66957/srsly-2.5.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e67b6bbacbfadea5e100266d2797f2d4cec9883ea4dc84a5537673850036a8d8", size = 656750, upload-time = "2026-03-23T11:56:10.708Z" }, - { url = "https://files.pythonhosted.org/packages/04/61/181c26370995f96f56f1b64b801e3ca1e0d703fc36506ae28606d62369fb/srsly-2.5.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:348c231b4477d8fe86603131d0f166d2feac9c372704dfc4398be71cc5b6fb07", size = 656746, upload-time = "2026-03-23T11:56:12.28Z" }, - { url = "https://files.pythonhosted.org/packages/77/c6/35876c78889f8ffe11ed3521644e666c3aef20ea31527b70f47456cf35c2/srsly-2.5.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b0938c2978c91ae1ef9c1f2ba35abb86330e198fb23469e356eba311e02233ee", size = 1155762, upload-time = "2026-03-23T11:56:14.075Z" }, - { url = "https://files.pythonhosted.org/packages/3e/da/40b71ca9906c8eb8f8feb6ac11d33dad458c85a56e1de764b96d402168a0/srsly-2.5.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5f6a837954429ecbe6dcdd27390d2fb4c7d01a3f99c9ffcf9ce66b2a6dd1b738", size = 1161092, upload-time = "2026-03-23T11:56:15.778Z" }, - { url = "https://files.pythonhosted.org/packages/dc/14/c0dd30cc8b93ce8137ff4766f743c882440ce49195fffc5d50eaeef311a6/srsly-2.5.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3576c125c486ce2958c2047e8858fe3cfc9ea877adfa05203b0986f9badee355", size = 1109984, upload-time = "2026-03-23T11:56:17.056Z" }, - { url = "https://files.pythonhosted.org/packages/08/f3/34354f183d8faafc631585571224b54d1b4b67e796972c36519c074ca355/srsly-2.5.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5fb59c42922e095d1ea36085c55bc16e2adb06a7bfe57b24d381e0194ae699f2", size = 1128409, upload-time = "2026-03-23T11:56:18.761Z" }, - { url = "https://files.pythonhosted.org/packages/a4/d9/5531f8a19492060b4e76e4ab06aca6f096fb5128fe18cc813d1772daf653/srsly-2.5.3-cp313-cp313-win_amd64.whl", hash = "sha256:111805927f05f5db440aeeacb85ce43da0b19ce7b2a09567a9ef8d30f3cc4d83", size = 650820, upload-time = "2026-03-23T11:56:20.096Z" }, - { url = "https://files.pythonhosted.org/packages/8e/8a/62fb7a971eca29e12f03fb9ddacb058548c14d33e5b5675ff0f85839cc7b/srsly-2.5.3-cp313-cp313-win_arm64.whl", hash = "sha256:0f106b0a700ab56e4a7c431b0f1444009ab6cb332edc7bbf6811c2a43f4722cb", size = 637278, upload-time = "2026-03-23T11:56:21.439Z" }, -] - [[package]] name = "sse-starlette" version = "3.4.4" @@ -6424,52 +6219,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl", hash = "sha256:6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55", size = 28926, upload-time = "2026-02-07T10:45:32.24Z" }, ] -[[package]] -name = "thinc" -version = "8.3.13" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "blis" }, - { name = "catalogue" }, - { name = "confection" }, - { name = "cymem" }, - { name = "murmurhash" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "preshed" }, - { name = "pydantic" }, - { name = "setuptools" }, - { name = "srsly" }, - { name = "wasabi" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/13/46/76df95f2c327f9a9cef30c1523bf285627897097163584dcf5f77b2ebce2/thinc-8.3.13.tar.gz", hash = "sha256:68e658549fc1eb3ff92aed5147fcbb9c15d6e9cc0e623b4d0998d16522ffb4f9", size = 194640, upload-time = "2026-03-23T07:22:36.41Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b4/72/ca06842a007e8c794e8c59462f242cdfd6167d7cc9d0155ad004b194b015/thinc-8.3.13-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:4565102638038a01a2193c7f5d41ccbd6233fbdcb1f1b184322a06add4f51f18", size = 844359, upload-time = "2026-03-23T07:21:43.017Z" }, - { url = "https://files.pythonhosted.org/packages/48/44/e6aef092f478d263f72eb3933b55a6f37ba97c6a0ea0a61d13fbf9bf0c19/thinc-8.3.13-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:859fbd9d9b16af5278da23589b4afbe2ab6b0dd615df4d3229b7c4e67cd3107e", size = 812089, upload-time = "2026-03-23T07:21:44.618Z" }, - { url = "https://files.pythonhosted.org/packages/ff/8a/9ce0424d456cd3580cc3a855b23a7ff86b81d5299fceb496a2f56f06c1c0/thinc-8.3.13-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a518d5c761a0f2341e530e867de133dc3ed814558365b2a68ec53b89c482a43f", size = 4101388, upload-time = "2026-03-23T07:21:46.135Z" }, - { url = "https://files.pythonhosted.org/packages/ad/51/ec91c0434bd9a1096ab874bbd6dc110c5089d7fc513137e6af59bd051eec/thinc-8.3.13-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:81337dfbee37f58f36c0c70f9a819dce1b32cdc13d959181e10de079621f6ac6", size = 4131972, upload-time = "2026-03-23T07:21:48.403Z" }, - { url = "https://files.pythonhosted.org/packages/ff/67/e30dea753c90cff5cb9e5feb34948fdb89a6774b84d849585b49e16a730e/thinc-8.3.13-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:fbc0ee16edd260c6a4a9e365ff36d0a682c9e7ca6d7b985682659ef2e3e73826", size = 5101283, upload-time = "2026-03-23T07:21:49.991Z" }, - { url = "https://files.pythonhosted.org/packages/00/e9/b7544eddababa16e548b26a96fff29eeb307ce938df5fa4af9371fe8ed5d/thinc-8.3.13-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:0355c37e40d1a9fc2a1b8e9c2e294d8586f6baa97bcac6b9002f2dddb4b82ae9", size = 5264488, upload-time = "2026-03-23T07:21:51.747Z" }, - { url = "https://files.pythonhosted.org/packages/4c/a9/49391a40d703efc0f7a451310373261835f71fd3e6e2e8cfc08ee02f78ad/thinc-8.3.13-cp311-cp311-win_amd64.whl", hash = "sha256:0a0fa13dcfe4b319c3a396432c1dbff30d3de37dbbdee559e76600ee2b9486df", size = 1795058, upload-time = "2026-03-23T07:21:53.424Z" }, - { url = "https://files.pythonhosted.org/packages/c1/31/fd5348d44beda12a3ee415cbba9ed4fd0b17ce65db1d473c38a29a8d6153/thinc-8.3.13-cp311-cp311-win_arm64.whl", hash = "sha256:cd8a2b714c061969eee65802965167a6ada1fe708d82fe176d98dcb95ebe182a", size = 1721215, upload-time = "2026-03-23T07:21:55.027Z" }, - { url = "https://files.pythonhosted.org/packages/3e/af/f7c1ebfe92eb5d27d7f2f3da67a11e2eb57bc30ab1553279af6dc65b65a8/thinc-8.3.13-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:77a41f66285321d20aaedaea1e87d7cd48dca6d2427bed1867ec7cba7109fc8d", size = 821097, upload-time = "2026-03-23T07:21:56.698Z" }, - { url = "https://files.pythonhosted.org/packages/45/8f/69d7338575d98df85d0b54c0f5fc277dba72587fe9ab846ecdd12a998bcb/thinc-8.3.13-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3710d318b4e5460cf366a6f7b5ddbefb5d39dbd4cfa408222750fdc6c27c4411", size = 791932, upload-time = "2026-03-23T07:21:58.38Z" }, - { url = "https://files.pythonhosted.org/packages/4b/a5/21d010c81e81e1589e5ccb4950e521804d13726e541e87f644c51815673b/thinc-8.3.13-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5a08c87143a6d20177652dca1ec0dc815d88216d8fc62594a57e8bc45bf5ed49", size = 3854219, upload-time = "2026-03-23T07:21:59.819Z" }, - { url = "https://files.pythonhosted.org/packages/f9/ff/6914bf370bd1d604d89e6dfb46b97d10cd9b00d42ff8c036283e92314a8c/thinc-8.3.13-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4b5ec9ff313819e7d8667794a3559463fa89ff45aaa73e3fd8d6273b1e0d7a7f", size = 3903307, upload-time = "2026-03-23T07:22:01.652Z" }, - { url = "https://files.pythonhosted.org/packages/f3/3d/5572b47fa155fb3388c071515b74024fa17a6efd1df9406da378f0aa84ef/thinc-8.3.13-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:5c9a48f2bc1e04f138240ed5f9b815a9141a5de26accd0f08fa0137fcefed258", size = 4836882, upload-time = "2026-03-23T07:22:03.565Z" }, - { url = "https://files.pythonhosted.org/packages/f0/f0/a8d77c7bac089697c6df302cc3c936a1ab36a4720deae889e6f1dbcbd0eb/thinc-8.3.13-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:79a29a44d76bd02f5ac0624268c6e42b3576ae472c791a8ae9c2d813ae789b59", size = 5033398, upload-time = "2026-03-23T07:22:05.045Z" }, - { url = "https://files.pythonhosted.org/packages/21/82/5651bb1f904d04220fc7670035ada921bf0638e2cff6444d67c12887a968/thinc-8.3.13-cp312-cp312-win_amd64.whl", hash = "sha256:ed1dc709ac4f2f03b710457889e4e02f05de51bc8456980c241d0b28798bc7cb", size = 1721248, upload-time = "2026-03-23T07:22:06.749Z" }, - { url = "https://files.pythonhosted.org/packages/94/8d/683703de021ffbe46833d722b70f49ffbbca8e5bd6876256977555d92d7d/thinc-8.3.13-cp312-cp312-win_arm64.whl", hash = "sha256:c6a049703a6011c8fe26ee41af7e70272145594140d82f79bb23de619c6a6525", size = 1645777, upload-time = "2026-03-23T07:22:08.104Z" }, - { url = "https://files.pythonhosted.org/packages/af/b9/7b46942176df459d1804a9e77b0976f7c56f3abf3ec7485d0e5f836a0382/thinc-8.3.13-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c2811dfd8d46d8b5d3b39051b23e64006b2994a5143b1978b436938018792af8", size = 817337, upload-time = "2026-03-23T07:22:09.538Z" }, - { url = "https://files.pythonhosted.org/packages/a7/79/53085a72cd8f4fc4e6e313d05ea5aa98e870684f4a0fb318a9875fc0a964/thinc-8.3.13-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:5593e6300cb1ebe0c0e546e9c9fb49e7c2627a0aa688795cd4f995a8b820d2ec", size = 788120, upload-time = "2026-03-23T07:22:11.215Z" }, - { url = "https://files.pythonhosted.org/packages/9e/3e/d61b462b16da95ac6885f95bb395e672040ee594833e571a6edcffd234f5/thinc-8.3.13-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f697174d3fb474966ce50b430bbafa101a6d2f7ffb559dac4b5c59389ef72d22", size = 3844666, upload-time = "2026-03-23T07:22:12.67Z" }, - { url = "https://files.pythonhosted.org/packages/78/4c/898cc654bb123734c71ec5a425c02ca34439517d01ce1c95a6563295580e/thinc-8.3.13-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9c7c5c104737b414c8c4ec578e67d78b6c859afe25cbc0684402e721415bd7f", size = 3890658, upload-time = "2026-03-23T07:22:14.668Z" }, - { url = "https://files.pythonhosted.org/packages/cd/56/1abdbf0a4ad628e8a05d6516fe0745969649d805367a3dccad8ee872981b/thinc-8.3.13-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7a99d0e242d1ccd23f9ae6bea7cd502f8626efa65c156b91d84581d0356696c3", size = 4819933, upload-time = "2026-03-23T07:22:16.85Z" }, - { url = "https://files.pythonhosted.org/packages/f1/22/b84dbdc6be5055bbdb2a7352e2c393f67e8593c137f1b83c82bf1e062b6e/thinc-8.3.13-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e676edd21a747afbe3e6b9f3fca8b962e36d146ded03b070cb0c28e2dfbe9499", size = 5018099, upload-time = "2026-03-23T07:22:18.356Z" }, - { url = "https://files.pythonhosted.org/packages/0f/a8/763cd7ba949334c9d2cddc92dadb68b344cb9546dc01b8d4a733dcaa16c1/thinc-8.3.13-cp313-cp313-win_amd64.whl", hash = "sha256:8ad40307f20e83f77af28ff5c6be0b86af7a8b251d1231c545508d2763157d8f", size = 1720309, upload-time = "2026-03-23T07:22:19.81Z" }, - { url = "https://files.pythonhosted.org/packages/f5/15/a11f7bb3cbc97dfecf32a90552f5a8f8a5c99316a99c6c17bdabf5baf256/thinc-8.3.13-cp313-cp313-win_arm64.whl", hash = "sha256:723949cab11d1925c15447928513a718276316cec6e0de28337cca0a62be0521", size = 1644606, upload-time = "2026-03-23T07:22:21.339Z" }, -] - [[package]] name = "threadpoolctl" version = "3.6.0" @@ -6528,21 +6277,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ce/34/d50e338631baaf65ec5396e70085e5de0b52b24b28db1ffbc1c6e82190dc/tinytag-2.2.1-py3-none-any.whl", hash = "sha256:ed8b1e6d25367937e3321e054f4974f9abfde1a3e0a538824c87da377130c2b6", size = 32927, upload-time = "2026-03-15T18:47:59.613Z" }, ] -[[package]] -name = "tldextract" -version = "5.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "filelock" }, - { name = "idna" }, - { name = "requests" }, - { name = "requests-file" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/65/7b/644fbbb49564a6cb124a8582013315a41148dba2f72209bba14a84242bf0/tldextract-5.3.1.tar.gz", hash = "sha256:a72756ca170b2510315076383ea2993478f7da6f897eef1f4a5400735d5057fb", size = 126105, upload-time = "2025-12-28T23:58:05.532Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/42/0e49d6d0aac449ca71952ec5bae764af009754fcb2e76a5cc097543747b3/tldextract-5.3.1-py3-none-any.whl", hash = "sha256:6bfe36d518de569c572062b788e16a659ccaceffc486d243af0484e8ecf432d9", size = 105886, upload-time = "2025-12-28T23:58:04.071Z" }, -] - [[package]] name = "tokenizers" version = "0.23.1" @@ -6887,18 +6621,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bf/c4/557dc082be035381b85fdb2b74e21d3d21b57750b74f2b47a32f3a639ff9/virtualenv-21.4.2-py3-none-any.whl", hash = "sha256:854210ca524a1a4d0d744734f4acbc721c3ffe163b85bbf5d56d14d5ae2f0fae", size = 7594079, upload-time = "2026-05-31T17:01:20.735Z" }, ] -[[package]] -name = "wasabi" -version = "1.1.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ac/f9/054e6e2f1071e963b5e746b48d1e3727470b2a490834d18ad92364929db3/wasabi-1.1.3.tar.gz", hash = "sha256:4bb3008f003809db0c3e28b4daf20906ea871a2bb43f9914197d540f4f2e0878", size = 30391, upload-time = "2024-05-31T16:56:18.99Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/06/7c/34330a89da55610daa5f245ddce5aab81244321101614751e7537f125133/wasabi-1.1.3-py3-none-any.whl", hash = "sha256:f76e16e8f7e79f8c4c8be49b4024ac725713ab10cd7f19350ad18a8e3f71728c", size = 27880, upload-time = "2024-05-31T16:56:16.699Z" }, -] - [[package]] name = "watchdog" version = "6.0.0" @@ -7015,26 +6737,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bd/6e/95b0e537de1f4d4301f76f944642c6da50d1511cc7b3d64dc418a66c7509/wcwidth-0.8.1-py3-none-any.whl", hash = "sha256:f453740b1e4a4f3291faa37944c555d71056c4da08d59809b307ef4feba695c8", size = 323092, upload-time = "2026-06-08T05:57:21.413Z" }, ] -[[package]] -name = "weasel" -version = "1.0.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cloudpathlib" }, - { name = "confection" }, - { name = "httpx" }, - { name = "packaging" }, - { name = "pydantic" }, - { name = "smart-open" }, - { name = "srsly" }, - { name = "typer" }, - { name = "wasabi" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ce/e5/e272bb9a045105a1fdf4b798d8086f5932a178f4d738f17a74f5c9e0ae9a/weasel-1.0.0.tar.gz", hash = "sha256:7b129b44c90cc543b760532974ca1e4eb30dad2aa2026f57bdce66354ae610fc", size = 38682, upload-time = "2026-03-20T08:10:25.266Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0a/07/57ebf7a6798b016c064bd0ca81b4c6a99daa4dc377b898bc7b41eb6b5af0/weasel-1.0.0-py3-none-any.whl", hash = "sha256:89518acee027f49d743126c3502d35e6dd14f5768be5c37c9af47c171b6005cc", size = 50713, upload-time = "2026-03-20T08:10:23.637Z" }, -] - [[package]] name = "websocket-client" version = "1.9.0"