diff --git a/deploy/helm/README.md b/deploy/helm/README.md index 84ff68640..1dab28cca 100644 --- a/deploy/helm/README.md +++ b/deploy/helm/README.md @@ -189,9 +189,14 @@ aiq-postgres-xxx 1/1 Running 0 30s ```bash kubectl port-forward -n ns-aiq svc/aiq-backend 8000:8000 & +curl http://localhost:8000/live curl http://localhost:8000/health ``` +The backend liveness probe uses `/live`, which checks only that the API process +responds. The readiness probe uses `/health`, which checks required dependencies +and can return HTTP 503 without causing Kubernetes to restart the process. + The backend API docs are available at `http://localhost:8000/docs` while the port-forward is active. ### Access the application diff --git a/deploy/helm/deployment-k8s/README.md b/deploy/helm/deployment-k8s/README.md index c436854eb..c4dfe2de4 100644 --- a/deploy/helm/deployment-k8s/README.md +++ b/deploy/helm/deployment-k8s/README.md @@ -68,9 +68,14 @@ aiq-postgres-xxx 1/1 Running 0 30s ```bash kubectl port-forward -n ns-aiq svc/aiq-backend 8000:8000 & +curl http://localhost:8000/live curl http://localhost:8000/health ``` +The backend liveness probe uses `/live`, which checks only that the API process +responds. The readiness probe uses `/health`, so database or content-encryption +outages remove the pod from service without triggering a restart loop. + ## Configuration, FRAG, Secrets, and Access These topics apply to all deployment methods and are documented in the [Helm README](../README.md): diff --git a/deploy/helm/deployment-k8s/values.yaml b/deploy/helm/deployment-k8s/values.yaml index 181da10e8..8ba558882 100644 --- a/deploy/helm/deployment-k8s/values.yaml +++ b/deploy/helm/deployment-k8s/values.yaml @@ -49,7 +49,7 @@ aiq: enabled: true livenessProbe: httpGet: - path: /health + path: /live port: 8000 initialDelaySeconds: 30 periodSeconds: 15 diff --git a/docs/source/deployment/content-encryption.md b/docs/source/deployment/content-encryption.md new file mode 100644 index 000000000..f99f59557 --- /dev/null +++ b/docs/source/deployment/content-encryption.md @@ -0,0 +1,216 @@ + + +# Async Job Content Encryption + +AI-Q can encrypt sensitive async job content before it is persisted by the +AI-Q async jobs API. This is application-level envelope encryption for final +report output and selected artifact event payload fields. + +Content encryption is disabled by default. Operators must explicitly set +`AIQ_CONTENT_ENCRYPTION=key` or `vault` on every API and worker process to +enable it. + +This feature is intentionally narrow in its first milestone. It protects the +serialized final output payload returned by +`GET /v1/jobs/async/job/{job_id}/report`, such as `{"report": "..."}`, and +the `content` field of `artifact.update` events for `output` and `file` +artifacts. + +## Scope and Limitations + +Encrypted when enabled: + +- `job_info.output` for jobs submitted through `/v1/jobs/async/submit` or + `aiq_api.jobs.submit.submit_agent_job`. +- `job_events.event_data` field `artifact.update.data.content` when the + artifact `data.type` is `output` or `file`. + +Still plaintext: + +- Other `job_events.event_data` fields, including artifact metadata, tool + events, heartbeat events, cancellation events, error events, citation source + artifacts, citation use artifacts, and todo artifacts. +- Job status, ownership metadata, timestamps, event type, and other control + plane fields. +- `job_info.error`. +- PostgreSQL notification payloads. +- `summaries.summary`. +- LangGraph checkpoints in `aiq_checkpoints`. +- Historical `job_info.output` rows written before encryption was enabled. +- Inline CLI and local NeMo Agent Toolkit runs that do not use the AI-Q async + API job runner. + +Because checkpoints, errors, citations, todos, and other event fields can +contain equivalent research content, this phase does not provide full +database-level job-content confidentiality. + +This feature protects the configured fields from storage readers and detects +ciphertext tampering within the authenticated job and field context. It is not +a database integrity or event-freshness control. A principal with database +write access can still delete or reorder rows, replay an entire encrypted row, +or swap encrypted artifact content between events in the same job when the +authenticated field path is identical. + +## Modes + +Set `AIQ_CONTENT_ENCRYPTION` on every API and worker process. + +| Mode | Behavior | +|------|----------| +| `off` | Default. Preserves existing plaintext behavior and never attempts to decrypt `aiqenc:` values. | +| `key` | Uses one operator-managed static 32-byte key to wrap per-job data encryption keys. Intended only for development, testing, or deployments that cannot use Vault. | +| `vault` | Uses HashiCorp Vault Transit to generate and wrap per-job data encryption keys. Recommended for production. | + +Encrypted values are stored as `aiqenc:` envelopes. Encrypted event fields are +stored as marker objects containing the envelope so the surrounding event JSON +remains inspectable. The envelope contains non-secret metadata, the wrapped data +encryption key, nonce, ciphertext, tag, algorithm, key id, and an AAD hint that +binds the value to either `job_info.output:{job_id}` or the encrypted +`job_events.event_data` field path for that job. + +## Static Key Configuration + +Static key mode requires a base64 or base64url value that decodes to exactly +32 raw bytes. + +```bash +AIQ_CONTENT_ENCRYPTION=key +AIQ_CONTENT_ENCRYPTION_KEY= +AIQ_CONTENT_ENCRYPTION_KEY_ID= +``` + +`AIQ_CONTENT_ENCRYPTION_KEY_ID` is optional, but it is cryptographic identity, +not cosmetic metadata: it is authenticated while wrapping each data encryption +key. If omitted, envelopes use `static-key` as the key id. Keep the configured +key id unchanged while encrypted jobs are retained. The first implementation +supports one active static key only; changing the key or key id makes existing +envelopes unreadable unless those jobs have expired or a future rewrap process +has migrated them. + +Invalid static-key configuration fails startup. + +## Vault Transit Configuration + +Vault mode uses AppRole authentication and Transit data keys. Token fallback is +not supported in the first implementation. + +```bash +AIQ_CONTENT_ENCRYPTION=vault +VAULT_ADDR= +VAULT_ROLE_ID= +VAULT_SECRET_ID= +AIQ_ENCRYPTION_TRANSIT_KEY= +VAULT_TRANSIT_MOUNT= +AIQ_CONTENT_ENCRYPTION_KEY_ID= +``` + +`VAULT_TRANSIT_MOUNT` defaults to `transit` if omitted. +`AIQ_CONTENT_ENCRYPTION_KEY_ID` is optional; if omitted, envelopes use +`/`. + +Set `VAULT_NAMESPACE=` only when your Vault deployment +requires a namespace. + +Missing Vault configuration fails startup. If Vault configuration is present +but Vault is temporarily unreachable, unauthorized, or otherwise operationally +unready, the API starts unhealthy instead of exiting. + +The application relies on Vault Transit versioned ciphertext for decrypting +after Transit key rotation. Do not disable or destroy old Transit key versions +until corresponding encrypted jobs have expired or have been rewrapped by a +future migration process. + +## Vault Client and Retry Behavior + +Vault mode uses the synchronous `hvac.Client`. API routes that may call Vault +offload that work to worker threads so Vault latency does not block the FastAPI +event loop. Each API or worker process keeps one process-local Vault client and +guards client creation, AppRole login, and Transit calls with a lock. This +serializes Vault Transit operations per process and avoids concurrent mutation +of the shared `hvac.Client`. + +Vault Transit operations use bounded retries for transient failures only: + +- request connection failures and timeouts, +- Vault rate limiting, +- Vault 5xx-style operational failures, and +- unauthorized responses after forcing one AppRole re-login. + +The retry policy does not retry missing configuration, invalid requests, +permission-denied responses, missing Transit paths, malformed Vault responses, +or decrypt denials. Tune `VAULT_TIMEOUT_SECONDS` to bound individual Vault +requests and `AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS` to control how +often health and submit readiness recheck Vault. + +## Rollout Behavior + +The feature defaults to `off`, so upgrading AI-Q does not change persisted job +content unless an operator explicitly enables encryption. The first +implementation is forward-only after enablement: + +- New `job_info.output` writes are encrypted after enablement. +- New `artifact.update.data.content` event fields are encrypted after + enablement for `output` and `file` artifacts. +- Existing plaintext `job_info.output` rows are intentionally unreadable while + `AIQ_CONTENT_ENCRYPTION=key` or `vault`. +- Existing plaintext event rows remain readable in encrypted modes. +- No historical plaintext backfill is included. +- No rewrap tooling is included. + +Migration, backfill, application-managed key rotation, and decrypt-on-rollback +tooling are outside the scope of this release. Enable encryption only for a new +or empty job history, after existing jobs have expired, or after operators +accept that old plaintext final-report rows in `job_info.output` cannot be read +in encrypted modes. + +Keep the encryption mode, static key and key id, or Vault Transit mount and key +name stable while encrypted jobs are retained. Switching the mode to `off` +does not decrypt existing values: report reads can expose the stored `aiqenc:` +value and encrypted event fields remain marker objects. Restore the original +encryption configuration to read those jobs. + +## Health and Failure Behavior + +`/live` reports process liveness without calling Vault or the database. `/health` +reports dependency readiness and includes encryption status. When encryption is +configured but unready, `/health` returns HTTP 503 and new async submissions are +rejected with HTTP 503, while `/live` remains successful so an orchestrator does +not restart an otherwise live API process during a dependency outage. + +Workers independently validate encryption before marking a job `RUNNING`. If +encryption is unavailable at worker startup, the job is marked `FAILURE` and +the agent does not run. + +Each submission also carries a non-secret encryption policy identity from the +API process to the worker. The identity covers the mode and key id, a static-key +fingerprint in `key` mode, or the Vault address, namespace, Transit mount, and +Transit key in `vault` mode. The worker fails the job before `RUNNING` if its +local policy does not match, including `key` or `vault` submissions received by +a worker configured with `off`. + +If final-report encryption, artifact event-content encryption, or encrypted +persistence fails after an agent has completed, the job is marked `FAILURE`. +Timer-triggered event flush failures are surfaced before the worker can mark +the job `SUCCESS`. The worker does not fall back to writing plaintext output. +When encryption is `off`, event persistence retains its legacy best-effort +behavior: database write failures are logged without changing the job status. + +Report reads fail closed: + +- Vault or crypto unavailability returns HTTP 503. +- Plaintext, malformed, or undecryptable `job_info.output` in encrypted mode + returns HTTP 500. +- Job access is authorized before decryption is attempted. + +## Cache + +Decrypt paths use an in-memory plaintext data encryption key cache per process. +The default TTL is 15 minutes with a maximum of 1024 entries. Set +`AIQ_CONTENT_ENCRYPTION_DEK_CACHE_TTL_SECONDS=0` to disable this cache. + +The readiness cache defaults to 60 seconds. Health checks and submit requests +reuse the cached state until it becomes stale. Set +`AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS` to override the default. diff --git a/docs/source/deployment/index.md b/docs/source/deployment/index.md index 5ce070422..0c938f256 100644 --- a/docs/source/deployment/index.md +++ b/docs/source/deployment/index.md @@ -33,6 +33,8 @@ All containerized deployments run the same three services: - **[Authentication](./authentication.md)** -- Enable OAuth/OIDC sign-in, configure backend JWT validation, and use AIQ user tokens in tools and MCP pass-through integrations. +- **[Async Job Content Encryption](./content-encryption.md)** -- Configure encryption at rest for async final reports and selected artifact event content, including Vault Transit and static-key modes. + - **[Observability](./observability.md)** -- Tracing and monitoring with Phoenix, LangSmith, Weave, and OpenTelemetry. - **[Production Considerations](./production.md)** -- Guidance on managed databases, horizontal scaling, security hardening, monitoring, and resource requirements. diff --git a/docs/source/deployment/kubernetes.md b/docs/source/deployment/kubernetes.md index 008ddeef4..b981c9a0d 100644 --- a/docs/source/deployment/kubernetes.md +++ b/docs/source/deployment/kubernetes.md @@ -154,9 +154,14 @@ Once all pods are running, verify the backend is responding: ```bash kubectl port-forward -n ns-aiq svc/aiq-backend 8000:8000 & +curl http://localhost:8000/live curl http://localhost:8000/health ``` +The chart uses `/live` for the liveness probe and `/health` for readiness. The +liveness endpoint checks only that the API process responds; database or content- +encryption outages make the pod unready without causing a restart loop. + The backend API docs are available at `http://localhost:8000/docs` while the port-forward is active. ### Access the application diff --git a/docs/source/deployment/production.md b/docs/source/deployment/production.md index 165599fbb..52d2efb53 100644 --- a/docs/source/deployment/production.md +++ b/docs/source/deployment/production.md @@ -115,11 +115,18 @@ Store API keys in `deploy/.env` and ensure the file is not committed to version ## Monitoring -### Health Endpoint +### Liveness and Readiness Endpoints -The backend exposes a health endpoint at `/health` for liveness and readiness probes. +The backend exposes separate probe endpoints: + +- `/live` checks only that the API process can respond. Use it for liveness probes. +- `/health` checks database and content-encryption dependencies. Use it for readiness probes. + +This separation keeps a temporary dependency outage from restarting a live API process while still removing an +unready instance from service traffic. ```bash +curl http://localhost:8000/live curl http://localhost:8000/health ``` diff --git a/docs/source/index.md b/docs/source/index.md index b57ff8fec..6efe71720 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -98,6 +98,7 @@ Overview <./deployment/index.md> Docker Compose <./deployment/docker-compose.md> Docker Build System <./deployment/docker-build.md> Authentication <./deployment/authentication.md> +Async Job Content Encryption <./deployment/content-encryption.md> Observability <./deployment/observability.md> Production <./deployment/production.md> Kubernetes <./deployment/kubernetes.md> diff --git a/docs/source/integration/rest-api.md b/docs/source/integration/rest-api.md index 268e023ad..670421e0f 100644 --- a/docs/source/integration/rest-api.md +++ b/docs/source/integration/rest-api.md @@ -35,7 +35,8 @@ Base path: `/v1/jobs/async` | `GET` | `/v1/jobs/async/job/{job_id}/state` | Get accumulated job artifacts | | `GET` | `/v1/jobs/async/job/{job_id}/report` | Get final research report | | `GET` | `/v1/data_sources` | List available data sources | -| `GET` | `/health` | Health check (includes Dask status) | +| `GET` | `/live` | Process liveness check (no dependency checks) | +| `GET` | `/health` | Dependency readiness check (database, Dask, and content encryption) | ### List Available Agents @@ -451,7 +452,23 @@ curl http://localhost:8000/v1/data_sources The `knowledge_layer` entry only appears when a knowledge retrieval function is configured. -### Health Check +### Liveness and Readiness Checks + +Use `/live` for process liveness probes. It returns success without checking the +database, Dask, or content-encryption dependencies. + +```bash +curl http://localhost:8000/live +``` + +```json +{ + "status": "alive" +} +``` + +Use `/health` for readiness checks. It returns HTTP 503 when a required +dependency is unavailable. ```bash curl http://localhost:8000/health @@ -461,8 +478,13 @@ curl http://localhost:8000/health ```json { - "status": "ok", - "dask_available": true + "status": "healthy", + "dask_available": true, + "db": "ok", + "encryption": { + "mode": "off", + "ready": true + } } ``` diff --git a/frontends/aiq_api/pyproject.toml b/frontends/aiq_api/pyproject.toml index a6a9611f3..883d67e57 100644 --- a/frontends/aiq_api/pyproject.toml +++ b/frontends/aiq_api/pyproject.toml @@ -47,6 +47,8 @@ dependencies = [ "python-multipart>=0.0.27,<0.1", # JWT validation for AuthMiddleware (RS256 via JWKS) "PyJWT[cryptography]>=2.8.0", + # Vault Transit support for async final-report encryption at rest. + "hvac>=2.3.0,<3", ] [project.optional-dependencies] diff --git a/frontends/aiq_api/src/aiq_api/auth/middleware.py b/frontends/aiq_api/src/aiq_api/auth/middleware.py index f635f4e2d..36e9f420f 100644 --- a/frontends/aiq_api/src/aiq_api/auth/middleware.py +++ b/frontends/aiq_api/src/aiq_api/auth/middleware.py @@ -182,6 +182,7 @@ def build_request_trace_tags( # Paths reachable from outside the cluster. Any external request whose path # does NOT match one of these receives 404. Prefix entries must end with "/". EXTERNAL_ALLOWED_PATHS: list[str] = [ + "/live", "/health", "/docs", "/redoc", @@ -197,7 +198,7 @@ def build_request_trace_tags( ] # External paths that require no token (monitoring, etc.) -AUTH_EXEMPT_PATHS: set[str] = {"/health", "/docs", "/redoc", "/openapi.json"} +AUTH_EXEMPT_PATHS: set[str] = {"/live", "/health", "/docs", "/redoc", "/openapi.json"} def _is_oauth_callback_path(path: str) -> bool: diff --git a/frontends/aiq_api/src/aiq_api/jobs/crypto.py b/frontends/aiq_api/src/aiq_api/jobs/crypto.py new file mode 100644 index 000000000..b6ca02679 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/jobs/crypto.py @@ -0,0 +1,1165 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +"""Content encryption for AI-Q async jobs. + +This module covers ``job_info.output`` and selected sensitive fields in +``job_events.event_data`` for the AI-Q async API. Checkpoint state, summaries, +and errors remain plaintext. +""" + +from __future__ import annotations + +import asyncio +import base64 +import hashlib +import json +import logging +import math +import os +import random +import threading +import time +from collections import OrderedDict +from dataclasses import dataclass +from dataclasses import field +from typing import Any + +from cryptography.exceptions import InvalidTag +from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from pydantic import BaseModel + +logger = logging.getLogger(__name__) + +ENVELOPE_PREFIX = "aiqenc:" +ENCRYPTED_FIELD_MARKER = "_aiq_encrypted" +ENCRYPTED_FIELD_VALUE = "value" +ENVELOPE_VERSION = 1 +CONTENT_ALGORITHM = "AES-256-GCM" +CONTENT_ENCRYPTION_POLICY_VERSION = 1 +DEK_BYTES = 32 +GCM_NONCE_BYTES = 12 +GCM_TAG_BYTES = 16 +DEFAULT_READINESS_TTL_SECONDS = 60.0 +DEFAULT_DEK_CACHE_TTL_SECONDS = 900.0 +DEFAULT_DEK_CACHE_MAX_ENTRIES = 1024 +DEFAULT_VAULT_TRANSIT_MOUNT = "transit" +DEFAULT_VAULT_TIMEOUT_SECONDS = 5.0 +_VAULT_ATTEMPTS = 2 +_VAULT_RETRY_BASE_SECONDS = 0.05 +_VAULT_RETRY_JITTER_SECONDS = 0.025 +_VAULT_AUTH_RETRY_EXCEPTION_NAMES = {"Unauthorized"} +_VAULT_RETRYABLE_EXCEPTION_NAMES = { + "BadGateway", + "InternalServerError", + "RateLimitExceeded", + "VaultDown", +} +_VAULT_NON_RETRYABLE_EXCEPTION_NAMES = { + "Forbidden", + "InvalidPath", + "InvalidRequest", + "ParamValidationError", + "PreconditionFailed", + "UnsupportedOperation", + "VaultNotInitialized", +} + + +class ContentEncryptionError(Exception): + """Base class for content-encryption failures.""" + + +class ContentEncryptionConfigError(ContentEncryptionError, ValueError): + """Invalid encryption configuration. This is a startup-hard failure.""" + + +class ContentEncryptionUnavailable(ContentEncryptionError): + """Encryption is configured but operationally unavailable.""" + + +class ContentEncryptionInvalidData(ContentEncryptionError): + """Persisted encrypted data is malformed, plaintext, or undecryptable.""" + + +class ContentEncryptionPlaintextViolation(ContentEncryptionInvalidData): + """Encrypted mode encountered plaintext where an envelope is required.""" + + +class ContentEncryptionPolicyMismatch(ContentEncryptionError): + """API and worker content-encryption policies do not match.""" + + +@dataclass(frozen=True) +class ContentEncryptionPolicyIdentity: + """Non-secret encryption identity propagated from the API to a worker.""" + + version: int + mode: str + key_id: str | None = None + static_key_fingerprint: str | None = field(default=None, repr=False) + vault_addr: str | None = None + vault_namespace: str | None = None + vault_transit_mount: str | None = None + vault_transit_key: str | None = None + + +@dataclass(frozen=True) +class ContentEncryptionConfig: + """Process-local content-encryption configuration parsed from env.""" + + mode: str + key_id: str | None = None + static_key: bytes | None = field(default=None, repr=False) + vault_addr: str | None = None + vault_namespace: str | None = None + vault_transit_mount: str = DEFAULT_VAULT_TRANSIT_MOUNT + vault_transit_key: str | None = None + vault_role_id: str | None = field(default=None, repr=False) + vault_secret_id: str | None = field(default=None, repr=False) + vault_timeout_seconds: float = DEFAULT_VAULT_TIMEOUT_SECONDS + readiness_ttl_seconds: float = DEFAULT_READINESS_TTL_SECONDS + dek_cache_ttl_seconds: float = DEFAULT_DEK_CACHE_TTL_SECONDS + dek_cache_max_entries: int = DEFAULT_DEK_CACHE_MAX_ENTRIES + + @property + def encrypted(self) -> bool: + return self.mode in {"key", "vault"} + + @property + def effective_key_id(self) -> str: + if self.key_id: + return self.key_id + if self.mode == "key": + return "static-key" + if self.mode == "vault": + return f"{self.vault_transit_mount}/{self.vault_transit_key}" + return "off" + + @property + def signature(self) -> tuple[Any, ...]: + return ( + self.mode, + self.key_id, + _secret_fingerprint(self.static_key), + self.vault_addr, + self.vault_namespace, + self.vault_transit_mount, + self.vault_transit_key, + _secret_fingerprint(self.vault_role_id), + _secret_fingerprint(self.vault_secret_id), + self.vault_timeout_seconds, + self.readiness_ttl_seconds, + self.dek_cache_ttl_seconds, + self.dek_cache_max_entries, + ) + + @property + def policy_identity(self) -> ContentEncryptionPolicyIdentity: + """Return the non-secret identity workers must match for submitted jobs.""" + + if self.mode == "key": + return ContentEncryptionPolicyIdentity( + version=CONTENT_ENCRYPTION_POLICY_VERSION, + mode=self.mode, + key_id=self.effective_key_id, + static_key_fingerprint=_secret_fingerprint(self.static_key), + ) + if self.mode == "vault": + return ContentEncryptionPolicyIdentity( + version=CONTENT_ENCRYPTION_POLICY_VERSION, + mode=self.mode, + key_id=self.effective_key_id, + vault_addr=self.vault_addr, + vault_namespace=self.vault_namespace, + vault_transit_mount=self.vault_transit_mount, + vault_transit_key=self.vault_transit_key, + ) + return ContentEncryptionPolicyIdentity( + version=CONTENT_ENCRYPTION_POLICY_VERSION, + mode="off", + ) + + +@dataclass(frozen=True) +class ContentEncryptionReadiness: + """Cached readiness state safe to expose through health responses.""" + + mode: str + ready: bool + checked_at: float | None = None + reason: str | None = None + exception_type: str | None = None + encrypt_ready: bool | None = None + decrypt_ready: bool | None = None + + def to_health_dict(self) -> dict[str, Any]: + result: dict[str, Any] = {"mode": self.mode, "ready": self.ready} + if self.encrypt_ready is not None: + result["encrypt_ready"] = self.encrypt_ready + if self.decrypt_ready is not None: + result["decrypt_ready"] = self.decrypt_ready + if self.reason: + result["reason"] = self.reason + if self.exception_type: + result["exception_type"] = self.exception_type + return result + + +@dataclass(frozen=True) +class WrappedDEK: + wrap: str + kid: str + wrapped_dek: str + wrap_nonce: str | None = None + wrap_tag: str | None = None + + +@dataclass +class JobContentCipher: + """Per-job encryption context used by the worker final-output write.""" + + manager: ContentEncryptionManager + job_id: str + dek: bytes | None = None + wrapped_dek: WrappedDEK | None = None + + def encrypt_output_json(self, output_json: str) -> str: + if not self.manager.config.encrypted: + return output_json + if self.dek is None or self.wrapped_dek is None: + raise ContentEncryptionUnavailable("encrypted job output cipher is not initialized") + return self.manager.encrypt_job_output_text(self.job_id, output_json, self.dek, self.wrapped_dek) + + def encrypt_event_field_json(self, field_path: str, value_json: str) -> str: + if not self.manager.config.encrypted: + return value_json + if self.dek is None or self.wrapped_dek is None: + raise ContentEncryptionUnavailable("encrypted job event cipher is not initialized") + return self.manager.encrypt_job_event_field_text( + self.job_id, + field_path, + value_json, + self.dek, + self.wrapped_dek, + ) + + +@dataclass +class _DEKCacheEntry: + dek: bytes + expires_at: float + + +@dataclass +class _DEKCache: + ttl_seconds: float + max_entries: int + _entries: OrderedDict[str, _DEKCacheEntry] = field(default_factory=OrderedDict) + _lock: threading.RLock = field(default_factory=threading.RLock) + + def get(self, cache_key: str) -> bytes | None: + if self.ttl_seconds <= 0: + return None + with self._lock: + now = time.monotonic() + entry = self._entries.get(cache_key) + if entry is None: + return None + if entry.expires_at <= now: + self._entries.pop(cache_key, None) + return None + self._entries.move_to_end(cache_key) + return entry.dek + + def put(self, cache_key: str, dek: bytes) -> None: + if self.ttl_seconds <= 0: + return + with self._lock: + now = time.monotonic() + self._entries[cache_key] = _DEKCacheEntry(dek=dek, expires_at=now + self.ttl_seconds) + self._entries.move_to_end(cache_key) + self._evict_locked(now) + + def _evict(self, now: float) -> None: + with self._lock: + self._evict_locked(now) + + def _evict_locked(self, now: float) -> None: + expired = [key for key, entry in self._entries.items() if entry.expires_at <= now] + for key in expired: + self._entries.pop(key, None) + while len(self._entries) > self.max_entries: + self._entries.popitem(last=False) + + +class ContentEncryptionManager: + """Encrypt/decrypt async final report payloads for one process config.""" + + def __init__(self, config: ContentEncryptionConfig): + self.config = config + self._lock = threading.RLock() + self._dek_cache = _DEKCache( + ttl_seconds=config.dek_cache_ttl_seconds, + max_entries=config.dek_cache_max_entries, + ) + self._readiness = ContentEncryptionReadiness(mode=config.mode, ready=not config.encrypted) + self._vault_client: _VaultTransitClient | None = None + + def get_readiness(self) -> ContentEncryptionReadiness: + with self._lock: + if self.config.mode == "off": + return ContentEncryptionReadiness(mode="off", ready=True) + return self._readiness + + def check_readiness(self, *, force: bool = False, operation: str = "readiness") -> ContentEncryptionReadiness: + with self._lock: + if self.config.mode == "off": + self._readiness = ContentEncryptionReadiness(mode="off", ready=True) + return self._readiness + if self.config.mode == "key": + self._readiness = ContentEncryptionReadiness(mode="key", ready=True, checked_at=time.monotonic()) + return self._readiness + if not force and self._readiness.checked_at is not None: + age = time.monotonic() - self._readiness.checked_at + if age < self.config.readiness_ttl_seconds: + return self._readiness + + self._readiness = self._check_vault_readiness(operation=operation) + return self._readiness + + def _check_vault_readiness(self, *, operation: str) -> ContentEncryptionReadiness: + dek: bytes | None = None + unwrapped_dek: bytes | None = None + try: + dek, wrapped = self._vault().generate_data_key(operation=f"{operation}_readiness_generate") + try: + unwrapped_dek = self._vault().unwrap_dek( + wrapped.wrapped_dek, + operation=f"{operation}_readiness_decrypt", + ) + except ContentEncryptionUnavailable as exc: + return self._failed_vault_readiness( + operation=operation, + reason="vault_decrypt_unavailable", + exception=exc, + encrypt_ready=True, + decrypt_ready=False, + ) + if unwrapped_dek != dek: + return self._failed_vault_readiness( + operation=operation, + reason="vault_readiness_dek_mismatch", + exception=None, + encrypt_ready=True, + decrypt_ready=False, + ) + return ContentEncryptionReadiness( + mode=self.config.mode, + ready=True, + checked_at=time.monotonic(), + encrypt_ready=True, + decrypt_ready=True, + ) + except ContentEncryptionUnavailable as exc: + return self._failed_vault_readiness( + operation=operation, + reason="vault_generate_unavailable", + exception=exc, + encrypt_ready=False, + decrypt_ready=False, + ) + finally: + if dek is not None: + _zero_bytes(dek) + if unwrapped_dek is not None: + _zero_bytes(unwrapped_dek) + + def _failed_vault_readiness( + self, + *, + operation: str, + reason: str, + exception: Exception | None, + encrypt_ready: bool, + decrypt_ready: bool, + ) -> ContentEncryptionReadiness: + readiness = ContentEncryptionReadiness( + mode=self.config.mode, + ready=False, + checked_at=time.monotonic(), + reason=reason, + exception_type=exception.__class__.__name__ if exception is not None else None, + encrypt_ready=encrypt_ready, + decrypt_ready=decrypt_ready, + ) + logger.warning( + "Content encryption readiness failed mode=%s operation=%s reason=%s exception=%s", + self.config.mode, + operation, + readiness.reason, + readiness.exception_type, + ) + return readiness + + def require_ready(self, *, operation: str) -> None: + readiness = self.check_readiness(force=False, operation=operation) + if self.config.encrypted and not readiness.ready: + raise ContentEncryptionUnavailable(readiness.reason or "content_encryption_unready") + + def create_job_cipher(self, job_id: str) -> JobContentCipher: + if self.config.mode == "off": + return JobContentCipher(manager=self, job_id=job_id) + if self.config.mode == "key": + dek = os.urandom(DEK_BYTES) + wrapped = self._wrap_dek_with_static_key(dek) + return JobContentCipher(manager=self, job_id=job_id, dek=dek, wrapped_dek=wrapped) + if self.config.mode == "vault": + dek, wrapped = self._vault().generate_data_key(operation="worker_job_cipher") + return JobContentCipher(manager=self, job_id=job_id, dek=dek, wrapped_dek=wrapped) + raise ContentEncryptionConfigError(f"Unsupported content encryption mode: {self.config.mode}") + + def encrypt_job_output_text(self, job_id: str, output_json: str, dek: bytes, wrapped_dek: WrappedDEK) -> str: + aad = job_output_aad(job_id) + return self._encrypt_text(output_json, dek, wrapped_dek, aad) + + def encrypt_job_event_field_text( + self, + job_id: str, + field_path: str, + value_json: str, + dek: bytes, + wrapped_dek: WrappedDEK, + ) -> str: + aad = job_event_field_aad(job_id, field_path) + return self._encrypt_text(value_json, dek, wrapped_dek, aad) + + def _encrypt_text(self, plaintext: str, dek: bytes, wrapped_dek: WrappedDEK, aad: str) -> str: + nonce = os.urandom(GCM_NONCE_BYTES) + encrypted = AESGCM(dek).encrypt(nonce, plaintext.encode("utf-8"), aad.encode("utf-8")) + ciphertext, tag = encrypted[:-GCM_TAG_BYTES], encrypted[-GCM_TAG_BYTES:] + envelope = { + "v": ENVELOPE_VERSION, + "alg": CONTENT_ALGORITHM, + "wrap": wrapped_dek.wrap, + "kid": wrapped_dek.kid, + "aad_hint": aad, + "wrapped_dek": wrapped_dek.wrapped_dek, + "nonce": _b64url_encode(nonce), + "ciphertext": _b64url_encode(ciphertext), + "tag": _b64url_encode(tag), + } + if wrapped_dek.wrap_nonce is not None: + envelope["wrap_nonce"] = wrapped_dek.wrap_nonce + if wrapped_dek.wrap_tag is not None: + envelope["wrap_tag"] = wrapped_dek.wrap_tag + return encode_envelope(envelope) + + def decrypt_job_output_text(self, job_id: str, stored_output: str) -> str: + aad = job_output_aad(job_id) + return self._decrypt_text(stored_output, aad) + + def decrypt_job_event_field_text(self, job_id: str, field_path: str, stored_value: str) -> str: + aad = job_event_field_aad(job_id, field_path) + return self._decrypt_text(stored_value, aad) + + def _decrypt_text(self, stored_value: str, aad: str) -> str: + envelope = decode_envelope(stored_value) + dek = self._unwrap_dek(envelope) + try: + nonce = _required_b64url(envelope, "nonce") + ciphertext = _required_b64url(envelope, "ciphertext") + tag = _required_b64url(envelope, "tag") + plaintext = AESGCM(dek).decrypt(nonce, ciphertext + tag, aad.encode("utf-8")) + except InvalidTag as exc: + raise ContentEncryptionInvalidData("encrypted job output failed authentication") from exc + except ValueError as exc: + raise ContentEncryptionInvalidData("encrypted job output is malformed") from exc + return plaintext.decode("utf-8") + + def _unwrap_dek(self, envelope: dict[str, Any]) -> bytes: + self._validate_envelope_metadata(envelope) + cache_key = _wrapped_dek_cache_key(envelope) + cached = self._dek_cache.get(cache_key) + if cached is not None: + return cached + + wrap = envelope["wrap"] + if wrap == "key": + dek = self._unwrap_dek_with_static_key(envelope) + elif wrap == "vault": + dek = self._vault().unwrap_dek(str(envelope["wrapped_dek"]), operation="decrypt_job_output") + else: + raise ContentEncryptionInvalidData("encrypted job output uses an unsupported wrap type") + + self._dek_cache.put(cache_key, dek) + return dek + + def _validate_envelope_metadata(self, envelope: dict[str, Any]) -> None: + if envelope.get("v") != ENVELOPE_VERSION: + raise ContentEncryptionInvalidData("encrypted job output uses an unsupported envelope version") + if envelope.get("alg") != CONTENT_ALGORITHM: + raise ContentEncryptionInvalidData("encrypted job output uses an unsupported content algorithm") + if envelope.get("wrap") not in {"key", "vault"}: + raise ContentEncryptionInvalidData("encrypted job output uses an unsupported wrap type") + if envelope.get("wrap") != self.config.mode: + raise ContentEncryptionInvalidData("encrypted job output wrap does not match configured mode") + if not isinstance(envelope.get("kid"), str) or not envelope["kid"]: + raise ContentEncryptionInvalidData("encrypted job output is missing a key id") + if not isinstance(envelope.get("wrapped_dek"), str) or not envelope["wrapped_dek"]: + raise ContentEncryptionInvalidData("encrypted job output is missing a wrapped DEK") + + def _wrap_dek_with_static_key(self, dek: bytes) -> WrappedDEK: + key = self.config.static_key + if key is None: + raise ContentEncryptionConfigError("AIQ_CONTENT_ENCRYPTION_KEY is required in key mode") + kid = self.config.effective_key_id + nonce = os.urandom(GCM_NONCE_BYTES) + encrypted = AESGCM(key).encrypt(nonce, dek, _kek_wrap_aad(kid)) + ciphertext, tag = encrypted[:-GCM_TAG_BYTES], encrypted[-GCM_TAG_BYTES:] + return WrappedDEK( + wrap="key", + kid=kid, + wrapped_dek=_b64url_encode(ciphertext), + wrap_nonce=_b64url_encode(nonce), + wrap_tag=_b64url_encode(tag), + ) + + def _unwrap_dek_with_static_key(self, envelope: dict[str, Any]) -> bytes: + key = self.config.static_key + if key is None: + raise ContentEncryptionConfigError("AIQ_CONTENT_ENCRYPTION_KEY is required in key mode") + if envelope.get("kid") != self.config.effective_key_id: + raise ContentEncryptionInvalidData("encrypted job output key id does not match configured key") + try: + nonce = _required_b64url(envelope, "wrap_nonce") + ciphertext = _required_b64url(envelope, "wrapped_dek") + tag = _required_b64url(envelope, "wrap_tag") + dek = AESGCM(key).decrypt(nonce, ciphertext + tag, _kek_wrap_aad(self.config.effective_key_id)) + except InvalidTag as exc: + raise ContentEncryptionInvalidData("encrypted job output DEK unwrap failed") from exc + except ValueError as exc: + raise ContentEncryptionInvalidData("encrypted job output DEK wrapper is malformed") from exc + if len(dek) != DEK_BYTES: + raise ContentEncryptionInvalidData("encrypted job output DEK has invalid length") + return dek + + def _vault(self) -> _VaultTransitClient: + with self._lock: + if self._vault_client is None: + self._vault_client = _VaultTransitClient(self.config) + return self._vault_client + + +class _VaultTransitClient: + """Small synchronous Vault Transit client with bounded retry.""" + + def __init__(self, config: ContentEncryptionConfig): + self._config = config + self._lock = threading.RLock() + self._client: Any | None = None + + def generate_data_key(self, *, operation: str) -> tuple[bytes, WrappedDEK]: + data = self._with_retry(lambda: self._generate_data_key_once(), operation=operation) + try: + plaintext = data["data"]["plaintext"] + ciphertext = data["data"]["ciphertext"] + except (KeyError, TypeError) as exc: + raise ContentEncryptionUnavailable("vault_datakey_response_invalid") from exc + + dek = _decode_vault_key(plaintext, field_name="vault plaintext data key") + if len(dek) != DEK_BYTES: + raise ContentEncryptionUnavailable("vault_datakey_invalid_length") + if not isinstance(ciphertext, str) or not ciphertext: + raise ContentEncryptionUnavailable("vault_datakey_missing_ciphertext") + return dek, WrappedDEK( + wrap="vault", + kid=self._config.effective_key_id, + wrapped_dek=ciphertext, + ) + + def unwrap_dek(self, wrapped_dek: str, *, operation: str) -> bytes: + data = self._with_retry(lambda: self._unwrap_dek_once(wrapped_dek), operation=operation) + try: + plaintext = data["data"]["plaintext"] + except (KeyError, TypeError) as exc: + raise ContentEncryptionUnavailable("vault_decrypt_response_invalid") from exc + dek = _decode_vault_key(plaintext, field_name="vault plaintext data key") + if len(dek) != DEK_BYTES: + raise ContentEncryptionUnavailable("vault_decrypt_invalid_length") + return dek + + def _generate_data_key_once(self) -> dict[str, Any]: + return self._authenticated_client().secrets.transit.generate_data_key( + name=self._required(self._config.vault_transit_key, "AIQ_ENCRYPTION_TRANSIT_KEY"), + key_type="plaintext", + bits=DEK_BYTES * 8, + mount_point=self._config.vault_transit_mount, + ) + + def _unwrap_dek_once(self, wrapped_dek: str) -> dict[str, Any]: + return self._authenticated_client().secrets.transit.decrypt_data( + name=self._required(self._config.vault_transit_key, "AIQ_ENCRYPTION_TRANSIT_KEY"), + ciphertext=wrapped_dek, + mount_point=self._config.vault_transit_mount, + ) + + def _with_retry(self, fn, *, operation: str) -> dict[str, Any]: + last_exc: Exception | None = None + for attempt in range(_VAULT_ATTEMPTS): + try: + with self._lock: + return fn() + except Exception as exc: # hvac exposes several operational exception classes. + if isinstance(exc, ContentEncryptionConfigError): + raise + last_exc = exc + if not _should_retry_vault_failure(exc): + break + if attempt + 1 < _VAULT_ATTEMPTS: + if _is_auth_failure(exc): + try: + self._login(force=True) + except Exception as login_exc: # AppRole login can fail transiently too. + if isinstance(login_exc, ContentEncryptionConfigError): + raise + last_exc = login_exc + if not _should_retry_vault_failure(login_exc): + break + _sleep_before_vault_retry(attempt) + continue + + assert last_exc is not None + logger.warning( + "Vault transit operation failed mode=vault operation=%s retryable=%s exception=%s", + operation, + _should_retry_vault_failure(last_exc), + last_exc.__class__.__name__, + ) + raise ContentEncryptionUnavailable(f"vault_{operation}_failed") from last_exc + + def _authenticated_client(self) -> Any: + with self._lock: + client = self._get_client() + if not client.is_authenticated(): + self._login(force=True) + return client + + def _get_client(self) -> Any: + with self._lock: + if self._client is not None: + return self._client + try: + import hvac + except ImportError as exc: + raise ContentEncryptionConfigError("Vault mode requires the hvac package") from exc + + self._client = hvac.Client( + url=self._required(self._config.vault_addr, "VAULT_ADDR"), + namespace=self._config.vault_namespace, + timeout=self._config.vault_timeout_seconds, + ) + self._login(force=True) + return self._client + + def _login(self, *, force: bool) -> None: + with self._lock: + client = self._get_client_without_login() + if not force and client.is_authenticated(): + return + try: + client.auth.approle.login( + role_id=self._required(self._config.vault_role_id, "VAULT_ROLE_ID"), + secret_id=self._required(self._config.vault_secret_id, "VAULT_SECRET_ID"), + ) + except Exception as exc: + logger.warning("Vault AppRole login failed mode=vault exception=%s", exc.__class__.__name__) + raise ContentEncryptionUnavailable("vault_approle_login_failed") from exc + + def _get_client_without_login(self) -> Any: + with self._lock: + if self._client is not None: + return self._client + try: + import hvac + except ImportError as exc: + raise ContentEncryptionConfigError("Vault mode requires the hvac package") from exc + self._client = hvac.Client( + url=self._required(self._config.vault_addr, "VAULT_ADDR"), + namespace=self._config.vault_namespace, + timeout=self._config.vault_timeout_seconds, + ) + return self._client + + @staticmethod + def _required(value: str | None, name: str) -> str: + if not value: + raise ContentEncryptionConfigError(f"{name} is required in vault mode") + return value + + +_manager: ContentEncryptionManager | None = None +_manager_signature: tuple[Any, ...] | None = None +_manager_lock = threading.RLock() + + +def reset_content_encryption_manager_for_tests() -> None: + global _manager + global _manager_signature + with _manager_lock: + _manager = None + _manager_signature = None + + +def get_content_encryption_config() -> ContentEncryptionConfig: + mode = os.environ.get("AIQ_CONTENT_ENCRYPTION", "off").strip().lower() + if mode not in {"off", "key", "vault"}: + raise ContentEncryptionConfigError("AIQ_CONTENT_ENCRYPTION must be one of: off, key, vault") + + readiness_ttl = _parse_non_negative_float( + os.environ.get("AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS"), + default=DEFAULT_READINESS_TTL_SECONDS, + name="AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS", + ) + dek_cache_ttl = _parse_non_negative_float( + os.environ.get("AIQ_CONTENT_ENCRYPTION_DEK_CACHE_TTL_SECONDS"), + default=DEFAULT_DEK_CACHE_TTL_SECONDS, + name="AIQ_CONTENT_ENCRYPTION_DEK_CACHE_TTL_SECONDS", + ) + + key_id = _empty_to_none(os.environ.get("AIQ_CONTENT_ENCRYPTION_KEY_ID")) + if mode == "off": + return ContentEncryptionConfig( + mode="off", + key_id=key_id, + readiness_ttl_seconds=readiness_ttl, + dek_cache_ttl_seconds=dek_cache_ttl, + ) + if mode == "key": + raw_key = os.environ.get("AIQ_CONTENT_ENCRYPTION_KEY") + if not raw_key: + raise ContentEncryptionConfigError("AIQ_CONTENT_ENCRYPTION_KEY is required in key mode") + key = _strict_base64_decode(raw_key, field_name="AIQ_CONTENT_ENCRYPTION_KEY") + if len(key) != DEK_BYTES: + raise ContentEncryptionConfigError("AIQ_CONTENT_ENCRYPTION_KEY must decode to exactly 32 bytes") + return ContentEncryptionConfig( + mode="key", + key_id=key_id, + static_key=key, + readiness_ttl_seconds=readiness_ttl, + dek_cache_ttl_seconds=dek_cache_ttl, + ) + + vault_addr = _empty_to_none(os.environ.get("VAULT_ADDR")) + vault_transit_key = _empty_to_none(os.environ.get("AIQ_ENCRYPTION_TRANSIT_KEY")) + vault_role_id = _empty_to_none(os.environ.get("VAULT_ROLE_ID")) + vault_secret_id = _empty_to_none(os.environ.get("VAULT_SECRET_ID")) + if vault_addr is None: + raise ContentEncryptionConfigError("VAULT_ADDR is required in vault mode") + if vault_transit_key is None: + raise ContentEncryptionConfigError("AIQ_ENCRYPTION_TRANSIT_KEY is required in vault mode") + if vault_role_id is None: + raise ContentEncryptionConfigError("VAULT_ROLE_ID is required in vault mode") + if vault_secret_id is None: + raise ContentEncryptionConfigError("VAULT_SECRET_ID is required in vault mode") + + return ContentEncryptionConfig( + mode="vault", + key_id=key_id, + vault_addr=vault_addr, + vault_namespace=_empty_to_none(os.environ.get("VAULT_NAMESPACE")), + vault_transit_mount=os.environ.get("VAULT_TRANSIT_MOUNT", DEFAULT_VAULT_TRANSIT_MOUNT).strip() + or DEFAULT_VAULT_TRANSIT_MOUNT, + vault_transit_key=vault_transit_key, + vault_role_id=vault_role_id, + vault_secret_id=vault_secret_id, + vault_timeout_seconds=_parse_positive_float( + os.environ.get("VAULT_TIMEOUT_SECONDS"), + default=DEFAULT_VAULT_TIMEOUT_SECONDS, + name="VAULT_TIMEOUT_SECONDS", + ), + readiness_ttl_seconds=readiness_ttl, + dek_cache_ttl_seconds=dek_cache_ttl, + ) + + +def get_content_encryption_manager() -> ContentEncryptionManager: + global _manager + global _manager_signature + config = get_content_encryption_config() + with _manager_lock: + if _manager is None or _manager_signature != config.signature: + _manager = ContentEncryptionManager(config) + _manager_signature = config.signature + return _manager + + +def validate_content_encryption_startup() -> ContentEncryptionReadiness: + """Validate config and eagerly check operational readiness. + + Configuration errors propagate and fail startup. Vault operational failures + are cached as unhealthy so the process can start but reject submissions. + """ + + manager = get_content_encryption_manager() + return manager.check_readiness(force=True, operation="api_startup") + + +async def validate_content_encryption_startup_async() -> ContentEncryptionReadiness: + """Validate startup readiness without blocking the FastAPI event loop.""" + + return await asyncio.to_thread(validate_content_encryption_startup) + + +def get_content_encryption_health() -> ContentEncryptionReadiness: + return get_content_encryption_manager().check_readiness(force=False, operation="health") + + +async def get_content_encryption_health_async() -> ContentEncryptionReadiness: + """Check encryption health without blocking the FastAPI event loop.""" + + return await asyncio.to_thread(get_content_encryption_health) + + +def require_content_encryption_ready_for_submission() -> None: + get_content_encryption_manager().require_ready(operation="submit") + + +async def require_content_encryption_ready_for_submission_async() -> None: + """Check submission readiness without blocking the FastAPI event loop.""" + + await asyncio.to_thread(require_content_encryption_ready_for_submission) + + +def get_content_encryption_policy_identity() -> ContentEncryptionPolicyIdentity: + """Return the process-local policy identity safe to send to a worker.""" + + return get_content_encryption_manager().config.policy_identity + + +def require_content_encryption_policy(expected: ContentEncryptionPolicyIdentity | None) -> None: + """Fail closed unless the worker policy exactly matches the submitter policy.""" + + if expected is None or get_content_encryption_policy_identity() != expected: + raise ContentEncryptionPolicyMismatch("worker content encryption policy does not match submission policy") + + +def create_job_content_cipher(job_id: str) -> JobContentCipher: + return get_content_encryption_manager().create_job_cipher(job_id) + + +async def update_job_output( + job_store: Any, + job_id: str, + status: Any, + *, + output: BaseModel | dict[str, Any] | list[Any] | str, + cipher: JobContentCipher, +) -> None: + """Update NAT JobStore output, encrypting only in encrypted modes.""" + + if not cipher.manager.config.encrypted: + await job_store.update_status(job_id, status, output=output) + return + + output_json = _serialize_output_json(output) + encrypted_output = cipher.encrypt_output_json(output_json) + await job_store.update_status(job_id, status, output=encrypted_output) + + +def read_job_output(job_id: str, stored_output: Any) -> Any: + """Read job output, decrypting only when encrypted mode is configured.""" + + if stored_output is None: + return None + + manager = get_content_encryption_manager() + if not manager.config.encrypted: + if isinstance(stored_output, str): + try: + return json.loads(stored_output) + except json.JSONDecodeError: + return stored_output + return stored_output + + if not isinstance(stored_output, str) or not stored_output.startswith(ENVELOPE_PREFIX): + logger.warning( + "Plaintext job output encountered in encrypted mode mode=%s job_id=%s operation=read_job_output", + manager.config.mode, + job_id, + ) + raise ContentEncryptionPlaintextViolation("job_info.output is plaintext in encrypted mode") + + plaintext_json = manager.decrypt_job_output_text(job_id, stored_output) + try: + return json.loads(plaintext_json) + except json.JSONDecodeError as exc: + raise ContentEncryptionInvalidData("decrypted job output is not valid JSON") from exc + + +async def read_job_output_async(job_id: str, stored_output: Any) -> Any: + """Read/decrypt job output without blocking the FastAPI event loop.""" + + return await asyncio.to_thread(read_job_output, job_id, stored_output) + + +def encrypt_event_field(job_id: str, field_path: str, value: Any, cipher: JobContentCipher | None) -> Any: + """Encrypt one event payload field, preserving plaintext behavior in off mode.""" + + if cipher is None or not cipher.manager.config.encrypted: + return value + + value_json = json.dumps(value) + encrypted_value = cipher.encrypt_event_field_json(field_path, value_json) + return { + ENCRYPTED_FIELD_MARKER: True, + ENCRYPTED_FIELD_VALUE: encrypted_value, + } + + +def decrypt_event_field(job_id: str, field_path: str, stored_value: Any) -> Any: + """Decrypt one event payload field when it uses the encrypted-field marker.""" + + if not is_encrypted_event_field(stored_value): + return stored_value + + manager = get_content_encryption_manager() + if not manager.config.encrypted: + return stored_value + + plaintext_json = manager.decrypt_job_event_field_text(job_id, field_path, stored_value[ENCRYPTED_FIELD_VALUE]) + try: + return json.loads(plaintext_json) + except json.JSONDecodeError as exc: + raise ContentEncryptionInvalidData("decrypted event field is not valid JSON") from exc + + +def is_encrypted_event_field(value: Any) -> bool: + return ( + isinstance(value, dict) + and value.get(ENCRYPTED_FIELD_MARKER) is True + and isinstance(value.get(ENCRYPTED_FIELD_VALUE), str) + and value[ENCRYPTED_FIELD_VALUE].startswith(ENVELOPE_PREFIX) + ) + + +def encode_envelope(envelope: dict[str, Any]) -> str: + data = json.dumps(envelope, separators=(",", ":"), sort_keys=True).encode("utf-8") + return ENVELOPE_PREFIX + _b64url_encode(data) + + +def decode_envelope(value: str) -> dict[str, Any]: + if not isinstance(value, str) or not value.startswith(ENVELOPE_PREFIX): + raise ContentEncryptionPlaintextViolation("job_info.output is not an aiqenc envelope") + encoded = value[len(ENVELOPE_PREFIX) :] + try: + raw = _b64url_decode(encoded) + envelope = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as exc: + raise ContentEncryptionInvalidData("encrypted job output envelope is malformed") from exc + if not isinstance(envelope, dict): + raise ContentEncryptionInvalidData("encrypted job output envelope is malformed") + return envelope + + +def job_output_aad(job_id: str) -> str: + return f"aiq:v1:job_info:output:{job_id}" + + +def job_event_field_aad(job_id: str, field_path: str) -> str: + return f"aiq:v1:job_events:event_data:{job_id}:{field_path}" + + +def _serialize_output_json(output: BaseModel | dict[str, Any] | list[Any] | str) -> str: + if isinstance(output, BaseModel): + return output.model_dump_json(round_trip=True) + if isinstance(output, str): + return output + return json.dumps(output) + + +def _kek_wrap_aad(key_id: str) -> bytes: + return f"aiq:kek-wrap:v1:{key_id}".encode() + + +def _wrapped_dek_cache_key(envelope: dict[str, Any]) -> str: + pieces = [ + str(envelope.get("wrap", "")), + str(envelope.get("kid", "")), + str(envelope.get("wrapped_dek", "")), + str(envelope.get("wrap_nonce", "")), + str(envelope.get("wrap_tag", "")), + ] + return hashlib.sha256("\0".join(pieces).encode("utf-8")).hexdigest() + + +def _required_b64url(envelope: dict[str, Any], field_name: str) -> bytes: + value = envelope.get(field_name) + if not isinstance(value, str) or not value: + raise ContentEncryptionInvalidData(f"encrypted job output is missing {field_name}") + return _b64url_decode(value) + + +def _b64url_encode(data: bytes) -> str: + return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=") + + +def _b64url_decode(data: str) -> bytes: + padding = "=" * (-len(data) % 4) + return base64.b64decode((data + padding).encode("ascii"), altchars=b"-_", validate=True) + + +def _strict_base64_decode(value: str, *, field_name: str) -> bytes: + if not isinstance(value, str) or not value: + raise ContentEncryptionConfigError(f"{field_name} must be base64 encoded") + try: + return _b64url_decode(value.strip()) + except (ValueError, UnicodeEncodeError) as exc: + raise ContentEncryptionConfigError(f"{field_name} must be base64 encoded") from exc + + +def _decode_vault_key(value: str, *, field_name: str) -> bytes: + if not isinstance(value, str) or not value: + raise ContentEncryptionUnavailable(f"{field_name} is missing") + try: + return _b64url_decode(value.strip()) + except (ValueError, UnicodeEncodeError) as exc: + raise ContentEncryptionUnavailable(f"{field_name} is not base64 encoded") from exc + + +def _parse_non_negative_float(value: str | None, *, default: float, name: str) -> float: + if value is None or value == "": + return default + try: + parsed = float(value) + except ValueError as exc: + raise ContentEncryptionConfigError(f"{name} must be a non-negative number") from exc + if not math.isfinite(parsed) or parsed < 0: + raise ContentEncryptionConfigError(f"{name} must be a non-negative number") + return parsed + + +def _parse_positive_float(value: str | None, *, default: float, name: str) -> float: + if value is None or value == "": + return default + try: + parsed = float(value) + except ValueError as exc: + raise ContentEncryptionConfigError(f"{name} must be a positive number") from exc + if not math.isfinite(parsed) or parsed <= 0: + raise ContentEncryptionConfigError(f"{name} must be a positive number") + return parsed + + +def _secret_fingerprint(value: bytes | str | None) -> str | None: + if value is None: + return None + if isinstance(value, bytes): + type_tag = b"bytes\0" + encoded = value + else: + type_tag = b"str\0" + encoded = value.encode("utf-8", errors="surrogatepass") + return hashlib.sha256(b"aiq-content-encryption-config\0" + type_tag + encoded).hexdigest() + + +def _empty_to_none(value: str | None) -> str | None: + if value is None: + return None + stripped = value.strip() + return stripped or None + + +def _is_auth_failure(exc: Exception) -> bool: + candidate = _vault_failure_candidate(exc) + return candidate.__class__.__name__ in _VAULT_AUTH_RETRY_EXCEPTION_NAMES + + +def _should_retry_vault_failure(exc: Exception) -> bool: + candidate = _vault_failure_candidate(exc) + if _is_auth_failure(candidate): + return True + if _is_retryable_network_failure(candidate): + return True + + status_code = _status_code_from_exception(candidate) + if status_code is not None: + return status_code == 429 or 500 <= status_code < 600 + + exception_name = candidate.__class__.__name__ + if exception_name in _VAULT_NON_RETRYABLE_EXCEPTION_NAMES: + return False + return exception_name in _VAULT_RETRYABLE_EXCEPTION_NAMES + + +def _vault_failure_candidate(exc: Exception) -> Exception: + if isinstance(exc, ContentEncryptionUnavailable) and isinstance(exc.__cause__, Exception): + return exc.__cause__ + return exc + + +def _is_retryable_network_failure(exc: Exception) -> bool: + try: + import requests + except ImportError: + requests = None + if requests is not None and isinstance( + exc, + ( + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + ), + ): + return True + + try: + import urllib3 + except ImportError: + urllib3 = None + if urllib3 is not None and isinstance( + exc, + ( + urllib3.exceptions.ConnectTimeoutError, + urllib3.exceptions.MaxRetryError, + urllib3.exceptions.NewConnectionError, + urllib3.exceptions.ReadTimeoutError, + urllib3.exceptions.TimeoutError, + ), + ): + return True + return False + + +def _status_code_from_exception(exc: Exception) -> int | None: + status_code = getattr(exc, "status_code", None) + if isinstance(status_code, int): + return status_code + + response = getattr(exc, "response", None) + response_status_code = getattr(response, "status_code", None) + if isinstance(response_status_code, int): + return response_status_code + return None + + +def _sleep_before_vault_retry(attempt: int) -> None: + delay = _VAULT_RETRY_BASE_SECONDS * (2**attempt) + jitter = random.uniform(0, _VAULT_RETRY_JITTER_SECONDS) + time.sleep(delay + jitter) + + +def _zero_bytes(_value: bytes) -> None: + # Python bytes are immutable, so this is only a lifetime boundary marker. + return None diff --git a/frontends/aiq_api/src/aiq_api/jobs/event_store.py b/frontends/aiq_api/src/aiq_api/jobs/event_store.py index de6174856..41e1f346e 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/event_store.py +++ b/frontends/aiq_api/src/aiq_api/jobs/event_store.py @@ -22,11 +22,14 @@ from __future__ import annotations import asyncio +import copy import logging import threading import time from typing import Any +from .crypto import ContentEncryptionError + logger = logging.getLogger(__name__) @@ -101,13 +104,70 @@ class EventStore: _cache_lock = threading.Lock() _tables_initialized: set[str] = set() - def __init__(self, db_url: str = "sqlite+aiosqlite:///./jobs.db", job_id: str | None = None): + def __init__( + self, + db_url: str = "sqlite+aiosqlite:///./jobs.db", + job_id: str | None = None, + content_cipher: Any | None = None, + ): self.db_url = db_url self.job_id = job_id + self._content_cipher = content_cipher + self._fail_closed = bool(content_cipher is not None and content_cipher.manager.config.encrypted) self._is_postgres = db_url.startswith("postgresql") self._sync_engine = self._get_or_create_sync_engine(db_url) self._ensure_table_sync() + def _prepare_event_for_storage(self, event: dict) -> dict: + """Encrypt sensitive event fields before persistence when content encryption is enabled.""" + + if self.job_id is None or self._content_cipher is None: + return event + if event.get("type") != "artifact.update": + return event + data = event.get("data") + if not isinstance(data, dict) or data.get("type") not in {"output", "file"} or "content" not in data: + return event + + from .crypto import encrypt_event_field + from .crypto import is_encrypted_event_field + + if is_encrypted_event_field(data["content"]): + return event + + encrypted_event = copy.deepcopy(event) + encrypted_event["data"]["content"] = encrypt_event_field( + self.job_id, + "artifact.update.data.content", + encrypted_event["data"].get("content"), + self._content_cipher, + ) + return encrypted_event + + @classmethod + def _prepare_event_for_return(cls, job_id: str, event: dict) -> dict: + """Decrypt sensitive event fields before returning events to API/SSE callers.""" + + if event.get("type") != "artifact.update": + return event + data = event.get("data") + if not isinstance(data, dict) or data.get("type") not in {"output", "file"} or "content" not in data: + return event + + from .crypto import decrypt_event_field + from .crypto import is_encrypted_event_field + + if not is_encrypted_event_field(data["content"]): + return event + + decrypted_event = copy.deepcopy(event) + decrypted_event["data"]["content"] = decrypt_event_field( + job_id, + "artifact.update.data.content", + decrypted_event["data"].get("content"), + ) + return decrypted_event + @classmethod def _get_or_create_sync_engine(cls, db_url: str): """Get or create a sync SQLAlchemy engine with TTL-based caching.""" @@ -343,7 +403,8 @@ def store(self, event: dict): from sqlalchemy import text event_type = event.get("type", "unknown") - event_json = json.dumps(event) + stored_event = self._prepare_event_for_storage(event) + event_json = json.dumps(stored_event) try: with self._sync_engine.connect() as conn: @@ -381,6 +442,8 @@ def store(self, event: dict): logger.debug("Stored event %s for job %s", event_type, self.job_id) except Exception as e: logger.warning("Failed to store event %s for job %s: %s", event_type, self.job_id, e) + if self._fail_closed: + raise def store_batch(self, events: list[dict]): """ @@ -398,11 +461,12 @@ def store_batch(self, events: list[dict]): rows = [] for event in events: + stored_event = self._prepare_event_for_storage(event) rows.append( { "job_id": self.job_id, - "event_type": event.get("type", "unknown"), - "event_data": json.dumps(event), + "event_type": stored_event.get("type", "unknown"), + "event_data": json.dumps(stored_event), } ) @@ -437,6 +501,8 @@ def store_batch(self, events: list[dict]): logger.debug("Stored batch of %d events for job %s", len(events), self.job_id) except Exception as e: logger.warning("Failed to store batch of %d events for job %s: %s", len(events), self.job_id, e) + if self._fail_closed: + raise @classmethod def _ensure_table_exists(cls, db_url: str): @@ -475,6 +541,30 @@ def _ensure_table_exists(cls, db_url: str): cls._tables_initialized.add(db_url) + @classmethod + def _prepare_event_rows_for_return(cls, job_id: str, rows: list[Any]) -> list[dict]: + """Deserialize and decrypt event rows in sync code, suitable for thread offload.""" + + import json + + events = [] + for row in rows: + try: + event = cls._prepare_event_for_return(job_id, json.loads(row[1])) + event["_id"] = row[0] + events.append(event) + except json.JSONDecodeError: + logger.warning("Malformed event data for job %s, event id %d", job_id, row[0]) + except ContentEncryptionError as e: + logger.warning( + "Encrypted event data failed for job %s, event id %d: %s", + job_id, + row[0], + e.__class__.__name__, + ) + raise + return events + @classmethod def get_events(cls, db_url: str, job_id: str, after_id: int = 0, limit: int = 100) -> list[dict]: """ @@ -489,8 +579,6 @@ def get_events(cls, db_url: str, job_id: str, after_id: int = 0, limit: int = 10 Returns: List of event dicts with '_id' field for cursor tracking """ - import json - from sqlalchemy import text try: @@ -505,15 +593,9 @@ def get_events(cls, db_url: str, job_id: str, after_id: int = 0, limit: int = 10 ), {"job_id": job_id, "after_id": after_id, "limit": limit}, ) - events = [] - for row in result: - try: - event = json.loads(row[1]) - event["_id"] = row[0] - events.append(event) - except json.JSONDecodeError: - logger.warning("Malformed event data for job %s, event id %d", job_id, row[0]) - return events + return cls._prepare_event_rows_for_return(job_id, list(result)) + except ContentEncryptionError: + raise except Exception as e: logger.warning("Failed to get events for job %s: %s", job_id, e) return [] @@ -525,8 +607,6 @@ async def get_events_async(cls, db_url: str, job_id: str, after_id: int = 0, lim Uses native async SQLAlchemy with psycopg for true async I/O. """ - import json - from sqlalchemy import text try: @@ -541,15 +621,9 @@ async def get_events_async(cls, db_url: str, job_id: str, after_id: int = 0, lim ), {"job_id": job_id, "after_id": after_id, "limit": limit}, ) - events = [] - for row in result: - try: - event = json.loads(row[1]) - event["_id"] = row[0] - events.append(event) - except json.JSONDecodeError: - logger.warning("Malformed event data for job %s, event id %d", job_id, row[0]) - return events + return await asyncio.to_thread(cls._prepare_event_rows_for_return, job_id, list(result)) + except ContentEncryptionError: + raise except Exception as e: logger.warning("Failed to get events async for job %s: %s", job_id, e) return await asyncio.get_running_loop().run_in_executor( @@ -580,18 +654,28 @@ def get_event_by_id(cls, db_url: str, event_id: int) -> dict | None: engine = cls._get_or_create_sync_engine(db_url) with engine.connect() as conn: result = conn.execute( - text("SELECT id, event_data FROM job_events WHERE id = :event_id"), + text("SELECT id, job_id, event_data FROM job_events WHERE id = :event_id"), {"event_id": event_id}, ) row = result.fetchone() if row: try: - event = json.loads(row[1]) + event = cls._prepare_event_for_return(row[1], json.loads(row[2])) event["_id"] = row[0] return event except json.JSONDecodeError: logger.warning("Malformed event data for event id %d", event_id) + except ContentEncryptionError as e: + logger.warning( + "Encrypted event data failed for job %s, event id %d: %s", + row[1], + event_id, + e.__class__.__name__, + ) + raise return None + except ContentEncryptionError: + raise except Exception as e: logger.warning("Failed to get event %d: %s", event_id, e) return None @@ -731,6 +815,7 @@ def __init__(self, event_store: EventStore): self._buffer: list[dict] = [] self._lock = threading.Lock() self._timer: threading.Timer | None = None + self._flush_error: Exception | None = None @property def job_id(self) -> str | None: @@ -740,30 +825,50 @@ def job_id(self) -> str | None: def store(self, event: dict): """Buffer an event; flush when batch is full or timer fires.""" with self._lock: + self._raise_flush_error_locked() self._buffer.append(event) if len(self._buffer) >= self.MAX_BATCH_SIZE: self._flush_locked() elif self._timer is None: - self._timer = threading.Timer(self.FLUSH_INTERVAL_MS / 1000, self._flush) + self._timer = threading.Timer(self.FLUSH_INTERVAL_MS / 1000, self._flush_from_timer) self._timer.daemon = True self._timer.start() + def _raise_flush_error_locked(self): + """Raise a previously captured background flush failure.""" + if self._flush_error is not None: + raise self._flush_error + def _flush_locked(self): """Flush while already holding the lock.""" if self._timer: self._timer.cancel() self._timer = None + self._raise_flush_error_locked() if not self._buffer: return batch = self._buffer[:] self._buffer.clear() - self._store.store_batch(batch) - - def _flush(self): - """Flush from timer callback (acquires lock).""" + try: + self._store.store_batch(batch) + except Exception as exc: + if self._flush_error is None: + self._flush_error = exc + raise + + def _flush_from_timer(self): + """Capture timer-thread failures for the owning job to observe.""" with self._lock: - self._flush_locked() + try: + self._flush_locked() + except Exception as exc: + logger.warning( + "Background event flush failed for job %s exception=%s", + self.job_id, + exc.__class__.__name__, + ) def flush(self): """Force flush all buffered events. Call before job completion.""" - self._flush() + with self._lock: + self._flush_locked() diff --git a/frontends/aiq_api/src/aiq_api/jobs/report_context.py b/frontends/aiq_api/src/aiq_api/jobs/report_context.py index 933415c86..409e50e20 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/report_context.py +++ b/frontends/aiq_api/src/aiq_api/jobs/report_context.py @@ -17,6 +17,7 @@ from aiq_agent.auth import Principal from .access import authorize_job_access +from .crypto import read_job_output_async from .event_store import EventStore _EVENT_SCAN_LIMIT = 10000 @@ -58,8 +59,8 @@ def _decode_job_output(output: Any) -> dict[str, Any]: return {} -def _extract_report_from_job_output(job: Any) -> str | None: - output = _decode_job_output(getattr(job, "output", None)) +def _extract_report_from_output(output: Any) -> str | None: + output = _decode_job_output(output) report = output.get("report") return report.strip() if isinstance(report, str) and report.strip() else None @@ -183,7 +184,8 @@ def _source_summary_markdown(sources: list[ReportContextSource]) -> str: async def resolve_report_context(job: Any, db_url: str, parent_job_id: str) -> ReportContext: """Build report context from a previously authorized parent job.""" - report = _extract_report_from_job_output(job) + output = await read_job_output_async(parent_job_id, getattr(job, "output", None)) + report = _extract_report_from_output(output) # Durable events are only needed when the report isn't already in job output, # or to reconstruct sources. Fetch the (potentially large) event log at most once. events: list[dict[str, Any]] | None = None diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index 40c00eb3d..6a73e0fa3 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -31,12 +31,16 @@ import uuid from collections.abc import Awaitable from collections.abc import Callable +from typing import TYPE_CHECKING from typing import Any from .callbacks import AgentEventCallback from .event_store import BatchingEventStore from .event_store import EventStore +if TYPE_CHECKING: + from .crypto import ContentEncryptionPolicyIdentity + logger = logging.getLogger(__name__) _DEEP_RESEARCH_FUNCTION_TYPE = "deep_research_agent" @@ -447,6 +451,7 @@ async def run_agent_job( available_documents: list[dict] | None = None, data_sources: list[str] | None = None, auth_token: str | None = None, + content_encryption_policy: ContentEncryptionPolicyIdentity | None = None, initial_files: dict[str, Any] | None = None, output_metadata: dict[str, Any] | None = None, owner_user_id: str | None = None, @@ -481,6 +486,8 @@ async def run_agent_job( data_sources: Optional list of allowed data sources to enforce in the worker. auth_token: Optional auth token propagated from the HTTP request for data sources that require authentication (requires_auth: true). + content_encryption_policy: Non-secret policy identity captured by the + submitting API process and required to match the worker configuration. 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 @@ -521,6 +528,7 @@ async def run_agent_job( std_logging.basicConfig(level=log_level) job_store: JobStore | None = None + job_output_cipher = None cancellation_monitor: CancellationMonitor | None = None event_store: EventStore | BatchingEventStore | None = None # Sandbox runtime is released on the terminal path; interrupted forces terminate() over close(). @@ -535,6 +543,28 @@ async def run_agent_job( try: job_store = JobStore(scheduler_address=scheduler_address, db_url=db_url) + try: + from .crypto import ContentEncryptionError + from .crypto import ContentEncryptionPolicyMismatch + from .crypto import create_job_content_cipher + from .crypto import require_content_encryption_policy + + require_content_encryption_policy(content_encryption_policy) + job_output_cipher = create_job_content_cipher(job_id) + except ContentEncryptionError as exc: + logger.warning( + "Job %s failed encryption policy/readiness before running exception=%s", + job_id, + exc.__class__.__name__, + ) + error = ( + "content encryption policy mismatch" + if isinstance(exc, ContentEncryptionPolicyMismatch) + else "content encryption unavailable" + ) + await job_store.update_status(job_id, JobStatus.FAILURE, error=error) + return + await job_store.update_status(job_id, JobStatus.RUNNING) cancellation_monitor = CancellationMonitor( @@ -687,7 +717,7 @@ async def run_agent_job( verbose = is_verbose(getattr(fn_config, "verbose", False)) callbacks = [VerboseTraceCallback()] if verbose else [] - raw_event_store = EventStore(db_url, job_id) + raw_event_store = EventStore(db_url, job_id, content_cipher=job_output_cipher) event_store = BatchingEventStore(raw_event_store) callbacks.append(AgentEventCallback(event_store)) callbacks.append(nat_profiler_callback) @@ -775,10 +805,28 @@ async def run_agent_job( # Extract report and update status inside the context manager # so the UI sees completion before exporter flush and cleanup report = _extract_result(result) + from .crypto import update_job_output + + if job_output_cipher is None: + raise RuntimeError("job output cipher was not initialized") # Apply caller metadata first, then set the canonical report last so a # stray "report" key in output_metadata can never overwrite the real report. output = {**(output_metadata or {}), "report": report} - await job_store.update_status(job_id, JobStatus.SUCCESS, output=output) + try: + await update_job_output( + job_store, + job_id, + JobStatus.SUCCESS, + output=output, + cipher=job_output_cipher, + ) + except Exception as exc: + logger.warning( + "Job %s encrypted output write failed exception=%s", + job_id, + exc.__class__.__name__, + ) + raise logger.info("Job %s completed (report: %d chars)", job_id, len(report)) except asyncio.CancelledError: @@ -795,14 +843,13 @@ async def run_agent_job( if event_store is None: event_store = BatchingEventStore(EventStore(db_url, job_id)) - event_store.store( + _store_terminal_event_best_effort( + event_store, { "type": "job.cancelled", "data": {"reason": "cancelled by user"}, - } + }, ) - if hasattr(event_store, "flush"): - event_store.flush() except Exception as e: logger.exception("Job %s failed: %s", job_id, type(e).__name__) @@ -812,22 +859,28 @@ async def run_agent_job( if event_store is None: event_store = BatchingEventStore(EventStore(db_url, job_id)) - event_store.store( + _store_terminal_event_best_effort( + event_store, { "type": "job.error", "data": { "error": str(e), "error_type": type(e).__name__, }, - } + }, ) - if hasattr(event_store, "flush"): - event_store.flush() finally: # Ensure terminal-path events are not left in the batch buffer. if event_store is not None and hasattr(event_store, "flush"): - event_store.flush() + try: + event_store.flush() + except Exception as exc: + logger.warning( + "Final event flush failed for job %s exception=%s", + job_id, + exc.__class__.__name__, + ) if cancellation_monitor: cancellation_monitor.stop() # Release the sandbox off the event loop so the SDK session close never blocks the Dask @@ -841,6 +894,21 @@ async def run_agent_job( job_auth_token.reset(_auth_token_reset) +def _store_terminal_event_best_effort(event_store, event: dict) -> None: + """Persist a terminal event without masking the job's terminal status.""" + try: + event_store.store(event) + if hasattr(event_store, "flush"): + event_store.flush() + except Exception as exc: + logger.warning( + "Failed to persist terminal event %s for job %s exception=%s", + event.get("type", "unknown"), + event_store.job_id, + exc.__class__.__name__, + ) + + def _teardown_sandbox(sandbox_runtime: Any | None, *, job_id: str, interrupted: bool) -> None: """Release sandbox resources on a terminal path (best-effort, never raises). @@ -1148,6 +1216,8 @@ async def run_deep_research( Preserved for backwards compatibility. New code should use run_agent_job directly. """ + from .crypto import get_content_encryption_policy_identity + await run_agent_job( configure_logging=configure_logging, log_level=log_level, @@ -1158,4 +1228,5 @@ async def run_deep_research( input_text=input_text, agent_class_path="aiq_agent.agents.deep_researcher.agent.DeepResearcherAgent", agent_config_name="deep_research_agent", + content_encryption_policy=get_content_encryption_policy_identity(), ) diff --git a/frontends/aiq_api/src/aiq_api/jobs/submit.py b/frontends/aiq_api/src/aiq_api/jobs/submit.py index 7f93e5bda..5dbadf3d8 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/submit.py +++ b/frontends/aiq_api/src/aiq_api/jobs/submit.py @@ -152,6 +152,7 @@ async def submit_agent_job( available_documents: list[dict] | None = None, data_sources: list[str] | None = None, auth_token: str | None = None, + skip_encryption_readiness_check: bool = False, initial_files: dict[str, Any] | None = None, output_metadata: dict[str, Any] | None = None, allow_internal: bool = False, @@ -173,8 +174,11 @@ async def submit_agent_job( data_sources: Optional list of allowed data sources to enforce in the worker. auth_token: Optional auth token to propagate to the Dask worker for data sources that require authentication. + skip_encryption_readiness_check: Skip the internal readiness check when + the caller already performed it off the event loop. initial_files: Optional DeepAgents virtual filesystem files to seed into worker state. output_metadata: Optional metadata persisted with the final job output. + allow_internal: Allow trusted callers to submit an internal-only agent. Returns: The job ID. @@ -193,6 +197,9 @@ async def submit_agent_job( """ from nat.front_ends.fastapi.async_jobs.job_store import JobStore + from .crypto import get_content_encryption_policy_identity + from .crypto import require_content_encryption_ready_for_submission_async + # Get agent configuration from registry agent_config = get_agent_config(agent_type) @@ -242,6 +249,10 @@ async def submit_agent_job( if not scheduler_address: raise RuntimeError("Async job submission requires NAT_DASK_SCHEDULER_ADDRESS to be set") + if not skip_encryption_readiness_check: + await require_content_encryption_ready_for_submission_async() + content_encryption_policy = get_content_encryption_policy_identity() + # Auto-capture auth token if not explicitly provided if auth_token is None: from aiq_agent.auth import get_auth_token @@ -307,6 +318,7 @@ async def _rollback_partial_submission() -> None: available_documents, data_sources, auth_token, + content_encryption_policy, initial_files, output_metadata, principal_user_id(principal), diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index 2d0ac8641..4263d20f6 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -42,6 +42,7 @@ from fastapi import FastAPI from fastapi import HTTPException from fastapi.responses import StreamingResponse +from fastapi.routing import APIRoute from pydantic import BaseModel from pydantic import ConfigDict from pydantic import Field @@ -64,6 +65,21 @@ logger = logging.getLogger(__name__) +def _remove_existing_health_routes(app: FastAPI) -> int: + """Remove existing GET /health routes before installing AI-Q readiness.""" + existing_routes = [ + route + for route in app.router.routes + if isinstance(route, APIRoute) and route.path == "/health" and "GET" in route.methods + ] + for route in existing_routes: + app.router.routes.remove(route) + if existing_routes: + logger.info("Replacing %d existing GET /health route(s) with AI-Q readiness", len(existing_routes)) + app.openapi_schema = None + return len(existing_routes) + + def _validate_artifact_store(db_url: str) -> None: """Validate configured artifact storage during API startup.""" from aiq_agent.agents.deep_researcher.sandbox.artifacts import build_artifact_store @@ -475,8 +491,14 @@ async def register_job_routes(app: FastAPI, builder: WorkflowBuilder, worker: Fa from ..jobs.access import authorize_job_access from ..jobs.access import ensure_job_access_table + from ..jobs.crypto import ContentEncryptionConfigError + from ..jobs.crypto import ContentEncryptionInvalidData + from ..jobs.crypto import ContentEncryptionUnavailable + from ..jobs.crypto import get_content_encryption_health_async + from ..jobs.crypto import read_job_output_async + from ..jobs.crypto import require_content_encryption_ready_for_submission_async + from ..jobs.crypto import validate_content_encryption_startup_async from ..jobs.event_store import EventStore - from ..jobs.report_context import _decode_job_output from ..jobs.report_context import report_output_metadata from ..jobs.report_context import resolve_report_context from ..jobs.report_context import to_initial_files @@ -498,6 +520,91 @@ async def register_job_routes(app: FastAPI, builder: WorkflowBuilder, worker: Fa set_active_mcp_auth_provider(mcp_auth_provider) register_mcp_auth_routes(app, mcp_auth_provider) + await validate_content_encryption_startup_async() + + @app.get( + "/live", + tags=["health"], + summary="Liveness check", + description="Returns success while the API process is running; does not check external dependencies.", + ) + async def liveness_check() -> dict[str, str]: + """Report process liveness without coupling restarts to dependency health.""" + + return {"status": "alive"} + + dask_available = getattr(worker, "_dask_available", False) + job_store = getattr(worker, "_job_store", None) + scheduler_address = getattr(worker, "_scheduler_address", None) or os.environ.get("NAT_DASK_SCHEDULER_ADDRESS") + db_url = getattr(worker, "_db_url", None) or os.environ.get("NAT_JOB_STORE_DB_URL", "sqlite:///./data/jobs.db") + config_path = getattr(worker, "_config_file_path", None) or os.environ.get("NAT_CONFIG_FILE", "") + log_level = getattr(worker, "_log_level", std_logging.INFO) + use_threads = getattr(worker, "_use_dask_threads", False) + front_end_config = getattr(worker, "_front_end_config", None) + default_expiry_seconds = getattr(front_end_config, "expiry_seconds", 86400) if front_end_config else 86400 + + # NAT registers its generic /health route first. Replace it before any + # async-job prerequisite early return so /health always means readiness. + _remove_existing_health_routes(app) + + @app.get( + "/health", + tags=["health"], + summary="Readiness check", + responses={503: {"description": "Async-job, database, or content-encryption dependency is unavailable"}}, + ) + async def health_check(): + """Readiness endpoint that validates async-job, DB, and encryption dependencies.""" + from fastapi.responses import JSONResponse + from sqlalchemy import text + + from ..jobs.event_store import EventStore + + result = {"status": "healthy", "dask_available": bool(dask_available), "db": "ok"} + if not dask_available or not job_store: + result["status"] = "degraded" + result["db"] = "unchecked" + result["reason"] = "async_jobs_unavailable" + return JSONResponse(status_code=503, content=result) + if not config_path: + result["status"] = "degraded" + result["db"] = "unchecked" + result["reason"] = "configuration_missing" + return JSONResponse(status_code=503, content=result) + + # Check DB connectivity by obtaining (or creating) the engine for the + # configured db_url and running a bounded ping. An empty async-engine + # cache is the normal fresh-process state and must never be treated as + # healthy: readiness must reflect the actual database. + try: + engine = EventStore._get_or_create_async_engine(db_url) + async with engine.connect() as conn: + await asyncio.wait_for(conn.execute(text("SELECT 1")), timeout=3.0) + except Exception: + logger.warning("Health check DB ping failed", exc_info=True) + result["status"] = "degraded" + result["db"] = "unreachable" + return JSONResponse(status_code=503, content=result) + + try: + encryption = await get_content_encryption_health_async() + result["encryption"] = encryption.to_health_dict() + if encryption.mode != "off" and not encryption.ready: + result["status"] = "degraded" + return JSONResponse(status_code=503, content=result) + except ContentEncryptionConfigError as exc: + logger.warning("Health check encryption config failed exception=%s", exc.__class__.__name__) + result["status"] = "degraded" + result["encryption"] = { + "mode": "invalid", + "ready": False, + "reason": "configuration_invalid", + "exception_type": exc.__class__.__name__, + } + return JSONResponse(status_code=503, content=result) + + return result + if not get_all_sources(): logger.warning( "No data sources registered. Add a 'data_sources' function with " @@ -547,9 +654,6 @@ async def list_data_sources() -> list[DataSource]: logger.info("Registered /v1/data_sources and /v1/jobs/async/agents routes") - dask_available = getattr(worker, "_dask_available", False) - job_store = getattr(worker, "_job_store", None) - if not dask_available or not job_store: logger.warning( "Dask not available - async job submission routes require NAT_DASK_SCHEDULER_ADDRESS" @@ -557,19 +661,10 @@ async def list_data_sources() -> list[DataSource]: ) return - scheduler_address = getattr(worker, "_scheduler_address", None) or os.environ.get("NAT_DASK_SCHEDULER_ADDRESS") - db_url = getattr(worker, "_db_url", None) or os.environ.get("NAT_JOB_STORE_DB_URL", "sqlite:///./data/jobs.db") - config_path = getattr(worker, "_config_file_path", None) or os.environ.get("NAT_CONFIG_FILE", "") - log_level = getattr(worker, "_log_level", std_logging.INFO) - use_threads = getattr(worker, "_use_dask_threads", False) - if not config_path: logger.error("Config file path not available - NAT_CONFIG_FILE not set") return - front_end_config = getattr(worker, "_front_end_config", None) - default_expiry_seconds = getattr(front_end_config, "expiry_seconds", 86400) if front_end_config else 86400 - logger.info( "Registering async job routes: scheduler=%s, db=%s, expiry=%ds", scheduler_address, @@ -580,34 +675,6 @@ async def list_data_sources() -> list[DataSource]: await loop.run_in_executor(None, ensure_job_access_table, db_url) await loop.run_in_executor(None, _validate_artifact_store, db_url) - @app.get("/health", tags=["health"], summary="Health check") - async def health_check(): - """Health check endpoint that validates DB connectivity.""" - from sqlalchemy import text - - from ..jobs.event_store import EventStore - - result = {"status": "ok", "dask_available": dask_available, "db": "ok"} - - # Check DB connectivity using any cached async engine - try: - cache = EventStore._async_engine_cache - if cache: - engine = next(iter(cache.values()))[0] - async with engine.connect() as conn: - await asyncio.wait_for(conn.execute(text("SELECT 1")), timeout=3.0) - else: - result["db"] = "no_engine" - except Exception: - logger.warning("Health check DB ping failed", exc_info=True) - result["status"] = "degraded" - result["db"] = "unreachable" - from fastapi.responses import JSONResponse - - return JSONResponse(status_code=503, content=result) - - return result - @app.post( "/v1/jobs/async/submit", response_model=JobStatusResponse, @@ -625,8 +692,12 @@ async def health_check(): ) }, 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"}, + 500: { + "description": ( + "Content encryption configuration is invalid or async job authorization persistence failed" + ) + }, + 503: {"description": "Content encryption, Dask scheduler, or sandbox capacity is unavailable"}, }, ) async def submit_job( @@ -644,6 +715,21 @@ async def submit_job( # Authenticate the caller (raises 401/403 if unverified). The returned principal # is also forwarded to submit_authorized_job(...) below for ownership recording. principal = require_verified_principal() + try: + await require_content_encryption_ready_for_submission_async() + except ContentEncryptionUnavailable as e: + logger.warning( + "Rejected async job submission because content encryption is unready: %s", + e.__class__.__name__, + ) + raise HTTPException(503, "Content encryption is not ready") + except ContentEncryptionConfigError as e: + logger.warning( + "Rejected async job submission because content encryption config is invalid: %s", + e.__class__.__name__, + ) + raise HTTPException(500, "Content encryption configuration is invalid") + validation_start = time.perf_counter() await _validate_data_sources_for_agent( builder=builder, @@ -686,7 +772,20 @@ async def submit_job( expiry_seconds=expiry, data_sources=req.data_sources, auth_token=auth_token, + skip_encryption_readiness_check=True, + ) + except ContentEncryptionUnavailable as e: + logger.warning( + "Failed to submit authorized job because content encryption is unready: %s", + e.__class__.__name__, + ) + raise HTTPException(503, "Content encryption is not ready") + except ContentEncryptionConfigError as e: + logger.warning( + "Failed to submit authorized job because content encryption config is invalid: %s", + e.__class__.__name__, ) + raise HTTPException(500, "Content encryption configuration is invalid") except JobIdConflictError: raise HTTPException(409, f"Job already exists: {req.job_id}") except McpAuthRequiredError as e: @@ -734,8 +833,8 @@ async def submit_job( 404: {"description": "Parent job not found"}, 409: {"description": "Parent job is incomplete, has no durable report, or the child job_id collides"}, 422: {"description": "Request validation failed (e.g. blank edit instruction)"}, - 500: {"description": "Failed to submit the report edit job"}, - 503: {"description": "Dask scheduler not available"}, + 500: {"description": "Parent report data is invalid or report edit submission failed"}, + 503: {"description": "Content encryption or Dask scheduler is unavailable"}, }, ) async def edit_job_report(job_id: str, req: ReportEditRequest) -> ReportEditResponse: @@ -745,7 +844,22 @@ async def edit_job_report(job_id: str, req: ReportEditRequest) -> ReportEditResp if getattr(parent_job, "status", None) != JobStatus.SUCCESS.value: raise HTTPException(409, f"Parent job is not complete: {job_id}") - context = await resolve_report_context(parent_job, db_url, job_id) + try: + context = await resolve_report_context(parent_job, db_url, job_id) + except ContentEncryptionUnavailable as e: + logger.warning( + "Parent report decrypt unavailable job_id=%s exception=%s", + job_id, + e.__class__.__name__, + ) + raise HTTPException(503, "Content encryption is unavailable") + except ContentEncryptionInvalidData as e: + logger.warning( + "Parent report persisted output invalid job_id=%s exception=%s", + job_id, + e.__class__.__name__, + ) + raise HTTPException(500, "Parent report data is invalid") expiry = req.expiry_seconds if req.expiry_seconds is not None else default_expiry_seconds from aiq_agent.auth import get_auth_token @@ -765,6 +879,21 @@ async def edit_job_report(job_id: str, req: ReportEditRequest) -> ReportEditResp output_metadata=report_output_metadata(job_id, "edit"), allow_internal=True, ) + except ContentEncryptionUnavailable as e: + logger.warning( + "Report edit submission rejected because content encryption is unready parent_job_id=%s exception=%s", + job_id, + e.__class__.__name__, + ) + raise HTTPException(503, "Content encryption is not ready") + except ContentEncryptionConfigError as e: + logger.warning( + "Report edit submission rejected because content encryption config is invalid " + "parent_job_id=%s exception=%s", + job_id, + e.__class__.__name__, + ) + raise HTTPException(500, "Content encryption configuration is invalid") except JobIdConflictError: raise HTTPException(409, f"Job already exists: {req.job_id}") except RuntimeError as e: @@ -890,7 +1019,22 @@ async def get_job_state(job_id: str) -> JobStateResponse: principal = require_verified_principal() await authorize_job_access(job_store, db_url, job_id, principal) - artifacts = await _get_job_artifacts(db_url, job_id) + try: + artifacts = await _get_job_artifacts(db_url, job_id) + except ContentEncryptionUnavailable as e: + logger.warning( + "Job state decrypt unavailable job_id=%s exception=%s", + job_id, + e.__class__.__name__, + ) + raise HTTPException(503, "Content encryption is unavailable") + except ContentEncryptionInvalidData as e: + logger.warning( + "Job state persisted event data invalid job_id=%s exception=%s", + job_id, + e.__class__.__name__, + ) + raise HTTPException(500, "Job state data is invalid") return JobStateResponse( job_id=job_id, has_state=artifacts is not None, @@ -980,7 +1124,26 @@ async def get_job_report(job_id: str) -> JobReportResponse: principal = require_verified_principal() job = await authorize_job_access(job_store, db_url, job_id, principal) - output = _decode_job_output(job.output) + output: dict[str, Any] = {} + if job.output is not None: + try: + decoded_output = await read_job_output_async(job_id, job.output) + except ContentEncryptionUnavailable as e: + logger.warning( + "Final report decrypt unavailable job_id=%s exception=%s", + job_id, + e.__class__.__name__, + ) + raise HTTPException(503, "Content encryption is unavailable") + except ContentEncryptionInvalidData as e: + logger.warning( + "Final report persisted output invalid job_id=%s exception=%s", + job_id, + e.__class__.__name__, + ) + raise HTTPException(500, "Final report data is invalid") + if isinstance(decoded_output, dict): + output = decoded_output report = output.get("report") return JobReportResponse( @@ -1451,6 +1614,7 @@ async def _get_job_artifacts(db_url: str, job_id: str) -> dict | None: Returns: Dict with 'tools', 'outputs', and 'sources' (counts), or None if no artifacts found. """ + from ..jobs.crypto import ContentEncryptionError from ..jobs.event_store import EventStore try: @@ -1487,6 +1651,8 @@ async def _get_job_artifacts(db_url: str, job_id: str) -> dict | None: } return result if tools or outputs or sources_found else None + except ContentEncryptionError: + raise except (KeyError, TypeError) as e: logger.warning("Failed to parse artifacts for job %s: %s", job_id, e) return None @@ -1502,12 +1668,20 @@ async def _sse_generator(job_store, job_id: str, db_url: str, start_event_id: in PostgreSQL: Uses LISTEN/NOTIFY for real-time push-based events (sub-10ms latency). SQLite: Uses polling (0.5s interval) since SQLite doesn't support pub-sub. """ + from ..jobs.crypto import ContentEncryptionInvalidData + from ..jobs.crypto import ContentEncryptionUnavailable from ..jobs.event_store import EventStore if EventStore.is_postgres(db_url): try: async for event in _sse_generator_postgres(job_store, job_id, db_url, start_event_id): yield event + except ContentEncryptionUnavailable as e: + logger.warning("SSE encrypted event decrypt unavailable for job %s: %s", job_id, e.__class__.__name__) + yield f"event: job.error\ndata: {json.dumps({'error': 'Content encryption is unavailable'})}\n\n" + except ContentEncryptionInvalidData as e: + logger.warning("SSE encrypted event data invalid for job %s: %s", job_id, e.__class__.__name__) + yield f"event: job.error\ndata: {json.dumps({'error': 'Job event data is invalid'})}\n\n" except Exception as e: logger.warning("Pub-sub failed, falling back to polling: %s", e) async for event in _sse_generator_polling(job_store, job_id, db_url, start_event_id): @@ -1532,6 +1706,8 @@ async def _sse_generator_postgres(job_store, job_id: str, db_url: str, start_eve from nat.front_ends.fastapi.async_jobs.job_store import JobStatus from ..jobs.connection_manager import get_connection_manager + from ..jobs.crypto import ContentEncryptionInvalidData + from ..jobs.crypto import ContentEncryptionUnavailable from ..jobs.event_store import EventStore connection_manager = get_connection_manager() @@ -1702,6 +1878,22 @@ def notification_handler(connection, pid, channel_name, payload): except asyncio.CancelledError: logger.info("SSE pub-sub stream cancelled for job %s", job_id) break + except ContentEncryptionUnavailable as e: + logger.warning( + "SSE pub-sub encrypted event decrypt unavailable for job %s: %s", + job_id, + e.__class__.__name__, + ) + yield format_sse("job.error", {"error": "Content encryption is unavailable"}) + break + except ContentEncryptionInvalidData as e: + logger.warning( + "SSE pub-sub encrypted event data invalid for job %s: %s", + job_id, + e.__class__.__name__, + ) + yield format_sse("job.error", {"error": "Job event data is invalid"}) + break except Exception as e: logger.exception("SSE pub-sub stream error for job %s: %s", job_id, e) yield format_sse("job.error", {"error": "Internal server error"}) @@ -1732,6 +1924,8 @@ async def _sse_generator_polling(job_store, job_id: str, db_url: str, start_even from nat.front_ends.fastapi.async_jobs.job_store import JobStatus from ..jobs.connection_manager import get_connection_manager + from ..jobs.crypto import ContentEncryptionInvalidData + from ..jobs.crypto import ContentEncryptionUnavailable from ..jobs.event_store import EventStore connection_manager = get_connection_manager() @@ -1857,6 +2051,14 @@ def format_sse(event_type: str, data: dict, event_id: int | None = None) -> str: except asyncio.CancelledError: logger.info("SSE stream cancelled for job %s", job_id) break + except ContentEncryptionUnavailable as e: + logger.warning("SSE encrypted event decrypt unavailable for job %s: %s", job_id, e.__class__.__name__) + yield format_sse("job.error", {"error": "Content encryption is unavailable"}) + break + except ContentEncryptionInvalidData as e: + logger.warning("SSE encrypted event data invalid for job %s: %s", job_id, e.__class__.__name__) + yield format_sse("job.error", {"error": "Job event data is invalid"}) + break except Exception as e: logger.exception("SSE stream error for job %s: %s", job_id, e) yield format_sse("job.error", {"error": "Internal server error"}) diff --git a/frontends/aiq_api/tests/test_auth.py b/frontends/aiq_api/tests/test_auth.py index 7c70e3a1f..8c1c2f9a9 100644 --- a/frontends/aiq_api/tests/test_auth.py +++ b/frontends/aiq_api/tests/test_auth.py @@ -893,6 +893,7 @@ def test_extract_token_from_cookie(self) -> None: def test_path_allowed_exact_and_prefix(self) -> None: mw = AuthMiddleware(MagicMock(), external_hostnames=set()) + assert mw._path_allowed("/live") is True assert mw._path_allowed("/health") is True assert mw._path_allowed("/v1/jobs/async/job/abc/result") is True assert mw._path_allowed("/v1/jobs/async/job") is True @@ -941,6 +942,19 @@ async def send(msg): class TestAuthMiddlewareExempt: + @pytest.mark.asyncio + async def test_liveness_exempt_without_auth(self, capture_asgi, external_host): + app, state = capture_asgi + mw = AuthMiddleware(app, require_auth=True, external_hostnames={external_host.decode()}) + + async def send(msg): + pass + + scope = _http_scope("/live", host=external_host) + await mw(scope, AsyncMock(), send) + + assert state["user"]["type"] == "anonymous" + @pytest.mark.asyncio async def test_health_exempt_without_auth(self, capture_asgi, external_host): app, state = capture_asgi diff --git a/frontends/aiq_api/tests/test_content_encryption.py b/frontends/aiq_api/tests/test_content_encryption.py new file mode 100644 index 000000000..753824e73 --- /dev/null +++ b/frontends/aiq_api/tests/test_content_encryption.py @@ -0,0 +1,871 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +from __future__ import annotations + +import asyncio +import base64 +import json +import os +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from hvac import exceptions as vault_exceptions +from requests import exceptions as requests_exceptions + +from aiq_api.jobs import crypto + + +def _static_key() -> str: + return base64.urlsafe_b64encode(bytes(range(32))).decode("ascii") + + +def _real_vault_env_present() -> bool: + return all( + os.environ.get(name) + for name in ( + "VAULT_ADDR", + "VAULT_ROLE_ID", + "VAULT_SECRET_ID", + "AIQ_ENCRYPTION_TRANSIT_KEY", + ) + ) + + +@pytest.fixture(autouse=True) +def clean_encryption_env(monkeypatch, request): + if request.node.name == "test_real_vault_transit_round_trip": + crypto.reset_content_encryption_manager_for_tests() + yield + crypto.reset_content_encryption_manager_for_tests() + return + + for name in ( + "AIQ_CONTENT_ENCRYPTION", + "AIQ_CONTENT_ENCRYPTION_KEY", + "AIQ_CONTENT_ENCRYPTION_KEY_ID", + "AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS", + "AIQ_CONTENT_ENCRYPTION_DEK_CACHE_TTL_SECONDS", + "VAULT_ADDR", + "VAULT_NAMESPACE", + "VAULT_TRANSIT_MOUNT", + "VAULT_ROLE_ID", + "VAULT_SECRET_ID", + "AIQ_ENCRYPTION_TRANSIT_KEY", + "VAULT_TIMEOUT_SECONDS", + ): + monkeypatch.delenv(name, raising=False) + crypto.reset_content_encryption_manager_for_tests() + yield + crypto.reset_content_encryption_manager_for_tests() + + +def _enable_static_key(monkeypatch, *, cache_ttl: str | None = None) -> None: + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "key") + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", _static_key()) + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY_ID", "test-key") + if cache_ttl is not None: + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_DEK_CACHE_TTL_SECONDS", cache_ttl) + crypto.reset_content_encryption_manager_for_tests() + + +def _vault_config() -> crypto.ContentEncryptionConfig: + return crypto.ContentEncryptionConfig( + mode="vault", + vault_addr="https://vault.example.com", + vault_transit_key="reports", + vault_role_id="role-id", + vault_secret_id="secret-id", + ) + + +def _vault_client_for_tests() -> crypto._VaultTransitClient: + return crypto._VaultTransitClient(_vault_config()) + + +def _disable_vault_retry_sleep(monkeypatch) -> list[float]: + sleeps = [] + monkeypatch.setattr(crypto.time, "sleep", sleeps.append) + monkeypatch.setattr(crypto.random, "uniform", lambda _start, _end: 0) + return sleeps + + +def test_content_encryption_defaults_to_off(): + config = crypto.get_content_encryption_config() + + assert config.mode == "off" + assert config.encrypted is False + + +def test_config_repr_and_signature_do_not_expose_credentials(): + static_key = b"static-key-credential" + role_id = "vault-role-credential" + secret_id = "vault-secret-credential" + config = crypto.ContentEncryptionConfig( + mode="vault", + static_key=static_key, + vault_role_id=role_id, + vault_secret_id=secret_id, + ) + + config_repr = repr(config) + signature_repr = repr(config.signature) + + assert "static_key=" not in config_repr + assert "vault_role_id=" not in config_repr + assert "vault_secret_id=" not in config_repr + for credential in (static_key.decode(), role_id, secret_id): + assert credential not in config_repr + assert credential not in signature_repr + + +@pytest.mark.parametrize( + ("credential_name", "replacement"), + [ + ("static_key", b"replacement-static-key"), + ("vault_role_id", "replacement-role-id"), + ("vault_secret_id", "replacement-secret-id"), + ], +) +def test_config_signature_changes_when_credentials_change(credential_name, replacement): + credentials = { + "static_key": b"original-static-key", + "vault_role_id": "original-role-id", + "vault_secret_id": "original-secret-id", + } + original = crypto.ContentEncryptionConfig(mode="vault", **credentials) + credentials[credential_name] = replacement + updated = crypto.ContentEncryptionConfig(mode="vault", **credentials) + + assert original.signature != updated.signature + + +def test_static_key_policy_identity_is_non_secret_and_key_specific(): + static_key = b"a" * crypto.DEK_BYTES + replacement_key = b"b" * crypto.DEK_BYTES + original = crypto.ContentEncryptionConfig(mode="key", key_id="reports", static_key=static_key).policy_identity + replacement = crypto.ContentEncryptionConfig( + mode="key", + key_id="reports", + static_key=replacement_key, + ).policy_identity + + assert original.mode == "key" + assert original.key_id == "reports" + assert original.static_key_fingerprint == crypto._secret_fingerprint(static_key) + assert original != replacement + assert static_key.decode() not in repr(original) + + +def test_vault_policy_identity_changes_with_transit_key_location(): + original = _vault_config().policy_identity + same_key_with_different_credentials = crypto.ContentEncryptionConfig( + mode="vault", + vault_addr="https://vault.example.com", + vault_transit_key="reports", + vault_role_id="different-role-id", + vault_secret_id="different-secret-id", + ).policy_identity + replacement = crypto.ContentEncryptionConfig( + mode="vault", + vault_addr="https://other-vault.example.com", + vault_transit_key="reports", + ).policy_identity + + assert original == same_key_with_different_credentials + assert original != replacement + assert not hasattr(original, "vault_role_id") + assert not hasattr(original, "vault_secret_id") + + +def test_worker_policy_identity_is_required_and_accepts_an_exact_match(): + expected = crypto.get_content_encryption_policy_identity() + + crypto.require_content_encryption_policy(expected) + with pytest.raises(crypto.ContentEncryptionPolicyMismatch, match="does not match"): + crypto.require_content_encryption_policy(None) + + +def test_secret_fingerprint_distinguishes_bytes_from_strings(): + assert crypto._secret_fingerprint(b"same-value") != crypto._secret_fingerprint("same-value") + + +def test_secret_fingerprint_supports_surrogate_escaped_strings(): + credential = "vault-role-\udcff" + + fingerprint = crypto._secret_fingerprint(credential) + + assert fingerprint == crypto._secret_fingerprint(credential) + assert credential not in fingerprint + + +@pytest.mark.parametrize( + ("credential_name", "replacement"), + [ + ("static_key", b"replacement-static-key"), + ("vault_role_id", "replacement-role-id"), + ("vault_secret_id", "replacement-secret-id"), + ], +) +def test_manager_cache_reuses_identical_config_and_recreates_for_changed_credentials( + monkeypatch, credential_name, replacement +): + credentials = { + "static_key": b"original-static-key", + "vault_role_id": "original-role-id", + "vault_secret_id": "original-secret-id", + } + configs = iter( + [ + crypto.ContentEncryptionConfig(mode="vault", **credentials), + crypto.ContentEncryptionConfig(mode="vault", **credentials), + crypto.ContentEncryptionConfig(mode="vault", **{**credentials, credential_name: replacement}), + ] + ) + monkeypatch.setattr(crypto, "get_content_encryption_config", lambda: next(configs)) + + original_manager = crypto.get_content_encryption_manager() + reused_manager = crypto.get_content_encryption_manager() + replacement_manager = crypto.get_content_encryption_manager() + + assert reused_manager is original_manager + assert replacement_manager is not original_manager + + +@pytest.mark.parametrize( + ("env_name", "mode"), + [ + ("AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS", "off"), + ("AIQ_CONTENT_ENCRYPTION_DEK_CACHE_TTL_SECONDS", "off"), + ("VAULT_TIMEOUT_SECONDS", "vault"), + ], +) +@pytest.mark.parametrize("non_finite_value", ["nan", "inf", "-inf"]) +def test_non_finite_numeric_config_is_rejected(monkeypatch, env_name, mode, non_finite_value): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", mode) + monkeypatch.setenv(env_name, non_finite_value) + if mode == "vault": + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.com") + monkeypatch.setenv("VAULT_ROLE_ID", "role-id") + monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") + monkeypatch.setenv("AIQ_ENCRYPTION_TRANSIT_KEY", "reports") + + with pytest.raises(crypto.ContentEncryptionConfigError): + crypto.get_content_encryption_config() + + +def test_static_key_envelope_round_trip(monkeypatch): + _enable_static_key(monkeypatch) + + cipher = crypto.create_job_content_cipher("job-1") + stored = cipher.encrypt_output_json('{"report":"secret report"}') + + assert stored.startswith(crypto.ENVELOPE_PREFIX) + assert "secret report" not in stored + assert crypto.read_job_output("job-1", stored) == {"report": "secret report"} + + +def test_aad_mismatch_fails(monkeypatch): + _enable_static_key(monkeypatch) + + stored = crypto.create_job_content_cipher("job-1").encrypt_output_json('{"report":"secret"}') + + with pytest.raises(crypto.ContentEncryptionInvalidData): + crypto.read_job_output("job-2", stored) + + +def test_tamper_fails(monkeypatch): + _enable_static_key(monkeypatch) + + stored = crypto.create_job_content_cipher("job-1").encrypt_output_json('{"report":"secret"}') + envelope = crypto.decode_envelope(stored) + padding = "=" * (-len(envelope["ciphertext"]) % 4) + ciphertext = bytearray(base64.urlsafe_b64decode(envelope["ciphertext"] + padding)) + ciphertext[0] ^= 0x01 + envelope["ciphertext"] = base64.urlsafe_b64encode(bytes(ciphertext)).decode("ascii").rstrip("=") + tampered = crypto.encode_envelope(envelope) + + with pytest.raises(crypto.ContentEncryptionInvalidData): + crypto.read_job_output("job-1", tampered) + + +def test_non_ascii_envelope_is_classified_as_invalid_data(): + with pytest.raises(crypto.ContentEncryptionInvalidData, match="envelope is malformed"): + crypto.decode_envelope(f"{crypto.ENVELOPE_PREFIX}\N{LATIN SMALL LETTER E WITH ACUTE}") + + +def test_plaintext_job_output_is_rejected_in_encrypted_mode(monkeypatch): + _enable_static_key(monkeypatch) + + with pytest.raises(crypto.ContentEncryptionPlaintextViolation): + crypto.read_job_output("job-1", '{"report":"plaintext"}') + + +def test_off_mode_preserves_current_behavior_and_does_not_decrypt_envelopes(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "off") + crypto.reset_content_encryption_manager_for_tests() + + assert crypto.read_job_output("job-1", '{"report":"plaintext"}') == {"report": "plaintext"} + assert crypto.read_job_output("job-1", "aiqenc:not-json") == "aiqenc:not-json" + + +@pytest.mark.asyncio +async def test_async_content_encryption_wrappers_run_off_event_loop(monkeypatch): + loop_thread = threading.get_ident() + call_threads = {} + + def fake_validate_startup(): + call_threads["startup"] = threading.get_ident() + return crypto.ContentEncryptionReadiness(mode="off", ready=True) + + def fake_health(): + call_threads["health"] = threading.get_ident() + return crypto.ContentEncryptionReadiness(mode="off", ready=True) + + def fake_require_ready(): + call_threads["readiness"] = threading.get_ident() + + def fake_read_output(job_id, stored_output): + call_threads["read_output"] = threading.get_ident() + return {"job_id": job_id, "stored_output": stored_output} + + monkeypatch.setattr(crypto, "validate_content_encryption_startup", fake_validate_startup) + monkeypatch.setattr(crypto, "get_content_encryption_health", fake_health) + monkeypatch.setattr(crypto, "require_content_encryption_ready_for_submission", fake_require_ready) + monkeypatch.setattr(crypto, "read_job_output", fake_read_output) + + assert await crypto.validate_content_encryption_startup_async() == crypto.ContentEncryptionReadiness( + mode="off", ready=True + ) + assert await crypto.get_content_encryption_health_async() == crypto.ContentEncryptionReadiness( + mode="off", ready=True + ) + await crypto.require_content_encryption_ready_for_submission_async() + assert await crypto.read_job_output_async("job-1", "stored") == {"job_id": "job-1", "stored_output": "stored"} + + assert set(call_threads) == {"startup", "health", "readiness", "read_output"} + assert all(thread_id != loop_thread for thread_id in call_threads.values()) + + +@pytest.mark.asyncio +async def test_async_content_encryption_wrapper_does_not_block_event_loop(monkeypatch): + import time + + def slow_read_output(_job_id, _stored_output): + time.sleep(0.1) + return {"report": "secret"} + + monkeypatch.setattr(crypto, "read_job_output", slow_read_output) + ticks = 0 + + async def ticker(): + nonlocal ticks + while ticks < 3: + await asyncio.sleep(0.01) + ticks += 1 + + read_task = asyncio.create_task(crypto.read_job_output_async("job-1", "stored")) + ticker_task = asyncio.create_task(ticker()) + + assert await read_task == {"report": "secret"} + await ticker_task + assert ticks >= 3 + + +@pytest.mark.asyncio +async def test_update_job_output_encrypts_entire_payload(monkeypatch): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + job_store = SimpleNamespace(update_status=AsyncMock()) + + await crypto.update_job_output(job_store, "job-1", "success", output={"report": "secret"}, cipher=cipher) + + stored = job_store.update_status.await_args.kwargs["output"] + assert stored.startswith(crypto.ENVELOPE_PREFIX) + assert "secret" not in stored + assert crypto.read_job_output("job-1", stored) == {"report": "secret"} + + +@pytest.mark.asyncio +async def test_sqlite_job_store_persists_encrypted_output(monkeypatch, tmp_path): + from nat.front_ends.fastapi.async_jobs.job_store import Base + from nat.front_ends.fastapi.async_jobs.job_store import JobStatus + from nat.front_ends.fastapi.async_jobs.job_store import JobStore + from nat.front_ends.fastapi.async_jobs.job_store import get_db_engine + + _enable_static_key(monkeypatch) + db_url = f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}" + engine = get_db_engine(db_url, use_async=True) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + job_store = JobStore(scheduler_address="tcp://localhost:8786", db_engine=engine) + await job_store._create_job(job_id="job-1") + cipher = crypto.create_job_content_cipher("job-1") + + await crypto.update_job_output(job_store, "job-1", JobStatus.SUCCESS, output={"report": "secret"}, cipher=cipher) + job = await job_store.get_job("job-1") + + assert job.output.startswith(crypto.ENVELOPE_PREFIX) + assert "secret" not in job.output + assert crypto.read_job_output("job-1", job.output) == {"report": "secret"} + await engine.dispose() + + +def test_static_key_invalid_config_fails_hard(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "key") + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", "not a 32 byte base64 key") + + with pytest.raises(crypto.ContentEncryptionConfigError): + crypto.get_content_encryption_config() + + +def test_large_report_round_trip(monkeypatch): + _enable_static_key(monkeypatch) + report = "large report\n" * 100_000 + + stored = crypto.create_job_content_cipher("job-large").encrypt_output_json(json.dumps({"report": report})) + + assert crypto.read_job_output("job-large", stored)["report"] == report + + +def test_dek_cache_reuses_unwrapped_dek(monkeypatch): + _enable_static_key(monkeypatch) + manager = crypto.get_content_encryption_manager() + stored = manager.create_job_cipher("job-1").encrypt_output_json('{"report":"secret"}') + calls = 0 + original = manager._unwrap_dek_with_static_key + + def counted_unwrap(envelope): + nonlocal calls + calls += 1 + return original(envelope) + + monkeypatch.setattr(manager, "_unwrap_dek_with_static_key", counted_unwrap) + + assert manager.decrypt_job_output_text("job-1", stored) == '{"report":"secret"}' + assert manager.decrypt_job_output_text("job-1", stored) == '{"report":"secret"}' + assert calls == 1 + + +def test_dek_cache_can_be_disabled(monkeypatch): + _enable_static_key(monkeypatch, cache_ttl="0") + manager = crypto.get_content_encryption_manager() + stored = manager.create_job_cipher("job-1").encrypt_output_json('{"report":"secret"}') + calls = 0 + original = manager._unwrap_dek_with_static_key + + def counted_unwrap(envelope): + nonlocal calls + calls += 1 + return original(envelope) + + monkeypatch.setattr(manager, "_unwrap_dek_with_static_key", counted_unwrap) + + assert manager.decrypt_job_output_text("job-1", stored) == '{"report":"secret"}' + assert manager.decrypt_job_output_text("job-1", stored) == '{"report":"secret"}' + assert calls == 2 + + +def test_dek_cache_evicts_lru_when_max_entries_is_exceeded(): + cache = crypto._DEKCache(ttl_seconds=60, max_entries=2) + + cache.put("first", b"1" * 32) + cache.put("second", b"2" * 32) + assert cache.get("first") == b"1" * 32 + cache.put("third", b"3" * 32) + + assert cache.get("second") is None + assert cache.get("first") == b"1" * 32 + assert cache.get("third") == b"3" * 32 + + +def test_dek_cache_is_thread_safe_under_concurrent_access(): + cache = crypto._DEKCache(ttl_seconds=60, max_entries=8) + + def exercise_cache(worker_id: int) -> None: + for index in range(500): + cache_key = f"key-{(worker_id + index) % 16}" + cache.put(cache_key, bytes([index % 256]) * crypto.DEK_BYTES) + cache.get(cache_key) + cache.get(f"key-{index % 16}") + + with ThreadPoolExecutor(max_workers=8) as executor: + futures = [executor.submit(exercise_cache, worker_id) for worker_id in range(8)] + for future in as_completed(futures): + future.result() + + cache.put("final", b"f" * crypto.DEK_BYTES) + assert cache.get("final") == b"f" * crypto.DEK_BYTES + + +def test_vault_client_initialization_is_thread_safe(monkeypatch): + clients = [] + login_calls = 0 + calls_lock = threading.Lock() + + class FakeAppRole: + def login(self, *, role_id, secret_id): + nonlocal login_calls + with calls_lock: + login_calls += 1 + assert role_id == "role-id" + assert secret_id == "secret-id" + clients[0].authenticated = True + + class FakeAuth: + def __init__(self): + self.approle = FakeAppRole() + + class FakeClient: + def __init__(self, *, url, namespace, timeout): + assert url == "https://vault.example.com" + assert namespace is None + assert timeout == crypto.DEFAULT_VAULT_TIMEOUT_SECONDS + self.auth = FakeAuth() + self.authenticated = False + clients.append(self) + + def is_authenticated(self): + return self.authenticated + + monkeypatch.setitem(sys.modules, "hvac", SimpleNamespace(Client=FakeClient)) + vault = _vault_client_for_tests() + + with ThreadPoolExecutor(max_workers=8) as executor: + futures = [executor.submit(vault._get_client) for _ in range(16)] + results = [future.result() for future in as_completed(futures)] + + assert len(clients) == 1 + assert login_calls == 1 + assert all(result is clients[0] for result in results) + + +def test_vault_retry_retries_transient_network_error(monkeypatch): + sleeps = _disable_vault_retry_sleep(monkeypatch) + vault = _vault_client_for_tests() + calls = 0 + + def operation(): + nonlocal calls + calls += 1 + if calls == 1: + raise requests_exceptions.Timeout("request timed out") + return {"data": {"ok": True}} + + assert vault._with_retry(operation, operation="unit") == {"data": {"ok": True}} + assert calls == 2 + assert len(sleeps) == 1 + + +def test_vault_retry_retries_transient_vault_error(monkeypatch): + sleeps = _disable_vault_retry_sleep(monkeypatch) + vault = _vault_client_for_tests() + calls = 0 + + def operation(): + nonlocal calls + calls += 1 + if calls == 1: + raise vault_exceptions.InternalServerError("vault server error") + return {"data": {"ok": True}} + + assert vault._with_retry(operation, operation="unit") == {"data": {"ok": True}} + assert calls == 2 + assert len(sleeps) == 1 + + +def test_vault_retry_fails_after_bounded_transient_attempts(monkeypatch): + sleeps = _disable_vault_retry_sleep(monkeypatch) + vault = _vault_client_for_tests() + calls = 0 + + def operation(): + nonlocal calls + calls += 1 + raise vault_exceptions.RateLimitExceeded("vault is busy") + + with pytest.raises(crypto.ContentEncryptionUnavailable, match="vault_unit_failed"): + vault._with_retry(operation, operation="unit") + + assert calls == crypto._VAULT_ATTEMPTS + assert len(sleeps) == crypto._VAULT_ATTEMPTS - 1 + + +def test_vault_retry_reauthenticates_after_unauthorized(monkeypatch): + sleeps = _disable_vault_retry_sleep(monkeypatch) + vault = _vault_client_for_tests() + calls = 0 + login_forces = [] + + def login(*, force): + login_forces.append(force) + + monkeypatch.setattr(vault, "_login", login) + + def operation(): + nonlocal calls + calls += 1 + if calls == 1: + raise vault_exceptions.Unauthorized("token expired") + return {"data": {"ok": True}} + + assert vault._with_retry(operation, operation="unit") == {"data": {"ok": True}} + assert calls == 2 + assert login_forces == [True] + assert len(sleeps) == 1 + + +@pytest.mark.parametrize("exc_cls", [vault_exceptions.Forbidden, vault_exceptions.InvalidRequest]) +def test_vault_retry_does_not_retry_permission_or_invalid_request_errors(monkeypatch, caplog, exc_cls): + sleeps = _disable_vault_retry_sleep(monkeypatch) + vault = _vault_client_for_tests() + calls = 0 + + def operation(): + nonlocal calls + calls += 1 + raise exc_cls("secret-value") + + with caplog.at_level("WARNING", logger=crypto.__name__): + with pytest.raises(crypto.ContentEncryptionUnavailable, match="vault_unit_failed"): + vault._with_retry(operation, operation="unit") + + assert calls == 1 + assert sleeps == [] + assert exc_cls.__name__ in caplog.text + assert "secret-value" not in caplog.text + + +def test_vault_missing_required_config_fails_startup(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "vault") + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.com") + monkeypatch.setenv("VAULT_ROLE_ID", "role-id") + monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") + + with pytest.raises(crypto.ContentEncryptionConfigError, match="AIQ_ENCRYPTION_TRANSIT_KEY"): + crypto.get_content_encryption_config() + + +def test_vault_operational_failure_starts_unhealthy_and_uses_readiness_cache(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "vault") + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.com") + monkeypatch.setenv("VAULT_ROLE_ID", "role-id") + monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") + monkeypatch.setenv("AIQ_ENCRYPTION_TRANSIT_KEY", "reports") + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS", "60") + calls = 0 + + class FailingVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + nonlocal calls + calls += 1 + raise crypto.ContentEncryptionUnavailable("vault down") + + monkeypatch.setattr(crypto, "_VaultTransitClient", FailingVault) + crypto.reset_content_encryption_manager_for_tests() + + startup = crypto.validate_content_encryption_startup() + health = crypto.get_content_encryption_health() + + assert startup.ready is False + assert health.ready is False + assert health.encrypt_ready is False + assert health.decrypt_ready is False + assert health.reason == "vault_generate_unavailable" + assert calls == 1 + + +def test_vault_readiness_requires_generate_and_unwrap(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "vault") + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.com") + monkeypatch.setenv("VAULT_ROLE_ID", "role-id") + monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") + monkeypatch.setenv("AIQ_ENCRYPTION_TRANSIT_KEY", "reports") + calls = {"generate": 0, "unwrap": 0} + dek = b"d" * crypto.DEK_BYTES + + class ReadyVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + calls["generate"] += 1 + assert operation == "api_startup_readiness_generate" + return dek, crypto.WrappedDEK(wrap="vault", kid="transit/reports", wrapped_dek="vault:v1:dek") + + def unwrap_dek(self, wrapped_dek, *, operation): + calls["unwrap"] += 1 + assert wrapped_dek == "vault:v1:dek" + assert operation == "api_startup_readiness_decrypt" + return dek + + monkeypatch.setattr(crypto, "_VaultTransitClient", ReadyVault) + crypto.reset_content_encryption_manager_for_tests() + + readiness = crypto.validate_content_encryption_startup() + + assert readiness.ready is True + assert readiness.encrypt_ready is True + assert readiness.decrypt_ready is True + assert calls == {"generate": 1, "unwrap": 1} + + +def test_vault_readiness_fails_when_unwrap_is_denied(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "vault") + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.com") + monkeypatch.setenv("VAULT_ROLE_ID", "role-id") + monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") + monkeypatch.setenv("AIQ_ENCRYPTION_TRANSIT_KEY", "reports") + + class DecryptDeniedVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + return b"d" * crypto.DEK_BYTES, crypto.WrappedDEK( + wrap="vault", kid="transit/reports", wrapped_dek="vault:v1:dek" + ) + + def unwrap_dek(self, wrapped_dek, *, operation): + raise crypto.ContentEncryptionUnavailable("decrypt denied") + + monkeypatch.setattr(crypto, "_VaultTransitClient", DecryptDeniedVault) + crypto.reset_content_encryption_manager_for_tests() + + readiness = crypto.validate_content_encryption_startup() + + assert readiness.ready is False + assert readiness.encrypt_ready is True + assert readiness.decrypt_ready is False + assert readiness.reason == "vault_decrypt_unavailable" + assert readiness.exception_type == "ContentEncryptionUnavailable" + + +def test_vault_readiness_fails_when_unwrapped_dek_differs(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "vault") + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.com") + monkeypatch.setenv("VAULT_ROLE_ID", "role-id") + monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") + monkeypatch.setenv("AIQ_ENCRYPTION_TRANSIT_KEY", "reports") + + class MismatchedVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + return b"d" * crypto.DEK_BYTES, crypto.WrappedDEK( + wrap="vault", kid="transit/reports", wrapped_dek="vault:v1:dek" + ) + + def unwrap_dek(self, wrapped_dek, *, operation): + return b"e" * crypto.DEK_BYTES + + monkeypatch.setattr(crypto, "_VaultTransitClient", MismatchedVault) + crypto.reset_content_encryption_manager_for_tests() + + readiness = crypto.validate_content_encryption_startup() + + assert readiness.ready is False + assert readiness.encrypt_ready is True + assert readiness.decrypt_ready is False + assert readiness.reason == "vault_readiness_dek_mismatch" + + +def test_vault_readiness_cache_avoids_repeated_generate_and_unwrap(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "vault") + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.com") + monkeypatch.setenv("VAULT_ROLE_ID", "role-id") + monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") + monkeypatch.setenv("AIQ_ENCRYPTION_TRANSIT_KEY", "reports") + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS", "60") + calls = {"generate": 0, "unwrap": 0} + dek = b"d" * crypto.DEK_BYTES + + class ReadyVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + calls["generate"] += 1 + return dek, crypto.WrappedDEK(wrap="vault", kid="transit/reports", wrapped_dek="vault:v1:dek") + + def unwrap_dek(self, wrapped_dek, *, operation): + calls["unwrap"] += 1 + return dek + + monkeypatch.setattr(crypto, "_VaultTransitClient", ReadyVault) + crypto.reset_content_encryption_manager_for_tests() + + startup = crypto.validate_content_encryption_startup() + health = crypto.get_content_encryption_health() + + assert startup.ready is True + assert health.ready is True + assert calls == {"generate": 1, "unwrap": 1} + + +def test_vault_readiness_rechecks_when_cache_is_stale(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "vault") + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.com") + monkeypatch.setenv("VAULT_ROLE_ID", "role-id") + monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") + monkeypatch.setenv("AIQ_ENCRYPTION_TRANSIT_KEY", "reports") + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS", "0") + calls = {"generate": 0, "unwrap": 0} + dek = b"d" * crypto.DEK_BYTES + + class ReadyVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + calls["generate"] += 1 + return dek, crypto.WrappedDEK(wrap="vault", kid="transit/reports", wrapped_dek="vault:v1:dek") + + def unwrap_dek(self, wrapped_dek, *, operation): + calls["unwrap"] += 1 + return dek + + monkeypatch.setattr(crypto, "_VaultTransitClient", ReadyVault) + crypto.reset_content_encryption_manager_for_tests() + + crypto.validate_content_encryption_startup() + crypto.get_content_encryption_health() + + assert calls == {"generate": 2, "unwrap": 2} + + +@pytest.mark.skipif(not _real_vault_env_present(), reason="real Vault Transit credentials are not configured") +def test_real_vault_transit_round_trip(monkeypatch): + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "vault") + crypto.reset_content_encryption_manager_for_tests() + + readiness = crypto.validate_content_encryption_startup() + stored = crypto.create_job_content_cipher("job-real-vault").encrypt_output_json('{"report":"secret"}') + + assert readiness.ready is True + assert readiness.encrypt_ready is True + assert readiness.decrypt_ready is True + assert stored.startswith(crypto.ENVELOPE_PREFIX) + assert "secret" not in stored + assert crypto.read_job_output("job-real-vault", stored) == {"report": "secret"} diff --git a/frontends/aiq_api/tests/test_content_encryption_routes.py b/frontends/aiq_api/tests/test_content_encryption_routes.py new file mode 100644 index 000000000..557a2359f --- /dev/null +++ b/frontends/aiq_api/tests/test_content_encryption_routes.py @@ -0,0 +1,702 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +from __future__ import annotations + +import base64 +import json +from datetime import UTC +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock +from unittest.mock import MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.routing import APIRoute +from fastapi.testclient import TestClient + +from aiq_agent.auth import Principal +from aiq_api.jobs import crypto +from aiq_api.registry import AgentConfig + + +def _static_key() -> str: + return base64.urlsafe_b64encode(bytes(range(32))).decode("ascii") + + +def _other_static_key() -> str: + return base64.urlsafe_b64encode(bytes(reversed(range(32)))).decode("ascii") + + +@pytest.fixture(autouse=True) +def clean_encryption_route_env(monkeypatch): + for name in ( + "AIQ_CONTENT_ENCRYPTION", + "AIQ_CONTENT_ENCRYPTION_KEY", + "AIQ_CONTENT_ENCRYPTION_KEY_ID", + "AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS", + "AIQ_CONTENT_ENCRYPTION_DEK_CACHE_TTL_SECONDS", + "VAULT_ADDR", + "VAULT_NAMESPACE", + "VAULT_TRANSIT_MOUNT", + "VAULT_ROLE_ID", + "VAULT_SECRET_ID", + "AIQ_ENCRYPTION_TRANSIT_KEY", + "VAULT_TIMEOUT_SECONDS", + "REQUIRE_AUTH", + ): + monkeypatch.delenv(name, raising=False) + crypto.reset_content_encryption_manager_for_tests() + yield + crypto.reset_content_encryption_manager_for_tests() + + +def _enable_static_key(monkeypatch) -> None: + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "key") + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", _static_key()) + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY_ID", "test-key") + crypto.reset_content_encryption_manager_for_tests() + + +def _enable_vault(monkeypatch) -> None: + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "vault") + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.com") + monkeypatch.setenv("VAULT_ROLE_ID", "role-id") + monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") + monkeypatch.setenv("AIQ_ENCRYPTION_TRANSIT_KEY", "reports") + crypto.reset_content_encryption_manager_for_tests() + + +async def _build_jobs_app(monkeypatch, tmp_path, *, job_output=None, submitted_job=None) -> FastAPI: + import aiq_api.routes.jobs as jobs_routes + from aiq_api.jobs import access + from aiq_api.jobs import event_store + from aiq_api.jobs import submit + + monkeypatch.setattr(jobs_routes, "_start_periodic_cleanup", MagicMock()) + + async def _no_op_reaper(*_args, **_kwargs): + return None + + monkeypatch.setattr(jobs_routes, "_reap_ghost_jobs", _no_op_reaper) + monkeypatch.setattr( + jobs_routes, + "require_verified_principal", + lambda: Principal(type="jwt", sub="user-1", email="user@example.com"), + ) + monkeypatch.setattr(event_store.EventStore, "_ensure_table_exists", MagicMock()) + + if submitted_job is not None: + # register_job_routes imports this helper after the patch and captures the mock. + monkeypatch.setattr(submit, "submit_agent_job", submitted_job) + + agent_config = AgentConfig( + class_path="aiq_agent.agents.deep_researcher.agent.DeepResearcherAgent", + config_name="deep_research_agent", + description="Test deep researcher", + ) + monkeypatch.setattr(jobs_routes, "get_agent_config", lambda _agent_type: agent_config) + + job = SimpleNamespace( + job_id="job-1", + status="success", + error=None, + output=job_output, + created_at=datetime.now(UTC), + ) + job_store = SimpleNamespace(get_job=AsyncMock(return_value=job), update_status=AsyncMock()) + db_url = f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}" + access._job_access_schema_initialized.clear() + + worker = SimpleNamespace( + _dask_available=True, + _job_store=job_store, + _scheduler_address="tcp://localhost:8786", + _db_url=db_url, + _config_file_path="config.yml", + _log_level=20, + _use_dask_threads=False, + _front_end_config=SimpleNamespace(expiry_seconds=86400), + ) + builder = MagicMock() + builder.get_function_config.return_value = SimpleNamespace(tools=[], exclude_tools=[]) + builder.get_tools = AsyncMock(return_value=[]) + + app = FastAPI() + await jobs_routes.register_job_routes(app, builder, worker) + return app + + +def _build_assembled_worker_app( + monkeypatch, tmp_path, *, dask_available: bool = True, db_url: str | None = None +) -> FastAPI: + """Build the AI-Q worker app while preserving NAT-before-AIQ route ordering.""" + import aiq_api.routes.jobs as jobs_routes + from aiq_api import plugin + from aiq_api.jobs import access + from nat.builder.workflow_builder import WorkflowBuilder + from nat.data_models.config import Config + from nat.data_models.config import GeneralConfig + from nat.front_ends.fastapi import fastapi_front_end_plugin_worker as worker_module + from nat.plugins.eval.fastapi import routes as eval_routes + from nat.plugins.mcp.client import fastapi_routes as mcp_routes + + monkeypatch.setenv("NAT_CONFIG_FILE", "config.yml") + monkeypatch.delenv("NAT_DASK_SCHEDULER_ADDRESS", raising=False) + monkeypatch.setattr(plugin, "_load_validators_from_entry_points", lambda: []) + monkeypatch.setattr(plugin.AIQAPIWorker, "_install_signal_handlers", lambda _self: None) + monkeypatch.setattr(jobs_routes, "_start_periodic_cleanup", MagicMock()) + monkeypatch.setattr(access, "ensure_job_access_table", MagicMock()) + monkeypatch.setattr( + worker_module.FastApiFrontEndPluginWorker, + "_create_session_manager", + AsyncMock(return_value=MagicMock()), + ) + monkeypatch.setattr(worker_module.FastApiFrontEndPluginWorker, "add_default_route", AsyncMock()) + for route_helper in ( + "add_authorization_route", + "add_execution_routes", + "add_monitor_route", + "add_static_files_route", + ): + monkeypatch.setattr(worker_module, route_helper, AsyncMock()) + monkeypatch.setattr(eval_routes, "add_evaluate_routes", AsyncMock()) + monkeypatch.setattr(mcp_routes, "add_mcp_client_tool_list_route", AsyncMock()) + + async def _no_op_reaper(*_args, **_kwargs): + return None + + monkeypatch.setattr(jobs_routes, "_reap_ghost_jobs", _no_op_reaper) + + builder = MagicMock() + builder.get_function_config.return_value = SimpleNamespace(tools=[], exclude_tools=[]) + builder.get_tools = AsyncMock(return_value=[]) + + class BuilderContext: + async def __aenter__(self): + return builder + + async def __aexit__(self, exc_type, exc, tb): + return False + + monkeypatch.setattr(WorkflowBuilder, "from_config", lambda _config: BuilderContext()) + + front_end = plugin.AIQAPIConfig(db_url=db_url or f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}") + worker = plugin.AIQAPIWorker(Config(general=GeneralConfig(front_end=front_end))) + worker._dask_available = dask_available + worker._job_store = object() if dask_available else None + worker._scheduler_address = "tcp://localhost:8786" + worker._db_url = front_end.db_url + return worker.build_app() + + +def _get_health_routes(app: FastAPI) -> list[APIRoute]: + return [ + route + for route in app.routes + if isinstance(route, APIRoute) and route.path == "/health" and "GET" in route.methods + ] + + +def _get_liveness_routes(app: FastAPI) -> list[APIRoute]: + return [ + route + for route in app.routes + if isinstance(route, APIRoute) and route.path == "/live" and "GET" in route.methods + ] + + +@pytest.mark.asyncio +async def test_health_returns_503_when_vault_readiness_failed(monkeypatch, tmp_path): + _enable_vault(monkeypatch) + + class FailingVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + raise crypto.ContentEncryptionUnavailable("vault down") + + monkeypatch.setattr(crypto, "_VaultTransitClient", FailingVault) + app = await _build_jobs_app(monkeypatch, tmp_path) + + with TestClient(app) as client: + response = client.get("/health") + liveness_response = client.get("/live") + + assert response.status_code == 503 + body = response.json() + assert body["encryption"]["mode"] == "vault" + assert body["encryption"]["ready"] is False + assert body["encryption"]["encrypt_ready"] is False + assert body["encryption"]["decrypt_ready"] is False + assert liveness_response.status_code == 200 + assert liveness_response.json() == {"status": "alive"} + + +def test_assembled_worker_serves_encryption_health_route(monkeypatch, tmp_path): + _enable_vault(monkeypatch) + + class FailingVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + raise crypto.ContentEncryptionUnavailable("vault down") + + monkeypatch.setattr(crypto, "_VaultTransitClient", FailingVault) + app = _build_assembled_worker_app(monkeypatch, tmp_path) + + with TestClient(app) as client: + response = client.get("/health") + liveness_response = client.get("/live") + health_routes = _get_health_routes(app) + liveness_routes = _get_liveness_routes(app) + + assert len(health_routes) == 1 + assert health_routes[0].endpoint.__module__ == "aiq_api.routes.jobs" + assert len(liveness_routes) == 1 + assert liveness_routes[0].endpoint.__module__ == "aiq_api.routes.jobs" + assert liveness_response.status_code == 200 + assert liveness_response.json() == {"status": "alive"} + assert response.status_code == 503 + body = response.json() + assert body["status"] == "degraded" + assert body["encryption"]["mode"] == "vault" + assert body["encryption"]["ready"] is False + + +def test_assembled_worker_keeps_liveness_when_async_jobs_are_unavailable(monkeypatch, tmp_path): + app = _build_assembled_worker_app(monkeypatch, tmp_path, dask_available=False) + + with TestClient(app) as client: + liveness_response = client.get("/live") + readiness_response = client.get("/health") + openapi = client.get("/openapi.json").json() + + assert liveness_response.status_code == 200 + assert liveness_response.json() == {"status": "alive"} + assert readiness_response.status_code == 503 + assert readiness_response.json() == { + "status": "degraded", + "dask_available": False, + "db": "unchecked", + "reason": "async_jobs_unavailable", + } + assert len(_get_liveness_routes(app)) == 1 + assert len(_get_health_routes(app)) == 1 + assert list(openapi["paths"]["/live"]) == ["get"] + assert list(openapi["paths"]["/health"]) == ["get"] + + +@pytest.mark.parametrize("mode", ["off", "key"]) +def test_assembled_worker_health_route_reports_ready_encryption(monkeypatch, tmp_path, mode): + if mode == "key": + _enable_static_key(monkeypatch) + app = _build_assembled_worker_app(monkeypatch, tmp_path) + + with TestClient(app) as client: + response = client.get("/health") + openapi = client.get("/openapi.json").json() + health_routes = _get_health_routes(app) + + assert len(health_routes) == 1 + assert health_routes[0].endpoint.__module__ == "aiq_api.routes.jobs" + assert response.status_code == 200 + body = response.json() + assert body["status"] == "healthy" + assert body["encryption"] == {"mode": mode, "ready": True} + assert list(openapi["paths"]["/live"]) == ["get"] + assert list(openapi["paths"]["/health"]) == ["get"] + + +def test_assembled_worker_health_returns_503_when_db_unreachable_on_fresh_process(monkeypatch, tmp_path): + """A fresh process with an empty async-engine cache must still ping the DB. + + Regression: the health check previously reported ``db: no_engine`` (HTTP 200) + whenever no async engine happened to be cached yet, so a fresh pod could pass + Helm readiness straight through a database outage. + """ + from sqlalchemy.exc import OperationalError + + from aiq_api.jobs.event_store import EventStore + + app = _build_assembled_worker_app(monkeypatch, tmp_path) + + with TestClient(app) as client: + # The process started cleanly, but by the time the readiness probe runs + # no async engine has been cached (fresh process) and the database has + # become unavailable. + EventStore._async_engine_cache.clear() + + def _unreachable(_db_url): + raise OperationalError("SELECT 1", {}, Exception("database is unavailable")) + + monkeypatch.setattr(EventStore, "_get_or_create_async_engine", _unreachable) + + response = client.get("/health") + liveness_response = client.get("/live") + + assert response.status_code == 503 + body = response.json() + assert body["status"] == "degraded" + assert body["db"] == "unreachable" + # Liveness must stay decoupled from dependency health. + assert liveness_response.status_code == 200 + assert liveness_response.json() == {"status": "alive"} + + +@pytest.mark.asyncio +async def test_health_returns_503_when_vault_decrypt_readiness_failed(monkeypatch, tmp_path): + _enable_vault(monkeypatch) + + class DecryptDeniedVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + return b"d" * crypto.DEK_BYTES, crypto.WrappedDEK( + wrap="vault", kid="transit/reports", wrapped_dek="vault:v1:dek" + ) + + def unwrap_dek(self, wrapped_dek, *, operation): + raise crypto.ContentEncryptionUnavailable("decrypt denied") + + monkeypatch.setattr(crypto, "_VaultTransitClient", DecryptDeniedVault) + app = await _build_jobs_app(monkeypatch, tmp_path) + + with TestClient(app) as client: + response = client.get("/health") + + assert response.status_code == 503 + body = response.json() + assert body["encryption"]["mode"] == "vault" + assert body["encryption"]["ready"] is False + assert body["encryption"]["encrypt_ready"] is True + assert body["encryption"]["decrypt_ready"] is False + assert body["encryption"]["reason"] == "vault_decrypt_unavailable" + + +@pytest.mark.asyncio +async def test_submit_rejects_when_encryption_readiness_failed(monkeypatch, tmp_path): + _enable_vault(monkeypatch) + submitted_job = AsyncMock(return_value="job-1") + + class FailingVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + raise crypto.ContentEncryptionUnavailable("vault down") + + monkeypatch.setattr(crypto, "_VaultTransitClient", FailingVault) + app = await _build_jobs_app(monkeypatch, tmp_path, submitted_job=submitted_job) + + with TestClient(app) as client: + response = client.post("/v1/jobs/async/submit", json={"agent_type": "deep_researcher", "input": "query"}) + + assert response.status_code == 503 + submitted_job.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_submit_rejects_when_vault_decrypt_readiness_failed(monkeypatch, tmp_path): + _enable_vault(monkeypatch) + submitted_job = AsyncMock(return_value="job-1") + + class DecryptDeniedVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + return b"d" * crypto.DEK_BYTES, crypto.WrappedDEK( + wrap="vault", kid="transit/reports", wrapped_dek="vault:v1:dek" + ) + + def unwrap_dek(self, wrapped_dek, *, operation): + raise crypto.ContentEncryptionUnavailable("decrypt denied") + + monkeypatch.setattr(crypto, "_VaultTransitClient", DecryptDeniedVault) + app = await _build_jobs_app(monkeypatch, tmp_path, submitted_job=submitted_job) + + with TestClient(app) as client: + response = client.post("/v1/jobs/async/submit", json={"agent_type": "deep_researcher", "input": "query"}) + + assert response.status_code == 503 + submitted_job.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_submit_uses_authorized_submission_when_encryption_is_ready(monkeypatch, tmp_path): + _enable_static_key(monkeypatch) + submitted_job = AsyncMock(return_value="job-1") + app = await _build_jobs_app(monkeypatch, tmp_path, submitted_job=submitted_job) + + with TestClient(app) as client: + response = client.post("/v1/jobs/async/submit", json={"agent_type": "deep_researcher", "input": "query"}) + + assert response.status_code == 200 + assert response.json()["job_id"] == "job-1" + submitted_job.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_submit_openapi_documents_encryption_failures(monkeypatch, tmp_path): + app = await _build_jobs_app(monkeypatch, tmp_path) + + responses = app.openapi()["paths"]["/v1/jobs/async/submit"]["post"]["responses"] + + assert responses["400"]["description"] == "Unknown agent type or invalid request" + assert responses["422"]["description"] == "One or more unknown or agent-unavailable data source IDs" + assert responses["500"]["description"] == ( + "Content encryption configuration is invalid or async job authorization persistence failed" + ) + assert responses["503"]["description"] == ("Content encryption, Dask scheduler, or sandbox capacity is unavailable") + + +@pytest.mark.asyncio +async def test_report_authorizes_before_decrypting(monkeypatch, tmp_path): + _enable_static_key(monkeypatch) + monkeypatch.setenv("REQUIRE_AUTH", "true") + app = await _build_jobs_app(monkeypatch, tmp_path, job_output='{"report":"plaintext"}') + + with TestClient(app) as client: + response = client.get("/v1/jobs/async/job/job-1/report") + + assert response.status_code == 404 + + +@pytest.mark.asyncio +async def test_report_returns_500_for_plaintext_violation(monkeypatch, tmp_path): + _enable_static_key(monkeypatch) + app = await _build_jobs_app(monkeypatch, tmp_path, job_output='{"report":"plaintext"}') + + with TestClient(app) as client: + response = client.get("/v1/jobs/async/job/job-1/report") + + assert response.status_code == 500 + assert response.json()["detail"] == "Final report data is invalid" + + +@pytest.mark.parametrize("job_output", ["", {}]) +@pytest.mark.asyncio +async def test_report_returns_500_for_falsey_plaintext_violation(monkeypatch, tmp_path, job_output): + _enable_static_key(monkeypatch) + app = await _build_jobs_app(monkeypatch, tmp_path, job_output=job_output) + + with TestClient(app) as client: + response = client.get("/v1/jobs/async/job/job-1/report") + + assert response.status_code == 500 + assert response.json()["detail"] == "Final report data is invalid" + + +@pytest.mark.asyncio +async def test_report_returns_500_for_malformed_envelope(monkeypatch, tmp_path): + _enable_static_key(monkeypatch) + app = await _build_jobs_app(monkeypatch, tmp_path, job_output="aiqenc:not-json") + + with TestClient(app) as client: + response = client.get("/v1/jobs/async/job/job-1/report") + + assert response.status_code == 500 + assert response.json()["detail"] == "Final report data is invalid" + + +@pytest.mark.asyncio +async def test_report_returns_503_when_vault_decrypt_is_unavailable(monkeypatch, tmp_path): + _enable_vault(monkeypatch) + + class UnwrapFailingVault: + def __init__(self, _config): + pass + + def generate_data_key(self, *, operation): + return b"0" * 32, crypto.WrappedDEK(wrap="vault", kid="transit/reports", wrapped_dek="vault:v1:dek") + + def unwrap_dek(self, wrapped_dek, *, operation): + raise crypto.ContentEncryptionUnavailable("vault down") + + envelope = crypto.encode_envelope( + { + "v": crypto.ENVELOPE_VERSION, + "alg": crypto.CONTENT_ALGORITHM, + "wrap": "vault", + "kid": "transit/reports", + "aad_hint": crypto.job_output_aad("job-1"), + "wrapped_dek": "vault:v1:dek", + "nonce": base64.urlsafe_b64encode(b"1" * 12).decode("ascii").rstrip("="), + "ciphertext": base64.urlsafe_b64encode(b"ciphertext").decode("ascii").rstrip("="), + "tag": base64.urlsafe_b64encode(b"2" * 16).decode("ascii").rstrip("="), + } + ) + monkeypatch.setattr(crypto, "_VaultTransitClient", UnwrapFailingVault) + app = await _build_jobs_app(monkeypatch, tmp_path, job_output=envelope) + + with TestClient(app) as client: + response = client.get("/v1/jobs/async/job/job-1/report") + + assert response.status_code == 503 + assert response.json()["detail"] == "Content encryption is unavailable" + + +@pytest.mark.asyncio +async def test_report_decrypts_encrypted_final_output(monkeypatch, tmp_path): + _enable_static_key(monkeypatch) + stored = crypto.create_job_content_cipher("job-1").encrypt_output_json( + json.dumps( + { + "report": "secret", + "parent_job_id": "parent-job", + "interaction_action": "edit", + "result_kind": "report", + } + ) + ) + app = await _build_jobs_app(monkeypatch, tmp_path, job_output=stored) + + with TestClient(app) as client: + response = client.get("/v1/jobs/async/job/job-1/report") + + assert response.status_code == 200 + assert response.json() == { + "job_id": "job-1", + "has_report": True, + "report": "secret", + "parent_job_id": "parent-job", + "interaction_action": "edit", + "result_kind": "report", + } + + +@pytest.mark.asyncio +async def test_state_returns_500_for_invalid_encrypted_event_data(monkeypatch, tmp_path): + from aiq_api.jobs.event_store import EventStore + + _enable_static_key(monkeypatch) + app = await _build_jobs_app(monkeypatch, tmp_path) + db_url = f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}" + cipher = crypto.create_job_content_cipher("job-1") + EventStore(db_url, "job-1", content_cipher=cipher).store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "secret report", + }, + } + ) + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", _other_static_key()) + crypto.reset_content_encryption_manager_for_tests() + + with TestClient(app) as client: + response = client.get("/v1/jobs/async/job/job-1/state") + + assert response.status_code == 500 + assert response.json()["detail"] == "Job state data is invalid" + + +@pytest.mark.asyncio +async def test_state_returns_503_when_event_decrypt_is_unavailable(monkeypatch, tmp_path): + from aiq_api.jobs.event_store import EventStore + + _enable_static_key(monkeypatch) + app = await _build_jobs_app(monkeypatch, tmp_path) + db_url = f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}" + cipher = crypto.create_job_content_cipher("job-1") + EventStore(db_url, "job-1", content_cipher=cipher).store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "secret report", + }, + } + ) + + def unavailable_decrypt(*_args, **_kwargs): + raise crypto.ContentEncryptionUnavailable("vault down") + + monkeypatch.setattr(crypto, "decrypt_event_field", unavailable_decrypt) + + with TestClient(app) as client: + response = client.get("/v1/jobs/async/job/job-1/state") + + assert response.status_code == 503 + assert response.json()["detail"] == "Content encryption is unavailable" + + +@pytest.mark.asyncio +async def test_sse_emits_job_error_for_invalid_encrypted_event_data(monkeypatch, tmp_path): + from aiq_api.jobs.event_store import EventStore + + _enable_static_key(monkeypatch) + app = await _build_jobs_app(monkeypatch, tmp_path) + db_url = f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}" + cipher = crypto.create_job_content_cipher("job-1") + EventStore(db_url, "job-1", content_cipher=cipher).store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "secret report", + }, + } + ) + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", _other_static_key()) + crypto.reset_content_encryption_manager_for_tests() + + with TestClient(app) as client: + with client.stream("GET", "/v1/jobs/async/job/job-1/stream") as response: + body = "".join(response.iter_text()) + + assert response.status_code == 200 + assert "event: job.error" in body + assert "Job event data is invalid" in body + assert "secret report" not in body + + +@pytest.mark.asyncio +async def test_sse_emits_job_error_when_event_decrypt_is_unavailable(monkeypatch, tmp_path): + from aiq_api.jobs.event_store import EventStore + + _enable_static_key(monkeypatch) + app = await _build_jobs_app(monkeypatch, tmp_path) + db_url = f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}" + cipher = crypto.create_job_content_cipher("job-1") + EventStore(db_url, "job-1", content_cipher=cipher).store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "secret report", + }, + } + ) + + def unavailable_decrypt(*_args, **_kwargs): + raise crypto.ContentEncryptionUnavailable("vault down") + + monkeypatch.setattr(crypto, "decrypt_event_field", unavailable_decrypt) + + with TestClient(app) as client: + with client.stream("GET", "/v1/jobs/async/job/job-1/stream") as response: + body = "".join(response.iter_text()) + + assert response.status_code == 200 + assert "event: job.error" in body + assert "Content encryption is unavailable" in body + assert "secret report" not in body diff --git a/frontends/aiq_api/tests/test_event_store_content_encryption.py b/frontends/aiq_api/tests/test_event_store_content_encryption.py new file mode 100644 index 000000000..d22f8bc5d --- /dev/null +++ b/frontends/aiq_api/tests/test_event_store_content_encryption.py @@ -0,0 +1,462 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +from __future__ import annotations + +import asyncio +import base64 +import json +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import as_completed + +import pytest +from sqlalchemy import text + +from aiq_api.jobs import crypto +from aiq_api.jobs.event_store import BatchingEventStore +from aiq_api.jobs.event_store import EventStore + + +def _static_key() -> str: + return base64.urlsafe_b64encode(bytes(range(32))).decode("ascii") + + +def _other_static_key() -> str: + return base64.urlsafe_b64encode(bytes(reversed(range(32)))).decode("ascii") + + +@pytest.fixture(autouse=True) +def clean_encryption_env(monkeypatch): + for name in ( + "AIQ_CONTENT_ENCRYPTION", + "AIQ_CONTENT_ENCRYPTION_KEY", + "AIQ_CONTENT_ENCRYPTION_KEY_ID", + "AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS", + "AIQ_CONTENT_ENCRYPTION_DEK_CACHE_TTL_SECONDS", + "VAULT_ADDR", + "VAULT_NAMESPACE", + "VAULT_TRANSIT_MOUNT", + "VAULT_ROLE_ID", + "VAULT_SECRET_ID", + "AIQ_ENCRYPTION_TRANSIT_KEY", + "VAULT_TIMEOUT_SECONDS", + ): + monkeypatch.delenv(name, raising=False) + EventStore._tables_initialized.clear() + crypto.reset_content_encryption_manager_for_tests() + yield + EventStore._tables_initialized.clear() + crypto.reset_content_encryption_manager_for_tests() + + +@pytest.fixture +def db_url(tmp_path): + return f"sqlite+aiosqlite:///{tmp_path / 'events.db'}" + + +def _enable_static_key(monkeypatch) -> None: + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "key") + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", _static_key()) + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY_ID", "test-key") + crypto.reset_content_encryption_manager_for_tests() + + +def _raw_event_data(db_url: str) -> dict: + engine = EventStore._get_or_create_sync_engine(db_url) + with engine.connect() as conn: + row = conn.execute(text("SELECT event_data FROM job_events ORDER BY id DESC LIMIT 1")).fetchone() + assert row is not None + return json.loads(row[0]) + + +def _latest_event_id(db_url: str) -> int: + engine = EventStore._get_or_create_sync_engine(db_url) + with engine.connect() as conn: + row = conn.execute(text("SELECT id FROM job_events ORDER BY id DESC LIMIT 1")).fetchone() + assert row is not None + return row[0] + + +def test_output_artifact_content_is_encrypted_at_rest_and_decrypted_on_read(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "secret report", + "output_category": "final_report", + }, + } + ) + + raw_event = _raw_event_data(db_url) + raw_content = raw_event["data"]["content"] + assert raw_event["data"]["type"] == "output" + assert raw_event["data"]["output_category"] == "final_report" + assert raw_content[crypto.ENCRYPTED_FIELD_MARKER] is True + assert raw_content[crypto.ENCRYPTED_FIELD_VALUE].startswith(crypto.ENVELOPE_PREFIX) + assert "secret report" not in json.dumps(raw_event) + + events = EventStore.get_events(db_url, "job-1") + assert events[0]["data"]["content"] == "secret report" + + +def test_file_artifact_content_is_encrypted_in_batch(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + + store.store_batch( + [ + { + "type": "artifact.update", + "data": { + "type": "file", + "content": "secret file content", + "file_path": "/shared/output.md", + }, + } + ] + ) + + raw_event = _raw_event_data(db_url) + assert raw_event["data"]["content"][crypto.ENCRYPTED_FIELD_VALUE].startswith(crypto.ENVELOPE_PREFIX) + assert "secret file content" not in json.dumps(raw_event) + + events = EventStore.get_events(db_url, "job-1") + assert events[0]["data"]["content"] == "secret file content" + + +@pytest.mark.parametrize("use_batch", [False, True]) +def test_encrypted_event_persistence_failure_is_raised(monkeypatch, db_url, use_batch): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + + def fail_connect(): + raise RuntimeError("transient database failure") + + monkeypatch.setattr(store._sync_engine, "connect", fail_connect) + event = {"type": "test.event", "data": {"value": "test"}} + + with pytest.raises(RuntimeError, match="transient database failure"): + if use_batch: + store.store_batch([event]) + else: + store.store(event) + + +@pytest.mark.parametrize("use_batch", [False, True]) +def test_off_mode_event_persistence_failure_remains_best_effort(monkeypatch, db_url, use_batch): + store = EventStore(db_url, "job-1", content_cipher=crypto.create_job_content_cipher("job-1")) + + def fail_connect(): + raise RuntimeError("transient database failure") + + monkeypatch.setattr(store._sync_engine, "connect", fail_connect) + event = {"type": "test.event", "data": {"value": "test"}} + + if use_batch: + store.store_batch([event]) + else: + store.store(event) + + +def test_background_encrypted_batch_failure_is_raised_by_foreground_flush(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + raw_store = EventStore(db_url, "job-1", content_cipher=cipher) + flush_attempted = threading.Event() + attempts = 0 + original_connect = raw_store._sync_engine.connect + + def fail_once_connect(): + nonlocal attempts + attempts += 1 + flush_attempted.set() + if attempts == 1: + raise RuntimeError("transient database failure") + return original_connect() + + monkeypatch.setattr(raw_store._sync_engine, "connect", fail_once_connect) + store = BatchingEventStore(raw_store) + store.FLUSH_INTERVAL_MS = 1 + + store.store( + { + "type": "artifact.update", + "data": {"type": "output", "content": "secret report"}, + } + ) + + assert flush_attempted.wait(timeout=1) + with pytest.raises(RuntimeError, match="transient database failure"): + store.flush() + with pytest.raises(RuntimeError, match="transient database failure"): + store.store({"type": "test.event", "data": {"value": "later"}}) + assert attempts == 1 + + +def test_non_sensitive_artifact_content_remains_plaintext(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + + store.store( + { + "type": "artifact.update", + "data": { + "type": "citation_source", + "content": "https://example.com/source", + "url": "https://example.com/source", + }, + } + ) + + raw_event = _raw_event_data(db_url) + assert raw_event["data"]["content"] == "https://example.com/source" + + +def test_plaintext_historical_event_rows_still_read_in_encrypted_mode(monkeypatch, db_url): + _enable_static_key(monkeypatch) + store = EventStore(db_url, "job-1") + + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "historical plaintext", + }, + } + ) + + events = EventStore.get_events(db_url, "job-1") + assert events[0]["data"]["content"] == "historical plaintext" + + +@pytest.mark.asyncio +async def test_async_event_reads_decrypt_content(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "async secret report", + }, + } + ) + + events = await EventStore.get_events_async(db_url, "job-1") + + assert events[0]["data"]["content"] == "async secret report" + + +@pytest.mark.asyncio +async def test_async_event_read_decrypts_off_event_loop(monkeypatch, db_url): + loop_thread = threading.get_ident() + decrypt_thread = None + store = EventStore(db_url, "job-1") + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": { + crypto.ENCRYPTED_FIELD_MARKER: True, + crypto.ENCRYPTED_FIELD_VALUE: f"{crypto.ENVELOPE_PREFIX}fake", + }, + }, + } + ) + + def fake_decrypt_event_field(_job_id, _field_path, _stored_value): + nonlocal decrypt_thread + decrypt_thread = threading.get_ident() + return "decrypted report" + + monkeypatch.setattr(crypto, "decrypt_event_field", fake_decrypt_event_field) + + events = await EventStore.get_events_async(db_url, "job-1") + + assert events[0]["data"]["content"] == "decrypted report" + assert decrypt_thread is not None + assert decrypt_thread != loop_thread + + +@pytest.mark.asyncio +async def test_async_event_read_slow_decrypt_does_not_block_event_loop(monkeypatch, db_url): + store = EventStore(db_url, "job-1") + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": { + crypto.ENCRYPTED_FIELD_MARKER: True, + crypto.ENCRYPTED_FIELD_VALUE: f"{crypto.ENVELOPE_PREFIX}fake", + }, + }, + } + ) + + def slow_decrypt_event_field(_job_id, _field_path, _stored_value): + time.sleep(0.1) + return "decrypted report" + + monkeypatch.setattr(crypto, "decrypt_event_field", slow_decrypt_event_field) + ticks = 0 + + async def ticker(): + nonlocal ticks + while ticks < 3: + await asyncio.sleep(0.01) + ticks += 1 + + read_task = asyncio.create_task(EventStore.get_events_async(db_url, "job-1")) + ticker_task = asyncio.create_task(ticker()) + + events = await read_task + await ticker_task + + assert events[0]["data"]["content"] == "decrypted report" + assert ticks >= 3 + + +def test_concurrent_event_reads_decrypt_consistently(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + for index in range(10): + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": f"secret report {index}", + }, + } + ) + event_id = _latest_event_id(db_url) + + def read_events() -> tuple[list[str], str]: + events = EventStore.get_events(db_url, "job-1", 0, 100) + event = EventStore.get_event_by_id(db_url, event_id) + assert event is not None + return [item["data"]["content"] for item in events], event["data"]["content"] + + with ThreadPoolExecutor(max_workers=8) as executor: + futures = [executor.submit(read_events) for _ in range(32)] + for future in as_completed(futures): + contents, latest_content = future.result() + assert contents == [f"secret report {index}" for index in range(10)] + assert latest_content == "secret report 9" + + +def test_encrypted_event_wrong_key_raises_instead_of_returning_empty(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "secret report", + }, + } + ) + + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", _other_static_key()) + crypto.reset_content_encryption_manager_for_tests() + + with pytest.raises(crypto.ContentEncryptionInvalidData): + EventStore.get_events(db_url, "job-1") + + +def test_encrypted_event_by_id_wrong_key_raises_instead_of_returning_none(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "secret report", + }, + } + ) + event_id = _latest_event_id(db_url) + + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", _other_static_key()) + crypto.reset_content_encryption_manager_for_tests() + + with pytest.raises(crypto.ContentEncryptionInvalidData): + EventStore.get_event_by_id(db_url, event_id) + + +@pytest.mark.asyncio +async def test_async_encrypted_event_wrong_key_raises_instead_of_fallback(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "secret report", + }, + } + ) + + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", _other_static_key()) + crypto.reset_content_encryption_manager_for_tests() + + with pytest.raises(crypto.ContentEncryptionInvalidData): + await EventStore.get_events_async(db_url, "job-1") + + +@pytest.mark.asyncio +async def test_async_encrypted_event_by_id_wrong_key_raises_instead_of_returning_none(monkeypatch, db_url): + _enable_static_key(monkeypatch) + cipher = crypto.create_job_content_cipher("job-1") + store = EventStore(db_url, "job-1", content_cipher=cipher) + store.store( + { + "type": "artifact.update", + "data": { + "type": "output", + "content": "secret report", + }, + } + ) + event_id = _latest_event_id(db_url) + + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY", _other_static_key()) + crypto.reset_content_encryption_manager_for_tests() + + with pytest.raises(crypto.ContentEncryptionInvalidData): + await EventStore.get_event_by_id_async(db_url, event_id) diff --git a/frontends/aiq_api/tests/test_report_context.py b/frontends/aiq_api/tests/test_report_context.py index dfb99d2a4..cc712641f 100644 --- a/frontends/aiq_api/tests/test_report_context.py +++ b/frontends/aiq_api/tests/test_report_context.py @@ -7,12 +7,10 @@ from fastapi import HTTPException -def test_extract_report_from_job_output_prefers_job_output(): - from aiq_api.jobs.report_context import _extract_report_from_job_output +def test_extract_report_from_output_prefers_job_output(): + from aiq_api.jobs.report_context import _extract_report_from_output - job = type("Job", (), {"output": {"report": "# Stored report"}})() - - assert _extract_report_from_job_output(job) == "# Stored report" + assert _extract_report_from_output({"report": "# Stored report"}) == "# Stored report" def test_report_from_events_prefers_final_report(): @@ -184,6 +182,40 @@ async def _events(_db_url: str, _job_id: str, _after_id: int, _limit: int): assert [s.url for s in ctx.sources] == ["https://example.com/path"] +@pytest.mark.asyncio +async def test_resolve_report_context_decrypts_encrypted_parent_output(monkeypatch): + """Report follow-up reads the encrypted parent output instead of relying on event fallback.""" + import base64 + import json + from unittest.mock import AsyncMock + + from aiq_api.jobs import crypto + from aiq_api.jobs import report_context + + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "key") + monkeypatch.setenv( + "AIQ_CONTENT_ENCRYPTION_KEY", + base64.urlsafe_b64encode(b"k" * crypto.DEK_BYTES).decode(), + ) + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY_ID", "test-key") + crypto.reset_content_encryption_manager_for_tests() + + report = "# Encrypted report\n\n## Sources\n\n[1] https://example.com/encrypted\n" + stored_output = crypto.create_job_content_cipher("job-1").encrypt_output_json(json.dumps({"report": report})) + job = type("Job", (), {"output": stored_output})() + get_events = AsyncMock(return_value=[]) + monkeypatch.setattr(report_context.EventStore, "get_events_async", get_events) + + try: + context = await report_context.resolve_report_context(job, "sqlite:///unused.db", "job-1") + finally: + crypto.reset_content_encryption_manager_for_tests() + + assert context.report_markdown == report.strip() + assert [source.url for source in context.sources] == ["https://example.com/encrypted"] + get_events.assert_not_awaited() + + @pytest.mark.asyncio async def test_resolve_report_context_scans_events_when_output_report_lacks_sources(monkeypatch): """Report in job.output but no inline ## Sources -> fall back to the event scan for enrichment.""" diff --git a/frontends/aiq_api/tests/test_report_edit.py b/frontends/aiq_api/tests/test_report_edit.py index 24fb7db92..edfb4a5ba 100644 --- a/frontends/aiq_api/tests/test_report_edit.py +++ b/frontends/aiq_api/tests/test_report_edit.py @@ -100,6 +100,131 @@ async def test_report_edit_authorizes_parent_and_submits_internal_child(report_e } +@pytest.mark.asyncio +async def test_report_edit_decrypts_parent_output_before_submitting_child(report_edit_app, monkeypatch): + """Encrypted parent reports remain available to the report-follow-up workflow.""" + import base64 + + from aiq_api.jobs import crypto + + app, parent_job, _authorize_job_access, submit_agent_job, _principal, _job_store = report_edit_app + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "key") + monkeypatch.setenv( + "AIQ_CONTENT_ENCRYPTION_KEY", + base64.urlsafe_b64encode(b"k" * crypto.DEK_BYTES).decode(), + ) + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY_ID", "test-key") + crypto.reset_content_encryption_manager_for_tests() + parent_report = "# Encrypted parent\n\n## Sources\n\n[1] https://example.com/source\n" + parent_job.output = crypto.create_job_content_cipher("parent-job-1").encrypt_output_json( + json.dumps({"report": parent_report}) + ) + + try: + with TestClient(app) as client: + response = client.post( + "/v1/jobs/async/job/parent-job-1/report/edit", + json={"input": "Shorten it."}, + ) + finally: + crypto.reset_content_encryption_manager_for_tests() + + assert response.status_code == 200 + assert submit_agent_job.await_args.kwargs["initial_files"]["/shared/original_report.md"] == parent_report.strip() + + +@pytest.mark.asyncio +async def test_report_edit_rejects_plaintext_parent_in_encrypted_mode(report_edit_app, monkeypatch): + """Report follow-up fails closed when encrypted mode encounters plaintext parent output.""" + import base64 + + from aiq_api.jobs import crypto + + app, parent_job, _authorize_job_access, submit_agent_job, _principal, _job_store = report_edit_app + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION", "key") + monkeypatch.setenv( + "AIQ_CONTENT_ENCRYPTION_KEY", + base64.urlsafe_b64encode(b"k" * crypto.DEK_BYTES).decode(), + ) + monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_KEY_ID", "test-key") + crypto.reset_content_encryption_manager_for_tests() + parent_job.output = {"report": "# Plaintext parent"} + + try: + with TestClient(app) as client: + response = client.post( + "/v1/jobs/async/job/parent-job-1/report/edit", + json={"input": "Shorten it."}, + ) + finally: + crypto.reset_content_encryption_manager_for_tests() + + assert response.status_code == 500 + assert response.json()["detail"] == "Parent report data is invalid" + submit_agent_job.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_report_edit_returns_503_when_parent_decrypt_is_unavailable(report_edit_app, monkeypatch): + from aiq_api.jobs import crypto + from aiq_api.jobs import report_context + + app, _parent_job, _authorize_job_access, submit_agent_job, _principal, _job_store = report_edit_app + monkeypatch.setattr( + report_context, + "read_job_output_async", + AsyncMock(side_effect=crypto.ContentEncryptionUnavailable("vault unavailable")), + ) + + with TestClient(app) as client: + response = client.post( + "/v1/jobs/async/job/parent-job-1/report/edit", + json={"input": "Shorten it."}, + ) + + assert response.status_code == 503 + assert response.json()["detail"] == "Content encryption is unavailable" + submit_agent_job.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_report_edit_returns_503_when_submit_encryption_is_unavailable(report_edit_app, caplog): + from aiq_api.jobs.crypto import ContentEncryptionUnavailable + + app, _parent_job, _authorize_job_access, submit_agent_job, _principal, _job_store = report_edit_app + submit_agent_job.side_effect = ContentEncryptionUnavailable("vault unavailable") + + with TestClient(app) as client: + response = client.post( + "/v1/jobs/async/job/parent-job-1/report/edit", + json={"input": "Shorten it."}, + ) + + assert response.status_code == 503 + assert response.json()["detail"] == "Content encryption is not ready" + assert "vault unavailable" not in caplog.text + submit_agent_job.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_report_edit_returns_500_when_submit_encryption_config_is_invalid(report_edit_app, caplog): + from aiq_api.jobs.crypto import ContentEncryptionConfigError + + app, _parent_job, _authorize_job_access, submit_agent_job, _principal, _job_store = report_edit_app + submit_agent_job.side_effect = ContentEncryptionConfigError("sensitive configuration detail") + + with TestClient(app) as client: + response = client.post( + "/v1/jobs/async/job/parent-job-1/report/edit", + json={"input": "Shorten it."}, + ) + + assert response.status_code == 500 + assert response.json()["detail"] == "Content encryption configuration is invalid" + assert "sensitive configuration detail" not in caplog.text + submit_agent_job.assert_awaited_once() + + @pytest.mark.asyncio async def test_report_edit_rejects_incomplete_parent(report_edit_app): app, parent_job, _authorize_job_access, submit_agent_job, _principal, _job_store = report_edit_app diff --git a/frontends/aiq_api/tests/test_submit_owner_user_id.py b/frontends/aiq_api/tests/test_submit_owner_user_id.py index 1890892c9..840b9294b 100644 --- a/frontends/aiq_api/tests/test_submit_owner_user_id.py +++ b/frontends/aiq_api/tests/test_submit_owner_user_id.py @@ -59,9 +59,10 @@ def test_submit_forwards_owner_user_id_as_last_arg(patched): 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. + # data_sources, auth_token, encryption policy, initial_files, output_metadata, owner_user_id. assert job_args[-1] == principal_user_id(principal) == "jwt:user-1" - assert job_args[-4] == "token-1" + assert job_args[-5] == "token-1" + assert job_args[-4].mode == "off" def test_context_user_id_binding_mechanism(): diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index e1d6da7d6..6f77d4476 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -421,6 +421,7 @@ async def test_submit_agent_job_passes_data_sources(self): { "NAT_DASK_SCHEDULER_ADDRESS": "tcp://localhost:8786", "NAT_JOB_STORE_DB_URL": "sqlite:///./test.db", + "AIQ_CONTENT_ENCRYPTION": "off", }, ): with patch("nat.front_ends.fastapi.async_jobs.job_store.JobStore", return_value=mock_job_store): @@ -436,7 +437,10 @@ async def test_submit_agent_job_passes_data_sources(self): assert result == "test-job-id" mock_job_store.submit_job.assert_called_once() job_args = mock_job_store.submit_job.call_args.kwargs["job_args"] - assert ["web_search"] in job_args + # Trailing worker args: available_documents, data_sources, auth_token, + # encryption policy, initial_files, output_metadata, principal_user_id. + assert job_args[-6] == ["web_search"] + assert job_args[-4].mode == "off" @pytest.mark.asyncio async def test_submit_agent_job_passes_initial_files_and_output_metadata(self): @@ -454,6 +458,7 @@ async def test_submit_agent_job_passes_initial_files_and_output_metadata(self): { "NAT_DASK_SCHEDULER_ADDRESS": "tcp://localhost:8786", "NAT_JOB_STORE_DB_URL": "sqlite:///./test.db", + "AIQ_CONTENT_ENCRYPTION": "off", }, ): with patch("nat.front_ends.fastapi.async_jobs.job_store.JobStore", return_value=mock_job_store): @@ -469,7 +474,8 @@ 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"] - # Trailing worker args: data_sources, auth_token, initial_files, output_metadata, principal_user_id + # Encryption policy precedes the upstream report-context arguments. + assert job_args[-4].mode == "off" assert job_args[-3] == initial_files assert job_args[-2] == output_metadata @@ -594,6 +600,299 @@ async def test_submit_rolls_back_when_job_access_persistence_fails(self): rollback_job_submission.assert_called_once_with("test-job-id", "sqlite:///./test.db") +class TestRunAgentJobEncryption: + """Tests for async worker encryption preflight behavior.""" + + @pytest.mark.asyncio + async def test_encryption_preflight_failure_marks_failure_before_running(self): + from aiq_api.jobs.crypto import ContentEncryptionConfig + from aiq_api.jobs.crypto import ContentEncryptionUnavailable + from aiq_api.jobs.runner import run_agent_job + from nat.front_ends.fastapi.async_jobs.job_store import JobStatus + + mock_job_store = MagicMock() + mock_job_store.update_status = AsyncMock() + + with patch("nat.front_ends.fastapi.async_jobs.job_store.JobStore", return_value=mock_job_store): + with patch( + "aiq_api.jobs.crypto.create_job_content_cipher", + side_effect=ContentEncryptionUnavailable("vault down"), + ): + await run_agent_job( + False, + 20, + "tcp://localhost:8786", + "sqlite:///./test.db", + "config.yml", + "job-1", + "input", + "aiq_agent.agents.deep_researcher.agent.DeepResearcherAgent", + "deep_research_agent", + content_encryption_policy=ContentEncryptionConfig(mode="off").policy_identity, + ) + + statuses = [call.args[1] for call in mock_job_store.update_status.await_args_list] + assert statuses == [JobStatus.FAILURE] + mock_job_store.update_status.assert_awaited_once_with( + "job-1", + JobStatus.FAILURE, + error="content encryption unavailable", + ) + + @pytest.mark.asyncio + async def test_worker_rejects_submission_policy_mismatch_before_running(self): + import base64 + + from aiq_api.jobs import crypto + from aiq_api.jobs.runner import run_agent_job + from nat.front_ends.fastapi.async_jobs.job_store import JobStatus + + with patch.dict( + "os.environ", + { + "AIQ_CONTENT_ENCRYPTION": "key", + "AIQ_CONTENT_ENCRYPTION_KEY": base64.urlsafe_b64encode(b"a" * crypto.DEK_BYTES).decode(), + "AIQ_CONTENT_ENCRYPTION_KEY_ID": "api-key", + }, + ): + crypto.reset_content_encryption_manager_for_tests() + api_policy = crypto.get_content_encryption_policy_identity() + mock_job_store = MagicMock() + mock_job_store.update_status = AsyncMock() + + with patch.dict("os.environ", {"AIQ_CONTENT_ENCRYPTION": "off"}): + crypto.reset_content_encryption_manager_for_tests() + with patch("nat.front_ends.fastapi.async_jobs.job_store.JobStore", return_value=mock_job_store): + with patch("aiq_api.jobs.crypto.create_job_content_cipher") as create_job_content_cipher: + await run_agent_job( + False, + 20, + "tcp://localhost:8786", + "sqlite:///./test.db", + "config.yml", + "job-1", + "input", + "aiq_agent.agents.deep_researcher.agent.DeepResearcherAgent", + "deep_research_agent", + content_encryption_policy=api_policy, + ) + crypto.reset_content_encryption_manager_for_tests() + + create_job_content_cipher.assert_not_called() + mock_job_store.update_status.assert_awaited_once_with( + "job-1", + JobStatus.FAILURE, + error="content encryption policy mismatch", + ) + + @pytest.mark.asyncio + async def test_final_output_encryption_failure_marks_failure_without_plaintext_write(self, tmp_path): + from types import SimpleNamespace + + from aiq_api.jobs.crypto import ContentEncryptionConfig + from aiq_api.jobs.crypto import ContentEncryptionUnavailable + from aiq_api.jobs.runner import run_agent_job + from nat.front_ends.fastapi.async_jobs.job_store import JobStatus + + class AsyncContext: + def __init__(self, value=None): + self.value = value + + async def __aenter__(self): + return self.value + + async def __aexit__(self, exc_type, exc, tb): + return False + + class FakeWorkflowBuilder: + _telemetry_exporters = {} + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + def get_function_config(self, _name): + return SimpleNamespace(tools=[], exclude_tools=[], verbose=False) + + async def get_tools(self, *, tool_names, wrapper_type): # noqa: ARG002 - mirrors NAT API + return [] + + class FakeExporterManager: + def start(self, *, context_state): + return AsyncContext() + + mock_job_store = MagicMock() + mock_job_store.update_status = AsyncMock() + update_job_output = AsyncMock(side_effect=ContentEncryptionUnavailable("encrypt failed")) + db_url = f"sqlite:///{tmp_path / 'test.db'}" + + config = SimpleNamespace(workflow=None, functions={}, middleware={}) + with patch("nat.front_ends.fastapi.async_jobs.job_store.JobStore", return_value=mock_job_store): + with patch("nat.runtime.loader.load_config", return_value=config): + with patch( + "nat.builder.workflow_builder.WorkflowBuilder.from_config", + return_value=FakeWorkflowBuilder(), + ): + with patch( + "nat.observability.exporter_manager.ExporterManager.from_exporters", + return_value=FakeExporterManager(), + ): + with patch("aiq_api.jobs.runner._load_agent_class", return_value=object): + with patch( + "aiq_api.jobs.runner._create_llm_provider", + AsyncMock(return_value=(object(), object())), + ): + with patch("aiq_api.jobs.runner._create_agent_instance", return_value=object()): + with patch( + "aiq_api.jobs.runner._run_agent", + AsyncMock(return_value="secret report"), + ): + with patch("aiq_api.jobs.crypto.update_job_output", update_job_output): + await run_agent_job( + False, + 20, + "tcp://localhost:8786", + db_url, + "config.yml", + "job-1", + "input", + "aiq_agent.agents.deep_researcher.agent.DeepResearcherAgent", + "deep_research_agent", + content_encryption_policy=ContentEncryptionConfig( + mode="off" + ).policy_identity, + output_metadata={ + "parent_job_id": "parent-job", + "interaction_action": "edit", + "report": "must not win", + }, + ) + + statuses = [call.args[1] for call in mock_job_store.update_status.await_args_list] + assert statuses == [JobStatus.RUNNING, JobStatus.FAILURE] + assert all("output" not in call.kwargs for call in mock_job_store.update_status.await_args_list) + update_job_output.assert_awaited_once() + assert update_job_output.await_args.kwargs["output"] == { + "parent_job_id": "parent-job", + "interaction_action": "edit", + "report": "secret report", + } + + @pytest.mark.asyncio + async def test_encrypted_event_flush_failure_marks_failure_before_success(self, tmp_path): + import base64 + from types import SimpleNamespace + + from aiq_api.jobs import crypto + from aiq_api.jobs.runner import run_agent_job + from nat.front_ends.fastapi.async_jobs.job_store import JobStatus + + class AsyncContext: + def __init__(self, value=None): + self.value = value + + async def __aenter__(self): + return self.value + + async def __aexit__(self, exc_type, exc, tb): + return False + + class FakeWorkflowBuilder: + _telemetry_exporters = {} + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + def get_function_config(self, _name): + return SimpleNamespace(tools=[], exclude_tools=[], verbose=False) + + async def get_tools(self, *, tool_names, wrapper_type): # noqa: ARG002 - mirrors NAT API + return [] + + class FakeExporterManager: + def start(self, *, context_state): + return AsyncContext() + + async def run_agent_with_event(*, event_store, **_kwargs): + event_store.store( + { + "type": "artifact.update", + "data": {"type": "output", "content": "secret report"}, + } + ) + return "secret report" + + mock_job_store = MagicMock() + mock_job_store.update_status = AsyncMock() + update_job_output = AsyncMock() + db_url = f"sqlite:///{tmp_path / 'test.db'}" + encryption_env = { + "AIQ_CONTENT_ENCRYPTION": "key", + "AIQ_CONTENT_ENCRYPTION_KEY": base64.urlsafe_b64encode(b"a" * crypto.DEK_BYTES).decode(), + "AIQ_CONTENT_ENCRYPTION_KEY_ID": "test-key", + } + + with patch.dict("os.environ", encryption_env): + crypto.reset_content_encryption_manager_for_tests() + policy = crypto.get_content_encryption_policy_identity() + config = SimpleNamespace(workflow=None, functions={}, middleware={}) + with patch("nat.front_ends.fastapi.async_jobs.job_store.JobStore", return_value=mock_job_store): + with patch("nat.runtime.loader.load_config", return_value=config): + with patch( + "nat.builder.workflow_builder.WorkflowBuilder.from_config", + return_value=FakeWorkflowBuilder(), + ): + with patch( + "nat.observability.exporter_manager.ExporterManager.from_exporters", + return_value=FakeExporterManager(), + ): + with patch("aiq_api.jobs.runner._load_agent_class", return_value=object): + with patch( + "aiq_api.jobs.runner._create_llm_provider", + AsyncMock(return_value=(object(), object())), + ): + with patch("aiq_api.jobs.runner._create_agent_instance", return_value=object()): + with patch( + "aiq_api.jobs.runner._run_agent", + side_effect=run_agent_with_event, + ): + with patch( + "aiq_api.jobs.event_store.EventStore.store_batch", + side_effect=RuntimeError("transient database failure"), + ): + with patch( + "aiq_api.jobs.crypto.update_job_output", + update_job_output, + ): + await run_agent_job( + False, + 20, + "tcp://localhost:8786", + db_url, + "config.yml", + "job-1", + "input", + "aiq_agent.agents.deep_researcher.agent.DeepResearcherAgent", + "deep_research_agent", + content_encryption_policy=policy, + ) + crypto.reset_content_encryption_manager_for_tests() + + statuses = [call.args[1] for call in mock_job_store.update_status.await_args_list] + assert statuses == [JobStatus.RUNNING, JobStatus.FAILURE] + mock_job_store.update_status.assert_awaited_with( + "job-1", + JobStatus.FAILURE, + error="transient database failure", + ) + update_job_output.assert_not_awaited() + + class TestEventStore: """Tests for the EventStore class.""" diff --git a/tests/deploy/test_helm_deployment_k8s.py b/tests/deploy/test_helm_deployment_k8s.py index 4a18f0a0e..bb67f164c 100644 --- a/tests/deploy/test_helm_deployment_k8s.py +++ b/tests/deploy/test_helm_deployment_k8s.py @@ -115,3 +115,17 @@ def test_default_chart_does_not_provision_or_configure_redis(): 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"}) + + +def test_backend_uses_separate_liveness_and_readiness_endpoints(): + manifests = render_chart() + + backend_deployment = next( + manifest + for manifest in manifests + if manifest.get("kind") == "Deployment" and manifest["metadata"]["name"] == "aiq-backend" + ) + backend_container = backend_deployment["spec"]["template"]["spec"]["containers"][0] + + assert backend_container["livenessProbe"]["httpGet"]["path"] == "/live" + assert backend_container["readinessProbe"]["httpGet"]["path"] == "/health" diff --git a/uv.lock b/uv.lock index 5a0264046..d7c53f9c0 100644 --- a/uv.lock +++ b/uv.lock @@ -342,6 +342,7 @@ dependencies = [ { name = "asyncpg" }, { name = "dask", extra = ["distributed"] }, { name = "fastapi" }, + { name = "hvac" }, { name = "langchain-core" }, { name = "pydantic" }, { name = "pyjwt" }, @@ -365,6 +366,7 @@ requires-dist = [ { name = "asyncpg", specifier = ">=0.29.0" }, { name = "dask", extras = ["distributed"], specifier = ">=2024.1.0" }, { name = "fastapi", specifier = ">=0.100.0" }, + { name = "hvac", specifier = ">=2.3.0,<3" }, { name = "langchain-core", specifier = ">=1.3.3,<2" }, { name = "psycopg", extras = ["binary"], marker = "extra == 'dev'", specifier = ">=3.0.0" }, { name = "psycopg", extras = ["binary"], marker = "extra == 'postgres'", specifier = ">=3.0.0" }, @@ -2141,6 +2143,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/49/79/621a7dbb80c70974f73a597275351ebe03ce5bc65cb5f8f4acb5859252bc/huggingface_hub-1.16.1-py3-none-any.whl", hash = "sha256:64340de934b9ce37857ef85a82de72f5629e8a270f9119eabb12bf495eb53c22", size = 668176, upload-time = "2026-05-21T18:39:58.596Z" }, ] +[[package]] +name = "hvac" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e2/57/b46c397fb3842cfb02a44609aa834c887f38dd75f290c2fc5a34da4b2fee/hvac-2.4.0.tar.gz", hash = "sha256:e0056ad9064e7923e874e6769015b032580b639e29246f5ab1044f7959c1c7e0", size = 332543, upload-time = "2025-10-30T12:57:47.512Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/55/33/71e45a6bd6875f44a26f99da31c63b6840123e88bedf2c0b1ce429b8be12/hvac-2.4.0-py3-none-any.whl", hash = "sha256:008db5efd8c2f77bd37d2368ea5f713edceae1c65f11fd608393179478649e0f", size = 155921, upload-time = "2025-10-30T12:57:46.253Z" }, +] + [[package]] name = "hyperframe" version = "6.1.0"