From 515f6081fe4e8bdd8be48121697af7d873911c8f Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Sun, 28 Jun 2026 15:49:05 -0400 Subject: [PATCH 01/10] feat(api): encrypt async job content artifacts Encrypt async final report output and selected artifact event content at rest, including static-key and Vault Transit modes, health/readiness checks, report read decryption, and coverage for worker, route, and event-store behavior. Signed-off-by: Tanner Leach --- deploy/Dockerfile | 1 + docs/source/deployment/content-encryption.md | 147 +++ docs/source/deployment/index.md | 2 + docs/source/index.md | 1 + frontends/aiq_api/pyproject.toml | 2 + frontends/aiq_api/src/aiq_api/jobs/crypto.py | 886 ++++++++++++++++++ .../aiq_api/src/aiq_api/jobs/event_store.py | 75 +- frontends/aiq_api/src/aiq_api/jobs/runner.py | 37 +- frontends/aiq_api/src/aiq_api/jobs/submit.py | 4 + frontends/aiq_api/src/aiq_api/routes/jobs.py | 76 +- .../aiq_api/tests/test_content_encryption.py | 316 +++++++ .../tests/test_content_encryption_routes.py | 265 ++++++ .../test_event_store_content_encryption.py | 185 ++++ tests/aiq_agent/jobs/test_runner.py | 118 +++ uv.lock | 14 + 15 files changed, 2116 insertions(+), 13 deletions(-) create mode 100644 docs/source/deployment/content-encryption.md create mode 100644 frontends/aiq_api/src/aiq_api/jobs/crypto.py create mode 100644 frontends/aiq_api/tests/test_content_encryption.py create mode 100644 frontends/aiq_api/tests/test_content_encryption_routes.py create mode 100644 frontends/aiq_api/tests/test_event_store_content_encryption.py diff --git a/deploy/Dockerfile b/deploy/Dockerfile index 0c5eb3916..57f314ec4 100644 --- a/deploy/Dockerfile +++ b/deploy/Dockerfile @@ -79,6 +79,7 @@ RUN uv pip install --no-deps -e . \ && uv pip install --no-deps -e ./sources/exa_web_search \ && uv pip install --no-deps -e "./sources/knowledge_layer[all]" \ && uv pip install --no-deps -e ./frontends/aiq_api \ + && uv pip install "hvac>=2.3.0,<3" \ && uv pip install "psycopg[binary]>=3.0.0" RUN /app/.venv/bin/python -c "import aiq_api; import knowledge_layer; print('✓ Base packages installed')" diff --git a/docs/source/deployment/content-encryption.md b/docs/source/deployment/content-encryption.md new file mode 100644 index 000000000..a0387a2f3 --- /dev/null +++ b/docs/source/deployment/content-encryption.md @@ -0,0 +1,147 @@ + + +# Async Final Report Encryption + +AI-Q can encrypt async job final reports before they are persisted in +`job_info.output`. This is application-level envelope encryption for the +AI-Q async jobs API. + +This feature is intentionally narrow in its first milestone. It protects only +the serialized final output payload returned by +`GET /v1/jobs/async/job/{job_id}/report`, such as `{"report": "..."}`. + +## 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`. + +Still plaintext: + +- `job_events.event_data`, including tool events, artifact updates, heartbeat + events, cancellation events, error events, and possible final-report + duplicates emitted through events. +- 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 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 events and checkpoints can contain equivalent research content, this +phase does not provide full database-level job-content confidentiality. + +## 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. 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 +`job_info.output:{job_id}`. + +## 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 metadata. If omitted, envelopes use +`static-key` as the key id. The first implementation supports one active static +key only; rotation requires jobs encrypted with the previous key to expire or a +future rewrap/backfill process. + +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. + +## Rollout Behavior + +The first implementation is forward-only: + +- New `job_info.output` writes are encrypted after enablement. +- Existing plaintext `job_info.output` rows are intentionally unreadable while + `AIQ_CONTENT_ENCRYPTION=key` or `vault`. +- No historical plaintext backfill is included. +- No rewrap tooling is included. + +Enable encryption only after operators accept that old plaintext final-report +rows cannot be read in encrypted modes until a future backfill exists. + +## Health and Failure Behavior + +`/health` includes encryption readiness. When encryption is configured but +unready, `/health` returns HTTP 503 and new async submissions are rejected with +HTTP 503. + +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. + +If final-report encryption or encrypted persistence fails after an agent has +completed, the job is marked `FAILURE`. The worker does not fall back to writing +plaintext output. + +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 050a2ade7..6876c2175 100644 --- a/docs/source/deployment/index.md +++ b/docs/source/deployment/index.md @@ -31,6 +31,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 Final Report Encryption](./content-encryption.md)** -- Configure encryption at rest for async job final reports stored in `job_info.output`, 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/index.md b/docs/source/index.md index 309d92dc5..d816e6fbe 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -97,6 +97,7 @@ Overview <./deployment/index.md> Docker Compose <./deployment/docker-compose.md> Docker Build System <./deployment/docker-build.md> Authentication <./deployment/authentication.md> +Async Final Report Encryption <./deployment/content-encryption.md> Observability <./deployment/observability.md> Production <./deployment/production.md> Kubernetes <./deployment/kubernetes.md> 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/jobs/crypto.py b/frontends/aiq_api/src/aiq_api/jobs/crypto.py new file mode 100644 index 000000000..d293dd038 --- /dev/null +++ b/frontends/aiq_api/src/aiq_api/jobs/crypto.py @@ -0,0 +1,886 @@ +# 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 base64 +import hashlib +import json +import logging +import os +import random +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" +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 + + +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.""" + + +@dataclass(frozen=True) +class ContentEncryptionConfig: + """Process-local content-encryption configuration parsed from env.""" + + mode: str + key_id: str | None = None + static_key: bytes | None = None + 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 = None + vault_secret_id: str | None = None + 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, + self.static_key, + self.vault_addr, + self.vault_namespace, + self.vault_transit_mount, + self.vault_transit_key, + self.vault_role_id, + self.vault_secret_id, + self.vault_timeout_seconds, + self.readiness_ttl_seconds, + self.dek_cache_ttl_seconds, + self.dek_cache_max_entries, + ) + + +@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 + + def to_health_dict(self) -> dict[str, Any]: + result: dict[str, Any] = {"mode": self.mode, "ready": self.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) + + def get(self, cache_key: str) -> bytes | None: + if self.ttl_seconds <= 0: + return None + 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 + 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(now) + + def _evict(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._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: + 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: + 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 + + try: + dek, _wrapped = self._vault().generate_data_key(operation=operation) + _zero_bytes(dek) + self._readiness = ContentEncryptionReadiness( + mode=self.config.mode, + ready=True, + checked_at=time.monotonic(), + ) + except ContentEncryptionUnavailable as exc: + self._readiness = ContentEncryptionReadiness( + mode=self.config.mode, + ready=False, + checked_at=time.monotonic(), + reason="vault_unavailable", + exception_type=exc.__class__.__name__, + ) + logger.warning( + "Content encryption readiness failed mode=%s operation=%s reason=%s exception=%s", + self.config.mode, + operation, + self._readiness.reason, + self._readiness.exception_type, + ) + return self._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: + 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._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: + return fn() + except Exception as exc: # hvac exposes several operational exception classes. + last_exc = exc + if _is_auth_failure(exc): + self._login(force=True) + if attempt + 1 < _VAULT_ATTEMPTS: + time.sleep(0.05 * (2**attempt) + random.uniform(0, 0.025)) + continue + + assert last_exc is not None + logger.warning( + "Vault transit operation failed mode=vault operation=%s exception=%s", + operation, + last_exc.__class__.__name__, + ) + raise ContentEncryptionUnavailable(f"vault_{operation}_failed") from last_exc + + def _authenticated_client(self) -> Any: + client = self._get_client() + if not client.is_authenticated(): + self._login(force=True) + return client + + def _get_client(self) -> Any: + 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: + 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: + 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 + + +def reset_content_encryption_manager_for_tests() -> None: + global _manager + global _manager_signature + _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() + 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") + + +def get_content_encryption_health() -> ContentEncryptionReadiness: + return get_content_encryption_manager().check_readiness(force=False, operation="health") + + +def require_content_encryption_ready_for_submission() -> None: + get_content_encryption_manager().require_ready(operation="submit") + + +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 + + +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 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 parsed <= 0: + raise ContentEncryptionConfigError(f"{name} must be a positive number") + return parsed + + +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: + return exc.__class__.__name__ in {"Forbidden", "Unauthorized"} + + +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..4dacf2179 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/event_store.py +++ b/frontends/aiq_api/src/aiq_api/jobs/event_store.py @@ -22,6 +22,7 @@ from __future__ import annotations import asyncio +import copy import logging import threading import time @@ -101,13 +102,69 @@ 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._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 +400,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: @@ -398,11 +456,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), } ) @@ -508,7 +567,7 @@ def get_events(cls, db_url: str, job_id: str, after_id: int = 0, limit: int = 10 events = [] for row in result: try: - event = json.loads(row[1]) + event = cls._prepare_event_for_return(job_id, json.loads(row[1])) event["_id"] = row[0] events.append(event) except json.JSONDecodeError: @@ -544,7 +603,7 @@ async def get_events_async(cls, db_url: str, job_id: str, after_id: int = 0, lim events = [] for row in result: try: - event = json.loads(row[1]) + event = cls._prepare_event_for_return(job_id, json.loads(row[1])) event["_id"] = row[0] events.append(event) except json.JSONDecodeError: @@ -580,13 +639,13 @@ 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: diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index 9c068e0de..d2f63734e 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -313,6 +313,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 logger.info( @@ -324,6 +325,20 @@ 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 create_job_content_cipher + + job_output_cipher = create_job_content_cipher(job_id) + except ContentEncryptionError as exc: + logger.warning( + "Job %s failed encryption readiness before running mode=encrypted exception=%s", + job_id, + exc.__class__.__name__, + ) + await job_store.update_status(job_id, JobStatus.FAILURE, error="content encryption unavailable") + return + await job_store.update_status(job_id, JobStatus.RUNNING) cancellation_monitor = CancellationMonitor( @@ -466,7 +481,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) @@ -524,7 +539,25 @@ 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) - await job_store.update_status(job_id, JobStatus.SUCCESS, output={"report": report}) + from .crypto import update_job_output + + if job_output_cipher is None: + raise RuntimeError("job output cipher was not initialized") + try: + await update_job_output( + job_store, + job_id, + JobStatus.SUCCESS, + output={"report": report}, + 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: diff --git a/frontends/aiq_api/src/aiq_api/jobs/submit.py b/frontends/aiq_api/src/aiq_api/jobs/submit.py index 2e876c222..4d27e34fb 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/submit.py +++ b/frontends/aiq_api/src/aiq_api/jobs/submit.py @@ -156,6 +156,8 @@ async def submit_agent_job( """ from nat.front_ends.fastapi.async_jobs.job_store import JobStore + from .crypto import require_content_encryption_ready_for_submission + # Get agent configuration from registry agent_config = get_agent_config(agent_type) @@ -200,6 +202,8 @@ async def submit_agent_job( if not scheduler_address: raise RuntimeError("Async job submission requires NAT_DASK_SCHEDULER_ADDRESS to be set") + require_content_encryption_ready_for_submission() + # Auto-capture auth token if not explicitly provided if auth_token is None: from aiq_agent.auth import get_auth_token diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index b26f9dedc..9e6451494 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -313,9 +313,18 @@ 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 + from ..jobs.crypto import read_job_output + from ..jobs.crypto import require_content_encryption_ready_for_submission + from ..jobs.crypto import validate_content_encryption_startup from ..jobs.event_store import EventStore from ..jobs.submit import submit_agent_job as submit_authorized_job + validate_content_encryption_startup() + if not get_all_sources(): logger.warning( "No data sources registered. Add a 'data_sources' function with " @@ -415,6 +424,27 @@ async def health_check(): return JSONResponse(status_code=503, content=result) + try: + encryption = get_content_encryption_health() + result["encryption"] = encryption.to_health_dict() + if encryption.mode != "off" and not encryption.ready: + result["status"] = "degraded" + from fastapi.responses import JSONResponse + + 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__, + } + from fastapi.responses import JSONResponse + + return JSONResponse(status_code=503, content=result) + return result @app.post( @@ -444,6 +474,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: + require_content_encryption_ready_for_submission() + 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, @@ -473,6 +518,18 @@ async def submit_job( data_sources=req.data_sources, auth_token=auth_token, ) + 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 RuntimeError as e: raise HTTPException(403, str(e)) except Exception as e: @@ -623,10 +680,23 @@ async def get_job_report(job_id: str) -> JobReportResponse: report = None if job.output: try: - output = json.loads(job.output) if isinstance(job.output, str) else job.output + output = read_job_output(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(output, dict): report = output.get("report") - except (json.JSONDecodeError, AttributeError): - pass return JobReportResponse(job_id=job_id, has_report=bool(report), report=report) 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..c57348776 --- /dev/null +++ b/frontends/aiq_api/tests/test_content_encryption.py @@ -0,0 +1,316 @@ +# 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 +import os +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +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 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_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_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_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 calls == 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 = 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() + + crypto.validate_content_encryption_startup() + crypto.get_content_encryption_health() + + assert calls == 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 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..b57cc15ca --- /dev/null +++ b/frontends/aiq_api/tests/test_content_encryption_routes.py @@ -0,0 +1,265 @@ +# 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.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") + + +@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: + 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 + + +@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") + + assert response.status_code == 503 + body = response.json() + assert body["encryption"]["mode"] == "vault" + assert body["encryption"]["ready"] is False + + +@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_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.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"})) + 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"} 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..5cad8c30f --- /dev/null +++ b/frontends/aiq_api/tests/test_event_store_content_encryption.py @@ -0,0 +1,185 @@ +# 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 + +import pytest +from sqlalchemy import text + +from aiq_api.jobs import crypto +from aiq_api.jobs.event_store import EventStore + + +def _static_key() -> str: + return base64.urlsafe_b64encode(bytes(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 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" + + +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" diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index 78721f56f..c2dd1934c 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -559,6 +559,124 @@ 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 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", + ) + + 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_final_output_encryption_failure_marks_failure_without_plaintext_write(self): + from types import SimpleNamespace + + 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")) + + 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=object()): + 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", + "sqlite:///./test.db", + "config.yml", + "job-1", + "input", + "aiq_agent.agents.deep_researcher.agent.DeepResearcherAgent", + "deep_research_agent", + ) + + 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"] == {"report": "secret report"} + + class TestEventStore: """Tests for the EventStore class.""" diff --git a/uv.lock b/uv.lock index 75724f38d..7ce8f8675 100644 --- a/uv.lock +++ b/uv.lock @@ -319,6 +319,7 @@ dependencies = [ { name = "asyncpg" }, { name = "dask", extra = ["distributed"] }, { name = "fastapi" }, + { name = "hvac" }, { name = "langchain-core" }, { name = "pydantic" }, { name = "pyjwt" }, @@ -342,6 +343,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" }, @@ -2074,6 +2076,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" From 7364481e1fc453f3451beddf08c4cc9d1c694672 Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Mon, 29 Jun 2026 17:07:28 -0400 Subject: [PATCH 02/10] docs(api): clarify opt-in encryption limitations Document the default-off behavior, exact encrypted field scope, Vault concurrency and retry model, forward-only rollout constraints, stable key identity requirement, and database-writer replay limitation. Signed-off-by: Tanner Leach --- docs/source/deployment/content-encryption.md | 111 ++++++++++++++----- docs/source/deployment/index.md | 2 +- docs/source/index.md | 2 +- 3 files changed, 86 insertions(+), 29 deletions(-) diff --git a/docs/source/deployment/content-encryption.md b/docs/source/deployment/content-encryption.md index a0387a2f3..aa5a49f31 100644 --- a/docs/source/deployment/content-encryption.md +++ b/docs/source/deployment/content-encryption.md @@ -3,15 +3,21 @@ SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# Async Final Report Encryption +# Async Job Content Encryption -AI-Q can encrypt async job final reports before they are persisted in -`job_info.output`. This is application-level envelope encryption for the -AI-Q async jobs API. +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. -This feature is intentionally narrow in its first milestone. It protects only -the serialized final output payload returned by -`GET /v1/jobs/async/job/{job_id}/report`, such as `{"report": "..."}`. +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 @@ -19,24 +25,34 @@ 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: -- `job_events.event_data`, including tool events, artifact updates, heartbeat - events, cancellation events, error events, and possible final-report - duplicates emitted through events. +- 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 rows written before encryption was enabled. +- 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 events and checkpoints can contain equivalent research content, this -phase does not provide full database-level job-content confidentiality. +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 @@ -48,10 +64,12 @@ Set `AIQ_CONTENT_ENCRYPTION` on every API and worker process. | `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. 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 -`job_info.output:{job_id}`. +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 @@ -64,10 +82,13 @@ AIQ_CONTENT_ENCRYPTION_KEY= AIQ_CONTENT_ENCRYPTION_KEY_ID= ``` -`AIQ_CONTENT_ENCRYPTION_KEY_ID` is optional metadata. If omitted, envelopes use -`static-key` as the key id. The first implementation supports one active static -key only; rotation requires jobs encrypted with the previous key to expire or a -future rewrap/backfill process. +`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. @@ -102,18 +123,54 @@ 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 first implementation is forward-only: +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. -Enable encryption only after operators accept that old plaintext final-report -rows cannot be read in encrypted modes until a future backfill exists. +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 @@ -125,9 +182,9 @@ 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. -If final-report encryption or encrypted persistence fails after an agent has -completed, the job is marked `FAILURE`. The worker does not fall back to writing -plaintext output. +If final-report encryption, artifact event-content encryption, or encrypted +persistence fails after an agent has completed, the job is marked `FAILURE`. +The worker does not fall back to writing plaintext output. Report reads fail closed: diff --git a/docs/source/deployment/index.md b/docs/source/deployment/index.md index 6876c2175..ac6e3b26c 100644 --- a/docs/source/deployment/index.md +++ b/docs/source/deployment/index.md @@ -31,7 +31,7 @@ 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 Final Report Encryption](./content-encryption.md)** -- Configure encryption at rest for async job final reports stored in `job_info.output`, including Vault Transit and static-key modes. +- **[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. diff --git a/docs/source/index.md b/docs/source/index.md index d816e6fbe..a312bfc79 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -97,7 +97,7 @@ Overview <./deployment/index.md> Docker Compose <./deployment/docker-compose.md> Docker Build System <./deployment/docker-build.md> Authentication <./deployment/authentication.md> -Async Final Report Encryption <./deployment/content-encryption.md> +Async Job Content Encryption <./deployment/content-encryption.md> Observability <./deployment/observability.md> Production <./deployment/production.md> Kubernetes <./deployment/kubernetes.md> From a0c58dc7c4bf8f8655f9e53e29736888767d7948 Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Mon, 29 Jun 2026 17:20:19 -0400 Subject: [PATCH 03/10] fix(api): harden async job content encryption Protect shared encryption state under concurrency, offload blocking Vault work from async routes, surface encrypted event failures, and harden Vault readiness and retry behavior. Add regression coverage for concurrency, async responsiveness, readiness, retry classification, explicit API and SSE failures, and default-off configuration. Signed-off-by: Tanner Leach --- frontends/aiq_api/src/aiq_api/jobs/crypto.py | 400 ++++++++++++----- .../aiq_api/src/aiq_api/jobs/event_store.py | 64 ++- frontends/aiq_api/src/aiq_api/jobs/submit.py | 8 +- frontends/aiq_api/src/aiq_api/routes/jobs.py | 73 +++- .../aiq_api/tests/test_content_encryption.py | 407 +++++++++++++++++- .../tests/test_content_encryption_routes.py | 184 ++++++++ .../test_event_store_content_encryption.py | 206 +++++++++ frontends/aiq_api/tests/test_job_access.py | 4 + 8 files changed, 1210 insertions(+), 136 deletions(-) diff --git a/frontends/aiq_api/src/aiq_api/jobs/crypto.py b/frontends/aiq_api/src/aiq_api/jobs/crypto.py index d293dd038..8ae4a0d63 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/crypto.py +++ b/frontends/aiq_api/src/aiq_api/jobs/crypto.py @@ -22,12 +22,14 @@ from __future__ import annotations +import asyncio import base64 import hashlib import json import logging import os import random +import threading import time from collections import OrderedDict from dataclasses import dataclass @@ -54,6 +56,24 @@ 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): @@ -136,9 +156,15 @@ class ContentEncryptionReadiness: 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: @@ -196,29 +222,36 @@ 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 - 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 + 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 - 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(now) + 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) @@ -231,6 +264,7 @@ class ContentEncryptionManager: 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, @@ -239,46 +273,100 @@ def __init__(self, config: ContentEncryptionConfig): self._vault_client: _VaultTransitClient | None = None def get_readiness(self) -> ContentEncryptionReadiness: - if self.config.mode == "off": - return ContentEncryptionReadiness(mode="off", ready=True) - return self._readiness + 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: - 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: + 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=operation) - _zero_bytes(dek) - self._readiness = ContentEncryptionReadiness( + 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: - self._readiness = ContentEncryptionReadiness( - mode=self.config.mode, - ready=False, - checked_at=time.monotonic(), - reason="vault_unavailable", - exception_type=exc.__class__.__name__, + return self._failed_vault_readiness( + operation=operation, + reason="vault_generate_unavailable", + exception=exc, + encrypt_ready=False, + decrypt_ready=False, ) - logger.warning( - "Content encryption readiness failed mode=%s operation=%s reason=%s exception=%s", - self.config.mode, - operation, - self._readiness.reason, - self._readiness.exception_type, - ) - return self._readiness + 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) @@ -423,9 +511,10 @@ def _unwrap_dek_with_static_key(self, envelope: dict[str, Any]) -> bytes: return dek def _vault(self) -> _VaultTransitClient: - if self._vault_client is None: - self._vault_client = _VaultTransitClient(self.config) - return self._vault_client + with self._lock: + if self._vault_client is None: + self._vault_client = _VaultTransitClient(self.config) + return self._vault_client class _VaultTransitClient: @@ -433,6 +522,7 @@ class _VaultTransitClient: 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]: @@ -484,71 +574,88 @@ def _with_retry(self, fn, *, operation: str) -> dict[str, Any]: last_exc: Exception | None = None for attempt in range(_VAULT_ATTEMPTS): try: - return fn() + with self._lock: + return fn() except Exception as exc: # hvac exposes several operational exception classes. + if isinstance(exc, ContentEncryptionConfigError): + raise last_exc = exc - if _is_auth_failure(exc): - self._login(force=True) + if not _should_retry_vault_failure(exc): + break if attempt + 1 < _VAULT_ATTEMPTS: - time.sleep(0.05 * (2**attempt) + random.uniform(0, 0.025)) + 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 exception=%s", + "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: - client = self._get_client() - if not client.is_authenticated(): - self._login(force=True) - return client + with self._lock: + client = self._get_client() + if not client.is_authenticated(): + self._login(force=True) + return client def _get_client(self) -> Any: - if self._client is not None: + 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 - 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: - 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 + 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: - if self._client is not None: + 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 - 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: @@ -559,13 +666,15 @@ def _required(value: str | None, name: str) -> str: _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 - _manager = None - _manager_signature = None + with _manager_lock: + _manager = None + _manager_signature = None def get_content_encryption_config() -> ContentEncryptionConfig: @@ -644,10 +753,11 @@ def get_content_encryption_manager() -> ContentEncryptionManager: global _manager global _manager_signature config = get_content_encryption_config() - if _manager is None or _manager_signature != config.signature: - _manager = ContentEncryptionManager(config) - _manager_signature = config.signature - return _manager + 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: @@ -661,14 +771,32 @@ def validate_content_encryption_startup() -> ContentEncryptionReadiness: 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 create_job_content_cipher(job_id: str) -> JobContentCipher: return get_content_encryption_manager().create_job_cipher(job_id) @@ -722,6 +850,12 @@ def read_job_output(job_id: str, stored_output: Any) -> Any: 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.""" @@ -878,7 +1012,81 @@ def _empty_to_none(value: str | None) -> str | None: def _is_auth_failure(exc: Exception) -> bool: - return exc.__class__.__name__ in {"Forbidden", "Unauthorized"} + 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: 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 4dacf2179..31e7fbe68 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/event_store.py +++ b/frontends/aiq_api/src/aiq_api/jobs/event_store.py @@ -28,6 +28,8 @@ import time from typing import Any +from .crypto import ContentEncryptionError + logger = logging.getLogger(__name__) @@ -534,6 +536,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]: """ @@ -548,8 +574,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: @@ -564,15 +588,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 = 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]) - 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 [] @@ -584,8 +602,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: @@ -600,15 +616,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 = 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]) - 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( @@ -650,7 +660,17 @@ def get_event_by_id(cls, db_url: str, event_id: int) -> dict | None: 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 diff --git a/frontends/aiq_api/src/aiq_api/jobs/submit.py b/frontends/aiq_api/src/aiq_api/jobs/submit.py index 4d27e34fb..7d9ac40d5 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/submit.py +++ b/frontends/aiq_api/src/aiq_api/jobs/submit.py @@ -120,6 +120,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, ) -> str: """ Submit an agent job to the Dask cluster. @@ -138,6 +139,8 @@ 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. Returns: The job ID. @@ -156,7 +159,7 @@ async def submit_agent_job( """ from nat.front_ends.fastapi.async_jobs.job_store import JobStore - from .crypto import require_content_encryption_ready_for_submission + from .crypto import require_content_encryption_ready_for_submission_async # Get agent configuration from registry agent_config = get_agent_config(agent_type) @@ -202,7 +205,8 @@ async def submit_agent_job( if not scheduler_address: raise RuntimeError("Async job submission requires NAT_DASK_SCHEDULER_ADDRESS to be set") - require_content_encryption_ready_for_submission() + if not skip_encryption_readiness_check: + await require_content_encryption_ready_for_submission_async() # Auto-capture auth token if not explicitly provided if auth_token is None: diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index 9e6451494..d34dbd6cb 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -316,14 +316,14 @@ async def register_job_routes(app: FastAPI, builder: WorkflowBuilder, worker: Fa from ..jobs.crypto import ContentEncryptionConfigError from ..jobs.crypto import ContentEncryptionInvalidData from ..jobs.crypto import ContentEncryptionUnavailable - from ..jobs.crypto import get_content_encryption_health - from ..jobs.crypto import read_job_output - from ..jobs.crypto import require_content_encryption_ready_for_submission - from ..jobs.crypto import validate_content_encryption_startup + 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.submit import submit_agent_job as submit_authorized_job - validate_content_encryption_startup() + await validate_content_encryption_startup_async() if not get_all_sources(): logger.warning( @@ -425,7 +425,7 @@ async def health_check(): return JSONResponse(status_code=503, content=result) try: - encryption = get_content_encryption_health() + encryption = await get_content_encryption_health_async() result["encryption"] = encryption.to_health_dict() if encryption.mode != "off" and not encryption.ready: result["status"] = "degraded" @@ -475,7 +475,7 @@ async def submit_job( # is also forwarded to submit_authorized_job(...) below for ownership recording. principal = require_verified_principal() try: - require_content_encryption_ready_for_submission() + await require_content_encryption_ready_for_submission_async() except ContentEncryptionUnavailable as e: logger.warning( "Rejected async job submission because content encryption is unready: %s", @@ -517,6 +517,7 @@ 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( @@ -656,7 +657,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, @@ -680,7 +696,7 @@ async def get_job_report(job_id: str) -> JobReportResponse: report = None if job.output: try: - output = read_job_output(job_id, job.output) + 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", @@ -1142,6 +1158,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: @@ -1178,6 +1195,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 @@ -1193,12 +1212,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): @@ -1222,6 +1249,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() @@ -1377,6 +1406,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"}) @@ -1406,6 +1451,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() @@ -1518,6 +1565,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_content_encryption.py b/frontends/aiq_api/tests/test_content_encryption.py index c57348776..f55e17748 100644 --- a/frontends/aiq_api/tests/test_content_encryption.py +++ b/frontends/aiq_api/tests/test_content_encryption.py @@ -15,13 +15,20 @@ 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 @@ -79,6 +86,34 @@ def _enable_static_key(monkeypatch, *, cache_ttl: str | None = None) -> None: 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_static_key_envelope_round_trip(monkeypatch): _enable_static_key(monkeypatch) @@ -129,6 +164,69 @@ def test_off_mode_preserves_current_behavior_and_does_not_decrypt_envelopes(monk 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) @@ -236,6 +334,163 @@ def test_dek_cache_evicts_lru_when_max_entries_is_exceeded(): 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") @@ -272,9 +527,141 @@ def generate_data_key(self, *, operation): 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") @@ -282,24 +669,28 @@ def test_vault_readiness_rechecks_when_cache_is_stale(monkeypatch): monkeypatch.setenv("VAULT_SECRET_ID", "secret-id") monkeypatch.setenv("AIQ_ENCRYPTION_TRANSIT_KEY", "reports") monkeypatch.setenv("AIQ_CONTENT_ENCRYPTION_READINESS_TTL_SECONDS", "0") - calls = 0 + calls = {"generate": 0, "unwrap": 0} + dek = b"d" * crypto.DEK_BYTES - class FailingVault: + class ReadyVault: def __init__(self, _config): pass def generate_data_key(self, *, operation): - nonlocal calls - calls += 1 - raise crypto.ContentEncryptionUnavailable("vault down") + calls["generate"] += 1 + return dek, crypto.WrappedDEK(wrap="vault", kid="transit/reports", wrapped_dek="vault:v1:dek") - monkeypatch.setattr(crypto, "_VaultTransitClient", FailingVault) + 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 == 2 + assert calls == {"generate": 2, "unwrap": 2} @pytest.mark.skipif(not _real_vault_env_present(), reason="real Vault Transit credentials are not configured") @@ -311,6 +702,8 @@ def test_real_vault_transit_round_trip(monkeypatch): 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 index b57cc15ca..f6c982756 100644 --- a/frontends/aiq_api/tests/test_content_encryption_routes.py +++ b/frontends/aiq_api/tests/test_content_encryption_routes.py @@ -36,6 +36,10 @@ 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 ( @@ -155,6 +159,39 @@ def generate_data_key(self, *, operation): 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 + + +@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 @@ -179,6 +216,33 @@ def generate_data_key(self, *, operation): 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_report_authorizes_before_decrypting(monkeypatch, tmp_path): _enable_static_key(monkeypatch) @@ -263,3 +327,123 @@ async def test_report_decrypts_encrypted_final_output(monkeypatch, tmp_path): assert response.status_code == 200 assert response.json() == {"job_id": "job-1", "has_report": True, "report": "secret"} + + +@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 index 5cad8c30f..711ef967d 100644 --- a/frontends/aiq_api/tests/test_event_store_content_encryption.py +++ b/frontends/aiq_api/tests/test_event_store_content_encryption.py @@ -15,8 +15,13 @@ 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 @@ -29,6 +34,10 @@ 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 ( @@ -73,6 +82,14 @@ def _raw_event_data(db_url: str) -> dict: 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") @@ -183,3 +200,192 @@ async def test_async_event_reads_decrypt_content(monkeypatch, db_url): 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_job_access.py b/frontends/aiq_api/tests/test_job_access.py index f1b379189..d971bda17 100644 --- a/frontends/aiq_api/tests/test_job_access.py +++ b/frontends/aiq_api/tests/test_job_access.py @@ -112,6 +112,10 @@ def test_cleanup_removes_expired_and_orphaned_access_rows(self, db_url): class TestAuthorizeJobAccess: + @pytest.fixture(autouse=True) + def require_auth(self, monkeypatch): + monkeypatch.setenv("REQUIRE_AUTH", "true") + def test_owner_can_access_job(self, db_url): _insert_job_info(db_url, "job-1") principal = Principal(type="jwt", sub="user-1") From 40e854bd07eca4031eda79d87ea8ac093753f3e6 Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Mon, 29 Jun 2026 21:35:05 -0400 Subject: [PATCH 04/10] fix(api): address content encryption review feedback Redact credential-bearing configuration fields and fingerprint manager signatures without retaining raw secrets. Reject non-finite timing values, clarify async submission response contracts, and isolate the encryption failure database test. Add regression coverage for credential handling, manager recreation, malformed envelopes, route submission mocks, and OpenAPI responses. Signed-off-by: Tanner Leach --- frontends/aiq_api/src/aiq_api/jobs/crypto.py | 29 +++-- frontends/aiq_api/src/aiq_api/routes/jobs.py | 3 +- .../aiq_api/tests/test_content_encryption.py | 116 ++++++++++++++++++ .../tests/test_content_encryption_routes.py | 27 ++++ tests/aiq_agent/jobs/test_runner.py | 5 +- 5 files changed, 169 insertions(+), 11 deletions(-) diff --git a/frontends/aiq_api/src/aiq_api/jobs/crypto.py b/frontends/aiq_api/src/aiq_api/jobs/crypto.py index 8ae4a0d63..37305aff2 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/crypto.py +++ b/frontends/aiq_api/src/aiq_api/jobs/crypto.py @@ -27,6 +27,7 @@ import hashlib import json import logging +import math import os import random import threading @@ -102,13 +103,13 @@ class ContentEncryptionConfig: mode: str key_id: str | None = None - static_key: bytes | 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 = None - vault_secret_id: 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 @@ -133,13 +134,13 @@ def signature(self) -> tuple[Any, ...]: return ( self.mode, self.key_id, - self.static_key, + _secret_fingerprint(self.static_key), self.vault_addr, self.vault_namespace, self.vault_transit_mount, self.vault_transit_key, - self.vault_role_id, - self.vault_secret_id, + _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, @@ -987,7 +988,7 @@ def _parse_non_negative_float(value: str | None, *, default: float, name: str) - parsed = float(value) except ValueError as exc: raise ContentEncryptionConfigError(f"{name} must be a non-negative number") from exc - if parsed < 0: + if not math.isfinite(parsed) or parsed < 0: raise ContentEncryptionConfigError(f"{name} must be a non-negative number") return parsed @@ -999,11 +1000,23 @@ def _parse_positive_float(value: str | None, *, default: float, name: str) -> fl parsed = float(value) except ValueError as exc: raise ContentEncryptionConfigError(f"{name} must be a positive number") from exc - if parsed <= 0: + 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 diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index d34dbd6cb..36035d224 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -458,7 +458,8 @@ async def health_check(): responses={ 400: {"description": "Unknown agent type or invalid request"}, 422: {"description": "One or more unknown or agent-unavailable data source IDs"}, - 503: {"description": "Dask scheduler not available"}, + 500: {"description": "Content encryption configuration is invalid or job submission failed"}, + 503: {"description": "Content encryption is not ready"}, }, ) async def submit_job( diff --git a/frontends/aiq_api/tests/test_content_encryption.py b/frontends/aiq_api/tests/test_content_encryption.py index f55e17748..5a119a218 100644 --- a/frontends/aiq_api/tests/test_content_encryption.py +++ b/frontends/aiq_api/tests/test_content_encryption.py @@ -114,6 +114,117 @@ def test_content_encryption_defaults_to_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_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) @@ -149,6 +260,11 @@ def test_tamper_fails(monkeypatch): 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) diff --git a/frontends/aiq_api/tests/test_content_encryption_routes.py b/frontends/aiq_api/tests/test_content_encryption_routes.py index f6c982756..6be69f74f 100644 --- a/frontends/aiq_api/tests/test_content_encryption_routes.py +++ b/frontends/aiq_api/tests/test_content_encryption_routes.py @@ -99,6 +99,7 @@ async def _no_op_reaper(*_args, **_kwargs): 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( @@ -243,6 +244,32 @@ def unwrap_dek(self, wrapped_dek, *, operation): 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 job submission failed" + assert responses["503"]["description"] == "Content encryption is not ready" + + @pytest.mark.asyncio async def test_report_authorizes_before_decrypting(monkeypatch, tmp_path): _enable_static_key(monkeypatch) diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index c2dd1934c..c6acd4fc0 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -597,7 +597,7 @@ async def test_encryption_preflight_failure_marks_failure_before_running(self): ) @pytest.mark.asyncio - async def test_final_output_encryption_failure_marks_failure_without_plaintext_write(self): + 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 ContentEncryptionUnavailable @@ -636,6 +636,7 @@ def start(self, *, context_state): 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'}" 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=object()): @@ -662,7 +663,7 @@ def start(self, *, context_state): False, 20, "tcp://localhost:8786", - "sqlite:///./test.db", + db_url, "config.yml", "job-1", "input", From 807ca2aebafeef51926a8669197a92654ed761ad Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Tue, 30 Jun 2026 10:08:58 -0400 Subject: [PATCH 05/10] fix(api): align async job submission error contract Document the submit endpoint's actual HTTP 500 failure cases and update the OpenAPI regression assertion to match. Signed-off-by: Tanner Leach --- frontends/aiq_api/src/aiq_api/routes/jobs.py | 6 +++++- frontends/aiq_api/tests/test_content_encryption_routes.py | 4 +++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index 36035d224..c9b1bdbf2 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -458,7 +458,11 @@ async def health_check(): responses={ 400: {"description": "Unknown agent type or invalid request"}, 422: {"description": "One or more unknown or agent-unavailable data source IDs"}, - 500: {"description": "Content encryption configuration is invalid or job submission failed"}, + 500: { + "description": ( + "Content encryption configuration is invalid or async job authorization persistence failed" + ) + }, 503: {"description": "Content encryption is not ready"}, }, ) diff --git a/frontends/aiq_api/tests/test_content_encryption_routes.py b/frontends/aiq_api/tests/test_content_encryption_routes.py index 6be69f74f..da22c58e0 100644 --- a/frontends/aiq_api/tests/test_content_encryption_routes.py +++ b/frontends/aiq_api/tests/test_content_encryption_routes.py @@ -266,7 +266,9 @@ async def test_submit_openapi_documents_encryption_failures(monkeypatch, tmp_pat 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 job submission failed" + assert responses["500"]["description"] == ( + "Content encryption configuration is invalid or async job authorization persistence failed" + ) assert responses["503"]["description"] == "Content encryption is not ready" From b6fd02bffc76049426faff46072a79d2fe886af0 Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Mon, 6 Jul 2026 11:41:59 -0400 Subject: [PATCH 06/10] fix(api): enforce worker encryption policy Propagate a non-secret encryption policy identity with async job submissions and fail workers before RUNNING when their local mode or key identity differs. Document the boundary and cover static-key, Vault, missing-policy, and API-to-worker mismatch behavior. Signed-off-by: Tanner Leach --- docs/source/deployment/content-encryption.md | 7 +++ frontends/aiq_api/src/aiq_api/jobs/crypto.py | 58 +++++++++++++++++++ frontends/aiq_api/src/aiq_api/jobs/runner.py | 22 ++++++- frontends/aiq_api/src/aiq_api/jobs/submit.py | 3 + .../aiq_api/tests/test_content_encryption.py | 46 +++++++++++++++ tests/aiq_agent/jobs/test_runner.py | 58 ++++++++++++++++++- 6 files changed, 190 insertions(+), 4 deletions(-) diff --git a/docs/source/deployment/content-encryption.md b/docs/source/deployment/content-encryption.md index aa5a49f31..9df3d1b32 100644 --- a/docs/source/deployment/content-encryption.md +++ b/docs/source/deployment/content-encryption.md @@ -182,6 +182,13 @@ 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`. The worker does not fall back to writing plaintext output. diff --git a/frontends/aiq_api/src/aiq_api/jobs/crypto.py b/frontends/aiq_api/src/aiq_api/jobs/crypto.py index 37305aff2..b6ca02679 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/crypto.py +++ b/frontends/aiq_api/src/aiq_api/jobs/crypto.py @@ -48,6 +48,7 @@ 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 @@ -97,6 +98,24 @@ 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.""" @@ -147,6 +166,32 @@ def signature(self) -> tuple[Any, ...]: 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: @@ -798,6 +843,19 @@ async def require_content_encryption_ready_for_submission_async() -> None: 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) diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index d2f63734e..7b9d03f6d 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -29,12 +29,16 @@ import importlib import logging import uuid +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__) @@ -247,6 +251,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, ): """ Dask task to run any registered agent with cancellation support and telemetry. @@ -278,6 +283,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. """ # Propagate auth token into the current async task's context so tools @@ -327,16 +334,24 @@ async def run_agent_job( 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 readiness before running mode=encrypted exception=%s", + "Job %s failed encryption policy/readiness before running exception=%s", job_id, exc.__class__.__name__, ) - await job_store.update_status(job_id, JobStatus.FAILURE, error="content encryption unavailable") + 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) @@ -843,6 +858,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, @@ -853,4 +870,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 7d9ac40d5..a517fb1e8 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/submit.py +++ b/frontends/aiq_api/src/aiq_api/jobs/submit.py @@ -159,6 +159,7 @@ 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 @@ -207,6 +208,7 @@ async def submit_agent_job( 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: @@ -242,6 +244,7 @@ async def submit_agent_job( available_documents, data_sources, auth_token, + content_encryption_policy, ], ) await loop.run_in_executor(None, create_job_access, resolved_job_id, principal, db_url) diff --git a/frontends/aiq_api/tests/test_content_encryption.py b/frontends/aiq_api/tests/test_content_encryption.py index 5a119a218..753824e73 100644 --- a/frontends/aiq_api/tests/test_content_encryption.py +++ b/frontends/aiq_api/tests/test_content_encryption.py @@ -157,6 +157,52 @@ def test_config_signature_changes_when_credentials_change(credential_name, repla 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") diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index c6acd4fc0..75366e699 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -420,6 +420,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): @@ -435,8 +436,9 @@ 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"] - # data_sources is second-to-last (auth_token is last) - assert job_args[-2] == ["web_search"] + # The policy identity is last so it crosses the same Dask boundary as the job. + assert job_args[-3] == ["web_search"] + assert job_args[-1].mode == "off" @pytest.mark.asyncio async def test_submit_with_custom_job_id(self): @@ -564,6 +566,7 @@ class TestRunAgentJobEncryption: @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 @@ -586,6 +589,7 @@ async def test_encryption_preflight_failure_marks_failure_before_running(self): "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] @@ -596,10 +600,57 @@ async def test_encryption_preflight_failure_marks_failure_before_running(self): 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 @@ -669,6 +720,9 @@ def start(self, *, context_state): "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] From 4b4f4f89320b83589b9993902ee9c9f483d200df Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Mon, 6 Jul 2026 14:20:06 -0400 Subject: [PATCH 07/10] fix(api): fail jobs on encrypted event write errors Propagate event persistence failures when content encryption is enabled and retain background batch flush errors for the runner to observe. Keep off-mode event writes best effort and add timer and job-state regression coverage. Signed-off-by: Tanner Leach --- docs/source/deployment/content-encryption.md | 5 +- .../aiq_api/src/aiq_api/jobs/event_store.py | 38 +++++- frontends/aiq_api/src/aiq_api/jobs/runner.py | 38 ++++-- .../test_event_store_content_encryption.py | 71 +++++++++++ tests/aiq_agent/jobs/test_runner.py | 111 ++++++++++++++++++ 5 files changed, 247 insertions(+), 16 deletions(-) diff --git a/docs/source/deployment/content-encryption.md b/docs/source/deployment/content-encryption.md index 9df3d1b32..b12c136fa 100644 --- a/docs/source/deployment/content-encryption.md +++ b/docs/source/deployment/content-encryption.md @@ -191,7 +191,10 @@ 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`. -The worker does not fall back to writing plaintext output. +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: 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 31e7fbe68..41e1f346e 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/event_store.py +++ b/frontends/aiq_api/src/aiq_api/jobs/event_store.py @@ -113,6 +113,7 @@ def __init__( 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() @@ -441,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]): """ @@ -498,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): @@ -810,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: @@ -819,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) + try: + self._store.store_batch(batch) + except Exception as exc: + if self._flush_error is None: + self._flush_error = exc + raise - def _flush(self): - """Flush from timer callback (acquires lock).""" + 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/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index 7b9d03f6d..99564ef39 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/frontends/aiq_api/src/aiq_api/jobs/runner.py @@ -588,14 +588,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__) @@ -605,22 +604,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() # Clean up job-scoped auth token @@ -630,6 +635,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 _create_agent_instance( agent_cls: type, llm_provider, diff --git a/frontends/aiq_api/tests/test_event_store_content_encryption.py b/frontends/aiq_api/tests/test_event_store_content_encryption.py index 711ef967d..d22f8bc5d 100644 --- a/frontends/aiq_api/tests/test_event_store_content_encryption.py +++ b/frontends/aiq_api/tests/test_event_store_content_encryption.py @@ -27,6 +27,7 @@ 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 @@ -144,6 +145,76 @@ def test_file_artifact_content_is_encrypted_in_batch(monkeypatch, db_url): 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") diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index 75366e699..0839bcc13 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -731,6 +731,117 @@ def start(self, *, context_state): update_job_output.assert_awaited_once() assert update_job_output.await_args.kwargs["output"] == {"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() + 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=object()): + 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.""" From 4a763c0660beb17176eaa9c377f12c40f5780f06 Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Mon, 6 Jul 2026 16:03:27 -0400 Subject: [PATCH 08/10] fix(api): serve encryption-aware health checks Replace NAT's generic health route with AI-Q readiness so runtime dispatch and OpenAPI expose database and encryption state. Preserve the healthy success response and cover assembled worker behavior for unavailable Vault, off mode, and static-key mode. Signed-off-by: Tanner Leach --- docs/source/integration/rest-api.md | 2 +- frontends/aiq_api/src/aiq_api/routes/jobs.py | 22 +++- .../tests/test_content_encryption_routes.py | 116 ++++++++++++++++++ 3 files changed, 138 insertions(+), 2 deletions(-) diff --git a/docs/source/integration/rest-api.md b/docs/source/integration/rest-api.md index 0b2de0f55..ec5fc3c8f 100644 --- a/docs/source/integration/rest-api.md +++ b/docs/source/integration/rest-api.md @@ -388,7 +388,7 @@ curl http://localhost:8000/health ```json { - "status": "ok", + "status": "healthy", "dask_available": true } ``` diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index c9b1bdbf2..b643d2749 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -40,6 +40,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 @@ -60,6 +61,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) + + class JobSubmitRequest(BaseModel): """Request to submit an async job.""" @@ -398,6 +414,10 @@ async def list_data_sources() -> list[DataSource]: ) await asyncio.get_running_loop().run_in_executor(None, ensure_job_access_table, db_url) + # NAT registers its generic liveness route first. Replace it so runtime + # dispatch and OpenAPI both use AI-Q's DB and encryption readiness checks. + _remove_existing_health_routes(app) + @app.get("/health", tags=["health"], summary="Health check") async def health_check(): """Health check endpoint that validates DB connectivity.""" @@ -405,7 +425,7 @@ async def health_check(): from ..jobs.event_store import EventStore - result = {"status": "ok", "dask_available": dask_available, "db": "ok"} + result = {"status": "healthy", "dask_available": dask_available, "db": "ok"} # Check DB connectivity using any cached async engine try: diff --git a/frontends/aiq_api/tests/test_content_encryption_routes.py b/frontends/aiq_api/tests/test_content_encryption_routes.py index da22c58e0..4fdd4c6a0 100644 --- a/frontends/aiq_api/tests/test_content_encryption_routes.py +++ b/frontends/aiq_api/tests/test_content_encryption_routes.py @@ -25,6 +25,7 @@ import pytest from fastapi import FastAPI +from fastapi.routing import APIRoute from fastapi.testclient import TestClient from aiq_agent.auth import Principal @@ -139,6 +140,75 @@ async def _no_op_reaper(*_args, **_kwargs): return app +def _build_assembled_worker_app(monkeypatch, tmp_path) -> 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=f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}") + worker = plugin.AIQAPIWorker(Config(general=GeneralConfig(front_end=front_end))) + worker._dask_available = True + worker._job_store = object() + 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 + ] + + @pytest.mark.asyncio async def test_health_returns_503_when_vault_readiness_failed(monkeypatch, tmp_path): _enable_vault(monkeypatch) @@ -164,6 +234,52 @@ def generate_data_key(self, *, operation): assert body["encryption"]["decrypt_ready"] is False +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") + 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 == 503 + body = response.json() + assert body["status"] == "degraded" + assert body["encryption"]["mode"] == "vault" + assert body["encryption"]["ready"] is False + + +@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"]["/health"]) == ["get"] + + @pytest.mark.asyncio async def test_health_returns_503_when_vault_decrypt_readiness_failed(monkeypatch, tmp_path): _enable_vault(monkeypatch) From 03347e15288ffc39c3bbf91ce07e3a7b62422794 Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Mon, 6 Jul 2026 21:20:53 -0400 Subject: [PATCH 09/10] fix(api): separate liveness from dependency readiness Add a process-only /live endpoint while retaining dependency checks on /health, and wire Helm probes to the correct paths. Map report-edit submission encryption outages to 503 and add regression coverage for health routing, auth exemptions, chart probes, and errors. Signed-off-by: Tanner Leach --- deploy/helm/README.md | 5 + deploy/helm/deployment-k8s/README.md | 5 + deploy/helm/deployment-k8s/values.yaml | 2 +- docs/source/deployment/content-encryption.md | 8 +- docs/source/deployment/kubernetes.md | 5 + docs/source/deployment/production.md | 11 +- docs/source/integration/rest-api.md | 28 ++- .../aiq_api/src/aiq_api/auth/middleware.py | 3 +- frontends/aiq_api/src/aiq_api/routes/jobs.py | 164 +++++++++++------- frontends/aiq_api/tests/test_auth.py | 14 ++ .../tests/test_content_encryption_routes.py | 47 ++++- frontends/aiq_api/tests/test_report_edit.py | 38 ++++ tests/deploy/test_helm_deployment_k8s.py | 14 ++ 13 files changed, 266 insertions(+), 78 deletions(-) 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 index b12c136fa..f99f59557 100644 --- a/docs/source/deployment/content-encryption.md +++ b/docs/source/deployment/content-encryption.md @@ -174,9 +174,11 @@ encryption configuration to read those jobs. ## Health and Failure Behavior -`/health` includes encryption readiness. When encryption is configured but -unready, `/health` returns HTTP 503 and new async submissions are rejected with -HTTP 503. +`/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 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/integration/rest-api.md b/docs/source/integration/rest-api.md index b8baa05ef..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 @@ -462,7 +479,12 @@ curl http://localhost:8000/health ```json { "status": "healthy", - "dask_available": true + "dask_available": true, + "db": "ok", + "encryption": { + "mode": "off", + "ready": true + } } ``` 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/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index 441d5cbeb..684699aac 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -522,6 +522,90 @@ async def register_job_routes(app: FastAPI, builder: WorkflowBuilder, worker: Fa 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 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" + 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 " @@ -571,9 +655,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" @@ -581,19 +662,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, @@ -604,59 +676,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) - # NAT registers its generic liveness route first. Replace it so runtime - # dispatch and OpenAPI both use AI-Q's DB and encryption readiness checks. - _remove_existing_health_routes(app) - - @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": "healthy", "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) - - 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" - from fastapi.responses import JSONResponse - - 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__, - } - from fastapi.responses import JSONResponse - - return JSONResponse(status_code=503, content=result) - - return result - @app.post( "/v1/jobs/async/submit", response_model=JobStatusResponse, @@ -861,6 +880,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: 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_routes.py b/frontends/aiq_api/tests/test_content_encryption_routes.py index 8d320201a..3e0276556 100644 --- a/frontends/aiq_api/tests/test_content_encryption_routes.py +++ b/frontends/aiq_api/tests/test_content_encryption_routes.py @@ -140,7 +140,7 @@ async def _no_op_reaper(*_args, **_kwargs): return app -def _build_assembled_worker_app(monkeypatch, tmp_path) -> FastAPI: +def _build_assembled_worker_app(monkeypatch, tmp_path, *, dask_available: bool = True) -> 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 @@ -194,8 +194,8 @@ async def __aexit__(self, exc_type, exc, tb): front_end = plugin.AIQAPIConfig(db_url=f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}") worker = plugin.AIQAPIWorker(Config(general=GeneralConfig(front_end=front_end))) - worker._dask_available = True - worker._job_store = object() + 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() @@ -209,6 +209,14 @@ def _get_health_routes(app: FastAPI) -> list[APIRoute]: ] +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) @@ -225,6 +233,7 @@ def generate_data_key(self, *, operation): with TestClient(app) as client: response = client.get("/health") + liveness_response = client.get("/live") assert response.status_code == 503 body = response.json() @@ -232,6 +241,8 @@ def generate_data_key(self, *, operation): 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): @@ -249,10 +260,16 @@ def generate_data_key(self, *, operation): 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" @@ -260,6 +277,29 @@ def generate_data_key(self, *, operation): 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": @@ -277,6 +317,7 @@ def test_assembled_worker_health_route_reports_ready_encryption(monkeypatch, tmp 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"] diff --git a/frontends/aiq_api/tests/test_report_edit.py b/frontends/aiq_api/tests/test_report_edit.py index 4092303e9..edfb4a5ba 100644 --- a/frontends/aiq_api/tests/test_report_edit.py +++ b/frontends/aiq_api/tests/test_report_edit.py @@ -187,6 +187,44 @@ async def test_report_edit_returns_503_when_parent_decrypt_is_unavailable(report 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/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" From 36d6c2361819d29c701d133bbcb906b2b657cba5 Mon Sep 17 00:00:00 2001 From: Tanner Leach Date: Tue, 7 Jul 2026 09:16:24 -0400 Subject: [PATCH 10/10] fix(api): probe database on every readiness check The /health readiness check only ran the SELECT 1 ping when an async engine happened to already be cached, and reported "db": "no_engine" with HTTP 200 otherwise. On a fresh process the async-engine cache is empty, so /health returned healthy without touching the database. Since Helm now gates readiness on /health, a fresh pod could pass readiness straight through a database outage. Obtain (or create) the engine for the configured db_url via EventStore._get_or_create_async_engine and always run the bounded ping; an absent cached engine is no longer treated as healthy. Add an assembled-app regression test asserting /health returns 503 when the database is unavailable and the async-engine cache is empty, while /live stays 200. Signed-off-by: Tanner Leach --- frontends/aiq_api/src/aiq_api/routes/jobs.py | 15 ++++--- .../tests/test_content_encryption_routes.py | 42 ++++++++++++++++++- 2 files changed, 47 insertions(+), 10 deletions(-) diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index 684699aac..4263d20f6 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/frontends/aiq_api/src/aiq_api/routes/jobs.py @@ -572,15 +572,14 @@ async def health_check(): result["reason"] = "configuration_missing" return JSONResponse(status_code=503, content=result) - # Check DB connectivity using any cached async engine + # 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: - 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" + 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" diff --git a/frontends/aiq_api/tests/test_content_encryption_routes.py b/frontends/aiq_api/tests/test_content_encryption_routes.py index 3e0276556..557a2359f 100644 --- a/frontends/aiq_api/tests/test_content_encryption_routes.py +++ b/frontends/aiq_api/tests/test_content_encryption_routes.py @@ -140,7 +140,9 @@ async def _no_op_reaper(*_args, **_kwargs): return app -def _build_assembled_worker_app(monkeypatch, tmp_path, *, dask_available: bool = True) -> FastAPI: +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 @@ -192,7 +194,7 @@ async def __aexit__(self, exc_type, exc, tb): monkeypatch.setattr(WorkflowBuilder, "from_config", lambda _config: BuilderContext()) - front_end = plugin.AIQAPIConfig(db_url=f"sqlite+aiosqlite:///{tmp_path / 'jobs.db'}") + 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 @@ -321,6 +323,42 @@ def test_assembled_worker_health_route_reports_ready_encryption(monkeypatch, tmp 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)