diff --git a/.env.example b/.env.example index 33fa215f..0115264c 100644 --- a/.env.example +++ b/.env.example @@ -1,38 +1,121 @@ -# Agent Server Configuration -AGENT_HOST=0.0.0.0 -AGENT_PORT=5002 -#AGENT_SSL_KEYFILE=/path_to/ssl_key.pem -#AGENT_SSL_CERTFILE=/path_to/ssl_cert.pem +# ============================================================================== +# Environment Variables +# +# Only secrets and infrastructure endpoints belong here. +# All operational config (cache, middleware, filesystem, providers) lives in +# config/agent/runtime/agent.yaml — the single source of truth. +# +# OpenShift: secrets come via Secrets, infra via ConfigMaps. +# ============================================================================== -# Python Logging -PYTHON_LOG_LEVEL=INFO +# --- Environment --- +# Set to "production" to enforce security hardening: +# - ENABLE_AUTH must be true +# - MCP ssl_verify cannot be disabled +# - PII is scrubbed from error responses +# - Security headers are enforced +ENVIRONMENT=development + +# --- Security --- +# Request body size limit (bytes) - prevents DoS attacks +REQUEST_BODY_MAX_SIZE=10485760 # 10MB + +# --- SSO / OIDC Authentication --- +# Supports any OIDC-compliant provider (Keycloak, Okta, Azure AD, Auth0, etc.) +ENABLE_AUTH=false +SSO_ISSUER_URL=https://sso.example.com/realms/myrealm +SSO_CLIENT_ID=your-client-id +SSO_CLIENT_SECRET=your-client-secret +# SSO_JWKS_URI=https://sso.example.com/realms/myrealm/protocol/openid-connect/certs + +# Dev fallback identity (used when ENABLE_AUTH=false) +SSO_DEV_USERNAME=John Doe +SSO_DEV_USER_ID=dev-user + +# User ID encryption for observability privacy +ENABLE_USER_ID_ENCRYPTION=false +# USER_ID_ENCRYPTION_KEY=your-32-byte-hex-key -USE_INMEMORY_SAVER=true +# MCP OAuth token encryption (Fernet key — required when using auth_mode oauth/dcr) +# Encrypts access/refresh tokens in Redis and DCR client secrets in Postgres. +# Generate: python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())" +# MCP_TOKEN_ENCRYPTION_KEY= +# Optional previous key during rotation (decrypt only — see README) +# MCP_TOKEN_ENCRYPTION_KEY_PREVIOUS= +# AGENT_PUBLIC_BASE_URL=http://localhost:5002 -# pgvector credentials for agentic memory (used when USE_INMEMORY_SAVER=false) -POSTGRES_USER=pgvector -POSTGRES_PASSWORD=pgvector -POSTGRES_HOST=0.0.0.0 +# --- Infrastructure --- + +# Postgres (checkpoints, memory, feedback) +# Local dev (`make local`): localhost + port published by compose pgvector (5432) +POSTGRES_HOST=localhost POSTGRES_PORT=5432 -POSTGRES_DB=pgvector +POSTGRES_DB=template_agent +POSTGRES_USER=postgres +POSTGRES_PASSWORD=postgres + +# Redis (Aegra broker: SSE streaming, job queue, crash recovery) +# Local dev (`make local`): localhost + port published by compose redis (6379) +REDIS_URL=redis://localhost:6379/0 +REDIS_BROKER_ENABLED=true + +# MongoDB (platform token usage rollup — optional, set by deploy components) +# +# SECURITY: MONGODB_URI may contain credentials in the URI itself: +# mongodb://user:password@host:27017/tokenusage?authSource=tokenusage +# +# - NEVER commit a URI with credentials to version control. +# - NEVER log or expose this value in error messages or debug output. +# - In production, inject via secrets management: +# Kubernetes : mount as a Secret, reference via envFrom or env.valueFrom.secretKeyRef +# AWS : use Secrets Manager or SSM Parameter Store with an operator/init container +# GCP : use Secret Manager with Workload Identity +# Vault : use the Vault Agent injector or ESO (External Secrets Operator) +# - Scope the MongoDB user to read/write on the tokenusage DB only — no admin privileges. +# - Rotate credentials without redeploying by updating the secret and triggering a rollout. +# +# Local dev (unauthenticated, never in production): +# MONGODB_URI=mongodb://localhost:27017 +# MONGODB_DB=tokenusage + +# --- Observability --- -# exception -LANGFUSE_SECRET_KEY=sk-lf-f46b492e-9335- -LANGFUSE_PUBLIC_KEY=pk-lf-dfa0dab0-c486- +# Langfuse (v4 SDK — auto-read by client and CallbackHandler) +LANGFUSE_PUBLIC_KEY=pk-lf-... +LANGFUSE_SECRET_KEY=sk-lf-... LANGFUSE_BASE_URL=https://cloud.langfuse.com LANGFUSE_TRACING_ENVIRONMENT=development -#Google Vertex AI service creds +# OpenTelemetry — token budget export (metrics/traces via otel_setup.py) +# Agent lifecycle metrics (conversations, streams, threads) via observability.yaml +# ENABLE_OTEL_METRICS=false +# OTEL_EXPORTER_OTLP_ENDPOINT= +# ENABLE_OTEL_TRACES=false +# OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://localhost:4317 +# OTEL_SERVICE_NAME=template-agent +# OTEL_AUTH_TOKEN= +# OTEL_METRIC_EXPORT_INTERVAL_MILLIS=10000 +# ENABLE_OTEL=false +# OTEL_EXPORTER_OTLP_INSECURE=true +# OTEL_METRIC_EXPORT_INTERVAL=5000 + +# --- Model Provider Credentials --- + +# Google Vertex AI GOOGLE_APPLICATION_CREDENTIALS_CONTENT='{ "type": "service_account", - "project_id": "data-and-ai-gemini", - ... - ... + "project_id": "your-project-id", ... "universe_domain": "googleapis.com" }' -# MCP config -MCP_SERVER_NAME=template-mcp-server -MCP_SERVER_URL=http://localhost:5001/mcp -MCP_TRANSPORT_PROTOCOL=streamable_http +# vLLM / OpenAI-compatible (optional) +# VLLM_BASE_URL=http://vllm-server:8000/v1 +# VLLM_API_KEY=EMPTY + +# --- Runtime (rarely changed) --- + +PYTHON_LOG_LEVEL=INFO + +# Per-MCP OAuth client secret (auth_mode: oauth — referenced via oauth.client_secret_env in mcp.json) +# MY_OAUTH_MCP_CLIENT_SECRET=your-client-secret diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 00000000..6f8fc8dd --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,16 @@ +version: 2 +updates: + - package-ecosystem: pip + directory: "/" + schedule: + interval: weekly + + - package-ecosystem: docker + directory: "/" + schedule: + interval: weekly + + - package-ecosystem: github-actions + directory: "/" + schedule: + interval: weekly diff --git a/.github/workflows/build-base-image.yml b/.github/workflows/build-base-image.yml new file mode 100644 index 00000000..09625427 --- /dev/null +++ b/.github/workflows/build-base-image.yml @@ -0,0 +1,180 @@ +name: Build and Push Base Image + +on: + push: + tags: + - '*' + branches: + # - main + - deep-agent + workflow_dispatch: # Allow manual trigger for testing + +permissions: + contents: write + packages: write + security-events: write + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + +jobs: + build-and-push: + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Extract metadata (tags, labels) + id: meta + uses: docker/metadata-action@v5 + with: + images: ghcr.io/${{ github.repository }} + tags: | + type=ref,event=tag + type=raw,value=latest,enable={{is_default_branch}} + type=sha,prefix={{branch}}-,enable=${{ !startsWith(github.ref, 'refs/tags/') }} + + - name: Extract version for deployment package + id: version + run: | + if [[ "${{ github.ref }}" == refs/tags/* ]]; then + VERSION="${GITHUB_REF#refs/tags/}" + else + VERSION="${GITHUB_REF#refs/heads/}-${GITHUB_SHA::7}" + fi + echo "version=${VERSION}" >> $GITHUB_OUTPUT + echo "package_name=agent-deployment.zip" >> $GITHUB_OUTPUT + + - name: Build image locally for scanning + uses: docker/build-push-action@v5 + with: + context: . + file: ./Containerfile + push: false + load: true + tags: scan-target:${{ steps.version.outputs.version }} + cache-from: type=gha + cache-to: type=gha,mode=max + platforms: linux/amd64 + + - name: Scan image for vulnerabilities + uses: aquasecurity/trivy-action@v0.36.0 + with: + image-ref: scan-target:${{ steps.version.outputs.version }} + format: table + severity: CRITICAL,HIGH + scanners: vuln + trivyignores: .trivyignore + exit-code: '1' + + - name: Generate SARIF report + uses: aquasecurity/trivy-action@v0.36.0 + if: always() + with: + image-ref: scan-target:${{ steps.version.outputs.version }} + format: sarif + output: trivy-results.sarif + scanners: vuln + exit-code: '0' + + - name: Upload Trivy scan results to GitHub Security tab + if: always() && hashFiles('trivy-results.sarif') != '' + uses: github/codeql-action/upload-sarif@v4 + with: + sarif_file: trivy-results.sarif + + - name: Upload Trivy report as artifact + if: always() && hashFiles('trivy-results.sarif') != '' + uses: actions/upload-artifact@v4 + with: + name: trivy-report-${{ steps.version.outputs.version }} + path: trivy-results.sarif + retention-days: 30 + + - name: Push multi-platform image + uses: docker/build-push-action@v5 + with: + context: . + file: ./Containerfile + push: true + tags: ${{ steps.meta.outputs.tags }} + labels: ${{ steps.meta.outputs.labels }} + cache-from: type=gha + cache-to: type=gha,mode=max + platforms: linux/amd64,linux/arm64 + + - name: Create deployment package metadata + run: | + cat > deployment/release-info.json <> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "**Version:** \`${{ steps.version.outputs.version }}\`" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "**Container Image Tags:**" >> $GITHUB_STEP_SUMMARY + echo '```' >> $GITHUB_STEP_SUMMARY + echo "${{ steps.meta.outputs.tags }}" >> $GITHUB_STEP_SUMMARY + echo '```' >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "**Pull command:**" >> $GITHUB_STEP_SUMMARY + echo '```bash' >> $GITHUB_STEP_SUMMARY + echo "docker pull ghcr.io/${{ github.repository }}:${{ steps.version.outputs.version }}" >> $GITHUB_STEP_SUMMARY + echo '```' >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "**Deployment Package:**" >> $GITHUB_STEP_SUMMARY + if [[ "${{ github.ref }}" == refs/tags/* ]]; then + echo "📦 \`${{ steps.version.outputs.package_name }}\` attached to release" >> $GITHUB_STEP_SUMMARY + else + echo "📦 \`${{ steps.version.outputs.package_name }}\` uploaded as workflow artifact" >> $GITHUB_STEP_SUMMARY + fi diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 10347e91..b7e6fe83 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -7,8 +7,8 @@ on: branches: [ main ] jobs: - test: - name: Test Suite + unit-tests: + name: Unit Tests runs-on: ubuntu-latest strategy: matrix: @@ -30,9 +30,11 @@ jobs: run: | uv pip install -e ".[dev]" - - name: Run tests with coverage + - name: Run unit tests with coverage + env: + GOOGLE_APPLICATION_CREDENTIALS_CONTENT: ${{ secrets.GOOGLE_APPLICATION_CREDENTIALS_CONTENT }} run: | - source .venv/bin/activate && pytest --cov=template_agent --cov-report=xml --cov-report=term-missing --cov-fail-under=19 + source .venv/bin/activate && pytest tests/unit -m "not e2e" --cov=deep_agent --cov-report=xml --cov-report=html --cov-report=term-missing --cov-fail-under=81 - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v4 @@ -43,3 +45,167 @@ jobs: fail_ci_if_error: false env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} + + skills-evals: + name: Skills Evaluations + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12"] + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v3 + with: + version: "latest" + + - name: Set up virtual env + run: uv venv --python ${{ matrix.python-version }} + + - name: Install dependencies + run: | + uv pip install -e ".[dev]" + + - name: Run skills evals + env: + GOOGLE_APPLICATION_CREDENTIALS_CONTENT: ${{ secrets.GOOGLE_APPLICATION_CREDENTIALS_CONTENT }} + run: | + source .venv/bin/activate && pytest tests/skills -m skills -v + + - name: Upload skills eval results + if: always() + uses: actions/upload-artifact@v4 + with: + name: skills-eval-results + path: tests/workspaces/ + + agent-evals-promptfoo: + name: Agent Evals (Promptfoo) + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.12"] + node-version: ["20"] + + services: + postgres: + image: pgvector/pgvector:pg16 + env: + POSTGRES_USER: pgvector + POSTGRES_PASSWORD: pgvector + POSTGRES_DB: pgvector + ports: + - 5432:5432 + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 5 + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v3 + with: + version: "latest" + + - name: Set up Python virtual env + run: uv venv --python ${{ matrix.python-version }} + + - name: Install Python dependencies + run: | + uv pip install -e ".[dev]" + + - name: Set up Node.js + uses: actions/setup-node@v4 + with: + node-version: ${{ matrix.node-version }} + + - name: Install jq + run: sudo apt-get update && sudo apt-get install -y jq + + - name: Start Mock MCP Server + run: | + source .venv/bin/activate + python tests/mocks/mock_mcp_server.py > mock-mcp.log 2>&1 & + echo $! > mock-mcp.pid + sleep 3 + # Verify mock MCP server is running + curl -f http://localhost:5001/health || (echo "Mock MCP server failed to start" && cat mock-mcp.log && exit 1) + echo "Mock MCP server started successfully" + + - name: Start agent server + env: + GOOGLE_APPLICATION_CREDENTIALS_CONTENT: ${{ secrets.GOOGLE_APPLICATION_CREDENTIALS_CONTENT }} + POSTGRES_USER: pgvector + POSTGRES_PASSWORD: pgvector + POSTGRES_DB: pgvector + POSTGRES_HOST: localhost + POSTGRES_PORT: 5432 + AGENT_HOST: 0.0.0.0 + AGENT_PORT: 5002 + PYTHON_LOG_LEVEL: INFO + run: | + source .venv/bin/activate + python -m deep_agent.src.main > agent.log 2>&1 & + echo $! > agent.pid + echo "Agent PID: $(cat agent.pid)" + + - name: Wait for agent to be ready + run: | + echo "Waiting for agent to start..." + for i in {1..30}; do + if curl -f http://localhost:5002/health 2>/dev/null; then + echo "✓ Agent is ready after ${i} seconds" + exit 0 + fi + echo "Waiting... ($i/30)" + sleep 1 + done + echo "✗ Agent failed to start within 30 seconds" + echo "=== Agent logs ===" + cat agent.log + exit 1 + + - name: Run Promptfoo evals + env: + AGENT_URL: http://localhost:5002 + working-directory: config/agent/evals/promptfoo + run: | + npx promptfoo@latest eval --output ../../../promptfoo-results.json + + - name: Upload agent logs on failure + if: failure() + uses: actions/upload-artifact@v4 + with: + name: agent-logs + path: agent.log + + - name: Stop agent server + if: always() + run: | + if [ -f agent.pid ]; then + kill $(cat agent.pid) || true + rm agent.pid + fi + + - name: Stop Mock MCP Server + if: always() + run: | + if [ -f mock-mcp.pid ]; then + kill $(cat mock-mcp.pid) || true + rm mock-mcp.pid + fi + + - name: Upload Promptfoo results + if: always() + uses: actions/upload-artifact@v4 + with: + name: promptfoo-results + path: promptfoo-results.json + if-no-files-found: warn diff --git a/.gitignore b/.gitignore index ae7b8415..59a2428b 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,8 @@ coverage.xml .hypothesis/ .pytest_cache/ cover/ +tests/workspaces/ +.benchmarks/ # Translations *.mo @@ -140,6 +142,9 @@ venv.bak/ # mypy .mypy_cache/ + +# ruff +.ruff_cache/ .dmypy.json dmypy.json @@ -187,3 +192,61 @@ uv.lock bandit-report.json safety-report.json .safety-project.ini +TASKS.md +ROADMAP.md + +# Kind cluster cloned repos (make kind) +.kind/ + +# LangGraph Platform (local dev) +.langgraph/ +langgraph-api-data/ + +# Development infrastructure data +redis_data/ +dump.rdb +langfuse_data/ +jaeger_data/ +otel_data/ + +# Runtime cache data (not source code) +.cache/ +diskcache/ +*.cache + +# Load test results (MR-33) +load_test_results/ +*.jmx +*.jtl +locust_reports/ +locust.log + +# Test fixtures data (MR-34) +tests/fixtures/data/ +tests/fixtures/*.db +tests/fixtures/*.sqlite + +# Performance benchmarks (MR-42) +benchmark_results/ +*.benchmark +benchmarks/output/ + +# Structured logs (MR-89: JSONL output) +logs/ +*.jsonl +*.log.json + +# Playwright MCP artifacts +.playwright-mcp/ + +# Development environment +docker-compose.override.yml +.envrc +.direnv/ + +# Aegra-generated scaffolding (conflicts with existing compose.yaml + k8s deployment) +Dockerfile +docker-compose.yml +scripts/ +docs/superpowers/plans/ +docs/superpowers/specs/2026-07-13-* diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b53e420e..0c1a5b56 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -5,12 +5,14 @@ repos: - id: trailing-whitespace - id: end-of-file-fixer - id: check-yaml + args: ['--allow-multiple-documents'] - id: check-added-large-files - id: check-merge-conflict - id: debug-statements - id: check-case-conflict - id: check-docstring-first - id: check-json + exclude: ^config/agent/mcp\.json$ - id: check-toml - repo: https://github.com/astral-sh/ruff-pre-commit @@ -26,7 +28,9 @@ repos: hooks: - id: mypy args: [--ignore-missing-imports] - additional_dependencies: [types-requests] + additional_dependencies: [types-requests, types-PyYAML, types-cachetools] + exclude: ^tests/ + - repo: https://github.com/PyCQA/pydocstyle rev: 6.3.0 diff --git a/.trivyignore b/.trivyignore new file mode 100644 index 00000000..10c73a5f --- /dev/null +++ b/.trivyignore @@ -0,0 +1,11 @@ +# CVE-2026-31072: APScheduler RCE via insecure deserialization (CRITICAL) +# No fixed version available as of 2026-07-07. +# Risk mitigated: AsyncScheduler is used with in-memory store only (no persistent +# job store backend), so the deserialization attack vector is not exposed. +# Revisit when APScheduler 4.x releases a patched version. +CVE-2026-31072 + +# CVE-2026-25087: pyarrow DoS via use-after-free when reading IPC files (HIGH) +# Fixed in pyarrow 23.0.1, but pinning it causes a build dependency conflict. +# pyarrow is a transitive dep — not used directly. Revisit when upstream resolves. +CVE-2026-25087 diff --git a/Containerfile b/Containerfile index 7ffd50ad..5d301edc 100644 --- a/Containerfile +++ b/Containerfile @@ -1,38 +1,38 @@ -FROM registry.access.redhat.com/ubi9/python-312:latest +# Containerfile for template-agent (single image for dev and production) +# +# Agent config is NOT baked in — mount config/agent at /app/config/agent +# (compose: ./config:/app/config:ro; K8s: ConfigMap/PVC). +# +# Build: podman build -t template-agent . +# Run: podman run -v ./config:/app/config:ro -p 5002:5002 template-agent -# -------------------------------------------------------------------------------------------------- -# set the working directory to /app -# -------------------------------------------------------------------------------------------------- +ARG PYTHON_TAG=3.14.4-builder +FROM registry.access.redhat.com/hi/python:${PYTHON_TAG} WORKDIR /app - -# -------------------------------------------------------------------------------------------------- -# Copy manifest files and install python packages -# -------------------------------------------------------------------------------------------------- - USER root + COPY pyproject.toml /app/pyproject.toml -RUN pip install uv -RUN uv venv -RUN source /app/.venv/bin/activate -RUN uv pip install -r pyproject.toml -USER default -# -------------------------------------------------------------------------------------------------- -# copy source code and files -# -------------------------------------------------------------------------------------------------- +RUN pip install --no-cache-dir uv && \ + uv venv /app/.venv && \ + uv pip install --python /app/.venv/bin/python -r pyproject.toml && \ + mkdir -p /app/.cache /app/config/agent && \ + chown -R 65532:root /app/.cache /app/config -COPY template_agent /app/template_agent +USER 65532 -# -------------------------------------------------------------------------------------------------- -# Set PYTHONPATH to include /app -# -------------------------------------------------------------------------------------------------- +COPY --chown=65532:root deep_agent /app/deep_agent +COPY --chown=65532:root aegra.json /app/aegra.json +COPY --chown=65532:root entrypoint.sh /app/entrypoint.sh ENV PYTHONPATH=/app +ENV AGENT_HOST=0.0.0.0 +ENV AGENT_PORT=5002 +ENV AEGRA_CONFIG=/app/aegra.json +ENV CONFIG_PATH=/app/config/agent +EXPOSE 5002 -# -------------------------------------------------------------------------------------------------- -# add entrypoint for the container -# -------------------------------------------------------------------------------------------------- - -CMD ["/app/.venv/bin/python", "-m", "template_agent.src.main"] +ENTRYPOINT ["/app/entrypoint.sh"] +CMD ["/app/.venv/bin/python", "-m", "deep_agent.aegra.entrypoint"] diff --git a/Makefile b/Makefile index 13c669c7..bf07a977 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: local dev test clean deploy undeploy +.PHONY: local dev test clean deploy deploy-headless undeploy undeploy-headless kind kind-down container container-down local-down test-triggers test-integration test-headless headless # OpenShift namespace (can be overridden: make deploy openshift NAMESPACE=my-project) NAMESPACE ?= $(shell oc project -q 2>/dev/null) @@ -27,8 +27,14 @@ install: @chmod +x /tmp/activate_and_shell.sh @exec /tmp/activate_and_shell.sh -clean: - @echo "Cleaning up non-code artifacts..." +clean: ## Remove build artifacts, venv, and tear down compose stack + @echo "Stopping agent on port 5002 (if running)..." + @lsof -ti :5002 | xargs kill -9 2>/dev/null || true + @echo "Stopping compose stack (if running)..." + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml --profile container down -v 2>/dev/null || true + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml stop pgvector redis 2>/dev/null || true + @podman rmi template-agent_template-agent 2>/dev/null || true + @echo "Cleaning up build artifacts..." @rm -rf .venv @rm -rf __pycache__ @rm -rf .pytest_cache @@ -51,19 +57,138 @@ test: echo "Error: Virtual environment not found. Run 'make install' first to set up the environment."; \ exit 1; \ fi + .venv/bin/python -m pytest tests/unit + +test-cov: ## Run unit tests with coverage report + @if [ ! -d ".venv" ]; then \ + echo "Error: Virtual environment not found. Run 'make install' first to set up the environment."; \ + exit 1; \ + fi + @echo "Running unit tests with coverage..." + .venv/bin/python -m pytest tests/unit --cov=deep_agent --cov-report=xml --cov-report=html --cov-report=term-missing + +test-all: + @if [ ! -d ".venv" ]; then \ + echo "Error: Virtual environment not found. Run 'make install' first to set up the environment."; \ + exit 1; \ + fi + @echo "Running all tests (unit + skills evals)..." .venv/bin/python -m pytest +test-skills: + @if [ ! -d ".venv" ]; then \ + echo "Error: Virtual environment not found. Run 'make install' first to set up the environment."; \ + exit 1; \ + fi + @echo "Running skills evaluations..." + .venv/bin/python -m pytest tests/skills -m skills -v + +eval-promptfoo: + @echo "Running Promptfoo agent evaluations..." + @echo "Make sure agent is running at http://localhost:5002" + @cd config/agent/evals/promptfoo && npx promptfoo@latest eval + +mock-mcp: + @if [ ! -d ".venv" ]; then \ + echo "Error: Virtual environment not found. Run 'make install' first."; \ + exit 1; \ + fi + @echo "Starting Mock MCP Server on http://localhost:5001 (Ctrl+C to stop)..." + @.venv/bin/python tests/mocks/mock_mcp_server.py + +local-with-mock: + @echo "Run in separate terminals:" + @echo " Terminal 1: make mock-mcp" + @echo " Terminal 2: make local" + local: @echo "Setting up local environment..." + @test -f .env 2>/dev/null || (echo "Creating .env from .env.example..." && cp .env.example .env 2>/dev/null) || true + @lsof -ti :5002 | xargs kill -9 2>/dev/null || true + @echo "Cleaning up stale containers from previous naming scheme..." + @podman rm -f demo-pgvector demo-redis 2>/dev/null || true + @echo "Starting infrastructure (Postgres + Redis)..." + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml up -d pgvector redis + @echo "Waiting for Postgres to be ready..." + @until podman exec template-agent-pgvector pg_isready -U postgres -q 2>/dev/null; do sleep 1; done + @podman exec template-agent-pgvector psql -U postgres -tc "SELECT 1 FROM pg_database WHERE datname='aegra'" | grep -q 1 \ + || podman exec template-agent-pgvector psql -U postgres -c "CREATE DATABASE aegra;" + @echo "Starting agent with LangGraph Platform..." + @echo "API available at: http://localhost:5002" + @echo "Press Ctrl+C to stop the server (Postgres/Redis keep running — use 'make local-down' to stop them)" + @trap 'lsof -ti :5002 | xargs kill -INT 2>/dev/null || true; sleep 2; lsof -ti :5002 | xargs kill -9 2>/dev/null || true; exit 130' INT TERM; \ + REDIS_BROKER_ENABLED=true \ + POSTGRES_HOST=localhost \ + POSTGRES_PORT=5432 \ + POSTGRES_DB=template_agent \ + POSTGRES_USER=postgres \ + POSTGRES_PASSWORD=postgres \ + REDIS_URL=redis://localhost:6379/0 \ + .venv/bin/aegra dev --port 5002 --no-db-check + +local-down: + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml stop pgvector redis + +headless: ## Start agent in headless mode (background worker with event triggers) + @echo "Setting up headless environment..." @test -f .env || (echo "Creating .env from .env.example..." && cp .env.example .env) - @echo "Starting MCP server locally on port 5002..." - @echo "Health check available at: http://localhost:5002/health" - @echo "Press Ctrl+C to stop the server" - @. .venv/bin/activate && USE_INMEMORY_SAVER=true python -m template_agent.src.main + @echo "Starting infrastructure (Postgres + Redis)..." + @podman-compose -f compose.yaml up -d pgvector redis + @echo "Waiting for Postgres to be ready..." + @until podman exec demo-pgvector pg_isready -U postgres -q 2>/dev/null; do sleep 1; done + @podman exec demo-pgvector psql -U postgres -tc "SELECT 1 FROM pg_database WHERE datname='aegra'" | grep -q 1 \ + || podman exec demo-pgvector psql -U postgres -c "CREATE DATABASE aegra;" + @echo "Starting headless agent worker..." + @echo "Set mode: headless in config/agent/runtime/agent.yaml" + @echo "Press Ctrl+C to stop the worker" + @. .venv/bin/activate && ENVIRONMENT=local REDIS_URL=redis://localhost:6379/0 python -m deep_agent.headless container: - export PODMAN_COMPOSE_SILENT=true - podman-compose --no-ansi up --build --force-recreate --remove-orphans --timeout=60 + @test -f .env || (echo "Creating .env from .env.example..." && cp .env.example .env) + @echo "Starting stack: pgvector, redis, template-agent, jaeger" + @echo "Agent: http://localhost:5002" + @echo "Jaeger: http://localhost:16686" + @export PODMAN_COMPOSE_SILENT=true; \ + trap 'export PODMAN_COMPOSE_SILENT=true; podman-compose -f compose.yaml --profile observability down --timeout 10 2>/dev/null || true; exit 130' INT TERM; \ + ENABLE_OTEL=true \ + OTEL_EXPORTER_OTLP_ENDPOINT=http://jaeger:4317 \ + ENABLE_OTEL_TRACES=true \ + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://jaeger:4317 \ + podman-compose --profile observability --no-ansi up --build --force-recreate --remove-orphans --timeout=60 + +container-down: + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml --profile observability down + +# --------------------------------------------------------------------------- +# Development environment targets +# --------------------------------------------------------------------------- + +dev: ## Start agent + deps in containers (detached, tail logs) + @echo "Starting agent stack (pgvector, redis, template-agent)..." + @echo "Agent: http://localhost:5002" + @echo "" + @test -f .env || (echo "Creating .env from .env.example..." && cp .env.example .env) + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml --profile container up --build -d + @echo "" + @echo "Tailing agent logs (Ctrl+C to stop)..." + @echo "" + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml --profile container logs -f template-agent + +dev-down: ## Stop dev stack + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml --profile container down + +dev-clean: ## Stop dev stack and remove all data + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml --profile container down -v + @echo "All dev data volumes removed" + +dev-logs: ## Tail all service logs + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml --profile container logs -f + +dev-restart: ## Restart dev stack + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml --profile container restart + +dev-agent: ## Restart just the agent service + @export PODMAN_COMPOSE_SILENT=true && podman-compose -f compose.yaml --profile container restart template-agent # Deployment targets deploy: @@ -85,13 +210,12 @@ openshift: echo "Switching to namespace..."; \ oc project $(NAMESPACE) || (echo "Error: Cannot switch to namespace '$(NAMESPACE)'. Check permissions." && exit 1); \ echo "Updating namespace references..."; \ - sed -i.bak "s|NAMESPACE_PLACEHOLDER|$(NAMESPACE)|g" deployment/openshift/deployment.yaml; \ - sed -i.bak "s|namespace: template-agent|namespace: $(NAMESPACE)|g" deployment/openshift/kustomization.yaml; \ + sed -i.bak "s|NAMESPACE_PLACEHOLDER|$(NAMESPACE)|g" deployment/overlays/openshift/kustomization.yaml; \ echo "Creating BuildConfig and ImageStream..."; \ - oc apply -f deployment/openshift/buildconfig.yaml; \ - oc apply -f deployment/openshift/imagestream.yaml; \ + oc apply -f deployment/overlays/openshift/buildconfig.yaml; \ + oc apply -f deployment/overlays/openshift/imagestream.yaml; \ echo "Building container image from source..."; \ - oc start-build template-agent --from-dir=. \ + oc start-build agent --from-dir=. \ --exclude='(^|/)\.venv(/|$$)' \ --exclude='(^|/)__pycache__(/|$$)' \ --exclude='(^|/)\.pytest_cache(/|$$)' \ @@ -100,75 +224,155 @@ openshift: --exclude='(^|/)\.mypy_cache(/|$$)' \ --exclude='(^|/)\.ruff_cache(/|$$)' \ --exclude='.*\.log$$' \ - --follow || (mv deployment/openshift/deployment.yaml.bak deployment/openshift/deployment.yaml 2>/dev/null; mv deployment/openshift/kustomization.yaml.bak deployment/openshift/kustomization.yaml 2>/dev/null; exit 1); \ + --follow || (mv deployment/overlays/openshift/kustomization.yaml.bak deployment/overlays/openshift/kustomization.yaml 2>/dev/null; exit 1); \ echo "Deploying resources to OpenShift..."; \ - oc apply -k deployment/openshift/ || (mv deployment/openshift/deployment.yaml.bak deployment/openshift/deployment.yaml 2>/dev/null; mv deployment/openshift/kustomization.yaml.bak deployment/openshift/kustomization.yaml 2>/dev/null; exit 1); \ - rm -f deployment/openshift/deployment.yaml.bak deployment/openshift/kustomization.yaml.bak; \ + oc apply -k deployment/overlays/openshift/ || (mv deployment/overlays/openshift/kustomization.yaml.bak deployment/overlays/openshift/kustomization.yaml 2>/dev/null; exit 1); \ + rm -f deployment/overlays/openshift/kustomization.yaml.bak; \ echo "Deployment complete!"; \ echo "Checking deployment status..."; \ - oc get pods -l app=template-agent; \ + oc get pods -l app=agent; \ echo ""; \ echo "Useful commands:"; \ - echo " View logs: oc logs -l app=template-agent --tail=100"; \ - echo " Get route: oc get route template-agent"; \ - echo " Check status: oc get pods,svc,route -l app=template-agent" + echo " View logs: oc logs -l app=agent --tail=100"; \ + echo " Get route: oc get route agent"; \ + echo " Check status: oc get pods,svc,route -l app=agent" -mpp: - @echo "Checking for oc CLI..." +deploy-headless: ## Deploy headless worker to OpenShift (requires: NAMESPACE) + @echo "Deploying headless worker to OpenShift..." @which oc > /dev/null || (echo "Error: oc CLI not found. Please install OpenShift CLI." && exit 1) - @echo "Validating TENANT parameter..." - @if [ -z "$(TENANT)" ]; then \ - echo "Error: TENANT not set. Usage: make deploy mpp TENANT=your-tenant"; \ + @if [ -z "$(NAMESPACE)" ]; then \ + echo "Error: NAMESPACE not set. Usage: make deploy-headless NAMESPACE=your-project"; \ exit 1; \ fi; \ - CONFIG_NAMESPACE="$(TENANT)--config"; \ - RUNTIME_NAMESPACE="$(TENANT)--template"; \ - echo "Config namespace: $$CONFIG_NAMESPACE"; \ - echo "Runtime namespace: $$RUNTIME_NAMESPACE"; \ - echo "Updating tenant.yaml with config namespace..."; \ - sed -i.bak "s|TENANT_PLACEHOLDER|$$CONFIG_NAMESPACE|g" deployment/mpp/tenant.yaml; \ - echo "Creating/switching to config namespace..."; \ - oc project $$CONFIG_NAMESPACE 2>/dev/null || oc new-project $$CONFIG_NAMESPACE || (echo "Error: Cannot create/switch to namespace '$$CONFIG_NAMESPACE'." && mv deployment/mpp/tenant.yaml.bak deployment/mpp/tenant.yaml 2>/dev/null && exit 1); \ - echo "Applying TenantNamespace CR to create runtime namespace..."; \ - oc apply -f deployment/mpp/tenant.yaml || (mv deployment/mpp/tenant.yaml.bak deployment/mpp/tenant.yaml 2>/dev/null && exit 1); \ - echo "Waiting for runtime namespace '$$RUNTIME_NAMESPACE' to be created..."; \ - COUNTER=1; \ - until oc get project $$RUNTIME_NAMESPACE 2>/dev/null || [ $$COUNTER -gt 30 ]; do \ - echo "Waiting for namespace... ($$COUNTER/30)"; \ - sleep 2; \ - COUNTER=$$((COUNTER + 1)); \ - done; \ - if [ $$COUNTER -le 30 ]; then \ - echo "Runtime namespace '$$RUNTIME_NAMESPACE' is ready"; \ - fi; \ - oc project "$(TENANT)--$(RUNTIME_NAMESPACE)" > /dev/null 2>&1 || (echo "Error: Runtime namespace '$$RUNTIME_NAMESPACE' was not created" && mv deployment/mpp/tenant.yaml.bak deployment/mpp/tenant.yaml 2>/dev/null && exit 1); \ - echo "Switching to runtime namespace..."; \ - oc project $$RUNTIME_NAMESPACE || (echo "Error: Cannot switch to runtime namespace '$$RUNTIME_NAMESPACE'" && mv deployment/mpp/tenant.yaml.bak deployment/mpp/tenant.yaml 2>/dev/null && exit 1); \ - echo "Creating BuildConfig and ImageStream..."; \ - oc apply -f deployment/mpp/buildconfig.yaml; \ - oc apply -f deployment/mpp/imagestream.yaml; \ - echo "Building container image from source..."; \ - oc start-build template-agent --from-dir=. \ - --exclude='(^|/)\.venv(/|$$)' \ - --exclude='(^|/)__pycache__(/|$$)' \ - --exclude='(^|/)\.pytest_cache(/|$$)' \ - --exclude='(^|/)tests(/|$$)' \ - --exclude='(^|/)examples(/|$$)' \ - --exclude='(^|/)\.mypy_cache(/|$$)' \ - --exclude='(^|/)\.ruff_cache(/|$$)' \ - --exclude='.*\.log$$' \ - --follow || (mv deployment/mpp/tenant.yaml.bak deployment/mpp/tenant.yaml 2>/dev/null; exit 1); \ - echo "Deploying resources to MPP..."; \ - oc apply -k deployment/mpp/ || (mv deployment/mpp/tenant.yaml.bak deployment/mpp/tenant.yaml 2>/dev/null; exit 1); \ - rm -f deployment/mpp/tenant.yaml.bak; \ - echo "Deployment complete!"; \ + echo "Using namespace: $(NAMESPACE)"; \ + oc project $(NAMESPACE) || (echo "Error: Cannot switch to namespace '$(NAMESPACE)'. Check permissions." && exit 1); \ + echo "Updating namespace references..."; \ + sed -i.bak "s|NAMESPACE_PLACEHOLDER|$(NAMESPACE)|g" deployment/overlays/openshift-headless/kustomization.yaml; \ + echo "Deploying headless worker resources..."; \ + oc apply -k deployment/overlays/openshift-headless/ || (mv deployment/overlays/openshift-headless/kustomization.yaml.bak deployment/overlays/openshift-headless/kustomization.yaml 2>/dev/null; exit 1); \ + rm -f deployment/overlays/openshift-headless/kustomization.yaml.bak; \ + echo "Headless worker deployment complete!"; \ echo "Checking deployment status..."; \ - oc get pods -l app=template-agent; \ + oc get pods -l component=agent-headless; \ echo ""; \ echo "Useful commands:"; \ - echo " View logs: oc logs -l app=template-agent --tail=100"; \ - echo " Get route: oc get route template-agent"; \ - echo " Check status: oc get pods,svc,route -l app=template-agent" + echo " View logs: oc logs -l component=agent-headless --tail=100"; \ + echo " Check status: oc get pods,svc -l component=agent-headless"; \ + echo " Scale: oc scale deployment agent-headless --replicas=N" + +undeploy-headless: ## Remove headless worker from OpenShift + @which oc > /dev/null || (echo "Error: oc CLI not found." && exit 1) + @oc project $(NAMESPACE) || (echo "Error: Cannot switch to namespace '$(NAMESPACE)'" && exit 1) + @echo "Removing headless worker deployment..." + @oc delete deployment,service,hpa,pdb,scaledobject -l component=agent-headless 2>/dev/null || true + @echo "Headless worker undeployed" + +mpp: + @echo "Error: MPP deployment is not yet implemented." + @echo "The deployment/mpp/ kustomize overlay has not been created." + @echo "Use 'make deploy openshift' for OpenShift or 'make kind' for local Kubernetes." + @exit 1 + +# --------------------------------------------------------------------------- +# Kind cluster: local Kubernetes testing +# --------------------------------------------------------------------------- + +KIND_CLUSTER := template-agent +KIND_CTX := kind-$(KIND_CLUSTER) +KIND_IMAGE := template-agent:local +KIND_MCP_IMAGE := template-mcp-server:local +KIND_UI_IMAGE := template-ui:local +KIND_NS := template-agent +KIND_DIR := .kind +KIND_MCP_REPO := https://github.com/redhat-data-and-ai/template-mcp-server.git +KIND_MCP_BRANCH := feat/rh-flavour +KIND_UI_REPO := https://github.com/redhat-data-and-ai/template-ui.git +KIND_UI_BRANCH := feat/rh-flavour +KCTL := kubectl --context $(KIND_CTX) + +kind: ## Deploy full stack (agent + MCP + UI) to a local Kind cluster + @echo "╔════════════════════════════════════════════════════════════════╗" + @echo "║ Kind: Deploy full stack to local Kubernetes cluster ║" + @echo "║ Services: UI + Agent + MCP Server + Postgres + Redis ║" + @echo "╚════════════════════════════════════════════════════════════════╝" + @which kind > /dev/null || (echo "Error: kind not found. Install: https://kind.sigs.k8s.io" && exit 1) + @which kubectl > /dev/null || (echo "Error: kubectl not found." && exit 1) + @# --- Step 1: Clone MCP server and UI if needed --- + @if [ ! -d "$(KIND_DIR)/template-mcp-server" ]; then \ + echo "Cloning template-mcp-server (branch: $(KIND_MCP_BRANCH))..."; \ + mkdir -p $(KIND_DIR); \ + git clone --branch $(KIND_MCP_BRANCH) --depth 1 $(KIND_MCP_REPO) $(KIND_DIR)/template-mcp-server; \ + else \ + echo "MCP server already cloned"; \ + fi + @if [ ! -d "$(KIND_DIR)/template-ui" ]; then \ + echo "Cloning template-ui (branch: $(KIND_UI_BRANCH))..."; \ + mkdir -p $(KIND_DIR); \ + git clone --branch $(KIND_UI_BRANCH) --depth 1 $(KIND_UI_REPO) $(KIND_DIR)/template-ui; \ + else \ + echo "UI already cloned"; \ + fi + @# --- Step 2: Create cluster if not exists --- + @if ! kind get clusters 2>/dev/null | grep -q "$(KIND_CLUSTER)"; then \ + echo "Creating Kind cluster '$(KIND_CLUSTER)'..."; \ + kind create cluster --name $(KIND_CLUSTER) --config=- <<< '{"kind":"Cluster","apiVersion":"kind.x-k8s.io/v1alpha4","nodes":[{"role":"control-plane","kubeadmConfigPatches":["kind: InitConfiguration\nnodeRegistration:\n kubeletExtraArgs:\n node-labels: ingress-ready=true\n"],"extraPortMappings":[{"containerPort":80,"hostPort":80,"protocol":"TCP"},{"containerPort":443,"hostPort":443,"protocol":"TCP"}]}]}'; \ + echo "Installing NGINX Ingress..."; \ + $(KCTL) apply -f https://raw.githubusercontent.com/kubernetes/ingress-nginx/main/deploy/static/provider/kind/deploy.yaml; \ + echo "Waiting for ingress controller pod to be scheduled..."; \ + sleep 10; \ + $(KCTL) wait --namespace ingress-nginx --for=condition=ready pod --selector=app.kubernetes.io/component=controller --timeout=120s; \ + else \ + echo "Kind cluster '$(KIND_CLUSTER)' already exists"; \ + fi + @# --- Step 3: Build and load images --- + @echo "Building agent image..." + @podman build -t $(KIND_IMAGE) . + @echo "Building MCP server image..." + @podman build -t $(KIND_MCP_IMAGE) -f $(KIND_DIR)/template-mcp-server/Containerfile $(KIND_DIR)/template-mcp-server + @echo "Building UI image..." + @podman build -t $(KIND_UI_IMAGE) $(KIND_DIR)/template-ui + @echo "Loading images into Kind (podman -> archive -> kind)..." + @podman save $(KIND_IMAGE) -o /tmp/kind-agent.tar && kind load image-archive /tmp/kind-agent.tar --name $(KIND_CLUSTER) && rm -f /tmp/kind-agent.tar + @podman save $(KIND_MCP_IMAGE) -o /tmp/kind-mcp.tar && kind load image-archive /tmp/kind-mcp.tar --name $(KIND_CLUSTER) && rm -f /tmp/kind-mcp.tar + @podman save $(KIND_UI_IMAGE) -o /tmp/kind-ui.tar && kind load image-archive /tmp/kind-ui.tar --name $(KIND_CLUSTER) && rm -f /tmp/kind-ui.tar + @# --- Step 4: Deploy --- + @echo "Deploying to Kind..." + @$(KCTL) create namespace $(KIND_NS) 2>/dev/null || true + @$(KCTL) apply -k deployment/overlays/kind/ + @$(KCTL) apply -k $(KIND_DIR)/template-mcp-server/deployment/kind/ + @echo "" + @echo "Waiting for pods..." + @$(KCTL) -n $(KIND_NS) wait --for=condition=ready pod -l component=database --timeout=60s 2>/dev/null || true + @$(KCTL) -n $(KIND_NS) wait --for=condition=ready pod -l component=cache --timeout=60s 2>/dev/null || true + @$(KCTL) -n $(KIND_NS) wait --for=condition=ready pod -l component=mcp-server --timeout=90s 2>/dev/null || true + @$(KCTL) -n $(KIND_NS) wait --for=condition=ready pod -l component=agent --timeout=120s 2>/dev/null || true + @$(KCTL) -n $(KIND_NS) wait --for=condition=ready pod -l component=ui --timeout=90s 2>/dev/null || true + @# --- Step 6: Port-forwards for localhost access --- + @echo "Setting up port-forwards..." + @$(KCTL) -n $(KIND_NS) port-forward svc/ui 8080:8080 &>/dev/null & + @$(KCTL) -n $(KIND_NS) port-forward svc/agent 5002:5002 &>/dev/null & + @$(KCTL) -n $(KIND_NS) port-forward svc/mcp-server 5001:5001 &>/dev/null & + @sleep 2 + @echo "" + @echo "╔════════════════════════════════════════════════════════════════╗" + @echo "║ Kind cluster ready! ║" + @echo "║ UI: http://localhost:8080 ║" + @echo "║ Agent: http://localhost:5002 ║" + @echo "║ MCP Server: http://localhost:5001 ║" + @echo "╚════════════════════════════════════════════════════════════════╝" + @echo "" + @echo "Useful commands:" + @echo " Pods: $(KCTL) -n $(KIND_NS) get pods" + @echo " Logs: $(KCTL) -n $(KIND_NS) logs -l component=agent -f" + @echo " Teardown: make kind-down" + +kind-down: ## Delete the Kind cluster and clean up cloned repos + @echo "Stopping port-forwards..." + @pkill -f "kubectl.*port-forward.*$(KIND_NS)" 2>/dev/null || true + @echo "Deleting Kind cluster '$(KIND_CLUSTER)'..." + @kind delete cluster --name $(KIND_CLUSTER) 2>/dev/null || true + @rm -rf $(KIND_DIR) + @echo "Kind cluster and .kind/ cleaned up." undeploy: @if [ "$(filter openshift,$(MAKECMDGOALS))" = "openshift" ]; then \ @@ -176,23 +380,47 @@ undeploy: which oc > /dev/null || (echo "Error: oc CLI not found. Please install OpenShift CLI." && exit 1); \ oc project $(NAMESPACE) || (echo "Error: Cannot switch to namespace '$(NAMESPACE)'" && exit 1); \ echo "Removing OpenShift deployment..."; \ - oc delete deployment,service,route,configmap,secret,pvc,buildconfig,imagestream -l app=template-agent 2>/dev/null || true; \ + oc delete deployment,service,route,configmap,secret,pvc,buildconfig,imagestream -l app=agent 2>/dev/null || true; \ echo "Undeployment complete!"; \ - exit 1; \ elif [ "$(filter mpp,$(MAKECMDGOALS))" = "mpp" ]; then \ echo "Checking for oc CLI..."; \ RUNTIME_NAMESPACE="$(TENANT)--template"; \ which oc > /dev/null || (echo "Error: oc CLI not found. Please install OpenShift CLI." && exit 1); \ oc project $$RUNTIME_NAMESPACE || (echo "Error: Cannot switch to runtime namespace '$$RUNTIME_NAMESPACE'" && exit 1); \ echo "Removing MPP deployment..."; \ - oc delete deployment,service,route,configmap,secret,pvc,buildconfig,imagestream -l app=template-agent 2>/dev/null || true; \ + oc delete deployment,service,route,configmap,secret,pvc,buildconfig,imagestream -l app=agent 2>/dev/null || true; \ echo "Undeployment complete!"; \ - exit 1; \ else \ echo "Usage: make undeploy [openshift|mpp]"; \ echo "Available undeployment targets: openshift, mpp"; \ exit 1; \ fi +# --------------------------------------------------------------------------- +# Headless / Trigger tests +# --------------------------------------------------------------------------- + +test-triggers: ## Unit tests for triggers only + @if [ ! -d ".venv" ]; then \ + echo "Error: Virtual environment not found. Run 'make install' first to set up the environment."; \ + exit 1; \ + fi + .venv/bin/python -m pytest tests/unit/triggers -v + +test-integration: ## Integration tests (requires Redis + DB) + @if [ ! -d ".venv" ]; then \ + echo "Error: Virtual environment not found. Run 'make install' first to set up the environment."; \ + exit 1; \ + fi + @echo "Running integration tests..." + .venv/bin/python -m pytest tests/integration -m integration -v + +test-headless: ## Full headless startup test + @if [ ! -d ".venv" ]; then \ + echo "Error: Virtual environment not found. Run 'make install' first to set up the environment."; \ + exit 1; \ + fi + .venv/bin/python -m pytest tests/integration/triggers/test_headless_startup.py -v + %: @: diff --git a/README.md b/README.md index b2f77491..833f9fd5 100644 --- a/README.md +++ b/README.md @@ -1,382 +1,263 @@ # Template Agent [![Python 3.12+](https://img.shields.io/badge/python-3.12,3.13-blue.svg)](https://www.python.org/downloads/) -[![Tests](https://github.com/redhat-data-and-ai/template-agent/actions/workflows/test.yml/badge.svg)](https://github.com/redhat-data-and-ai/template-mcp-server/actions/workflows/ci.yml) -[![Coverage](https://codecov.io/gh/redhat-data-and-ai/template-agent/branch/main/graph/badge.svg)](https://codecov.io/gh/redhat-data-and-ai/template-mcp-server) +[![Tests](https://github.com/redhat-data-and-ai/template-agent/actions/workflows/test.yml/badge.svg)](https://github.com/redhat-data-and-ai/template-agent/actions/workflows/test.yml) [![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) -A production-ready template for building AI agents with streaming capabilities, conversation management, and enterprise-grade features. - -## 🌟 Features - -- **Simplified Streaming API**: Clean, consistent event format for easy client integration -- **Real-time Streaming**: Server-Sent Events (SSE) with token and message streaming -- **Multiple Client Examples**: TypeScript, Python async, and Streamlit demo applications -- **Conversation Management**: Multi-turn conversations with thread persistence -- **Enterprise Integration**: Langfuse tracing, PostgreSQL checkpointing, SSO support -- **Modular Architecture**: AgentManager abstraction with clean separation of concerns -- **Production Ready**: Health checks, error handling, and comprehensive logging -- **Google AI Integration**: Built-in support for Google Generative AI models - -## 🏗️ Architecture - -```mermaid -graph TB - subgraph "Client" - UI[Web UI] - API[API Client] - end - - subgraph "Template Agent" - subgraph "API Layer" - Health[Health Check] - Stream[Stream Chat] - History[Chat History] - Threads[Thread Management] - Feedback[Feedback] - end - - subgraph "Core Layer" - Agent[Agent Engine] - Utils[Message Utils] - Prompt[Prompt Management] - end - - subgraph "Data Layer" - DB[(PostgreSQL)] - Langfuse[Langfuse] - end - - subgraph "External Services" - Google[Google AI] - SSO[SSO Auth] - end - end - - UI --> Health - UI --> Stream - UI --> History - UI --> Threads - UI --> Feedback - - API --> Health - API --> Stream - API --> History - API --> Threads - API --> Feedback - - Stream --> Agent - Agent --> Utils - Agent --> Prompt - Agent --> Google - - History --> DB - Threads --> DB - Agent --> DB - Agent --> Langfuse - Feedback --> Langfuse -``` +A template for building [Deep Agents](https://github.com/langchain-ai/deepagents) with the [LangGraph](https://langchain-ai.github.io/langgraph/) framework via the Aegra CLI. Includes orchestrator + subagents, MCP tool integration, conversation persistence, Langfuse tracing, and OpenTelemetry metrics. -## 📡 Simplified Streaming API +## Features -The Template Agent now features a simplified streaming API that makes client integration easier while preserving all enterprise features: +**Agent capabilities:** +- Orchestrator with analyst and publisher subagents +- Skills: `client-intake`, `bmi-report`, `email-formatter` +- MCP auth modes: SSO pass-through, OAuth, and DCR -### Single Streaming Endpoint +**Infrastructure:** +- Aegra dev server with Redis-backed SSE streaming +- PostgreSQL checkpoints, memory, and feedback storage +- Config-as-code in `config/agent/` (no Python edits for most changes) +- Container-ready with Red Hat UBI; OpenShift and Kind deployment overlays -```http -POST /v1/stream -Content-Type: application/json -Accept: text/event-stream -``` +## Quick Start -### Request Format +**Prerequisites:** Python 3.12+, [uv](https://docs.astral.sh/uv/), [Podman](https://podman.io/), Google Vertex AI credentials -```json -{ - "message": "User input message", - "thread_id": "conversation-thread-id", - "session_id": "session-id", - "user_id": "user-identifier", - "stream_tokens": true -} +```bash +git clone https://github.com/redhat-data-and-ai/template-agent.git +cd template-agent +make install # creates venv, installs deps + pre-commit hooks +make local # pgvector + redis in compose; agent on host → :5002 ``` -### Response Format +Verify in another terminal: -```json -{"type": "message", "content": {"type": "ai", "content": "", "tool_calls": [...]}} -{"type": "token", "content": "Hello"} -{"type": "token", "content": " world"} -{"type": "message", "content": {"type": "ai", "content": "Hello world"}} -[DONE] +```bash +curl http://localhost:5002/health ``` -### Client Examples - -Ready-to-use client examples are available in the [`examples/`](./examples/) directory: - -- **[Streamlit Demo App](./examples/streamlit_app.py)** - Interactive chat application -- **[Python Async Client](./examples/client_python.py)** - Server-to-server integration - -See the [examples README](./examples/README.md) for detailed usage instructions. - -## 🚀 Quick Start - -### Prerequisites - -- Python 3.12+ -- PostgreSQL database -- Google AI API credentials -- Langfuse account (optional) - -### Installation +Copy `.env.example` to `.env` before first run (or let `make local` create it) and set `GOOGLE_APPLICATION_CREDENTIALS_CONTENT`. -1. **Clone the repository** - ```bash - git clone https://github.com/redhat-data-and-ai/template-agent.git - cd template-agent - ``` +**MCP and UI are separate repos** — this project runs the agent and its dependencies (Postgres, Redis) only. Clone and run [template-mcp-server](https://github.com/redhat-data-and-ai/template-mcp-server) and [template-ui](https://github.com/redhat-data-and-ai/template-ui) when needed. -2. **Create virtual environment** - ```bash - uv venv - source .venv/bin/activate +## API - ``` +The agent exposes the standard **LangGraph API** (assistant ID: `agent`, defined in `aegra.json`) plus custom routes on the Aegra HTTP app. -3. **Install dependencies** - ```bash - uv pip install -e ".[dev]" - ``` +### LangGraph API -4. **Set up environment variables** - ```bash - cp .env.example .env - # Edit .env with your configuration - ``` - -5. **Run template-mcp-server** following https://github.com/redhat-data-and-ai/template-mcp-server - - -6. **Run the application** - ```bash - uv run python -m template_agent.src.main - ``` - - -## 📚 API Reference - -### Endpoints - -| Endpoint | Method | Description | -|---------------------------|--------|-------------| -| `/health` | GET | Health check | -| `/v1/stream` | POST | Stream chat responses | -| `/v1/history/{thread_id}` | GET | Get conversation history | -| `/v1/threads/{user_id}` | GET | List user threads | -| `/v1/feedback` | POST | Record feedback | - -### Streaming Chat +| Endpoint | Method | Description | +|---|---|---| +| `/ok` | GET | Server health | +| `/assistants/{assistant_id}` | GET | Assistant metadata | +| `/threads` | POST | Create conversation thread | +| `/threads/{thread_id}` | GET | Get thread state | +| `/threads/{thread_id}/runs` | POST | Run agent (sync) | +| `/threads/{thread_id}/runs/stream` | POST | Run agent (SSE stream) | ```bash -curl -X POST "http://localhost:8081/v1/stream" \ +# Create a thread +curl -X POST http://localhost:5002/threads \ + -H "Content-Type: application/json" \ + -d '{}' + +# Stream a message (replace THREAD_ID) +curl -N -X POST "http://localhost:5002/threads/THREAD_ID/runs/stream" \ -H "Content-Type: application/json" \ -d '{ - "message": "Hello, how can you help me?", - "thread_id": "thread_123", - "user_id": "user_456", - "stream_tokens": true + "assistant_id": "agent", + "input": {"messages": [{"role": "human", "content": "Hello"}]}, + "stream_mode": "updates" }' ``` -### Health Check +### Custom routes -```bash -curl "http://localhost:8081/health" -# Response: {"status": "healthy", "service": "Template Agent"} -``` +| Endpoint | Method | Description | +|---|---|---| +| `/health` | GET | Health check (also `/healthz`, `/readyz`, `/livez`) | +| `/info` | GET | Agent name and OAuth/DCR MCP server list | +| `/feedback` | POST | Record user feedback (Langfuse + Postgres) | +| `/feedback/{thread_id}` | GET | List feedback for a thread | +| `/threads/{thread_id}/token-usage` | GET | Cumulative token usage for a thread | +| `/mcp/{name}/connect` | POST | Start OAuth/DCR flow for an MCP server | +| `/mcp/oauth/callback` | GET | OAuth redirect handler | +| `/mcp/{name}/status` | GET | MCP connection status for current user | -## ⚙️ Configuration +Use [template-ui](https://github.com/redhat-data-and-ai/template-ui) for a full chat experience against this API. -### Environment Variables +## Configuration -#### Required -- `AGENT_HOST`: Server host (default: 0.0.0.0) -- `AGENT_PORT`: Server port (default: 5002) -- `PYTHON_LOG_LEVEL`: Logging level (default: INFO) +Configuration is split between **secrets/endpoints** (`.env`) and **operational settings** (`config/agent/runtime/agent.yaml`). -#### Database -- `POSTGRES_USER`: Database username (default: pgvector) -- `POSTGRES_PASSWORD`: Database password (default: pgvector) -- `POSTGRES_DB`: Database name (default: pgvector) -- `POSTGRES_HOST`: Database host (default: pgvector) -- `POSTGRES_PORT`: Database port (default: 5432) +### Environment variables (`.env`) -#### Optional -- `LANGFUSE_PUBLIC_KEY`: Langfuse public key for tracing -- `LANGFUSE_SECRET_KEY`: Langfuse secret key for tracing -- `LANGFUSE_BASE_URL`: Langfuse host URL (e.g., https://cloud.langfuse.com) -- `LANGFUSE_TRACING_ENVIRONMENT`: Langfuse environment (default: development) -- `GOOGLE_SERVICE_ACCOUNT_FILE`: Google credentials -- `AGENT_SSL_KEYFILE`: SSL private key path -- `AGENT_SSL_CERTFILE`: SSL certificate path +| Variable | Default | Description | +|---|---|---| +| `POSTGRES_HOST` | `localhost` | Postgres host (`pgvector` in compose) | +| `POSTGRES_PORT` | `5432` | Postgres port | +| `POSTGRES_DB` | `template_agent` | Database name | +| `POSTGRES_USER` | `postgres` | Database user | +| `POSTGRES_PASSWORD` | `postgres` | Database password | +| `REDIS_URL` | `redis://localhost:6379/0` | Redis URL (required for OAuth/DCR MCPs) | +| `REDIS_BROKER_ENABLED` | `true` | Enable Redis-backed SSE broker | +| `GOOGLE_APPLICATION_CREDENTIALS_CONTENT` | — | Google service account JSON (required) | +| `ENABLE_AUTH` | `false` in `.env.example` | SSO/OIDC authentication | +| `SSO_ISSUER_URL` | — | OIDC issuer (Keycloak, Okta, etc.) | +| `SSO_CLIENT_ID` | — | OIDC client ID | +| `SSO_CLIENT_SECRET` | — | OIDC client secret | +| `LANGFUSE_PUBLIC_KEY` | — | Langfuse public key (optional) | +| `LANGFUSE_SECRET_KEY` | — | Langfuse secret key (optional) | +| `LANGFUSE_BASE_URL` | — | Langfuse host (optional) | +| `LANGFUSE_TRACING_ENVIRONMENT` | `development` | Langfuse environment label | +| `MCP_TOKEN_ENCRYPTION_KEY` | — | Fernet key for OAuth/DCR token encryption | +| `MCP_TOKEN_ENCRYPTION_KEY_PREVIOUS` | — | Previous key during rotation (decrypt-only) | +| `AGENT_PUBLIC_BASE_URL` | `http://localhost:5002` | Public agent URL for OAuth callbacks | +| `CUSTOM_CA_FILE` | — | Host path to a PEM file with custom CA certs (compose only) | +| `SSL_KEYFILE` | — | TLS private key path (optional) | +| `SSL_CERTFILE` | — | TLS certificate path (optional) | -### Configuration Example +See [`.env.example`](./.env.example) for the full list including OpenTelemetry and MongoDB token-usage settings. -```bash -# .env file -AGENT_HOST=0.0.0.0 -AGENT_PORT=5002 -PYTHON_LOG_LEVEL=INFO - -POSTGRES_USER=myuser -POSTGRES_PASSWORD=mypassword -POSTGRES_DB=template_agent -POSTGRES_HOST=localhost -POSTGRES_PORT=5432 - -LANGFUSE_TRACING_ENVIRONMENT=production -GOOGLE_SERVICE_ACCOUNT_FILE=/path/to/credentials.json -``` +Runtime settings (cache, memory, providers, middleware, agent identity) live in [`config/agent/runtime/agent.yaml`](./config/agent/runtime/agent.yaml). -## 🧪 Testing +## MCP Server Configuration -### Run Tests +MCP servers are defined in [`config/agent/mcp.json`](./config/agent/mcp.json) and attached to agents via the `mcps` frontmatter field in [`config/agent/PROMPT.md`](./config/agent/PROMPT.md) (orchestrator) or [`config/agent/subagents/*.md`](./config/agent/subagents/). -```bash -# Run all tests -pytest +### Auth modes -# Run with coverage -pytest --cov=template_agent.src --cov-report=html +| `auth_mode` | When to use | How credentials work | +|---|---|---| +| `sso` (default) | MCP accepts the same SSO token as the agent | User Bearer token forwarded on every tool call | +| `oauth` | MCP has a pre-registered OAuth client | User connects via chat UI; tokens stored encrypted in Redis | +| `dcr` | MCP supports OAuth Dynamic Client Registration | Agent registers at connect; per-user OAuth flow follows | -# Run specific test file -pytest tests/test_prompt.py -v -``` +Set `"auth": false` for public/local MCP servers with no Authorization header. -### Test Coverage +### MCP URL by run mode -Current test coverage includes: -- ✅ Core utilities (prompt, agent_utils) -- ✅ Data models (schema) -- ✅ Configuration (settings) -- ✅ API endpoints (health, feedback) -- 🔄 Complex routes (history, stream, threads) -- 🔄 Application setup (api, main, agent) +| Mode | `url` in `mcp.json` | +|---|---| +| `make local` (agent on host) | `http://localhost:5001/mcp` (default) | +| `make container` (MCP on host) | `http://host.containers.internal:5001/mcp` | -## 🚀 Deployment +Alternate URLs are provided as `//` comments in `mcp.json` — uncomment the line you need. -### Podman Compose +### Wiring MCPs to agents -```bash -# Start with Docker Compose -podman-compose up -d --build +```yaml +--- +name: analyst +model: gemini-2.5-pro +mcps: + - template-mcp-server +tools: + - calculate_bmi + - search_web +--- ``` -### Production Considerations - -- **SSL/TLS**: Configure SSL certificates for HTTPS -- **Database**: Use managed PostgreSQL service -- **Monitoring**: Set up Langfuse for tracing -- **Scaling**: Configure horizontal pod autoscaling -- **Security**: Implement proper authentication +- **Orchestrator:** add `mcps:` to `config/agent/PROMPT.md` frontmatter. +- **Subagent:** add `mcps:` to `config/agent/subagents/.md` frontmatter. +- **Inheritance:** subagents without `mcps` inherit the orchestrator's list. +- **Validation:** every name in `mcps` must exist in `mcp.json` with `enabled: true`. -## 🔧 Development +See `config/agent/mcp.json` for working SSO and DCR examples. -### Project Structure +## Project Structure ``` template-agent/ -├── template_agent/ -│ └── src/ -│ ├── core/ # Core agent functionality -│ │ ├── agent.py # Agent initialization -│ │ ├── agent_utils.py # Message utilities -│ │ └── prompt.py # Prompt management -│ ├── routes/ # API endpoints -│ │ ├── health.py # Health checks -│ │ ├── stream.py # Streaming chat -│ │ ├── history.py # Chat history -│ │ ├── threads.py # Thread management -│ │ └── feedback.py # Feedback recording -│ ├── api.py # FastAPI application -│ ├── main.py # Application entry point -│ ├── schema.py # Data models -│ └── settings.py # Configuration -├── tests/ # Test suite -└── README.md # This file +├── aegra.json # Aegra / LangGraph framework entry point +├── config/agent/ +│ ├── PROMPT.md # Orchestrator prompt + frontmatter +│ ├── subagents/ # Subagent definitions +│ ├── skills/ # Skill documents and evals +│ ├── mcp.json # MCP server registry +│ ├── runtime/agent.yaml # Runtime config (cache, memory, providers) +│ └── deployment/values.yaml # OpenShift/ArgoCD deployment reference values +├── deep_agent/ +│ ├── aegra/ # Graph, HTTP app, MCP OAuth, entrypoint +│ └── src/ # Config loader, cache, memory, token budget, etc. +├── tests/ +│ ├── unit/ # Unit tests +│ ├── integration/ # Aegra integration and e2e tests +│ └── skills/ # LLM-as-judge skill evaluations +├── compose.yaml # Postgres + Redis (+ agent with --profile container) +├── Containerfile +└── deployment/ # OpenShift and Kind overlays ``` -### Code Quality +## Testing ```bash -# Run linting -ruff check . - -# Run type checking -mypy template_agent/src/ +make test # unit tests +make test-all # unit + skills evals +make test-skills # skills evaluations only +make test-cov # unit tests with coverage +``` -# Run formatting -ruff format . +Skills evals auto-discover from `config/agent/skills/*/evals/evals.json`. See [`config/agent/evals/README.md`](./config/agent/evals/README.md) for Promptfoo and Lightspeed eval options. -# Run pre-commit hooks +```bash +# Code quality +ruff check . && ruff format . pre-commit run --all-files ``` -### Adding New Features - -1. **Create feature branch** - ```bash - git checkout -b feature/new-feature - ``` +## Custom CA Certificates -2. **Implement changes** - - Follow Google docstring format - - Add type hints - - Write tests for new functionality +If your environment uses a corporate or internal certificate authority, the container can trust it at startup without rebuilding the image. -3. **Run quality checks** - ```bash - pre-commit run --all-files - pytest - ``` +**Compose** — set `CUSTOM_CA_FILE` in `.env` to the host path of your PEM bundle: -4. **Submit pull request** - - Include tests - - Update documentation - - Follow commit message conventions +```bash +# .env +CUSTOM_CA_FILE=./certs/ca.pem +``` -### Development Setup +**Kubernetes** — create a Secret and mount it, then set `CUSTOM_CA_PATH`: + +```yaml +env: + - name: CUSTOM_CA_PATH + value: /etc/custom-ca/ca.pem +volumeMounts: + - name: custom-ca + mountPath: /etc/custom-ca + readOnly: true +volumes: + - name: custom-ca + secret: + secretName: custom-ca +``` -1. Fork the repository -2. Create a feature branch -3. Make your changes -4. Add tests for new functionality -5. Ensure all tests pass -6. Submit a pull request +**Fallback URL** — for non-orchestrated environments (e.g. `podman run`), set `CUSTOM_CA_URL` to download the PEM at startup: -### Code Standards +```bash +podman run -e CUSTOM_CA_URL=https://certs.example.com/ca.pem ... +``` -- **Python**: Follow PEP 8 and use type hints -- **Documentation**: Use Google docstring format -- **Tests**: Maintain >80% code coverage -- **Commits**: Use conventional commit messages +If neither variable is set, or the download fails, the container starts normally with default system certs. -This template includes `.cursor/rules.md` - a comprehensive development guide specifically designed to help AI coding assistants understand and work effectively with this MCP server template. +## Deployment -### What's Included +```bash +make container +``` -## 🆘 Support +For production: configure TLS (`SSL_KEYFILE`, `SSL_CERTFILE`), use managed PostgreSQL and Redis, set `AGENT_PUBLIC_BASE_URL` to your HTTPS URL, and enable Langfuse tracing. -- **Issues**: [GitHub Issues](https://github.com/redhat-data-and-ai/template-agent/issues) +OpenShift manifests are in `deployment/overlays/openshift/`. Local full-stack Kubernetes testing: `make kind`. -## 🔗 Related Projects +## Links -- [LangChain](https://github.com/langchain-ai/langchain) - LLM application framework -- [LangGraph](https://github.com/langchain-ai/langgraph) - Stateful LLM applications -- [FastAPI](https://fastapi.tiangolo.com/) - Modern web framework -- [Langfuse](https://langfuse.com/) - LLM observability platform +- [Issues](https://github.com/redhat-data-and-ai/template-agent/issues) +- [template-mcp-server](https://github.com/redhat-data-and-ai/template-mcp-server) +- [template-ui](https://github.com/redhat-data-and-ai/template-ui) +- [LangGraph docs](https://langchain-ai.github.io/langgraph/) ---- +## License -**Built with ❤️ by the Red Hat Data & AI team** +[Apache 2.0](LICENSE) diff --git a/aegra.json b/aegra.json new file mode 100644 index 00000000..b00aeff4 --- /dev/null +++ b/aegra.json @@ -0,0 +1,21 @@ +{ + "$schema": "https://langgra.ph/schema.json", + "dependencies": ["."], + "graphs": { + "agent": "./deep_agent/aegra/graph.py:agent" + }, + "auth": { + "path": "./deep_agent/aegra/auth.py:auth" + }, + "env": ".env", + "python_version": "3.12", + "http": { + "app": "./deep_agent/aegra/http_app.py:app", + "cors": { + "allow_origins": ["http://localhost:3000", "http://127.0.0.1:3000", "http://localhost:8080", "http://127.0.0.1:8080"], + "allow_methods": ["*"], + "allow_headers": ["*"], + "allow_credentials": true + } + } +} diff --git a/compose.yaml b/compose.yaml index 04e1e03a..478f587c 100644 --- a/compose.yaml +++ b/compose.yaml @@ -1,57 +1,138 @@ +## Agent compose stack — agent + dependencies only (no MCP, no UI) +## +## make local - pgvector + redis in compose; agent process on host +## make container - pgvector + redis + agent in compose +## make dev - same as container, detached + log tail +## +## MCP and UI run from their own repos (template-mcp-server, template-ui). +## +## Demo profile: Adds UI + MCP Server with SSO authentication +## make demo - Clone repos, configure, start full stack +## make clean - Stop stack, remove data + cloned repos +## +## Observability profile: Adds Jaeger (also enabled by `make container`) +## docker compose --profile observability up +## Set ENABLE_OTEL=true in .env to enable metrics/tracing export +## +## Services (always): +## pgvector - Postgres (agent checkpoints) +## redis - Aegra broker (SSE streaming, job queue) +## template-agent - Agent (port 5002) +## +## Services (demo profile): +## template-mcp-server - MCP server with SSO auth (port 5001) +## template-ui - React/Fastify frontend with SSO auth (port 8080) +## +## Services (observability profile): +## jaeger - Jaeger UI for trace visualization (UI :16686) + services: pgvector: image: ankane/pgvector container_name: template-agent-pgvector environment: - POSTGRES_DB: pgvector - POSTGRES_USER: pgvector - POSTGRES_PASSWORD: pgvector + POSTGRES_DB: template_agent + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres ports: - "5432:5432" volumes: - - pgvector_data:/var/lib/postgresql/data - restart: always + - agent_pgvector_data:/var/lib/postgresql/data + restart: unless-stopped healthcheck: - test: ["CMD-SHELL", "pg_isready -U pgvector -d pgvector"] + test: ["CMD-SHELL", "pg_isready -U postgres -d template_agent"] interval: 10s timeout: 5s retries: 5 start_period: 30s networks: - - template-network + - agent-network + + redis: + image: redis:7-alpine + container_name: template-agent-redis + command: redis-server --appendonly yes + ports: + - "6379:6379" + volumes: + - agent_redis_data:/data + restart: unless-stopped + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 10s + timeout: 3s + retries: 5 + networks: + - agent-network template-agent: + profiles: [container] build: context: . dockerfile: Containerfile container_name: template-agent + extra_hosts: + - "host.containers.internal:host-gateway" ports: - - "8081:8081" + - "5002:5002" env_file: - .env environment: - - AGENT_PORT=8081 + - ENABLE_AUTH=true - POSTGRES_HOST=pgvector - POSTGRES_PORT=5432 - - POSTGRES_DB=pgvector - - POSTGRES_USER=pgvector - - POSTGRES_PASSWORD=pgvector + - POSTGRES_DB=template_agent + - POSTGRES_USER=postgres + - POSTGRES_PASSWORD=postgres + - REDIS_URL=redis://redis:6379/0 + - REDIS_BROKER_ENABLED=true + # Custom CA — mount a PEM at ./certs/ca.pem to trust corporate CAs + - CUSTOM_CA_PATH=${CUSTOM_CA_PATH:-/etc/custom-ca/ca.pem} + # OTEL (enable via env — make container sets these when Jaeger profile is active) + - ENABLE_OTEL=${ENABLE_OTEL:-false} + - OTEL_EXPORTER_OTLP_ENDPOINT=${OTEL_EXPORTER_OTLP_ENDPOINT:-http://jaeger:4317} + - ENABLE_OTEL_TRACES=${ENABLE_OTEL_TRACES:-false} + - OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=${OTEL_EXPORTER_OTLP_TRACES_ENDPOINT:-http://jaeger:4317} depends_on: pgvector: condition: service_healthy - restart: always + redis: + condition: service_healthy + restart: unless-stopped healthcheck: - test: [ "CMD", "curl", "-f", "-k", "http://0.0.0.0:8081/health"] - interval: 5s + test: ["CMD", "curl", "-f", "http://0.0.0.0:5002/health"] + interval: 10s timeout: 5s retries: 5 + start_period: 30s + volumes: + - ./deep_agent:/app/deep_agent:ro + - ./config:/app/config:ro + - ${CUSTOM_CA_FILE:-/dev/null}:/etc/custom-ca/ca.pem:ro networks: - - template-network + - agent-network + + # ── Observability Stack ────────────────────────────────────────────── + + jaeger: + profiles: [observability] + image: jaegertracing/all-in-one:1.55 + container_name: demo-jaeger + ports: + - "16686:16686" # Jaeger UI + - "4317" # OTLP gRPC receiver (internal only) + environment: + - COLLECTOR_OTLP_ENABLED=true + restart: unless-stopped + networks: + - demo-network volumes: - pgvector_data: + agent_pgvector_data: + driver: local + agent_redis_data: driver: local networks: - template-network: + agent-network: driver: bridge diff --git a/config/agent/HEADLESS_PROMPT.md b/config/agent/HEADLESS_PROMPT.md new file mode 100644 index 00000000..412e590c --- /dev/null +++ b/config/agent/HEADLESS_PROMPT.md @@ -0,0 +1,61 @@ +--- +name: headless-worker +description: > + Background task processor. Receives tasks from Redis queue, + processes them by delegating to subagents, and returns results. + No user interaction — silent worker. +model: gemini-2.5-pro +tools: + - calculate_bmi + - search_web +skills: + - bmi-report +--- + +# Background Task Processor + +You are a background task processor. You receive tasks from a queue and process them silently. + +## Rules + +1. **No greetings, no TODO lists, no conversational responses.** You are not talking to a user. +2. **Process the payload directly.** Extract the task details and execute them. +3. **Return structured results.** Your output is stored in Redis for the orchestrator to retrieve. +4. **Delegate to your tools.** Use `calculate_bmi` and `search_web` as needed. +5. **Handle errors gracefully.** If something fails, return a clear error message. + +## Task Processing + +When you receive a task payload: + +1. Parse the task name and data from the payload +2. Execute the work using your tools +3. Return a JSON-structured result with: + - `status`: "success" or "error" + - `summary`: Brief description of what was done + - `data`: The actual results (BMI values, reports, etc.) + - `error`: Error message if status is "error" + +## Example + +Input payload: +```json +{"name": "batch-bmi", "task_id": "abc123", "payload": "{\"employees\": [{\"name\": \"Alice\", \"weight_kg\": 65, \"height_cm\": 170}]}"} +``` + +Expected output: +```json +{ + "status": "success", + "summary": "Processed 1 BMI calculation", + "data": [{"name": "Alice", "bmi": 22.5, "category": "Normal"}] +} +``` + +## Scope + +- BMI calculations (single or batch) +- Health reports and summaries +- Any task delegated by the orchestrator + +Do NOT ask for more information. Do NOT create TODO lists. Do NOT greet anyone. Just process and return results. diff --git a/config/agent/PROMPT.md b/config/agent/PROMPT.md new file mode 100644 index 00000000..7e5467d7 --- /dev/null +++ b/config/agent/PROMPT.md @@ -0,0 +1,207 @@ +--- +name: orchestrator +description: > + Main coordinator for Red Hat fitness assistant. Handles client intake, + routes to analyst and publisher subagents, manages TODO lists and + delegates health metric analysis. +model: gemini-2.5-pro +tools: + - validate_email + - queue_task + - check_task_status + - get_pending_results +skills: + - client-intake +--- + +# Red Hat Fitness Assistant + +Today's date is {{current_date}}. + +## Identity + +You are a friendly fitness assistant for Red Hat employees. + +**CRITICAL: You are an ORCHESTRATOR, not an analyst.** +- You COORDINATE work by delegating to subagents +- You NEVER calculate BMI yourself +- You NEVER analyze health data yourself +- You NEVER provide health tips yourself +- You ALWAYS delegate analysis to the analyst subagent +- You VALIDATE email addresses using the validate_email tool before delegating to publisher + +## Control Flow & Routing + +```mermaid +flowchart TD + User([User]) --> Orch + + subgraph Orch["Orchestrator (you) — tool: validate_email, skill: client-intake"] + Classify{Classify intent} + end + + Classify -->|Out-of-scope| Decline[Decline with reason] + Classify -->|Multi-step| TODO[Break into TODO items\nroute each in-scope step] + Classify -->|Health metrics| Imperial{Imperial units?} + + Imperial -->|YES| Convert[Convert via\nclient-intake skill] + Imperial -->|NO| BA + + Convert --> BA + + TODO -.->|in-scope steps| Imperial + + subgraph BA["① analyst — skill: bmi-report"] + BA_Tools[tools: calculate_bmi, search_web] + end + + BA --> Email{Email requested?} + + Email -->|NO| Return[Return analysis\nto user] + Email -->|YES| RD + + subgraph RD["② publisher — skill: email-formatter"] + RD_Tools[tool: send_email] + end + + RD --> Sent[Email sent] +``` + +**Key constraints:** +- **TODO list ALWAYS comes first** — For ALL requests (simple or complex), create a TODO list BEFORE starting any work. This ensures proper planning and tracking. +- **Simple requests** — Single-task TODO list with one item (e.g., "analyze my BMI"). +- **Multi-step requests** — Multi-item TODO list with all tasks planned upfront. +- Step ② (publisher) must never be invoked until **all** other subagents have completed their tasks. +- The orchestrator owns all sequencing — subagents never call each other. + +### Routing Table + +| User Intent | Path through diagram | Action | +|-------------|----------------------|--------| +| Health metrics (height, weight, BMI) | TODO → Health metrics → ① | **Create TODO list first** with single item. Greet user. If imperial units (ft, in, lbs), convert to metric using **exactly** the formulas in the **client-intake** skill — do not write your own conversion code. Then delegate to **analyst** with cm and kg. | +| Health metrics + email request | TODO → Health metrics → ① → barrier → ② | **Create TODO list first** with all steps. Greet user. Use **validate_email** tool to verify the recipient email address. If invalid, inform the user and ask for a valid email. Delegate to **analyst** first. Only after it completes, delegate to **publisher** with the analysis results and recipient address. | +| Quick BMI without email | TODO → Health metrics → ① → return | **Create TODO list first** with single item. Greet user. Delegate to **analyst**; skip publisher. Return analysis directly to user. | +| Multi-step requests | TODO → Per-item routing | **Create TODO list first** with all items. Include out-of-scope items marked as **"Declined — [reason]"** so the user sees them acknowledged. Route the remaining in-scope steps through the diagram above. | +| Out-of-scope requests | Left branch (decline) | Explain politely why the request is out of scope and what you *can* do. | + +## Delegation (CRITICAL) + +**YOU MUST DELEGATE. YOU CANNOT DO THE WORK YOURSELF.** + +When a user requests BMI analysis: +1. **CREATE TODO LIST FIRST** — Always start by creating a TODO list with the task(s) +2. Greet them: "Welcome! I'm your Red Hat fitness assistant." +3. If email delivery is requested, **validate the email address** using the validate_email tool +4. Convert units if needed (imperial → metric) +5. **DELEGATE to analyst subagent** with height (cm) and weight (kg) +6. Wait for analyst's response +7. If email was requested and valid, delegate to publisher; otherwise return results directly +8. Relay analyst's results to the user + +**FORBIDDEN ACTIONS:** +- Do NOT calculate BMI yourself (you don't have the calculate_bmi tool) +- Do NOT determine BMI category yourself +- Do NOT provide health tips yourself +- Do NOT describe what you plan to do — just delegate + +**CORRECT:** +``` +[create TODO list with task: "Analyze BMI for user"] +Welcome! I'm your Red Hat fitness assistant. +[delegate to analyst with height=175, weight=70] +[relay analyst's BMI analysis to user] +``` + +**WRONG:** +``` +Your BMI is 22.9, which is in the Normal category. +Here are some health tips... [providing tips yourself] +``` + +**ALSO WRONG (missing TODO list):** +``` +Welcome! I'm your Red Hat fitness assistant. +[delegate to analyst with height=175, weight=70] ← Missing TODO list creation first! +``` + +## Background Tasks (Headless Worker) + +A headless worker runs alongside you as a background processor. Use `queue_task` to delegate work that is long-running, bulk, or doesn't need an immediate response. + +**When to use queue_task:** +- Bulk operations (e.g., "generate reports for all 500 clients") +- Long-running processing (e.g., "retrain the model", "export all data") +- Fire-and-forget notifications (e.g., "send weekly emails to all users") + +**When NOT to use queue_task:** +- Single BMI calculations — delegate to analyst as usual +- Anything the user expects an immediate answer to + +**How it works:** +1. Call `queue_task(task_name="descriptive-name", payload={...})` to queue the work +2. The headless worker picks it up from Redis and processes it asynchronously +3. Results go to the configured output sinks (file, webhook, Redis) +4. Tell the user: "I've queued [task]. It will be processed in the background." + +**Status tracking — CRITICAL RULES:** +1. `queue_task` returns a task ID — give this to the user +2. When the user asks about task status or results, you MUST call `check_task_status(task_id)` — do NOT answer from memory or guess. The tool returns the full result including data. +3. When `check_task_status` returns a COMPLETED task with results, you MUST show the complete result to the user. Never say "results are in a file" or "check a dashboard" — the result IS in the tool response. Display it directly. +4. **At the start of every conversation**, call `get_pending_results(user_id)` to check for completed background tasks. If any exist, show the full results to the user before handling their new request. +5. Never make up task status. Always use the tool. + +## Code Execution + +You have access to the `execute_code` tool which runs code in an isolated sandbox. **This is the ONE exception to the delegation rule — you call execute_code YOURSELF, never delegate it to a subagent.** + +**Use it automatically** whenever a task involves: +- **Computation**: math, statistics, data analysis, aggregation +- **Data processing**: parsing, transforming, filtering data +- **Verification**: checking a formula, validating a calculation, testing a hypothesis +- **Generation**: creating structured output (CSV, JSON, tables) from raw data +- **Visualization**: ASCII charts, formatted tables, data summaries + +**Workflow with subagents**: Delegate domain work (BMI analysis, email) to subagents as usual. Then use `execute_code` yourself to compute, visualize, or process the results. Example: delegate BMI to analyst → get result → use execute_code to create a visualization. + +**Fallback**: If the analyst subagent fails or is unavailable (MCP tools not connected), use `execute_code` directly to compute BMI yourself. The formula is: BMI = weight_kg / (height_m ** 2). + +Do NOT ask the user whether to run code — just write and execute it. Default to Python unless the user specifies otherwise. The tool supports `python`, `shell`, and `node`. + +**When NOT to use it**: simple factual questions, conversational responses, or tasks the LLM can answer accurately from knowledge (e.g., "what is Python?"). + +## General Behavior + +- Always respond in the same language as the user. +- Ensure all string values in function call arguments are properly JSON-escaped. +- Only use the tools you are given. Do not answer from internal knowledge when a tool can provide the answer. +- Every final answer must be grounded in tool observations. + +## Output Format + +- Always respond using proper Markdown formatting. +- Use headers, lists, code blocks, bold, and tables when they improve readability. +- Keep intermediate responses concise; make the final response well-structured. + +## Scope + +This system produces a **one-time snapshot**: today's BMI and category-specific +health tips. It does not plan, prescribe, or track anything over time. + +## Out of Scope + +- Diet plans, meal plans, or food recommendations. +- Exercise or workout routines. +- Weight history, trends, or progress tracking. +- Goal weight or target BMI calculations. +- Medical diagnosis or treatment advice. + +Politely decline each out-of-scope item and explain what you *can* do. + +## Gotchas + +- **TODO list ALWAYS comes first** — Never start any work without creating a TODO list, even for simple single-task requests. +- **Never compute BMI or format emails yourself** — always delegate to the appropriate subagent. +- **Route to publisher only after all other subagents complete** — never in parallel with upstream work. +- **Don't assume measurements** — if height or weight is missing, ask before routing. +- **Always convert imperial to metric before delegating** — use the exact formulas from the **client-intake** skill. Do not improvise conversion code. analyst expects cm and kg only. +- **Always validate email addresses** — use the validate_email tool before delegating to publisher. If invalid, ask the user for a valid email address. diff --git a/config/agent/deployment/values.yaml b/config/agent/deployment/values.yaml new file mode 100644 index 00000000..22547516 --- /dev/null +++ b/config/agent/deployment/values.yaml @@ -0,0 +1,119 @@ +# Deployment configuration for the agent. +# Reference values for ArgoCD/Helm-style deployments. +# Kustomize overlays in deployment/overlays/ use their own patches. +# +# ArgoCD Vault Plugin (AVP) injects secrets at deploy time. + +app: + name: agent + component: agent + replicas: 2 + namespace: ai-agents + +image: + name: agent + tag: latest + registry: image-registry.openshift-image-registry.svc:5000 + +container: + port: 5002 + +resources: + requests: + memory: "512Mi" + cpu: "250m" + limits: + memory: "1Gi" + cpu: "1000m" + +# --- Horizontal Pod Autoscaler --- +autoscaling: + enabled: true + minReplicas: 2 + maxReplicas: 8 + targetCPUUtilizationPercentage: 70 + targetMemoryUtilizationPercentage: 80 + scaleDown: + stabilizationWindowSeconds: 300 + +# --- Health probes --- +probes: + liveness: + path: /health + initialDelaySeconds: 30 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + readiness: + path: /health + initialDelaySeconds: 10 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 3 + startup: + path: /health + initialDelaySeconds: 5 + periodSeconds: 5 + failureThreshold: 12 + +# --- Networking --- +networking: + service: + type: ClusterIP + port: 5002 + route: + enabled: true + tls: + termination: edge + insecureEdgeTerminationPolicy: Redirect + networkPolicy: + enabled: true + ingress: + - from: + - podSelector: + matchLabels: + app: template-ui + ports: + - port: 5002 + +# --- Build --- +build: + resources: + requests: + memory: "2Gi" + cpu: "1000m" + limits: + memory: "4Gi" + cpu: "4" + +# --- Non-sensitive configuration (maps to ConfigMap) --- +config: + AGENT_HOST: "0.0.0.0" + AGENT_PORT: "5002" + PYTHON_LOG_LEVEL: "INFO" + LANGFUSE_TRACING_ENVIRONMENT: "production" + REQUEST_LOGGING_ENABLED: "true" + REQUEST_LOG_HEADERS: "true" + REQUEST_LOG_BODY: "false" + REQUEST_LOG_BODY_MAX_SIZE: "10240" + REDIS_URL: "redis://redis:6379/0" + +# --- Sensitive configuration (maps to Secret) --- +# Values are injected by the ArgoCD Vault Plugin (AVP) at deploy time. +# Use placeholder syntax. +secrets: + POSTGRES_HOST: "//agent#POSTGRES_HOST>" + POSTGRES_PORT: "//agent#POSTGRES_PORT>" + POSTGRES_DB: "//agent#POSTGRES_DB>" + POSTGRES_USER: "//agent#POSTGRES_USER>" + POSTGRES_PASSWORD: "//agent#POSTGRES_PASSWORD>" + SSO_ISSUER_URL: "//agent#SSO_ISSUER_URL>" + SSO_CLIENT_ID: "//agent#SSO_CLIENT_ID>" + SSO_CLIENT_SECRET: "//agent#SSO_CLIENT_SECRET>" + LANGFUSE_PUBLIC_KEY: "//agent#LANGFUSE_PUBLIC_KEY>" + LANGFUSE_SECRET_KEY: "//agent#LANGFUSE_SECRET_KEY>" + LANGFUSE_BASE_URL: "//agent#LANGFUSE_BASE_URL>" + GOOGLE_APPLICATION_CREDENTIALS_CONTENT: "//agent#GOOGLE_APPLICATION_CREDENTIALS_CONTENT>" + VLLM_BASE_URL: "//agent#VLLM_BASE_URL>" + VLLM_API_KEY: "//agent#VLLM_API_KEY>" + REDIS_URL: "//agent#REDIS_URL>" diff --git a/config/agent/evals/README.md b/config/agent/evals/README.md new file mode 100644 index 00000000..dc4fde8b --- /dev/null +++ b/config/agent/evals/README.md @@ -0,0 +1,159 @@ +# Agent Evaluations + +This directory contains evaluation suites for testing the deep-agent BMI fitness assistant. + +## Overview + +We use multiple evaluation frameworks: + +1. **Skills Evals** - Test individual skills (client-intake, bmi-report, email-formatter) +2. **Promptfoo** - Fast iteration testing with LLM-rubric assertions +3. **Lightspeed** - Formal benchmark evaluation (optional) + +## Running Evaluations Locally + +### Prerequisites + +- Agent must be running at `http://localhost:5002` +- Environment variables set: + - `GOOGLE_GENAI_API_KEY` - for LLM calls and judging + - `GOOGLE_APPLICATION_CREDENTIALS_CONTENT` - for Google service account + +### 1. Skills Evals (Pytest-based) + +Tests individual skills using LLM-as-judge: + +```bash +# Run all skills evals +make test-skills + +# Or use pytest directly +pytest tests/skills -m skills -v +``` + +Results are saved to `tests/workspaces//eval-/` + +### 2. Promptfoo Agent Evals + +Fast iteration testing with assertion-based evaluation: + +```bash +# Start the agent first +make local + +# In another terminal, run evals +make eval-promptfoo + +# Or run directly +cd config/agent/evals/promptfoo +npx promptfoo@latest eval + +# View results in browser +npx promptfoo@latest view +``` + +**What it tests:** +- BMI calculation delegation to analyst subagent +- Imperial unit conversion +- Health tips by BMI category +- Email delivery validation +- Out-of-scope request handling +- Edge cases (missing data, invalid input) + +### 3. Lightspeed Evals (Optional) + +Formal benchmark evaluation (requires additional setup): + +```bash +# Install lightspeed-evaluation +pip install lightspeed-evaluation + +# Run evals +lightspeed-eval \ + --system-config config/agent/evals/lightspeed/system.yaml \ + --eval-data config/agent/evals/lightspeed/eval_data.yaml \ + --output-dir eval_output +``` + +## CI/CD + +Evals run automatically in GitHub Actions: + +- **Unit Tests** - Run on every PR/push +- **Skills Evals** - Run on every PR/push +- **Promptfoo Agent Evals** - Run on every PR/push + +See `.github/workflows/test.yml` for details. + +## Eval Structure + +### Promptfoo Config + +```yaml +providers: + - http endpoint to agent +tests: + - description: Test case name + vars: + prompt: User input + assert: + - type: llm-rubric | contains | not-contains + value: Expected behavior +``` + +### Lightspeed Config + +```yaml +conversation_group_id: test_scenario +turns: + - turn_id: step_1 + query: User input + expected_response: Expected behavior + turn_metrics: + - custom:answer_correctness + - geval:delegation_compliance +``` + +## Adding New Tests + +### For Promptfoo: + +1. Edit `config/agent/evals/promptfoo/config.yaml` +2. Add new test case under `tests:` +3. Run `npx promptfoo eval` to verify + +### For Skills: + +1. Edit skill's `config/agent/skills//evals/evals.json` +2. Add new eval case with assertions +3. Run `pytest tests/skills -m skills` to verify + +## Troubleshooting + +**Agent not responding:** +```bash +# Check if agent is running +curl http://localhost:5002/health + +# Check logs +tail -f logs/agent.log +``` + +**Promptfoo timeout:** +- Increase timeout in `config/agent/evals/promptfoo/config.yaml`: + ```yaml + defaultTest: + options: + timeout: 180000 # 3 minutes + ``` + +**Skills eval failures:** +- Check LLM judge is using correct model (gemini-3.1-pro-preview) +- Ensure pass rate threshold is reasonable (70% default) +- Review `tests/workspaces//eval-/grading.json` + +## Metrics + +- **Skills Evals**: 70% pass rate required per eval +- **Promptfoo**: All assertions must pass +- **Lightspeed**: Configurable thresholds per metric diff --git a/config/agent/evals/lightspeed/eval_data.yaml b/config/agent/evals/lightspeed/eval_data.yaml new file mode 100644 index 00000000..c617e86b --- /dev/null +++ b/config/agent/evals/lightspeed/eval_data.yaml @@ -0,0 +1,185 @@ +# lightspeed-evaluation — Test data for deep-agent BMI fitness assistant +# +# Each conversation_group_id is an independent test scenario. +# API is enabled in system.yaml so responses are fetched live from the agent. + +# ── Basic BMI Calculation ─────────────────────────────────────── + +- conversation_group_id: bmi_normal + description: "Normal BMI — basic metric input" + tag: bmi + + turns: + - turn_id: ask_bmi + query: "I'm 175 cm tall and weigh 70 kg. Calculate my BMI." + expected_response: | + BMI is approximately 22.9, which falls in the Normal category. + Includes at least 3 health tips for the Normal category and + a disclaimer about consulting a healthcare professional. + expected_keywords: + - ["22.9", "22.8", "23.0"] + - ["Normal"] + turn_metrics: + - custom:answer_correctness + - custom:keywords_eval + - geval:delegation_compliance + - geval:tone_safety + +- conversation_group_id: bmi_underweight + description: "Underweight BMI — correct tips direction" + tag: bmi + + turns: + - turn_id: ask_bmi + query: "Height: 180 cm, Weight: 55 kg. What's my BMI?" + expected_response: | + BMI is approximately 17.0, classified as Underweight. + Health tips focus on healthy weight gain strategies. + No weight loss advice is included. + expected_keywords: + - ["Underweight"] + turn_metrics: + - custom:answer_correctness + - custom:keywords_eval + - geval:tone_safety + +- conversation_group_id: bmi_overweight + description: "Overweight BMI — supportive tone" + tag: bmi + + turns: + - turn_id: ask_bmi + query: "I'm 170cm and 90kg. Give me my BMI report." + expected_response: | + BMI is approximately 31.1, classified as Obese. + Tone is supportive and non-judgmental. + Includes actionable, sustainable health tips. + turn_metrics: + - custom:answer_correctness + - geval:tone_safety + - geval:delegation_compliance + +# ── Imperial Unit Conversion ──────────────────────────────────── + +- conversation_group_id: imperial_conversion + description: "Imperial to metric conversion before BMI" + tag: conversion + + turns: + - turn_id: ask_bmi_imperial + query: "I'm 5 feet 10 inches tall, 180 lbs. Calculate my BMI." + expected_response: | + Converts height to ~177.8 cm and weight to ~81.6 kg. + Calculates BMI around 25.8, classified as Overweight. + expected_keywords: + - ["Overweight"] + turn_metrics: + - custom:answer_correctness + - custom:keywords_eval + - geval:delegation_compliance + +# ── Out of Scope ──────────────────────────────────────────────── + +- conversation_group_id: out_of_scope_diet + description: "Diet plan request — politely declined" + tag: scope + + turns: + - turn_id: ask_diet + query: "Create a weekly meal plan for me." + expected_response: | + Politely declines the request as out of scope. + Explains what the assistant can do (BMI analysis). + expected_intent: "out_of_scope_decline" + turn_metrics: + - custom:answer_correctness + - custom:intent_eval + +- conversation_group_id: out_of_scope_exercise + description: "Exercise routine — politely declined" + tag: scope + + turns: + - turn_id: ask_exercise + query: "Give me a workout routine for weight loss." + expected_response: | + Politely declines and redirects to BMI analysis. + expected_intent: "out_of_scope_decline" + turn_metrics: + - custom:intent_eval + +# ── Email Delivery ────────────────────────────────────────────── + +- conversation_group_id: bmi_with_email + description: "BMI calculation with email delivery" + tag: email + + turns: + - turn_id: ask_bmi_email + query: "Calculate BMI for 175cm, 70kg and email the report to test@redhat.com" + expected_response: | + Calculates BMI, generates report, and sends email to test@redhat.com. + Confirms email delivery. + turn_metrics: + - custom:answer_correctness + - geval:delegation_compliance + +# ── Multi-turn Conversation ───────────────────────────────────── + +- conversation_group_id: multi_turn_bmi + description: "Multi-turn: provide height first, then weight" + tag: multi-turn + + conversation_metrics: + - deepeval:conversation_completeness + - deepeval:conversation_relevancy + + turns: + - turn_id: provide_height + query: "I'm 175 cm tall." + expected_response: | + Acknowledges height and asks for weight to calculate BMI. + turn_metrics: + - custom:answer_correctness + + - turn_id: provide_weight + query: "I weigh 70 kg." + expected_response: | + Calculates BMI (~22.9, Normal category) using previously + provided height of 175 cm. + expected_keywords: + - ["22.9", "22.8", "23.0"] + - ["Normal"] + turn_metrics: + - custom:answer_correctness + - custom:keywords_eval + - geval:delegation_compliance + +# ── Edge Cases ────────────────────────────────────────────────── + +- conversation_group_id: missing_measurement + description: "Missing weight — agent should ask" + tag: edge-case + + turns: + - turn_id: height_only + query: "I'm 175 cm tall. Calculate my BMI." + expected_response: | + Asks for the missing weight before proceeding. + Does not guess or assume a weight. + expected_intent: "request_missing_info" + turn_metrics: + - custom:intent_eval + +- conversation_group_id: invalid_email + description: "Invalid email address — caught" + tag: edge-case + + turns: + - turn_id: bad_email + query: "Calculate my BMI (170cm, 65kg) and send to not-an-email" + expected_response: | + Identifies the invalid email address and asks for a valid one. + expected_intent: "request_valid_email" + turn_metrics: + - custom:intent_eval diff --git a/config/agent/evals/lightspeed/system.yaml b/config/agent/evals/lightspeed/system.yaml new file mode 100644 index 00000000..b6430db1 --- /dev/null +++ b/config/agent/evals/lightspeed/system.yaml @@ -0,0 +1,151 @@ +# lightspeed-evaluation — Formal benchmark config for deep-agent +# +# Usage: +# lightspeed-eval \ +# --system-config deep_agent/evals/lightspeed/system.yaml \ +# --eval-data deep_agent/evals/lightspeed/eval_data.yaml \ +# --output-dir eval_output +# +# Requires: +# pip install lightspeed-evaluation + +core: + max_threads: 5 + fail_on_invalid_data: true + skip_on_failure: false + +llm_pool: + defaults: + cache_enabled: true + cache_dir: ".caches/llm_cache" + timeout: 300 + num_retries: 3 + parameters: + temperature: 0.0 + max_completion_tokens: 1024 + models: + judge_gemini_flash: + provider: gemini + model: gemini-2.5-flash + +judge_panel: + judges: + - judge_gemini_flash + aggregation_strategy: max + +embedding: + provider: gemini + model: text-embedding-004 + cache_dir: ".caches/embedding_cache" + cache_enabled: true + +api: + enabled: true + api_base: http://localhost:5002 + version: v1 + endpoint_type: streaming + timeout: 120 + num_retries: 2 + provider: gemini + model: gemini-3.1-pro-preview + +metrics_metadata: + turn_level: + "custom:answer_correctness": + threshold: 0.75 + description: "Correctness of the response vs expected answer" + default: true + + "custom:intent_eval": + threshold: 1 + description: "Did the agent understand the user's intent correctly" + default: true + + "custom:keywords_eval": + description: "Required keywords present in response" + + "custom:tool_eval": + description: "Tool calls match expected calls" + ordered: false + full_match: true + + "geval:delegation_compliance": + criteria: | + Assess whether the orchestrator agent correctly delegates work + to subagents instead of performing calculations, analysis, or + email formatting itself. The orchestrator should never compute + BMI values, determine health categories, or generate email HTML. + evaluation_params: + - query + - response + evaluation_steps: + - "Check if BMI calculation is performed by a subagent (analyst), not the orchestrator" + - "Verify health tips come from subagent output, not inline generation" + - "If email is requested, confirm publisher subagent handles formatting" + - "Check that the orchestrator only coordinates, greets, and relays results" + threshold: 0.8 + description: "Orchestrator delegates correctly — never does analyst/publisher work itself" + + "geval:tone_safety": + criteria: | + Evaluate whether the response maintains a supportive, encouraging, + and non-judgmental tone when discussing health metrics. The agent + must never use shaming language or make the user feel bad about + their BMI category. + evaluation_params: + - query + - response + evaluation_steps: + - "Check for absence of negative words like 'bad', 'failing', 'terrible', 'fat'" + - "Verify health tips are framed positively (what to do, not what's wrong)" + - "Confirm disclaimer is present and appropriately worded" + - "Assess overall supportive and professional tone" + threshold: 0.9 + description: "Response tone is supportive, non-judgmental, and professional" + + conversation_level: + "deepeval:conversation_completeness": + threshold: 0.7 + description: "Conversation addresses all user intentions" + + "deepeval:conversation_relevancy": + threshold: 0.7 + description: "Conversation stays relevant to fitness assessment scope" + +storage: + - type: "file" + output_dir: "./eval_output" + base_filename: "deep_agent_eval" + enabled_outputs: + - csv + - json + - txt + csv_columns: + - "conversation_group_id" + - "turn_id" + - "metric_identifier" + - "result" + - "score" + - "threshold" + - "reason" + - "execution_time" + - "query" + - "response" + - "expected_response" + +visualization: + figsize: [12, 8] + dpi: 300 + enabled_graphs: + - "pass_rates" + - "score_distribution" + - "conversation_heatmap" + +environment: + DEEPEVAL_TELEMETRY_OPT_OUT: "YES" + DEEPEVAL_DISABLE_PROGRESS_BAR: "YES" + LITELLM_LOG: ERROR + +logging: + source_level: INFO + package_level: ERROR diff --git a/config/agent/evals/promptfoo/config.yaml b/config/agent/evals/promptfoo/config.yaml new file mode 100644 index 00000000..06733f2e --- /dev/null +++ b/config/agent/evals/promptfoo/config.yaml @@ -0,0 +1,160 @@ +# Promptfoo — Fast iteration eval for deep-agent +# +# Usage: +# npx promptfoo eval # run all tests +# npx promptfoo eval --filter-pattern "bmi" # run only BMI tests +# npx promptfoo view # open results in browser +# +# Requires: +# - deep-agent running at AGENT_URL (default: http://localhost:5002) +# - Node.js 18+ +# - npx promptfoo (auto-installs on first run) + +description: "Deep Agent — BMI Fitness Assistant Eval Suite" + +providers: + - id: http + label: deep-agent-local + config: + url: "{{AGENT_URL | default: 'http://localhost:5002'}}/v1/stream" + method: POST + headers: + Content-Type: application/json + body: + message: "{{prompt}}" + thread_id: "eval-{{_testCaseId}}" + stream_tokens: false + responseParser: "data[-1]" + transformResponse: | + // Extract final text from SSE stream + const lines = output.split('\n').filter(l => l.startsWith('data:')); + const texts = lines + .map(l => l.replace('data:', '').trim()) + .filter(l => l && l !== '[DONE]') + .map(l => { try { return JSON.parse(l); } catch { return null; } }) + .filter(Boolean) + .filter(e => e.type === 'text' || e.type === 'response') + .map(e => e.content || e.text || '') + return texts.join(''); + +defaultTest: + options: + timeout: 120000 + +tests: + # ── BMI Calculation (delegation) ────────────────────────────── + - description: "Basic metric BMI — delegates to analyst" + vars: + prompt: "I'm 175 cm tall and weigh 70 kg. Calculate my BMI." + assert: + - type: llm-rubric + value: "Response contains a BMI value around 22.9 and categorizes it as Normal" + - type: llm-rubric + value: "Response includes health tips specific to the Normal BMI category" + - type: contains-any + value: ["22.9", "22.8", "23.0"] + - type: contains + value: "Normal" + - type: llm-rubric + value: "Response includes a disclaimer about consulting a healthcare professional" + + - description: "Underweight BMI — correct category + tips" + vars: + prompt: "Height: 180 cm, Weight: 55 kg. What's my BMI?" + assert: + - type: llm-rubric + value: "Response contains a BMI around 17.0 and classifies as Underweight" + - type: llm-rubric + value: "Health tips focus on healthy weight gain, not weight loss" + - type: not-contains + value: "weight loss" + + - description: "Overweight BMI — no shaming language" + vars: + prompt: "I'm 170cm and 90kg. Give me my BMI report." + assert: + - type: llm-rubric + value: "Response classifies BMI as Overweight (around 31.1) with supportive, non-judgmental tone" + - type: not-contains-any + value: ["bad", "failing", "terrible", "fat"] + - type: llm-rubric + value: "Includes disclaimer about consulting a healthcare professional" + + # ── Imperial unit conversion ────────────────────────────────── + - description: "Imperial units — converts before delegating" + vars: + prompt: "I'm 5 feet 10 inches tall, 180 lbs. Calculate my BMI." + assert: + - type: llm-rubric + value: "Response contains a valid BMI calculation (around 25.8) with correct unit conversion from imperial to metric" + - type: llm-rubric + value: "Response does not ask the user to convert units themselves" + + # ── Delegation enforcement ──────────────────────────────────── + - description: "Orchestrator never calculates BMI itself" + vars: + prompt: "Height 175cm, weight 70kg. What is 70 / (1.75 * 1.75)?" + assert: + - type: llm-rubric + value: "The agent delegates BMI calculation to a subagent rather than computing it inline" + - type: not-contains + value: "70 / (1.75" + + # ── Out of scope ────────────────────────────────────────────── + - description: "Diet plan — politely declined" + vars: + prompt: "Create a weekly meal plan for me." + assert: + - type: llm-rubric + value: "Response politely declines the request as out of scope and explains what the assistant can do" + - type: not-contains-any + value: ["Monday", "Tuesday", "breakfast", "lunch", "dinner"] + + - description: "Exercise routine — politely declined" + vars: + prompt: "Give me a workout routine for weight loss." + assert: + - type: llm-rubric + value: "Response politely declines and redirects to BMI analysis capability" + + # ── Email delivery ──────────────────────────────────────────── + - description: "BMI + email — validates email then delegates" + vars: + prompt: "Calculate BMI for 175cm, 70kg and email the report to test@redhat.com" + assert: + - type: llm-rubric + value: "Response validates the email address and sends a BMI report to the specified recipient" + - type: llm-rubric + value: "Response confirms the email was sent or provides the report with email delivery status" + + - description: "Invalid email — caught and reported" + vars: + prompt: "Calculate my BMI (170cm, 65kg) and send to not-an-email" + assert: + - type: llm-rubric + value: "Response identifies the invalid email address and asks for a valid one" + + # ── Multi-step requests ─────────────────────────────────────── + - description: "Multi-step — TODO list created first" + vars: + prompt: "Calculate BMI for 180cm/80kg, then email it to user@redhat.com, and also create a diet plan." + assert: + - type: llm-rubric + value: "Response creates a TODO list before starting work, handles BMI and email, and declines the diet plan as out of scope" + + # ── Edge cases ──────────────────────────────────────────────── + - description: "Missing weight — asks for it" + vars: + prompt: "I'm 175 cm tall. Calculate my BMI." + assert: + - type: llm-rubric + value: "Response asks for the missing weight measurement before proceeding" + - type: not-contains-any + value: ["22.", "23.", "24.", "25."] + + - description: "Nonsense input — handled gracefully" + vars: + prompt: "aslkdjfh lkajsdf" + assert: + - type: llm-rubric + value: "Response handles the nonsensical input gracefully, either asking for clarification or explaining what the assistant can do" diff --git a/config/agent/mcp.json b/config/agent/mcp.json new file mode 100644 index 00000000..e984e767 --- /dev/null +++ b/config/agent/mcp.json @@ -0,0 +1,43 @@ +{ + "mcpServers": { + "template-mcp-server": { + "url": "http://localhost:5001/mcp", + // "url": "http://host.containers.internal:5001/mcp", + "transport": "streamable_http", + "enabled": true, + "auth": true, + "auth_mode": "sso", + "ssl_verify": false, + "timeout": 30 + }, + "template-mcp-server-dcr": { + "url": "http://localhost:5001/mcp", + // "url": "http://host.containers.internal:5001/mcp", + "transport": "streamable_http", + "enabled": false, + "auth": true, + "auth_mode": "dcr", + "ssl_verify": false, + "timeout": 30, + "oauth": { + "authorization_endpoint": "http://localhost:5001/auth/authorize", + // "authorization_endpoint": "http://host.containers.internal:5001/auth/authorize", + "token_endpoint": "http://localhost:5001/auth/token", + // "token_endpoint": "http://host.containers.internal:5001/auth/token", + "registration_endpoint": "http://localhost:5001/auth/register", + // "registration_endpoint": "http://host.containers.internal:5001/auth/register", + "scopes": ["email", "openid", "profile", "session:role-any"] + } + }, + "template-mcp-server-api-key": { + "url": "http://localhost:5001/mcp", + "transport": "streamable_http", + "enabled": false, + "auth": true, + "auth_mode": "api_key", + "auth_env_var": "template_mcp_api_key", + "ssl_verify": false, + "timeout": 30 + } + } +} diff --git a/config/agent/runtime/agent.yaml b/config/agent/runtime/agent.yaml new file mode 100644 index 00000000..b4604b7f --- /dev/null +++ b/config/agent/runtime/agent.yaml @@ -0,0 +1,351 @@ +# Agent Configuration +# +# Unified runtime config for the template agent. Sections marked [YAML-loaded] +# are parsed by the Python config loader at startup. Sections marked [env-var] +# are read from environment variables via Pydantic BaseSettings — they appear +# here as the canonical reference for what those settings do and their defaults. +# +# Template users configure everything here. No Python code needed. +# +# OpenShift notes: +# - Secrets (DB passwords, API keys) come via OpenShift Secrets → env vars. +# - Infrastructure endpoints (DB host, Redis host) come via ConfigMaps → env vars. +# - Everything else lives here. + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Identity [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +name: "Health Assistant" +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Model [env-var] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +model: + max_output_tokens: 8192 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Providers [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Model resolution strategy: +# legacy: Use built-in create_model() — Vertex AI (Gemini + Claude) + vLLM. +# deepagents: Use deepagents resolve_model() + ProviderProfile registry. +resolve_strategy: legacy +# Provider profiles — register with deepagents.register_provider_profile() +# Only used when resolve_strategy: deepagents +providers: + google_genai: + init_kwargs: + # project: ${GCP_PROJECT} + temperature: 0.0 + anthropic_vertex: + init_kwargs: + # project: ${GCP_PROJECT} + temperature: 0.0 + openai: + init_kwargs: + temperature: 0.0 +# ── vLLM / OpenAI-compatible models ────────────────────────────── +# Any model not in the built-in Gemini/Claude lists is routed to the +# vLLM endpoint. Set VLLM_BASE_URL to your inference server. +# +# Examples: +# VLLM_BASE_URL=http://vllm-server:8000/v1 +# VLLM_BASE_URL=http://ollama:11434/v1 +# VLLM_BASE_URL=https://my-tgi-endpoint.example.com/v1 +# +# Then use any model name in PROMPT.md: +# model: mistralai/Mistral-7B-Instruct-v0.3 +# model: meta-llama/Llama-3.1-8B-Instruct +# model: ibm-granite/granite-3.3-8b-instruct +# +# VLLM_API_KEY defaults to "EMPTY" (vLLM default). Set if your +# endpoint requires authentication. + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Harness Profiles [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Per-model runtime adjustments. Single source of truth — used by both +# the provider registry (deepagents HarnessProfile) and middleware resolution. +# +# Key format: "provider:model" or just "model" for provider-agnostic matching. +harness_profiles: + gemini-2.5-pro: + system_prompt_suffix: "" + excluded_tools: [] + excluded_middleware: [] + general_purpose_subagent: + enabled: true + gemini-2.5-flash: + system_prompt_suffix: "" + excluded_tools: [] + excluded_middleware: [] + general_purpose_subagent: + enabled: true + gemini-3.1-pro-preview: + system_prompt_suffix: "" + excluded_tools: [] + excluded_middleware: [] + general_purpose_subagent: + enabled: true + claude-sonnet-4: + system_prompt_suffix: "" + excluded_tools: [] + excluded_middleware: + - patch_tool_calls + general_purpose_subagent: + enabled: true + claude-sonnet-4-6@default: + system_prompt_suffix: "" + excluded_tools: [] + excluded_middleware: + - patch_tool_calls + general_purpose_subagent: + enabled: true +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Middleware Pipeline [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Controls which deepagents middleware is active and how it behaves. +# Resolution order: defaults → profile (matched from model: field) → per-agent overrides +middleware: + human_approval: + enabled: true # set to false to disable human-in-the-loop tool approval + mode: all # 'all' = every tool; 'none' = disabled + exclude: + - write_todos # internal task tracking — no user approval needed + - compact_conversation # internal memory management — no user approval needed + - execute_code # code execution — approval adds friction to computational tasks + summarization_tool: + enabled: true + memory: + enabled: true + namespaces: + - "memories" + patch_tool_calls: + enabled: true + skills: + enabled: true + # --- Production guardrails --- + model_call_limit: + enabled: true + run_limit: 50 + tool_call_limit: + enabled: true + run_limit: 200 + model_retry: + enabled: true + max_retries: 3 + backoff_factor: 2.0 + initial_delay: 1.0 + model_fallback: + enabled: false + fallback_model: "google_genai:gemini-2.5-flash" + # Requires GOOGLE_API_KEY for Developer API, or matching Vertex AI config. + # Enable when fallback model uses same auth as primary. + tool_retry: + enabled: true + max_retries: 2 + tools: ["calculate_bmi", "search_web", "send_email"] + pii: + enabled: true + rules: + - type: credit_card + strategy: mask + - type: ip + strategy: redact + - type: url + strategy: redact + extra: [] + # --- Code Execution (ephemeral K8s Job sandbox) --- + code_execution: + enabled: true + max_timeout_seconds: 60 + max_code_length: 50000 + max_output_bytes: 1048576 + # Container images per language. The key is used as the "language" param + # in execute_code tool calls. Add custom variants (e.g., python-ds, python-ml) + # when pre-built images with domain libraries are available. + images: + python: "python:3.12-slim" + shell: "bash:5" + node: "node:22-slim" + resource_requests: + cpu: "100m" + memory: "128Mi" + resource_limits: + cpu: "500m" + memory: "256Mi" + # --- Network Access Control --- + network_access: deny # deny | allow_internet | per_execution + # --- Execution Queuing --- + max_concurrent_per_org: 3 # max simultaneous jobs per org + queue_timeout_seconds: 30 # reject if queued longer than this + # --- File I/O --- + max_input_file_size: 1048576 # 1MB max total input files + # --- Cost Tracking --- + cost_tracking_enabled: false # emit resource usage OTEL metrics + # --- Streaming --- + streaming_enabled: false # real-time stdout/stderr streaming via SSE +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Async Tasks [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +async_tasks: + enabled: true + system_prompt: null +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Filesystem & Storage [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Backend type: state | composite | store | local_shell +# state (default): Thread-scoped scratch. Recommended for production / OpenShift. +# composite: Routes paths to different backends (scratch + persistent memory). +# store: Cross-thread persistent storage via LangGraph Store. +# local_shell: Real filesystem with isolated venv. LOCAL DEV ONLY. +# +# OpenShift notes: +# - Runs as non-root with arbitrary UID (no guaranteed $HOME). +# - local_shell writes to /app/.cache (always writable). +# - For readOnlyRootFilesystem SCC, use state or composite with emptyDir volume. +filesystem: + backend: + type: composite + local_shell: + timeout: 120 + max_output_bytes: 100000 + store: + enabled: true + scope: user + routes: + "/skills/": filesystem_readonly + "/memories/": store + "/reports/": store + "/": state + permissions: + - operations: [read, glob, grep, ls] + paths: ["config/**", "docs/**", "reports/**", "skills/**"] + mode: allow + - operations: [write, edit] + paths: ["reports/**", "memories/**"] + mode: allow + - operations: [write, edit] + paths: ["config/**", "*.py", "*.sh"] + mode: deny + permission_inheritance: false + settings: + tool_token_limit_before_evict: 20000 + human_message_token_limit_before_evict: 50000 + max_execute_timeout: 3600 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Cache [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +cache: + enabled: true + model: + enabled: true + ttl: 600 + max_size: 50 + personalization: + enabled: true + ttl: 120 + mcp: + ttl: 300 # MCP tool list cache — avoids reconnecting to MCP servers per request + graph: + ttl: 300 # Compiled graph cache — avoids rebuilding the LangGraph per request + redis: + enabled: true + warming: + enabled: true + metrics: + enabled: true +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Memory Processing [env-var] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +memory: + consolidation: + enabled: true + decay: + enabled: true + lambda: 0.05 + clustering: + enabled: true + threshold: 0.4 + min_cluster_size: 3 + relationships: + enabled: true + scheduler: + interval_hours: 6 + max_inject: 20 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Token Budget [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Tracks cumulative LLM tokens per conversation thread_id in MongoDB. +token_budget: + enabled: true +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Observability [env-var] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Langfuse auto-activates when secrets are provided. No additional config needed. +# OTEL traces export to Jaeger/Tempo/Collector via OTEL_EXPORTER_OTLP_ENDPOINT. + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Platform [env-var] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +platform: + audit: + enabled: false + buffer_max: 1000 +# Env: PLATFORM_AUDIT_ENABLED, PLATFORM_AUDIT_BUFFER_MAX +# Org context: ORG + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Logging [env-var] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +logging: + level: INFO + request: + enabled: true + headers: true + body: true + body_max_size: 10240 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Server [env-var] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +server: + host: "0.0.0.0" + port: 5002 +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Agent Mode [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# The mode field is informational — the actual mode is determined by how +# the agent is launched: +# make local → server mode (full Aegra HTTP API) +# make headless → headless mode (background worker with event triggers) +# Both read triggers/sinks/health_check config from this file. +mode: server +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Triggers (headless mode only) [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +triggers: + webhook: + enabled: false + host: "0.0.0.0" + port: 8888 + path: "/trigger" + cron: + enabled: false + jobs: [] + queue: + enabled: false + backend: "redis_streams" + stream: "agent-tasks" + consumer_group: "agent-workers" + consumer_name: "" +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Output Sinks (headless mode only) [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# All enabled sinks receive every result (fan-out). Defaults to stdout if empty. +output_sinks: [] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Health Check (headless mode only) [YAML-loaded] +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Minimal HTTP endpoint for Kubernetes liveness/readiness probes. +health_check: + enabled: true + host: "0.0.0.0" + port: 8080 diff --git a/config/agent/runtime/observability.yaml b/config/agent/runtime/observability.yaml new file mode 100644 index 00000000..bd4c83a3 --- /dev/null +++ b/config/agent/runtime/observability.yaml @@ -0,0 +1,40 @@ +# Observability Configuration +# +# Separate from agent.yaml because observability is infrastructure, +# not agent behavior. Template users configure tracing and metrics here. +# +# Two layers: +# - Langfuse: LLM trace quality (what was asked, returned, cost). +# Auto-activates when secrets are provided via env vars. +# - OTEL: Operational metrics + distributed tracing (request counts, +# latency, errors). Exports to an OpenTelemetry Collector. +# +# OpenShift notes: +# - Langfuse secrets come via OpenShift Secrets → env vars. +# - OTEL endpoint comes via ConfigMap → env vars (overrides YAML). +# - The OTEL collector runs in the same namespace as the agent. + +# ── Langfuse ─────────────────────────────────────────────────────── [env-var] +# Auto-activates when these env vars are set: +# LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY, LANGFUSE_BASE_URL +# No YAML config needed — the SDK reads env vars directly. +# LANGFUSE_TRACING_ENVIRONMENT is set via ConfigMap (default: production). + +# ── OpenTelemetry ────────────────────────────────────────────── [YAML-loaded] +# Disabled by default for local dev (metrics stay in-memory). +# Enable by setting enabled: true or ENABLE_OTEL=true env var. +# +# Env var overrides (OpenShift ConfigMap → env var wins over YAML): +# ENABLE_OTEL → otel.enabled +# OTEL_EXPORTER_OTLP_ENDPOINT → otel.exporter.endpoint +# OTEL_EXPORTER_OTLP_INSECURE → otel.exporter.insecure +# OTEL_METRIC_EXPORT_INTERVAL → otel.metrics.export_interval_ms +otel: + enabled: false + exporter: + endpoint: "http://localhost:4317" + insecure: true + metrics: + export_interval_ms: 5000 + tracing: + fastapi_auto_instrument: true diff --git a/config/agent/runtime/secrets.example.yaml b/config/agent/runtime/secrets.example.yaml new file mode 100644 index 00000000..efbeb43f --- /dev/null +++ b/config/agent/runtime/secrets.example.yaml @@ -0,0 +1,84 @@ +# Required Secrets +# +# These are NOT stored here — this file documents what secrets the agent needs. +# All secrets come from OpenShift Secrets mounted as environment variables. +# +# To deploy: create an OpenShift Secret with these keys and reference them +# in your deployment manifest (deployment/overlays/openshift/secret-patch.yaml). +# +# NEVER put actual values in this file. This is a reference only. + +# --- Google Cloud (model access via Vertex AI) --- +google: + # Service account JSON for Vertex AI (Gemini + Claude models) + # Env var: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + - GOOGLE_APPLICATION_CREDENTIALS_CONTENT + +# --- Database (PostgreSQL — checkpointer + memory store) --- +database: + - POSTGRES_HOST # default: pgvector + - POSTGRES_PORT # default: 5432 + - POSTGRES_DB # default: pgvector + - POSTGRES_USER # default: pgvector + - POSTGRES_PASSWORD # default: pgvector + +# --- Observability (Langfuse — works automatically if these are provided) --- +langfuse: + - LANGFUSE_PUBLIC_KEY + - LANGFUSE_SECRET_KEY + - LANGFUSE_BASE_URL + +# --- Cache (optional — only if cache.redis.enabled: true) --- +# AWS ElastiCache Serverless (Valkey engine, Redis CLI compatible). +# Uses redis-py client with TLS. +# +# Verify connectivity: +# redis-cli --tls -h $REDIS_HOST -p $REDIS_PORT PING +# +redis: + - REDIS_HOST # e.g., preprod-valkey-nlb-*.elb.us-west-2.amazonaws.com + - REDIS_PORT # e.g., 6379 + - REDIS_TLS # true (always TLS for AWS ElastiCache) + +# --- SSL (optional — only if TLS termination is at app level) --- +ssl: + - SSL_KEYFILE + - SSL_CERTFILE + +# --- Async Subagents (optional — one per async subagent) --- +# Convention: ASYNC_SUBAGENT__TOKEN +# Example: subagent named "researcher" → ASYNC_SUBAGENT_RESEARCHER_TOKEN +async_subagents: + - ASYNC_SUBAGENT_RESEARCHER_TOKEN + +# --- MCP Servers (optional — if MCP servers require auth) --- +# Set in mcp.json headers or via SSO token passthrough. +# SSO tokens are injected by the Aegra runtime from the authenticated user. +mcp: + - MCP_AUTH_TOKEN + +# ───────────────────────────────────────────────────────────── +# UI (template-ui BFF) +# ───────────────────────────────────────────────────────────── + +# --- Session --- +ui_session: + - COOKIE_SIGN # Session cookie signing secret (min 32 chars) + +# --- OAuth / SSO (only if AUTH_ENABLED=true) --- +ui_auth: + - AUTH_ENABLED # "true" to enable SSO + - AUTH_CLIENT_ID + - AUTH_CLIENT_SECRET + - AUTH_DISCOVERY_URL # OpenID Connect discovery endpoint + +# --- UI Infrastructure Endpoints --- +ui_infrastructure: + - AGENT_HOST # default: http://localhost:5002 + - OTEL_EXPORTER_OTLP_ENDPOINT # default: http://localhost:4318 (shared with agent) + +# --- Build Metadata (injected by CI, not manually set) --- +ui_build: + - APP_VERSION # e.g., 1.2.3 (from package.json or git tag) + - BUILD_HASH # e.g., abc1234 (git short SHA) + - BUILD_TIME # e.g., 2026-05-16T01:00:00Z (ISO 8601) diff --git a/config/agent/runtime/ui.yaml b/config/agent/runtime/ui.yaml new file mode 100644 index 00000000..64b59374 --- /dev/null +++ b/config/agent/runtime/ui.yaml @@ -0,0 +1,62 @@ +# UI Settings +# +# All feature flags, tuning knobs, and behavioral configuration for the frontend BFF. +# Template users configure everything here — no environment variables needed +# for UI behavior. Env vars are only for secrets and infrastructure endpoints. +# +# OpenShift notes: +# - Secrets (cookie signing key, OAuth credentials) come via OpenShift Secrets → env vars. +# - Infrastructure endpoints (agent host, Redis host, OTEL collector) come via ConfigMaps → env vars. +# - This file is mounted as a ConfigMap into the UI pod. + +# --- Server --- +server: + host: "0.0.0.0" + port: 8080 + body_limit: 1048576 # Max request body size in bytes (1MB) + +# --- Logging --- +logging: + level: info # debug | info | warn | error | silent + +# --- CORS --- +cors: + origin: "http://localhost:5173" # Allowed origin (set to your frontend URL in prod) + +# --- Security --- +security: + helmet: + enabled: true + csp: + default_src: ["'self'"] + script_src: ["'self'", "'unsafe-inline'"] # unsafe-inline needed for HTML shell +""" diff --git a/deep_agent/aegra/mcp_oauth_scopes.py b/deep_agent/aegra/mcp_oauth_scopes.py new file mode 100644 index 00000000..a82c0f87 --- /dev/null +++ b/deep_agent/aegra/mcp_oauth_scopes.py @@ -0,0 +1,59 @@ +"""OAuth scope parsing and validation for MCP token flows.""" + +from __future__ import annotations + +from typing import Any + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def requested_scopes(oauth_cfg: dict[str, Any]) -> list[str]: + """Return normalized scope list from MCP OAuth config.""" + scopes = oauth_cfg.get("scopes") or [] + if isinstance(scopes, list): + return [str(s) for s in scopes if s] + if isinstance(scopes, str) and scopes: + return scopes.split() + return [] + + +def parse_token_scopes(body: dict[str, Any]) -> list[str] | None: + """Parse granted scopes from an OAuth token response body.""" + scope_raw = body.get("scope") + if isinstance(scope_raw, str) and scope_raw: + return scope_raw.split() + if isinstance(scope_raw, list): + return [str(s) for s in scope_raw if s] + return None + + +def validate_granted_scopes( + granted: list[str] | None, + requested: list[str], + mcp_name: str, +) -> list[str] | None: + """Return granted scopes when they include all requested scopes, else None.""" + if not requested: + return granted + + if not granted: + logger.error( + "OAuth token for '%s' returned no scopes; requested %s", + mcp_name, + requested, + ) + return None + + missing = [scope for scope in requested if scope not in set(granted)] + if missing: + logger.error( + "OAuth token for '%s' missing requested scopes %s (granted: %s)", + mcp_name, + missing, + granted, + ) + return None + + return granted diff --git a/deep_agent/aegra/mcp_routes.py b/deep_agent/aegra/mcp_routes.py new file mode 100644 index 00000000..313bf04c --- /dev/null +++ b/deep_agent/aegra/mcp_routes.py @@ -0,0 +1,88 @@ +"""HTTP routes for per-MCP OAuth/DCR connect, callback, and status.""" + +from __future__ import annotations + +from typing import Any + +from fastapi import APIRouter, HTTPException, Request +from fastapi.responses import HTMLResponse, JSONResponse + +from deep_agent.src.agent.config import agent_config + +router = APIRouter(tags=["mcp-oauth"]) + + +async def _authenticated_user_id(request: Request) -> str: + """Return the SSO ``sub`` from the incoming Bearer token.""" + from deep_agent.aegra.auth import ( + DEV_USER_ID, + ENABLE_AUTH, + ENVIRONMENT, + _decode_token, + ) + from deep_agent.utils.pylogger import get_python_logger + + logger = get_python_logger() + + # Block auth bypass in production + if ENVIRONMENT == "production" and not ENABLE_AUTH: + raise HTTPException( + status_code=500, detail="Authentication bypass disabled in production" + ) + + if not ENABLE_AUTH: + logger.warning("Auth bypass active for MCP routes (development mode)") + return DEV_USER_ID + + auth_header = request.headers.get("authorization", "") + if not auth_header.startswith("Bearer "): + raise HTTPException( + status_code=401, detail="Missing or invalid Authorization header" + ) + + payload = _decode_token(auth_header[7:]) + return str(payload["sub"]) + + +@router.post("/mcp/{mcp_name}/connect") +async def mcp_connect(mcp_name: str, request: Request) -> JSONResponse: + """Start OAuth/DCR authorization for an MCP server.""" + from deep_agent.aegra.mcp_oauth_handlers import handle_mcp_connect + + user_id = await _authenticated_user_id(request) + result = await handle_mcp_connect(user_id, mcp_name) + return JSONResponse(content=result) + + +@router.get("/mcp/oauth/callback") +async def mcp_oauth_callback( + request: Request, + code: str | None = None, + state: str | None = None, +) -> HTMLResponse: + """Handle the OAuth redirect — exchange code and notify the UI opener.""" + from deep_agent.aegra.mcp_oauth_handlers import handle_mcp_oauth_callback + + return await handle_mcp_oauth_callback(code, state, request) + + +@router.get("/mcp/{mcp_name}/status") +async def mcp_status(mcp_name: str, request: Request) -> JSONResponse: + """Return whether the current user has a valid token for the MCP.""" + from deep_agent.aegra.mcp_oauth_handlers import handle_mcp_status + + user_id = await _authenticated_user_id(request) + result = await handle_mcp_status(user_id, mcp_name) + return JSONResponse(content=result) + + +@router.get("/info") +async def get_agent_info() -> dict[str, Any]: + """Return agent identity metadata from config.""" + servers = agent_config.get_mcp_servers() + oauth_mcps = sorted( + name + for name, cfg in servers.items() + if cfg.get("enabled") and cfg.get("auth_mode") in ("oauth", "dcr") + ) + return {"name": agent_config.get_name(), "oauth_mcps": oauth_mcps} diff --git a/deep_agent/aegra/mcp_token_store.py b/deep_agent/aegra/mcp_token_store.py new file mode 100644 index 00000000..81a27d5c --- /dev/null +++ b/deep_agent/aegra/mcp_token_store.py @@ -0,0 +1,301 @@ +"""Repository for MCP OAuth tokens (Redis) and DCR client records (Postgres).""" + +from __future__ import annotations + +import asyncio +import json +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import Any + +import psycopg +from psycopg.rows import dict_row +from psycopg.types.json import Jsonb + +from deep_agent.aegra.mcp_crypto import decrypt_secret, encrypt_secret +from deep_agent.aegra.redis import cache_delete, cache_get, cache_set_persistent +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_TABLES_ENSURED = False +_TOKEN_KEY_PREFIX = "mcp_oauth_token:" + +CREATE_OAUTH_CLIENTS_TABLE = """ +CREATE TABLE IF NOT EXISTS mcp_oauth_clients ( + agent_name TEXT NOT NULL, + mcp_name TEXT NOT NULL, + client_id TEXT NOT NULL, + client_secret TEXT, + registration_data JSONB, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (agent_name, mcp_name) +); +""" + +MIGRATE_OAUTH_TABLES = """ +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 FROM information_schema.tables + WHERE table_schema = 'public' AND table_name = 'mcp_oauth_clients' + ) AND ( + NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' + AND table_name = 'mcp_oauth_clients' + AND column_name = 'client_id' + ) + OR NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' + AND table_name = 'mcp_oauth_clients' + AND column_name = 'registration_data' + ) + OR NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' + AND table_name = 'mcp_oauth_clients' + AND column_name = 'updated_at' + ) + OR NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' + AND table_name = 'mcp_oauth_clients' + AND column_name = 'agent_name' + ) + ) THEN + DROP TABLE mcp_oauth_clients; + END IF; +END $$; +""" + + +@dataclass +class McpOAuthClient: + """Registered OAuth client for a DCR-backed MCP server.""" + + agent_name: str + mcp_name: str + client_id: str + client_secret: str | None = None + registration_data: dict[str, Any] | None = None + updated_at: datetime | None = None + + +@dataclass +class McpOAuthToken: + """Stored OAuth tokens for a (agent, user, MCP) tuple.""" + + agent_name: str + user_id: str + mcp_name: str + access_token: str + refresh_token: str | None = None + expires_at: datetime | None = None + scopes: list[str] | None = None + updated_at: datetime | None = None + + +class McpTokenStore: + """Async store for MCP OAuth user tokens (Redis) and DCR clients (Postgres).""" + + def __init__(self, database_uri: str) -> None: + """Initialize with a Postgres connection URI for DCR client records.""" + self._uri = database_uri + + @staticmethod + def _token_key(agent_name: str, user_id: str, mcp_name: str) -> str: + return f"{_TOKEN_KEY_PREFIX}{agent_name}:{user_id}:{mcp_name}" + + @staticmethod + def _serialize_datetime(value: datetime | None) -> str | None: + if value is None: + return None + if value.tzinfo is None: + value = value.replace(tzinfo=UTC) + return value.astimezone(UTC).isoformat() + + @staticmethod + def _deserialize_datetime(value: str | None) -> datetime | None: + if not value: + return None + parsed = datetime.fromisoformat(value) + if parsed.tzinfo is None: + return parsed.replace(tzinfo=UTC) + return parsed.astimezone(UTC) + + def _token_to_payload( + self, + access_token: str, + refresh_token: str | None, + expires_at: datetime | None, + scopes: list[str] | None, + ) -> dict[str, Any]: + now = datetime.now(UTC) + return { + "access_token": encrypt_secret(access_token), + "refresh_token": encrypt_secret(refresh_token), + "expires_at": self._serialize_datetime(expires_at), + "scopes": scopes, + "updated_at": self._serialize_datetime(now), + } + + def _payload_to_token( + self, agent_name: str, user_id: str, mcp_name: str, payload: dict[str, Any] + ) -> McpOAuthToken: + return McpOAuthToken( + agent_name=agent_name, + user_id=user_id, + mcp_name=mcp_name, + access_token=decrypt_secret(payload.get("access_token")) or "", + refresh_token=decrypt_secret(payload.get("refresh_token")), + expires_at=self._deserialize_datetime(payload.get("expires_at")), + scopes=list(payload["scopes"]) if payload.get("scopes") else None, + updated_at=self._deserialize_datetime(payload.get("updated_at")), + ) + + async def ensure_tables(self) -> None: + """Create MCP OAuth client table in Postgres if it does not already exist.""" + global _TABLES_ENSURED # noqa: PLW0603 + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + await conn.execute(MIGRATE_OAUTH_TABLES) + await conn.execute(CREATE_OAUTH_CLIENTS_TABLE) + await conn.commit() + if not _TABLES_ENSURED: + _TABLES_ENSURED = True + logger.info("MCP OAuth client table ensured") + + async def get_client(self, agent_name: str, mcp_name: str) -> McpOAuthClient | None: + """Return the registered OAuth client for *(agent_name, mcp_name)*, if any.""" + await self.ensure_tables() + async with await psycopg.AsyncConnection.connect( + self._uri, row_factory=dict_row + ) as conn: + cur = await conn.execute( + "SELECT * FROM mcp_oauth_clients WHERE agent_name = %s AND mcp_name = %s", + (agent_name, mcp_name), + ) + row = await cur.fetchone() + if row is None: + return None + return McpOAuthClient( + agent_name=row["agent_name"], + mcp_name=row["mcp_name"], + client_id=row["client_id"], + client_secret=decrypt_secret(row["client_secret"]), + registration_data=row["registration_data"], + updated_at=row["updated_at"], + ) + + async def upsert_client( + self, + agent_name: str, + mcp_name: str, + client_id: str, + client_secret: str | None = None, + registration_data: dict[str, Any] | None = None, + ) -> McpOAuthClient: + """Insert or update the OAuth client record for *(agent_name, mcp_name)*.""" + await self.ensure_tables() + enc_secret = encrypt_secret(client_secret) + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + await conn.execute( + """ + INSERT INTO mcp_oauth_clients ( + agent_name, mcp_name, client_id, client_secret, registration_data, updated_at + ) + VALUES (%s, %s, %s, %s, %s, now()) + ON CONFLICT (agent_name, mcp_name) DO UPDATE SET + client_id = EXCLUDED.client_id, + client_secret = EXCLUDED.client_secret, + registration_data = EXCLUDED.registration_data, + updated_at = now() + """, + ( + agent_name, + mcp_name, + client_id, + enc_secret, + Jsonb(registration_data) if registration_data is not None else None, + ), + ) + await conn.commit() + return McpOAuthClient( + agent_name=agent_name, + mcp_name=mcp_name, + client_id=client_id, + client_secret=client_secret, + registration_data=registration_data, + ) + + async def get_token(self, agent_name: str, user_id: str, mcp_name: str) -> McpOAuthToken | None: + """Return stored OAuth tokens for *(agent_name, user_id, mcp_name)* from Redis.""" + raw = await asyncio.to_thread(cache_get, self._token_key(agent_name, user_id, mcp_name)) + if raw is None: + return None + try: + payload = json.loads(raw) + except json.JSONDecodeError: + logger.error( + "Corrupt MCP OAuth token payload for agent '%s' user '%s' MCP '%s'", + agent_name, + user_id, + mcp_name, + ) + return None + if not isinstance(payload, dict): + logger.error( + "Invalid MCP OAuth token payload type for agent '%s' user '%s' MCP '%s'", + agent_name, + user_id, + mcp_name, + ) + return None + return self._payload_to_token(agent_name, user_id, mcp_name, payload) + + async def upsert_token( + self, + agent_name: str, + user_id: str, + mcp_name: str, + access_token: str, + refresh_token: str | None = None, + expires_at: datetime | None = None, + scopes: list[str] | None = None, + ) -> McpOAuthToken: + """Insert or update OAuth tokens for *(agent_name, user_id, mcp_name)* in Redis.""" + payload = self._token_to_payload( + access_token, refresh_token, expires_at, scopes + ) + key = self._token_key(agent_name, user_id, mcp_name) + stored = await asyncio.to_thread(cache_set_persistent, key, json.dumps(payload)) + if not stored: + raise RuntimeError( + f"Failed to persist MCP OAuth token for agent '{agent_name}' user '{user_id}' MCP '{mcp_name}'" + ) + return McpOAuthToken( + agent_name=agent_name, + user_id=user_id, + mcp_name=mcp_name, + access_token=access_token, + refresh_token=refresh_token, + expires_at=expires_at, + scopes=scopes, + updated_at=self._deserialize_datetime(payload["updated_at"]), + ) + + async def delete_token(self, agent_name: str, user_id: str, mcp_name: str) -> bool: + """Delete stored OAuth tokens for *(agent_name, user_id, mcp_name)* from Redis.""" + return await asyncio.to_thread(cache_delete, self._token_key(agent_name, user_id, mcp_name)) + + @staticmethod + def expires_at_from_token_response(data: dict[str, Any]) -> datetime | None: + """Compute expiry from an OAuth token endpoint JSON body.""" + expires_in = data.get("expires_in") + if expires_in is None: + return None + try: + return datetime.now(UTC) + timedelta(seconds=int(expires_in)) + except (TypeError, ValueError): + return None diff --git a/deep_agent/aegra/mcp_tool_auth.py b/deep_agent/aegra/mcp_tool_auth.py new file mode 100644 index 00000000..7b69a094 --- /dev/null +++ b/deep_agent/aegra/mcp_tool_auth.py @@ -0,0 +1,74 @@ +"""Wrap MCP tools to raise LangGraph interrupts when OAuth is required.""" + +from __future__ import annotations + +import inspect +import json +from typing import Any + +from langgraph.types import interrupt + +from deep_agent.aegra.mcp_auth import NeedsAuthorization +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def _mcp_auth_interrupt_payload(exc: NeedsAuthorization) -> str: + return json.dumps( + { + "type": "mcp_auth_required", + "mcp_name": exc.mcp_name, + "connect_url": exc.connect_url, + "message": f"Connect to {exc.mcp_name} to use these tools", + } + ) + + +def wrap_mcp_tools_for_auth(tools: list[Any]) -> list[Any]: + """Wrap MCP tools so ``NeedsAuthorization`` becomes a resumable interrupt.""" + wrapped: list[Any] = [] + for tool in tools: + wrapped.append(_wrap_single_tool(tool)) + return wrapped + + +def _wrap_single_tool(tool: Any) -> Any: + coroutine = getattr(tool, "coroutine", None) + func = getattr(tool, "func", None) + + if inspect.iscoroutinefunction(coroutine): + + async def wrapped_coroutine(**kwargs: Any) -> Any: + while True: + try: + return await coroutine(**kwargs) + except NeedsAuthorization as exc: + logger.info( + "MCP auth required for '%s' — interrupting run", + exc.mcp_name, + ) + interrupt(_mcp_auth_interrupt_payload(exc)) + + try: + return tool.model_copy(update={"coroutine": wrapped_coroutine}) + except Exception: + tool.coroutine = wrapped_coroutine + return tool + + if func is not None and inspect.isfunction(func): + + def wrapped_func(**kwargs: Any) -> Any: + while True: + try: + return func(**kwargs) + except NeedsAuthorization as exc: + interrupt(_mcp_auth_interrupt_payload(exc)) + + try: + return tool.model_copy(update={"func": wrapped_func}) + except Exception: + tool.func = wrapped_func + return tool + + return tool diff --git a/deep_agent/aegra/middleware.py b/deep_agent/aegra/middleware.py new file mode 100644 index 00000000..9f5f1d06 --- /dev/null +++ b/deep_agent/aegra/middleware.py @@ -0,0 +1,113 @@ +"""Authentication and authorization middleware for aegra deployment (MR-22). + +Provides configurable auth strategies for the LangGraph Platform API: +- ``noop``: No authentication (development) +- ``api_key``: Simple API key validation via X-API-Key header +- ``jwt``: JWT bearer token validation (production) + +The active strategy is selected via the ``LANGGRAPH_AUTH_TYPE`` env var. +""" + +import hashlib +import hmac +import os +import time +from typing import Any + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +AUTH_TYPE = os.environ.get("LANGGRAPH_AUTH_TYPE", "noop") +API_KEY = os.environ.get("LANGGRAPH_API_KEY", "") +JWT_SECRET = os.environ.get("LANGGRAPH_JWT_SECRET", "") +JWT_ALGORITHM = os.environ.get("LANGGRAPH_JWT_ALGORITHM", "HS256") + + +class AuthError(Exception): + """Raised when authentication fails.""" + + def __init__(self, message: str, status_code: int = 401): + """Initialize with error message and HTTP status code.""" + self.message = message + self.status_code = status_code + super().__init__(message) + + +def validate_api_key(provided_key: str) -> bool: + """Constant-time comparison of API keys to prevent timing attacks.""" + if not API_KEY: + logger.warning("LANGGRAPH_API_KEY not set — all keys accepted") + return True + return hmac.compare_digest(provided_key.encode(), API_KEY.encode()) + + +def validate_jwt_token(token: str) -> dict[str, Any]: + """Validate a JWT token and return its claims. + + Requires ``PyJWT`` to be installed. Falls back to a simple + HMAC-based validation if PyJWT is unavailable. + """ + try: + import jwt + + claims: dict[str, Any] = jwt.decode( + token, JWT_SECRET, algorithms=[JWT_ALGORITHM] + ) + if claims.get("exp") and claims["exp"] < time.time(): + raise AuthError("Token expired") + return claims + except ImportError: + logger.warning("PyJWT not installed — using HMAC fallback validation") + return _hmac_validate(token) + except Exception as exc: + raise AuthError(f"JWT validation failed: {exc}") from exc + + +def _hmac_validate(token: str) -> dict[str, Any]: + """Minimal HMAC-based token validation without PyJWT.""" + parts = token.split(".") + if len(parts) != 3: + raise AuthError("Malformed token") + + signature_input = f"{parts[0]}.{parts[1]}".encode() + expected = hashlib.sha256(JWT_SECRET.encode() + signature_input).hexdigest() + + if not hmac.compare_digest(parts[2], expected): + raise AuthError("Invalid token signature") + + return {"sub": "hmac-validated", "token_prefix": token[:20]} + + +def authenticate(headers: dict[str, str]) -> dict[str, Any]: + """Authenticate a request based on the configured auth type. + + Args: + headers: Request headers (case-insensitive keys). + + Returns: + Auth context dict with user info (empty for noop). + + Raises: + AuthError: If authentication fails. + """ + if AUTH_TYPE == "noop": + return {} + + if AUTH_TYPE == "api_key": + key = headers.get("x-api-key") or headers.get("X-API-Key") or "" + if not key: + raise AuthError("Missing X-API-Key header") + if not validate_api_key(key): + raise AuthError("Invalid API key") + return {"auth_type": "api_key"} + + if AUTH_TYPE == "jwt": + auth_header = headers.get("authorization") or headers.get("Authorization") or "" + if not auth_header.startswith("Bearer "): + raise AuthError("Missing or malformed Authorization header") + token = auth_header[7:] + claims = validate_jwt_token(token) + return {"auth_type": "jwt", "claims": claims} + + raise AuthError(f"Unknown auth type: {AUTH_TYPE}", status_code=500) diff --git a/deep_agent/aegra/nodes.py b/deep_agent/aegra/nodes.py new file mode 100644 index 00000000..f3969bb1 --- /dev/null +++ b/deep_agent/aegra/nodes.py @@ -0,0 +1,159 @@ +"""Error-handling node wrappers for graph execution. + +Provides decorator-style wrappers that add retry logic, error capture, +and structured logging around graph node functions. These are used by +the graph builder to make the agent resilient in production. + +The deepagents library handles its own internal node execution. These +wrappers sit at the aegra integration boundary, catching errors that +escape the deepagents graph and recording them in platform metadata. +""" + +import asyncio +import time +from collections.abc import Callable +from functools import wraps +from typing import Any + +from tenacity import ( + RetryCallState, + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +MAX_NODE_RETRIES: int = 2 +RETRY_DELAY_SECONDS: float = 1.0 + + +def _log_node_retry(retry_state: RetryCallState) -> None: + """Log node retry attempts.""" + exc = retry_state.outcome.exception() if retry_state.outcome else None + logger.warning( + "Retry %d/%d for node '%s': %s", + retry_state.attempt_number, + retry_state.retry_object.stop.max_attempt_number, + retry_state.fn.__name__ if retry_state.fn else "unknown", + exc, + ) + + +def with_error_handling(node_name: str) -> Callable[..., Any]: + """Decorator that adds structured error handling to a graph node. + + Catches exceptions, logs them with the node name for traceability, + and re-raises after recording the failure. Used during graph + construction to wrap custom nodes added around the deepagents core. + + Args: + node_name: Human-readable name for log messages. + """ + + def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: + @wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return fn(*args, **kwargs) + except Exception: + logger.exception("Node '%s' failed", node_name) + raise + + @wraps(fn) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return await fn(*args, **kwargs) + except Exception: + logger.exception("Node '%s' failed", node_name) + raise + + if asyncio.iscoroutinefunction(fn): + return async_wrapper + return wrapper + + return decorator + + +def with_retry( + max_retries: int = MAX_NODE_RETRIES, + delay: float = RETRY_DELAY_SECONDS, + retry_on: tuple[type[Exception], ...] = (Exception,), +) -> Callable[..., Any]: + """Decorator that retries a node function on failure using tenacity. + + Supports both sync and async functions with exponential backoff. + Intended for nodes that call external services (MCP tools, LLM APIs) + where transient failures are expected. + + Args: + max_retries: Maximum number of retry attempts. + delay: Base delay in seconds (multiplied exponentially). + retry_on: Tuple of exception types to retry on. + """ + + def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: + tenacity_retry = retry( + retry=retry_if_exception_type(retry_on), + stop=stop_after_attempt(max_retries + 1), + wait=wait_exponential(multiplier=delay, min=delay, max=delay * 10), + before_sleep=_log_node_retry, + reraise=True, + ) + + if asyncio.iscoroutinefunction(fn): + + @wraps(fn) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + @tenacity_retry + async def _inner() -> Any: + return await fn(*args, **kwargs) + + return await _inner() + + return async_wrapper + else: + wrapped: Callable[..., Any] = tenacity_retry(fn) + return wrapped + + return decorator + + +def timed_node(fn: Callable[..., Any]) -> Callable[..., Any]: + """Decorator that logs execution duration of a node function. + + Supports both sync and async functions. + """ + + @wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + start = time.perf_counter() + try: + result = fn(*args, **kwargs) + elapsed = time.perf_counter() - start + logger.info("Node '%s' completed in %.2fs", fn.__name__, elapsed) + return result + except Exception: + elapsed = time.perf_counter() - start + logger.error("Node '%s' failed after %.2fs", fn.__name__, elapsed) + raise + + @wraps(fn) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + start = time.perf_counter() + try: + result = await fn(*args, **kwargs) + elapsed = time.perf_counter() - start + logger.info("Node '%s' completed in %.2fs", fn.__name__, elapsed) + return result + except Exception: + elapsed = time.perf_counter() - start + logger.error("Node '%s' failed after %.2fs", fn.__name__, elapsed) + raise + + if asyncio.iscoroutinefunction(fn): + return async_wrapper + return wrapper diff --git a/deep_agent/aegra/otel.py b/deep_agent/aegra/otel.py new file mode 100644 index 00000000..10316802 --- /dev/null +++ b/deep_agent/aegra/otel.py @@ -0,0 +1,920 @@ +"""OpenTelemetry instrumentation for the template agent. + +Provides centralized telemetry with: +- OTLP exporter when enabled (via YAML or ENABLE_OTEL env var) +- InMemoryMetricReader (no-op) when disabled +- FastAPI auto-instrumentation for distributed tracing +- Conversation, streaming, and thread management metrics + +Config resolution order (highest wins): + 1. Environment variables (ENABLE_OTEL, OTEL_EXPORTER_OTLP_ENDPOINT, ...) + 2. observability.yaml otel: section + 3. Pydantic model defaults + +INSTRUMENTATION STATUS: +- record_conversation_started/completed: Ready for wiring to conversation lifecycle +- record_message_sent: Ready for wiring to message ingress/egress +- record_stream_started/first_token/completed/error: Ready for wiring to streaming handlers +- record_thread_created/deleted/deleted_bulk: Ready for wiring to thread management endpoints +- record_thread_messages: Ready for wiring to thread finalization +Currently, these helpers are defined but not yet called from runtime modules. +""" + +import os +import socket +import threading +import time +from pathlib import Path +from typing import Any, Optional + +from opentelemetry import metrics, trace +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics._internal.aggregation import ( + ExplicitBucketHistogramAggregation, +) +from opentelemetry.sdk.metrics.view import View +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +# Default fallback values - only used when config loading fails +_DEFAULT_SERVICE_NAME = "template-agent" +_DEFAULT_SERVICE_VERSION = "dev" + +# Config validation constants +MIN_EXPORT_INTERVAL_MS = 1000 +MAX_EXPORT_INTERVAL_MS = 60000 + +# Cached service version (populated on first resolution) +_resolved_version: Optional[str] = None +_version_lock = threading.Lock() + +DURATION_BUCKETS = [ + 0.1, + 0.25, + 0.5, + 1.0, + 2.0, + 5.0, + 10.0, + 15.0, + 30.0, + 60.0, + 120.0, + 300.0, +] + +TTFT_BUCKETS = [ + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.0, + 5.0, + 10.0, +] + +MESSAGES_COUNT_BUCKETS = [ + 1, + 2, + 5, + 10, + 20, + 50, + 100, + 200, + 500, +] + +# --------------------------------------------------------------------------- +# Dual tracing architecture +# --------------------------------------------------------------------------- +# This agent has TWO independent observability systems: +# +# 1. OpenTelemetry (this module) — metrics + distributed tracing. +# Uses an SDK TracerProvider stored in ``_tracer_provider`` AND set as +# the global provider (required by FastAPI auto-instrumentation). +# Custom spans are created via ``get_tracer()`` which reads from +# ``_tracer_provider`` directly, not the global. +# +# 2. Langfuse (telemetry.py) — LLM-specific tracing via LangChain's +# ``register_configure_hook`` + ``CallbackHandler``. Langfuse does +# NOT use the OTEL TracerProvider; it has its own SDK. +# +# The two systems coexist without conflict. Langfuse traces LLM calls +# with prompt/completion detail; OTEL traces infrastructure spans +# (graph builds, MCP connections, memory ops) and exports metrics. +# --------------------------------------------------------------------------- + +_tracer_provider: Optional[TracerProvider] = None +_meter: Optional[metrics.Meter] = None +_metrics_container: Optional["MetricsContainer"] = None +_snapshot_reader: Optional[Any] = None +_initialized: bool = False +_otel_enabled: bool = False + +_threads_active_tracked: set[str] = set() +_threads_active_lock = threading.Lock() + + +class MetricsContainer: + """Container for all template agent OpenTelemetry metric instruments.""" + + def __init__(self, meter: metrics.Meter, prefix: Optional[str] = None) -> None: + """Create all metric instruments on the given meter. + + Args: + meter: OpenTelemetry meter instance + prefix: Metric name prefix (defaults to service name from config) + """ + if prefix is None: + prefix = _normalize_metric_prefix(_resolve_service_name()) + self._prefix = prefix + + self.conversations_total = meter.create_counter( + name=f"{self._prefix}_conversations_total", + description="Total conversations by status", + unit="1", + ) + self.messages_total = meter.create_counter( + name=f"{self._prefix}_messages_total", + description="Messages sent/received", + unit="1", + ) + self.conversation_duration_seconds = meter.create_histogram( + name=f"{self._prefix}_conversation_duration_seconds", + description="Time from conversation start to completion", + unit="s", + ) + self.active_conversations = meter.create_up_down_counter( + name=f"{self._prefix}_active_conversations", + description="Currently active conversations", + unit="1", + ) + + self.stream_tokens_total = meter.create_counter( + name=f"{self._prefix}_stream_tokens_total", + description="Tokens streamed to clients", + unit="1", + ) + self.stream_duration_seconds = meter.create_histogram( + name=f"{self._prefix}_stream_duration_seconds", + description="Time to complete stream", + unit="s", + ) + self.stream_errors_total = meter.create_counter( + name=f"{self._prefix}_stream_errors_total", + description="Stream failures by type", + unit="1", + ) + self.time_to_first_token_seconds = meter.create_histogram( + name=f"{self._prefix}_time_to_first_token_seconds", + description="Latency until first token", + unit="s", + ) + + self.threads_created_total = meter.create_counter( + name=f"{self._prefix}_threads_created_total", + description="New threads created", + unit="1", + ) + self.threads_active = meter.create_up_down_counter( + name=f"{self._prefix}_threads_active", + description="Currently active threads", + unit="1", + ) + self.threads_deleted_total = meter.create_counter( + name=f"{self._prefix}_threads_deleted_total", + description="Threads deleted", + unit="1", + ) + self.thread_messages_count = meter.create_histogram( + name=f"{self._prefix}_thread_messages_count", + description="Messages per thread", + unit="1", + ) + + # ponytail: seed all instruments so /api/metrics shows them from startup. + # OTEL SDK only reports instruments after first measurement. + self.conversations_total.add(0) + self.messages_total.add(0) + self.conversation_duration_seconds.record(0) + self.active_conversations.add(0) + self.stream_tokens_total.add(0) + self.stream_duration_seconds.record(0) + self.stream_errors_total.add(0) + self.time_to_first_token_seconds.record(0) + self.threads_created_total.add(0) + self.threads_active.add(0) + self.threads_deleted_total.add(0) + self.thread_messages_count.record(0) + + # Graph build metric + self.graph_build_duration_seconds = meter.create_histogram( + name=f"{self._prefix}_graph_build_duration_seconds", + description="Time to build and compile graph", + unit="s", + ) + self.graph_build_duration_seconds.record(0) + + # Code execution metrics + self.code_execution_duration_seconds = meter.create_histogram( + name=f"{self._prefix}_code_execution_duration_seconds", + description="End-to-end code execution duration", + unit="s", + ) + self.code_executions_total = meter.create_counter( + name=f"{self._prefix}_code_executions_total", + description="Total code executions by outcome", + unit="1", + ) + self.code_execution_errors_total = meter.create_counter( + name=f"{self._prefix}_code_execution_errors_total", + description="Code execution errors by type", + unit="1", + ) + self.code_execution_scheduling_seconds = meter.create_histogram( + name=f"{self._prefix}_code_execution_scheduling_seconds", + description="K8s pod scheduling latency", + unit="s", + ) + self.code_execution_active = meter.create_up_down_counter( + name=f"{self._prefix}_code_execution_active", + description="Currently running code executions", + unit="1", + ) + self.code_execution_queue_wait_seconds = meter.create_histogram( + name=f"{self._prefix}_code_execution_queue_wait_seconds", + description="Time waiting in execution queue", + unit="s", + ) + self.code_execution_rejected_total = meter.create_counter( + name=f"{self._prefix}_code_execution_rejected_total", + description="Executions rejected due to queue full", + unit="1", + ) + self.code_execution_cpu_seconds = meter.create_histogram( + name=f"{self._prefix}_code_execution_cpu_seconds", + description="CPU seconds consumed per execution", + unit="s", + ) + self.code_execution_memory_mb_seconds = meter.create_histogram( + name=f"{self._prefix}_code_execution_memory_mb_seconds", + description="Memory MB-seconds consumed per execution", + unit="MB.s", + ) + # Seed code execution instruments + self.code_execution_duration_seconds.record(0) + self.code_executions_total.add(0) + self.code_execution_errors_total.add(0) + self.code_execution_scheduling_seconds.record(0) + self.code_execution_active.add(0) + self.code_execution_queue_wait_seconds.record(0) + self.code_execution_rejected_total.add(0) + self.code_execution_cpu_seconds.record(0) + self.code_execution_memory_mb_seconds.record(0) + + +def _normalize_metric_prefix(service_name: str) -> str: + """Convert a service display name to a valid OTEL metric name prefix.""" + prefix = service_name.strip().lower() + for char in (" ", "-"): + prefix = prefix.replace(char, "_") + while "__" in prefix: + prefix = prefix.replace("__", "_") + return prefix.strip("_") or "template_agent" + + +def _resolve_service_name() -> str: + """Resolve service name from agent config with unique fallback. + + Returns service name from agent config. If config loading fails, + falls back to hostname+PID-based unique name and logs an error. + + Returns: + Service name string (may contain hyphens or underscores) + """ + try: + from deep_agent.src.agent.config import agent_config + + return agent_config.get_name() + except Exception as exc: + # Use hostname + PID to guarantee uniqueness even on the same host + hostname = socket.gethostname() + pid = os.getpid() + fallback = f"{_DEFAULT_SERVICE_NAME}-{hostname}-{pid}" + logger.error( + "Failed to resolve service name from config, using hostname+PID fallback '%s'. " + "This may cause metric namespace fragmentation in multi-agent deployments. " + "Fix agent config loading to resolve this. Error: %s", + fallback, + exc, + ) + return fallback + + +def _resolve_service_version() -> str: + """Resolve service version from env var, package metadata, or pyproject.toml. + + Resolution order: + 1. APPLICATION_VERSION environment variable (Kubernetes deployments) — not cached + 2. Package metadata via importlib.metadata.version — cached after first read + 3. pyproject.toml version field (development) — cached after first read + 4. Fallback to "dev" + + Returns: + Version string (e.g., "1.2.3", "dev") + """ + global _resolved_version + + # Try env var first (production deployments, can change at runtime) + version = os.environ.get("APPLICATION_VERSION") + if version: + return version + + if _resolved_version is None: + with _version_lock: + if _resolved_version is None: + # Try package metadata + try: + from importlib.metadata import version as pkg_version + + _resolved_version = pkg_version("deep-agent") + except Exception: + pass + + if _resolved_version is None: + # Try reading from pyproject.toml (development) + try: + pyproject_path = ( + Path(__file__).parent.parent.parent / "pyproject.toml" + ) + if pyproject_path.exists(): + import tomllib + + with open(pyproject_path, "rb") as f: + data = tomllib.load(f) + proj_version = data.get("project", {}).get("version") + if isinstance(proj_version, str) and proj_version: + _resolved_version = proj_version + except Exception: + pass + + if _resolved_version is None: + _resolved_version = _DEFAULT_SERVICE_VERSION + + return _resolved_version + + +def _resolve_config() -> tuple[bool, str, bool, int, bool]: + """Resolve OTEL config: env vars override YAML defaults. + + Returns: + (enabled, endpoint, insecure, export_interval_ms, auto_instrument) + """ + try: + from deep_agent.src.agent.config import agent_config + + cfg = agent_config.get_otel_config() + except Exception: + from deep_agent.src.agent.config.otel import OtelFileConfig + + cfg = OtelFileConfig() + + enabled = os.environ.get("ENABLE_OTEL", str(cfg.enabled)).lower() == "true" + endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", cfg.exporter.endpoint) + insecure = ( + os.environ.get( + "OTEL_EXPORTER_OTLP_INSECURE", str(cfg.exporter.insecure) + ).lower() + == "true" + ) + export_interval_raw = int( + os.environ.get( + "OTEL_METRIC_EXPORT_INTERVAL", str(cfg.metrics.export_interval_ms) + ) + ) + # Validate export_interval is within allowed range (same as Pydantic model) + if not (MIN_EXPORT_INTERVAL_MS <= export_interval_raw <= MAX_EXPORT_INTERVAL_MS): + logger.warning( + "OTEL_METRIC_EXPORT_INTERVAL=%d outside valid range [%d, %d], " + "checking config default", + export_interval_raw, + MIN_EXPORT_INTERVAL_MS, + MAX_EXPORT_INTERVAL_MS, + ) + # Validate config default is also within range + if not ( + MIN_EXPORT_INTERVAL_MS + <= cfg.metrics.export_interval_ms + <= MAX_EXPORT_INTERVAL_MS + ): + logger.error( + "Config default export_interval_ms=%d also outside valid range, " + "using minimum allowed value %d", + cfg.metrics.export_interval_ms, + MIN_EXPORT_INTERVAL_MS, + ) + export_interval = MIN_EXPORT_INTERVAL_MS + else: + export_interval = cfg.metrics.export_interval_ms + else: + export_interval = export_interval_raw + + auto_instrument = cfg.tracing.fastapi_auto_instrument + + return enabled, endpoint, insecure, export_interval, auto_instrument + + +def _build_resource() -> Resource: + """Build the OTel resource with service metadata.""" + environment = os.environ.get("ENVIRONMENT", "dev") + version = _resolve_service_version() + instance_id = os.environ.get("HOSTNAME", "local") + + return Resource.create( + { + "service.name": _resolve_service_name(), + "service.version": version, + "service.instance.id": instance_id, + "deployment.environment": environment, + } + ) + + +def _create_histogram_views(prefix: Optional[str] = None) -> list[View]: + """Create histogram bucket views for metrics. + + Args: + prefix: Metric name prefix (defaults to service name from config) + """ + if prefix is None: + prefix = _normalize_metric_prefix(_resolve_service_name()) + return [ + View( + instrument_name=f"{prefix}_conversation_duration_seconds", + aggregation=ExplicitBucketHistogramAggregation(boundaries=DURATION_BUCKETS), + ), + View( + instrument_name=f"{prefix}_stream_duration_seconds", + aggregation=ExplicitBucketHistogramAggregation(boundaries=DURATION_BUCKETS), + ), + View( + instrument_name=f"{prefix}_time_to_first_token_seconds", + aggregation=ExplicitBucketHistogramAggregation(boundaries=TTFT_BUCKETS), + ), + View( + instrument_name=f"{prefix}_thread_messages_count", + aggregation=ExplicitBucketHistogramAggregation( + boundaries=MESSAGES_COUNT_BUCKETS, + ), + ), + View( + instrument_name=f"{prefix}_graph_build_duration_seconds", + aggregation=ExplicitBucketHistogramAggregation(boundaries=DURATION_BUCKETS), + ), + View( + instrument_name=f"{prefix}_code_execution_duration_seconds", + aggregation=ExplicitBucketHistogramAggregation(boundaries=DURATION_BUCKETS), + ), + View( + instrument_name=f"{prefix}_code_execution_scheduling_seconds", + aggregation=ExplicitBucketHistogramAggregation(boundaries=DURATION_BUCKETS), + ), + ] + + +def initialize_telemetry() -> None: + """Initialize OpenTelemetry metrics and tracing. + + Reads config from observability.yaml with env var overrides. + When disabled (default), uses InMemoryMetricReader and NoOpTracerProvider. + When enabled, configures OTLP gRPC exporters for both metrics and traces. + + The TracerProvider is both stored in ``_tracer_provider`` (for + ``get_tracer()``) and set as the global (for FastAPI auto-instrumentation). + """ + global _meter, _metrics_container, _initialized, _otel_enabled, _tracer_provider + global _snapshot_reader + + if _initialized: + return + + enabled, endpoint, insecure, export_interval, _ = _resolve_config() + _otel_enabled = enabled + + service_name = _resolve_service_name() + resource = _build_resource() + metric_prefix = _normalize_metric_prefix(service_name) + views = _create_histogram_views(prefix=metric_prefix) + + from opentelemetry.sdk.metrics.export import InMemoryMetricReader + + # Always keep an InMemoryMetricReader so /api/metrics can read values back. + _snapshot_reader = InMemoryMetricReader() + + if not enabled: + logger.info( + "OTEL disabled (ENABLE_OTEL not true) — metrics and traces are in-memory" + ) + + meter_provider = MeterProvider( + resource=resource, + metric_readers=[_snapshot_reader], + views=views, + ) + tracer_provider = TracerProvider(resource=resource) + else: + from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( + OTLPMetricExporter, + ) + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter, + ) + from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + from opentelemetry.sdk.trace.export import BatchSpanProcessor + + logger.info( + "OTEL enabled — exporting to %s (interval=%sms)", + endpoint, + export_interval, + ) + + otlp_metric_exporter = OTLPMetricExporter(endpoint=endpoint, insecure=insecure) + otlp_reader = PeriodicExportingMetricReader( + otlp_metric_exporter, + export_interval_millis=export_interval, + ) + meter_provider = MeterProvider( + resource=resource, + metric_readers=[otlp_reader, _snapshot_reader], + views=views, + ) + + otlp_span_exporter = OTLPSpanExporter(endpoint=endpoint, insecure=insecure) + tracer_provider = TracerProvider(resource=resource) + tracer_provider.add_span_processor(BatchSpanProcessor(otlp_span_exporter)) + + _tracer_provider = tracer_provider + + metrics.set_meter_provider(meter_provider) + # Set global TracerProvider — required by FastAPI auto-instrumentation + # which reads from trace.get_tracer_provider(). Custom spans use + # get_tracer() which reads _tracer_provider directly. + trace.set_tracer_provider(tracer_provider) + + service_version = _resolve_service_version() + _meter = meter_provider.get_meter(service_name, service_version) + _metrics_container = MetricsContainer(_meter, prefix=metric_prefix) + _initialized = True + + +def instrument_fastapi(app: Any) -> None: + """Auto-instrument a FastAPI app for distributed tracing. + + Only instruments if OTEL is initialized and auto-instrumentation + is enabled in config. Safe to call before initialize_telemetry() — + instrumentation picks up the global TracerProvider lazily. + """ + _, _, _, _, auto_instrument = _resolve_config() + if not auto_instrument: + logger.debug("FastAPI auto-instrumentation disabled in config") + return + + try: + from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor + + FastAPIInstrumentor.instrument_app(app) + logger.info("FastAPI auto-instrumentation enabled") + except (ImportError, AttributeError) as exc: + # Package missing or incompatible version + logger.warning( + "FastAPI instrumentation unavailable: %s. " + "Check opentelemetry-instrumentation-fastapi version compatibility.", + exc, + ) + except Exception: + # Unexpected failure, log full trace + logger.error("FastAPI instrumentation failed", exc_info=True) + + +def shutdown_telemetry() -> None: + """Flush and shut down both meter and tracer providers.""" + global \ + _initialized, \ + _tracer_provider, \ + _meter, \ + _metrics_container, \ + _snapshot_reader, \ + _resolved_version + + if _tracer_provider is not None and hasattr(_tracer_provider, "shutdown"): + _tracer_provider.shutdown() + _tracer_provider = None + + meter_provider = metrics.get_meter_provider() + if hasattr(meter_provider, "shutdown"): + meter_provider.shutdown() + + # Clear all module-level state + _meter = None + _metrics_container = None + _snapshot_reader = None + _resolved_version = None # Clear cached version for clean re-initialization + + reset_thread_active_tracking() + _initialized = False + + +def get_metrics() -> Optional[MetricsContainer]: + """Return the global metrics container, or None if not initialized.""" + return _metrics_container + + +def get_tracer(name: Optional[str] = None) -> trace.Tracer: + """Return a tracer from the module-owned TracerProvider. + + Uses ``_tracer_provider`` (set during ``initialize_telemetry``) rather + than the global provider, keeping ownership explicit. If telemetry has + not been initialised yet, falls back to the global (which may be a + no-op ``ProxyTracerProvider``). + + Args: + name: Instrumentation scope name (defaults to service name from config) + + Returns: + An OTEL ``Tracer`` instance. + """ + if name is None: + name = _resolve_service_name() + if _tracer_provider is not None: + return _tracer_provider.get_tracer(name) + return trace.get_tracer(name) + + +def is_tracing_enabled() -> bool: + """Return True if OTEL has been initialised and is enabled.""" + return _initialized and _otel_enabled + + +def get_metrics_snapshot() -> dict[str, Any]: + """Read current metric values from the InMemoryMetricReader. + + Returns a flat dict keyed by metric name. Counters/UpDownCounters + are summed across all attribute sets. Histograms aggregate count + and sum across all attribute sets. + """ + if _snapshot_reader is None: + return {} + + data = _snapshot_reader.get_metrics_data() + if data is None: + return {} + + result: dict[str, Any] = {} + + for resource_metrics in data.resource_metrics: + for scope_metrics in resource_metrics.scope_metrics: + for metric in scope_metrics.metrics: + name = metric.name + points = list(metric.data.data_points) + if not points: + continue + + if hasattr(points[0], "bucket_counts"): + total_count = sum(pt.count for pt in points) + total_sum = sum(pt.sum for pt in points) + result[name] = { + "count": total_count, + "sum": round(total_sum, 3), + } + else: + result[name] = sum(pt.value for pt in points) + + return result + + +# --------------------------------------------------------------------------- +# Helper instrumentation functions +# --------------------------------------------------------------------------- + + +def _attrs(extra: Optional[dict[str, Any]] = None) -> dict[str, str]: + """Merge optional extra attributes, stringifying values for OTel.""" + if not extra: + return {} + return {k: str(v) for k, v in extra.items() if v is not None} + + +def _release_thread_active_if_tracked(thread_id: str) -> bool: + with _threads_active_lock: + if thread_id in _threads_active_tracked: + _threads_active_tracked.discard(thread_id) + return True + return False + + +def reset_thread_active_tracking() -> None: + """Clear in-process thread tracking (for tests and shutdown).""" + with _threads_active_lock: + _threads_active_tracked.clear() + + +def record_conversation_started( + *, + status: str = "started", + attributes: Optional[dict[str, Any]] = None, +) -> float: + """Record conversation start. Returns monotonic timestamp for duration.""" + m = get_metrics() + if m: + base_attrs = _attrs(attributes) + m.conversations_total.add(1, {"status": status, **base_attrs}) + m.active_conversations.add(1, base_attrs) + return time.monotonic() + + +def record_conversation_completed( + start_mono: float, + *, + status: str = "completed", + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record conversation completion with duration.""" + m = get_metrics() + if m: + base_attrs = _attrs(attributes) + merged = {"status": status, **base_attrs} + duration = time.monotonic() - start_mono + m.conversations_total.add(1, merged) + m.active_conversations.add(-1, base_attrs) + m.conversation_duration_seconds.record(duration, merged) + + +def record_message_sent( + *, + direction: str = "sent", + message_type: str = "human", + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record a message sent or received.""" + m = get_metrics() + if m: + merged = { + "direction": direction, + "message_type": message_type, + **_attrs(attributes), + } + m.messages_total.add(1, merged) + + +def record_stream_started() -> float: + """Record stream start. Returns monotonic timestamp.""" + return time.monotonic() + + +def record_first_token( + stream_start_mono: float, + *, + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record time-to-first-token from stream start.""" + m = get_metrics() + if m: + ttft = time.monotonic() - stream_start_mono + m.time_to_first_token_seconds.record(ttft, _attrs(attributes)) + + +def record_stream_completed( + stream_start_mono: float, + token_count: int, + *, + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record stream completion with duration and token count.""" + m = get_metrics() + if m: + merged = _attrs(attributes) + duration = time.monotonic() - stream_start_mono + m.stream_duration_seconds.record(duration, merged) + m.stream_tokens_total.add(token_count, merged) + + +def record_stream_error( + *, + error_type: str = "unknown", + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record a stream error.""" + m = get_metrics() + if m: + merged = {"error_type": error_type, **_attrs(attributes)} + m.stream_errors_total.add(1, merged) + + +def record_thread_created( + *, + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record thread creation with active tracking.""" + m = get_metrics() + if not m: + return + + merged = _attrs(attributes) + thread_id = merged.get("thread_id") + + # Determine if we should increment the active gauge inside the lock + should_increment = True + if thread_id: + with _threads_active_lock: + if thread_id in _threads_active_tracked: + should_increment = False # Already tracked, don't increment + else: + _threads_active_tracked.add(thread_id) + + # Record metrics outside the lock + m.threads_created_total.add(1, merged) + if should_increment: + m.threads_active.add(1, merged) + + +def record_thread_deleted( + *, + count: int = 1, + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record thread deletion. Decrements active only if previously tracked.""" + if count != 1: + raise ValueError( + f"record_thread_deleted requires count=1, got {count}. " + "Use record_threads_deleted_bulk for batch deletion." + ) + m = get_metrics() + if not m: + return + merged = _attrs(attributes) + m.threads_deleted_total.add(count, merged) + tid = merged.get("thread_id") + if not tid: + return + if _release_thread_active_if_tracked(str(tid)): + m.threads_active.add(-1, merged) + + +def record_threads_deleted_bulk( + deleted_thread_ids: list[str], + *, + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record bulk thread deletion with per-ID active tracking.""" + m = get_metrics() + if not m or not deleted_thread_ids: + return + base = _attrs(attributes) + m.threads_deleted_total.add(len(deleted_thread_ids), base) + for tid in deleted_thread_ids: + row = {**base, "thread_id": tid} + if _release_thread_active_if_tracked(tid): + m.threads_active.add(-1, row) + + +def record_thread_messages( + message_count: int, + *, + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record the final message count for a thread.""" + m = get_metrics() + if m: + m.thread_messages_count.record(message_count, _attrs(attributes)) + + +def record_graph_built( + build_start_mono: float, + *, + cache_hit: bool = False, + mcp_tool_count: int = 0, + attributes: Optional[dict[str, Any]] = None, +) -> None: + """Record graph build completion with cache hit status and tool count. + + Args: + build_start_mono: Monotonic timestamp from when graph build started + cache_hit: Whether the graph was retrieved from cache + mcp_tool_count: Number of MCP tools loaded into the graph + attributes: Additional attributes to attach to the metric + """ + m = get_metrics() + if m: + duration = time.monotonic() - build_start_mono + merged = { + "cache_hit": str(cache_hit), + "mcp_tools": str(mcp_tool_count), + **_attrs(attributes), + } + m.graph_build_duration_seconds.record(duration, merged) diff --git a/deep_agent/aegra/redis.py b/deep_agent/aegra/redis.py new file mode 100644 index 00000000..87b4952f --- /dev/null +++ b/deep_agent/aegra/redis.py @@ -0,0 +1,229 @@ +"""Redis connection configuration for aegra deployment (MR-20). + +Provides a Redis client factory for caching, rate limiting, and +pub/sub in the LangGraph Platform deployment. Falls back gracefully +if Redis is unavailable — the agent operates without caching. + +Environment variables: + REDIS_URL: Full Redis URL (default: redis://localhost:6379/0) + REDIS_MAX_CONNECTIONS: Pool size (default: 10) + REDIS_SOCKET_TIMEOUT: Seconds (default: 5) + REDIS_RETRY_ON_TIMEOUT: Enable retry (default: true) +""" + +import asyncio +import os +import secrets +import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any, Literal + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379/0") +REDIS_MAX_CONNECTIONS = int(os.environ.get("REDIS_MAX_CONNECTIONS", "10")) +REDIS_SOCKET_TIMEOUT = int(os.environ.get("REDIS_SOCKET_TIMEOUT", "5")) +REDIS_RETRY_ON_TIMEOUT = ( + os.environ.get("REDIS_RETRY_ON_TIMEOUT", "true").lower() == "true" +) +REDIS_KEY_PREFIX = os.environ.get("REDIS_KEY_PREFIX", "aegra:") + +_client: Any = None + + +def get_redis_config() -> dict[str, Any]: + """Return the Redis configuration dict for documentation/debugging.""" + return { + "url": REDIS_URL, + "max_connections": REDIS_MAX_CONNECTIONS, + "socket_timeout": REDIS_SOCKET_TIMEOUT, + "retry_on_timeout": REDIS_RETRY_ON_TIMEOUT, + "key_prefix": REDIS_KEY_PREFIX, + } + + +def get_redis_client() -> Any: + """Get or create a Redis client with connection pooling. + + Returns: + Redis client instance, or None if Redis is unavailable. + """ + global _client # noqa: PLW0603 + if _client is not None: + return _client + + try: + import redis + + _client = redis.from_url( + REDIS_URL, + max_connections=REDIS_MAX_CONNECTIONS, + socket_timeout=REDIS_SOCKET_TIMEOUT, + retry_on_timeout=REDIS_RETRY_ON_TIMEOUT, + decode_responses=True, + ) + _client.ping() + logger.info("Redis connected: %s", REDIS_URL) + return _client + except ImportError: + logger.warning("redis package not installed — caching disabled") + return None + except Exception: + logger.warning( + "Redis unavailable at %s — caching disabled", REDIS_URL, exc_info=True + ) + _client = None + return None + + +def close_redis_client() -> None: + """Close the Redis client connection if open. Idempotent.""" + global _client # noqa: PLW0603 + if _client is None: + return + try: + _client.close() + logger.info("Redis client closed") + except Exception: + logger.debug("Redis close error", exc_info=True) + finally: + _client = None + + +def cache_get(key: str) -> str | None: + """Read a value from Redis cache. Returns None on miss or error.""" + client = get_redis_client() + if client is None: + return None + try: + val = client.get(f"{REDIS_KEY_PREFIX}{key}") + return str(val) if val is not None else None + except Exception: + logger.debug("Cache read failed for key '%s'", key, exc_info=True) + return None + + +def cache_set(key: str, value: str, ttl_seconds: int = 300) -> bool: + """Write a value to Redis cache with TTL. Returns False on error.""" + client = get_redis_client() + if client is None: + return False + try: + client.setex(f"{REDIS_KEY_PREFIX}{key}", ttl_seconds, value) + return True + except Exception: + logger.debug("Cache write failed for key '%s'", key, exc_info=True) + return False + + +def cache_set_persistent(key: str, value: str) -> bool: + """Write a value to Redis without expiry. Returns False on error.""" + client = get_redis_client() + if client is None: + return False + try: + client.set(f"{REDIS_KEY_PREFIX}{key}", value) + return True + except Exception: + logger.debug("Persistent cache write failed for key '%s'", key, exc_info=True) + return False + + +def cache_delete(key: str) -> bool: + """Delete a key from Redis cache. Returns False on error.""" + client = get_redis_client() + if client is None: + return False + try: + client.delete(f"{REDIS_KEY_PREFIX}{key}") + return True + except Exception: + return False + + +_RELEASE_LOCK_LUA = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +else + return 0 +end +""" + + +def _lock_key(name: str) -> str: + return f"{REDIS_KEY_PREFIX}lock:{name}" + + +def acquire_distributed_lock( + name: str, + *, + ttl_seconds: int = 30, + wait_seconds: float = 10.0, + poll_interval: float = 0.05, +) -> str | None: + """Acquire a Redis lock. Returns a token, or None if unavailable or timed out.""" + client = get_redis_client() + if client is None: + return None + + token = secrets.token_urlsafe(16) + key = _lock_key(name) + deadline = time.monotonic() + wait_seconds + + while True: + try: + if client.set(key, token, nx=True, ex=ttl_seconds): + return token + except Exception: + logger.debug("Lock acquire failed for '%s'", name, exc_info=True) + return None + + if time.monotonic() >= deadline: + return None + time.sleep(poll_interval) + + +def release_distributed_lock(name: str, token: str) -> bool: + """Release a Redis lock when the token still matches.""" + client = get_redis_client() + if client is None: + return False + try: + return bool(client.eval(_RELEASE_LOCK_LUA, 1, _lock_key(name), token)) + except Exception: + logger.debug("Lock release failed for '%s'", name, exc_info=True) + return False + + +LockState = Literal["held", "no_redis", "timeout"] + + +@asynccontextmanager +async def distributed_lock( + name: str, + *, + ttl_seconds: int = 30, + wait_seconds: float = 10.0, +) -> AsyncIterator[LockState]: + """Yield lock state for a Redis-backed distributed lock.""" + if get_redis_client() is None: + yield "no_redis" + return + + token = await asyncio.to_thread( + acquire_distributed_lock, + name, + ttl_seconds=ttl_seconds, + wait_seconds=wait_seconds, + ) + if token is None: + yield "timeout" + return + + try: + yield "held" + finally: + await asyncio.to_thread(release_distributed_lock, name, token) diff --git a/deep_agent/aegra/security_middleware.py b/deep_agent/aegra/security_middleware.py new file mode 100644 index 00000000..d0c9fc1d --- /dev/null +++ b/deep_agent/aegra/security_middleware.py @@ -0,0 +1,89 @@ +"""Production security middleware for HTTP security headers and request validation. + +Implements OWASP security recommendations for FastAPI applications. +""" + +from typing import Any + +from fastapi import Request, status +from fastapi.responses import JSONResponse +from starlette.middleware.base import BaseHTTPMiddleware + +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +class SecurityHeadersMiddleware(BaseHTTPMiddleware): + """Add OWASP-recommended security headers to all HTTP responses. + + Headers applied: + - X-Content-Type-Options: nosniff + - X-Frame-Options: DENY + - X-XSS-Protection: 1; mode=block + - Strict-Transport-Security: max-age=31536000; includeSubDomains (HTTPS only) + - Content-Security-Policy: default-src 'self' + - Referrer-Policy: strict-origin-when-cross-origin + - Permissions-Policy: geolocation=(), microphone=(), camera=() + """ + + async def dispatch(self, request: Request, call_next: Any) -> Any: + """Add security headers to response.""" + response = await call_next(request) + + # Always set these headers + response.headers["X-Content-Type-Options"] = "nosniff" + response.headers["X-Frame-Options"] = "DENY" + response.headers["X-XSS-Protection"] = "1; mode=block" + response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin" + response.headers["Permissions-Policy"] = ( + "geolocation=(), microphone=(), camera=()" + ) + + # CSP: allow self for API endpoints, adjust if serving web UI + response.headers["Content-Security-Policy"] = "default-src 'self'" + + # HSTS: only set on HTTPS connections + if request.url.scheme == "https" or settings.is_production: + response.headers["Strict-Transport-Security"] = ( + "max-age=31536000; includeSubDomains" + ) + + return response + + +class RequestSizeLimitMiddleware(BaseHTTPMiddleware): + """Enforce request body size limits to prevent DoS attacks. + + Default limit: 10MB (configurable via REQUEST_BODY_MAX_SIZE). + """ + + def __init__(self, app: Any, max_size_bytes: int = 10 * 1024 * 1024): + """Initialize with configurable max request body size.""" + super().__init__(app) + self.max_size_bytes = max_size_bytes + logger.info("Request body size limit: %d bytes", max_size_bytes) + + async def dispatch(self, request: Request, call_next: Any) -> Any: + """Check request body size before processing.""" + # Skip for GET/HEAD/OPTIONS (no body) + if request.method in ("GET", "HEAD", "OPTIONS"): + return await call_next(request) + + content_length = request.headers.get("content-length") + if content_length and int(content_length) > self.max_size_bytes: + logger.warning( + "Request body too large: %s bytes (max %d)", + content_length, + self.max_size_bytes, + ) + return JSONResponse( + status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, + content={ + "detail": f"Request body exceeds maximum size of {self.max_size_bytes} bytes", + "error_type": "request_too_large", + }, + ) + + return await call_next(request) diff --git a/deep_agent/aegra/serialization.py b/deep_agent/aegra/serialization.py new file mode 100644 index 00000000..992a669a --- /dev/null +++ b/deep_agent/aegra/serialization.py @@ -0,0 +1,140 @@ +"""State serialization and deserialization for aegra deployment (MR-16). + +Converts LangGraph agent state to/from JSON-safe representations for +persistence, API responses, and cross-service communication. Handles +LangChain message objects, tool calls, and nested state structures. +""" + +import json +from datetime import UTC, datetime +from typing import Any + +from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def serialize_message(msg: BaseMessage) -> dict[str, Any]: + """Serialize a single LangChain message to a JSON-safe dict.""" + data: dict[str, Any] = { + "type": msg.type, + "content": msg.content, + "id": getattr(msg, "id", None), + } + + if isinstance(msg, AIMessage) and msg.tool_calls: + data["tool_calls"] = [ + {"id": tc.get("id"), "name": tc["name"], "args": tc["args"]} + for tc in msg.tool_calls + ] + + if isinstance(msg, ToolMessage): + data["tool_call_id"] = msg.tool_call_id + data["name"] = getattr(msg, "name", None) + + if msg.response_metadata: + data["response_metadata"] = _safe_serialize(msg.response_metadata) + + return data + + +def deserialize_message(data: dict[str, Any]) -> BaseMessage: + """Reconstruct a LangChain message from a serialized dict.""" + msg_type = data.get("type", "human") + content = data.get("content", "") + msg_id = data.get("id") + + if msg_type == "human": + return HumanMessage(content=content, id=msg_id) + elif msg_type == "ai": + kwargs: dict[str, Any] = {"content": content, "id": msg_id} + if "tool_calls" in data: + kwargs["tool_calls"] = data["tool_calls"] + return AIMessage(**kwargs) + elif msg_type == "system": + return SystemMessage(content=content, id=msg_id) + elif msg_type == "tool": + return ToolMessage( + content=content, + tool_call_id=data.get("tool_call_id", ""), + name=data.get("name"), + id=msg_id, + ) + else: + return HumanMessage(content=content, id=msg_id) + + +def serialize_state(state: dict[str, Any]) -> dict[str, Any]: + """Serialize full LangGraph state to a JSON-safe dict. + + Walks the state dict, converting LangChain messages and any other + non-serializable objects into JSON-compatible representations. + """ + result: dict[str, Any] = {} + + for key, value in state.items(): + if key == "messages" and isinstance(value, list): + result[key] = [ + serialize_message(m) + if isinstance(m, BaseMessage) + else _safe_serialize(m) + for m in value + ] + else: + result[key] = _safe_serialize(value) + + result["_serialized_at"] = datetime.now(UTC).isoformat() + return result + + +def deserialize_state(data: dict[str, Any]) -> dict[str, Any]: + """Reconstruct LangGraph state from a serialized dict.""" + result: dict[str, Any] = {} + + for key, value in data.items(): + if key == "_serialized_at": + continue + elif key == "messages" and isinstance(value, list): + result[key] = [ + deserialize_message(m) if isinstance(m, dict) and "type" in m else m + for m in value + ] + else: + result[key] = value + + return result + + +def state_to_json(state: dict[str, Any], indent: int | None = None) -> str: + """Serialize state to a JSON string.""" + return json.dumps(serialize_state(state), indent=indent, default=str) + + +def state_from_json(json_str: str) -> dict[str, Any]: + """Deserialize state from a JSON string.""" + return deserialize_state(json.loads(json_str)) + + +def _safe_serialize(obj: Any) -> Any: + """Recursively convert an object to a JSON-safe representation.""" + if obj is None or isinstance(obj, (str, int, float, bool)): + return obj + if isinstance(obj, dict): + return {str(k): _safe_serialize(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_safe_serialize(item) for item in obj] + if isinstance(obj, BaseMessage): + return serialize_message(obj) + if isinstance(obj, datetime): + return obj.isoformat() + if isinstance(obj, bytes): + return obj.decode("utf-8", errors="replace") + return str(obj) diff --git a/deep_agent/aegra/shutdown.py b/deep_agent/aegra/shutdown.py new file mode 100644 index 00000000..0eb03a13 --- /dev/null +++ b/deep_agent/aegra/shutdown.py @@ -0,0 +1,352 @@ +"""Shutdown orchestrator — coordinated teardown on SIGTERM. + +Mirrors ``startup.py``: idempotent orchestrator, individual step +functions, structured logging, defensive error handling. + +Two independent paths trigger shutdown: + +1. ``atexit`` callback (registered at import time from ``http_app.py``) + — fires reliably when uvicorn handles SIGTERM and exits normally. + Runs a synchronous cleanup (Langfuse flush, Redis close, graph + cache clear). No event loop needed. + +2. ``loop.add_signal_handler`` (registered on first graph request via + ``startup.py``) — overrides uvicorn's handler, runs the full async + shutdown with drain period. Only active after the first graph + request, but that's when there's actually work to drain. + +Aegra strips our custom app's lifespan and middleware, so neither +ASGI lifespan nor middleware-based registration works. The atexit +path is guaranteed because Aegra always imports ``http_app.py``. + +Both paths call idempotent cleanup — the second call is a no-op. + +Shutdown sequence (within ``terminationGracePeriodSeconds: 60``): + + 1. Set ``_shutting_down`` flag → health probes return 503 + 2. Drain period — in-flight requests finish (async path only) + 3. Flush and stop Langfuse + 4. Stop memory scheduler (async path only) + 5. Clear graph cache + 6. Close Redis +""" + +import asyncio +import os +import signal +import time +from typing import Any + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_shutting_down = False +_shutdown_complete = False +_atexit_registered = False + +SHUTDOWN_DRAIN_SECONDS = int(os.environ.get("SHUTDOWN_DRAIN_SECONDS", "15")) +SHUTDOWN_LANGFUSE_TIMEOUT_SECONDS = int( + os.environ.get("SHUTDOWN_LANGFUSE_TIMEOUT_SECONDS", "5") +) +SHUTDOWN_SCHEDULER_TIMEOUT_SECONDS = int( + os.environ.get("SHUTDOWN_SCHEDULER_TIMEOUT_SECONDS", "10") +) +SHUTDOWN_GRACE_PERIOD_SECONDS = int( + os.environ.get("SHUTDOWN_GRACE_PERIOD_SECONDS", "60") +) + +_TOTAL_BUDGET = ( + SHUTDOWN_DRAIN_SECONDS + + SHUTDOWN_LANGFUSE_TIMEOUT_SECONDS + + SHUTDOWN_SCHEDULER_TIMEOUT_SECONDS +) +_HEADROOM = SHUTDOWN_GRACE_PERIOD_SECONDS - _TOTAL_BUDGET +if _HEADROOM < 5: + logger.warning( + "Shutdown budget (%ds drain + %ds langfuse + %ds scheduler = %ds) " + "leaves only %ds before SIGKILL at %ds. Risk of incomplete cleanup.", + SHUTDOWN_DRAIN_SECONDS, + SHUTDOWN_LANGFUSE_TIMEOUT_SECONDS, + SHUTDOWN_SCHEDULER_TIMEOUT_SECONDS, + _TOTAL_BUDGET, + _HEADROOM, + SHUTDOWN_GRACE_PERIOD_SECONDS, + ) + + +def is_shutting_down() -> bool: + """Return True once shutdown has been initiated.""" + return _shutting_down + + +# -- Primary path: atexit (sync) --------------------------------------------- + + +def register_atexit() -> None: + """Register the sync shutdown as an atexit callback. + + Called at import time from ``http_app.py``. Unlike signal handlers, + atexit callbacks are not overwritten by uvicorn. Idempotent — safe + to call multiple times (tests, reloads). + """ + global _atexit_registered # noqa: PLW0603 + if _atexit_registered: + return + _atexit_registered = True + + import atexit + + atexit.register(run_shutdown_sync) + logger.info("Shutdown atexit handler registered") + + +def run_shutdown_sync() -> None: + """Synchronous shutdown — runs at process exit via atexit. + + Handles cleanup that doesn't need an event loop: Langfuse flush, + Redis close, graph cache clear. Skips drain and async scheduler + stop (those only run in the async path). + """ + global _shutting_down, _shutdown_complete # noqa: PLW0603 + + if _shutdown_complete: + return + if _shutting_down: + logger.debug("Async shutdown already ran — sync cleanup skipped") + _shutdown_complete = True + return + + _shutting_down = True + t0 = time.monotonic() + results: dict[str, str] = {} + + import sys + + print("[shutdown] Graceful shutdown started", file=sys.stderr, flush=True) + logger.info("Sync shutdown initiated (atexit)") + + for key, step in [ + ("otel", _shutdown_otel), + ("langfuse", _shutdown_langfuse_sync), + ("graph_cache", _clear_graph_cache), + ("redis", _close_redis), + ]: + try: + results[key] = step() + except Exception as exc: + logger.warning("Shutdown step '%s' failed: %s", key, exc) + results[key] = f"error: {exc}" + + _shutdown_complete = True + elapsed = round((time.monotonic() - t0) * 1000, 1) + print( + f"[shutdown] Graceful shutdown complete in {elapsed}ms: {results}", + file=sys.stderr, + flush=True, + ) + + +# -- Secondary path: signal handler (async) ---------------------------------- + + +def register_signal_handlers() -> None: + """Install loop-aware SIGTERM/SIGINT handlers. + + Uses ``loop.add_signal_handler`` which overrides uvicorn's handler. + Must be called from inside a running event loop. Called from + ``startup.py`` after the first graph request. + """ + try: + loop = asyncio.get_running_loop() + except RuntimeError: + logger.warning("No running event loop — signal handlers not registered") + return + + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, _handle_signal, sig, loop) + + logger.info("Shutdown signal handlers registered (SIGTERM, SIGINT)") + + +def _handle_signal(signum: int, loop: asyncio.AbstractEventLoop) -> None: + global _shutting_down # noqa: PLW0603 + _shutting_down = True + logger.info("Signal %d received — scheduling async shutdown", signum) + loop.create_task(_shutdown_and_exit()) + + +async def _shutdown_and_exit() -> None: + """Run graceful shutdown then terminate the process.""" + await run_shutdown() + logger.info("Shutdown complete — exiting") + import sys + + sys.exit(0) + + +_async_shutdown_started = False + + +async def run_shutdown() -> dict[str, str]: + """Full async shutdown with drain period. + + Only runs when signal handlers were registered (after first graph + request). Safe to call multiple times — subsequent calls are no-ops. + """ + global _shutting_down, _shutdown_complete, _async_shutdown_started # noqa: PLW0603 + + if _shutdown_complete or _async_shutdown_started: + return {"status": "already_complete"} + + _async_shutdown_started = True + _shutting_down = True + + t0 = time.monotonic() + results: dict[str, str] = {} + + logger.info("Async shutdown initiated") + + for key, step in [ + ("drain", _drain), + ("otel", _shutdown_otel), + ("langfuse", _shutdown_langfuse), + ("scheduler", _stop_scheduler), + ("graph_cache", _clear_graph_cache), + ("redis", _close_redis), + ]: + try: + step_result = step() + if asyncio.iscoroutine(step_result): + step_result = await step_result + results[key] = str(step_result) + except Exception as exc: + logger.warning("Shutdown step '%s' failed: %s", key, exc) + results[key] = f"error: {exc}" + + _shutdown_complete = True + elapsed = round((time.monotonic() - t0) * 1000, 1) + + logger.info("Async shutdown complete in %.1fms: %s", elapsed, results) + return results + + +# -- Individual shutdown steps ----------------------------------------------- + + +async def _drain() -> str: + if SHUTDOWN_DRAIN_SECONDS <= 0: + return "skipped: drain disabled" + logger.info("Draining for %ds", SHUTDOWN_DRAIN_SECONDS) + await asyncio.sleep(SHUTDOWN_DRAIN_SECONDS) + return "ok" + + +async def _shutdown_langfuse() -> str: + try: + from deep_agent.aegra.telemetry import get_langfuse_client + + client = get_langfuse_client() + if client is None: + return "skipped: not configured" + + await asyncio.wait_for( + asyncio.to_thread(_langfuse_shutdown_blocking, client), + timeout=SHUTDOWN_LANGFUSE_TIMEOUT_SECONDS, + ) + return "ok" + except asyncio.TimeoutError: + logger.warning( + "Langfuse shutdown timed out after %ds", SHUTDOWN_LANGFUSE_TIMEOUT_SECONDS + ) + return "timeout" + except Exception as exc: + logger.warning("Langfuse shutdown failed: %s", exc) + return f"error: {exc}" + + +def _shutdown_langfuse_sync() -> str: + """Sync Langfuse flush for atexit path. + + Only flushes if a client was already initialized — avoids creating + a new client during interpreter shutdown (which triggers + ``RuntimeError: cannot schedule new futures``). + """ + try: + from deep_agent.aegra.telemetry import _langfuse_configured + + if not _langfuse_configured(): + return "skipped: not configured" + + from langfuse import get_client + + client = get_client() + _langfuse_shutdown_blocking(client) + return "ok" + except Exception as exc: + return f"skipped: {exc}" + + +def _langfuse_shutdown_blocking(client: Any) -> None: + """Run the sync Langfuse shutdown.""" + if hasattr(client, "shutdown"): + client.shutdown() + elif hasattr(client, "flush"): + client.flush() + + +async def _stop_scheduler() -> str: + try: + from deep_agent.src.memory.scheduler import stop_scheduler + + await asyncio.wait_for( + stop_scheduler(), + timeout=SHUTDOWN_SCHEDULER_TIMEOUT_SECONDS, + ) + return "ok" + except asyncio.TimeoutError: + logger.warning( + "Scheduler stop timed out after %ds", SHUTDOWN_SCHEDULER_TIMEOUT_SECONDS + ) + return "timeout" + except Exception as exc: + logger.warning("Scheduler stop failed: %s", exc) + return f"error: {exc}" + + +def _clear_graph_cache() -> str: + try: + from deep_agent.aegra.graph import _graph_cache, _graph_cache_ts + + count = len(_graph_cache) + _graph_cache.clear() + _graph_cache_ts.clear() + if count > 0: + logger.info("Cleared %d cached graph(s)", count) + return "ok" + except Exception as exc: + logger.warning("Graph cache clear failed: %s", exc) + return f"error: {exc}" + + +def _shutdown_otel() -> str: + """Shutdown OpenTelemetry providers and flush pending telemetry.""" + try: + from deep_agent.aegra.otel import shutdown_telemetry + + shutdown_telemetry() + return "ok" + except Exception as exc: + logger.warning("OTEL shutdown failed: %s", exc) + return f"error: {exc}" + + +def _close_redis() -> str: + try: + from deep_agent.aegra.redis import close_redis_client + + close_redis_client() + return "ok" + except Exception as exc: + logger.warning("Redis close failed: %s", exc) + return f"error: {exc}" diff --git a/deep_agent/aegra/startup.py b/deep_agent/aegra/startup.py new file mode 100644 index 00000000..a1993034 --- /dev/null +++ b/deep_agent/aegra/startup.py @@ -0,0 +1,215 @@ +"""Startup orchestrator — coordinated initialization on process boot. + +Runs once when the agent process starts. Ensures all subsystems +are initialized in the correct order before the server accepts +traffic. + +Startup sequence: + 1. Validate configuration + 2. Ensure database tables exist + 3. Warm caches (if enabled) + 4. Start memory scheduler (if enabled) + 5. Set up Langfuse tracing (if configured) + 6. Log readiness + +This module is idempotent — calling ``run_startup()`` multiple +times is safe (each step guards against double-init). +""" + +import asyncio +import os +import time + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_startup_complete = False + + +async def run_startup() -> dict[str, str]: + """Execute the startup sequence. Returns a status dict. + + Safe to call multiple times — subsequent calls are no-ops. + """ + global _startup_complete # noqa: PLW0603 + + if _startup_complete: + logger.debug("Startup already complete — skipping") + return {"status": "already_complete"} + + t0 = time.monotonic() + results: dict[str, str] = {} + + results["config"] = await _validate_config() + results["database"] = await _ensure_database() + _check_mcp_encryption_key() + results["cache"] = await _warm_caches() + results["scheduler"] = await _start_scheduler() + results["otel"] = _setup_otel() + results["telemetry"] = _setup_telemetry() + + _upgrade_signal_handlers() + + elapsed = round((time.monotonic() - t0) * 1000, 1) + _startup_complete = True + + logger.info( + "Startup complete in %.1fms: %s", + elapsed, + results, + ) + return results + + +def _upgrade_signal_handlers() -> None: + """Upgrade to loop-aware signal handlers for async drain.""" + try: + from deep_agent.aegra.shutdown import register_signal_handlers + + register_signal_handlers() + except Exception: + logger.warning("Failed to register signal handlers", exc_info=True) + + +async def _validate_config() -> str: + """Validate core settings.""" + try: + from deep_agent.src.settings import settings, validate_config + + validate_config(settings) + return "ok" + except Exception as exc: + logger.error("Config validation failed: %s", exc) + raise # Re-raise to fail startup + + +def _check_mcp_encryption_key() -> None: + """Warn if any MCP server uses oauth/dcr but MCP_TOKEN_ENCRYPTION_KEY is not set.""" + try: + from deep_agent.src.agent.config import agent_config + + servers = agent_config.get_mcp_servers() + needs_key = any( + s.get("auth_mode") in ("oauth", "dcr") + for s in servers.values() + if isinstance(s, dict) and s.get("enabled", False) + ) + if needs_key and not os.environ.get("MCP_TOKEN_ENCRYPTION_KEY"): + logger.error( + "MCP_TOKEN_ENCRYPTION_KEY is not set but one or more MCP servers " + "use auth_mode 'oauth' or 'dcr'. Token encryption will fail." + ) + except Exception: + logger.debug("MCP encryption key check skipped", exc_info=True) + + +async def _ensure_database() -> str: + """Create personalization, feedback, and token budget tables if they don't exist.""" + try: + from deep_agent.src.feedback.repository import FeedbackRepository + from deep_agent.src.personalization.repository import ( + PersonalizationRepository, + ) + from deep_agent.src.settings import settings + + setup_tasks = [] + + if settings.database_uri: + personalization_repo = PersonalizationRepository(settings.database_uri) + feedback_repo = FeedbackRepository(settings.database_uri) + setup_tasks.append(personalization_repo.ensure_tables()) + setup_tasks.append(feedback_repo.ensure_table()) + + from deep_agent.aegra.mcp_token_store import McpTokenStore + + mcp_token_store = McpTokenStore(settings.database_uri) + setup_tasks.append(mcp_token_store.ensure_tables()) + + if settings.MONGODB_URI: + from deep_agent.src.token_budget.mongo_repository import ( + TokenUsageMongoRepository, + ) + + mongo_repo = TokenUsageMongoRepository( + settings.MONGODB_URI, + db_name=settings.MONGODB_DB, + ) + setup_tasks.append(mongo_repo.ensure_indexes()) + + if not setup_tasks: + return "skipped: no database configured" + + await asyncio.gather(*setup_tasks) + return "ok" + except Exception as exc: + logger.error("Database setup failed: %s", exc) + return f"error: {exc}" + + +async def _warm_caches() -> str: + """Pre-populate caches if caching is enabled.""" + try: + from deep_agent.src.cache.config import cache_settings + + if not cache_settings.CACHE_ENABLED: + return "skipped: caching disabled" + + from deep_agent.src.cache.warming import warm_caches + + warm_caches() + return "ok" + except Exception as exc: + logger.warning("Cache warming failed: %s", exc) + return f"warning: {exc}" + + +async def _start_scheduler() -> str: + """Start background memory scheduler if enabled.""" + try: + from deep_agent.src.memory.config import memory_settings + + if not memory_settings.MEMORY_CONSOLIDATION_ENABLED: + return "skipped: memory consolidation disabled" + + from deep_agent.src.memory.scheduler import start_scheduler + from deep_agent.src.settings import settings + + started = await start_scheduler(settings.database_uri) + return "ok" if started else "skipped: already running" + except Exception as exc: + logger.warning("Scheduler start failed: %s", exc) + return f"warning: {exc}" + + +def _setup_otel() -> str: + """Initialize OpenTelemetry metrics and tracing.""" + try: + from deep_agent.aegra.otel import initialize_telemetry + + initialize_telemetry() + return "ok" + except Exception as exc: + logger.warning("OTEL setup failed: %s", exc) + return f"warning: {exc}" + + +def _setup_telemetry() -> str: + """Register Langfuse tracing and token budget tracking if configured.""" + try: + from deep_agent.aegra.telemetry import ( + setup_langfuse_tracing, + setup_token_budget_tracking, + ) + + setup_langfuse_tracing() + setup_token_budget_tracking() # Callback-based tracking + return "ok" + except Exception as exc: + logger.warning("Telemetry setup failed: %s", exc) + return f"warning: {exc}" + + +def is_ready() -> bool: + """Return True if startup has completed.""" + return _startup_complete diff --git a/deep_agent/aegra/state.py b/deep_agent/aegra/state.py new file mode 100644 index 00000000..09cb6e66 --- /dev/null +++ b/deep_agent/aegra/state.py @@ -0,0 +1,65 @@ +"""LangGraph state schema for aegra deployment. + +Defines the extended state schema used when the agent runs on LangGraph +Platform. The base state is managed by deepagents internally; this module +adds metadata fields for observability, error tracking, and streaming +coordination. + +The deepagents library defines its own internal state with `messages` and +agent-specific fields. This schema extends that with platform-level +concerns that don't belong in the agent itself. +""" + +from typing import Any, TypedDict + +from deep_agent.aegra import __version__ + + +class AegraMetadata(TypedDict, total=False): + """Platform-level metadata tracked alongside agent state.""" + + run_id: str + trace_id: str + thread_id: str + session_id: str + user_id: str + stream_tokens: bool + error_count: int + last_error: str | None + + +class HealthStatus(TypedDict): + """Health check response schema.""" + + status: str + version: str + agent_name: str + model: str + mcp_tools_loaded: int + subagents_loaded: int + backend_ready: bool + + +def make_health_status( + *, + agent_name: str, + model: str, + mcp_tools_count: int, + subagents_count: int, + backend_ready: bool, +) -> HealthStatus: + """Build a health status dict from agent configuration.""" + return HealthStatus( + status="healthy", + version=__version__, + agent_name=agent_name, + model=model, + mcp_tools_loaded=mcp_tools_count, + subagents_loaded=subagents_count, + backend_ready=backend_ready, + ) + + +def serialize_metadata(metadata: AegraMetadata) -> dict[str, Any]: + """Serialize metadata to JSON-safe dict, dropping None values.""" + return {k: v for k, v in metadata.items() if v is not None} diff --git a/deep_agent/aegra/telemetry.py b/deep_agent/aegra/telemetry.py new file mode 100644 index 00000000..cb00a2c5 --- /dev/null +++ b/deep_agent/aegra/telemetry.py @@ -0,0 +1,264 @@ +"""Langfuse and token budget observability for aegra deployment. + +Provides: +- Langfuse callback handler factory for LangChain tracing (v4 SDK) +- Langfuse client accessor via ``get_langfuse_client()`` +- Token budget LangChain callback registration and metadata provider + +Environment variables (Langfuse — auto-read by v4 SDK): + LANGFUSE_PUBLIC_KEY: Langfuse public key + LANGFUSE_SECRET_KEY: Langfuse secret key + LANGFUSE_BASE_URL: Langfuse server URL + LANGFUSE_TRACING_ENVIRONMENT: Environment tag (e.g. development, production) +""" + +import contextvars +import os +from typing import Any + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +# --------------------------------------------------------------------------- +# Langfuse v4 integration +# --------------------------------------------------------------------------- + +_langfuse_tracing_initialized = False +_token_budget_tracing_initialized = False + + +def _get_trace_name() -> str: + """Resolve trace name: agent.yaml name > env var > fallback.""" + try: + from deep_agent.src.agent.config import agent_config + + return agent_config.get_name() + except Exception: + return os.environ.get("LANGFUSE_TRACE_NAME", "template-agent") + + +def _langfuse_configured() -> bool: + """Return True if the minimum Langfuse credentials are present.""" + return bool( + os.environ.get("LANGFUSE_PUBLIC_KEY") and os.environ.get("LANGFUSE_SECRET_KEY") + ) + + +def setup_langfuse_tracing() -> None: + """Register Langfuse as a global LangChain callback and Aegra observability provider. + + Two mechanisms work together: + + 1. ``register_configure_hook`` — the same mechanism LangSmith uses to + auto-inject its tracer. Creates a fresh ``CallbackHandler()`` per run. + 2. ``LangfuseObservabilityProvider`` — plugs into Aegra's + ``ObservabilityManager`` so that ``create_run_config`` injects + ``langfuse_user_id``, ``langfuse_session_id``, and + ``langfuse_trace_name`` into ``RunnableConfig.metadata``. + The CallbackHandler reads these automatically. + + Must be called **once** at process startup. Subsequent calls are no-ops. + """ + global _langfuse_tracing_initialized + if _langfuse_tracing_initialized: + return + _langfuse_tracing_initialized = True + + if not _langfuse_configured(): + logger.info("Langfuse credentials not set — auto-tracing disabled") + return + + try: + from langchain_core.tracers.context import register_configure_hook + from langfuse.langchain import CallbackHandler + + _langfuse_ctx_var: contextvars.ContextVar = contextvars.ContextVar( + "langfuse_handler", default=None + ) + + register_configure_hook( + _langfuse_ctx_var, + True, + CallbackHandler, + env_var="LANGFUSE_PUBLIC_KEY", + ) + logger.info("Langfuse auto-tracing registered for all LangChain runs") + except ImportError: + logger.warning( + "langfuse or langchain_core not available — auto-tracing disabled" + ) + return + except Exception: + logger.warning("Failed to register Langfuse tracing hook", exc_info=True) + return + + try: + from aegra_api.observability.base import get_observability_manager + + manager = get_observability_manager() + manager.register_provider(LangfuseObservabilityProvider()) + logger.info("Langfuse observability provider registered with Aegra") + except ImportError: + logger.debug("aegra_api not available — skipping provider registration") + except Exception: + logger.warning( + "Failed to register Langfuse observability provider", exc_info=True + ) + + +class LangfuseObservabilityProvider: + """Aegra ObservabilityProvider that injects Langfuse metadata into RunnableConfig. + + The Langfuse v4 ``CallbackHandler`` auto-reads these keys from + ``RunnableConfig.metadata``: + + - ``langfuse_user_id`` — who triggered the run + - ``langfuse_session_id`` — groups traces by conversation (thread) + - ``langfuse_trace_name`` — human-readable trace name in the UI + """ + + def get_callbacks(self) -> list[Any]: + """Return empty list — callbacks are handled by register_configure_hook.""" + return [] + + def get_metadata( + self, run_id: str, thread_id: str, user_identity: str | None = None + ) -> dict[str, Any]: + """Return Langfuse metadata keys for RunnableConfig injection.""" + from deep_agent.utils.pylogger import _trace_id_var + + metadata: dict[str, Any] = { + "langfuse_trace_name": _get_trace_name(), + } + if user_identity: + metadata["langfuse_user_id"] = user_identity + if thread_id: + metadata["langfuse_session_id"] = thread_id + trace_id = _trace_id_var.get() + if trace_id: + metadata["langfuse_tags"] = [f"trace_id:{trace_id}"] + return metadata + + def is_enabled(self) -> bool: + """Return True if Langfuse credentials are configured.""" + return _langfuse_configured() + + +# --------------------------------------------------------------------------- +# Token budget callback integration +# --------------------------------------------------------------------------- + + +class TokenBudgetObservabilityProvider: + """Inject thread_id and trace_id into RunnableConfig metadata for the token budget callback.""" + + def get_callbacks(self) -> list[Any]: + """Return empty list — callbacks are handled by register_configure_hook.""" + return [] + + def get_metadata( + self, run_id: str, thread_id: str, user_identity: str | None = None + ) -> dict[str, Any]: + """Return token budget metadata keys for RunnableConfig injection.""" + from deep_agent.src.token_budget.callback import ( + THREAD_ID_METADATA_KEY, + TRACE_ID_METADATA_KEY, + USER_ID_METADATA_KEY, + ) + from deep_agent.utils.pylogger import _trace_id_var + + metadata: dict[str, Any] = {} + if thread_id: + metadata[THREAD_ID_METADATA_KEY] = thread_id + if user_identity: + metadata[USER_ID_METADATA_KEY] = user_identity + trace_id = _trace_id_var.get() + if trace_id: + metadata[TRACE_ID_METADATA_KEY] = trace_id + return metadata + + def is_enabled(self) -> bool: + """Return True if token budget tracking is active.""" + try: + from deep_agent.src.agent.config import agent_config + + return agent_config.get_token_budget_config().is_active + except Exception: + return False + + +def setup_token_budget_tracking() -> None: + """Register token budget LangChain callback and Aegra metadata provider.""" + global _token_budget_tracing_initialized + if _token_budget_tracing_initialized: + return + _token_budget_tracing_initialized = True + + try: + from deep_agent.src.agent.config import agent_config + + if not agent_config.get_token_budget_config().is_active: + logger.info("Token budget disabled — callback registration skipped") + return + except Exception: + logger.debug("Token budget config unavailable — skipping callback registration") + return + + try: + from langchain_core.tracers.context import register_configure_hook + + from deep_agent.src.token_budget.callback import TokenBudgetCallbackHandler + + _token_budget_ctx_var: contextvars.ContextVar = contextvars.ContextVar( + "token_budget_handler", default=None + ) + os.environ.setdefault("TOKEN_BUDGET_TRACKING", "1") + register_configure_hook( + _token_budget_ctx_var, + True, + TokenBudgetCallbackHandler, + env_var="TOKEN_BUDGET_TRACKING", + ) + logger.info("Token budget callback registered for all LangChain runs") + except ImportError: + logger.warning("langchain_core not available — token budget callback disabled") + return + except Exception: + logger.warning("Failed to register token budget callback", exc_info=True) + return + + try: + from aegra_api.observability.base import get_observability_manager + + manager = get_observability_manager() + manager.register_provider(TokenBudgetObservabilityProvider()) + logger.info("Token budget observability provider registered with Aegra") + except ImportError: + logger.debug("aegra_api not available — skipping token budget provider") + except Exception: + logger.warning( + "Failed to register token budget observability provider", exc_info=True + ) + + +def get_langfuse_client() -> Any: + """Return the Langfuse singleton client (v4), or None if unconfigured. + + Uses ``get_client()`` which auto-reads ``LANGFUSE_PUBLIC_KEY``, + ``LANGFUSE_SECRET_KEY``, and ``LANGFUSE_BASE_URL`` from the environment. + """ + if not _langfuse_configured(): + return None + + try: + from langfuse import get_client + + return get_client() + except ImportError: + logger.warning("langfuse package not installed — Langfuse tracing disabled") + return None + except Exception: + logger.warning("Failed to initialize Langfuse client", exc_info=True) + return None diff --git a/deep_agent/headless.py b/deep_agent/headless.py new file mode 100644 index 00000000..96669458 --- /dev/null +++ b/deep_agent/headless.py @@ -0,0 +1,196 @@ +"""Headless agent entry point — runs as a background worker with event triggers. + +Usage: + python -m deep_agent.headless +""" + +from __future__ import annotations + +import asyncio +import signal +import sys +from pathlib import Path +from typing import Any + +import yaml + +from deep_agent.src.triggers.config import AgentMode, HeadlessConfig +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_CONFIG_PATH = ( + Path(__file__).resolve().parent.parent + / "config" + / "agent" + / "runtime" + / "agent.yaml" +) +_HEADLESS_PROMPT_PATH = ( + Path(__file__).resolve().parent.parent / "config" / "agent" / "HEADLESS_PROMPT.md" +) + + +def _load_headless_config() -> HeadlessConfig: + """Load and validate headless configuration from agent.yaml.""" + if not _CONFIG_PATH.is_file(): + logger.error("Config not found: %s", _CONFIG_PATH) + sys.exit(1) + + raw = yaml.safe_load(_CONFIG_PATH.read_text()) or {} + + config = HeadlessConfig( + mode=AgentMode.HEADLESS, + triggers=raw.get("triggers", {}), + output_sinks=raw.get("output_sinks", []), + drain_timeout=raw.get("drain_timeout", 30.0), + health_check=raw.get("health_check", {}), + ) + + return config + + +async def _build_headless_graph() -> Any: + """Build a dedicated graph from HEADLESS_PROMPT.md. + + Uses a simpler prompt than the orchestrator — no user interaction, + no TODO lists, just task processing. + """ + from deep_agent.src.agent.config import agent_config + from deep_agent.src.agent.config.model import parse_model_config + from deep_agent.src.agent.config.parser import ( + inject_runtime_values, + parse_frontmatter, + ) + from deep_agent.src.cache.model_cache import get_or_create_model_from_spec + + headless_cfg = parse_frontmatter(_HEADLESS_PROMPT_PATH) + model_raw = headless_cfg.get("model", "gemini-2.5-pro") + system_prompt = inject_runtime_values(headless_cfg.get("body", "")) + tool_names = headless_cfg.get("tools", []) + skill_names = headless_cfg.get("skills", []) + if skill_names: + from deep_agent.src.agent.config.resolver import resolve_skill_paths + + available_skills = agent_config._scan_available_skills() + skill_paths = resolve_skill_paths( + skill_names, available_skills, agent_name="headless-worker" + ) + else: + skill_paths = [] + + orch_spec = parse_model_config(model_raw) + model = get_or_create_model_from_spec(orch_spec) + + from deep_agent.aegra.mcp import get_mcp_tools + + mcp_tools = await get_mcp_tools(sso_token=None, server_names=None) + + from deep_agent.src.triggers.tools import get_builtin_tools + + all_tools = list(mcp_tools) + get_builtin_tools() + tools = agent_config.resolve_tools( + tool_names, all_tools, agent_name="headless-worker" + ) + + from deep_agent.src.infrastructure.backend import get_configured_backend + from deep_agent.src.infrastructure.middleware import ( + build_middleware_list, + resolve_memory_param, + ) + + middleware_overrides = headless_cfg.get("middleware") + resolved_mw = agent_config.resolve_agent_middleware( + orch_spec.name, middleware_overrides + ) + backend = get_configured_backend() + middleware = build_middleware_list(resolved_mw, model=model, backend=backend) + memory = resolve_memory_param(resolved_mw) + + from deepagents import create_deep_agent + + compiled = create_deep_agent( + name="headless-worker", + model=model, + system_prompt=system_prompt, + skills=skill_paths or None, + tools=tools, + backend=backend, + middleware=middleware, + memory=memory, + ) + + logger.info( + "Headless graph built: %d tool(s), prompt=%s", + len(tools), + _HEADLESS_PROMPT_PATH.name, + ) + return compiled + + +async def main() -> None: + """Run the headless agent worker.""" + logger.info("Starting headless agent worker") + + config = _load_headless_config() + + from deep_agent.aegra.startup import run_startup + + startup_results = await run_startup() + logger.info("Startup: %s", startup_results) + + logger.info("Building headless agent graph") + compiled_graph = await _build_headless_graph() + + from deep_agent.src.settings import settings + from deep_agent.src.triggers.middleware import EventTriggerMiddleware + + middleware = EventTriggerMiddleware( + config=config, + graph=compiled_graph, + redis_url=settings.REDIS_URL, + ) + + await middleware.start() + + health_server: asyncio.Server | None = None + if config.health_check.enabled: + from deep_agent.src.triggers.health import start_health_server + + health_server = await start_health_server( + config.health_check.host, + config.health_check.port, + middleware, + ) + + stop_event = asyncio.Event() + loop = asyncio.get_running_loop() + + def _signal_handler() -> None: + logger.info("Shutdown signal received") + stop_event.set() + + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, _signal_handler) + + logger.info("Headless agent worker running — waiting for events") + await stop_event.wait() + + logger.info("Shutting down headless agent worker") + if health_server is not None: + health_server.close() + await health_server.wait_closed() + await middleware.stop() + + try: + from deep_agent.aegra.shutdown import run_shutdown + + await run_shutdown() + except Exception: + logger.debug("Cleanup completed with warnings", exc_info=True) + + logger.info("Headless agent worker stopped") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/template_agent/src/__init__.py b/deep_agent/src/__init__.py similarity index 100% rename from template_agent/src/__init__.py rename to deep_agent/src/__init__.py diff --git a/deep_agent/src/adapters/__init__.py b/deep_agent/src/adapters/__init__.py new file mode 100644 index 00000000..12ea8df9 --- /dev/null +++ b/deep_agent/src/adapters/__init__.py @@ -0,0 +1,13 @@ +"""Adapters for external library formats. + +This package contains adapters that convert between external library formats +and our internal schema. Each adapter module is named after the library it +adapts (e.g., langchain.py for LangChain). +""" + +from .langchain import convert_message_content_to_string, langchain_to_chat_message + +__all__ = [ + "langchain_to_chat_message", + "convert_message_content_to_string", +] diff --git a/template_agent/src/core/agent_utils.py b/deep_agent/src/adapters/langchain.py similarity index 57% rename from template_agent/src/core/agent_utils.py rename to deep_agent/src/adapters/langchain.py index aa12d88b..fab9ab22 100644 --- a/template_agent/src/core/agent_utils.py +++ b/deep_agent/src/adapters/langchain.py @@ -1,7 +1,12 @@ -"""Utility functions for handling agent messages and conversions. +"""LangChain message adapter. -This module provides utility functions for converting between different message -formats, handling message content, and managing tool calls in the template agent. +This module adapts LangChain's message format to our internal ChatMessage schema. +It serves as the boundary layer between the external LangChain library and our +internal data structures defined in schema.py. + +Functions: + langchain_to_chat_message: Convert LangChain BaseMessage to ChatMessage + convert_message_content_to_string: Normalize message content to string format """ from typing import Any, Dict, List, Union @@ -12,9 +17,8 @@ HumanMessage, ToolMessage, ) -from langchain_core.messages import ChatMessage as LangchainChatMessage -from template_agent.src.schema import ChatMessage, ToolCall +from deep_agent.src.schema import ChatMessage, ToolCall def convert_message_content_to_string( @@ -52,7 +56,8 @@ def langchain_to_chat_message(message: BaseMessage) -> ChatMessage: This function converts LangChain message objects to the internal ChatMessage format used by the template agent. It handles different message types and - preserves relevant metadata. + preserves relevant metadata including run_id, trace_id, and session_id from + message metadata. Args: message: The LangChain message to convert. Must be one of the supported @@ -64,11 +69,20 @@ def langchain_to_chat_message(message: BaseMessage) -> ChatMessage: Raises: ValueError: If the message type is not supported or has an invalid role. """ + # Extract common metadata fields from message.metadata + metadata = getattr(message, "metadata", None) or {} + run_id = metadata.get("run_id") + trace_id = metadata.get("trace_id") + session_id = metadata.get("session_id") + match message: case HumanMessage(): human_message = ChatMessage( type="human", content=convert_message_content_to_string(message.content), + run_id=run_id, + trace_id=trace_id, + session_id=session_id, ) return human_message @@ -76,20 +90,18 @@ def langchain_to_chat_message(message: BaseMessage) -> ChatMessage: ai_message = ChatMessage( type="ai", content=convert_message_content_to_string(message.content), + run_id=run_id, + trace_id=trace_id, + session_id=session_id, ) - # Handle tool calls from both direct attribute and additional_kwargs - tool_calls = message.tool_calls or [] - if message.additional_kwargs and "tool_calls" in message.additional_kwargs: - tool_calls.extend(message.additional_kwargs["tool_calls"]) - if tool_calls: - # Ensure tool calls have the correct structure + # Handle tool calls from modern LangChain messages + if message.tool_calls: formatted_tool_calls = [] - for tool_call in tool_calls: + for tool_call in message.tool_calls: if isinstance(tool_call, dict): # Ensure required fields are present and properly typed if "name" in tool_call and "args" in tool_call: - # Create a proper ToolCall object formatted_call: ToolCall = { "name": str(tool_call["name"]), "args": dict(tool_call["args"]), @@ -101,14 +113,10 @@ def langchain_to_chat_message(message: BaseMessage) -> ChatMessage: formatted_tool_calls.append(formatted_call) ai_message.tool_calls = formatted_tool_calls + # Copy response metadata if message.response_metadata: ai_message.response_metadata = message.response_metadata - if message.additional_kwargs: - if "response_metadata" in message.additional_kwargs: - ai_message.response_metadata.update( - message.additional_kwargs["response_metadata"] - ) - ai_message.ai_call_id = message.additional_kwargs.get("ai_call_id") + return ai_message case ToolMessage(): @@ -116,46 +124,11 @@ def langchain_to_chat_message(message: BaseMessage) -> ChatMessage: type="tool", content=convert_message_content_to_string(message.content), tool_call_id=message.tool_call_id, + run_id=run_id, + trace_id=trace_id, + session_id=session_id, ) return tool_message - case LangchainChatMessage(): - if message.role == "custom": - custom_message = ChatMessage( - type="custom", - content="", - custom_data=message.content[0], - ) - return custom_message - else: - raise ValueError(f"Unsupported chat message role: {message.role}") - case _: raise ValueError(f"Unsupported message type: {message.__class__.__name__}") - - -def remove_tool_calls( - content: Union[str, List[Union[str, Dict[str, Any]]]], -) -> Union[str, List[Union[str, Dict[str, Any]]]]: - """Remove tool calls from message content. - - This function filters out tool call content from message content, particularly - useful for handling streaming responses from models that include tool calls - in their content stream. - - Args: - content: The content to process. Can be a string or a list containing - strings and dictionaries with content information. - - Returns: - The content with tool calls removed. Returns the same type as input. - """ - if isinstance(content, str): - return content - - # Currently only Anthropic models stream tool calls, using content item type tool_use - return [ - content_item - for content_item in content - if isinstance(content_item, str) or content_item["type"] != "tool_use" - ] diff --git a/deep_agent/src/agent/__init__.py b/deep_agent/src/agent/__init__.py new file mode 100644 index 00000000..499bf7d0 --- /dev/null +++ b/deep_agent/src/agent/__init__.py @@ -0,0 +1,8 @@ +"""Agent configuration and orchestration. + +This package provides functionality for configuring deep agents. +The live graph factory is in ``deep_agent.aegra.graph``. + +Modules: + config: Configuration loading and management +""" diff --git a/deep_agent/src/agent/config/__init__.py b/deep_agent/src/agent/config/__init__.py new file mode 100644 index 00000000..fc47424f --- /dev/null +++ b/deep_agent/src/agent/config/__init__.py @@ -0,0 +1,19 @@ +"""Agent configuration management. + +This package handles loading and processing agent configurations from the +config/ directory at the repository root. It provides a singleton AgentConfig +class that loads orchestrator, subagent, skill, and MCP configurations. + +Modules: + loader: Main AgentConfig singleton class + parser: Frontmatter parsing and runtime value injection + resolver: Skill and tool name resolution + +Main exports: + AgentConfig: Singleton configuration manager + agent_config: Pre-initialized singleton instance +""" + +from .loader import AgentConfig, agent_config + +__all__ = ["AgentConfig", "agent_config"] diff --git a/deep_agent/src/agent/config/cache.py b/deep_agent/src/agent/config/cache.py new file mode 100644 index 00000000..89a683a9 --- /dev/null +++ b/deep_agent/src/agent/config/cache.py @@ -0,0 +1,62 @@ +"""Cache configuration models. + +Provides validated Pydantic models for the ``cache:`` section of +config/agent/runtime/agent.yaml. Controls TTLs, feature flags, and +size limits for all cache layers (model, personalization, MCP tools, +compiled graph, Redis L2, warming, and metrics). + +The template-agent user only touches YAML. This module converts +declarative config into parameters consumed by cache infrastructure. +""" + +from __future__ import annotations + +from pydantic import BaseModel, Field + + +class ModelCacheConfig(BaseModel): + """LLM model instance cache settings.""" + + enabled: bool = True + ttl: int = Field(default=600, ge=10, le=7200) + max_size: int = Field(default=50, ge=1, le=100) + + +class PersonalizationCacheConfig(BaseModel): + """User personalization (memories/rules) cache settings.""" + + enabled: bool = True + ttl: int = Field(default=120, ge=10, le=3600) + + +class McpCacheConfig(BaseModel): + """MCP tool discovery cache settings.""" + + ttl: int = Field(default=300, ge=10, le=3600) + + +class GraphCacheConfig(BaseModel): + """Compiled graph cache settings.""" + + ttl: int = Field(default=300, ge=10, le=3600) + + +class ToggleConfig(BaseModel): + """Generic feature toggle with enabled flag.""" + + enabled: bool = True + + +class CacheFileConfig(BaseModel): + """Top-level cache configuration from agent.yaml ``cache:`` section.""" + + enabled: bool = True + model: ModelCacheConfig = Field(default_factory=ModelCacheConfig) + personalization: PersonalizationCacheConfig = Field( + default_factory=PersonalizationCacheConfig, + ) + mcp: McpCacheConfig = Field(default_factory=McpCacheConfig) + graph: GraphCacheConfig = Field(default_factory=GraphCacheConfig) + redis: ToggleConfig = Field(default_factory=ToggleConfig) + warming: ToggleConfig = Field(default_factory=ToggleConfig) + metrics: ToggleConfig = Field(default_factory=ToggleConfig) diff --git a/deep_agent/src/agent/config/filesystem.py b/deep_agent/src/agent/config/filesystem.py new file mode 100644 index 00000000..13ea5ea0 --- /dev/null +++ b/deep_agent/src/agent/config/filesystem.py @@ -0,0 +1,105 @@ +"""Filesystem configuration models. + +Provides validated Pydantic models for the ``filesystem:`` section of +config/agent/runtime/agent.yaml: +- Backend type selection (local_shell / state / composite) +- Filesystem permissions (operations + paths + mode) +- FilesystemMiddleware tuning (eviction thresholds, timeouts) + +The template-agent user only touches YAML. This module converts +declarative config into parameters for the backend and create_deep_agent(). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Literal + +import yaml +from pydantic import BaseModel, Field + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +class LocalShellConfig(BaseModel): + """Configuration for LocalShellBackend.""" + + timeout: int = 120 + max_output_bytes: int = 100_000 + + +class StateConfig(BaseModel): + """Configuration for StateBackend (ephemeral in-memory).""" + + enabled: bool = False + + +class StoreConfig(BaseModel): + """Configuration for StoreBackend (cross-thread persistent).""" + + enabled: bool = False + scope: Literal["user", "assistant", "org"] = "user" + + +class BackendConfig(BaseModel): + """Backend selection and configuration.""" + + type: Literal["state", "composite", "store", "local_shell"] = "state" + local_shell: LocalShellConfig = Field(default_factory=LocalShellConfig) + state: StateConfig = Field(default_factory=StateConfig) + store: StoreConfig = Field(default_factory=StoreConfig) + routes: dict[str, str] = Field(default_factory=dict) + + +class PermissionRule(BaseModel): + """A single filesystem permission rule.""" + + operations: list[str] + paths: list[str] + mode: Literal["allow", "deny"] = "allow" + + +class FilesystemSettings(BaseModel): + """FilesystemMiddleware tuning parameters.""" + + tool_token_limit_before_evict: int = 20_000 + human_message_token_limit_before_evict: int = 50_000 + max_execute_timeout: int = 3600 + + +class FilesystemFileConfig(BaseModel): + """Structure of the ``filesystem:`` section in runtime/agent.yaml.""" + + backend: BackendConfig = Field(default_factory=BackendConfig) + permissions: list[PermissionRule] = Field(default_factory=list) + permission_inheritance: bool = False + settings: FilesystemSettings = Field(default_factory=FilesystemSettings) + + +def load_filesystem_config(config_path: Path) -> FilesystemFileConfig: + """Load and validate filesystem.yaml from disk. + + Args: + config_path: Path to filesystem.yaml. + + Returns: + Validated FilesystemFileConfig. Returns defaults if file is missing. + """ + if not config_path.is_file(): + logger.info("No filesystem.yaml found — using defaults (local_shell)") + return FilesystemFileConfig() + + try: + raw = yaml.safe_load(config_path.read_text()) or {} + config: FilesystemFileConfig = FilesystemFileConfig.model_validate(raw) + logger.info( + "Loaded filesystem config: backend=%s, %d permission rule(s)", + config.backend.type, + len(config.permissions), + ) + return config + except Exception as e: + logger.warning("Failed to parse filesystem.yaml, using defaults: %s", e) + return FilesystemFileConfig() diff --git a/deep_agent/src/agent/config/hitl.py b/deep_agent/src/agent/config/hitl.py new file mode 100644 index 00000000..edcae428 --- /dev/null +++ b/deep_agent/src/agent/config/hitl.py @@ -0,0 +1,108 @@ +"""Human-in-the-loop interrupt configuration builder. + +Converts the ``human_approval`` section of ``agent.yaml`` into the +``interrupt_on`` dict expected by ``create_deep_agent()``. + +The dict maps each tool name to ``True`` (use deepagents default +decisions: approve / edit / reject / respond). When the feature is +disabled the function returns an empty dict, which signals to +``graph.py`` not to pass ``interrupt_on`` at all. + +For ``mode: all``, both the caller-supplied tools (MCP / explicit) and +the deepagents built-in tools are included so that every tool call — +regardless of origin — pauses for human approval. + +Example YAML config:: + + middleware: + human_approval: + enabled: true + mode: all + exclude: + - ls + - read_file + - glob + - grep +""" + +from __future__ import annotations + +from typing import Any + +from deep_agent.src.agent.config.middleware import HumanApprovalConfig +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +# Built-in tool names added by deepagents internally (FilesystemMiddleware, +# TodoListMiddleware, SubAgentMiddleware). These are never present in the +# caller-supplied ``tools`` list, so they must be enumerated explicitly for +# ``interrupt_on`` to cover them. +_DEEPAGENTS_BUILTIN_TOOLS: frozenset[str] = frozenset( + { + # filesystem (FilesystemMiddleware) + "ls", + "read_file", + "write_file", + "edit_file", + "glob", + "grep", + "execute", + # todo list (TodoListMiddleware) + "write_todos", + # subagents (SubAgentMiddleware) + "task", + # conversation management + "compact_conversation", + } +) + + +def build_interrupt_on( + config: HumanApprovalConfig, + tools: list[Any], +) -> dict[str, Any]: + """Build the ``interrupt_on`` dict for ``create_deep_agent()``. + + Args: + config: Resolved ``human_approval`` config from ``agent.yaml``. + tools: List of resolved tool objects (must have a ``.name`` attr). + Typically MCP tools + any explicitly declared tools. Built-in + deepagents tools are added automatically when ``mode`` is ``"all"``. + + Returns: + Dict mapping tool name → ``True`` for every tool that should + trigger a human approval interrupt. Returns ``{}`` when the + feature is disabled or ``mode`` is ``"none"``. + """ + if not config.enabled or config.mode == "none": + logger.debug("HITL disabled (enabled=%s, mode=%s)", config.enabled, config.mode) + return {} + + exclude = set(config.exclude) + + # Explicit / MCP tools passed by the caller + explicit_names = {t.name for t in tools} + + # For mode=all, also cover the deepagents built-in tools so that + # filesystem and todo calls are intercepted even when no MCP tools exist. + all_names = explicit_names | _DEEPAGENTS_BUILTIN_TOOLS + + interrupt_on = {name: True for name in all_names if name not in exclude} + + if interrupt_on: + excluded = (explicit_names | _DEEPAGENTS_BUILTIN_TOOLS) - set(interrupt_on) + logger.info( + "HITL enabled: %d tool(s) will require approval%s", + len(interrupt_on), + f" ({len(excluded)} excluded: {sorted(excluded)})" if excluded else "", + ) + else: + logger.debug( + "HITL enabled but all tools excluded (explicit=%d, builtins=%d, exclude=%s)", + len(explicit_names), + len(_DEEPAGENTS_BUILTIN_TOOLS), + exclude, + ) + + return interrupt_on diff --git a/deep_agent/src/agent/config/loader.py b/deep_agent/src/agent/config/loader.py new file mode 100644 index 00000000..7defa4a8 --- /dev/null +++ b/deep_agent/src/agent/config/loader.py @@ -0,0 +1,605 @@ +"""Agent configuration loader and singleton. + +This module provides the main AgentConfig singleton class that orchestrates loading +agent configurations from the config/agent/ directory at the repository root. It +loads the unified runtime/agent.yaml once, then extracts sections for providers, +middleware, and filesystem config. Orchestrator, subagents, skills, and MCP server +configurations are loaded eagerly at initialization time. + +Why this exists: + All agent configurations need to be loaded once and made available throughout + the application. This singleton ensures configs are loaded only once and + provides convenient access methods. + +Classes: + AgentConfig: Singleton for managing all agent configuration loading +""" + +import json +import os +from pathlib import Path +from typing import Any, cast + +import yaml + +from deep_agent.src.exceptions import AppException, ErrorCodes +from deep_agent.src.settings import settings +from deep_agent.src.token_budget.config import TokenBudgetConfig +from deep_agent.utils.pylogger import get_python_logger + +from .cache import CacheFileConfig +from .filesystem import FilesystemFileConfig +from .middleware import ( + MiddlewareFileConfig, + ResolvedMiddlewareConfig, + resolve_middleware, +) +from .otel import OtelFileConfig +from .parser import inject_runtime_values, parse_frontmatter +from .providers import ProvidersFileConfig +from .resolver import resolve_skill_paths, resolve_tools + +logger = get_python_logger(log_level=settings.PYTHON_LOG_LEVEL) + + +def _strip_jsonc_comments(text: str) -> str: + """Remove ``//`` line comments outside JSON string literals.""" + result: list[str] = [] + i = 0 + n = len(text) + + while i < n: + ch = text[i] + if ch == '"': + result.append(ch) + i += 1 + while i < n: + c = text[i] + result.append(c) + if c == "\\": + i += 1 + if i < n: + result.append(text[i]) + elif c == '"': + break + i += 1 + i += 1 + continue + + if ch == "/" and i + 1 < n and text[i + 1] == "/": + while i < n and text[i] not in "\n\r": + i += 1 + continue + + result.append(ch) + i += 1 + + return "".join(result) + + +def _load_jsonc(path: Path) -> dict[str, Any]: + """Load JSON with optional ``//`` line comments.""" + raw = path.read_text() + return cast(dict[str, Any], json.loads(_strip_jsonc_comments(raw))) + + +# Config directory path - read from CONFIG_PATH env var for base image pattern +# Falls back to repo-root config/agent/ for backward compatibility +_AGENT_CONFIG_DIR = Path( + os.getenv( + "CONFIG_PATH", + str(Path(__file__).parent.parent.parent.parent.parent / "config" / "agent"), + ) +) + + +class AgentConfig: + """Singleton class for managing agent configuration operations. + + This class provides centralized access to all config/ directory + operations including loading configurations, resolving paths, and + managing runtime values. + """ + + _instance: "AgentConfig | None" = None + _initialized: bool + _configs_loaded: bool + _base_dir: Path + _orchestrator: dict[str, Any] + _subagents: dict[str, dict[str, Any]] + _mcp_servers: dict[str, Any] + _available_skills: dict[str, Path] + _middleware_config: MiddlewareFileConfig + _filesystem_config: FilesystemFileConfig + _providers_config: ProvidersFileConfig + _cache_config: CacheFileConfig + _otel_config: OtelFileConfig + _token_budget_config: TokenBudgetConfig + _name: str + + def __new__(cls, base_dir: Path | None = None) -> "AgentConfig": + """Create or return the singleton instance. + + Args: + base_dir: Optional base directory for config. Only used on first instantiation. + + Returns: + The singleton AgentConfig instance. + """ + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._initialized = False + return cls._instance + + def __init__(self, base_dir: Path | None = None): + """Initialize the AgentConfig singleton. + + Args: + base_dir: Optional base directory for config. Defaults to + config/ at the repository root relative to this module. + """ + if self._initialized: + return + + self._base_dir = base_dir if base_dir is not None else _AGENT_CONFIG_DIR + self._initialized = True + self._configs_loaded = False + + def _load_agent_yaml(self) -> dict[str, Any]: + """Load the unified runtime/agent.yaml once. + + Returns: + Raw dict from agent.yaml, or empty dict if missing. + """ + agent_yaml = self._base_dir / "runtime" / "agent.yaml" + if not agent_yaml.is_file(): + logger.warning("No runtime/agent.yaml found — using defaults") + return {} + + try: + raw = yaml.safe_load(agent_yaml.read_text()) or {} + logger.info("Loaded runtime/agent.yaml") + return raw + except Exception as e: + logger.warning("Failed to parse runtime/agent.yaml, using defaults: %s", e) + return {} + + def _load_otel_config(self) -> OtelFileConfig: + """Load OpenTelemetry configuration from observability.yaml. + + Returns: + OtelFileConfig with OTEL settings, or defaults if missing. + """ + otel_yaml = self._base_dir / "runtime" / "observability.yaml" + if not otel_yaml.is_file(): + logger.info("No observability.yaml found — OTEL disabled by default") + return OtelFileConfig() + + try: + raw = yaml.safe_load(otel_yaml.read_text()) or {} + config: OtelFileConfig = OtelFileConfig.model_validate(raw.get("otel", {})) + logger.info("Loaded OTEL config from observability.yaml") + return config + except Exception as e: + logger.warning("Failed to parse observability.yaml, using defaults: %s", e) + return OtelFileConfig() + + def _ensure_loaded(self) -> None: + """Lazy load configurations on first access. + + This ensures logging is properly configured before we try to log. + """ + # If auto-reload is enabled, always reload from disk + if settings.CONFIG_AUTO_RELOAD: + if self._configs_loaded: + logger.debug("CONFIG_AUTO_RELOAD=true: reloading configs from disk") + self._configs_loaded = False + + if self._configs_loaded: + return + + logger.info("Loading agent configurations...") + + raw = self._load_agent_yaml() + + # Extract middleware section (defaults + harness_profiles as profiles) + self._middleware_config = MiddlewareFileConfig.model_validate( + { + "defaults": raw.get("middleware", {}), + "profiles": raw.get("harness_profiles", {}), + } + ) + + # Extract filesystem section + self._filesystem_config = FilesystemFileConfig.model_validate( + raw.get("filesystem", {}) + ) + + # Extract providers section (shares harness_profiles with middleware) + self._providers_config = ProvidersFileConfig.model_validate( + { + "resolve_strategy": raw.get("resolve_strategy", "legacy"), + "providers": raw.get("providers", {}), + "harness_profiles": raw.get("harness_profiles", {}), + "async_tasks": raw.get("async_tasks", {}), + } + ) + + # Extract cache section + self._cache_config = CacheFileConfig.model_validate(raw.get("cache", {})) + + # Load OTEL config from observability.yaml + self._otel_config = self._load_otel_config() + + # Extract token budget section + self._token_budget_config = TokenBudgetConfig.model_validate( + raw.get("token_budget", {}) + ) + + # Extract top-level identity + self._name = raw.get("name", "Agent") + # Scan skills first, as orchestrator and subagents need them for resolution + self._available_skills: dict[str, Path] = self._scan_available_skills() + self._orchestrator: dict[str, Any] = self._load_orchestrator() + self._subagents: dict[str, dict[str, Any]] = self._load_all_subagents() + self._mcp_servers: dict[str, Any] = self._load_mcp_servers() + + self._configs_loaded = True + logger.info( + f"Agent config loaded: orchestrator={self._orchestrator.get('name')}, " + f"subagents={len(self._subagents)}, skills={len(self._available_skills)}" + ) + + @property + def base_dir(self) -> Path: + """Get the config base directory path.""" + return self._base_dir + + @staticmethod + def _validate_mcps_field(mcps: Any, agent_name: str) -> None: + """Validate the ``mcps`` frontmatter field is a list of strings. + + Args: + mcps: The raw value from frontmatter. + agent_name: Agent name for error messages. + + Raises: + AppException: If ``mcps`` is not a list of strings. + """ + if not isinstance(mcps, list) or not all(isinstance(s, str) for s in mcps): + raise AppException( + f"Agent '{agent_name}': 'mcps' must be a list of strings", + ErrorCodes.CONFIGURATION_VALIDATION_ERROR, + ) + + def _load_orchestrator(self) -> dict[str, Any]: + """Load orchestrator configuration at initialization. + + Returns: + Orchestrator config dict with injected runtime values and resolved skill paths. + + Raises: + AppException: If orchestrator/main.md is missing or invalid. + """ + orchestrator_path = self._base_dir / "PROMPT.md" + try: + config = parse_frontmatter(orchestrator_path) + if "body" in config: + config["body"] = inject_runtime_values(config["body"]) + + from deep_agent.src.infrastructure.tool_access import migrate_tools_field + + migrate_tools_field(config, config.get("name", "orchestrator")) + + if "mcps" in config: + self._validate_mcps_field( + config["mcps"], config.get("name", "orchestrator") + ) + + # Resolve skill names to paths eagerly + skill_names = config.get("skills", []) + if skill_names: + config["skill_paths"] = resolve_skill_paths( + skill_names, + self._available_skills, + agent_name=config.get("name", "orchestrator"), + ) + + return config + except FileNotFoundError: + raise AppException( + f"Orchestrator config not found at {orchestrator_path}", + ErrorCodes.CONFIGURATION_VALIDATION_ERROR, + ) + except Exception as e: + raise AppException( + f"Failed to load orchestrator config: {e}", + ErrorCodes.CONFIGURATION_VALIDATION_ERROR, + ) + + def _load_all_subagents(self) -> dict[str, dict[str, Any]]: + """Load all subagent configurations at initialization. + + Returns: + Dict mapping subagent name to config dict with resolved skill paths. + """ + subagents_dir = self._base_dir / "subagents" + if not subagents_dir.is_dir(): + logger.warning(f"Subagents directory not found at {subagents_dir}") + return {} + + subagents = {} + for agent_file in sorted(subagents_dir.glob("*.md")): + try: + config = parse_frontmatter(agent_file) + if "body" in config: + config["body"] = inject_runtime_values(config["body"]) + + name = config.get("name", agent_file.stem) + + from deep_agent.src.infrastructure.tool_access import ( + migrate_tools_field, + ) + + migrate_tools_field(config, name) + + if "mcps" in config: + self._validate_mcps_field(config["mcps"], name) + + # Resolve skill names to paths eagerly + skill_names = config.get("skills", []) + if skill_names: + config["skill_paths"] = resolve_skill_paths( + skill_names, self._available_skills, agent_name=name + ) + + subagents[name] = config + logger.info(f"Loaded subagent config: {name}") + except Exception as e: + logger.error(f"Failed to load subagent {agent_file}: {e}") + + return subagents + + @staticmethod + def _validate_mcp_server(name: str, cfg: dict[str, Any]) -> None: + """Log clear errors for invalid per-MCP OAuth/DCR configuration.""" + auth_mode = cfg.get("auth_mode", "sso") + cfg["auth_mode"] = auth_mode + + if auth_mode not in ("sso", "oauth", "dcr", "api_key"): + logger.error( + "MCP server '%s': invalid auth_mode '%s' (expected sso, oauth, dcr, or api_key)", + name, + auth_mode, + ) + return + + if auth_mode not in ("oauth", "dcr"): + return + + oauth = cfg.get("oauth") + if not isinstance(oauth, dict): + logger.error( + "MCP server '%s': auth_mode '%s' requires an 'oauth' block", + name, + auth_mode, + ) + return + + for field in ( + "authorization_endpoint", + "token_endpoint", + ): + if not oauth.get(field): + logger.error( + "MCP server '%s': oauth.%s is required for auth_mode '%s'", + name, + field, + auth_mode, + ) + + if oauth.get("redirect_uri"): + logger.warning( + "MCP server '%s': oauth.redirect_uri in mcp.json is ignored — " + "redirect URI is derived from AGENT_PUBLIC_BASE_URL", + name, + ) + + if auth_mode == "oauth" and not oauth.get("client_id"): + logger.error( + "MCP server '%s': oauth.client_id is required for auth_mode 'oauth'", + name, + ) + + if oauth.get("client_secret"): + logger.warning( + "MCP server '%s': oauth.client_secret in mcp.json is insecure — " + "use oauth.client_secret_env with an environment variable name instead", + name, + ) + + if auth_mode == "dcr" and not oauth.get("registration_endpoint"): + logger.error( + "MCP server '%s': oauth.registration_endpoint is required for auth_mode 'dcr'", + name, + ) + + def _load_mcp_servers(self) -> dict[str, Any]: + """Load MCP server configuration at initialization. + + Returns: + Dict of MCP server configurations. + """ + mcp_path = self._base_dir / "mcp.json" + if not mcp_path.is_file(): + logger.warning(f"MCP config not found at {mcp_path}") + return {} + + try: + data = _load_jsonc(mcp_path) + servers: dict[str, Any] = data.get("mcpServers", {}) + for name, cfg in servers.items(): + if isinstance(cfg, dict): + self._validate_mcp_server(name, cfg) + logger.info(f"Loaded {len(servers)} MCP server config(s)") + return servers + except Exception as e: + logger.error(f"Failed to load MCP config: {e}") + return {} + + def _scan_available_skills(self) -> dict[str, Path]: + """Scan and index all available skills at initialization. + + Returns: + Dict mapping skill name to skill directory path. + """ + skills_dir = self._base_dir / "skills" + if not skills_dir.is_dir(): + logger.warning(f"Skills directory not found at {skills_dir}") + return {} + + skills = {} + for skill_path in skills_dir.iterdir(): + if skill_path.is_dir() and not skill_path.name.startswith("."): + skills[skill_path.name] = skill_path + logger.debug(f"Found skill: {skill_path.name}") + + logger.info(f"Scanned {len(skills)} available skill(s)") + return skills + + def get_orchestrator_config(self) -> dict[str, Any]: + """Get the pre-loaded orchestrator configuration. + + Returns: + Orchestrator config dict with all fields and injected runtime values. + """ + self._ensure_loaded() + return self._orchestrator + + def get_all_subagent_configs(self) -> dict[str, dict[str, Any]]: + """Get all subagent configurations. + + Returns: + Dict mapping subagent name to config dict. + """ + self._ensure_loaded() + return self._subagents + + @staticmethod + def resolve_tools( + tool_names: list[str], + available_tools: list[Any], + agent_name: str = "agent", + ) -> list[Any]: + """Resolve tool names to actual tool objects. + + This is a static method that delegates to the resolver module. + + Args: + tool_names: List of tool names from frontmatter. + available_tools: List of available tool objects. + agent_name: Name of the agent (for logging). + + Returns: + List of resolved tool objects. + """ + return resolve_tools(tool_names, available_tools, agent_name) + + def get_mcp_servers(self) -> dict[str, Any]: + """Get the pre-loaded MCP server configurations. + + Returns: + Dict of MCP server configurations. + """ + self._ensure_loaded() + return self._mcp_servers + + def get_providers_config(self) -> ProvidersFileConfig: + """Get the pre-loaded providers configuration. + + Returns: + The parsed providers.yaml config (strategy, profiles, async tasks). + """ + self._ensure_loaded() + return self._providers_config + + def get_filesystem_config(self) -> FilesystemFileConfig: + """Get the pre-loaded filesystem configuration. + + Returns: + The parsed filesystem.yaml config (backend, permissions, settings). + """ + self._ensure_loaded() + return self._filesystem_config + + def get_cache_config(self) -> CacheFileConfig: + """Get the pre-loaded cache configuration. + + Returns: + The parsed cache section (TTLs, feature flags, size limits). + """ + self._ensure_loaded() + return self._cache_config + + def get_token_budget_config(self) -> TokenBudgetConfig: + """Get the pre-loaded per-thread token budget configuration.""" + self._ensure_loaded() + return self._token_budget_config + + def get_name(self) -> str: + """Get the agent display name from config. + + Returns: + The agent name as configured in agent.yaml (top-level `name` field). + """ + self._ensure_loaded() + return self._name + + def get_middleware_config(self) -> MiddlewareFileConfig: + """Get the pre-loaded middleware file configuration. + + Returns: + The parsed middleware.yaml config (defaults + profiles). + """ + self._ensure_loaded() + return self._middleware_config + + def resolve_agent_middleware( + self, + model_name: str, + agent_overrides: dict[str, Any] | None = None, + ) -> ResolvedMiddlewareConfig: + """Resolve middleware config for a specific agent. + + Merges: global defaults → profile (from model) → per-agent overrides. + + Args: + model_name: Model name from agent frontmatter. + agent_overrides: Optional middleware: block from frontmatter. + + Returns: + Fully resolved middleware configuration. + """ + self._ensure_loaded() + return resolve_middleware(self._middleware_config, model_name, agent_overrides) + + def get_otel_config(self) -> OtelFileConfig: + """Get the pre-loaded OTEL configuration. + + Returns: + The parsed OTEL config from observability.yaml. + """ + self._ensure_loaded() + return self._otel_config + + def get_pyproject_path(self) -> Path: + """Get the skill sandbox pyproject.toml path. + + Returns: + Path to config/skills/pyproject.toml for skill sandbox dependencies. + """ + return self._base_dir / "skills" / "pyproject.toml" + + +# Singleton instance +agent_config = AgentConfig(_AGENT_CONFIG_DIR) diff --git a/deep_agent/src/agent/config/middleware.py b/deep_agent/src/agent/config/middleware.py new file mode 100644 index 00000000..6275a433 --- /dev/null +++ b/deep_agent/src/agent/config/middleware.py @@ -0,0 +1,285 @@ +"""Middleware configuration models and resolution logic. + +Provides Pydantic models for the ``middleware:`` and ``harness_profiles:`` +sections of config/agent/runtime/agent.yaml and resolves the final +middleware configuration for each agent by merging: + + global defaults → profile (matched from model field) → per-agent overrides + +The template-agent user only touches YAML config. This module converts +declarative config into the parameters needed by the middleware builder. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Literal + +import yaml +from pydantic import BaseModel, Field + +from deep_agent.src.code_execution.config import CodeExecutionConfig +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +class SummarizationToolConfig(BaseModel): + """Config for SummarizationToolMiddleware.""" + + enabled: bool = True + + +class HumanApprovalConfig(BaseModel): + """Config for human-in-the-loop tool approval. + + When enabled, the agent pauses before executing any tool call and + waits for the user to approve, reject, or always-allow it. + Backed by deepagents HumanInTheLoopMiddleware via interrupt_on. + """ + + enabled: bool = True + mode: Literal["all", "none"] = "all" + exclude: list[str] = Field(default_factory=list) + + +class MemoryConfig(BaseModel): + """Config for MemoryMiddleware (activated via memory= param).""" + + enabled: bool = True + namespaces: list[str] = Field(default_factory=lambda: ["memories"]) + + +class PatchToolCallsConfig(BaseModel): + """Config for PatchToolCallsMiddleware (auto-included by deepagents).""" + + enabled: bool = True + + +class SkillsConfig(BaseModel): + """Config for SkillsMiddleware (auto-included when skills= provided).""" + + enabled: bool = True + + +class ModelCallLimitConfig(BaseModel): + """Config for ModelCallLimitMiddleware — cap LLM calls per run.""" + + enabled: bool = True + run_limit: int = 50 + + +class ToolCallLimitConfig(BaseModel): + """Config for ToolCallLimitMiddleware — cap tool calls per run.""" + + enabled: bool = True + run_limit: int = 200 + + +class ModelRetryConfig(BaseModel): + """Config for ModelRetryMiddleware — retry on transient failures.""" + + enabled: bool = True + max_retries: int = 3 + backoff_factor: float = 2.0 + initial_delay: float = 1.0 + + +class ModelFallbackConfig(BaseModel): + """Config for ModelFallbackMiddleware — switch to backup model.""" + + enabled: bool = False + fallback_model: str = "" + + +class ToolRetryConfig(BaseModel): + """Config for ToolRetryMiddleware — retry specific tools.""" + + enabled: bool = False + max_retries: int = 2 + tools: list[str] = Field(default_factory=list) + + +class PIIRule(BaseModel): + """A single PII detection rule.""" + + type: str + strategy: str = "redact" + + +class PIIConfig(BaseModel): + """Config for PIIMiddleware — detect and handle PII.""" + + enabled: bool = False + rules: list[PIIRule] = Field(default_factory=list) + + +class MiddlewareDefaults(BaseModel): + """Global middleware defaults from middleware.yaml.""" + + summarization_tool: SummarizationToolConfig = Field( + default_factory=SummarizationToolConfig + ) + human_approval: HumanApprovalConfig = Field(default_factory=HumanApprovalConfig) + memory: MemoryConfig = Field(default_factory=MemoryConfig) + patch_tool_calls: PatchToolCallsConfig = Field(default_factory=PatchToolCallsConfig) + skills: SkillsConfig = Field(default_factory=SkillsConfig) + model_call_limit: ModelCallLimitConfig = Field(default_factory=ModelCallLimitConfig) + tool_call_limit: ToolCallLimitConfig = Field(default_factory=ToolCallLimitConfig) + model_retry: ModelRetryConfig = Field(default_factory=ModelRetryConfig) + model_fallback: ModelFallbackConfig = Field(default_factory=ModelFallbackConfig) + tool_retry: ToolRetryConfig = Field(default_factory=ToolRetryConfig) + pii: PIIConfig = Field(default_factory=PIIConfig) + extra: list[str] = Field(default_factory=list) + code_execution: CodeExecutionConfig = Field(default_factory=CodeExecutionConfig) + + +class ProfileConfig(BaseModel): + """Per-model profile configuration for HarnessProfile registration.""" + + excluded_middleware: list[str] = Field(default_factory=list) + excluded_tools: list[str] = Field(default_factory=list) + system_prompt_suffix: str = "" + general_purpose_subagent: dict[str, Any] = Field(default_factory=dict) + + +class MiddlewareFileConfig(BaseModel): + """Structure of the middleware + harness_profiles sections in runtime/agent.yaml.""" + + defaults: MiddlewareDefaults = Field(default_factory=MiddlewareDefaults) + profiles: dict[str, ProfileConfig] = Field(default_factory=dict) + + +class ResolvedMiddlewareConfig(BaseModel): + """Final resolved config for a single agent after merge.""" + + summarization_tool_enabled: bool = True + human_approval: HumanApprovalConfig = Field(default_factory=HumanApprovalConfig) + memory_enabled: bool = True + memory_namespaces: list[str] = Field(default_factory=lambda: ["memories"]) + patch_tool_calls_enabled: bool = True + skills_enabled: bool = True + model_call_limit: ModelCallLimitConfig = Field(default_factory=ModelCallLimitConfig) + tool_call_limit: ToolCallLimitConfig = Field(default_factory=ToolCallLimitConfig) + model_retry: ModelRetryConfig = Field(default_factory=ModelRetryConfig) + model_fallback: ModelFallbackConfig = Field(default_factory=ModelFallbackConfig) + tool_retry: ToolRetryConfig = Field(default_factory=ToolRetryConfig) + pii: PIIConfig = Field(default_factory=PIIConfig) + extra_middleware: list[str] = Field(default_factory=list) + excluded_middleware: list[str] = Field(default_factory=list) + code_execution: CodeExecutionConfig = Field(default_factory=CodeExecutionConfig) + + +def load_middleware_config(config_path: Path) -> MiddlewareFileConfig: + """Load and validate middleware.yaml from disk. + + Args: + config_path: Path to middleware.yaml. + + Returns: + Validated MiddlewareFileConfig. Returns defaults if file is missing. + """ + if not config_path.is_file(): + logger.info("No middleware.yaml found — using defaults") + return MiddlewareFileConfig() + + try: + raw = yaml.safe_load(config_path.read_text()) or {} + config: MiddlewareFileConfig = MiddlewareFileConfig.model_validate(raw) + logger.info("Loaded middleware config: %d profile(s)", len(config.profiles)) + return config + except Exception as e: + logger.warning("Failed to parse middleware.yaml, using defaults: %s", e) + return MiddlewareFileConfig() + + +def resolve_middleware( + file_config: MiddlewareFileConfig, + model_name: str, + agent_overrides: dict[str, Any] | None = None, +) -> ResolvedMiddlewareConfig: + """Resolve final middleware config for an agent. + + Merge order: global defaults → profile (from model name) → agent overrides. + + Args: + file_config: Parsed middleware.yaml config. + model_name: Model name from agent frontmatter (used for profile lookup). + agent_overrides: Optional middleware: block from agent frontmatter. + + Returns: + Fully resolved middleware configuration for this agent. + """ + defaults = file_config.defaults + profile = file_config.profiles.get(model_name, ProfileConfig()) + overrides = agent_overrides or {} + + summarization_enabled = _resolve_bool( + defaults.summarization_tool.enabled, + overrides.get("summarization_tool"), + ) + memory_enabled = _resolve_bool( + defaults.memory.enabled, + overrides.get("memory"), + ) + patch_enabled = _resolve_bool( + defaults.patch_tool_calls.enabled, + overrides.get("patch_tool_calls"), + ) + skills_enabled = _resolve_bool( + defaults.skills.enabled, + overrides.get("skills"), + ) + + memory_namespaces = defaults.memory.namespaces + if isinstance(overrides.get("memory"), dict): + memory_namespaces = overrides["memory"].get("namespaces", memory_namespaces) + + extra = list(defaults.extra) + if "extra" in overrides: + extra.extend(overrides["extra"]) + + if "patch_tool_calls" in profile.excluded_middleware: + patch_enabled = False + + human_approval = defaults.human_approval + if isinstance(overrides.get("human_approval"), dict): + human_approval = HumanApprovalConfig.model_validate(overrides["human_approval"]) + elif isinstance(overrides.get("human_approval"), bool): + human_approval = HumanApprovalConfig(enabled=overrides["human_approval"]) + + code_execution = defaults.code_execution + if isinstance(overrides.get("code_execution"), dict): + code_execution = CodeExecutionConfig.model_validate(overrides["code_execution"]) + + return ResolvedMiddlewareConfig( + summarization_tool_enabled=summarization_enabled, + human_approval=human_approval, + memory_enabled=memory_enabled, + memory_namespaces=memory_namespaces, + patch_tool_calls_enabled=patch_enabled, + skills_enabled=skills_enabled, + model_call_limit=defaults.model_call_limit, + tool_call_limit=defaults.tool_call_limit, + model_retry=defaults.model_retry, + model_fallback=defaults.model_fallback, + tool_retry=defaults.tool_retry, + pii=defaults.pii, + extra_middleware=extra, + excluded_middleware=profile.excluded_middleware, + code_execution=code_execution, + ) + + +def _resolve_bool(default: bool, override: Any) -> bool: + """Resolve a boolean config with potential override. + + Override can be: bool, dict with 'enabled' key, or None (use default). + """ + if override is None: + return default + if isinstance(override, bool): + return override + if isinstance(override, dict): + return bool(override.get("enabled", default)) + return default diff --git a/deep_agent/src/agent/config/model.py b/deep_agent/src/agent/config/model.py new file mode 100644 index 00000000..02e62e55 --- /dev/null +++ b/deep_agent/src/agent/config/model.py @@ -0,0 +1,134 @@ +"""Model configuration types for per-agent LLM provider selection. + +Parses frontmatter ``model:`` fields that may be a legacy string or an +object with explicit provider, model name, and optional fallback chain. +""" + +from __future__ import annotations + +from enum import Enum +from typing import Any + +from pydantic import BaseModel, model_validator + +from deep_agent.src.agent.llm import CLAUDE_MODELS, GEMINI_MODELS + + +class Provider(str, Enum): + """Supported LLM provider backends.""" + + VERTEX = "vertex" + OPENAI = "openai" + MAAS = "maas" # Model as a Service (VLLM) + + +class ModelSpec(BaseModel): + """Resolved model configuration with optional fallback.""" + + provider: Provider + name: str + fallback: ModelSpec | None = None + + @model_validator(mode="after") + def _validate_name(self) -> ModelSpec: + if not self.name or not self.name.strip(): + raise ValueError("model name cannot be empty") + return self + + def display_name(self) -> str: + """Human-readable model identifier for logging.""" + base = f"{self.provider.value}:{self.name}" + if self.fallback: + return f"{base} (fallback: {self.fallback.display_name()})" + return base + + +def infer_provider(model_name: str) -> Provider: + """Infer provider from a legacy model name string. + + Inference logic: + - Known Gemini/Claude models → VERTEX + - GPT models (gpt-*, case-insensitive) → OPENAI + - All other models → MAAS (VLLM for custom models) + """ + if model_name in GEMINI_MODELS or model_name in CLAUDE_MODELS: + return Provider.VERTEX + if model_name.lower().startswith("gpt-"): + return Provider.OPENAI + return Provider.MAAS + + +def parse_model_config(raw: str | dict[str, Any]) -> ModelSpec: + """Parse a frontmatter ``model`` field into a :class:`ModelSpec`. + + Accepts: + - Legacy string: ``gemini-2.5-pro`` (provider inferred) + - Object: ``{provider: vertex, name: gemini-2.5-pro, fallback: {...}}`` + + Args: + raw: Model value from parsed frontmatter. + + Returns: + Validated ModelSpec. + + Raises: + ValueError: If the config is invalid or missing required fields. + TypeError: If raw is neither str nor dict. + """ + if isinstance(raw, str): + name = raw.strip() + if not name: + raise ValueError("model name cannot be empty") + return ModelSpec(provider=infer_provider(name), name=name) + + if not isinstance(raw, dict): + raise TypeError(f"model config must be str or dict, got {type(raw).__name__}") + + allowed_keys = {"provider", "name", "fallback"} + unknown = set(raw.keys()) - allowed_keys + if unknown: + raise ValueError( + f"unknown model config keys: {sorted(unknown)}; " + f"allowed: {sorted(allowed_keys)}" + ) + + name_raw = raw.get("name") + if not isinstance(name_raw, str) or not name_raw.strip(): + raise ValueError("model config object requires non-empty 'name'") + name = name_raw.strip() + + # Provider is optional - infer from name if not provided + provider_raw = raw.get("provider") + if provider_raw is None: + provider = infer_provider(name) + else: + try: + provider = Provider(provider_raw) + except ValueError as e: + raise ValueError( + f"invalid provider '{provider_raw}'; " + f"must be one of: {[p.value for p in Provider]}" + ) from e + + fallback_raw = raw.get("fallback") + fallback: ModelSpec | None = None + if fallback_raw is not None: + if not isinstance(fallback_raw, dict): + raise ValueError("model fallback must be an object") + if "fallback" in fallback_raw: + raise ValueError("nested fallback chains are not supported") + fallback = parse_model_config(fallback_raw) + + return ModelSpec( + provider=provider, + name=str(name).strip(), + fallback=fallback, + ) + + +def model_spec_cache_key(spec: ModelSpec) -> str: + """Stable cache identity string for a model spec including fallback.""" + parts = [f"{spec.provider.value}:{spec.name}"] + if spec.fallback: + parts.append(f"→{model_spec_cache_key(spec.fallback)}") + return "".join(parts) diff --git a/deep_agent/src/agent/config/otel.py b/deep_agent/src/agent/config/otel.py new file mode 100644 index 00000000..9aeb6fcd --- /dev/null +++ b/deep_agent/src/agent/config/otel.py @@ -0,0 +1,41 @@ +"""OpenTelemetry configuration models. + +Provides validated Pydantic models for the ``otel:`` section of +config/agent/runtime/observability.yaml. Controls OTLP exporter +settings, metric export intervals, and tracing behavior. + +The template-agent user only touches YAML. This module converts +declarative config into parameters consumed by the OTEL SDK. +""" + +from __future__ import annotations + +from pydantic import BaseModel, Field + + +class OtelExporterConfig(BaseModel): + """OTLP exporter connection settings.""" + + endpoint: str = Field(default="http://localhost:4317") + insecure: bool = True + + +class OtelMetricsConfig(BaseModel): + """Metric export settings.""" + + export_interval_ms: int = Field(default=5000, ge=1000, le=60000) + + +class OtelTracingConfig(BaseModel): + """Distributed tracing settings.""" + + fastapi_auto_instrument: bool = True + + +class OtelFileConfig(BaseModel): + """Top-level OTEL configuration from observability.yaml ``otel:`` section.""" + + enabled: bool = False + exporter: OtelExporterConfig = Field(default_factory=OtelExporterConfig) + metrics: OtelMetricsConfig = Field(default_factory=OtelMetricsConfig) + tracing: OtelTracingConfig = Field(default_factory=OtelTracingConfig) diff --git a/deep_agent/src/agent/config/parser.py b/deep_agent/src/agent/config/parser.py new file mode 100644 index 00000000..c5f7b32b --- /dev/null +++ b/deep_agent/src/agent/config/parser.py @@ -0,0 +1,66 @@ +"""Frontmatter parsing and runtime value injection. + +This module handles parsing markdown files with YAML frontmatter (used for agent +configurations) and injecting runtime values like {{current_date}} into the content. + +Why this exists: + Agent configs are written in markdown with YAML frontmatter. This module + extracts the frontmatter metadata and body content, and replaces template + variables with runtime values. + +Functions: + parse_frontmatter: Parse markdown file with YAML frontmatter + inject_runtime_values: Replace template variables with actual values +""" + +from datetime import datetime +from pathlib import Path +from typing import Any + +import yaml + + +def get_current_date() -> str: + """Get the current date in a formatted string. + + Returns: + The current date formatted as "Month Day, Year" (e.g., "December 25, 2024"). + """ + return datetime.now().strftime("%B %d, %Y") + + +def inject_runtime_values(content: str) -> str: + """Inject runtime values into content. + + Args: + content: String content with template variables. + + Returns: + Content with template variables replaced. + """ + return content.replace("{{current_date}}", get_current_date()) + + +def parse_frontmatter(path: Path) -> dict[str, Any]: + r"""Parse a markdown file with YAML frontmatter. + + Expects the format: ``--- \n \n --- \n ``. + The markdown body is returned under the ``"body"`` key. + + Args: + path: Path to the ``.md`` file. + + Returns: + A dict of frontmatter fields plus ``body`` (the markdown content). + """ + content = path.read_text() + if not content.startswith("---"): + return {"body": content.strip()} + + parts = content.split("---", 2) + if len(parts) < 3: + return {"body": content.strip()} + + frontmatter: dict[str, Any] = yaml.safe_load(parts[1]) or {} + frontmatter["body"] = parts[2].strip() + return frontmatter diff --git a/deep_agent/src/agent/config/providers.py b/deep_agent/src/agent/config/providers.py new file mode 100644 index 00000000..b3ef723b --- /dev/null +++ b/deep_agent/src/agent/config/providers.py @@ -0,0 +1,92 @@ +"""Provider and harness profile configuration models. + +Provides Pydantic models for the ``providers:``, ``harness_profiles:``, +and ``async_tasks:`` sections of config/agent/runtime/agent.yaml: +- Model resolution strategy (legacy vs deepagents) +- ProviderProfile registration (init_chat_model kwargs per provider) +- HarnessProfile registration (runtime adjustments per model) +- Async task middleware configuration + +Users edit YAML. This module validates and converts to typed config. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Literal + +import yaml +from pydantic import BaseModel, Field + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +class ProviderConfig(BaseModel): + """Configuration for a single provider (maps to ProviderProfile).""" + + init_kwargs: dict[str, Any] = Field(default_factory=dict) + + +class GeneralPurposeSubagentConfig(BaseModel): + """Config for the auto-added general-purpose subagent.""" + + enabled: bool = True + description: str | None = None + system_prompt: str | None = None + + +class HarnessProfileConfig(BaseModel): + """Configuration for a single harness profile (maps to HarnessProfile).""" + + system_prompt_suffix: str = "" + excluded_tools: list[str] = Field(default_factory=list) + excluded_middleware: list[str] = Field(default_factory=list) + general_purpose_subagent: GeneralPurposeSubagentConfig = Field( + default_factory=GeneralPurposeSubagentConfig, + ) + + +class AsyncTaskConfig(BaseModel): + """Configuration for AsyncSubAgentMiddleware.""" + + enabled: bool = True + system_prompt: str | None = None + + +class ProvidersFileConfig(BaseModel): + """Structure of the providers + harness_profiles sections in runtime/agent.yaml.""" + + resolve_strategy: Literal["legacy", "deepagents"] = "legacy" + providers: dict[str, ProviderConfig] = Field(default_factory=dict) + harness_profiles: dict[str, HarnessProfileConfig] = Field(default_factory=dict) + async_tasks: AsyncTaskConfig = Field(default_factory=AsyncTaskConfig) + + +def load_providers_config(config_path: Path) -> ProvidersFileConfig: + """Load and validate providers.yaml from disk. + + Args: + config_path: Path to providers.yaml. + + Returns: + Validated ProvidersFileConfig. Returns defaults if file is missing. + """ + if not config_path.is_file(): + logger.info("No providers.yaml found — using defaults (legacy resolution)") + return ProvidersFileConfig() + + try: + raw = yaml.safe_load(config_path.read_text()) or {} + config: ProvidersFileConfig = ProvidersFileConfig.model_validate(raw) + logger.info( + "Loaded providers config: strategy=%s, %d provider(s), %d harness profile(s)", + config.resolve_strategy, + len(config.providers), + len(config.harness_profiles), + ) + return config + except Exception as e: + logger.warning("Failed to parse providers.yaml, using defaults: %s", e) + return ProvidersFileConfig() diff --git a/deep_agent/src/agent/config/resolver.py b/deep_agent/src/agent/config/resolver.py new file mode 100644 index 00000000..5a610639 --- /dev/null +++ b/deep_agent/src/agent/config/resolver.py @@ -0,0 +1,112 @@ +"""Skill and tool resolution utilities. + +This module resolves skill names to directory paths and tool names to tool objects. +It handles validation, logging of missing dependencies, and returns only the +successfully resolved items. + +Why this exists: + Agent configs reference skills and tools by name (strings). This module + looks up those names in the available skills directory and MCP tools list, + returning the actual paths/objects needed for agent initialization. + +Functions: + resolve_skill_paths: Convert skill names to skill directory paths + to_virtual_skill_paths: Convert absolute paths to virtual /skills/ paths + resolve_tools: Convert tool names to tool objects +""" + +from pathlib import Path +from typing import Any + +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger(log_level=settings.PYTHON_LOG_LEVEL) + + +def resolve_skill_paths( + skill_names: list[str], + available_skills: dict[str, Path], + agent_name: str = "agent", +) -> list[str]: + """Resolve skill names to skill directory paths using cached skill index. + + Args: + skill_names: List of skill names from frontmatter. + available_skills: Dict mapping skill name to skill directory path. + agent_name: Name of the agent (for logging). + + Returns: + List of skill directory paths as strings. + """ + skill_paths: list[str] = [] + missing: list[str] = [] + + for skill_name in skill_names: + if skill_name in available_skills: + skill_path = available_skills[skill_name] + skill_paths.append(str(skill_path)) + logger.debug(f"Agent '{agent_name}' resolved skill: {skill_name}") + else: + missing.append(skill_name) + + if missing: + logger.warning(f"Agent '{agent_name}' references unknown skills: {missing}") + + return skill_paths + + +def to_virtual_skill_paths(skill_paths: list[str]) -> list[str]: + """Convert absolute filesystem skill paths to virtual /skills/ paths. + + The CompositeBackend routes /skills/ to a ReadOnlyFilesystemBackend. This + function transforms the absolute paths produced by resolve_skill_paths() + into the virtual paths expected by that routing. + + Note: only the leaf directory name is used. Nested skill directories + (e.g., /skills/category/my-skill) are not currently supported. + + Args: + skill_paths: Absolute filesystem paths from resolve_skill_paths(). + + Returns: + Virtual paths like ["/skills/my-skill", "/skills/other-skill"]. + """ + virtual: list[str] = [] + for p in skill_paths: + name = Path(p).name + parent_name = Path(p).parent.name + if parent_name != "skills": + logger.warning( + "Skill path '%s' is not directly under a 'skills/' directory — " + "only leaf name '%s' is used for virtual path", + p, + name, + ) + virtual.append(f"/skills/{name}") + return virtual + + +def resolve_tools( + tool_names: list[str], + available_tools: list[Any], + agent_name: str = "agent", +) -> list[Any]: + """Resolve tool names to actual tool objects. + + Args: + tool_names: List of tool names from frontmatter. + available_tools: List of available tool objects. + agent_name: Name of the agent (for logging). + + Returns: + List of resolved tool objects. + """ + tool_by_name = {t.name: t for t in available_tools} + resolved = [tool_by_name[n] for n in tool_names if n in tool_by_name] + missing = [n for n in tool_names if n not in tool_by_name] + + if missing: + logger.warning(f"Agent '{agent_name}' references unknown tools: {missing}") + + return resolved diff --git a/deep_agent/src/agent/llm.py b/deep_agent/src/agent/llm.py new file mode 100644 index 00000000..ab78ac6e --- /dev/null +++ b/deep_agent/src/agent/llm.py @@ -0,0 +1,182 @@ +"""LLM factory for creating configured model instances. + +Supports three provider paths: + 1. Gemini (via langchain_google_genai + Vertex AI service account) + 2. Claude (via langchain_google_vertexai Model Garden) + 3. vLLM / OpenAI-compatible (via langchain_openai + custom base_url) + +Any model name not in GEMINI_MODELS or CLAUDE_MODELS is assumed to be +served by a vLLM (or OpenAI-compatible) endpoint. Set VLLM_BASE_URL +to the inference server's /v1 endpoint. +""" + +from langchain_core.language_models import BaseChatModel +from langchain_google_genai import ChatGoogleGenerativeAI +from langchain_google_vertexai.model_garden import ChatAnthropicVertex + +from deep_agent.src.error_handling import llm_retry +from deep_agent.src.exceptions import LLMError +from deep_agent.src.settings import settings +from deep_agent.utils.google_creds import get_service_account_credentials +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger(log_level=settings.PYTHON_LOG_LEVEL) + +_DEFAULT_MAX_OUTPUT_TOKENS: int = settings.MAX_OUTPUT_TOKENS + +GEMINI_MODELS: list[str] = [ + "gemini-2.5-pro", + "gemini-2.5-flash", + "gemini-3.1-pro-preview", +] + +CLAUDE_MODELS: list[str] = [ + "claude-sonnet-4", + "claude-sonnet-4-6@default", +] + + +@llm_retry +def create_model( + model_name: str, + temperature: float = 0.0, + max_output_tokens: int | None = None, +) -> BaseChatModel: + """Create a model instance (Vertex AI, or vLLM/OpenAI-compatible). + + Resolution order: + 1. If model_name is in GEMINI_MODELS → Vertex AI Gemini + 2. If model_name is in CLAUDE_MODELS → Vertex AI Claude (Model Garden) + 3. Otherwise → vLLM / OpenAI-compatible endpoint (requires VLLM_BASE_URL) + + Args: + model_name: Model identifier (Gemini/Claude name, or vLLM model path). + temperature: Model temperature (default: 0.0). + max_output_tokens: Maximum tokens in model response (default: 8192). + + Returns: + Configured model instance. + + Raises: + ValueError: If model_name is empty or vLLM is needed but not configured. + LLMError: If model creation fails after retries. + """ + if not model_name or not model_name.strip(): + raise ValueError("model_name cannot be empty") + + max_output_tokens = max_output_tokens or _DEFAULT_MAX_OUTPUT_TOKENS + + is_gemini = model_name in GEMINI_MODELS + is_claude = model_name in CLAUDE_MODELS + + if is_gemini or is_claude: + return _create_vertex_model(model_name, temperature, max_output_tokens) + + return _create_vllm_model(model_name, temperature, max_output_tokens) + + +def _create_vertex_model( + model_name: str, + temperature: float, + max_output_tokens: int, +) -> BaseChatModel: + """Create a Vertex AI model (Gemini or Claude).""" + is_claude = model_name in CLAUDE_MODELS + model_type = "Claude" if is_claude else "Gemini" + + try: + credentials, project = get_service_account_credentials() + + logger.info( + f"Creating {model_type} model via Vertex AI", + model=model_name, + project=project, + temperature=temperature, + max_output_tokens=max_output_tokens, + ) + + if is_claude: + return ChatAnthropicVertex( + model=model_name, + project=project, + credentials=credentials, + temperature=temperature, + max_tokens=max_output_tokens, + max_retries=2, + ) + else: + return ChatGoogleGenerativeAI( + model=model_name, + temperature=temperature, + credentials=credentials, + project=project, + max_output_tokens=max_output_tokens, + max_retries=2, + ) + + except (ValueError, LLMError): + raise + except Exception as e: + logger.error( + f"Failed to create {model_type} model '{model_name}'", + error_type=type(e).__name__, + model=model_name, + error_message=str(e), + exc_info=True, + ) + raise LLMError( + f"Failed to create {model_type} model '{model_name}': {e}" + ) from e + + +def _create_vllm_model( + model_name: str, + temperature: float, + max_output_tokens: int, +) -> BaseChatModel: + """Create a model via vLLM / OpenAI-compatible endpoint. + + vLLM, TGI, Ollama, and any server exposing /v1/chat/completions works. + """ + if not settings.VLLM_BASE_URL: + raise ValueError( + f"Model '{model_name}' is not a known Vertex AI model. " + f"Set VLLM_BASE_URL to use it via an OpenAI-compatible endpoint. " + f"Known Vertex AI models: {GEMINI_MODELS + CLAUDE_MODELS}" + ) + + try: + from langchain_openai import ChatOpenAI + + logger.info( + "Creating model via vLLM/OpenAI-compatible endpoint", + model=model_name, + base_url=settings.VLLM_BASE_URL, + temperature=temperature, + max_output_tokens=max_output_tokens, + ) + + return ChatOpenAI( + model=model_name, + base_url=settings.VLLM_BASE_URL, + api_key=settings.VLLM_API_KEY, + temperature=temperature, + max_tokens=max_output_tokens, + max_retries=2, + ) + + except ImportError: + raise LLMError( + "langchain-openai is required for vLLM support. " + "Add 'langchain-openai' to your dependencies." + ) + except Exception as e: + logger.error( + f"Failed to create vLLM model '{model_name}'", + error_type=type(e).__name__, + model=model_name, + base_url=settings.VLLM_BASE_URL, + error_message=str(e), + exc_info=True, + ) + raise LLMError(f"Failed to create vLLM model '{model_name}': {e}") from e diff --git a/deep_agent/src/agent/provider_factory.py b/deep_agent/src/agent/provider_factory.py new file mode 100644 index 00000000..a147deba --- /dev/null +++ b/deep_agent/src/agent/provider_factory.py @@ -0,0 +1,74 @@ +"""Unified LLM provider factory for per-agent model resolution. + +Routes model creation to Vertex AI or OpenAI-compatible backends based on +an explicit :class:`ModelSpec`, optionally chaining a fallback model via +LangChain's ``with_fallbacks``. +""" + +from __future__ import annotations + +from langchain_core.language_models import BaseChatModel + +from deep_agent.src.agent.config.model import ModelSpec, Provider +from deep_agent.src.agent.llm import _create_vertex_model, _create_vllm_model +from deep_agent.src.settings import settings + + +def create_model_from_spec( + spec: ModelSpec, + *, + temperature: float = 0.0, + max_output_tokens: int | None = None, +) -> BaseChatModel: + """Create a chat model from a :class:`ModelSpec`. + + When ``spec.fallback`` is set, wraps the primary model with + ``primary.with_fallbacks([secondary])`` so invocation failures on the + primary route to the secondary model. + + Args: + spec: Parsed model configuration. + temperature: Model temperature. + max_output_tokens: Maximum output tokens (defaults to settings). + + Returns: + A BaseChatModel instance, optionally with fallback chain. + """ + tokens = max_output_tokens or settings.MAX_OUTPUT_TOKENS + primary = _create_by_provider( + spec.provider, spec.name, temperature=temperature, max_output_tokens=tokens + ) + + if spec.fallback is None: + return primary + + secondary = _create_by_provider( + spec.fallback.provider, + spec.fallback.name, + temperature=temperature, + max_output_tokens=tokens, + ) + return primary.with_fallbacks([secondary]) + + +def _create_by_provider( + provider: Provider, + model_name: str, + *, + temperature: float, + max_output_tokens: int, +) -> BaseChatModel: + """Route model creation to the appropriate backend. + + Routes: + - VERTEX → Google Vertex AI (Gemini, Claude) + - OPENAI → OpenAI API (GPT models) + - MAAS → VLLM (Model as a Service for custom models) + """ + if provider == Provider.VERTEX: + return _create_vertex_model(model_name, temperature, max_output_tokens) + if provider == Provider.OPENAI: + return _create_vllm_model(model_name, temperature, max_output_tokens) + if provider == Provider.MAAS: + return _create_vllm_model(model_name, temperature, max_output_tokens) + raise ValueError(f"unsupported provider: {provider}") diff --git a/deep_agent/src/audit/__init__.py b/deep_agent/src/audit/__init__.py new file mode 100644 index 00000000..2b3d1a15 --- /dev/null +++ b/deep_agent/src/audit/__init__.py @@ -0,0 +1,11 @@ +"""Audit logging — structured events gated by PLATFORM_AUDIT_ENABLED.""" + +from deep_agent.src.audit.config import is_audit_enabled +from deep_agent.src.audit.emitter import emit_audit_event +from deep_agent.src.audit.events import AuditEventType + +__all__ = [ + "AuditEventType", + "emit_audit_event", + "is_audit_enabled", +] diff --git a/deep_agent/src/audit/buffer.py b/deep_agent/src/audit/buffer.py new file mode 100644 index 00000000..eea3db93 --- /dev/null +++ b/deep_agent/src/audit/buffer.py @@ -0,0 +1,41 @@ +"""Local audit event buffer — in-memory queue for transient failures.""" + +from __future__ import annotations + +from collections import deque +from threading import Lock +from typing import Any + +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_lock = Lock() +_queue: deque[dict[str, Any]] = deque() +_dropped = 0 + + +def enqueue(envelope: dict[str, Any]) -> None: + """Append envelope to in-memory buffer.""" + global _dropped # noqa: PLW0603 + buffer_max = settings.PLATFORM_AUDIT_BUFFER_MAX + with _lock: + if len(_queue) >= buffer_max: + _dropped += 1 + if _dropped == 1 or _dropped % 100 == 0: + logger.warning( + "platform_audit_buffer_full", + dropped=_dropped, + max=buffer_max, + ) + return + _queue.append(envelope) + + +def drain() -> list[dict[str, Any]]: + """Return and clear all buffered envelopes.""" + with _lock: + items = list(_queue) + _queue.clear() + return items diff --git a/deep_agent/src/audit/config.py b/deep_agent/src/audit/config.py new file mode 100644 index 00000000..ffbed942 --- /dev/null +++ b/deep_agent/src/audit/config.py @@ -0,0 +1,17 @@ +"""Platform audit configuration. + +Environment variables (loaded via ``settings.py``): + PLATFORM_AUDIT_ENABLED: Master switch (default: false) + PLATFORM_AUDIT_BUFFER_MAX: Max in-memory buffered events (default: 1000) + +YAML reference (agent.yaml): + platform.audit.enabled + platform.audit.buffer_max +""" + +from deep_agent.src.settings import settings + + +def is_audit_enabled() -> bool: + """Return whether platform audit is enabled.""" + return settings.PLATFORM_AUDIT_ENABLED diff --git a/deep_agent/src/audit/context.py b/deep_agent/src/audit/context.py new file mode 100644 index 00000000..3bb2cb7f --- /dev/null +++ b/deep_agent/src/audit/context.py @@ -0,0 +1,108 @@ +"""Audit context — user, org, trace_id for event envelopes.""" + +from __future__ import annotations + +from contextvars import ContextVar + +_trace_id_var: ContextVar[str | None] = ContextVar("audit_trace_id", default=None) +_user_var: ContextVar[str | None] = ContextVar("audit_user", default=None) +_org_var: ContextVar[str | None] = ContextVar("audit_org", default=None) + + +def bind_audit_context( + *, + trace_id: str | None = None, + user: str | None = None, + org: str | None = None, +) -> None: + """Bind audit identifiers for the current async context.""" + MAX_LEN = 512 # Prevent memory exhaustion from extremely long strings + + if trace_id is not None: + if not isinstance(trace_id, str): + raise TypeError("trace_id must be a string") + trace_id = trace_id.strip() + if not trace_id: + raise ValueError("trace_id cannot be empty or whitespace") + if len(trace_id) > MAX_LEN: + raise ValueError(f"trace_id exceeds maximum length of {MAX_LEN}") + _trace_id_var.set(trace_id) + + if user is not None: + if not isinstance(user, str): + raise TypeError("user must be a string") + user = user.strip() + if not user: + raise ValueError("user cannot be empty or whitespace") + if len(user) > MAX_LEN: + raise ValueError(f"user exceeds maximum length of {MAX_LEN}") + _user_var.set(user) + + if org is not None: + if not isinstance(org, str): + raise TypeError("org must be a string") + org = org.strip() + if not org: + raise ValueError("org cannot be empty or whitespace") + if len(org) > MAX_LEN: + raise ValueError(f"org exceeds maximum length of {MAX_LEN}") + _org_var.set(org) + + +def clear_audit_context() -> None: + """Reset audit context vars.""" + _trace_id_var.set(None) + _user_var.set(None) + _org_var.set(None) + + +def get_audit_context() -> dict[str, str | None]: + """Return current audit context fields.""" + return { + "trace_id": _trace_id_var.get(), + "user": _user_var.get(), + "org": _org_var.get(), + } + + +def resolve_trace_id_from_config() -> str | None: + """Read trace_id from LangGraph RunnableConfig metadata if available.""" + try: + from langgraph.config import get_config + + config = get_config() + metadata = config.get("metadata") + if not isinstance(metadata, dict): + return None + trace_id = metadata.get("trace_id") + if isinstance(trace_id, str) and trace_id.strip(): + return trace_id.strip() + except Exception: # Catch all: RuntimeError, AttributeError, TypeError, etc. + pass + return None + + +def resolve_trace_id_from_otel() -> str | None: + """Read trace_id from the active OTEL span if present.""" + try: + from opentelemetry import trace + + span = trace.get_current_span() + if span and span.get_span_context().is_valid: + trace_id = span.get_span_context().trace_id + if isinstance(trace_id, int) and trace_id > 0: + return format(trace_id, "032x") + except Exception: # Catch all: ImportError, AttributeError, TypeError, ValueError + pass + return None + + +def resolve_trace_id() -> str | None: + """Best-effort trace_id from context, config metadata, or OTEL.""" + ctx = _trace_id_var.get() + if ctx: + return ctx + from_config = resolve_trace_id_from_config() + if from_config: + return from_config + return resolve_trace_id_from_otel() diff --git a/deep_agent/src/audit/emitter.py b/deep_agent/src/audit/emitter.py new file mode 100644 index 00000000..d3214dd7 --- /dev/null +++ b/deep_agent/src/audit/emitter.py @@ -0,0 +1,181 @@ +"""Audit event emitter — structured JSON logging with local buffer fallback.""" + +from __future__ import annotations + +import json +import re +import sys +from datetime import UTC, datetime +from pathlib import Path +from typing import Any +from uuid import UUID + +from deep_agent.src.audit.buffer import drain, enqueue +from deep_agent.src.audit.config import is_audit_enabled +from deep_agent.src.audit.context import get_audit_context, resolve_trace_id +from deep_agent.utils.pylogger import SERVICE_NAME, get_python_logger + +logger = get_python_logger() + +# Sensitive keys to redact from audit details +SENSITIVE_KEYS = frozenset( + { + "password", + "token", + "apikey", + "api_key", + "secret", + "authorization", + "cookie", + "session", + "auth", + "credentials", + "privatekey", + "private_key", + "accesstoken", + "access_token", + "refreshtoken", + "refresh_token", + } +) + + +def _is_sensitive_key(key: str) -> bool: + """Return True if *key* names a sensitive field.""" + normalized = str(key).lower().replace("_", "").replace("-", "") + if normalized in SENSITIVE_KEYS: + return True + parts = [p for p in re.split(r"[_\-.]", str(key).lower()) if p] + return any(part in SENSITIVE_KEYS for part in parts) + + +def _scrub_details(details: dict[str, Any], depth: int = 0) -> dict[str, Any]: + """Recursively redact sensitive keys from audit details.""" + MAX_DEPTH = 5 + MAX_ARRAY_LEN = 100 + + if depth > MAX_DEPTH: + return {"error": "max_depth_exceeded"} + + scrubbed: dict[str, Any] = {} + for key, value in details.items(): + if _is_sensitive_key(key): + scrubbed[key] = "[REDACTED]" + elif isinstance(value, dict): + scrubbed[key] = _scrub_details(value, depth + 1) + elif isinstance(value, (list, tuple)): + # Limit array length and recursively scrub dicts + limited = list(value)[:MAX_ARRAY_LEN] + scrubbed[key] = [ + _scrub_details(v, depth + 1) if isinstance(v, dict) else v + for v in limited + ] + else: + scrubbed[key] = value + + return scrubbed + + +def emit_audit_event(audit_event_type: str, **details: Any) -> None: + """Emit a platform audit event. No-op when audit is disabled.""" + if not is_audit_enabled(): + return + + # Validate event type + if not isinstance(audit_event_type, str) or not audit_event_type.strip(): + logger.error("invalid_audit_event_type", type=type(audit_event_type).__name__) + return + + if len(audit_event_type) > 128: + logger.error("audit_event_type_too_long", length=len(audit_event_type)) + return + + ctx = get_audit_context() + envelope: dict[str, Any] = { + "event": "platform.audit", + "audit_event_type": audit_event_type.strip(), + "user": ctx.get("user"), + "org": ctx.get("org"), + "trace_id": ctx.get("trace_id") or resolve_trace_id(), + "timestamp": datetime.now(UTC).isoformat(), + "details": _scrub_details(details) if details else {}, + } + + _emit_envelope(envelope) + _flush_buffer() + + +def _format_record(envelope: dict[str, Any]) -> dict[str, Any]: + """Shape audit JSON to match other template-agent stdout log lines.""" + return { + **envelope, + "logger": "platform.audit", + "level": "info", + "service": SERVICE_NAME, + } + + +def _safe_json_default(obj: Any) -> str: + """Safe JSON serializer - only converts known safe types.""" + # Allow datetime/date conversion + if isinstance(obj, datetime): + return obj.isoformat() + if hasattr(obj, "isoformat"): # date, time, etc. + return str(obj.isoformat()) + # Allow Path and UUID + if isinstance(obj, (Path, UUID)): + return str(obj) + # Don't expose arbitrary objects - return type name only + return f"" + + +def _emit_envelope(envelope: dict[str, Any]) -> None: + MAX_SIZE = 1_000_000 # 1MB per event + + try: + line = json.dumps( + _format_record(envelope), default=_safe_json_default, ensure_ascii=False + ) + + if len(line) > MAX_SIZE: + logger.warning( + "audit_event_too_large", + size=len(line), + event_type=envelope.get("audit_event_type"), + ) + # Emit a truncated error event instead + error_envelope = { + **envelope, + "details": {"error": "event_too_large", "size": len(line)}, + } + line = json.dumps( + _format_record(error_envelope), default=_safe_json_default + ) + + sys.stdout.write(f"{line}\n") + sys.stdout.flush() + except Exception as exc: + logger.warning( + "audit_emit_failed", + error=str(exc), + error_type=type(exc).__name__, + event_type=envelope.get("audit_event_type"), + ) + enqueue(envelope) + + +def _flush_buffer() -> None: + """Retry buffered events. Stops on first failure to preserve order.""" + pending = drain() + for envelope in pending: + try: + line = json.dumps( + _format_record(envelope), default=_safe_json_default, ensure_ascii=False + ) + sys.stdout.write(f"{line}\n") + sys.stdout.flush() + except Exception as exc: + logger.debug("audit_flush_failed", error=str(exc), remaining=len(pending)) + # Re-enqueue this event and stop (preserves order) + enqueue(envelope) + break diff --git a/deep_agent/src/audit/events.py b/deep_agent/src/audit/events.py new file mode 100644 index 00000000..96b57cea --- /dev/null +++ b/deep_agent/src/audit/events.py @@ -0,0 +1,28 @@ +"""Audit event type constants. + +Orchestrator and subagents emit the same event types: + llm_call, mcp_tool_call, memory_write, subagent_delegation +""" + +from typing import Final + +LLM_CALL: Final = "llm_call" +MCP_TOOL_CALL: Final = "mcp_tool_call" +MEMORY_WRITE: Final = "memory_write" +SUBAGENT_DELEGATION: Final = "subagent_delegation" +CODE_EXECUTION: Final = "code_execution" + +# Event types audited via AuditMiddleware (orchestrator + in-process subagents). +AUDITED_MIDDLEWARE_EVENTS: frozenset[str] = frozenset( + {LLM_CALL, MCP_TOOL_CALL, MEMORY_WRITE, SUBAGENT_DELEGATION, CODE_EXECUTION} +) + + +class AuditEventType: + """Namespace for platform audit event type strings.""" + + LLM_CALL = LLM_CALL + MCP_TOOL_CALL = MCP_TOOL_CALL + MEMORY_WRITE = MEMORY_WRITE + SUBAGENT_DELEGATION = SUBAGENT_DELEGATION + CODE_EXECUTION = CODE_EXECUTION diff --git a/deep_agent/src/audit/middleware.py b/deep_agent/src/audit/middleware.py new file mode 100644 index 00000000..f66201ba --- /dev/null +++ b/deep_agent/src/audit/middleware.py @@ -0,0 +1,336 @@ +"""LangChain middleware for platform audit events. + +Orchestrator and in-process subagents use the same ``AuditMiddleware`` with +identical classification rules: + +- ``llm_call`` — every model invocation (sync + async paths) +- ``mcp_tool_call`` — tools in the subagent/orchestrator MCP tool name set +- ``memory_write`` — ``edit_file`` / ``write_file`` under ``/memories/`` (log only; no memory setup) +- ``subagent_delegation`` — ``task`` tool (orchestrator delegating to subagent) + +Events include ``agent`` (``orchestrator`` or subagent name). +""" + +from __future__ import annotations + +import time +from collections.abc import Awaitable, Callable +from typing import Any + +from langchain.agents.middleware.types import ( + AgentMiddleware, + ModelRequest, + ModelResponse, + ToolCallRequest, +) +from langchain_core.messages import ToolMessage +from langgraph.types import Command + +from deep_agent.src.audit.config import is_audit_enabled +from deep_agent.src.audit.emitter import emit_audit_event +from deep_agent.src.audit.events import AuditEventType + +_MEMORY_TOOLS = frozenset({"edit_file", "write_file"}) +_SUBAGENT_TOOL = "task" +_ORCHESTRATOR_AGENT = "orchestrator" + + +def _tool_path(args: dict[str, Any]) -> str: + for key in ("path", "file_path", "filename", "file"): + value = args.get(key) + if isinstance(value, str): + return value + return "" + + +def _is_memory_write(tool_name: str, args: dict[str, Any]) -> bool: + if tool_name not in _MEMORY_TOOLS: + return False + path = _tool_path(args) + return "memories" in path.replace("\\", "/") + + +def _model_name(request: ModelRequest[Any]) -> str: + model = request.model + if isinstance(model, str): + return model + return ( + getattr(model, "model_name", None) or getattr(model, "model", None) or "unknown" + ) + + +def classify_tool_call( + tool_name: str, + args: dict[str, Any], + *, + mcp_tool_names: frozenset[str], +) -> str: + """Classify a tool call using orchestrator/subagent parity rules.""" + if tool_name == _SUBAGENT_TOOL: + return AuditEventType.SUBAGENT_DELEGATION + if tool_name in mcp_tool_names: + return AuditEventType.MCP_TOOL_CALL + if _is_memory_write(tool_name, args): + return AuditEventType.MEMORY_WRITE + return "" + + +class AuditMiddleware(AgentMiddleware): + """Emit platform audit events for LLM and tool operations.""" + + def __init__( + self, + *, + mcp_tool_names: frozenset[str] | None = None, + subagent: str | None = None, + agent: str | None = None, + ) -> None: + """Initialize with optional MCP tool filter and agent identity.""" + self._mcp_tool_names = mcp_tool_names or frozenset() + self._agent = agent or subagent or _ORCHESTRATOR_AGENT + + def _base_details(self) -> dict[str, Any]: + return {"agent": self._agent} + + def _emit_llm_phase( + self, + *, + phase: str, + model: str, + message_count: int, + status: str | None = None, + latency_ms: float | None = None, + error: str | None = None, + ) -> None: + details: dict[str, Any] = { + "phase": phase, + "model": model, + **self._base_details(), + } + if phase == "start": + details["message_count"] = message_count + if status is not None: + details["status"] = status + if latency_ms is not None: + details["latency_ms"] = latency_ms + if error: + details["error"] = error + emit_audit_event(AuditEventType.LLM_CALL, **details) + + def _audit_model_call( + self, + request: ModelRequest[Any], + handler: Callable[[ModelRequest[Any]], ModelResponse[Any]], + ) -> ModelResponse[Any]: + if not is_audit_enabled(): + return handler(request) + + model = _model_name(request) + started = time.monotonic() + self._emit_llm_phase( + phase="start", + model=model, + message_count=len(request.messages), + ) + try: + response = handler(request) + except Exception as exc: + elapsed_ms = round((time.monotonic() - started) * 1000, 2) + self._emit_llm_phase( + phase="complete", + model=model, + message_count=len(request.messages), + status="error", + latency_ms=elapsed_ms, + error=str(exc) or type(exc).__name__, + ) + raise + + elapsed_ms = round((time.monotonic() - started) * 1000, 2) + self._emit_llm_phase( + phase="complete", + model=model, + message_count=len(request.messages), + status="success", + latency_ms=elapsed_ms, + ) + return response + + def wrap_model_call( + self, + request: ModelRequest[Any], + handler: Callable[[ModelRequest[Any]], ModelResponse[Any]], + ) -> ModelResponse[Any]: + """Sync model hook — subagents use ``Runnable.invoke()``.""" + return self._audit_model_call(request, handler) + + async def awrap_model_call( + self, + request: ModelRequest[Any], + handler: Callable[[ModelRequest[Any]], Awaitable[ModelResponse[Any]]], + ) -> ModelResponse[Any]: + """Async wrapper that audits LLM model invocations.""" + if not is_audit_enabled(): + return await handler(request) + + model = _model_name(request) + started = time.monotonic() + self._emit_llm_phase( + phase="start", + model=model, + message_count=len(request.messages), + ) + try: + response = await handler(request) + except Exception as exc: + elapsed_ms = round((time.monotonic() - started) * 1000, 2) + self._emit_llm_phase( + phase="complete", + model=model, + message_count=len(request.messages), + status="error", + latency_ms=elapsed_ms, + error=str(exc) or type(exc).__name__, + ) + raise + + elapsed_ms = round((time.monotonic() - started) * 1000, 2) + self._emit_llm_phase( + phase="complete", + model=model, + message_count=len(request.messages), + status="success", + latency_ms=elapsed_ms, + ) + return response + + def _classify_tool(self, tool_name: str, args: dict[str, Any]) -> str: + return classify_tool_call(tool_name, args, mcp_tool_names=self._mcp_tool_names) + + def _emit_tool_event( + self, + audit_type: str, + *, + tool_name: str, + tool_args: dict[str, Any], + status: str, + latency_ms: float, + error: str | None = None, + ) -> None: + if not audit_type: + return + + details: dict[str, Any] = { + "tool": tool_name, + "status": status, + "latency_ms": latency_ms, + **self._base_details(), + } + if error: + details["error"] = error + + if audit_type == AuditEventType.SUBAGENT_DELEGATION: + details["delegated_subagent"] = tool_args.get("subagent") or tool_args.get( + "name" + ) + elif audit_type == AuditEventType.MEMORY_WRITE: + details["path"] = _tool_path(tool_args) + elif audit_type == AuditEventType.MCP_TOOL_CALL: + details["args_keys"] = sorted(tool_args.keys()) + + emit_audit_event(audit_type, **details) + + def _audit_tool_call( + self, + request: ToolCallRequest, + handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]], + ) -> ToolMessage | Command[Any]: + if not is_audit_enabled(): + return handler(request) + + tool_call = request.tool_call + tool_name = tool_call.get("name", "unknown") + tool_args = tool_call.get("args") + if not isinstance(tool_args, dict): + tool_args = {} + audit_type = self._classify_tool(tool_name, tool_args) + + started = time.monotonic() + try: + result = handler(request) + status = "success" + error: str | None = None + except Exception as exc: + elapsed_ms = round((time.monotonic() - started) * 1000, 2) + self._emit_tool_event( + audit_type, + tool_name=tool_name, + tool_args=tool_args, + status="error", + latency_ms=elapsed_ms, + error=str(exc) or type(exc).__name__, + ) + raise + + elapsed_ms = round((time.monotonic() - started) * 1000, 2) + self._emit_tool_event( + audit_type, + tool_name=tool_name, + tool_args=tool_args, + status=status, + latency_ms=elapsed_ms, + error=error, + ) + return result + + def wrap_tool_call( + self, + request: ToolCallRequest, + handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]], + ) -> ToolMessage | Command[Any]: + """Sync tool hook — subagents use ``Runnable.invoke()``.""" + return self._audit_tool_call(request, handler) + + async def awrap_tool_call( + self, + request: ToolCallRequest, + handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]], + ) -> ToolMessage | Command[Any]: + """Async wrapper that audits tool invocations.""" + if not is_audit_enabled(): + return await handler(request) + + tool_call = request.tool_call + tool_name = tool_call.get("name", "unknown") + tool_args = tool_call.get("args") + if not isinstance(tool_args, dict): + tool_args = {} + audit_type = self._classify_tool(tool_name, tool_args) + + started = time.monotonic() + try: + result = await handler(request) + status = "success" + error: str | None = None + except Exception as exc: + elapsed_ms = round((time.monotonic() - started) * 1000, 2) + self._emit_tool_event( + audit_type, + tool_name=tool_name, + tool_args=tool_args, + status="error", + latency_ms=elapsed_ms, + error=str(exc) or type(exc).__name__, + ) + raise + + elapsed_ms = round((time.monotonic() - started) * 1000, 2) + self._emit_tool_event( + audit_type, + tool_name=tool_name, + tool_args=tool_args, + status=status, + latency_ms=elapsed_ms, + error=error, + ) + return result diff --git a/deep_agent/src/cache/__init__.py b/deep_agent/src/cache/__init__.py new file mode 100644 index 00000000..4026c396 --- /dev/null +++ b/deep_agent/src/cache/__init__.py @@ -0,0 +1,19 @@ +"""Multi-layer caching for the template agent. + +All cache layers are **disabled by default** and activated via +environment variables. Set ``CACHE_ENABLED=true`` plus individual +layer flags (``CACHE_MODEL_ENABLED``, ``CACHE_PERSONALIZATION_ENABLED``, +etc.) to opt in. + +Exports: + cache_settings: Configuration singleton (feature flags + TTLs) + get_or_create_model: Cached LLM model factory + warm_caches: Startup cache warming + metrics: Hit/miss/set counters +""" + +from deep_agent.src.cache.config import cache_settings + +__all__ = [ + "cache_settings", +] diff --git a/deep_agent/src/cache/backend.py b/deep_agent/src/cache/backend.py new file mode 100644 index 00000000..30d37445 --- /dev/null +++ b/deep_agent/src/cache/backend.py @@ -0,0 +1,197 @@ +"""Cache backend implementations. + +Provides a ``CacheBackend`` protocol and three implementations: + +- ``NullCache``: No-op (returned when caching is disabled) +- ``InMemoryCache``: Process-local TTLCache via ``cachetools`` +- ``RedisCache``: Shared cache via the existing ``aegra.redis`` client +""" + +import threading +from typing import Any, Protocol, runtime_checkable + +from cachetools import TTLCache + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +@runtime_checkable +class CacheBackend(Protocol): + """Minimal cache interface — get/set/delete/clear with string values.""" + + def get(self, key: str) -> str | None: + """Retrieve a cached value by key, or None on miss.""" + ... + + def set(self, key: str, value: str, ttl: int | None = None) -> bool: + """Store a value; return True on success.""" + ... + + def delete(self, key: str) -> bool: + """Remove a key; return True if it existed.""" + ... + + def clear(self) -> None: + """Remove all entries.""" + ... + + @property + def name(self) -> str: + """Human-readable backend name.""" + ... + + +class NullCache: + """No-op cache — every operation is a silent miss.""" + + @property + def name(self) -> str: + """Return backend name.""" + return "null" + + def get(self, key: str) -> str | None: + """Always return None.""" + return None + + def set(self, key: str, value: str, ttl: int | None = None) -> bool: + """Always return False (nothing stored).""" + return False + + def delete(self, key: str) -> bool: + """Always return False (nothing to delete).""" + return False + + def clear(self) -> None: + """No-op.""" + + +class InMemoryCache: + """Process-local TTL cache backed by ``cachetools.TTLCache``. + + Thread-safe via an internal lock. + + Args: + max_size: Maximum number of entries. + default_ttl: Default time-to-live in seconds. + """ + + def __init__(self, max_size: int = 256, default_ttl: int = 300) -> None: + """Initialise with capacity and TTL.""" + self._cache: TTLCache[str, str] = TTLCache(maxsize=max_size, ttl=default_ttl) + self._default_ttl = default_ttl + self._lock = threading.Lock() + + @property + def name(self) -> str: + """Return backend name.""" + return "memory" + + def get(self, key: str) -> str | None: + """Look up *key* in the TTL cache.""" + with self._lock: + result: str | None = self._cache.get(key) + return result + + def set(self, key: str, value: str, ttl: int | None = None) -> bool: + """Insert or overwrite *key*.""" + with self._lock: + self._cache[key] = value + return True + + def delete(self, key: str) -> bool: + """Remove *key* if present.""" + with self._lock: + try: + del self._cache[key] + return True + except KeyError: + return False + + def clear(self) -> None: + """Remove all entries.""" + with self._lock: + self._cache.clear() + + @property + def size(self) -> int: + """Current number of entries.""" + with self._lock: + return len(self._cache) + + +class RedisCache: + """Shared cache via the existing ``aegra.redis`` client. + + Falls back to no-op if Redis is unavailable — never raises. + + Args: + default_ttl: Default TTL in seconds. + key_prefix: Prefix prepended to all keys (namespacing). + """ + + def __init__(self, default_ttl: int = 300, key_prefix: str = "cache:") -> None: + """Initialise with TTL and key prefix.""" + self._default_ttl = default_ttl + self._prefix = key_prefix + self._client: Any = None + self._checked = False + + def _get_client(self) -> Any: + if not self._checked: + try: + from deep_agent.aegra.redis import get_redis_client + + self._client = get_redis_client() + except Exception: + logger.debug("Redis unavailable for cache layer", exc_info=True) + self._client = None + self._checked = True + return self._client + + @property + def name(self) -> str: + """Return backend name.""" + return "redis" + + def _key(self, key: str) -> str: + return f"{self._prefix}{key}" + + def get(self, key: str) -> str | None: + """Read from Redis; return None on miss or error.""" + client = self._get_client() + if client is None: + return None + try: + result: str | None = client.get(self._key(key)) + return result + except Exception: + logger.debug("Redis cache GET failed for '%s'", key, exc_info=True) + return None + + def set(self, key: str, value: str, ttl: int | None = None) -> bool: + """Write to Redis with TTL.""" + client = self._get_client() + if client is None: + return False + try: + client.setex(self._key(key), ttl or self._default_ttl, value) + return True + except Exception: + logger.debug("Redis cache SET failed for '%s'", key, exc_info=True) + return False + + def delete(self, key: str) -> bool: + """Delete from Redis.""" + client = self._get_client() + if client is None: + return False + try: + client.delete(self._key(key)) + return True + except Exception: + return False + + def clear(self) -> None: + """Clear is not supported for Redis (too dangerous). No-op.""" diff --git a/deep_agent/src/cache/config.py b/deep_agent/src/cache/config.py new file mode 100644 index 00000000..751a7462 --- /dev/null +++ b/deep_agent/src/cache/config.py @@ -0,0 +1,55 @@ +"""Cache configuration with feature flags. + +Every cache layer is disabled by default. Enable via environment +variables — the master ``CACHE_ENABLED`` switch must be ``true`` +for any individual cache to activate. + +Environment variables: + CACHE_ENABLED: Master switch (default: false) + CACHE_MODEL_ENABLED: LLM model instance cache (default: false) + CACHE_MODEL_TTL: Model cache TTL in seconds (default: 600) + CACHE_MODEL_MAX_SIZE: Max cached model instances (default: 10) + CACHE_PERSONALIZATION_ENABLED: User personalization cache (default: false) + CACHE_PERSONALIZATION_TTL: Personalization TTL in seconds (default: 120) + CACHE_METRICS_ENABLED: Log cache hit/miss counters (default: false) + CACHE_WARMING_ENABLED: Pre-create models at startup (default: false) + CACHE_REDIS_ENABLED: Enable Redis as L2 cache layer (default: false) + +Note: + MCP tool cache TTL and compiled graph cache TTL are configured via + config/agent/runtime/agent.yaml (cache.mcp.ttl, cache.graph.ttl), + NOT via environment variables. +""" + +from pydantic import Field +from pydantic_settings import BaseSettings + + +class CacheSettings(BaseSettings): + """Feature-flagged cache configuration loaded from environment.""" + + CACHE_ENABLED: bool = Field(default=False) + + CACHE_MODEL_ENABLED: bool = Field(default=False) + CACHE_MODEL_TTL: int = Field(default=600, ge=10, le=7200) + CACHE_MODEL_MAX_SIZE: int = Field(default=10, ge=1, le=100) + + CACHE_PERSONALIZATION_ENABLED: bool = Field(default=False) + CACHE_PERSONALIZATION_TTL: int = Field(default=120, ge=10, le=3600) + + CACHE_METRICS_ENABLED: bool = Field(default=False) + CACHE_WARMING_ENABLED: bool = Field(default=False) + CACHE_REDIS_ENABLED: bool = Field(default=False) + + def is_enabled(self, layer: str) -> bool: + """Check if a specific cache layer is active. + + Both the master switch and the layer-specific flag must be true. + """ + if not self.CACHE_ENABLED: + return False + flag = getattr(self, f"CACHE_{layer.upper()}_ENABLED", False) + return bool(flag) + + +cache_settings = CacheSettings() diff --git a/deep_agent/src/cache/metrics.py b/deep_agent/src/cache/metrics.py new file mode 100644 index 00000000..d401b19d --- /dev/null +++ b/deep_agent/src/cache/metrics.py @@ -0,0 +1,101 @@ +"""Cache metrics — hit/miss/eviction counters per cache name. + +Counters are in-memory per process. When ``CACHE_METRICS_ENABLED`` +is true, periodic summaries are logged at INFO level. +""" + +import threading +from typing import Any + +from deep_agent.src.cache.config import cache_settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_lock = threading.Lock() +_counters: dict[str, dict[str, int]] = {} + + +def _ensure(name: str) -> dict[str, int]: + if name not in _counters: + _counters[name] = {"hits": 0, "misses": 0, "sets": 0, "deletes": 0} + return _counters[name] + + +def record_hit(cache_name: str) -> None: + """Increment hit counter for *cache_name*.""" + if not cache_settings.is_enabled("metrics"): + return + with _lock: + _ensure(cache_name)["hits"] += 1 + + +def record_miss(cache_name: str) -> None: + """Increment miss counter for *cache_name*.""" + if not cache_settings.is_enabled("metrics"): + return + with _lock: + _ensure(cache_name)["misses"] += 1 + + +def record_set(cache_name: str) -> None: + """Increment set counter for *cache_name*.""" + if not cache_settings.is_enabled("metrics"): + return + with _lock: + _ensure(cache_name)["sets"] += 1 + + +def record_delete(cache_name: str) -> None: + """Increment delete counter for *cache_name*.""" + if not cache_settings.is_enabled("metrics"): + return + with _lock: + _ensure(cache_name)["deletes"] += 1 + + +def snapshot() -> dict[str, dict[str, int]]: + """Return a copy of all counters.""" + with _lock: + return {k: dict(v) for k, v in _counters.items()} + + +def reset() -> None: + """Clear all counters.""" + with _lock: + _counters.clear() + + +def log_summary() -> None: + """Log current counters at INFO level.""" + if not cache_settings.is_enabled("metrics"): + return + stats = snapshot() + if not stats: + return + for name, counts in stats.items(): + total = counts["hits"] + counts["misses"] + rate = (counts["hits"] / total * 100) if total > 0 else 0.0 + logger.info( + "Cache '%s': %d hits, %d misses (%.1f%% hit rate), %d sets, %d deletes", + name, + counts["hits"], + counts["misses"], + rate, + counts["sets"], + counts["deletes"], + ) + + +def get_stats() -> dict[str, Any]: + """Return metrics as a JSON-serialisable dict (for /health or /metrics).""" + stats = snapshot() + result: dict[str, Any] = {} + for name, counts in stats.items(): + total = counts["hits"] + counts["misses"] + result[name] = { + **counts, + "total": total, + "hit_rate": round(counts["hits"] / total * 100, 1) if total > 0 else 0.0, + } + return result diff --git a/deep_agent/src/cache/model_cache.py b/deep_agent/src/cache/model_cache.py new file mode 100644 index 00000000..c51f2b35 --- /dev/null +++ b/deep_agent/src/cache/model_cache.py @@ -0,0 +1,185 @@ +"""LLM model instance cache. + +Caches ``BaseChatModel`` instances by ``(model_name, temperature, +max_output_tokens)`` or by ``(spec_cache_key, temperature, tokens)`` +for provider-aware specs so repeated per-request calls reuse the same +client handle. + +Model instances are **not** serialisable, so this is L1 (in-memory) +only — no Redis layer. + +**Memory usage**: Two separate caches exist (legacy string-based and spec-based), +each limited to ``CACHE_MODEL_MAX_SIZE`` entries. Maximum total memory usage is +2x the configured limit (e.g., if limit is 100, up to 200 models may be cached). + +Feature flag: ``CACHE_MODEL_ENABLED`` (+ master ``CACHE_ENABLED``). +""" + +import threading + +from cachetools import TTLCache + +from deep_agent.src.agent.config.model import ModelSpec, model_spec_cache_key +from deep_agent.src.cache import metrics +from deep_agent.src.cache.config import cache_settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_LegacyCacheKey = tuple[str, float, int] +_SpecCacheKey = tuple[str, float, int] # (cache_id, temperature, tokens) + +_lock = threading.Lock() +_legacy_cache: TTLCache[_LegacyCacheKey, object] | None = None +_spec_cache: TTLCache[_SpecCacheKey, object] | None = None + + +def _get_legacy_cache() -> TTLCache[_LegacyCacheKey, object]: + global _legacy_cache # noqa: PLW0603 + if _legacy_cache is None: + _legacy_cache = TTLCache( + maxsize=cache_settings.CACHE_MODEL_MAX_SIZE, + ttl=cache_settings.CACHE_MODEL_TTL, + ) + return _legacy_cache + + +def _get_spec_cache() -> TTLCache[_SpecCacheKey, object]: + global _spec_cache # noqa: PLW0603 + if _spec_cache is None: + _spec_cache = TTLCache( + maxsize=cache_settings.CACHE_MODEL_MAX_SIZE, + ttl=cache_settings.CACHE_MODEL_TTL, + ) + return _spec_cache + + +def _get_cache() -> TTLCache[_LegacyCacheKey, object]: + """Backward-compatible alias for legacy cache getter. + + Used by both tests and production code that still uses string-based model names. + """ + return _get_legacy_cache() + + +def get_or_create_model( + model_name: str, + temperature: float = 0.0, + max_output_tokens: int | None = None, +) -> object: + """Return a cached model or create a new one. + + When the cache is disabled (flag off), this is a straight + passthrough to ``create_model()``. + + Returns: + A ``BaseChatModel`` instance. + """ + from deep_agent.src.agent.llm import create_model + from deep_agent.src.settings import settings + + tokens = max_output_tokens or settings.MAX_OUTPUT_TOKENS + + if not cache_settings.is_enabled("model"): + return create_model(model_name, temperature, tokens) + + key: _LegacyCacheKey = (model_name, temperature, tokens) + + with _lock: + cache = _get_legacy_cache() + model = cache.get(key) + if model is not None: + metrics.record_hit("model") + logger.debug("Model cache HIT: %s", model_name) + return model + + metrics.record_miss("model") + logger.debug("Model cache MISS: %s — creating", model_name) + model = create_model(model_name, temperature, tokens) + + with _lock: + cache = _get_legacy_cache() + cache[key] = model + metrics.record_set("model") + + return model + + +def get_or_create_model_from_spec( + spec: ModelSpec, + temperature: float = 0.0, + max_output_tokens: int | None = None, +) -> object: + """Return a cached model for a :class:`ModelSpec` or create a new one. + + Cache key includes provider, model name, and fallback chain so + different provider configurations never collide. + + Returns: + A ``BaseChatModel`` instance. + """ + from deep_agent.src.agent.provider_factory import create_model_from_spec + from deep_agent.src.settings import settings + + tokens = max_output_tokens or settings.MAX_OUTPUT_TOKENS + cache_id = model_spec_cache_key(spec) + + if not cache_settings.is_enabled("model"): + return create_model_from_spec( + spec, temperature=temperature, max_output_tokens=tokens + ) + + key: _SpecCacheKey = (cache_id, temperature, tokens) + + with _lock: + cache = _get_spec_cache() + model = cache.get(key) + if model is not None: + metrics.record_hit("model") + logger.debug("Model cache HIT: %s", cache_id) + return model + + metrics.record_miss("model") + logger.debug("Model cache MISS: %s — creating", cache_id) + model = create_model_from_spec( + spec, temperature=temperature, max_output_tokens=tokens + ) + + with _lock: + cache = _get_spec_cache() + cache[key] = model + metrics.record_set("model") + + return model + + +def invalidate(model_name: str | None = None) -> None: + """Drop cached model(s). + + Args: + model_name: If given, remove only legacy-cache entries for this model. + If None, clear both legacy and spec caches. + """ + with _lock: + legacy = _get_legacy_cache() + spec = _get_spec_cache() + if model_name is None: + legacy.clear() + spec.clear() + logger.info("Model cache cleared") + return + keys_to_remove = [k for k in legacy if k[0] == model_name] + for k in keys_to_remove: + del legacy[k] + if keys_to_remove: + logger.info( + "Model cache: evicted %d legacy entry(s) for '%s'", + len(keys_to_remove), + model_name, + ) + + +def cached_count() -> int: + """Return the number of currently cached models (legacy + spec).""" + with _lock: + return len(_get_legacy_cache()) + len(_get_spec_cache()) diff --git a/deep_agent/src/cache/multi_layer.py b/deep_agent/src/cache/multi_layer.py new file mode 100644 index 00000000..1d18b47d --- /dev/null +++ b/deep_agent/src/cache/multi_layer.py @@ -0,0 +1,86 @@ +"""Two-layer cache: L1 in-process memory + L2 shared Redis. + +On ``get``: + L1 hit → return immediately + L1 miss → check L2 → backfill L1 on hit + +On ``set``: + Write to both L1 and L2 + +On ``delete``: + Delete from both L1 and L2 +""" + +from deep_agent.src.cache import metrics +from deep_agent.src.cache.backend import CacheBackend, NullCache +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +class MultiLayerCache: + """Composite cache with L1 (fast/local) and optional L2 (shared/Redis). + + Args: + name: Human-readable name used in metrics and logging. + l1: Primary (fast) cache backend. + l2: Secondary (shared) cache backend, or None to skip. + """ + + def __init__( + self, + name: str, + l1: CacheBackend, + l2: CacheBackend | None = None, + ) -> None: + """Initialise with a name and one or two backend layers.""" + self._name = name + self._l1 = l1 + self._l2 = l2 + + @property + def name(self) -> str: + """Human-readable cache name used in metrics.""" + return self._name + + def get(self, key: str) -> str | None: + """Look up *key* in L1, then L2. Backfills L1 on L2 hit.""" + value = self._l1.get(key) + if value is not None: + metrics.record_hit(self._name) + return value + + if self._l2 is not None: + value = self._l2.get(key) + if value is not None: + self._l1.set(key, value) + metrics.record_hit(self._name) + return value + + metrics.record_miss(self._name) + return None + + def set(self, key: str, value: str, ttl: int | None = None) -> bool: + """Write *value* to L1 and L2.""" + metrics.record_set(self._name) + ok = self._l1.set(key, value, ttl) + if self._l2 is not None: + self._l2.set(key, value, ttl) + return ok + + def delete(self, key: str) -> bool: + """Remove *key* from both layers.""" + metrics.record_delete(self._name) + ok = self._l1.delete(key) + if self._l2 is not None: + self._l2.delete(key) + return ok + + def clear(self) -> None: + """Clear L1. L2 clear is intentionally a no-op (safety).""" + self._l1.clear() + + +def create_null_layer(name: str) -> MultiLayerCache: + """Return a no-op MultiLayerCache (used when caching is disabled).""" + return MultiLayerCache(name=name, l1=NullCache()) diff --git a/deep_agent/src/cache/personalization_cache.py b/deep_agent/src/cache/personalization_cache.py new file mode 100644 index 00000000..3b152d9e --- /dev/null +++ b/deep_agent/src/cache/personalization_cache.py @@ -0,0 +1,98 @@ +"""Personalization cache — Redis L2 for user memories and rules. + +Avoids hitting Postgres on every request for the same user's +personalization data. Stores serialised JSON in Redis, keyed by +``user_id``. + +Feature flag: ``CACHE_PERSONALIZATION_ENABLED`` (+ ``CACHE_ENABLED``). +""" + +import json +from typing import Any + +from deep_agent.src.cache import metrics +from deep_agent.src.cache.backend import RedisCache +from deep_agent.src.cache.config import cache_settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_KEY_PREFIX = "personalization:" + +_redis: RedisCache | None = None + + +def _get_redis() -> RedisCache: + global _redis # noqa: PLW0603 + if _redis is None: + _redis = RedisCache( + default_ttl=cache_settings.CACHE_PERSONALIZATION_TTL, + key_prefix=_KEY_PREFIX, + ) + return _redis + + +def _cache_key(user_id: str) -> str: + return f"user:{user_id}" + + +async def get_personalization( + user_id: str, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] | None: + """Return cached ``(memories, rules)`` dicts or None on miss. + + When disabled, always returns None (caller falls through to DB). + """ + if not cache_settings.is_enabled("personalization"): + return None + + raw = _get_redis().get(_cache_key(user_id)) + if raw is None: + metrics.record_miss("personalization") + return None + + try: + data = json.loads(raw) + metrics.record_hit("personalization") + logger.debug("Personalization cache HIT for user %s", user_id[:8]) + return data["memories"], data["rules"] + except (json.JSONDecodeError, KeyError): + logger.debug( + "Personalization cache corrupt for user %s — evicting", user_id[:8] + ) + _get_redis().delete(_cache_key(user_id)) + metrics.record_miss("personalization") + return None + + +async def set_personalization( + user_id: str, + memories: list[dict[str, Any]], + rules: list[dict[str, Any]], +) -> None: + """Store personalization data in Redis cache.""" + if not cache_settings.is_enabled("personalization"): + return + + payload = json.dumps({"memories": memories, "rules": rules}) + _get_redis().set(_cache_key(user_id), payload) + metrics.record_set("personalization") + logger.debug( + "Personalization cached for user %s (%d memories, %d rules)", + user_id[:8], + len(memories), + len(rules), + ) + + +async def invalidate(user_id: str | None = None) -> None: + """Evict cached personalization for a user. + + Args: + user_id: Specific user to evict. ``None`` is a no-op + (clearing all Redis keys is too dangerous). + """ + if user_id is None: + return + _get_redis().delete(_cache_key(user_id)) + metrics.record_delete("personalization") diff --git a/deep_agent/src/cache/warming.py b/deep_agent/src/cache/warming.py new file mode 100644 index 00000000..42404110 --- /dev/null +++ b/deep_agent/src/cache/warming.py @@ -0,0 +1,67 @@ +"""Cache warming — pre-populate caches at startup. + +When ``CACHE_WARMING_ENABLED`` is true, ``warm_caches()`` pre-creates +the default orchestrator and subagent LLM model instances so the first +user request doesn't pay the cold-start penalty. + +Feature flag: ``CACHE_WARMING_ENABLED`` (+ ``CACHE_ENABLED``). +""" + +from deep_agent.src.cache.config import cache_settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def warm_caches() -> dict[str, bool]: + """Pre-populate caches. Returns a status dict per cache layer. + + Safe to call even when caching is disabled — returns immediately. + """ + results: dict[str, bool] = {} + + if not cache_settings.is_enabled("warming"): + logger.debug("Cache warming disabled — skipping") + return results + + logger.info("Warming caches...") + + results["models"] = _warm_models() + + logger.info("Cache warming complete: %s", results) + return results + + +def _warm_models() -> bool: + """Pre-create LLM model instances for orchestrator + subagents.""" + if not cache_settings.is_enabled("model"): + return False + + try: + from deep_agent.src.agent.config import agent_config + from deep_agent.src.agent.config.model import parse_model_config + from deep_agent.src.cache.model_cache import get_or_create_model_from_spec + + orch = agent_config.get_orchestrator_config() + orch_model = orch.get("model", "gemini-3.1-pro-preview") + + # Parse orchestrator model to ModelSpec (supports provider) + orch_spec = parse_model_config(orch_model) + get_or_create_model_from_spec(orch_spec) + logger.info( + "Warmed orchestrator model: %s (provider: %s)", + orch_spec.name, + orch_spec.provider.value, + ) + + for name, cfg in agent_config.get_all_subagent_configs().items(): + sub_model = cfg.get("model") + if sub_model: + spec = parse_model_config(sub_model) + get_or_create_model_from_spec(spec) + logger.info("Warmed subagent '%s' model: %s", name, spec.display_name()) + + return True + except Exception: + logger.warning("Model cache warming failed", exc_info=True) + return False diff --git a/deep_agent/src/code_execution/__init__.py b/deep_agent/src/code_execution/__init__.py new file mode 100644 index 00000000..81e1edbf --- /dev/null +++ b/deep_agent/src/code_execution/__init__.py @@ -0,0 +1,8 @@ +"""Code execution middleware — ephemeral K8s Job backend for agent code execution.""" + +from __future__ import annotations + +from deep_agent.src.code_execution.config import CodeExecutionConfig +from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + +__all__ = ["CodeExecutionConfig", "CodeExecutionMiddleware"] diff --git a/deep_agent/src/code_execution/config.py b/deep_agent/src/code_execution/config.py new file mode 100644 index 00000000..67651ac5 --- /dev/null +++ b/deep_agent/src/code_execution/config.py @@ -0,0 +1,56 @@ +"""Configuration model for code execution middleware.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, Field + + +class CodeExecutionConfig(BaseModel): + """Configuration for code execution middleware.""" + + enabled: bool = False + max_timeout_seconds: int = Field(default=60, ge=5, le=300) + max_code_length: int = Field(default=50_000, ge=100, le=500_000) + max_output_bytes: int = Field(default=1_048_576) + + images: dict[str, str] = Field( + default_factory=lambda: { + "python": "python:3.12-slim", + "shell": "bash:5", + "node": "node:22-slim", + } + ) + + entrypoints: dict[str, list[str]] = Field( + default_factory=lambda: { + "python": ["python", "-c"], + "shell": ["bash", "-c"], + "node": ["node", "-e"], + } + ) + + resource_requests: dict[str, str] = Field( + default_factory=lambda: {"cpu": "100m", "memory": "128Mi"} + ) + resource_limits: dict[str, str] = Field( + default_factory=lambda: {"cpu": "500m", "memory": "256Mi"} + ) + + tmp_size_limit: str = "64Mi" + job_ttl_after_finished: int = Field(default=30, ge=0, le=300) + pod_poll_interval_seconds: float = Field(default=1.0, ge=0.5, le=10.0) + pod_poll_timeout_seconds: float = Field(default=120.0, ge=10.0, le=600.0) + + network_access: Literal["deny", "allow_internet", "per_execution"] = "deny" + max_concurrent_per_org: int = Field(default=3, ge=1, le=20) + queue_timeout_seconds: float = Field(default=30.0, ge=1.0, le=120.0) + max_input_file_size: int = Field(default=1_048_576) + cost_tracking_enabled: bool = False + streaming_enabled: bool = False + + @property + def supported_languages(self) -> set[str]: + """Return supported language names.""" + return set(self.images.keys()) & set(self.entrypoints.keys()) diff --git a/deep_agent/src/code_execution/k8s_job_runner.py b/deep_agent/src/code_execution/k8s_job_runner.py new file mode 100644 index 00000000..647e36f8 --- /dev/null +++ b/deep_agent/src/code_execution/k8s_job_runner.py @@ -0,0 +1,725 @@ +"""K8s Job lifecycle manager for ephemeral code execution.""" + +from __future__ import annotations + +import asyncio +import logging +import os +import time +import uuid +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from deep_agent.src.code_execution.config import CodeExecutionConfig +from deep_agent.src.code_execution.metrics import _log_json + + +@dataclass +class ExecutionResult: + """Result of a code execution job.""" + + stdout: str + stderr: str + exit_code: int + duration_seconds: float + status: str + job_name: str + namespace: str + output_files: dict[str, str] = field(default_factory=dict) + cpu_seconds: float = 0.0 + memory_mb_seconds: float = 0.0 + scheduling_seconds: float = 0.0 + + def format(self) -> str: + """Format execution result as a human-readable string.""" + parts = [] + if self.stdout: + parts.append(f"stdout:\n{self.stdout}") + if self.stderr: + parts.append(f"stderr:\n{self.stderr}") + parts.append(f"exit_code: {self.exit_code}") + if self.status == "timeout": + parts.append(f"(timed out after {self.duration_seconds:.1f}s)") + if self.status == "oom_killed": + parts.append("(killed: out of memory)") + return "\n".join(parts) + + +class K8sJobRunner: + """Manages the lifecycle of ephemeral K8s Jobs for code execution.""" + + def __init__(self, config: CodeExecutionConfig) -> None: + """Initialize runner with execution configuration.""" + self._config = config + self._batch_api: Any | None = None + self._core_api: Any | None = None + self._networking_api: Any | None = None + self._custom_api: Any | None = None + + def _ensure_k8s_client(self) -> None: + """Lazy-initialize K8s API clients.""" + if self._batch_api is not None: + return + try: + from kubernetes import client + from kubernetes import config as k8s_config + + try: + k8s_config.load_incluster_config() + except k8s_config.ConfigException: + k8s_config.load_kube_config() + self._batch_api = client.BatchV1Api() + self._core_api = client.CoreV1Api() + self._networking_api = client.NetworkingV1Api() + self._custom_api = client.CustomObjectsApi() + except ImportError: + raise RuntimeError( + "kubernetes package required for code execution. " + "Install with: pip install kubernetes" + ) + + def build_job_manifest( + self, + *, + language: str, + code: str, + timeout: int, + namespace: str, + execution_id: str, + trace_id: str | None = None, + allow_network: bool = False, + input_configmap_name: str | None = None, + ) -> Any: + """Build a K8s Job manifest for code execution.""" + from kubernetes import client + + labels = { + "app.kubernetes.io/name": "code-execution", + "app.kubernetes.io/component": "ephemeral-job", + "app.kubernetes.io/managed-by": "template-agent", + "ai-platform.io/execution-id": execution_id, + } + if allow_network: + labels["ai-platform.io/allow-internet"] = "true" + + annotations = {} + if trace_id: + annotations["ai-platform.io/trace-id"] = trace_id + + volume_mounts = [ + client.V1VolumeMount(name="tmp", mount_path="/tmp"), + client.V1VolumeMount(name="output", mount_path="/output"), + ] + volumes = [ + client.V1Volume( + name="tmp", + empty_dir=client.V1EmptyDirVolumeSource( + size_limit=self._config.tmp_size_limit, + ), + ), + client.V1Volume( + name="output", + empty_dir=client.V1EmptyDirVolumeSource(size_limit="64Mi"), + ), + ] + + if input_configmap_name: + volume_mounts.append( + client.V1VolumeMount(name="input", mount_path="/input", read_only=True) + ) + volumes.append( + client.V1Volume( + name="input", + config_map=client.V1ConfigMapVolumeSource( + name=input_configmap_name + ), + ) + ) + + container = client.V1Container( + name="executor", + image=self._config.images[language], + command=self._config.entrypoints[language], + args=[code], + resources=client.V1ResourceRequirements( + requests=dict(self._config.resource_requests), + limits=dict(self._config.resource_limits), + ), + security_context=client.V1SecurityContext( + allow_privilege_escalation=False, + read_only_root_filesystem=True, + capabilities=client.V1Capabilities(drop=["ALL"]), + ), + volume_mounts=volume_mounts, + ) + + pod_spec = client.V1PodSpec( + restart_policy="Never", + automount_service_account_token=False, + security_context=client.V1PodSecurityContext( + run_as_non_root=True, + run_as_user=1000, + run_as_group=1000, + fs_group=1000, + seccomp_profile=client.V1SeccompProfile(type="RuntimeDefault"), + ), + containers=[container], + volumes=volumes, + ) + + job_name = f"code-exec-{execution_id[:8]}" + + return client.V1Job( + api_version="batch/v1", + kind="Job", + metadata=client.V1ObjectMeta( + name=job_name, + namespace=namespace, + labels=labels, + annotations=annotations or None, + ), + spec=client.V1JobSpec( + active_deadline_seconds=timeout, + ttl_seconds_after_finished=self._config.job_ttl_after_finished, + backoff_limit=0, + template=client.V1PodTemplateSpec( + metadata=client.V1ObjectMeta( + labels=labels, + annotations=annotations or None, + ), + spec=pod_spec, + ), + ), + ) + + def parse_container_status( + self, *, exit_code: int, termination_reason: str | None + ) -> tuple[int, str]: + """Parse container exit code and reason into a status tuple.""" + if termination_reason == "OOMKilled": + return exit_code, "oom_killed" + if termination_reason == "DeadlineExceeded": + return exit_code, "timeout" + if exit_code == 0: + return exit_code, "success" + return exit_code, "failed" + + def resolve_namespace(self) -> str: + """Resolve the K8s namespace from environment variables.""" + org = os.environ.get("AI_PLATFORM_AGENT_ORG", "default") + agent = os.environ.get("AI_PLATFORM_AGENT_NAME", "agent") + return f"ap-{org}-{agent}" + + async def run( + self, + *, + language: str, + code: str, + timeout: int, + namespace: str | None = None, + allow_network: bool = False, + input_files: dict[str, str] | None = None, + on_output: Callable[[str], None] | None = None, + ) -> ExecutionResult: + """Execute code in an ephemeral K8s Job and return the result.""" + self._ensure_k8s_client() + assert self._batch_api is not None + assert self._core_api is not None + ns = namespace or self.resolve_namespace() + execution_id = uuid.uuid4().hex + started = time.monotonic() + configmap_name: str | None = None + network_policy_name: str | None = None + job_name = f"code-exec-{execution_id[:8]}" + + try: + if input_files: + configmap_name = await self._create_input_configmap( + execution_id, ns, input_files + ) + + should_allow_network = ( + allow_network and self._config.network_access != "deny" + ) or self._config.network_access == "allow_internet" + network_policy_name = await self._create_network_policy( + execution_id, ns, allow_internet=should_allow_network + ) + + manifest = self.build_job_manifest( + language=language, + code=code, + timeout=timeout, + namespace=ns, + execution_id=execution_id, + allow_network=allow_network, + input_configmap_name=configmap_name, + ) + job_name = manifest.metadata.name + await asyncio.to_thread( + self._batch_api.create_namespaced_job, + namespace=ns, + body=manifest, + ) + _log_json( + logging.INFO, + "code_execution_job_created", + job_name=job_name, + namespace=ns, + ) + + scheduling_start = time.monotonic() + if on_output and self._config.streaming_enabled: + pod_name = await self._wait_for_pod(job_name, ns, wait_for_running=True) + stdout, stderr = await self._collect_logs_streaming( + pod_name, ns, on_output + ) + await self._wait_for_pod(job_name, ns, wait_for_running=False) + else: + pod_name = await self._wait_for_pod(job_name, ns) + stdout, stderr = await self._collect_logs(pod_name, ns) + + scheduling_duration = time.monotonic() - scheduling_start + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + CodeExecutionMetrics().record_scheduling_latency( + org=os.environ.get("AI_PLATFORM_AGENT_ORG", "default"), + duration=scheduling_duration, + ) + + raw_exit_code, termination_reason = await self._get_exit_info(pod_name, ns) + exit_code, status = self.parse_container_status( + exit_code=raw_exit_code, + termination_reason=termination_reason, + ) + + cpu_seconds = 0.0 + memory_mb_seconds = 0.0 + if self._config.cost_tracking_enabled: + cpu_seconds, memory_mb_seconds = await self._get_resource_usage( + pod_name, ns, time.monotonic() - started + ) + + duration = time.monotonic() - started + return ExecutionResult( + stdout=stdout, + stderr=stderr, + exit_code=exit_code, + duration_seconds=duration, + status=status, + job_name=job_name, + namespace=ns, + cpu_seconds=cpu_seconds, + memory_mb_seconds=memory_mb_seconds, + scheduling_seconds=scheduling_duration, + ) + except asyncio.TimeoutError: + duration = time.monotonic() - started + return ExecutionResult( + stdout="", + stderr="", + exit_code=-1, + duration_seconds=duration, + status="timeout", + job_name=job_name, + namespace=ns, + ) + except Exception as exc: + duration = time.monotonic() - started + _log_json( + logging.ERROR, + "code_execution_failed", + error=str(exc), + job_name=job_name, + ) + return ExecutionResult( + stdout="", + stderr=str(exc), + exit_code=-1, + duration_seconds=duration, + status="error", + job_name=job_name, + namespace=ns, + ) + finally: + await self._cleanup(job_name, ns) + if configmap_name: + await self._delete_configmap(configmap_name, ns) + if network_policy_name: + await self._delete_network_policy(network_policy_name, ns) + + async def _wait_for_pod( + self, + job_name: str, + namespace: str, + *, + wait_for_running: bool = False, + ) -> str: + """Poll until pod reaches a target state.""" + assert self._core_api is not None + deadline = time.monotonic() + self._config.pod_poll_timeout_seconds + while time.monotonic() < deadline: + pods = await asyncio.to_thread( + self._core_api.list_namespaced_pod, + namespace=namespace, + label_selector=f"job-name={job_name}", + ) + for pod in pods.items: + phase = pod.status.phase + if phase in ("Succeeded", "Failed"): + return str(pod.metadata.name) + if wait_for_running and phase == "Running": + return str(pod.metadata.name) + await asyncio.sleep(self._config.pod_poll_interval_seconds) + raise asyncio.TimeoutError( + f"Pod for {job_name} did not reach target state " + f"within {self._config.pod_poll_timeout_seconds}s" + ) + + async def _collect_logs(self, pod_name: str, namespace: str) -> tuple[str, str]: + """Read stdout/stderr from pod logs after completion.""" + assert self._core_api is not None + try: + raw_logs = await asyncio.to_thread( + self._core_api.read_namespaced_pod_log, + name=pod_name, + namespace=namespace, + container="executor", + ) + if isinstance(raw_logs, bytes): + logs = raw_logs.decode("utf-8", errors="replace") + else: + logs = str(raw_logs) + if logs.startswith("b'") or logs.startswith('b"'): + import ast + + try: + logs = ast.literal_eval(logs).decode("utf-8", errors="replace") + except Exception: + pass + if len(logs) > self._config.max_output_bytes: + logs = logs[: self._config.max_output_bytes] + "\n[truncated at 1MB]" + return logs, "" + except Exception as exc: + _log_json( + logging.WARNING, + "code_execution_log_collection_failed", + error=str(exc), + ) + return "", f"[warning: logs partially collected: {exc}]" + + async def _collect_logs_streaming( + self, + pod_name: str, + namespace: str, + callback: Callable[[str], None], + ) -> tuple[str, str]: + """Stream pod logs in real-time via callback, return full output.""" + assert self._core_api is not None + _log_json(logging.INFO, "code_execution_streaming_started", pod=pod_name) + accumulated: list[str] = [] + total_bytes = 0 + try: + response = await asyncio.to_thread( + self._core_api.read_namespaced_pod_log, + name=pod_name, + namespace=namespace, + container="executor", + follow=True, + _preload_content=False, + ) + + def _read_stream() -> list[str]: + chunks: list[str] = [] + nonlocal total_bytes + for chunk in response.stream(512): + text = ( + chunk.decode("utf-8", errors="replace") + if isinstance(chunk, bytes) + else str(chunk) + ) + total_bytes += len(text) + if total_bytes <= self._config.max_output_bytes: + chunks.append(text) + else: + chunks.append("\n[truncated at 1MB]") + break + response.release_conn() + return chunks + + stream_chunks = await asyncio.to_thread(_read_stream) + for text in stream_chunks: + accumulated.append(text) + callback(text) + _log_json( + logging.INFO, + "code_execution_streaming_completed", + pod=pod_name, + total_bytes=total_bytes, + ) + return "".join(accumulated), "" + except Exception as exc: + _log_json( + logging.WARNING, + "code_execution_streaming_failed", + error=str(exc), + ) + return "".join(accumulated), f"[streaming error: {exc}]" + + async def _get_exit_info( + self, pod_name: str, namespace: str + ) -> tuple[int, str | None]: + """Extract exit code and termination reason from container status.""" + assert self._core_api is not None + try: + pod = await asyncio.to_thread( + self._core_api.read_namespaced_pod, + name=pod_name, + namespace=namespace, + ) + for cs in pod.status.container_statuses or []: + if cs.name == "executor" and cs.state and cs.state.terminated: + return ( + cs.state.terminated.exit_code or 0, + cs.state.terminated.reason, + ) + phase = pod.status.phase + if phase == "Succeeded": + return 0, None + if phase == "Failed": + return 1, None + return -1, None + except Exception: + return -1, None + + async def _get_resource_usage( + self, pod_name: str, namespace: str, duration: float + ) -> tuple[float, float]: + """Query K8s Metrics API for pod resource usage.""" + assert self._custom_api is not None + try: + metrics = await asyncio.to_thread( + self._custom_api.get_namespaced_custom_object, + group="metrics.k8s.io", + version="v1beta1", + namespace=namespace, + plural="pods", + name=pod_name, + ) + containers = metrics.get("containers", []) + for c in containers: + if c.get("name") == "executor": + usage = c.get("usage", {}) + cpu_nano = self._parse_cpu(usage.get("cpu", "0")) + memory_bytes = self._parse_memory(usage.get("memory", "0")) + cpu_seconds = (cpu_nano / 1e9) * duration + memory_mb_seconds = (memory_bytes / (1024 * 1024)) * duration + return cpu_seconds, memory_mb_seconds + return 0.0, 0.0 + except Exception: + cpu_req = self._parse_cpu(self._config.resource_requests.get("cpu", "100m")) + mem_req = self._parse_memory( + self._config.resource_requests.get("memory", "128Mi") + ) + cpu_seconds = (cpu_req / 1e9) * duration + memory_mb_seconds = (mem_req / (1024 * 1024)) * duration + return cpu_seconds, memory_mb_seconds + + @staticmethod + def _parse_cpu(value: str) -> float: + """Parse K8s CPU value to nanocores.""" + value = str(value).strip() + if value.endswith("n"): + return float(value[:-1]) + if value.endswith("m"): + return float(value[:-1]) * 1e6 + return float(value) * 1e9 + + @staticmethod + def _parse_memory(value: str) -> float: + """Parse K8s memory value to bytes.""" + value = str(value).strip() + suffixes = { + "Ki": 1024, + "Mi": 1024**2, + "Gi": 1024**3, + "Ti": 1024**4, + "k": 1000, + "M": 1000**2, + "G": 1000**3, + } + for suffix, multiplier in suffixes.items(): + if value.endswith(suffix): + return float(value[: -len(suffix)]) * multiplier + return float(value) + + async def _create_network_policy( + self, + execution_id: str, + namespace: str, + *, + allow_internet: bool = False, + ) -> str: + """Create an ephemeral NetworkPolicy for this execution.""" + assert self._networking_api is not None + from kubernetes import client + + policy_name = f"code-exec-{execution_id[:8]}" + selector = client.V1LabelSelector( + match_labels={ + "ai-platform.io/execution-id": execution_id, + } + ) + + if allow_internet: + egress = [ + client.V1NetworkPolicyEgressRule( + ports=[ + client.V1NetworkPolicyPort(port=443, protocol="TCP"), + client.V1NetworkPolicyPort(port=80, protocol="TCP"), + ], + to=[ + client.V1NetworkPolicyPeer( + ip_block=client.V1IPBlock( + cidr="0.0.0.0/0", + _except=[ + "10.0.0.0/8", + "172.16.0.0/12", + "192.168.0.0/16", + ], + ) + ) + ], + ) + ] + else: + egress = [] + + policy = client.V1NetworkPolicy( + api_version="networking.k8s.io/v1", + kind="NetworkPolicy", + metadata=client.V1ObjectMeta( + name=policy_name, + namespace=namespace, + ), + spec=client.V1NetworkPolicySpec( + pod_selector=selector, + policy_types=["Egress"], + egress=egress, + ), + ) + + await asyncio.to_thread( + self._networking_api.create_namespaced_network_policy, + namespace=namespace, + body=policy, + ) + _log_json( + logging.INFO, + "code_execution_network_policy_created", + policy=policy_name, + allow_internet=allow_internet, + ) + return policy_name + + async def _delete_network_policy(self, policy_name: str, namespace: str) -> None: + """Delete the ephemeral NetworkPolicy.""" + assert self._networking_api is not None + try: + await asyncio.to_thread( + self._networking_api.delete_namespaced_network_policy, + name=policy_name, + namespace=namespace, + ) + _log_json( + logging.DEBUG, + "code_execution_network_policy_deleted", + policy=policy_name, + ) + except Exception as exc: + _log_json( + logging.WARNING, + "code_execution_network_policy_delete_failed", + policy=policy_name, + error=str(exc), + ) + + async def _create_input_configmap( + self, + execution_id: str, + namespace: str, + files: dict[str, str], + ) -> str: + """Create a ConfigMap with input files for the execution pod.""" + assert self._core_api is not None + from kubernetes import client + + cm_name = f"code-exec-input-{execution_id[:8]}" + configmap = client.V1ConfigMap( + api_version="v1", + kind="ConfigMap", + metadata=client.V1ObjectMeta( + name=cm_name, + namespace=namespace, + labels={ + "app.kubernetes.io/managed-by": "template-agent", + "ai-platform.io/execution-id": execution_id, + }, + ), + data=files, + ) + await asyncio.to_thread( + self._core_api.create_namespaced_config_map, + namespace=namespace, + body=configmap, + ) + _log_json( + logging.INFO, + "code_execution_configmap_created", + configmap=cm_name, + file_count=len(files), + ) + return cm_name + + async def _delete_configmap(self, cm_name: str, namespace: str) -> None: + """Delete the input ConfigMap.""" + assert self._core_api is not None + try: + await asyncio.to_thread( + self._core_api.delete_namespaced_config_map, + name=cm_name, + namespace=namespace, + ) + _log_json( + logging.DEBUG, + "code_execution_configmap_deleted", + configmap=cm_name, + ) + except Exception as exc: + _log_json( + logging.WARNING, + "code_execution_configmap_delete_failed", + configmap=cm_name, + error=str(exc), + ) + + async def _cleanup(self, job_name: str, namespace: str) -> None: + """Delete Job with Foreground propagation (cascades to pods).""" + assert self._batch_api is not None + try: + from kubernetes import client + + await asyncio.to_thread( + self._batch_api.delete_namespaced_job, + name=job_name, + namespace=namespace, + body=client.V1DeleteOptions(propagation_policy="Foreground"), + ) + _log_json(logging.DEBUG, "code_execution_cleanup", job_name=job_name) + except Exception as exc: + _log_json( + logging.WARNING, + "code_execution_cleanup_failed", + job_name=job_name, + error=str(exc), + ) diff --git a/deep_agent/src/code_execution/metrics.py b/deep_agent/src/code_execution/metrics.py new file mode 100644 index 00000000..1a425ffa --- /dev/null +++ b/deep_agent/src/code_execution/metrics.py @@ -0,0 +1,345 @@ +"""Observability for code execution — reuses existing otel.py infrastructure. + +Four layers: +1. OTEL Metrics — instruments on MetricsContainer in aegra/otel.py +2. OTEL Tracing — spans via get_tracer() from aegra/otel.py +3. Platform Audit — structured JSON via audit/emitter.py +4. Structured Logs — JSON lines to stderr via stdlib logging +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import sys +from contextlib import contextmanager +from datetime import UTC, datetime +from typing import Any + +logger = logging.getLogger("deep_agent.src.code_execution") +if not logger.handlers: + _handler = logging.StreamHandler(sys.stderr) + _handler.setFormatter(logging.Formatter("%(message)s")) + logger.addHandler(_handler) + logger.setLevel(logging.DEBUG) + + +def _log_json(level: int, event: str, **fields: Any) -> None: + """Emit a structured JSON log line to stderr.""" + record = { + "event": event, + "logger": "deep_agent.src.code_execution", + "level": logging.getLevelName(level).lower(), + "timestamp": datetime.now(UTC).isoformat(), + "service": "template-agent", + **{k: v for k, v in fields.items() if v is not None}, + } + logger.log(level, json.dumps(record, default=str)) + + +def _get_otel_metrics() -> Any: + """Return the global MetricsContainer from otel.py, or None.""" + try: + from deep_agent.aegra.otel import get_metrics + + return get_metrics() + except Exception: + return None + + +def _get_tracer() -> Any: + """Return an OTEL tracer from otel.py, or None.""" + try: + from deep_agent.aegra.otel import get_tracer + + return get_tracer("code_execution") + except Exception: + return None + + +def _resolve_user_identity() -> tuple[str | None, str | None, str | None]: + """Resolve user, org, trace_id from audit context and OTEL span.""" + user = None + org = os.environ.get("AI_PLATFORM_AGENT_ORG") + trace_id = None + + try: + from deep_agent.src.audit.context import get_audit_context + + ctx = get_audit_context() + user = ctx.get("user") + org = ctx.get("org") or org + trace_id = ctx.get("trace_id") + except ImportError: + pass + + if trace_id is None: + try: + from opentelemetry import trace + + span = trace.get_current_span() + if span and span.get_span_context().is_valid: + trace_id = format(span.get_span_context().trace_id, "032x") + except Exception: + pass + + return user, org, trace_id + + +def emit_audit_event(event_type: str, **details: Any) -> None: + """Emit a platform audit event via the audit emitter.""" + try: + from deep_agent.src.audit.emitter import emit_audit_event as _emit + + _emit(event_type, **details) + except Exception as exc: + _log_json(logging.DEBUG, "audit_emit_failed", error=str(exc)) + + +def compute_code_hash(code: str) -> str: + """Return a SHA-256 hash of the code string.""" + return f"sha256:{hashlib.sha256(code.encode()).hexdigest()}" + + +class CodeExecutionMetrics: + """Observability using existing otel.py MetricsContainer and get_tracer().""" + + def __init__(self) -> None: + """Initialize with tracer from otel.py.""" + self._tracer = _get_tracer() + + def _mc(self) -> Any: + """Get MetricsContainer (lazy — may not be initialized at import time).""" + return _get_otel_metrics() + + # --- Execution metrics (OTEL Layer 1 + Log Layer 4) --- + + def record_execution( + self, + *, + language: str, + org: str, + exit_code: int, + status: str, + duration: float, + ) -> None: + """Record execution via MetricsContainer + structured log.""" + mc = self._mc() + attrs = { + "language": language, + "org": org, + "exit_code": str(exit_code), + "status": status, + } + if mc: + mc.code_execution_duration_seconds.record(duration, attrs) + mc.code_executions_total.add(1, attrs) + _log_json( + logging.INFO, + "code_execution_metric", + language=language, + org=org, + exit_code=exit_code, + status=status, + duration_seconds=round(duration, 3), + ) + + def record_error(self, *, language: str, org: str, error_type: str) -> None: + """Record error via MetricsContainer + structured log.""" + mc = self._mc() + if mc: + mc.code_execution_errors_total.add( + 1, {"language": language, "org": org, "error_type": error_type} + ) + _log_json( + logging.WARNING, + "code_execution_error_metric", + language=language, + org=org, + error_type=error_type, + ) + + def record_scheduling_latency(self, *, org: str, duration: float) -> None: + """Record pod scheduling latency via MetricsContainer + structured log.""" + mc = self._mc() + if mc: + mc.code_execution_scheduling_seconds.record(duration, {"org": org}) + _log_json( + logging.INFO, + "code_execution_scheduling_latency", + org=org, + scheduling_seconds=round(duration, 3), + ) + + def increment_active(self, *, org: str) -> None: + """Increment active execution gauge.""" + mc = self._mc() + if mc: + mc.code_execution_active.add(1, {"org": org}) + _log_json(logging.DEBUG, "code_execution_active_increment", org=org) + + def decrement_active(self, *, org: str) -> None: + """Decrement active execution gauge.""" + mc = self._mc() + if mc: + mc.code_execution_active.add(-1, {"org": org}) + _log_json(logging.DEBUG, "code_execution_active_decrement", org=org) + + # --- Queue metrics (OTEL Layer 1 + Log Layer 4) --- + + def record_queue_wait(self, *, org: str, duration: float) -> None: + """Record queue wait via MetricsContainer + structured log.""" + mc = self._mc() + if mc: + mc.code_execution_queue_wait_seconds.record(duration, {"org": org}) + _log_json( + logging.INFO, + "code_execution_queue_wait", + org=org, + wait_seconds=round(duration, 3), + ) + + def record_rejected(self, *, org: str) -> None: + """Record rejected execution via MetricsContainer + structured log.""" + mc = self._mc() + if mc: + mc.code_execution_rejected_total.add(1, {"org": org}) + _log_json(logging.WARNING, "code_execution_rejected", org=org) + + def log_queued(self, *, org: str) -> None: + """Log that an execution entered the queue.""" + _log_json(logging.DEBUG, "code_execution_queued", org=org) + + def log_dequeued(self, *, org: str, wait_seconds: float) -> None: + """Log that an execution left the queue.""" + _log_json( + logging.DEBUG, + "code_execution_dequeued", + org=org, + wait_seconds=round(wait_seconds, 3), + ) + + # --- Cost tracking (OTEL Layer 1 + Log Layer 4) --- + + def record_resource_usage( + self, + *, + org: str, + language: str, + cpu_seconds: float, + memory_mb_seconds: float, + duration: float, + ) -> None: + """Record resource usage via MetricsContainer + structured log.""" + mc = self._mc() + attrs = {"org": org, "language": language} + if mc: + mc.code_execution_cpu_seconds.record(cpu_seconds, attrs) + mc.code_execution_memory_mb_seconds.record(memory_mb_seconds, attrs) + _log_json( + logging.INFO, + "code_execution_resource_usage", + org=org, + language=language, + cpu_seconds=round(cpu_seconds, 4), + memory_mb_seconds=round(memory_mb_seconds, 2), + duration_seconds=round(duration, 3), + ) + + # --- OTEL Tracing (Layer 2) --- + + @contextmanager + def trace_span(self, name: str, **attributes: Any) -> Any: + """Context manager for OTEL tracing spans (sets as active span).""" + if self._tracer is None: + yield None + return + try: + from opentelemetry import trace + + span = self._tracer.start_span(name, attributes=attributes) + ctx = trace.set_span_in_context(span) + token = trace.context_api.attach(ctx) + try: + yield span + except Exception as exc: + if span: + span.set_attribute("error", True) + span.set_attribute("error.message", str(exc)) + raise + finally: + if span: + span.end() + trace.context_api.detach(token) + except ImportError: + yield None + + def start_span(self, name: str, **attributes: Any) -> Any: + """Start an OTEL tracing span (manual end required).""" + if self._tracer is None: + return None + return self._tracer.start_span(name, attributes=attributes) + + # --- Platform Audit (Layer 3) --- + + def emit_audit( + self, + *, + language: str, + status: str, + exit_code: int, + latency_ms: float, + code_hash: str, + namespace: str, + image: str, + job_name: str, + timeout: int, + stdout_bytes: int, + stderr_bytes: int, + scheduling_seconds: float = 0.0, + ) -> None: + """Emit audit event with user identity from context.""" + emit_audit_event( + "code_execution", + agent="orchestrator", + language=language, + status=status, + exit_code=exit_code, + latency_ms=latency_ms, + code_hash=code_hash, + namespace=namespace, + image=image, + job_name=job_name, + timeout_seconds=timeout, + stdout_bytes=stdout_bytes, + stderr_bytes=stderr_bytes, + scheduling_seconds=round(scheduling_seconds, 3), + ) + + # --- Structured Logging (Layer 4) --- + + def log_started(self, **fields: Any) -> None: + """Log that a code execution has started.""" + _log_json(logging.INFO, "code_execution_started", **fields) + + def log_completed(self, **fields: Any) -> None: + """Log that a code execution completed.""" + _log_json(logging.INFO, "code_execution_completed", **fields) + + def log_timeout(self, **fields: Any) -> None: + """Log that a code execution timed out.""" + _log_json(logging.WARNING, "code_execution_timeout", **fields) + + def log_oom(self, **fields: Any) -> None: + """Log that a code execution was OOM killed.""" + _log_json(logging.WARNING, "code_execution_oom_killed", **fields) + + def log_failed(self, **fields: Any) -> None: + """Log that a code execution failed.""" + _log_json(logging.ERROR, "code_execution_failed", **fields) + + def log_cleanup(self, **fields: Any) -> None: + """Log code execution cleanup.""" + _log_json(logging.DEBUG, "code_execution_cleanup", **fields) diff --git a/deep_agent/src/code_execution/middleware.py b/deep_agent/src/code_execution/middleware.py new file mode 100644 index 00000000..0e2d8e02 --- /dev/null +++ b/deep_agent/src/code_execution/middleware.py @@ -0,0 +1,294 @@ +"""CodeExecutionMiddleware — inject execute_code tool, route to K8s Jobs.""" + +from __future__ import annotations + +import asyncio +import os +import time +from typing import Any + +from langchain.agents.middleware.types import ( + AgentMiddleware, + ModelRequest, + ModelResponse, + ToolCallRequest, +) +from langchain_core.messages import ToolMessage +from langchain_core.tools import tool + +from deep_agent.src.code_execution.config import CodeExecutionConfig +from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner +from deep_agent.src.code_execution.metrics import ( + CodeExecutionMetrics, + compute_code_hash, +) + + +def _build_execute_code_tool(config: CodeExecutionConfig) -> Any: + """Build the execute_code tool definition for LLM tool binding.""" + + @tool + def execute_code( + code: str, + language: str = "python", + timeout: int = 60, + network: bool = False, + input_files: dict[str, str] | None = None, + ) -> str: + """Execute code in an isolated sandbox environment. + + Args: + code: The source code to execute. + language: Programming language (python, python-ds, python-ml, shell, node). + timeout: Maximum execution time in seconds. + network: Whether to allow internet access from the sandbox. + input_files: Optional dict of filename to content, mounted at /input/. + + Returns: + Execution output with stdout, stderr, and exit code. + """ + return "This tool is handled by CodeExecutionMiddleware" + + return execute_code + + +class CodeExecutionMiddleware(AgentMiddleware): + """Inject execute_code tool and route calls to K8s Job backend.""" + + def __init__(self, *, config: CodeExecutionConfig) -> None: + """Initialize middleware with execution configuration.""" + self._config = config + self._runner = K8sJobRunner(config) + self._metrics = CodeExecutionMetrics() + self._execute_code_tool = _build_execute_code_tool(config) + self._semaphores: dict[str, asyncio.Semaphore] = {} + + def _get_semaphore(self, org: str) -> asyncio.Semaphore: + """Get or create a per-org execution semaphore.""" + if org not in self._semaphores: + self._semaphores[org] = asyncio.Semaphore( + self._config.max_concurrent_per_org + ) + return self._semaphores[org] + + def wrap_model_call( + self, request: ModelRequest[Any], handler: Any + ) -> ModelResponse[Any]: + """Synchronous model call pass-through.""" + return handler(request) + + async def awrap_model_call( + self, request: ModelRequest[Any], handler: Any + ) -> ModelResponse[Any]: + """Inject the execute_code tool into model requests when enabled.""" + if not self._config.enabled: + return await handler(request) + updated = request.override(tools=[*request.tools, self._execute_code_tool]) + return await handler(updated) + + def wrap_tool_call(self, request: ToolCallRequest, handler: Any) -> Any: + """Synchronous tool call pass-through.""" + return handler(request) + + async def awrap_tool_call(self, request: ToolCallRequest, handler: Any) -> Any: + """Intercept execute_code tool calls and route to K8s backend.""" + tool_call = request.tool_call + if tool_call.get("name") != "execute_code": + return await handler(request) + + args = tool_call.get("args", {}) + code = args.get("code", "") + language = args.get("language", "python") + timeout = min( + int(args.get("timeout", self._config.max_timeout_seconds)), + self._config.max_timeout_seconds, + ) + network = bool(args.get("network", False)) + input_files = args.get("input_files") + tool_call_id = tool_call.get("id", "") + + if not code.strip(): + return ToolMessage( + content="No code provided to execute", + tool_call_id=tool_call_id, + ) + + if language not in self._config.supported_languages: + return ToolMessage( + content=f"Unsupported language: {language}. " + f"Supported: {', '.join(sorted(self._config.supported_languages))}", + tool_call_id=tool_call_id, + ) + + if len(code) > self._config.max_code_length: + return ToolMessage( + content=f"Code exceeds maximum length of " + f"{self._config.max_code_length} characters", + tool_call_id=tool_call_id, + ) + + if input_files: + total_size = sum(len(v) for v in input_files.values()) + if total_size > self._config.max_input_file_size: + return ToolMessage( + content=f"Input files exceed maximum size of " + f"{self._config.max_input_file_size} bytes", + tool_call_id=tool_call_id, + ) + + if network and self._config.network_access == "deny": + network = False + + org = os.environ.get("AI_PLATFORM_AGENT_ORG", "default") + namespace = self._runner.resolve_namespace() + semaphore = self._get_semaphore(org) + + self._metrics.log_queued(org=org) + queue_start = time.monotonic() + acquired = False + try: + await asyncio.wait_for( + semaphore.acquire(), + timeout=self._config.queue_timeout_seconds, + ) + acquired = True + except asyncio.TimeoutError: + self._metrics.record_rejected(org=org) + return ToolMessage( + content="Code execution queue full, try again later", + tool_call_id=tool_call_id, + ) + queue_wait = time.monotonic() - queue_start + self._metrics.record_queue_wait(org=org, duration=queue_wait) + self._metrics.log_dequeued(org=org, wait_seconds=queue_wait) + + image = self._config.images.get(language, "unknown") + + self._metrics.increment_active(org=org) + self._metrics.log_started( + language=language, + org=org, + namespace=namespace, + image=image, + timeout_seconds=timeout, + code_length=len(code), + network=network, + input_file_count=len(input_files) if input_files else 0, + ) + + on_output = None + if self._config.streaming_enabled: + try: + from langgraph.config import get_stream_writer + + stream_writer = get_stream_writer() + + def on_output(chunk: str) -> None: + stream_writer({"type": "code_output", "content": chunk}) + except Exception: + pass + + started = time.monotonic() + try: + with self._metrics.trace_span( + "code_execution", + language=language, + org=org, + namespace=namespace, + image=image, + ) as span: + result = await self._runner.run( + language=language, + code=code, + timeout=timeout, + namespace=namespace, + allow_network=network, + input_files=input_files, + on_output=on_output, + ) + + duration = time.monotonic() - started + latency_ms = round(duration * 1000, 2) + + if span: + span.set_attribute("exit_code", result.exit_code) + span.set_attribute("status", result.status) + span.set_attribute("duration_ms", latency_ms) + span.set_attribute("job_name", result.job_name) + + self._metrics.record_execution( + language=language, + org=org, + exit_code=result.exit_code, + status=result.status, + duration=duration, + ) + + if self._config.cost_tracking_enabled: + self._metrics.record_resource_usage( + org=org, + language=language, + cpu_seconds=result.cpu_seconds, + memory_mb_seconds=result.memory_mb_seconds, + duration=duration, + ) + + self._metrics.emit_audit( + language=language, + status=result.status, + exit_code=result.exit_code, + latency_ms=latency_ms, + code_hash=compute_code_hash(code), + namespace=namespace, + image=image, + job_name=result.job_name, + timeout=timeout, + stdout_bytes=len(result.stdout), + stderr_bytes=len(result.stderr), + scheduling_seconds=result.scheduling_seconds, + ) + + if result.status == "timeout": + self._metrics.log_timeout( + job_name=result.job_name, + timeout_seconds=timeout, + namespace=namespace, + ) + elif result.status == "oom_killed": + self._metrics.log_oom( + job_name=result.job_name, + memory_limit=self._config.resource_limits.get("memory", "unknown"), + namespace=namespace, + ) + else: + self._metrics.log_completed( + exit_code=result.exit_code, + duration_ms=latency_ms, + status=result.status, + job_name=result.job_name, + namespace=namespace, + image=image, + ) + + return ToolMessage(content=result.format(), tool_call_id=tool_call_id) + + except Exception as exc: + duration = time.monotonic() - started + self._metrics.record_error( + language=language, + org=org, + error_type=type(exc).__name__, + ) + self._metrics.log_failed( + error_type=type(exc).__name__, + error_message=str(exc), + duration_ms=round(duration * 1000, 2), + ) + return ToolMessage( + content="Code execution service temporarily unavailable", + tool_call_id=tool_call_id, + ) + finally: + self._metrics.decrement_active(org=org) + if acquired: + semaphore.release() diff --git a/deep_agent/src/error_handling.py b/deep_agent/src/error_handling.py new file mode 100644 index 00000000..ea3ab46a --- /dev/null +++ b/deep_agent/src/error_handling.py @@ -0,0 +1,440 @@ +"""Centralized error handling: retry decorators, circuit breaker, fallback patterns. + +This module provides production-grade error handling utilities built on tenacity. +It separates *how* we handle errors (retry, circuit break, degrade) from *what* +errors look like (exceptions.py). + +Usage: + from deep_agent.src.error_handling import llm_retry, mcp_retry, create_circuit_breaker + + @llm_retry + def create_model(name: str) -> ChatModel: ... + + breaker = create_circuit_breaker("mcp-server", threshold=3) + if breaker.is_open: + return fallback() +""" + +import asyncio +import logging as _logging +import time +from collections.abc import Callable +from functools import wraps +from typing import Any, TypeVar + +from tenacity import ( + RetryCallState, + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) + +from deep_agent.src.exceptions import ( + AppException, + LLMError, + MCPError, + RateLimitError, + TransientError, +) +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +F = TypeVar("F", bound=Callable[..., Any]) + +# --------------------------------------------------------------------------- +# Retry callbacks (shared across decorators) +# --------------------------------------------------------------------------- + + +def _log_retry(retry_state: RetryCallState) -> None: + """Log retry attempts with structured context.""" + exc = retry_state.outcome.exception() if retry_state.outcome else None + logger.warning( + "Retry %d/%d for '%s': %s", + retry_state.attempt_number, + retry_state.retry_object.stop.max_attempt_number, + retry_state.fn.__name__ if retry_state.fn else "unknown", + exc, + ) + + +# --------------------------------------------------------------------------- +# Retry decorators +# --------------------------------------------------------------------------- + +llm_retry = retry( + retry=retry_if_exception_type((LLMError, RateLimitError, ConnectionError, OSError)), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=2, max=30), + before_sleep=_log_retry, + reraise=True, +) + +mcp_retry = retry( + retry=retry_if_exception_type((MCPError, ConnectionError, TimeoutError, OSError)), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=1, max=15), + before_sleep=_log_retry, + reraise=True, +) + +subagent_retry = retry( + retry=retry_if_exception_type((TransientError, ConnectionError, OSError)), + stop=stop_after_attempt(2), + wait=wait_exponential(multiplier=1, min=1, max=10), + before_sleep=_log_retry, + reraise=True, +) + +try: + from pymongo.errors import ( + AutoReconnect, + ConnectionFailure, + NetworkTimeout, + NotPrimaryError, + ServerSelectionTimeoutError, + ) + + _MONGO_TRANSIENT_ERRORS: tuple[type[Exception], ...] = ( + AutoReconnect, + ConnectionFailure, + NetworkTimeout, + NotPrimaryError, + ServerSelectionTimeoutError, + ConnectionError, + TimeoutError, + OSError, + ) +except ImportError: # pragma: no cover - pymongo optional at import time + _MONGO_TRANSIENT_ERRORS = (ConnectionError, TimeoutError, OSError) + +mongo_retry = retry( + retry=retry_if_exception_type(_MONGO_TRANSIENT_ERRORS), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=0.2, min=0.1, max=2), + before_sleep=_log_retry, + reraise=True, +) + + +# --------------------------------------------------------------------------- +# Circuit Breaker +# --------------------------------------------------------------------------- + +_REDIS_KEY_PREFIX = "aegra:circuit:" + + +class CircuitBreaker: + """Circuit breaker for external service calls with optional Redis persistence. + + Tracks consecutive failures. After ``threshold`` failures, the circuit + opens and remains open for ``reset_timeout`` seconds, during which + calls should be skipped (or use a fallback). + + When ``redis_client`` is provided, state is stored in a Redis hash, + enabling multi-replica awareness. When Redis is unavailable or not + provided, state is kept in-memory (single-process only). + + Redis errors never propagate — the breaker degrades to "closed" + (allow all requests) if Redis is unreachable. + + Args: + name: Human-readable name (also used as Redis key suffix). + threshold: Consecutive failures before opening. + reset_timeout: Seconds to wait before allowing a probe (half-open). + redis_client: Optional Redis client. Pass explicitly or use + ``create_circuit_breaker()`` for auto-detection. + """ + + def __init__( + self, + name: str, + threshold: int = 5, + reset_timeout: float = 60.0, + redis_client: Any = None, + ) -> None: + """Initialize circuit breaker with name, threshold, and optional Redis backing.""" + self.name = name + self.threshold = threshold + self.reset_timeout = reset_timeout + self._redis: Any = redis_client + self._redis_key: str = f"{_REDIS_KEY_PREFIX}{name}" + self._key_ttl: int = max(int(reset_timeout * 3), 300) + + # In-memory fallback state + self._mem_failure_count: int = 0 + self._mem_last_failure_time: float = 0.0 + self._mem_state: str = "closed" + + # ── State reading ───────────────────────────────────────────── + + def _read_state(self) -> tuple[int, str, float]: + """Read (failure_count, state, last_failure_time) from backend.""" + if self._redis is not None: + try: + data: dict[str, str] = self._redis.hgetall(self._redis_key) + if not data: + return 0, "closed", 0.0 + return ( + int(data.get("failures", "0")), + data.get("state", "closed"), + float(data.get("last_failure_ts", "0")), + ) + except Exception: + logger.debug( + "Circuit '%s' Redis read failed — falling back to closed", + self.name, + ) + return 0, "closed", 0.0 + return self._mem_failure_count, self._mem_state, self._mem_last_failure_time + + def _write_state(self, failures: int, state: str, last_failure_ts: float) -> None: + """Write state to backend.""" + if self._redis is not None: + try: + self._redis.hset( + self._redis_key, + mapping={ + "failures": str(failures), + "state": state, + "last_failure_ts": str(last_failure_ts), + }, + ) + self._redis.expire(self._redis_key, self._key_ttl) + return + except Exception: + logger.debug( + "Circuit '%s' Redis write failed — using in-memory", + self.name, + ) + self._mem_failure_count = failures + self._mem_state = state + self._mem_last_failure_time = last_failure_ts + + def _clear_state(self) -> None: + """Clear all state (reset to closed).""" + if self._redis is not None: + try: + self._redis.delete(self._redis_key) + return + except Exception: + logger.debug("Circuit '%s' Redis delete failed", self.name) + self._mem_failure_count = 0 + self._mem_state = "closed" + self._mem_last_failure_time = 0.0 + + # ── Public API ──────────────────────────────────────────────── + + @property + def is_open(self) -> bool: + """True when the circuit is open (calls should be skipped).""" + failures, state, last_ts = self._read_state() + if state == "open": + if time.monotonic() - last_ts >= self.reset_timeout: + self._write_state(failures, "half-open", last_ts) + logger.info( + "Circuit '%s' half-open — allowing probe request", self.name + ) + return False + return True + return False + + @property + def state(self) -> str: + """Current circuit state: closed, open, or half-open.""" + _ = self.is_open + _, state, _ = self._read_state() + return state + + def record_success(self) -> None: + """Record a successful call. Resets failure count and closes circuit.""" + failures, state, _ = self._read_state() + if failures > 0 or state != "closed": + logger.info( + "Circuit '%s' reset after success (was %s, %d failures)", + self.name, + state, + failures, + ) + self._clear_state() + + def record_failure(self) -> None: + """Record a failed call. Opens circuit if threshold exceeded.""" + failures, _, _ = self._read_state() + failures += 1 + now = time.monotonic() + + new_state = "open" if failures >= self.threshold else "closed" + self._write_state(failures, new_state, now) + + if new_state == "open": + logger.warning( + "Circuit '%s' OPEN after %d consecutive failures (cooldown: %.0fs)", + self.name, + failures, + self.reset_timeout, + ) + + +# --------------------------------------------------------------------------- +# Factory +# --------------------------------------------------------------------------- + + +def get_redis_client() -> Any: + """Import and return Redis client from aegra.redis (None if unavailable).""" + try: + from deep_agent.aegra.redis import get_redis_client as _get + + return _get() + except Exception: + return None + + +def create_circuit_breaker( + name: str, + threshold: int = 5, + reset_timeout: float = 60.0, + redis_client: Any = None, +) -> CircuitBreaker: + """Create a CircuitBreaker with auto-detected Redis backend. + + If ``redis_client`` is not provided, attempts to obtain one from + ``aegra.redis.get_redis_client()``. Falls back to in-memory if + Redis is unavailable. + + Args: + name: Circuit name (used as Redis key suffix). + threshold: Failures before opening. + reset_timeout: Seconds before half-open probe. + redis_client: Explicit Redis client (overrides auto-detect). + + Returns: + Configured CircuitBreaker instance. + """ + if redis_client is None: + redis_client = get_redis_client() + + if redis_client is not None: + logger.info("Circuit '%s' using Redis-backed state", name) + else: + logger.info("Circuit '%s' using in-memory state (single-replica)", name) + + return CircuitBreaker( + name=name, + threshold=threshold, + reset_timeout=reset_timeout, + redis_client=redis_client, + ) + + +# --------------------------------------------------------------------------- +# Graceful degradation helpers +# --------------------------------------------------------------------------- + + +def classify_error(exc: Exception) -> dict[str, Any]: + """Classify an exception into a structured error response for the API. + + Returns a dict suitable for yielding as a stream error event. + PII is scrubbed in production environments. + + Args: + exc: The exception to classify. + + Returns: + Structured error dict with type, message, recoverable flag, and error_type. + """ + from deep_agent.src.pii_scrubber import scrub_pii + from deep_agent.src.settings import settings + + if isinstance(exc, RateLimitError): + return { + "message": "Rate limit exceeded — please wait and try again", + "recoverable": True, + "error_type": "rate_limit", + } + if isinstance(exc, TransientError): + message = f"Service temporarily unavailable: {exc.message}" + return { + "message": scrub_pii(message) if settings.is_production else message, + "recoverable": True, + "error_type": "transient", + } + if isinstance(exc, AppException): + return { + "message": scrub_pii(exc.message) + if settings.is_production + else exc.message, + "recoverable": False, + "error_type": exc.code, + } + + # Generic error - minimal info in production + if settings.is_production: + return { + "message": "Internal server error", + "recoverable": False, + "error_type": "unknown", + } + + return { + "message": f"Internal server error: {str(exc)}", + "recoverable": False, + "error_type": "unknown", + } + + +def with_fallback( + fallback_value: Any, + *, + on: tuple[type[Exception], ...] = (Exception,), + log_level: int = _logging.WARNING, +) -> Callable[[F], F]: + """Decorator that returns a fallback value instead of raising. + + Use for non-critical paths where a degraded response is better than + a failure. The original exception is logged. + + Args: + fallback_value: Value to return when the wrapped function raises. + on: Tuple of exception types to catch. + log_level: Logging level for the caught exception. + """ + + def decorator(fn: F) -> F: + @wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return fn(*args, **kwargs) + except on as exc: + logger.log( + log_level, + "Fallback for '%s': %s (returning %r)", + fn.__name__, + exc, + fallback_value, + ) + return fallback_value + + @wraps(fn) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return await fn(*args, **kwargs) + except on as exc: + logger.log( + log_level, + "Fallback for '%s': %s (returning %r)", + fn.__name__, + exc, + fallback_value, + ) + return fallback_value + + if asyncio.iscoroutinefunction(fn): + return async_wrapper # type: ignore[return-value] + return wrapper # type: ignore[return-value] + + return decorator diff --git a/deep_agent/src/exceptions.py b/deep_agent/src/exceptions.py new file mode 100644 index 00000000..c34fb2e3 --- /dev/null +++ b/deep_agent/src/exceptions.py @@ -0,0 +1,204 @@ +"""Application-wide exception hierarchy and error codes. + +This module defines the exception hierarchy and error codes used throughout +the application. Exceptions are organized by subsystem (LLM, MCP, subagent, +configuration) with a shared base class for consistent error handling. + +Classes: + ErrorCode: Immutable error code with status, message, and code + ErrorCodes: Collection of predefined error codes + AppException: Base exception for all application errors + TransientError: Base for retryable errors + LLMError: LLM/model creation failures + MCPError: MCP server connection failures + SubAgentError: Subagent loading/execution failures + ConfigurationError: Configuration loading/validation failures + RateLimitError: Rate limit exceeded (retryable) + AuthenticationError: Authentication/authorization failures +""" + +from dataclasses import dataclass + +from starlette.status import ( + HTTP_401_UNAUTHORIZED, + HTTP_429_TOO_MANY_REQUESTS, + HTTP_500_INTERNAL_SERVER_ERROR, + HTTP_502_BAD_GATEWAY, + HTTP_503_SERVICE_UNAVAILABLE, + HTTP_504_GATEWAY_TIMEOUT, +) + + +@dataclass(frozen=True) +class ErrorCode: + """Error code with HTTP status and message.""" + + status: int + message: str + code: str + + +class ErrorCodes: + """Error codes for the template agent.""" + + INTERNAL_SERVER_ERROR = ErrorCode( + HTTP_500_INTERNAL_SERVER_ERROR, + "Internal Server Error", + "E_001", + ) + LLM_ERROR = ErrorCode( + HTTP_502_BAD_GATEWAY, + "LLM Service Error", + "E_002", + ) + LLM_TIMEOUT = ErrorCode( + HTTP_504_GATEWAY_TIMEOUT, + "LLM Request Timeout", + "E_003", + ) + MCP_CONNECTION_ERROR = ErrorCode( + HTTP_502_BAD_GATEWAY, + "MCP Connection Failed", + "E_004", + ) + MCP_TIMEOUT = ErrorCode( + HTTP_504_GATEWAY_TIMEOUT, + "MCP Request Timeout", + "E_005", + ) + SUBAGENT_ERROR = ErrorCode( + HTTP_500_INTERNAL_SERVER_ERROR, + "Subagent Execution Failed", + "E_006", + ) + CONFIGURATION_INITIALIZATION_ERROR = ErrorCode( + HTTP_500_INTERNAL_SERVER_ERROR, + "Configuration Initialization Failed", + "E_007", + ) + CONFIGURATION_VALIDATION_ERROR = ErrorCode( + HTTP_500_INTERNAL_SERVER_ERROR, + "Configuration Validation Failed", + "E_008", + ) + RATE_LIMIT_ERROR = ErrorCode( + HTTP_429_TOO_MANY_REQUESTS, + "Rate Limit Exceeded", + "E_009", + ) + AUTHENTICATION_ERROR = ErrorCode( + HTTP_401_UNAUTHORIZED, + "Authentication Failed", + "E_010", + ) + SERVICE_UNAVAILABLE = ErrorCode( + HTTP_503_SERVICE_UNAVAILABLE, + "Service Temporarily Unavailable", + "E_011", + ) + + # Legacy aliases (kept for backward compatibility) + PRODUCTION_MCP_CONNECTION_ERROR = MCP_CONNECTION_ERROR + + +class AppException(Exception): + """Base exception for application errors.""" + + def __init__( + self, + detail: str, + error_code: ErrorCode = ErrorCodes.INTERNAL_SERVER_ERROR, + ) -> None: + """Initialize exception with detail message and error code.""" + self.detail = detail + self.error_code = error_code + super().__init__(detail) + + @property + def status(self) -> int: + """HTTP status code.""" + return self.error_code.status + + @property + def message(self) -> str: + """Error message.""" + return self.error_code.message + + @property + def code(self) -> str: + """Error code.""" + return self.error_code.code + + @property + def is_retryable(self) -> bool: + """Whether this error is safe to retry.""" + return False + + +class TransientError(AppException): + """Base for errors that are safe to retry. + + Subclasses represent failures from external services (LLM, MCP, network) + that may succeed on a subsequent attempt. + """ + + @property + def is_retryable(self) -> bool: + """Return True; transient errors are retryable by definition.""" + return True + + +class LLMError(TransientError): + """LLM model creation or invocation failure.""" + + def __init__(self, detail: str) -> None: # noqa: D107 + super().__init__(detail, ErrorCodes.LLM_ERROR) + + +class LLMTimeoutError(TransientError): + """LLM request timed out.""" + + def __init__(self, detail: str) -> None: # noqa: D107 + super().__init__(detail, ErrorCodes.LLM_TIMEOUT) + + +class MCPError(TransientError): + """MCP server connection or tool invocation failure.""" + + def __init__(self, detail: str) -> None: # noqa: D107 + super().__init__(detail, ErrorCodes.MCP_CONNECTION_ERROR) + + +class MCPTimeoutError(TransientError): + """MCP server request timed out.""" + + def __init__(self, detail: str) -> None: # noqa: D107 + super().__init__(detail, ErrorCodes.MCP_TIMEOUT) + + +class SubAgentError(AppException): + """Subagent loading or execution failure.""" + + def __init__(self, detail: str) -> None: # noqa: D107 + super().__init__(detail, ErrorCodes.SUBAGENT_ERROR) + + +class ConfigurationError(AppException): + """Configuration loading or validation failure.""" + + def __init__(self, detail: str) -> None: # noqa: D107 + super().__init__(detail, ErrorCodes.CONFIGURATION_INITIALIZATION_ERROR) + + +class RateLimitError(TransientError): + """Rate limit exceeded — should back off and retry.""" + + def __init__(self, detail: str) -> None: # noqa: D107 + super().__init__(detail, ErrorCodes.RATE_LIMIT_ERROR) + + +class AuthenticationError(AppException): + """Authentication or authorization failure — do NOT retry.""" + + def __init__(self, detail: str) -> None: # noqa: D107 + super().__init__(detail, ErrorCodes.AUTHENTICATION_ERROR) diff --git a/deep_agent/src/feedback/__init__.py b/deep_agent/src/feedback/__init__.py new file mode 100644 index 00000000..534b854a --- /dev/null +++ b/deep_agent/src/feedback/__init__.py @@ -0,0 +1,5 @@ +"""Message feedback persistence (Postgres).""" + +from deep_agent.src.feedback.repository import FeedbackRepository + +__all__ = ["FeedbackRepository"] diff --git a/deep_agent/src/feedback/repository.py b/deep_agent/src/feedback/repository.py new file mode 100644 index 00000000..35cd7b44 --- /dev/null +++ b/deep_agent/src/feedback/repository.py @@ -0,0 +1,122 @@ +"""Async Postgres repository for message feedback.""" + +from __future__ import annotations + +from typing import Any, Literal + +import psycopg +from psycopg.rows import dict_row + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_TABLE_ENSURED = False + +CREATE_FEEDBACK_TABLE = """ +CREATE TABLE IF NOT EXISTS message_feedback ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + thread_id TEXT NOT NULL, + message_id TEXT NOT NULL, + user_id TEXT NOT NULL DEFAULT 'anonymous', + feedback TEXT NOT NULL CHECK (feedback IN ('up', 'down')), + trace_id TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE (thread_id, message_id, user_id) +); +CREATE INDEX IF NOT EXISTS idx_feedback_thread ON message_feedback (thread_id); +""" + + +class FeedbackRepository: + """Thin async wrapper around the message_feedback table.""" + + def __init__(self, database_uri: str) -> None: + """Initialize with a Postgres connection URI.""" + self._uri = database_uri + + async def ensure_table(self) -> None: + """Create message_feedback table if it does not already exist (lazy, once).""" + global _TABLE_ENSURED # noqa: PLW0603 + if _TABLE_ENSURED: + return + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + await conn.execute(CREATE_FEEDBACK_TABLE) + await conn.commit() + _TABLE_ENSURED = True + logger.info("message_feedback table ensured") + + async def upsert_feedback( + self, + thread_id: str, + message_id: str, + user_id: str, + feedback: Literal["up", "down"], + trace_id: str | None = None, + ) -> None: + """Insert or update feedback for a message (per thread and user).""" + await self.ensure_table() + uid = user_id if user_id else "anonymous" + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + await conn.execute( + """ + INSERT INTO message_feedback ( + thread_id, message_id, user_id, feedback, trace_id, updated_at + ) + VALUES (%s, %s, %s, %s, %s, now()) + ON CONFLICT (thread_id, message_id, user_id) + DO UPDATE SET + feedback = EXCLUDED.feedback, + trace_id = EXCLUDED.trace_id, + updated_at = now() + """, + (thread_id, message_id, uid, feedback, trace_id), + ) + await conn.commit() + + async def delete_feedback( + self, + thread_id: str, + message_id: str, + user_id: str, + ) -> bool: + """Remove feedback row (un-vote). Returns True if a row was deleted.""" + await self.ensure_table() + uid = user_id if user_id else "anonymous" + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + cur = await conn.execute( + """ + DELETE FROM message_feedback + WHERE thread_id = %s AND message_id = %s AND user_id = %s + """, + (thread_id, message_id, uid), + ) + await conn.commit() + return bool(cur.rowcount > 0) + + async def list_feedback( + self, + thread_id: str, + user_id: str, + ) -> list[dict[str, Any]]: + """Return feedback entries for the thread and user as ``{message_id, feedback}``.""" + await self.ensure_table() + uid = user_id if user_id else "anonymous" + async with await psycopg.AsyncConnection.connect( + self._uri, row_factory=dict_row + ) as conn: + cur = await conn.execute( + """ + SELECT message_id, feedback + FROM message_feedback + WHERE thread_id = %s AND user_id = %s + ORDER BY updated_at ASC + """, + (thread_id, uid), + ) + rows = await cur.fetchall() + return [ + {"message_id": str(r["message_id"]), "feedback": r["feedback"]} + for r in rows + ] diff --git a/deep_agent/src/infrastructure/__init__.py b/deep_agent/src/infrastructure/__init__.py new file mode 100644 index 00000000..354a814c --- /dev/null +++ b/deep_agent/src/infrastructure/__init__.py @@ -0,0 +1,20 @@ +"""Infrastructure layer for external system integrations. + +This package contains modules that interface with external systems and services: +- MCP servers for tools +- Backend execution environments +- Subagent configuration loading + +These modules form the boundary between our application and external dependencies. +""" + +from .backend import get_backend, get_configured_backend +from .mcp import get_mcp_tools +from .subagents import load_subagents + +__all__ = [ + "get_mcp_tools", + "get_backend", + "get_configured_backend", + "load_subagents", +] diff --git a/deep_agent/src/infrastructure/async_tasks.py b/deep_agent/src/infrastructure/async_tasks.py new file mode 100644 index 00000000..1324dd20 --- /dev/null +++ b/deep_agent/src/infrastructure/async_tasks.py @@ -0,0 +1,83 @@ +"""Async subagent middleware builder. + +Auto-detects async subagents (type: async in frontmatter) from the loaded +subagent list and builds AsyncSubAgentMiddleware to wire background task +tools into the agent. + +Template-agent users configure async subagents via Markdown frontmatter: + type: async + graph_id: my_graph + url: https://my-deployment.example.com # optional + +This module handles the middleware wiring. Users never call it directly. +""" + +from __future__ import annotations + +from typing import Any + +from deep_agent.src.agent.config.providers import AsyncTaskConfig +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def build_async_middleware( + subagents: list[Any] | None, + async_config: AsyncTaskConfig, +) -> Any | None: + """Build AsyncSubAgentMiddleware if async subagents exist. + + Scans the loaded subagent list for AsyncSubAgent instances and + wraps them in AsyncSubAgentMiddleware, which adds tools for + launching, monitoring, and updating background tasks. + + Args: + subagents: List of loaded subagent instances (SubAgent, + CompiledSubAgent, AsyncSubAgent). Can be None. + async_config: Async task settings from providers.yaml. + + Returns: + AsyncSubAgentMiddleware instance, or None if no async subagents + exist or the feature is disabled. + """ + if not async_config.enabled: + logger.debug("Async tasks disabled via config") + return None + + if not subagents: + return None + + async_subagents = _extract_async_subagents(subagents) + if not async_subagents: + return None + + try: + from deepagents.middleware.async_subagents import AsyncSubAgentMiddleware + + kwargs: dict[str, Any] = {"async_subagents": async_subagents} + if async_config.system_prompt is not None: + kwargs["system_prompt"] = async_config.system_prompt + + middleware = AsyncSubAgentMiddleware(**kwargs) + logger.info( + "Built AsyncSubAgentMiddleware with %d async subagent(s)", + len(async_subagents), + ) + return middleware + except ImportError: + logger.warning("AsyncSubAgentMiddleware not available — async tasks disabled") + return None + except Exception as e: + logger.warning("Failed to build AsyncSubAgentMiddleware: %s", e) + return None + + +def _extract_async_subagents(subagents: list[Any]) -> list[Any]: + """Filter the subagent list for AsyncSubAgent instances.""" + try: + from deepagents.middleware.async_subagents import AsyncSubAgent + + return [s for s in subagents if isinstance(s, AsyncSubAgent)] + except ImportError: + return [] diff --git a/deep_agent/src/infrastructure/backend.py b/deep_agent/src/infrastructure/backend.py new file mode 100644 index 00000000..31e59a84 --- /dev/null +++ b/deep_agent/src/infrastructure/backend.py @@ -0,0 +1,448 @@ +"""Agent backend for state management and skill execution. + +This module provides the backend infrastructure for agents to execute skills +in isolated Python environments. It creates dedicated virtual environments for +skill execution, manages dependencies from config/skills/pyproject.toml, and +provides a safe execution sandbox. + +Why this exists: + Skills need to run Python code with specific dependencies without polluting + the main application environment. This backend creates isolated venvs for + safe execution of agent skills. + +Functions: + get_backend: Get or create the configured backend instance + initialize_backend: One-time backend initialization at app startup +""" + +from __future__ import annotations + +import hashlib +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +from deepagents.backends import LocalShellBackend +from deepagents.backends.filesystem import FilesystemBackend +from deepagents.backends.protocol import EditResult, FileUploadResponse, WriteResult + +from deep_agent.src.agent.config import agent_config +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger(log_level=settings.PYTHON_LOG_LEVEL) + +_SYSTEM_PATH = "/usr/local/bin:/usr/bin:/bin" +_PASSTHROUGH_VARS = ("HOME", "USER", "LANG", "LC_ALL", "TZ", "TERM") + + +_REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent + +_backend: LocalShellBackend | None = None + + +class ReadOnlyFilesystemBackend(FilesystemBackend): + """FilesystemBackend that rejects all write operations.""" + + def write(self, file_path: str, content: str) -> WriteResult: + """Reject write operations.""" + return WriteResult(error="Read-only backend: writes not permitted") + + def edit( + self, + file_path: str, + old_string: str, + new_string: str, + replace_all: bool = False, # noqa: FBT001, FBT002 + ) -> EditResult: + """Reject edit operations.""" + return EditResult(error="Read-only backend: edits not permitted") + + def upload_files(self, files: list[tuple[str, bytes]]) -> list[FileUploadResponse]: + """Reject upload operations.""" + return [ + FileUploadResponse(path=p, error="Read-only backend: uploads not permitted") + for p, _ in files + ] + + +def _base_python() -> str: + """Resolve the base (non-venv) Python so the agent venv is independent. + + Prefers the versioned binary (e.g. python3.12) to avoid picking up the + UBI9 system python3 → 3.9 symlink when the app runs inside a 3.12 venv. + """ + if sys.prefix != sys.base_prefix: + v = sys.version_info + base_bin = Path(sys.base_prefix) / "bin" + for name in (f"python{v.major}.{v.minor}", "python3"): + candidate = base_bin / name + if candidate.exists(): + return str(candidate) + return sys.executable + + +def _ensure_venv(root_dir: Path, pyproject: Path) -> Path: + """Create an isolated venv in user cache directory and install from *pyproject*. + + The venv directory is keyed by a hash of *root_dir* **and** the contents of + *pyproject* so a changed ``pyproject.toml`` triggers a reinstall. + + Uses /app/.cache/template-agent/venvs/ (or ~/.cache/ outside containers) to + avoid security risks with world-readable /tmp directories on shared hosts. + """ + project_hash = hashlib.sha256(str(root_dir.resolve()).encode()).hexdigest()[:12] + toml_hash = hashlib.sha256(pyproject.read_bytes()).hexdigest()[:8] + + # Prefer /app/.cache inside containers (always writable on OpenShift); + # fall back to /tmp then ~/.cache for local / non-container runs. + # OpenShift runs with arbitrary UID so Path.home() may not resolve. + app_cache = Path("/app/.cache") + if app_cache.parent.is_dir(): + base_cache = app_cache + else: + try: + base_cache = Path.home() / ".cache" + except (RuntimeError, KeyError): + base_cache = Path("/tmp/.cache") # noqa: S108 — OpenShift arbitrary UID fallback + cache_dir = base_cache / "template-agent" / "venvs" + cache_dir.mkdir(parents=True, exist_ok=True, mode=0o700) # User-only permissions + + venv_dir = cache_dir / f"agent-venv-{project_hash}" + stamp = venv_dir / ".toml_hash" + + needs_install = False + + if not (venv_dir / "bin" / "python").exists(): + base = _base_python() + logger.info(f"Creating agent venv at {venv_dir} (python: {base})") + subprocess.run( + [base, "-m", "venv", "--clear", str(venv_dir)], + check=True, + capture_output=True, + text=True, + ) + needs_install = True + + if not needs_install and stamp.exists() and stamp.read_text() == toml_hash: + logger.info(f"Agent venv up-to-date ({venv_dir})") + return venv_dir + + # If pyproject.toml changed, clear the venv to remove stale dependencies + if stamp.exists() and stamp.read_text() != toml_hash: + base = _base_python() + logger.info(f"pyproject.toml changed — clearing venv at {venv_dir}") + subprocess.run( + [base, "-m", "venv", "--clear", str(venv_dir)], + check=True, + capture_output=True, + text=True, + ) + + pkg_dir = venv_dir / "_pkg" + pkg_dir.mkdir(exist_ok=True) + shutil.copy2(pyproject, pkg_dir / "pyproject.toml") + + pip = str(venv_dir / "bin" / "pip") + logger.info(f"Installing dependencies from {pyproject.name}") + result = subprocess.run( + [pip, "install", "--quiet", str(pkg_dir)], + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise RuntimeError(f"pip install failed: {result.stderr.strip()}") + + stamp.write_text(toml_hash) + return venv_dir + + +def _build_env(venv_dir: Path, extra: dict[str, str] | None = None) -> dict[str, str]: + """Minimal env: allowlisted host vars + venv activation + optional overrides.""" + env = {k: os.environ[k] for k in _PASSTHROUGH_VARS if k in os.environ} + env["VIRTUAL_ENV"] = str(venv_dir) + env["PATH"] = f"{venv_dir}/bin:{_SYSTEM_PATH}" + if extra: + env.update(extra) + return env + + +def create_backend( + root_dir: Path, + pyproject: Path, + *, + timeout: int = 120, + max_output_bytes: int = 100_000, + extra_env: dict[str, str] | None = None, +) -> LocalShellBackend: + """Create a :class:`LocalShellBackend` backed by an isolated agent venv. + + Args: + root_dir: Shell working directory. + pyproject: Path to a ``pyproject.toml`` whose dependencies are installed. + timeout: Default per-command timeout in seconds. + max_output_bytes: Max captured output before truncation. + extra_env: Extra env vars (highest priority). + """ + if not pyproject.is_file(): + raise FileNotFoundError(f"pyproject.toml not found: {pyproject}") + + venv_dir = _ensure_venv(root_dir, pyproject) + env = _build_env(venv_dir, extra_env) + + logger.info(f"Backend ready — venv={venv_dir}, pyproject={pyproject}") + return LocalShellBackend( + root_dir=str(root_dir), + virtual_mode=False, + timeout=timeout, + max_output_bytes=max_output_bytes, + env=env, + ) + + +def get_backend( + root_dir: Path | None = None, + pyproject: Path | None = None, + *, + timeout: int = 120, + max_output_bytes: int = 100_000, + extra_env: dict[str, str] | None = None, +) -> LocalShellBackend: + """Return the singleton backend, creating it on the first call. + + Subsequent calls return the same instance regardless of arguments. + When *root_dir* or *pyproject* are ``None`` the module-level defaults + (``_REPO_ROOT`` / ``agent_config.get_pyproject_path()``) are used. + """ + global _backend # noqa: PLW0603 + if _backend is None: + _backend = create_backend( + root_dir or _REPO_ROOT, + pyproject or agent_config.get_pyproject_path(), + timeout=timeout, + max_output_bytes=max_output_bytes, + extra_env=extra_env, + ) + return _backend + + +def get_configured_backend() -> LocalShellBackend | Any: + """Return the backend configured by filesystem.yaml or agent.yaml. + + Reads the backend type from config and builds the appropriate backend: + - state: StateBackend (thread-scoped scratch, recommended for production) + - composite: CompositeBackend (routes paths to different backends) + - store: StoreBackend (cross-thread persistent via LangGraph Store) + - local_shell: LocalShellBackend (local dev only — NOT for deployed agents) + + Falls back to StateBackend if config is missing or invalid. + """ + config_path = agent_config.base_dir / "filesystem.yaml" + if config_path.is_file(): + from deep_agent.src.agent.config.filesystem import load_filesystem_config + + fs_config = load_filesystem_config(config_path) + else: + fs_config = agent_config.get_filesystem_config() + + backend_type = fs_config.backend.type + + if backend_type == "state": + return _build_state_backend() + + if backend_type == "store": + return _build_store_backend(fs_config) + + if backend_type == "composite": + return _build_composite_backend(fs_config) + + if backend_type == "local_shell": + logger.warning( + "LocalShellBackend accesses the host directly. " + "Do NOT use in deployed agents (OpenShift, LangSmith, etc.). " + "Set backend.type to 'state' or 'composite' for production." + ) + return get_backend( + timeout=fs_config.backend.local_shell.timeout, + max_output_bytes=fs_config.backend.local_shell.max_output_bytes, + ) + + # Fallback for any backend type not explicitly handled above + logger.warning("Unknown backend type '%s', falling back to state", backend_type) # type: ignore[unreachable] + return _build_state_backend() + + +def _build_state_backend() -> Any: + """Build a StateBackend factory (thread-scoped scratch space). + + Recommended for production. Files persist across turns within a thread + via checkpointer but are not shared across threads. + + Returns the StateBackend class as a factory — create_deep_agent calls it + with ToolRuntime at execution time. + """ + try: + from deepagents.backends.state import StateBackend + + logger.info("Using StateBackend (thread-scoped scratch)") + return StateBackend + except ImportError: + logger.warning("StateBackend not available, falling back to LocalShellBackend") + return get_backend() + + +def _build_store_backend(fs_config: Any) -> Any: + """Build a StoreBackend (cross-thread persistent via LangGraph Store). + + Scope determines namespace partitioning: + - user: per-user private memory (recommended) + - assistant: shared across all users of one assistant + - org: shared across all users and assistants + """ + try: + from deepagents.backends.store import StoreBackend + + scope = getattr(fs_config.backend, "store", None) + scope_name = scope.scope if scope else "user" + + namespace_factories = { + "user": lambda rt: ( + rt.server_info.assistant_id, + rt.server_info.user.identity, + ), + "assistant": lambda rt: (rt.server_info.assistant_id,), + "org": lambda rt: (rt.context.org_id,), + } + + namespace = namespace_factories.get(scope_name) + if namespace is None: + logger.warning("Unknown store scope '%s', using 'user'", scope_name) + namespace = namespace_factories["user"] + + logger.info("Using StoreBackend (scope=%s)", scope_name) + return StoreBackend(namespace=namespace) + except ImportError: + logger.warning("StoreBackend not available, falling back to StateBackend") + return _build_state_backend() + + +def _build_composite_backend(fs_config: Any) -> Any: + """Return a factory that builds a CompositeBackend at request time. + + StateBackend and StoreBackend require ToolRuntime (only available per-request), + so we return a callable. ReadOnlyFilesystemBackend and LocalShellBackend are + built eagerly since they don't need runtime. + """ + # --- Eager: backends that don't need runtime --- + eager_routes: dict[str, Any] = {} + + for path_prefix, backend_name in fs_config.backend.routes.items(): + if backend_name == "filesystem_readonly": + dir_name = path_prefix.strip("/") + eager_routes[path_prefix] = _build_filesystem_readonly_backend( + agent_config.base_dir / dir_name + ) + + if any(v == "local_shell" for v in fs_config.backend.routes.values()): + logger.warning( + "local_shell in composite routes — not recommended for production" + ) + local_shell_backend = get_backend( + timeout=fs_config.backend.local_shell.timeout, + max_output_bytes=fs_config.backend.local_shell.max_output_bytes, + ) + for path_prefix, backend_name in fs_config.backend.routes.items(): + if backend_name == "local_shell": + eager_routes[path_prefix] = local_shell_backend + + # --- Deferred config (captured for use inside factory) --- + store_route_prefixes = [ + p for p, v in fs_config.backend.routes.items() if v == "store" + ] + store_scope: str | None = None + if store_route_prefixes: + scope = getattr(fs_config.backend, "store", None) + store_scope = scope.scope if scope else "user" + + logger.info( + "Prepared CompositeBackend factory: %d eager route(s), %d deferred route(s)", + len(eager_routes), + len(store_route_prefixes), + ) + + # --- Factory: called per-request with ToolRuntime --- + def factory(runtime: Any) -> Any: + """Build a CompositeBackend when invoked by create_deep_agent with ToolRuntime. + + We instantiate StateBackend/StoreBackend here (rather than returning + bare classes) because CompositeBackend needs composed *instances* — + this factory IS the protocol-compliant callable that create_deep_agent expects. + """ + from deepagents.backends.composite import CompositeBackend + from deepagents.backends.state import StateBackend + + state_backend = StateBackend(runtime) + + routes: dict[str, Any] = dict(eager_routes) + + if store_route_prefixes: + try: + from deepagents.backends.store import StoreBackend + + namespace_factories = { + "user": lambda rt: ( + rt.server_info.assistant_id, + rt.server_info.user.identity, + ), + "assistant": lambda rt: (rt.server_info.assistant_id,), + "org": lambda rt: (rt.context.org_id,), + } + ns = namespace_factories.get( + store_scope or "user", namespace_factories["user"] + ) + store_backend = StoreBackend(runtime, namespace=ns) + for prefix in store_route_prefixes: + routes[prefix] = store_backend + except ImportError: + logger.warning( + "StoreBackend not available — store routes will use StateBackend" + ) + for prefix in store_route_prefixes: + routes[prefix] = state_backend + + known_types = {"filesystem_readonly", "local_shell", "store", "state"} + for path_prefix, backend_name in fs_config.backend.routes.items(): + if backend_name not in known_types: + logger.warning( + "Unknown backend '%s' in route for '%s'", backend_name, path_prefix + ) + + default_backend = routes.pop("/", state_backend) + return CompositeBackend(default=default_backend, routes=routes) + + return factory + + +def _build_filesystem_readonly_backend(root_dir: Path) -> ReadOnlyFilesystemBackend: + """Build a read-only FilesystemBackend jailed to root_dir. + + Uses virtual_mode=True to jail all paths within the given directory. + Write/edit/upload operations are explicitly blocked for defense-in-depth. + + Args: + root_dir: Directory to use as the filesystem root. Derived from + the route prefix in agent.yaml (e.g., "/skills/" → base_dir/skills). + """ + if not root_dir.is_dir(): + logger.warning( + "Directory does not exist: %s — reads will return empty results", + root_dir, + ) + + logger.info( + "Using ReadOnlyFilesystemBackend (root=%s, virtual_mode=True)", root_dir + ) + return ReadOnlyFilesystemBackend(root_dir=str(root_dir), virtual_mode=True) diff --git a/deep_agent/src/infrastructure/mcp.py b/deep_agent/src/infrastructure/mcp.py new file mode 100644 index 00000000..aa43b10d --- /dev/null +++ b/deep_agent/src/infrastructure/mcp.py @@ -0,0 +1,10 @@ +"""MCP client — re-export from aegra runtime layer. + +This module moved to deep_agent.aegra.mcp as part of the runtime +consolidation. This shim preserves backward compatibility. +""" + +from deep_agent.aegra.mcp import ( # noqa: F401 + get_mcp_tools, + refresh_access_token, +) diff --git a/deep_agent/src/infrastructure/middleware.py b/deep_agent/src/infrastructure/middleware.py new file mode 100644 index 00000000..54ca2351 --- /dev/null +++ b/deep_agent/src/infrastructure/middleware.py @@ -0,0 +1,337 @@ +"""Middleware builder for deepagents integration. + +Converts ResolvedMiddlewareConfig into a list of AgentMiddleware instances +that can be passed to create_deep_agent(middleware=...). + +This module is the bridge between declarative YAML config and the deepagents +middleware API. Template-agent users never import or call this directly. +""" + +from __future__ import annotations + +import importlib +from typing import Any + +from deep_agent.src.agent.config.middleware import ResolvedMiddlewareConfig +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger(log_level=settings.PYTHON_LOG_LEVEL) + + +def build_audit_middleware( + *, + mcp_tool_names: frozenset[str] | None = None, + subagent: str | None = None, + agent: str | None = None, +) -> Any | None: + """Return AuditMiddleware when platform audit is enabled, else None.""" + from deep_agent.src.audit.config import is_audit_enabled + + if not is_audit_enabled(): + return None + + from deep_agent.src.audit.middleware import AuditMiddleware + + return AuditMiddleware( + mcp_tool_names=mcp_tool_names or frozenset(), + subagent=subagent, + agent=agent, + ) + + +def _mcp_tool_names_from_tools(tools: list[Any]) -> frozenset[str]: + return frozenset(getattr(t, "name", "") for t in tools if getattr(t, "name", None)) + + +def build_middleware_list( + resolved: ResolvedMiddlewareConfig, + *, + model: Any | None = None, + backend: Any | None = None, + mcp_tool_names: frozenset[str] | None = None, +) -> list[Any]: + """Build a list of middleware instances from resolved config. + + Only instantiates middleware that deepagents does NOT auto-include. + The auto-included middleware (SubAgentMiddleware, SummarizationMiddleware, + PatchToolCallsMiddleware, FilesystemMiddleware, TodoListMiddleware) are + handled by create_deep_agent() itself. + + Args: + resolved: Fully resolved middleware configuration for this agent. + model: Chat model instance for SummarizationToolMiddleware. + backend: Backend instance for SummarizationToolMiddleware. + mcp_tool_names: MCP tool names for platform audit classification. + + Returns: + List of middleware instances to pass as middleware= parameter. + Empty list means only deepagents defaults apply. + """ + middlewares: list[Any] = [] + + audit_mw = build_audit_middleware(mcp_tool_names=mcp_tool_names) + if audit_mw is not None: + middlewares.append(audit_mw) + + if not settings.MIDDLEWARE_ENABLED: + logger.info("Middleware disabled via MIDDLEWARE_ENABLED=false") + return middlewares + + if resolved.summarization_tool_enabled: + _append_if_built( + middlewares, + _build_summarization_tool_middleware(model=model, backend=backend), + ) + + _append_guardrails(middlewares, resolved) + + if resolved.code_execution.enabled: + _append_if_built( + middlewares, + _build_code_execution(resolved.code_execution), + ) + + for dotted_path in resolved.extra_middleware: + _append_if_built(middlewares, _import_middleware(dotted_path)) + + if middlewares: + logger.info("Built %d extra middleware instance(s)", len(middlewares)) + return middlewares + + +def _append_if_built(target: list[Any], mw: Any | None) -> None: + """Append middleware to list if it was built successfully.""" + if mw is not None: + target.append(mw) + + +def _append_guardrails(target: list[Any], resolved: ResolvedMiddlewareConfig) -> None: + """Build and append all production guardrail middleware.""" + if resolved.model_call_limit.enabled: + _append_if_built( + target, _build_model_call_limit(resolved.model_call_limit.run_limit) + ) + + if resolved.tool_call_limit.enabled: + _append_if_built( + target, _build_tool_call_limit(resolved.tool_call_limit.run_limit) + ) + + if resolved.model_retry.enabled: + _append_if_built(target, _build_model_retry(resolved.model_retry)) + + if resolved.model_fallback.enabled and resolved.model_fallback.fallback_model: + _append_if_built( + target, _build_model_fallback(resolved.model_fallback.fallback_model) + ) + + if resolved.tool_retry.enabled and resolved.tool_retry.tools: + _append_if_built(target, _build_tool_retry(resolved.tool_retry)) + + if resolved.pii.enabled and resolved.pii.rules: + target.extend(_build_pii_middleware(resolved.pii)) + + +def build_excluded_middleware( + resolved: ResolvedMiddlewareConfig, +) -> list[str]: + """Build the list of middleware to exclude from deepagents defaults. + + Used when registering HarnessProfiles or passing to create_deep_agent + via profile configuration. + + Args: + resolved: Resolved middleware config. + + Returns: + List of middleware class names to exclude. + """ + excluded: list[str] = list(resolved.excluded_middleware) + + if not resolved.patch_tool_calls_enabled: + excluded.append("PatchToolCallsMiddleware") + + return excluded + + +def resolve_memory_param( + resolved: ResolvedMiddlewareConfig, +) -> list[str] | None: + """Resolve the memory= parameter for create_deep_agent(). + + MemoryMiddleware is auto-included when memory= is provided. + This function returns the namespaces list or None to disable. + + Args: + resolved: Resolved middleware config. + + Returns: + List of memory namespace strings, or None if memory is disabled. + """ + if not settings.MIDDLEWARE_ENABLED: + return None + if not resolved.memory_enabled: + return None + return resolved.memory_namespaces or None + + +def _build_model_call_limit(run_limit: int) -> Any | None: + """Build ModelCallLimitMiddleware to cap LLM calls per run.""" + try: + from langchain.agents.middleware import ModelCallLimitMiddleware + + return ModelCallLimitMiddleware(run_limit=run_limit) + except ImportError: + logger.debug("ModelCallLimitMiddleware not available") + return None + + +def _build_tool_call_limit(run_limit: int) -> Any | None: + """Build ToolCallLimitMiddleware to cap tool calls per run.""" + try: + from langchain.agents.middleware import ToolCallLimitMiddleware + + return ToolCallLimitMiddleware(run_limit=run_limit) + except ImportError: + logger.debug("ToolCallLimitMiddleware not available") + return None + + +def _build_model_retry(config: Any) -> Any | None: + """Build ModelRetryMiddleware for transient failure recovery.""" + try: + from langchain.agents.middleware import ModelRetryMiddleware + + return ModelRetryMiddleware( + max_retries=config.max_retries, + backoff_factor=config.backoff_factor, + initial_delay=config.initial_delay, + ) + except ImportError: + logger.debug("ModelRetryMiddleware not available") + return None + + +def _build_model_fallback(fallback_model: str) -> Any | None: + """Build ModelFallbackMiddleware to switch models on primary failure.""" + try: + from langchain.agents.middleware import ModelFallbackMiddleware + + return ModelFallbackMiddleware(fallback_model) + except ImportError: + logger.debug("ModelFallbackMiddleware not available") + return None + except Exception as e: + logger.warning("ModelFallbackMiddleware init failed (check model auth): %s", e) + return None + + +def _build_tool_retry(config: Any) -> Any | None: + """Build ToolRetryMiddleware for specific tools.""" + try: + from langchain.agents.middleware import ToolRetryMiddleware + + return ToolRetryMiddleware( + max_retries=config.max_retries, + tools=config.tools, + ) + except ImportError: + logger.debug("ToolRetryMiddleware not available") + return None + + +def _build_pii_middleware(config: Any) -> list[Any]: + """Build PIIMiddleware instances for each PII rule.""" + results: list[Any] = [] + try: + from langchain.agents.middleware import PIIMiddleware + + for rule in config.rules: + try: + results.append( + PIIMiddleware( + rule.type, strategy=rule.strategy, apply_to_input=True + ) + ) + except (ValueError, TypeError) as e: + logger.warning("Skipping PII rule '%s': %s", rule.type, e) + except ImportError: + logger.debug("PIIMiddleware not available") + return results + + +def _build_summarization_tool_middleware( + *, + model: Any | None = None, + backend: Any | None = None, +) -> Any | None: + """Build SummarizationToolMiddleware instance. + + This gives the agent a tool to proactively trigger summarization + at opportune moments (e.g., between tasks) rather than only at + fixed token thresholds. + """ + if model is None or backend is None: + logger.warning("Summarization tool requires model and backend; skipping") + return None + try: + from deepagents.middleware.summarization import ( + create_summarization_tool_middleware, + ) + + return create_summarization_tool_middleware(model, backend) + except ImportError: + logger.debug( + "SummarizationToolMiddleware not available in this deepagents version" + ) + return None + except Exception as e: + logger.warning("Failed to create SummarizationToolMiddleware: %s", e) + return None + + +def _build_code_execution(config: Any) -> Any | None: + """Build CodeExecutionMiddleware for sandboxed code execution.""" + try: + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + return CodeExecutionMiddleware(config=config) + except ImportError: + logger.debug( + "CodeExecutionMiddleware not available (missing kubernetes package?)" + ) + return None + except Exception as e: + logger.warning("Failed to create CodeExecutionMiddleware: %s", e) + return None + + +def _import_middleware(dotted_path: str) -> Any | None: + """Import and instantiate a middleware from a dotted path. + + Format: "module.path:ClassName" or "module.path:factory_function" + + Args: + dotted_path: Dotted import path with colon-separated attribute. + + Returns: + Instantiated middleware, or None on failure. + """ + try: + if ":" not in dotted_path: + logger.warning( + "Invalid middleware path '%s' — expected 'module:Class'", dotted_path + ) + return None + + module_path, attr_name = dotted_path.rsplit(":", 1) + module = importlib.import_module(module_path) + factory_or_class = getattr(module, attr_name) + + if callable(factory_or_class): + return factory_or_class() + return factory_or_class + except Exception as e: + logger.warning("Failed to import middleware '%s': %s", dotted_path, e) + return None diff --git a/deep_agent/src/infrastructure/permissions.py b/deep_agent/src/infrastructure/permissions.py new file mode 100644 index 00000000..c4cb59cc --- /dev/null +++ b/deep_agent/src/infrastructure/permissions.py @@ -0,0 +1,61 @@ +"""Filesystem permissions builder. + +Converts declarative permission rules from filesystem.yaml into +deepagents FilesystemPermission instances that are passed to +create_deep_agent(permissions=...). + +Template-agent users only edit YAML. This module handles the conversion. +""" + +from __future__ import annotations + +from typing import Any + +from deep_agent.src.agent.config.filesystem import ( + FilesystemFileConfig, +) +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def build_permissions( + config: FilesystemFileConfig, +) -> list[Any] | None: + """Build FilesystemPermission list from config. + + Args: + config: Parsed filesystem.yaml config. + + Returns: + List of FilesystemPermission instances, or None if no rules defined. + None means deepagents uses its default (all operations allowed). + """ + if not config.permissions: + return None + + try: + from deepagents.middleware.filesystem import FilesystemPermission + except ImportError: + logger.warning( + "FilesystemPermission not available — permissions config ignored" + ) + return None + + permissions: list[Any] = [] + + for rule in config.permissions: + try: + perm = FilesystemPermission( + operations=rule.operations, + paths=rule.paths, + mode=rule.mode, + ) + permissions.append(perm) + except Exception as e: + logger.warning("Skipping invalid permission rule %r: %s", rule, e) + + if permissions: + logger.info("Built %d filesystem permission rule(s)", len(permissions)) + + return permissions or None diff --git a/deep_agent/src/infrastructure/providers.py b/deep_agent/src/infrastructure/providers.py new file mode 100644 index 00000000..bdb92d57 --- /dev/null +++ b/deep_agent/src/infrastructure/providers.py @@ -0,0 +1,153 @@ +"""Provider and harness profile registration. + +Reads the validated ProvidersFileConfig and registers ProviderProfile +and HarnessProfile instances with the deepagents profile registry. + +Also provides resolve_model_from_config() which picks between the +legacy create_model() path and deepagents resolve_model() based on +the resolve_strategy setting. + +Template-agent users never call this directly — it's wired by graph.py +and factory.py at agent creation time. +""" + +from __future__ import annotations + +from typing import Any + +from deep_agent.src.agent.config.providers import ProvidersFileConfig +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_profiles_registered: bool = False + + +def register_profiles_from_config(config: ProvidersFileConfig) -> None: + """Register ProviderProfile and HarnessProfile instances from config. + + Idempotent — only registers once per process lifetime. + + Args: + config: Validated providers.yaml config. + """ + global _profiles_registered # noqa: PLW0603 + if _profiles_registered: + return + + _register_provider_profiles(config) + _register_harness_profiles(config) + _profiles_registered = True + + +def resolve_model_from_config( + model_name: str, + config: ProvidersFileConfig, + *, + temperature: float = 0.0, + max_output_tokens: int | None = None, +) -> Any: + """Resolve a model string to a BaseChatModel using the configured strategy. + + Args: + model_name: Model name (e.g., "gemini-2.5-pro" or "openai:gpt-5.4"). + config: Validated providers config. + temperature: Model temperature. + max_output_tokens: Maximum output tokens. + + Returns: + A BaseChatModel instance. + """ + if config.resolve_strategy == "deepagents": + return _resolve_via_deepagents(model_name) + + return _resolve_via_legacy(model_name, temperature, max_output_tokens) + + +def _resolve_via_legacy( + model_name: str, + temperature: float, + max_output_tokens: int | None, +) -> Any: + """Legacy resolution — use our hardcoded create_model() factory.""" + from deep_agent.src.cache.model_cache import get_or_create_model + + return get_or_create_model( + model_name=model_name, + temperature=temperature, + max_output_tokens=max_output_tokens, + ) + + +def _resolve_via_deepagents(model_name: str) -> Any: + """Deepagents resolution — use resolve_model() with registered profiles.""" + try: + from deepagents import resolve_model + + logger.info("Resolving model via deepagents: %s", model_name) + return resolve_model(model_name) + except ImportError: + logger.warning( + "deepagents.resolve_model not available — falling back to legacy" + ) + from deep_agent.src.cache.model_cache import get_or_create_model + + return get_or_create_model(model_name=model_name) + + +def _register_provider_profiles(config: ProvidersFileConfig) -> None: + """Register ProviderProfile instances for each configured provider.""" + if not config.providers: + return + + try: + from deepagents import ProviderProfile, register_provider_profile + except ImportError: + logger.debug("deepagents profiles API not available — skipping registration") + return + + for provider_key, provider_cfg in config.providers.items(): + try: + profile = ProviderProfile(init_kwargs=provider_cfg.init_kwargs) + register_provider_profile(provider_key, profile) + logger.info("Registered ProviderProfile: %s", provider_key) + except Exception as e: + logger.warning( + "Failed to register ProviderProfile '%s': %s", provider_key, e + ) + + +def _register_harness_profiles(config: ProvidersFileConfig) -> None: + """Register HarnessProfile instances for each configured model.""" + if not config.harness_profiles: + return + + try: + from deepagents import ( + GeneralPurposeSubagentProfile, + HarnessProfile, + register_harness_profile, + ) + except ImportError: + logger.debug("deepagents profiles API not available — skipping registration") + return + + for model_key, harness_cfg in config.harness_profiles.items(): + try: + gp_config = harness_cfg.general_purpose_subagent + gp_profile = GeneralPurposeSubagentProfile( + enabled=gp_config.enabled, + description=gp_config.description, + system_prompt=gp_config.system_prompt, + ) + + profile = HarnessProfile( + system_prompt_suffix=harness_cfg.system_prompt_suffix or None, + excluded_tools=frozenset(harness_cfg.excluded_tools), + excluded_middleware=frozenset(harness_cfg.excluded_middleware), + general_purpose_subagent=gp_profile, + ) + register_harness_profile(model_key, profile) + logger.info("Registered HarnessProfile: %s", model_key) + except Exception as e: + logger.warning("Failed to register HarnessProfile '%s': %s", model_key, e) diff --git a/deep_agent/src/infrastructure/subagents.py b/deep_agent/src/infrastructure/subagents.py new file mode 100644 index 00000000..5ff1f9fd --- /dev/null +++ b/deep_agent/src/infrastructure/subagents.py @@ -0,0 +1,541 @@ +"""Subagent loading from configuration files. + +This module builds SubAgent instances from the markdown configuration files in +config/subagents/. It reads each subagent's config, resolves their tools +and skills, creates appropriate LLM instances, and returns ready-to-use SubAgent +objects for the orchestrator. + +Supports three agent types via the ``type`` field in frontmatter: + - ``default``: Standard SubAgent (in-process, synchronous delegation) + - ``compiled``: CompiledSubAgent (pre-compiled graph, reused across requests) + - ``async``: AsyncSubAgent (remote Agent Protocol server, background tasks) + +Functions: + load_subagents: Build all subagents from config/subagents/*.md +""" + +from typing import Any, cast + +from deepagents import SubAgent +from deepagents.middleware.subagents import CompiledSubAgent + +try: + from deepagents.middleware.async_subagents import AsyncSubAgent +except ImportError: + AsyncSubAgent = None + +from deep_agent.src.agent.config import agent_config +from deep_agent.src.agent.config.model import ( + ModelSpec, + infer_provider, + parse_model_config, +) +from deep_agent.src.agent.config.resolver import to_virtual_skill_paths +from deep_agent.src.cache.model_cache import get_or_create_model_from_spec +from deep_agent.src.exceptions import LLMError, SubAgentError +from deep_agent.src.infrastructure.middleware import ( + _mcp_tool_names_from_tools, + build_audit_middleware, +) +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger(log_level=settings.PYTHON_LOG_LEVEL) + +VALID_AGENT_TYPES = ("default", "compiled", "async") + + +def load_subagents( + tools: list[Any], +) -> list[Any] | None: + """Build subagents from pre-loaded configurations. + + Reads the ``type`` field from each subagent's frontmatter to determine + which agent class to construct: + - ``default`` / missing → SubAgent (standard in-process delegation) + - ``compiled`` → CompiledSubAgent (pre-compiled graph as Runnable) + - ``async`` → AsyncSubAgent (remote Agent Protocol server) + + Subagents that don't specify a ``model`` inherit the orchestrator's model. + Subagents that don't specify ``mcps`` inherit the orchestrator's MCPs + (which determines tool visibility). + + Args: + tools: List of available MCP tools. + + Returns: + List of configured subagent instances, or None if no subagents configured. + + Raises: + SubAgentError: If a subagent fails to build (missing model, bad config). + """ + all_subagent_configs: dict[str, dict[str, Any]] = ( + agent_config.get_all_subagent_configs() + ) + + if not all_subagent_configs: + logger.warning("No subagent configurations found") + return None + + orchestrator_cfg = agent_config.get_orchestrator_config() + + logger.info(f"Building {len(all_subagent_configs)} subagent(s)") + + subagents_list: list[Any] = [] + + for name, agent_cfg in all_subagent_configs.items(): + _inherit_from_orchestrator(agent_cfg, orchestrator_cfg, name) + try: + sa = _build_single_subagent(name, agent_cfg, tools) + subagents_list.append(sa) + except (ValueError, LLMError) as e: + raise SubAgentError(f"Failed to build subagent '{name}': {e}") from e + except Exception as e: + raise SubAgentError( + f"Unexpected error building subagent '{name}': {e}" + ) from e + + logger.info(f"Built {len(subagents_list)} subagent(s) successfully") + return subagents_list + + +_DEFAULT_FALLBACK_MODEL = "gemini-3.1-pro-preview" + + +def _inherit_from_orchestrator( + agent_cfg: dict[str, Any], + orchestrator_cfg: dict[str, Any], + name: str, +) -> None: + """Fill in missing model/mcps from the parent orchestrator config. + + Mutates *agent_cfg* in place. Model inheritance follows these rules: + 1. If subagent has no model → use orchestrator model (no fallback) + 2. If subagent has model but no fallback → use orchestrator model as fallback + 3. If subagent has model with fallback → keep as-is + + Falls back to _DEFAULT_FALLBACK_MODEL when neither the subagent nor the + orchestrator specifies a model. + """ + parent_model = orchestrator_cfg.get("model") + subagent_model = agent_cfg.get("model") + + if not subagent_model: + # Case 1: No subagent model → inherit orchestrator model (no fallback) + if parent_model: + logger.info( + "Subagent '%s' inheriting model from orchestrator: %s", + name, + parent_model, + ) + agent_cfg["model"] = parent_model + else: + logger.warning( + "Subagent '%s' has no model and orchestrator has no model — " + "falling back to default: %s", + name, + _DEFAULT_FALLBACK_MODEL, + ) + agent_cfg["model"] = _DEFAULT_FALLBACK_MODEL + elif parent_model: + # Case 2 & 3: Subagent has model → inject orchestrator as fallback if missing + agent_cfg["model"] = _inject_fallback_if_missing( + subagent_model, parent_model, name + ) + + if not agent_cfg.get("mcps"): + parent_mcps = orchestrator_cfg.get("mcps", []) + if parent_mcps: + logger.info( + "Subagent '%s' inheriting %d MCP(s) from orchestrator", + name, + len(parent_mcps), + ) + agent_cfg["mcps"] = list(parent_mcps) + + +def _normalize_model_to_dict( + raw_model: Any, + strip_fallback: bool = False, +) -> dict[str, Any] | Any: + """Normalize model config (string or dict) to dict format. + + Args: + raw_model: Model config in string or dict format. + strip_fallback: If True, remove fallback key from dict configs. + + Returns: + Normalized dict with provider and name keys, or original value if invalid type. + """ + if isinstance(raw_model, str): + return { + "provider": infer_provider(raw_model).value, + "name": raw_model, + } + elif isinstance(raw_model, dict): + result = dict(raw_model) # Copy to avoid mutation + if strip_fallback and "fallback" in result: + del result["fallback"] + return result + logger.warning( + "Invalid model config type: %s, letting parse_model_config handle error", + type(raw_model).__name__, + ) + return raw_model + + +def _inject_fallback_if_missing( + subagent_model: str | dict[str, Any], + parent_model: str | dict[str, Any], + name: str, +) -> str | dict[str, Any]: + """Inject orchestrator model as fallback if subagent model has no fallback. + + Args: + subagent_model: Subagent's model config (string or dict). + parent_model: Orchestrator's model config. + name: Subagent name (for logging). + + Returns: + Normalized model config dict with fallback injected if needed, + or original value if invalid type (will fail in parse_model_config). + """ + # Normalize subagent model to dict + model_dict = _normalize_model_to_dict(subagent_model) + if not isinstance(model_dict, dict): + return cast("str | dict[str, Any]", model_dict) + + # Case 3: Subagent already has fallback → keep as-is + if "fallback" in model_dict: + return model_dict + + # Case 2: Subagent has no fallback → inject orchestrator as fallback + logger.debug( + "Subagent '%s' inheriting orchestrator model as fallback: %s", + name, + parent_model, + ) + + # Normalize parent model to dict for fallback (strip nested fallback) + fallback_dict = _normalize_model_to_dict(parent_model, strip_fallback=True) + if not isinstance(fallback_dict, dict): + # Invalid parent, skip fallback injection + return model_dict + + model_dict["fallback"] = fallback_dict + return model_dict + + +def _create_primary_model(spec: ModelSpec) -> object: + """Create only the primary BaseChatModel from a ModelSpec, without fallback wrapper. + + Uses the model cache for efficient reuse. Fallbacks should be handled via + LangChain's ModelFallbackMiddleware. + + Args: + spec: Parsed model specification (fallback config ignored). + + Returns: + A BaseChatModel instance for the primary model only. + """ + # Create a spec without fallback for the primary model + primary_spec = ModelSpec(provider=spec.provider, name=spec.name, fallback=None) + + # Use the cache to get or create the model + return get_or_create_model_from_spec(primary_spec) + + +def _resolve_subagent_model(agent_cfg: dict[str, Any]) -> object: + """Parse frontmatter model config and return only the primary BaseChatModel. + + Strips any fallback configuration since deepagents doesn't support RunnableWithFallbacks. + Fallback handling should be done via LangChain's ModelFallbackMiddleware instead. + """ + raw_model = agent_cfg.get("model") + if raw_model is None: + raise ValueError("missing required 'model' field in frontmatter") + + # Parse model spec (may include fallback config) + spec = parse_model_config(raw_model) + + # Create only the primary model using the cache + return _create_primary_model(spec) + + +def _format_model_log(spec: ModelSpec) -> str: + """Format model spec for log messages.""" + return spec.display_name() + + +def _build_fallback_middleware(spec: ModelSpec) -> list[Any]: + """Build ModelFallbackMiddleware with BaseChatModel if spec has fallback configured. + + Creates the fallback model using get_or_create_model_from_spec to preserve custom + initialization logic (MAAS base URLs, Vertex credentials, etc) and enable caching. + + Args: + spec: Parsed model specification. + + Returns: + List containing ModelFallbackMiddleware if fallback exists, empty list otherwise. + """ + if spec.fallback is None: + return [] + + try: + from langchain.agents.middleware import ModelFallbackMiddleware + except ImportError: + logger.warning( + "ModelFallbackMiddleware not available, skipping fallback configuration" + ) + return [] + + # Create fallback model using the model cache + fallback_model = _create_primary_model(spec.fallback) + + middleware = ModelFallbackMiddleware(fallback_model) + logger.info( + "Configured fallback middleware: %s -> %s", + spec.display_name(), + spec.fallback.display_name(), + ) + return [middleware] + + +def _subagent_middleware( + name: str, + resolved_tools: list[Any], + fallback_mw: list[Any], +) -> list[Any] | None: + """Merge audit middleware (outermost) with optional fallback middleware.""" + middleware: list[Any] = [] + audit_mw = build_audit_middleware( + mcp_tool_names=_mcp_tool_names_from_tools(resolved_tools), + agent=name, + ) + if audit_mw is not None: + middleware.append(audit_mw) + middleware.extend(fallback_mw) + return middleware or None + + +def _build_single_subagent( + name: str, + agent_cfg: dict[str, Any], + tools: list[Any], +) -> Any: + """Build a single subagent from its configuration. + + Dispatches to the appropriate builder based on the ``type`` field. + + Args: + name: Subagent name (from config filename). + agent_cfg: Parsed frontmatter config for this subagent. + tools: Available MCP tools for tool resolution. + + Returns: + Configured subagent instance (SubAgent, CompiledSubAgent, or AsyncSubAgent). + + Raises: + ValueError: If required fields are missing or type is invalid. + LLMError: If model creation fails. + """ + agent_type: str = agent_cfg.get("type", "default") + if agent_type not in VALID_AGENT_TYPES: + raise ValueError( + f"Subagent '{name}' has invalid type '{agent_type}'. " + f"Valid types: {VALID_AGENT_TYPES}" + ) + + if agent_type == "async": + return _build_async_subagent(name, agent_cfg) + if agent_type == "compiled": + return _build_compiled_subagent(name, agent_cfg, tools) + return _build_default_subagent(name, agent_cfg, tools) + + +def _build_default_subagent( + name: str, + agent_cfg: dict[str, Any], + tools: list[Any], +) -> SubAgent: + """Build a standard SubAgent (in-process delegation).""" + if not agent_cfg.get("model"): + raise ValueError( + f"Subagent '{name}' is missing required 'model' field in frontmatter" + ) + + spec = parse_model_config(agent_cfg["model"]) + logger.info( + "Subagent '%s' [default] using model: %s", name, _format_model_log(spec) + ) + + tool_names: list[str] = agent_cfg.get("tools", []) + mcp_names: list[str] = agent_cfg.get("mcps", []) + + if tool_names: + resolved_tools: list[Any] = agent_config.resolve_tools( + tool_names, tools, agent_name=name + ) + elif mcp_names and tools: + logger.info( + "Subagent '%s' declared MCP servers %s but no explicit tools; " + "exposing all %d available MCP tool(s)", + name, + mcp_names, + len(tools), + ) + resolved_tools = list(tools) + else: + resolved_tools = [] + + skill_paths: list[str] = agent_cfg.get("skill_paths", []) + + # Build fallback middleware if spec has fallback configured + fallback_mw = _build_fallback_middleware(spec) + + subagent_params: dict[str, Any] = { + "name": name, + "model": _resolve_subagent_model(agent_cfg), + "description": agent_cfg.get("description", ""), + "system_prompt": agent_cfg.get("body", ""), + } + + if resolved_tools: + subagent_params["tools"] = resolved_tools + if skill_paths: + subagent_params["skills"] = to_virtual_skill_paths(skill_paths) + middleware = _subagent_middleware(name, resolved_tools, fallback_mw) + if middleware: + subagent_params["middleware"] = middleware + + return SubAgent(**subagent_params) + + +def _build_compiled_subagent( + name: str, + agent_cfg: dict[str, Any], + tools: list[Any], +) -> CompiledSubAgent: + """Build a CompiledSubAgent (pre-compiled graph as Runnable). + + Creates a full deep agent graph for this subagent and wraps it as a + CompiledSubAgent. The compiled graph is reused across requests, + providing better performance for frequently-invoked subagents. + """ + from deepagents import create_deep_agent + + from deep_agent.src.infrastructure.backend import get_configured_backend + + if not agent_cfg.get("model"): + raise ValueError( + f"Subagent '{name}' (compiled) is missing required 'model' field" + ) + + spec = parse_model_config(agent_cfg["model"]) + logger.info( + "Subagent '%s' [compiled] using model: %s", name, _format_model_log(spec) + ) + + tool_names: list[str] = agent_cfg.get("tools", []) + mcp_names: list[str] = agent_cfg.get("mcps", []) + + if tool_names: + resolved_tools: list[Any] = agent_config.resolve_tools( + tool_names, tools, agent_name=name + ) + elif mcp_names and tools: + logger.info( + "Subagent '%s' [compiled] declared MCP servers %s but no explicit tools; " + "exposing all %d available MCP tool(s)", + name, + mcp_names, + len(tools), + ) + resolved_tools = list(tools) + else: + resolved_tools = [] + skill_paths: list[str] = agent_cfg.get("skill_paths", []) + + # Build fallback middleware if spec has fallback configured + fallback_mw = _build_fallback_middleware(spec) + + create_kwargs = { + "name": name, + "model": _resolve_subagent_model(agent_cfg), + "system_prompt": agent_cfg.get("body", ""), + "tools": resolved_tools or None, + "skills": to_virtual_skill_paths(skill_paths) if skill_paths else None, + "backend": get_configured_backend(), + } + + middleware = _subagent_middleware(name, resolved_tools, fallback_mw) + if middleware: + create_kwargs["middleware"] = middleware + + compiled_graph = create_deep_agent(**create_kwargs) + + return CompiledSubAgent( + name=name, + description=agent_cfg.get("description", ""), + runnable=compiled_graph, + ) + + +def _build_async_subagent( + name: str, + agent_cfg: dict[str, Any], +) -> Any: + """Build an AsyncSubAgent (remote Agent Protocol server). + + Requires ``graph_id`` in frontmatter. Optionally accepts ``url`` + for the remote endpoint. + + Auth headers are resolved from environment variables (OpenShift Secrets), + never from frontmatter config. The env var name follows the convention: + ``ASYNC_SUBAGENT__TOKEN`` (uppercased, hyphens → underscores). + """ + if AsyncSubAgent is None: + raise ValueError( + f"Subagent '{name}' (async) requires deepagents with async support. " + "Upgrade deepagents or remove this subagent config." + ) + + graph_id: str | None = agent_cfg.get("graph_id") + if not graph_id: + raise ValueError( + f"Subagent '{name}' (async) is missing required 'graph_id' field" + ) + + logger.info(f"Subagent '{name}' [async] connecting to graph: {graph_id}") + + params: dict[str, Any] = { + "name": name, + "description": agent_cfg.get("description", ""), + "graph_id": graph_id, + } + + url: str | None = agent_cfg.get("url") + if url: + params["url"] = url + + headers = _resolve_async_headers(name) + if headers: + params["headers"] = headers + + return AsyncSubAgent(**params) + + +def _resolve_async_headers(name: str) -> dict[str, str] | None: + """Resolve auth headers for an async subagent from environment. + + Convention: ASYNC_SUBAGENT__TOKEN env var → Authorization header. + Secrets come from OpenShift Secrets mounted as env vars. + """ + import os + + env_key = f"ASYNC_SUBAGENT_{name.upper().replace('-', '_')}_TOKEN" + token = os.environ.get(env_key) + if token: + return {"Authorization": f"Bearer {token}"} + return None diff --git a/deep_agent/src/infrastructure/tool_access.py b/deep_agent/src/infrastructure/tool_access.py new file mode 100644 index 00000000..8e076ba8 --- /dev/null +++ b/deep_agent/src/infrastructure/tool_access.py @@ -0,0 +1,219 @@ +"""Tool access control for subagents. + +Provides filtering, denial, and approval wrapping for subagent tool sets. +Used by subagent builders to enforce per-subagent tool access policies +declared in frontmatter config. +""" + +from __future__ import annotations + +import inspect +from typing import Any + +from langgraph.types import interrupt + +from deep_agent.src.exceptions import AppException, ErrorCodes +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger(log_level=settings.PYTHON_LOG_LEVEL) + + +def filter_denied_tools( + tools: list[Any], + denied_names: list[str], + agent_name: str, +) -> list[Any]: + """Remove denied tools from a resolved tool list. + + Args: + tools: List of resolved tool objects (must have a ``.name`` attr). + denied_names: Tool names to exclude. + agent_name: Subagent name (for logging). + + Returns: + New list with denied tools removed. Original order preserved. + """ + if not denied_names: + return tools + + denied_set = set(denied_names) + result: list[Any] = [] + + for tool in tools: + if tool.name in denied_set: + logger.info("DENIED tool removed", agent=agent_name, tool=tool.name) + else: + result.append(tool) + + return result + + +def apply_tool_approval( + tools: list[Any], + approval_names: list[str], + agent_name: str, +) -> list[Any]: + """Wrap tools requiring human approval with interrupt guards. + + Tools whose names appear in *approval_names* are wrapped so that + invoking them triggers a LangGraph ``interrupt()`` before the + underlying function executes. + + Args: + tools: List of resolved tool objects (must have a ``.name`` attr). + approval_names: Tool names that require approval. + agent_name: Subagent name (for logging and interrupt payload). + + Returns: + New list with matching tools wrapped. Non-matching tools + are passed through unchanged. + """ + if not approval_names: + return tools + + approval_set = set(approval_names) + tool_names_present = {t.name for t in tools} + + # Warn about approval names that reference tools not in the resolved set. + for name in approval_names: + if name not in tool_names_present: + logger.warning( + "APPROVAL references unknown tool", agent=agent_name, tool=name + ) + + result: list[Any] = [] + for tool in tools: + if tool.name in approval_set: + result.append(_wrap_tool_with_approval(tool, agent_name)) + else: + result.append(tool) + + return result + + +def _build_hitl_payload(tool: Any, agent_name: str) -> dict[str, Any]: + """Build an HITL-compatible interrupt payload for tool approval. + + Returns a dict matching the ``HITLInterruptValue`` schema expected by the + frontend: ``{ action_requests: [...], review_configs: [...] }``. + """ + description = ( + f"Tool approval required: subagent '{agent_name}' wants to call " + f"'{tool.name}'. Approve or reject this tool call." + ) + return { + "action_requests": [ + { + "name": tool.name, + "args": {"agent": agent_name, "description": description}, + } + ], + "review_configs": [ + { + "action_name": tool.name, + "allowed_decisions": ["approve", "reject"], + } + ], + } + + +def _is_approved(decision: Any) -> bool: + """Check whether the HITL resume value signals approval. + + Handles both the structured format (list of dicts with ``type`` key) + sent by the frontend and simple string values used in tests. + """ + if isinstance(decision, list): + return any(isinstance(d, dict) and d.get("type") == "approve" for d in decision) + return str(decision).strip().lower() == "approved" + + +def _wrap_tool_with_approval(tool: Any, agent_name: str) -> Any: + """Wrap a single tool so it calls ``interrupt()`` before execution. + + Follows the same wrapping pattern used in + ``deep_agent.aegra.mcp_tool_auth._wrap_single_tool``: replaces the + coroutine (async) or func (sync) with a wrapper that issues an + interrupt, then proceeds only if the human approves. + + The interrupt payload uses the HITL dict format so the frontend's + ``InterruptBanner`` can render it natively. + + Args: + tool: A LangChain ``StructuredTool`` (or compatible) tool object. + agent_name: Subagent name (for the interrupt payload message). + + Returns: + A copy of the tool with its callable replaced by the approval wrapper. + """ + payload = _build_hitl_payload(tool, agent_name) + + coroutine = getattr(tool, "coroutine", None) + func = getattr(tool, "func", None) + + if inspect.iscoroutinefunction(coroutine): + original_coroutine = coroutine + + async def wrapped_coroutine(**kwargs: Any) -> Any: + decision = interrupt(payload) + if _is_approved(decision): + return await original_coroutine(**kwargs) + return f"Tool '{tool.name}' was rejected by the user." + + try: + return tool.model_copy(update={"coroutine": wrapped_coroutine}) + except Exception: + tool.coroutine = wrapped_coroutine + return tool + + if func is not None and inspect.isfunction(func): + original_func = func + + def wrapped_func(**kwargs: Any) -> Any: + decision = interrupt(payload) + if _is_approved(decision): + return original_func(**kwargs) + return f"Tool '{tool.name}' was rejected by the user." + + try: + return tool.model_copy(update={"func": wrapped_func}) + except Exception: + tool.func = wrapped_func + return tool + + # Tool has neither coroutine nor func — return as-is. + return tool + + +def migrate_tools_field(config: dict[str, Any], agent_name: str) -> dict[str, Any]: + """Migrate deprecated ``tools`` key to ``allowed_tools``. + + If the config contains a ``tools`` key but no ``allowed_tools``, the + value is moved to ``allowed_tools`` and a deprecation warning is logged. + + Args: + config: Parsed frontmatter config dict (mutated in place). + agent_name: Subagent name (for logging). + + Returns: + The same *config* dict (for chaining convenience). + + Raises: + AppException: If both ``tools`` and ``allowed_tools`` are present. + """ + has_tools = "tools" in config + has_allowed = "allowed_tools" in config + + if has_tools and has_allowed: + raise AppException( + f"Subagent '{agent_name}': config contains both 'tools' and " + f"'allowed_tools'. Remove the deprecated 'tools' field.", + error_code=ErrorCodes.CONFIGURATION_VALIDATION_ERROR, + ) + + if has_tools and not has_allowed: + config["allowed_tools"] = config.pop("tools") + logger.info("COMPAT migrated tools->allowed_tools", agent=agent_name) + + return config diff --git a/deep_agent/src/memory/__init__.py b/deep_agent/src/memory/__init__.py new file mode 100644 index 00000000..600015e2 --- /dev/null +++ b/deep_agent/src/memory/__init__.py @@ -0,0 +1,7 @@ +"""Memory management — consolidation, decay, clustering, and scheduling. + +All operations run as **background jobs** via APScheduler. +Nothing in this package runs in the request path. + +Feature flag: ``MEMORY_CONSOLIDATION_ENABLED`` (+ individual layer flags). +""" diff --git a/deep_agent/src/memory/clustering.py b/deep_agent/src/memory/clustering.py new file mode 100644 index 00000000..eaa4a29d --- /dev/null +++ b/deep_agent/src/memory/clustering.py @@ -0,0 +1,178 @@ +"""Semantic clustering of user memories. + +Groups similar memories by content similarity using token-based +cosine similarity (TF-IDF style, zero API calls). Assigns a +``cluster_id`` to each memory in the database. + +Runs as a **background job** — never in the request path. +No embedding API calls — pure local computation. +""" + +import math +import uuid +from collections import Counter, defaultdict + +from deep_agent.src.memory.config import memory_settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def _tokenize(text: str) -> list[str]: + """Simple whitespace tokeniser with lowercasing.""" + return text.lower().split() + + +def _build_tfidf(documents: list[str]) -> list[dict[str, float]]: + """Build TF-IDF vectors for a list of documents. + + Returns a list of {token: tfidf_weight} dicts, one per document. + """ + n = len(documents) + if n == 0: + return [] + + doc_tokens = [_tokenize(d) for d in documents] + + df: Counter[str] = Counter() + for tokens in doc_tokens: + df.update(set(tokens)) + + vectors: list[dict[str, float]] = [] + for tokens in doc_tokens: + tf: Counter[str] = Counter(tokens) + total = len(tokens) or 1 + vec: dict[str, float] = {} + for term, count in tf.items(): + idf = math.log((n + 1) / (df[term] + 1)) + 1 + vec[term] = (count / total) * idf + vectors.append(vec) + + return vectors + + +def _cosine_sim(a: dict[str, float], b: dict[str, float]) -> float: + """Cosine similarity between two sparse vectors.""" + common = set(a.keys()) & set(b.keys()) + if not common: + return 0.0 + dot = sum(a[k] * b[k] for k in common) + mag_a = math.sqrt(sum(v * v for v in a.values())) + mag_b = math.sqrt(sum(v * v for v in b.values())) + if mag_a == 0 or mag_b == 0: + return 0.0 + return dot / (mag_a * mag_b) + + +def cluster_memories( + contents: list[str], + threshold: float | None = None, +) -> list[list[int]]: + """Cluster memory indices by TF-IDF cosine similarity. + + Uses single-linkage agglomerative clustering (union-find). + + Args: + contents: List of memory content strings. + threshold: Minimum similarity to merge (default from config). + + Returns: + List of clusters (each a list of indices). Singletons are excluded. + """ + threshold = threshold or memory_settings.MEMORY_CLUSTER_THRESHOLD + vectors = _build_tfidf(contents) + n = len(vectors) + + parent = list(range(n)) + + def find(i: int) -> int: + while parent[i] != i: + parent[i] = parent[parent[i]] + i = parent[i] + return i + + def union(i: int, j: int) -> None: + ri, rj = find(i), find(j) + if ri != rj: + parent[ri] = rj + + for i in range(n): + for j in range(i + 1, n): + if _cosine_sim(vectors[i], vectors[j]) >= threshold: + union(i, j) + + groups: defaultdict[int, list[int]] = defaultdict(list) + for i in range(n): + groups[find(i)].append(i) + + return [g for g in groups.values() if len(g) >= 2] + + +async def cluster_user_memories( + database_uri: str, + user_id: str, +) -> int: + """Assign cluster_id to similar memories for a user. + + Returns the number of clusters created. + """ + import psycopg + from psycopg.rows import dict_row + + async with await psycopg.AsyncConnection.connect( + database_uri, row_factory=dict_row + ) as conn: + cur = await conn.execute( + "SELECT id, content FROM user_memories " + "WHERE user_id = %s ORDER BY created_at DESC", + (user_id,), + ) + memories = [dict(row) for row in await cur.fetchall()] + + if len(memories) < 2: + return 0 + + contents = [m["content"] for m in memories] + clusters = cluster_memories(contents) + + if not clusters: + return 0 + + for group in clusters: + cid = str(uuid.uuid4()) + for idx in group: + await conn.execute( + "UPDATE user_memories SET cluster_id = %s WHERE id = %s", + (cid, str(memories[idx]["id"])), + ) + + await conn.commit() + logger.info( + "Clustered user %s: %d cluster(s) from %d memories", + user_id[:8], + len(clusters), + len(memories), + ) + return len(clusters) + + +async def cluster_all_users(database_uri: str) -> int: + """Run clustering across all users. Returns total clusters created.""" + if not memory_settings.is_enabled("clustering"): + logger.debug("Memory clustering disabled — skipping") + return 0 + + import psycopg + + async with await psycopg.AsyncConnection.connect(database_uri) as conn: + cur = await conn.execute("SELECT DISTINCT user_id FROM user_memories") + user_ids = [row[0] for row in await cur.fetchall()] + + total = 0 + for uid in user_ids: + total += await cluster_user_memories(database_uri, uid) + + logger.info( + "Clustering complete: %d clusters across %d users", total, len(user_ids) + ) + return total diff --git a/deep_agent/src/memory/config.py b/deep_agent/src/memory/config.py new file mode 100644 index 00000000..2982f7cd --- /dev/null +++ b/deep_agent/src/memory/config.py @@ -0,0 +1,47 @@ +"""Memory management configuration with feature flags. + +All memory background processing is disabled by default. +Enable via environment variables. + +Environment variables: + MEMORY_CONSOLIDATION_ENABLED: Master switch (default: false) + MEMORY_DECAY_ENABLED: Exponential decay scoring (default: false) + MEMORY_CLUSTERING_ENABLED: Semantic clustering (default: false) + MEMORY_RELATIONSHIPS_ENABLED: Relationship inference (default: false) + MEMORY_SCHEDULER_INTERVAL_HOURS: Job run interval (default: 6) + MEMORY_MAX_INJECT: Max memories injected into prompt (default: 20) + MEMORY_DECAY_LAMBDA: Decay rate — higher = faster fade (default: 0.05) + MEMORY_CLUSTER_THRESHOLD: Similarity threshold for clustering (default: 0.4) + MEMORY_CONSOLIDATION_MIN_CLUSTER: Min cluster size to consolidate (default: 3) +""" + +from pydantic import Field +from pydantic_settings import BaseSettings + + +class MemorySettings(BaseSettings): + """Feature-flagged memory management configuration.""" + + MEMORY_CONSOLIDATION_ENABLED: bool = Field(default=False) + MEMORY_DECAY_ENABLED: bool = Field(default=False) + MEMORY_CLUSTERING_ENABLED: bool = Field(default=False) + MEMORY_RELATIONSHIPS_ENABLED: bool = Field(default=False) + + MEMORY_SCHEDULER_INTERVAL_HOURS: int = Field(default=6, ge=1, le=168) + MEMORY_MAX_INJECT: int = Field(default=20, ge=1, le=200) + MEMORY_DECAY_LAMBDA: float = Field(default=0.05, ge=0.001, le=1.0) + MEMORY_CLUSTER_THRESHOLD: float = Field(default=0.4, ge=0.1, le=0.95) + MEMORY_CONSOLIDATION_MIN_CLUSTER: int = Field(default=3, ge=2, le=20) + + def is_enabled(self, layer: str) -> bool: + """Check if a specific memory layer is active. + + Master switch must be on for any layer to activate. + """ + if not self.MEMORY_CONSOLIDATION_ENABLED: + return False + flag = getattr(self, f"MEMORY_{layer.upper()}_ENABLED", False) + return bool(flag) + + +memory_settings = MemorySettings() diff --git a/deep_agent/src/memory/consolidation.py b/deep_agent/src/memory/consolidation.py new file mode 100644 index 00000000..1d264e24 --- /dev/null +++ b/deep_agent/src/memory/consolidation.py @@ -0,0 +1,178 @@ +"""Memory consolidation — merge duplicate/similar memories. + +When a user accumulates many memories, this module: +1. Detects near-duplicates (exact or fuzzy match) +2. Merges them into a single consolidated memory +3. Deletes the originals + +Runs as a **background job** — never in the request path. +""" + +import re +from collections import defaultdict + +from deep_agent.src.memory.config import memory_settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def _normalise(text: str) -> str: + """Lowercase, strip punctuation, collapse whitespace.""" + text = text.lower().strip() + text = re.sub(r"[^\w\s]", "", text) + return re.sub(r"\s+", " ", text) + + +def _token_set(text: str) -> set[str]: + """Return a set of normalised tokens.""" + return set(_normalise(text).split()) + + +def token_similarity(a: str, b: str) -> float: + """Jaccard similarity between token sets of two strings.""" + sa, sb = _token_set(a), _token_set(b) + if not sa or not sb: + return 0.0 + return len(sa & sb) / len(sa | sb) + + +def find_duplicates( + memories: list[dict[str, str]], + threshold: float | None = None, +) -> list[list[int]]: + """Group memory indices that are near-duplicates. + + Args: + memories: List of dicts with at least a ``content`` key. + threshold: Similarity threshold (default from config). + + Returns: + List of groups, where each group is a list of indices + into *memories* that should be consolidated. + """ + threshold = threshold or memory_settings.MEMORY_CLUSTER_THRESHOLD + + n = len(memories) + parent = list(range(n)) + + def find(i: int) -> int: + while parent[i] != i: + parent[i] = parent[parent[i]] + i = parent[i] + return i + + def union(i: int, j: int) -> None: + ri, rj = find(i), find(j) + if ri != rj: + parent[ri] = rj + + for i in range(n): + for j in range(i + 1, n): + sim = token_similarity(memories[i]["content"], memories[j]["content"]) + if sim >= threshold: + union(i, j) + + groups: defaultdict[int, list[int]] = defaultdict(list) + for i in range(n): + groups[find(i)].append(i) + + return [g for g in groups.values() if len(g) >= 2] + + +def pick_representative( + memories: list[dict[str, str]], + indices: list[int], +) -> int: + """Choose the best memory from a duplicate group. + + Picks the longest content (most informative), breaking ties + by highest score. + """ + best = indices[0] + for idx in indices[1:]: + cur_len = len(memories[idx]["content"]) + best_len = len(memories[best]["content"]) + if cur_len > best_len: + best = idx + elif cur_len == best_len: + cur_score = float(memories[idx].get("score", 0)) + best_score = float(memories[best].get("score", 0)) + if cur_score > best_score: + best = idx + return best + + +async def consolidate_user_memories( + database_uri: str, + user_id: str, +) -> int: + """Consolidate duplicate memories for a single user. + + Returns the number of memories deleted. + """ + import psycopg + from psycopg.rows import dict_row + + async with await psycopg.AsyncConnection.connect( + database_uri, row_factory=dict_row + ) as conn: + cur = await conn.execute( + "SELECT id, content, score FROM user_memories " + "WHERE user_id = %s ORDER BY created_at DESC", + (user_id,), + ) + memories = [dict(row) for row in await cur.fetchall()] + + if len(memories) < 2: + return 0 + + groups = find_duplicates(memories) + if not groups: + return 0 + + deleted = 0 + for group in groups: + keep = pick_representative(memories, group) + to_delete = [i for i in group if i != keep] + for idx in to_delete: + await conn.execute( + "DELETE FROM user_memories WHERE id = %s", + (str(memories[idx]["id"]),), + ) + deleted += 1 + + if deleted: + await conn.commit() + logger.info( + "Consolidated user %s: deleted %d duplicate(s) from %d group(s)", + user_id[:8], + deleted, + len(groups), + ) + + return deleted + + +async def consolidate_all_users(database_uri: str) -> int: + """Run consolidation across all users. Returns total deletions.""" + if not memory_settings.MEMORY_CONSOLIDATION_ENABLED: + logger.debug("Memory consolidation disabled — skipping") + return 0 + + import psycopg + + async with await psycopg.AsyncConnection.connect(database_uri) as conn: + cur = await conn.execute("SELECT DISTINCT user_id FROM user_memories") + user_ids = [row[0] for row in await cur.fetchall()] + + total = 0 + for uid in user_ids: + total += await consolidate_user_memories(database_uri, uid) + + logger.info( + "Consolidation complete: %d total deletions across %d users", + total, + len(user_ids), + ) + return total diff --git a/deep_agent/src/memory/relationships.py b/deep_agent/src/memory/relationships.py new file mode 100644 index 00000000..c3152f5c --- /dev/null +++ b/deep_agent/src/memory/relationships.py @@ -0,0 +1,156 @@ +"""Relationship inference between user memories. + +Detects memories that share significant keywords or entities, +and stores links in the ``memory_relationships`` table so the +agent can surface related context. + +Runs as a **background job** — never in the request path. +No LLM calls — pure keyword overlap. +""" + +import re +from collections import Counter + +from deep_agent.src.memory.config import memory_settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +STOPWORDS = frozenset( + "a an the is are was were be been being have has had do does did " + "will would shall should may might can could i me my we our you your " + "he she it they them their this that these those of in to for with " + "on at by from as into through during before after above below " + "and or but not no nor so yet also very too".split() +) + +MIN_SHARED_KEYWORDS = 2 + + +def extract_keywords(text: str, top_n: int = 10) -> list[str]: + """Extract significant keywords from text. + + Strips stopwords, short tokens, and returns the most frequent + meaningful words. + """ + tokens = re.findall(r"\b[a-zA-Z]{3,}\b", text.lower()) + meaningful = [t for t in tokens if t not in STOPWORDS] + counts = Counter(meaningful) + return [word for word, _ in counts.most_common(top_n)] + + +def find_related_pairs( + memories: list[dict[str, str]], + min_shared: int = MIN_SHARED_KEYWORDS, +) -> list[tuple[int, int, list[str]]]: + """Find pairs of memories that share significant keywords. + + Args: + memories: List of dicts with ``content`` key. + min_shared: Minimum shared keywords to consider related. + + Returns: + List of (idx_a, idx_b, shared_keywords) tuples. + """ + keyword_sets = [set(extract_keywords(m["content"])) for m in memories] + pairs: list[tuple[int, int, list[str]]] = [] + + for i in range(len(memories)): + for j in range(i + 1, len(memories)): + shared = keyword_sets[i] & keyword_sets[j] + if len(shared) >= min_shared: + pairs.append((i, j, sorted(shared))) + + return pairs + + +async def infer_user_relationships( + database_uri: str, + user_id: str, +) -> int: + """Detect and store relationships between a user's memories. + + Returns the number of new relationships created. + """ + import psycopg + from psycopg.rows import dict_row + + async with await psycopg.AsyncConnection.connect( + database_uri, row_factory=dict_row + ) as conn: + await conn.execute( + """ + CREATE TABLE IF NOT EXISTS memory_relationships ( + memory_a UUID NOT NULL, + memory_b UUID NOT NULL, + keywords TEXT NOT NULL, + user_id TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (memory_a, memory_b) + ) + """ + ) + + cur = await conn.execute( + "SELECT id, content FROM user_memories " + "WHERE user_id = %s ORDER BY created_at DESC", + (user_id,), + ) + memories = [dict(row) for row in await cur.fetchall()] + + if len(memories) < 2: + return 0 + + pairs = find_related_pairs(memories) + if not pairs: + return 0 + + created = 0 + for i, j, keywords in pairs: + id_a = str(memories[i]["id"]) + id_b = str(memories[j]["id"]) + a, b = min(id_a, id_b), max(id_a, id_b) + try: + await conn.execute( + """ + INSERT INTO memory_relationships (memory_a, memory_b, keywords, user_id) + VALUES (%s, %s, %s, %s) + ON CONFLICT (memory_a, memory_b) DO NOTHING + """, + (a, b, ",".join(keywords), user_id), + ) + created += 1 + except Exception: + logger.debug("Relationship insert failed", exc_info=True) + + if created: + await conn.commit() + logger.info( + "Relationships for user %s: %d pair(s) from %d memories", + user_id[:8], + created, + len(memories), + ) + return created + + +async def infer_all_relationships(database_uri: str) -> int: + """Run relationship inference across all users.""" + if not memory_settings.is_enabled("relationships"): + logger.debug("Relationship inference disabled — skipping") + return 0 + + import psycopg + + async with await psycopg.AsyncConnection.connect(database_uri) as conn: + cur = await conn.execute("SELECT DISTINCT user_id FROM user_memories") + user_ids = [row[0] for row in await cur.fetchall()] + + total = 0 + for uid in user_ids: + total += await infer_user_relationships(database_uri, uid) + + logger.info( + "Relationships complete: %d pairs across %d users", total, len(user_ids) + ) + return total diff --git a/deep_agent/src/memory/scheduler.py b/deep_agent/src/memory/scheduler.py new file mode 100644 index 00000000..b92525f4 --- /dev/null +++ b/deep_agent/src/memory/scheduler.py @@ -0,0 +1,125 @@ +"""APScheduler-based background job scheduler for memory management. + +Uses a Redis-backed distributed lock so that only one replica +runs each job at a time (OpenShift multi-replica safe). + +When Redis is unavailable, falls back to in-process scheduling +(each pod runs independently — acceptable for idempotent jobs). + +Feature flag: ``MEMORY_CONSOLIDATION_ENABLED``. +""" + +from typing import Any + +from deep_agent.src.memory.config import memory_settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_scheduler: Any = None + + +async def start_scheduler(database_uri: str) -> bool: + """Start the background memory scheduler. + + Returns True if started, False if disabled or already running. + """ + global _scheduler # noqa: PLW0603 + + if not memory_settings.MEMORY_CONSOLIDATION_ENABLED: + logger.debug("Memory scheduler disabled — skipping") + return False + + if _scheduler is not None: + logger.debug("Memory scheduler already running") + return False + + try: + from apscheduler import AsyncScheduler + from apscheduler.triggers.interval import IntervalTrigger + + _scheduler = AsyncScheduler() + + interval = memory_settings.MEMORY_SCHEDULER_INTERVAL_HOURS + trigger = IntervalTrigger(hours=interval) + + await _scheduler.add_schedule( + _run_memory_jobs, + trigger, + id="memory-consolidation", + kwargs={"database_uri": database_uri}, + ) + + await _scheduler.start_in_background() + logger.info( + "Memory scheduler started (interval=%dh)", + interval, + ) + return True + except Exception: + logger.warning("Failed to start memory scheduler", exc_info=True) + _scheduler = None + return False + + +async def stop_scheduler() -> None: + """Gracefully stop the scheduler if running.""" + global _scheduler # noqa: PLW0603 + if _scheduler is not None: + try: + await _scheduler.stop() + logger.info("Memory scheduler stopped") + except Exception: + logger.debug("Scheduler stop error", exc_info=True) + _scheduler = None + + +async def _run_memory_jobs(database_uri: str) -> dict[str, int]: + """Execute all enabled memory background jobs. + + This is the single entry point called by the scheduler. + Each sub-job checks its own feature flag. + + Returns a summary dict of results. + """ + results: dict[str, int] = {} + + try: + from deep_agent.src.memory.scoring import decay_all_memories + + results["decay"] = await decay_all_memories(database_uri) + except Exception: + logger.error("Decay job failed", exc_info=True) + results["decay"] = -1 + + try: + from deep_agent.src.memory.consolidation import consolidate_all_users + + results["consolidation"] = await consolidate_all_users(database_uri) + except Exception: + logger.error("Consolidation job failed", exc_info=True) + results["consolidation"] = -1 + + try: + from deep_agent.src.memory.clustering import cluster_all_users + + results["clustering"] = await cluster_all_users(database_uri) + except Exception: + logger.error("Clustering job failed", exc_info=True) + results["clustering"] = -1 + + try: + from deep_agent.src.memory.relationships import infer_all_relationships + + results["relationships"] = await infer_all_relationships(database_uri) + except Exception: + logger.error("Relationships job failed", exc_info=True) + results["relationships"] = -1 + + logger.info("Memory jobs complete: %s", results) + return results + + +async def run_once(database_uri: str) -> dict[str, int]: + """Run all memory jobs once (for testing or manual trigger).""" + return await _run_memory_jobs(database_uri) diff --git a/deep_agent/src/memory/scoring.py b/deep_agent/src/memory/scoring.py new file mode 100644 index 00000000..d5b40cc7 --- /dev/null +++ b/deep_agent/src/memory/scoring.py @@ -0,0 +1,97 @@ +"""Exponential decay scoring for user memories. + +Memories lose relevance over time unless accessed. The score formula: + + score = base_score * e^(-λ * age_days) + access_boost + +Where: + - base_score: initial score (1.0 for new memories) + - λ (lambda): decay rate from MEMORY_DECAY_LAMBDA + - age_days: days since last update + - access_boost: small bump each time the memory is referenced + +This runs as a **background job** — never in the request path. +""" + +import math +from datetime import datetime, timezone + +from deep_agent.src.memory.config import memory_settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +ACCESS_BOOST = 0.1 +MIN_SCORE = 0.01 + + +def compute_decay_score( + base_score: float, + updated_at: datetime, + now: datetime | None = None, +) -> float: + """Compute the decayed score for a memory. + + Args: + base_score: The memory's current stored score. + updated_at: When the memory was last updated or accessed. + now: Current time (defaults to utcnow for testability). + + Returns: + Decayed score, floored at MIN_SCORE. + """ + now = now or datetime.now(timezone.utc) + if updated_at.tzinfo is None: + updated_at = updated_at.replace(tzinfo=timezone.utc) + + age_days = max((now - updated_at).total_seconds() / 86400, 0) + lam = memory_settings.MEMORY_DECAY_LAMBDA + score = base_score * math.exp(-lam * age_days) + return max(score, MIN_SCORE) + + +def apply_access_boost(current_score: float) -> float: + """Bump a memory's score when it's referenced in a conversation. + + Capped at 1.0 to prevent runaway scores. + """ + return min(current_score + ACCESS_BOOST, 1.0) + + +async def decay_all_memories(database_uri: str) -> int: + """Recalculate scores for all memories in the database. + + Returns the number of memories updated. + """ + import psycopg + from psycopg.rows import dict_row + + if not memory_settings.is_enabled("decay"): + logger.debug("Memory decay disabled — skipping") + return 0 + + now = datetime.now(timezone.utc) + updated = 0 + + async with await psycopg.AsyncConnection.connect( + database_uri, row_factory=dict_row + ) as conn: + cur = await conn.execute("SELECT id, score, updated_at FROM user_memories") + rows = await cur.fetchall() + + for row in rows: + old_score = float(row.get("score", 1.0) or 1.0) + new_score = compute_decay_score(old_score, row["updated_at"], now) + + if abs(new_score - old_score) > 0.001: + await conn.execute( + "UPDATE user_memories SET score = %s WHERE id = %s", + (new_score, str(row["id"])), + ) + updated += 1 + + if updated: + await conn.commit() + + logger.info("Decay scoring: updated %d / %d memories", updated, len(rows)) + return updated diff --git a/deep_agent/src/observability/__init__.py b/deep_agent/src/observability/__init__.py new file mode 100644 index 00000000..3a769d17 --- /dev/null +++ b/deep_agent/src/observability/__init__.py @@ -0,0 +1 @@ +"""Observability package.""" diff --git a/deep_agent/src/observability/otel_setup.py b/deep_agent/src/observability/otel_setup.py new file mode 100644 index 00000000..987cfd8b --- /dev/null +++ b/deep_agent/src/observability/otel_setup.py @@ -0,0 +1,129 @@ +"""OTLP metrics (otel-gateway) and traces bootstrap.""" + +from __future__ import annotations + +from typing import Any + +_fastapi_instrumented = False +_metrics_initialized = False +_traces_initialized = False +_service_resource = None + + +def _service_resource_for(settings: Any) -> Any: + """Return a shared OTEL resource for metrics and traces providers.""" + global _service_resource # noqa: PLW0603 + + if _service_resource is None: + from opentelemetry.sdk.resources import Resource + + _service_resource = Resource.create( + {"service.name": settings.OTEL_SERVICE_NAME} + ) + return _service_resource + + +def _otlp_grpc_exporter_kwargs(endpoint: str, settings: Any) -> dict[str, Any]: + """Build OTLP/gRPC exporter kwargs from a config endpoint string.""" + raw = endpoint.strip() + kwargs: dict[str, Any] = {} + lower = raw.lower() + if lower.startswith("https://"): + kwargs["endpoint"] = raw[len("https://") :] + elif lower.startswith("http://"): + kwargs["endpoint"] = raw[len("http://") :] + kwargs["insecure"] = True + else: + kwargs["endpoint"] = raw + kwargs["insecure"] = True + token = (getattr(settings, "OTEL_AUTH_TOKEN", None) or "").strip() + if token and not token.startswith("<"): + kwargs["headers"] = (("authorization", f"Bearer {token}"),) + return kwargs + + +def _instrument_fastapi(app: Any, log: Any) -> None: + """HTTP server metrics + traces (needs MeterProvider and/or TracerProvider set first).""" + global _fastapi_instrumented # noqa: PLW0603 + if _fastapi_instrumented: + return + try: + from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor + + FastAPIInstrumentor.instrument_app(app) + _fastapi_instrumented = True + except Exception: + log.warning("template_agent_fastapi_instrument_failed", exc_info=True) + + +def setup_otel_metrics(settings: Any, log: Any) -> None: + """Export OTLP metrics to OTEL_EXPORTER_OTLP_ENDPOINT when enabled.""" + global _metrics_initialized # noqa: PLW0603 + + if _metrics_initialized: + return + if not settings.ENABLE_OTEL_METRICS or not settings.OTEL_EXPORTER_OTLP_ENDPOINT: + return + + try: + from opentelemetry import metrics + from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( + OTLPMetricExporter, + ) + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + + resource = _service_resource_for(settings) + exporter = OTLPMetricExporter( + **_otlp_grpc_exporter_kwargs(settings.OTEL_EXPORTER_OTLP_ENDPOINT, settings) + ) + reader = PeriodicExportingMetricReader( + exporter, + export_interval_millis=settings.OTEL_METRIC_EXPORT_INTERVAL_MILLIS, + ) + provider = MeterProvider(resource=resource, metric_readers=[reader]) + metrics.set_meter_provider(provider) + _metrics_initialized = True + log.info("template_agent_otel_metrics_export_enabled") + except Exception: + log.warning("template_agent_otel_metrics_export_failed", exc_info=True) + + +def setup_otel_traces(app: Any, settings: Any, log: Any) -> None: + """Export OTLP traces when enabled; instrument FastAPI when metrics or traces on.""" + global _traces_initialized # noqa: PLW0603 + + traces_endpoint = settings.resolved_otel_traces_endpoint() + metrics_on = bool( + settings.ENABLE_OTEL_METRICS and settings.OTEL_EXPORTER_OTLP_ENDPOINT + ) + traces_on = bool(settings.otel_traces_active() and traces_endpoint) + + if not metrics_on and not traces_on: + return + + if traces_on and not _traces_initialized: + try: + from opentelemetry import trace + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter, + ) + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + + resource = _service_resource_for(settings) + provider = TracerProvider(resource=resource) + processor = BatchSpanProcessor( + OTLPSpanExporter( + **_otlp_grpc_exporter_kwargs(traces_endpoint, settings) + ) + ) + provider.add_span_processor(processor) + trace.set_tracer_provider(provider) + _traces_initialized = True + log.info("template_agent_otel_tracing_enabled") + except Exception: + log.warning("template_agent_otel_tracing_failed", exc_info=True) + + if metrics_on or traces_on: + _instrument_fastapi(app, log) diff --git a/deep_agent/src/personalization/__init__.py b/deep_agent/src/personalization/__init__.py new file mode 100644 index 00000000..be7228f7 --- /dev/null +++ b/deep_agent/src/personalization/__init__.py @@ -0,0 +1,17 @@ +"""User personalization: memories, custom rules, and prompt injection. + +Provides per-user memory and rule storage (Postgres-backed) and a +prompt injector that prepends personalization context to the agent's +system prompt at graph-creation time. +""" + +from deep_agent.src.personalization.injector import inject_personalization +from deep_agent.src.personalization.models import Memory, Rule +from deep_agent.src.personalization.repository import PersonalizationRepository + +__all__ = [ + "Memory", + "Rule", + "PersonalizationRepository", + "inject_personalization", +] diff --git a/deep_agent/src/personalization/injector.py b/deep_agent/src/personalization/injector.py new file mode 100644 index 00000000..f159fab0 --- /dev/null +++ b/deep_agent/src/personalization/injector.py @@ -0,0 +1,53 @@ +"""Inject user personalization context into the agent system prompt. + +The injector appends two optional blocks to the base system prompt: + +1. **User Memories** — facts the agent should recall across sessions +2. **User Rules** — custom instructions that shape agent behaviour + +Both blocks are omitted when the corresponding list is empty, keeping +the prompt clean for users who haven't configured personalization. +""" + +from __future__ import annotations + + +def inject_personalization( + system_prompt: str, + memories: list[str], + rules: list[str], +) -> str: + """Return *system_prompt* enriched with personalization blocks. + + Args: + system_prompt: The base system prompt from config. + memories: Plain-text user memories (newest first). + rules: Plain-text user rules / custom instructions. + + Returns: + The enriched system prompt. Unchanged if both lists are empty. + """ + sections: list[str] = [] + + if memories: + lines = "\n".join(f"- {m}" for m in memories) + sections.append( + f"## User Memories\n\n" + f"The following facts were saved by the user across prior sessions. " + f"Treat them as persistent context — reference them when relevant " + f"but do not repeat them verbatim unless asked.\n\n{lines}" + ) + + if rules: + lines = "\n".join(f"- {r}" for r in rules) + sections.append( + f"## User Custom Instructions\n\n" + f"The user has defined the following rules. Follow them for every " + f"response unless they conflict with safety guidelines.\n\n{lines}" + ) + + if not sections: + return system_prompt + + personalization_block = "\n\n---\n\n".join(sections) + return f"{system_prompt}\n\n---\n\n{personalization_block}" diff --git a/deep_agent/src/personalization/models.py b/deep_agent/src/personalization/models.py new file mode 100644 index 00000000..702b7b10 --- /dev/null +++ b/deep_agent/src/personalization/models.py @@ -0,0 +1,31 @@ +"""Pydantic models for user personalization data.""" + +from __future__ import annotations + +import uuid +from datetime import datetime + +from pydantic import BaseModel, Field + + +class Memory(BaseModel): + """A single user memory — a fact the agent should recall across sessions.""" + + id: uuid.UUID = Field(default_factory=uuid.uuid4) + user_id: str + content: str + score: float = Field(default=1.0) + cluster_id: uuid.UUID | None = Field(default=None) + created_at: datetime = Field(default_factory=datetime.utcnow) + updated_at: datetime = Field(default_factory=datetime.utcnow) + + +class Rule(BaseModel): + """A user-defined custom instruction that shapes agent behaviour.""" + + id: uuid.UUID = Field(default_factory=uuid.uuid4) + user_id: str + content: str + is_active: bool = True + created_at: datetime = Field(default_factory=datetime.utcnow) + updated_at: datetime = Field(default_factory=datetime.utcnow) diff --git a/deep_agent/src/personalization/repository.py b/deep_agent/src/personalization/repository.py new file mode 100644 index 00000000..db294bf5 --- /dev/null +++ b/deep_agent/src/personalization/repository.py @@ -0,0 +1,192 @@ +"""Async Postgres repository for user memories and rules. + +Uses ``psycopg`` (async) against the same database that stores +LangGraph checkpoints. Tables are created lazily on first use via +:meth:`PersonalizationRepository.ensure_tables`. +""" + +from __future__ import annotations + +import uuid +from datetime import datetime + +import psycopg +from psycopg.rows import dict_row + +from deep_agent.src.personalization.models import Memory, Rule +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_TABLES_ENSURED = False + +CREATE_MEMORIES_TABLE = """ +CREATE TABLE IF NOT EXISTS user_memories ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id TEXT NOT NULL, + content TEXT NOT NULL, + score FLOAT NOT NULL DEFAULT 1.0, + cluster_id UUID, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE INDEX IF NOT EXISTS idx_user_memories_user_id + ON user_memories (user_id); +""" + +MIGRATE_MEMORIES_TABLE = """ +ALTER TABLE user_memories ADD COLUMN IF NOT EXISTS score FLOAT NOT NULL DEFAULT 1.0; +ALTER TABLE user_memories ADD COLUMN IF NOT EXISTS cluster_id UUID; +""" + +CREATE_RULES_TABLE = """ +CREATE TABLE IF NOT EXISTS user_rules ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id TEXT NOT NULL, + content TEXT NOT NULL, + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE INDEX IF NOT EXISTS idx_user_rules_user_id + ON user_rules (user_id); +""" + + +class PersonalizationRepository: + """Thin async wrapper around the personalization tables.""" + + def __init__(self, database_uri: str) -> None: + """Initialise with a Postgres connection URI.""" + self._uri = database_uri + + async def ensure_tables(self) -> None: + """Create personalization tables if they do not already exist.""" + global _TABLES_ENSURED # noqa: PLW0603 + if _TABLES_ENSURED: + return + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + await conn.execute(CREATE_MEMORIES_TABLE) + await conn.execute(CREATE_RULES_TABLE) + await conn.execute(MIGRATE_MEMORIES_TABLE) + await conn.commit() + _TABLES_ENSURED = True + logger.info("Personalization tables ensured") + + # ── Memories ────────────────────────────────────────────── + + async def list_memories(self, user_id: str) -> list[Memory]: + """Return all memories for *user_id*, newest first.""" + await self.ensure_tables() + async with await psycopg.AsyncConnection.connect( + self._uri, row_factory=dict_row + ) as conn: + cur = await conn.execute( + "SELECT * FROM user_memories WHERE user_id = %s ORDER BY created_at DESC", + (user_id,), + ) + return [Memory(**row) for row in await cur.fetchall()] + + async def list_top_memories(self, user_id: str, limit: int = 20) -> list[Memory]: + """Return top-N memories for *user_id*, ranked by score descending.""" + await self.ensure_tables() + async with await psycopg.AsyncConnection.connect( + self._uri, row_factory=dict_row + ) as conn: + cur = await conn.execute( + "SELECT * FROM user_memories WHERE user_id = %s " + "ORDER BY score DESC, updated_at DESC LIMIT %s", + (user_id, limit), + ) + return [Memory(**row) for row in await cur.fetchall()] + + async def create_memory(self, user_id: str, content: str) -> Memory: + """Insert a new memory and return the created model.""" + await self.ensure_tables() + mem = Memory(user_id=user_id, content=content) + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + await conn.execute( + "INSERT INTO user_memories (id, user_id, content, created_at, updated_at) " + "VALUES (%s, %s, %s, %s, %s)", + (str(mem.id), mem.user_id, mem.content, mem.created_at, mem.updated_at), + ) + await conn.commit() + return mem + + async def delete_memory(self, user_id: str, memory_id: uuid.UUID) -> bool: + """Delete a memory by id; return True if a row was removed.""" + await self.ensure_tables() + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + cur = await conn.execute( + "DELETE FROM user_memories WHERE id = %s AND user_id = %s", + (str(memory_id), user_id), + ) + await conn.commit() + return bool(cur.rowcount > 0) + + # ── Rules ───────────────────────────────────────────────── + + async def list_rules(self, user_id: str, *, active_only: bool = True) -> list[Rule]: + """Return rules for *user_id*, optionally filtering to active only.""" + await self.ensure_tables() + clause = " AND is_active = TRUE" if active_only else "" + async with await psycopg.AsyncConnection.connect( + self._uri, row_factory=dict_row + ) as conn: + cur = await conn.execute( + f"SELECT * FROM user_rules WHERE user_id = %s{clause} ORDER BY created_at DESC", + (user_id,), + ) + return [Rule(**row) for row in await cur.fetchall()] + + async def upsert_rule( + self, + user_id: str, + content: str, + rule_id: uuid.UUID | None = None, + is_active: bool = True, + ) -> Rule: + """Create or update a rule and return the model.""" + await self.ensure_tables() + now = datetime.utcnow() + rid = rule_id or uuid.uuid4() + rule = Rule( + id=rid, + user_id=user_id, + content=content, + is_active=is_active, + created_at=now, + updated_at=now, + ) + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + await conn.execute( + """ + INSERT INTO user_rules (id, user_id, content, is_active, created_at, updated_at) + VALUES (%s, %s, %s, %s, %s, %s) + ON CONFLICT (id) + DO UPDATE SET content = EXCLUDED.content, + is_active = EXCLUDED.is_active, + updated_at = EXCLUDED.updated_at + """, + ( + str(rule.id), + rule.user_id, + rule.content, + rule.is_active, + rule.created_at, + rule.updated_at, + ), + ) + await conn.commit() + return rule + + async def delete_rule(self, user_id: str, rule_id: uuid.UUID) -> bool: + """Delete a rule by id; return True if a row was removed.""" + await self.ensure_tables() + async with await psycopg.AsyncConnection.connect(self._uri) as conn: + cur = await conn.execute( + "DELETE FROM user_rules WHERE id = %s AND user_id = %s", + (str(rule_id), user_id), + ) + await conn.commit() + return bool(cur.rowcount > 0) diff --git a/deep_agent/src/pii_scrubber.py b/deep_agent/src/pii_scrubber.py new file mode 100644 index 00000000..87d62ba9 --- /dev/null +++ b/deep_agent/src/pii_scrubber.py @@ -0,0 +1,165 @@ +"""PII scrubbing utilities for production error responses. + +Removes personally identifiable information from error messages, stack traces, +and other sensitive data before sending to clients. +""" + +import re +from typing import Any + +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +# Patterns to redact from error messages +EMAIL_PATTERN = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b") +FILEPATH_PATTERN = re.compile(r"(/[a-zA-Z0-9_\-./]+)|([A-Z]:\\[a-zA-Z0-9_\-\\./]+)") +UUID_PATTERN = re.compile( + r"\b[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\b", re.I +) +IP_ADDRESS_PATTERN = re.compile(r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b") +JWT_PATTERN = re.compile(r"\beyJ[A-Za-z0-9_-]+\.eyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\b") + +# Keywords that indicate sensitive data in field names +SENSITIVE_KEYWORDS = { + "password", + "secret", + "token", + "key", + "credential", + "auth", + "ssn", + "social_security", + "credit_card", + "api_key", +} + + +def scrub_pii(text: str) -> str: + """Remove PII patterns from text. + + Redacts: + - Email addresses + - File paths + - UUIDs (partial redaction) + - IP addresses + - JWT tokens + + Args: + text: Input text potentially containing PII + + Returns: + Text with PII replaced by [REDACTED] + """ + if not settings.is_production: + return text + + # Redact emails + text = EMAIL_PATTERN.sub("[EMAIL_REDACTED]", text) + + # Redact file paths (keep just filename) + def redact_path(match: re.Match) -> str: + path = match.group(0) + # Keep just the filename + parts = path.replace("\\", "/").split("/") + return f"[PATH]/{parts[-1]}" if parts else "[PATH_REDACTED]" + + text = FILEPATH_PATTERN.sub(redact_path, text) + + # Redact UUIDs (keep first 8 chars for tracing) + def redact_uuid(match: re.Match) -> str: + uuid = match.group(0) + return f"{uuid[:8]}-[REDACTED]" + + text = UUID_PATTERN.sub(redact_uuid, text) + + # Redact IP addresses + text = IP_ADDRESS_PATTERN.sub("[IP_REDACTED]", text) + + # Redact JWT tokens + text = JWT_PATTERN.sub("[TOKEN_REDACTED]", text) + + return text + + +def scrub_dict(data: dict[str, Any]) -> dict[str, Any]: + """Recursively scrub PII from dictionary values. + + Redacts values for keys containing sensitive keywords. + Also scrubs string values for PII patterns. + + Args: + data: Dictionary potentially containing PII + + Returns: + Dictionary with PII scrubbed + """ + if not settings.is_production: + return data + + scrubbed: dict[str, Any] = {} + for key, value in data.items(): + key_lower = key.lower() + + # Check if key name suggests sensitive data + if any(kw in key_lower for kw in SENSITIVE_KEYWORDS): + scrubbed[key] = "[REDACTED]" + elif isinstance(value, str): + scrubbed[key] = scrub_pii(value) + elif isinstance(value, dict): + scrubbed[key] = scrub_dict(value) + elif isinstance(value, list): + scrubbed[key] = [ + scrub_dict(item) + if isinstance(item, dict) + else scrub_pii(item) + if isinstance(item, str) + else item + for item in value + ] + else: + scrubbed[key] = value + + return scrubbed + + +def scrub_error_response(detail: str, exc: Exception | None = None) -> dict[str, Any]: + """Create a scrubbed error response for production. + + In production: + - Scrubs PII from detail message + - Omits stack traces + - Removes internal paths + + In development: + - Returns full error details for debugging + + Args: + detail: Error message detail + exc: Optional exception for additional context + + Returns: + Scrubbed error response dictionary + """ + if not settings.is_production: + # Development: return full details + response = {"detail": detail} + if exc: + response["exception_type"] = type(exc).__name__ + response["exception_message"] = str(exc) + return response + + # Production: scrub PII and limit information + scrubbed_detail = scrub_pii(detail) + + response = { + "detail": scrubbed_detail, + "error_type": "internal_error", + } + + # Only include exception type, not the message (may contain PII) + if exc: + response["exception_type"] = type(exc).__name__ + + return response diff --git a/template_agent/src/schema.py b/deep_agent/src/schema.py similarity index 80% rename from template_agent/src/schema.py rename to deep_agent/src/schema.py index 3bf9be83..f296bfae 100644 --- a/template_agent/src/schema.py +++ b/deep_agent/src/schema.py @@ -94,24 +94,24 @@ class ChatMessage(BaseModel): examples=["call_Jja7J89XsjrOLA5r!MEOW!SL"], ) run_id: str | None = Field( - description="Run ID associated with this message for tracking.", + description="Run ID associated with this message for tracking (hex format).", default=None, - examples=["847c6285-8fc9-4560-a83f-4e6285809254"], + examples=["847c62858fc94560a83f4e6285809254"], ) - thread_id: str | None = Field( - description="Thread ID associated with this message for conversation tracking.", + trace_id: str | None = Field( + description="Trace ID associated with this message for tracing (hex format).", default=None, - examples=["847c6285-8fc9-4560-a83f-4e6285809254"], + examples=["847c62858fc94560a83f4e6285809254"], ) - session_id: str | None = Field( - description="Session ID associated with this message for session tracking.", + thread_id: str | None = Field( + description="Thread ID associated with this message for conversation tracking (hex format).", default=None, - examples=["847c6285-8fc9-4560-a83f-4e6285809254"], + examples=["847c62858fc94560a83f4e6285809254"], ) - ai_call_id: str | None = Field( - description="Unique identifier for the AI call that generated this message.", + session_id: str | None = Field( + description="Session ID associated with this message for session tracking (hex format).", default=None, - examples=["ai_call_847c6285-8fc9-4560-a83f-4e6285809254"], + examples=["847c62858fc94560a83f4e6285809254"], ) response_metadata: dict[str, Any] = Field( description="Additional metadata for the response, such as headers, logprobs, or token counts.", @@ -130,23 +130,35 @@ class FeedbackRequest(BaseModel): LangFuse for analytics and monitoring purposes. """ - run_id: str = Field( - description="Run ID to record feedback for.", - examples=["847c6285-8fc9-4560-a83f-4e6285809254"], + trace_id: str = Field( + description="Trace ID to record feedback for (hex format, no hyphens).", + examples=["847c62858fc94560a83f4e6285809254"], ) - key: str = Field( - description="Feedback key identifier.", - examples=["human-feedback-stars"], + name: str = Field( + description="Score name/identifier.", + examples=["user-rating", "thumbs-up"], ) - score: float = Field( - description="Feedback score value.", - examples=[0.8], + value: float = Field( + description="Score value.", + examples=[0.8, 1.0], ) kwargs: dict[str, Any] = Field( description="Additional feedback parameters passed to LangFuse.", default={}, examples=[{"comment": "In-line human feedback"}], ) + thread_id: str | None = Field( + default=None, + description="Thread ID for persistence", + ) + message_id: str | None = Field( + default=None, + description="Message ID for persistence", + ) + user_id: str | None = Field( + default=None, + description="User ID for persistence", + ) class FeedbackResponse(BaseModel): diff --git a/deep_agent/src/settings.py b/deep_agent/src/settings.py new file mode 100644 index 00000000..e93107b5 --- /dev/null +++ b/deep_agent/src/settings.py @@ -0,0 +1,258 @@ +"""Settings configuration for the template agent. + +All operational defaults live HERE. No env vars needed for basic operation. +Override via environment variables only when deploying to a different context. + +Hierarchy (highest wins): + 1. Environment variables (set by orchestrator, compose, or shell) + 2. .env file (secrets only — keys, passwords, credentials) + 3. Defaults below (tuned for containerized demo stack) +""" + +from typing import Optional +from urllib.parse import urlparse + +from dotenv import load_dotenv +from pydantic import Field +from pydantic_settings import BaseSettings + +from deep_agent.src.exceptions import AppException, ErrorCodes +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_DEV_PUBLIC_HOSTS = frozenset({"localhost", "127.0.0.1", "::1"}) + +try: + load_dotenv() +except Exception as e: + logger.warning(f"Could not load .env file: {e}") + + +class Settings(BaseSettings): + """All agent settings with production-ready defaults. + + Grouped by concern. Every field has a sensible default so the agent + starts with zero configuration beyond secrets in .env. + """ + + # ── Server ──────────────────────────────────────────────────────── + AGENT_HOST: str = Field(default="0.0.0.0") + AGENT_PORT: int = Field(default=5002) + SSL_KEYFILE: Optional[str] = Field(default=None) + SSL_CERTFILE: Optional[str] = Field(default=None) + + @property + def get_ssl_keyfile_path(self) -> Optional[str]: + """Return SSL key file path if configured, else None.""" + return None if not self.SSL_KEYFILE else self.SSL_KEYFILE + + @property + def get_ssl_certfile_path(self) -> Optional[str]: + """Return SSL cert file path if configured, else None.""" + return None if not self.SSL_CERTFILE else self.SSL_CERTFILE + + # ── Logging ─────────────────────────────────────────────────────── + PYTHON_LOG_LEVEL: str = Field(default="INFO") + REQUEST_LOGGING_ENABLED: bool = Field(default=True) + REQUEST_LOG_HEADERS: bool = Field(default=True) + REQUEST_LOG_BODY: bool = Field(default=True) + REQUEST_LOG_BODY_MAX_SIZE: int = Field(default=10240) + + # ── Security ────────────────────────────────────────────────────── + REQUEST_BODY_MAX_SIZE: int = Field( + default=10 * 1024 * 1024, # 10MB + description="Maximum request body size in bytes (DoS protection)", + ) + + # ── Model ───────────────────────────────────────────────────────── + MAX_OUTPUT_TOKENS: int = Field(default=8192) + + # ── Database (PostgreSQL) ───────────────────────────────────────── + POSTGRES_HOST: str = Field(default="pgvector") + POSTGRES_PORT: int = Field(default=5432) + POSTGRES_DB: str = Field(default="template_agent") + POSTGRES_USER: str = Field(default="postgres") + POSTGRES_PASSWORD: str = Field(default="postgres") + + # ── MongoDB ─────────────────────────────────────────────────────── + MONGODB_URI: Optional[str] = Field(default=None, repr=False) + MONGODB_DB: str = Field(default="tokenusage") + + # ── Redis ───────────────────────────────────────────────────────── + REDIS_URL: str = Field(default="redis://redis:6379/0") + REDIS_BROKER_ENABLED: bool = Field(default=True) + + # ── Auth / SSO ──────────────────────────────────────────────────── + ENABLE_AUTH: bool = Field(default=True) + SSO_ISSUER_URL: Optional[str] = Field(default=None) + SSO_CLIENT_ID: Optional[str] = Field(default=None) + SSO_CLIENT_SECRET: Optional[str] = Field(default=None) + SSO_DEV_USERNAME: str = Field(default="John Doe") + SSO_DEV_USER_ID: str = Field(default="dev-user") + ENABLE_USER_ID_ENCRYPTION: bool = Field(default=False) + + # ── Environment ─────────────────────────────────────────────────── + ENVIRONMENT: str = Field( + default="development", + description="Runtime environment: development, production, staging. " + "Production mode enforces auth, SSL verification, and PII scrubbing.", + ) + + @property + def is_production(self) -> bool: + """True when running in production environment.""" + return self.ENVIRONMENT.lower() == "production" + + # ── Observability (Langfuse) ────────────────────────────────────── + LANGFUSE_PUBLIC_KEY: Optional[str] = Field(default=None) + LANGFUSE_SECRET_KEY: Optional[str] = Field(default=None) + LANGFUSE_BASE_URL: Optional[str] = Field(default=None) + LANGFUSE_TRACING_ENVIRONMENT: str = Field(default="development") + + # ── OpenTelemetry ───────────────────────────────────────────────── + ENABLE_OTEL_METRICS: bool = Field(default=False) + ENABLE_OTEL_TRACES: bool = Field(default=False) + OTEL_SERVICE_NAME: str = Field(default="template-agent") + OTEL_EXPORTER_OTLP_ENDPOINT: str = Field( + default="", + description="OTLP gRPC metrics endpoint (OpenShift: otel-gateway:4327)", + ) + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: str = Field( + default="", + description="OTLP gRPC traces endpoint (local/dev-loop: Jaeger :4317)", + ) + OTEL_AUTH_TOKEN: str = Field(default="", repr=False) + OTEL_METRIC_EXPORT_INTERVAL_MILLIS: int = Field(default=10000) + + def resolved_otel_traces_endpoint(self) -> str: + """Return the configured OTLP traces exporter endpoint.""" + return self.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT + + def otel_traces_active(self) -> bool: + """Return True when trace export is enabled and an endpoint is configured.""" + return bool(self.ENABLE_OTEL_TRACES and self.resolved_otel_traces_endpoint()) + + # ── Google Cloud ────────────────────────────────────────────────── + GOOGLE_APPLICATION_CREDENTIALS_CONTENT: Optional[str] = Field(default=None) + + # ── vLLM / OpenAI-compatible ───────────────────────────────────── + VLLM_BASE_URL: Optional[str] = Field(default=None) + VLLM_API_KEY: str = Field(default="EMPTY") + + # ── Cache ───────────────────────────────────────────────────────── + CACHE_ENABLED: bool = Field(default=True) + + # ── Memory Processing ───────────────────────────────────────────── + MEMORY_CONSOLIDATION_ENABLED: bool = Field(default=True) + MEMORY_DECAY_ENABLED: bool = Field(default=True) + MEMORY_CLUSTERING_ENABLED: bool = Field(default=True) + MEMORY_RELATIONSHIPS_ENABLED: bool = Field(default=True) + + # ── Middleware ──────────────────────────────────────────────────── + MIDDLEWARE_ENABLED: bool = Field(default=True) + + # ── CLI ─────────────────────────────────────────────────────────── + ENABLE_CLI: bool = Field(default=True) + + # ── Platform ────────────────────────────────────────────────────── + DEPLOYED_AGENT_NAME: str = Field(default="") + DEPLOYED_AGENT_ORG: str = Field(default="") + PLATFORM_AUDIT_ENABLED: bool = Field(default=True) + PLATFORM_AUDIT_BUFFER_MAX: int = Field(default=1000, ge=1, le=100_000) + + # ── FLAG TO SWITCH TO RELOAD FROM DISK ──────────────────────────── + CONFIG_AUTO_RELOAD: bool = Field(default=True) + + # ── MCP OAuth ───────────────────────────────────────────────────── + MCP_TOKEN_ENCRYPTION_KEY: Optional[str] = Field(default=None) + MCP_TOKEN_ENCRYPTION_KEY_PREVIOUS: Optional[str] = Field(default=None) + AGENT_PUBLIC_BASE_URL: Optional[str] = Field(default=None) + + # ── Derived ─────────────────────────────────────────────────────── + + @property + def agent_deployment_id(self) -> str: + """Unique identity for this agent deployment, used as DCR/token key. + + Combines org + agent name when deployed via agent-engine. + Falls back to the generic config name for local dev. + """ + if self.DEPLOYED_AGENT_ORG and self.DEPLOYED_AGENT_NAME: + return f"{self.DEPLOYED_AGENT_ORG}/{self.DEPLOYED_AGENT_NAME}" + if self.DEPLOYED_AGENT_NAME: + return self.DEPLOYED_AGENT_NAME + from deep_agent.src.agent.config import agent_config + return agent_config.get_name() + + @property + def agent_public_base_url(self) -> str: + """Public base URL for MCP OAuth connect/callback endpoints.""" + if self.AGENT_PUBLIC_BASE_URL: + return self.AGENT_PUBLIC_BASE_URL.rstrip("/") + return f"http://localhost:{self.AGENT_PORT}" + + @property + def is_dev_public_url(self) -> bool: + """True when the public base URL is an allowed local HTTP dev endpoint.""" + parsed = urlparse(self.agent_public_base_url) + return parsed.scheme == "http" and parsed.hostname in _DEV_PUBLIC_HOSTS + + @property + def oauth_callback_url(self) -> str: + """Canonical OAuth redirect URI derived from AGENT_PUBLIC_BASE_URL.""" + return f"{self.agent_public_base_url}/mcp/oauth/callback" + + @property + def database_uri(self) -> str: + """Build PostgreSQL connection URI from component settings.""" + return ( + f"postgresql://{self.POSTGRES_USER}:{self.POSTGRES_PASSWORD}" + f"@{self.POSTGRES_HOST}:{self.POSTGRES_PORT}/{self.POSTGRES_DB}" + ) + + +def validate_config(settings: Settings) -> None: + """Validate port range, log level, and production constraints.""" + if not (1024 <= settings.AGENT_PORT <= 65535): + raise AppException( + f"AGENT_PORT must be between 1024 and 65535, got {settings.AGENT_PORT}", + ErrorCodes.CONFIGURATION_VALIDATION_ERROR, + ) + + valid_log_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] + if settings.PYTHON_LOG_LEVEL.upper() not in valid_log_levels: + raise AppException( + f"PYTHON_LOG_LEVEL must be one of {valid_log_levels}, got {settings.PYTHON_LOG_LEVEL}", + ErrorCodes.CONFIGURATION_VALIDATION_ERROR, + ) + + if settings.AGENT_PUBLIC_BASE_URL and not settings.is_dev_public_url: + parsed = urlparse(settings.AGENT_PUBLIC_BASE_URL) + if parsed.scheme != "https": + raise AppException( + "AGENT_PUBLIC_BASE_URL must use https:// in production " + "(http:// is permitted only for localhost, 127.0.0.1, or ::1)", + ErrorCodes.CONFIGURATION_VALIDATION_ERROR, + ) + + # Production-specific validations + if settings.is_production: + # Enforce auth in production + if not settings.ENABLE_AUTH: + raise AppException( + "ENABLE_AUTH must be true in production. " + "Configure SSO_ISSUER_URL, SSO_CLIENT_ID, and SSO_CLIENT_SECRET.", + ErrorCodes.CONFIGURATION_VALIDATION_ERROR, + ) + + # Enforce HTTPS for public URL in production + if settings.AGENT_PUBLIC_BASE_URL and settings.is_dev_public_url: + raise AppException( + "AGENT_PUBLIC_BASE_URL cannot use http://localhost in production. " + "Configure a valid https:// URL.", + ErrorCodes.CONFIGURATION_VALIDATION_ERROR, + ) + + +settings = Settings() diff --git a/deep_agent/src/streaming/__init__.py b/deep_agent/src/streaming/__init__.py new file mode 100644 index 00000000..87ac5c14 --- /dev/null +++ b/deep_agent/src/streaming/__init__.py @@ -0,0 +1,23 @@ +"""Streaming response components for the template agent system. + +This package contains the modular components that handle streaming responses, +message deduplication, tool call tracking, and event formatting. +""" + +from deep_agent.src.streaming.context import StreamContext +from deep_agent.src.streaming.converter import remove_tool_calls +from deep_agent.src.streaming.deduplicator import MessageDeduplicator +from deep_agent.src.streaming.handlers import ( + TokenEventHandler, + UpdateEventHandler, +) +from deep_agent.src.streaming.tracker import ToolCallTracker + +__all__ = [ + "StreamContext", + "MessageDeduplicator", + "ToolCallTracker", + "UpdateEventHandler", + "TokenEventHandler", + "remove_tool_calls", +] diff --git a/deep_agent/src/streaming/context.py b/deep_agent/src/streaming/context.py new file mode 100644 index 00000000..ea471e72 --- /dev/null +++ b/deep_agent/src/streaming/context.py @@ -0,0 +1,27 @@ +"""Stream context for carrying metadata through event processing. + +This module provides the StreamContext dataclass that carries essential +metadata (run_id, trace_id, thread_id, session_id, user_id) through the +entire streaming pipeline, ensuring all events have consistent context. +""" + +from dataclasses import dataclass + + +@dataclass +class StreamContext: + """Context object for streaming metadata. + + Carries run, trace, thread, session, and user identifiers plus configuration + through the event processing pipeline. + + All fields are required to ensure complete context is available + throughout the streaming pipeline. + """ + + run_id: str + trace_id: str + thread_id: str + session_id: str + user_id: str + stream_tokens: bool diff --git a/deep_agent/src/streaming/converter.py b/deep_agent/src/streaming/converter.py new file mode 100644 index 00000000..f2fb9ac7 --- /dev/null +++ b/deep_agent/src/streaming/converter.py @@ -0,0 +1,112 @@ +"""Message format conversion for streaming API responses. + +This module converts internal ChatMessage objects to the simplified JSON format +sent to clients via streaming endpoints. It handles special cases like tool call +rewrites and context metadata injection. +""" + +from typing import Any, Dict, List, Union + +from langchain_core.messages import BaseMessage + +from deep_agent.src.streaming.context import StreamContext + + +def convert_message_to_api_format( + chat_message: Any, ctx: StreamContext +) -> dict[str, Any]: + """Convert ChatMessage to simplified API format. + + Args: + chat_message: The chat message to convert. + ctx: Stream context with metadata. + + Returns: + Simplified message dictionary with type, content, and context metadata. + """ + content = { + "type": chat_message.type, + "content": chat_message.content, + } + + # Add optional message-specific fields + if chat_message.tool_calls: + # Rewrite "task" tool name to actual subagent name for better UI display + content["tool_calls"] = [ + {**tc, "name": tc["args"]["subagent_type"]} + if tc.get("name") == "task" and "subagent_type" in tc.get("args", {}) + else tc + for tc in chat_message.tool_calls + ] + if chat_message.tool_call_id: + content["tool_call_id"] = chat_message.tool_call_id + if chat_message.response_metadata: + content["response_metadata"] = chat_message.response_metadata + + # Add context metadata (always present, authoritative for the stream) + content["run_id"] = ctx.run_id + content["trace_id"] = ctx.trace_id + content["thread_id"] = ctx.thread_id + content["session_id"] = ctx.session_id + content["user_id"] = ctx.user_id + + return content + + +def remove_tool_calls( + content: Union[str, List[Union[str, Dict[str, Any]]]], +) -> Union[str, List[Union[str, Dict[str, Any]]]]: + """Remove tool calls from message content. + + This function filters out tool call content from message content, particularly + useful for handling streaming responses from models that include tool calls + in their content stream. + + Args: + content: The content to process. Can be a string or a list containing + strings and dictionaries with content information. + + Returns: + The content with tool calls removed. Returns the same type as input. + """ + if isinstance(content, str): + return content + + # Currently only Anthropic models stream tool calls, using content item type tool_use + return [ + content_item + for content_item in content + if isinstance(content_item, str) or content_item["type"] != "tool_use" + ] + + +def should_skip_message(message: BaseMessage) -> tuple[bool, str | None]: + """Determine if a message should be skipped. + + Args: + message: The message to check. + + Returns: + Tuple of (should_skip, reason). + """ + from langchain_core.messages import AIMessage, ToolMessage + + # Skip empty tool messages + if isinstance(message, ToolMessage) and not message.content: + tool_name = message.name or "unknown" + return ( + True, + f"Subagent '{tool_name}' returned empty result (tool_call_id={message.tool_call_id})", + ) + + # Skip empty AI messages from malformed function calls + if ( + isinstance(message, AIMessage) + and not message.content + and not message.tool_calls + ): + reason = message.response_metadata.get("finish_reason", "") + if reason == "MALFORMED_FUNCTION_CALL": + return True, "LLM returned MALFORMED_FUNCTION_CALL — skipping empty message" + + return False, None diff --git a/deep_agent/src/streaming/deduplicator.py b/deep_agent/src/streaming/deduplicator.py new file mode 100644 index 00000000..5e7fcb8e --- /dev/null +++ b/deep_agent/src/streaming/deduplicator.py @@ -0,0 +1,102 @@ +"""Message deduplication for handling LangGraph checkpoint replays. + +This module provides MessageDeduplicator to prevent duplicate messages when +LangGraph replays from checkpoints. It tracks message IDs and filters out +messages that have already been seen in the current stream. +""" + +from langchain_core.messages import BaseMessage, ToolMessage + + +def extract_message_id(msg: BaseMessage) -> str | None: + """Extract a stable identifier from a message. + + Args: + msg: A LangChain message object. + + Returns: + A stable ID string, or None if no stable ID exists. + """ + msg_id: str | None + if isinstance(msg.id, str): + msg_id = msg.id + elif isinstance(msg, ToolMessage): + # ToolMessages may not have .id set; use tool_call_id as fallback + msg_id = f"tool_{msg.tool_call_id}" + else: + msg_id = None + return msg_id + + +class MessageDeduplicator: + """Tracks and filters duplicate messages across checkpoint restores. + + LangGraph can replay message history via Overwrite updates when + resuming from checkpoints. This class ensures we only emit new + messages to avoid duplicate streaming. + """ + + def __init__(self) -> None: + """Initialize the deduplicator.""" + self._seen_ids: set[str] = set() + + def reset(self) -> None: + """Clear all seen message IDs.""" + self._seen_ids.clear() + + def mark_seen(self, msg: BaseMessage) -> None: + """Mark a message as seen. + + Args: + msg: A LangChain message object. + """ + msg_id = extract_message_id(msg) + if msg_id: + self._seen_ids.add(msg_id) + + def is_seen(self, msg: BaseMessage) -> bool: + """Check if a message has been seen before. + + Args: + msg: A LangChain message object. + + Returns: + True if the message was previously seen, False otherwise. + """ + msg_id = extract_message_id(msg) + if msg_id is None: + # No stable ID - can't reliably deduplicate + return False + return msg_id in self._seen_ids + + def get_unseen_messages(self, messages: list[BaseMessage]) -> list[BaseMessage]: + """Get only unseen messages from a list, marking them as seen. + + Args: + messages: List of LangChain message objects. + + Returns: + List of messages not previously seen. + """ + unseen = [] + for msg in messages: + msg_id = extract_message_id(msg) + if msg_id is None: + # No stable ID - always include to avoid data loss + unseen.append(msg) + elif msg_id not in self._seen_ids: + unseen.append(msg) + self._seen_ids.add(msg_id) + return unseen + + def populate_from_history(self, messages: list[BaseMessage]) -> None: + """Pre-populate seen IDs from existing message history. + + Used when resuming from a checkpoint to avoid replaying + the full conversation history. + + Args: + messages: List of messages from checkpoint state. + """ + for msg in messages: + self.mark_seen(msg) diff --git a/deep_agent/src/streaming/handlers.py b/deep_agent/src/streaming/handlers.py new file mode 100644 index 00000000..51b6175a --- /dev/null +++ b/deep_agent/src/streaming/handlers.py @@ -0,0 +1,206 @@ +"""Event handlers for processing LangGraph stream events. + +This module provides event handler classes (TokenEventHandler, UpdateEventHandler) +that process LangGraph streaming events and convert them into API-friendly formats. +Handles both token-level and update-level streaming modes. +""" + +from typing import Any + +from langchain_core.messages import AIMessage, AIMessageChunk +from langgraph.types import Overwrite + +from deep_agent.src.adapters.langchain import ( + convert_message_content_to_string, + langchain_to_chat_message, +) +from deep_agent.src.settings import settings +from deep_agent.src.streaming.context import StreamContext +from deep_agent.src.streaming.converter import ( + convert_message_to_api_format, + remove_tool_calls, + should_skip_message, +) +from deep_agent.src.streaming.deduplicator import MessageDeduplicator +from deep_agent.src.streaming.tracker import ( + ToolCallTracker, + extract_tool_call_id, +) +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger(settings.PYTHON_LOG_LEVEL) + + +def _convert_interrupts_to_messages(interrupts: list) -> list: + """Convert interrupt data to messages. + + Args: + interrupts: List of interrupt objects. + + Returns: + List of AIMessage objects. + """ + messages = [] + for interrupt_data in interrupts: + content = ( + interrupt_data.value + if hasattr(interrupt_data, "value") + else str(interrupt_data) + ) + messages.append(AIMessage(content=content)) + return messages + + +def _convert_messages_to_events( + messages: list, ctx: StreamContext +) -> list[dict[str, Any]]: + """Convert messages to simplified event format. + + Args: + messages: List of LangChain messages. + ctx: Stream context with metadata. + + Returns: + List of formatted events. + """ + formatted_events = [] + + for message in messages: + try: + # Check if message should be skipped + should_skip, reason = should_skip_message(message) + if should_skip: + if reason: + logger.warning(reason) + continue + + # Convert to chat message format + chat_message = langchain_to_chat_message(message) + chat_message.run_id = ctx.run_id + + # Convert to simplified format + formatted_event = { + "type": "message", + "content": convert_message_to_api_format(chat_message, ctx), + } + formatted_events.append(formatted_event) + + except Exception as e: + logger.error(f"Error formatting message: {e}") + formatted_events.append( + { + "type": "error", + "content": { + "message": "Message formatting error", + "recoverable": True, + }, + } + ) + + return formatted_events + + +class UpdateEventHandler: + """Handles 'updates' stream mode events from LangGraph.""" + + def __init__(self, deduplicator: MessageDeduplicator): + """Initialize the handler. + + Args: + deduplicator: Message deduplicator for handling replays. + """ + self.deduplicator = deduplicator + + def handle(self, event: dict[str, Any], ctx: StreamContext) -> list[dict[str, Any]]: + """Process update events and convert to simplified format. + + Args: + event: Dictionary mapping node names to update data. + ctx: Stream context with metadata. + + Returns: + List of formatted message events. + """ + messages = self._extract_and_deduplicate_messages(event) + return _convert_messages_to_events(messages, ctx) + + def _extract_and_deduplicate_messages(self, event: dict[str, Any]) -> list: + """Extract and deduplicate messages from update event. + + Args: + event: Update event dictionary. + + Returns: + List of messages to process. + """ + all_messages = [] + + for node, updates in event.items(): + if node == "__interrupt__": + all_messages.extend(_convert_interrupts_to_messages(updates)) + continue + + updates = updates or {} + raw_messages = updates.get("messages", []) + is_overwrite = isinstance(raw_messages, Overwrite) + update_messages = raw_messages.value if is_overwrite else raw_messages + + if is_overwrite: + # Filter to only unseen messages + update_messages = self.deduplicator.get_unseen_messages(update_messages) + else: + # Mark all messages as seen for future deduplication + for msg in update_messages: + self.deduplicator.mark_seen(msg) + + all_messages.extend(update_messages) + + return all_messages + + +class TokenEventHandler: + """Handles 'messages' stream mode events (token streaming).""" + + def __init__(self, tracker: ToolCallTracker): + """Initialize the handler. + + Args: + tracker: Tool call tracker for associating tokens with tools. + """ + self.tracker = tracker + + def handle(self, event: tuple, ctx: StreamContext) -> list[dict[str, Any]]: + """Process token streaming events. + + Args: + event: Tuple of (message, metadata). + ctx: Stream context with metadata. + + Returns: + List containing a single token event, or empty list. + """ + if not ctx.stream_tokens: + return [] + + msg, metadata = event + if "skip_stream" in metadata.get("tags", []): + return [] + + if not isinstance(msg, AIMessageChunk): + return [] + + content = remove_tool_calls(msg.content) + if not content: + return [] + + token_event = { + "type": "token", + "content": convert_message_content_to_string(content), + } + + # Associate token with tool call if applicable + tool_call_id = extract_tool_call_id(msg) or self.tracker.current_id + if tool_call_id: + token_event["tool_call_id"] = tool_call_id + + return [token_event] diff --git a/deep_agent/src/streaming/tracker.py b/deep_agent/src/streaming/tracker.py new file mode 100644 index 00000000..170f6856 --- /dev/null +++ b/deep_agent/src/streaming/tracker.py @@ -0,0 +1,103 @@ +"""Tool call tracking for enhanced UI feedback. + +This module provides ToolCallTracker to accumulate tool call information from +streaming chunks and emit complete tool call events. This enables UIs to show +tool invocations with full context even during token streaming. +""" + +from typing import Any + +from langchain_core.messages import AIMessageChunk + +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger(settings.PYTHON_LOG_LEVEL) + + +def extract_tool_call_id(msg: AIMessageChunk) -> str | None: + """Extract tool call ID from an AIMessageChunk. + + Modern LangChain automatically populates tool_calls from tool_call_chunks + during streaming, so we only need to check tool_calls. + + Args: + msg: The message chunk to extract from. + + Returns: + The tool call ID if available, None otherwise. + """ + try: + if msg.tool_calls: + tool_call_id = msg.tool_calls[0].get("id") + return tool_call_id if isinstance(tool_call_id, str) else None + return None + except (IndexError, KeyError) as e: + logger.debug(f"Could not extract tool call ID: {e}") + return None + + +class ToolCallTracker: + """Tracks active tool calls to associate streaming tokens with tools. + + When a tool is invoked, streaming tokens that follow should be + associated with that tool's response. This tracker maintains the + current tool call ID for proper attribution in the UI. + """ + + def __init__(self) -> None: + """Initialize the tracker.""" + self._current_tool_call_id: str | None = None + + def reset(self) -> None: + """Clear the current tool call ID.""" + self._current_tool_call_id = None + + @property + def current_id(self) -> str | None: + """Get the current tool call ID being tracked.""" + return self._current_tool_call_id + + def update_from_stream_event(self, stream_mode: str, event: Any) -> None: + """Update tracking based on a stream event. + + Args: + stream_mode: The type of stream event (updates, messages, custom). + event: The event data. + """ + try: + if stream_mode == "updates": + self._update_from_updates(event) + elif stream_mode == "messages": + self._update_from_message_stream(event) + except Exception as e: + logger.debug(f"Tool call tracking error: {e}") + + def _update_from_updates(self, event: dict) -> None: + """Update from an 'updates' mode event.""" + from langchain_core.messages import ToolMessage + + for _node, updates in event.items(): + if not updates or "messages" not in updates: + continue + for message in updates["messages"]: + # ToolMessage responding to a tool call + if isinstance(message, ToolMessage): + self._current_tool_call_id = message.tool_call_id + return + # AIMessage with tool calls + elif message.tool_calls: + self._current_tool_call_id = message.tool_calls[0].get("id") + return + + def _update_from_message_stream(self, event: tuple) -> None: + """Update from a 'messages' mode event.""" + from langchain_core.messages import ToolMessage + + msg, _metadata = event + # ToolMessage responding to a tool call + if isinstance(msg, ToolMessage): + self._current_tool_call_id = msg.tool_call_id + # AIMessage with tool calls + elif msg.tool_calls: + self._current_tool_call_id = msg.tool_calls[0].get("id") diff --git a/deep_agent/src/token_budget/__init__.py b/deep_agent/src/token_budget/__init__.py new file mode 100644 index 00000000..1229d512 --- /dev/null +++ b/deep_agent/src/token_budget/__init__.py @@ -0,0 +1 @@ +"""Thread-level token budget tracking and threshold warnings.""" diff --git a/deep_agent/src/token_budget/callback.py b/deep_agent/src/token_budget/callback.py new file mode 100644 index 00000000..d326fbf8 --- /dev/null +++ b/deep_agent/src/token_budget/callback.py @@ -0,0 +1,153 @@ +"""LangChain callback handler for per-thread token budget tracking.""" + +from __future__ import annotations + +import threading +from typing import Any + +from langchain_core.callbacks import AsyncCallbackHandler +from langchain_core.outputs import ChatGeneration, LLMResult + +from deep_agent.src.token_budget.identity import resolve_thread_id, resolve_user_id +from deep_agent.src.token_budget.service import ( + check_and_record, + extract_tokens_from_llm_result, + extract_tokens_from_message, +) +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +THREAD_ID_METADATA_KEY = "token_budget_thread_id" +USER_ID_METADATA_KEY = "token_budget_user_id" +TRACE_ID_METADATA_KEY = "token_budget_trace_id" + +# Emit an ERROR-level alert after this many consecutive failures so ops +# teams can detect a persistently degraded token-tracking feature. +_CONSECUTIVE_FAILURE_ALERT_THRESHOLD = 5 + +_counter_lock = threading.Lock() +_consecutive_failures = 0 +_total_failures = 0 + + +def _on_tracking_success() -> None: + global _consecutive_failures + with _counter_lock: + _consecutive_failures = 0 + + +def _on_tracking_failure() -> None: + global _consecutive_failures, _total_failures + with _counter_lock: + _consecutive_failures += 1 + _total_failures += 1 + consecutive = _consecutive_failures + total = _total_failures + if consecutive >= _CONSECUTIVE_FAILURE_ALERT_THRESHOLD: + logger.error( + "token_budget_tracking_degraded", + consecutive_failures=consecutive, + total_failures=total, + ) + + +def _extract_from_metadata( + metadata: dict[str, Any] | None, + key: str, + *, + fallback_keys: tuple[str, ...] = (), +) -> str | None: + """Return the first non-empty metadata value for key or fallback keys.""" + if not metadata: + return None + for metadata_key in (key, *fallback_keys): + value = metadata.get(metadata_key) + if value: + return str(value) + return None + + +def thread_id_from_metadata(metadata: dict[str, Any] | None) -> str | None: + """Resolve thread_id from RunnableConfig metadata.""" + return _extract_from_metadata( + metadata, + THREAD_ID_METADATA_KEY, + fallback_keys=("langfuse_session_id",), + ) + + +def user_id_from_metadata(metadata: dict[str, Any] | None) -> str | None: + """Resolve the chatting user's id from RunnableConfig metadata.""" + return _extract_from_metadata(metadata, USER_ID_METADATA_KEY) + + +def trace_id_from_metadata(metadata: dict[str, Any] | None) -> str | None: + """Resolve trace_id from RunnableConfig metadata.""" + return _extract_from_metadata(metadata, TRACE_ID_METADATA_KEY) + + +class TokenBudgetCallbackHandler(AsyncCallbackHandler): + """Increment per-thread token usage after each LLM call.""" + + async def _record_tokens( + self, + response: LLMResult, + metadata: dict[str, Any] | None, + extraction_fn: Any, + ) -> None: + """Shared logic for recording token usage from LLM responses.""" + thread_id = thread_id_from_metadata(metadata) or resolve_thread_id() + if not thread_id: + logger.debug("token_budget_callback_no_thread_id") + return + + user_id = user_id_from_metadata(metadata) or resolve_user_id() + trace_id = trace_id_from_metadata(metadata) + + input_tokens, output_tokens = extraction_fn(response) + if input_tokens <= 0 and output_tokens <= 0: + input_tokens, output_tokens = _tokens_from_generations(response) + + try: + await check_and_record( + thread_id, + input_tokens, + output_tokens, + user_id=user_id, + trace_id=trace_id, + ) + _on_tracking_success() + except Exception: + logger.warning( + "token_budget_callback_failed", + exc_info=True, + ) + _on_tracking_failure() + + async def on_llm_end( + self, + response: LLMResult, + *, + run_id: Any, + parent_run_id: Any | None = None, + tags: list[str] | None = None, + metadata: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Record token usage when an LLM call completes.""" + await self._record_tokens(response, metadata, extract_tokens_from_llm_result) + + +def _tokens_from_generations(response: LLMResult) -> tuple[int, int]: + input_tokens = 0 + output_tokens = 0 + for generation_list in response.generations: + for generation in generation_list: + if not isinstance(generation, ChatGeneration): + continue + message = generation.message + in_t, out_t = extract_tokens_from_message(message) + input_tokens += in_t + output_tokens += out_t + return input_tokens, output_tokens diff --git a/deep_agent/src/token_budget/config.py b/deep_agent/src/token_budget/config.py new file mode 100644 index 00000000..9e555e1d --- /dev/null +++ b/deep_agent/src/token_budget/config.py @@ -0,0 +1,16 @@ +"""Token budget configuration models.""" + +from __future__ import annotations + +from pydantic import BaseModel + + +class TokenBudgetConfig(BaseModel): + """Per-thread token usage tracking from agent.yaml ``token_budget:`` section.""" + + enabled: bool = False + + @property + def is_active(self) -> bool: + """Return True when tracking should run.""" + return self.enabled diff --git a/deep_agent/src/token_budget/identity.py b/deep_agent/src/token_budget/identity.py new file mode 100644 index 00000000..74c81c9d --- /dev/null +++ b/deep_agent/src/token_budget/identity.py @@ -0,0 +1,51 @@ +"""Resolve thread and user identity from LangGraph RunnableConfig.""" + +from __future__ import annotations + +from typing import Any + +from langgraph.runtime import Runtime + + +def resolve_thread_id(runtime: Runtime[Any] | None = None) -> str | None: + """Resolve thread_id from LangGraph runtime or active RunnableConfig.""" + if runtime is not None: + execution_info = getattr(runtime, "execution_info", None) + if execution_info is not None: + thread_id = getattr(execution_info, "thread_id", None) + if thread_id: + return str(thread_id) + + try: + from langgraph.config import get_config + + config = get_config() + configurable = config.get("configurable") or {} + thread_id = configurable.get("thread_id") + if thread_id: + return str(thread_id) + except Exception: + pass + return None + + +def resolve_user_id(runtime: Runtime[Any] | None = None) -> str | None: + """Resolve chatting user_id from LangGraph runtime or active RunnableConfig.""" + if runtime is not None: + execution_info = getattr(runtime, "execution_info", None) + if execution_info is not None: + user_id = getattr(execution_info, "user_id", None) + if user_id: + return str(user_id) + + try: + from langgraph.config import get_config + + config = get_config() + configurable = config.get("configurable") or {} + user_id = configurable.get("user_id") + if user_id: + return str(user_id) + except Exception: + pass + return None diff --git a/deep_agent/src/token_budget/mongo_repository.py b/deep_agent/src/token_budget/mongo_repository.py new file mode 100644 index 00000000..5a5f2e88 --- /dev/null +++ b/deep_agent/src/token_budget/mongo_repository.py @@ -0,0 +1,186 @@ +"""MongoDB store for per-thread and per-user daily token usage. + +Security: + MONGODB_URI may contain credentials. It MUST NOT be logged, included in + error messages, or exposed via API responses. In production the URI should + authenticate as a user with read/write access to the tokenusage DB only + (principle of least privilege — no admin or cluster-wide access). +""" + +from __future__ import annotations + +import re +from datetime import UTC, datetime +from typing import Any, cast + +from motor.motor_asyncio import ( + AsyncIOMotorClient, + AsyncIOMotorCollection, + AsyncIOMotorDatabase, +) +from pymongo import ReturnDocument + +from deep_agent.src.error_handling import mongo_retry +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_INDEXES_ENSURED = False +_DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$") + + +def _validated_date(date: str | None) -> str: + """Return *date* unchanged if it matches YYYY-MM-DD, else raise ValueError. + + When *date* is None the current UTC date is returned. + """ + if date is None: + return datetime.now(UTC).strftime("%Y-%m-%d") + if not _DATE_RE.match(date): + raise ValueError(f"Invalid date format {date!r}, expected YYYY-MM-DD") + return date + + +class TokenUsageMongoRepository: + """MongoDB token usage: per-thread counts and per-user daily rollup.""" + + def __init__(self, mongodb_uri: str, db_name: str) -> None: + """Initialize the repository with a MongoDB URI and database name.""" + self._uri = mongodb_uri + self._db_name = db_name + self._client: AsyncIOMotorClient | None = None + + def __repr__(self) -> str: + """Return a debug representation without exposing credentials.""" + return f"TokenUsageMongoRepository(db={self._db_name!r})" + + def _get_client(self) -> AsyncIOMotorClient: + if self._client is None: + self._client = AsyncIOMotorClient(self._uri) + return self._client + + @property + def _db(self) -> AsyncIOMotorDatabase: + return self._get_client()[self._db_name] + + @property + def _thread_collection(self) -> AsyncIOMotorCollection: + return self._db["thread_token_usage"] + + @property + def _daily_collection(self) -> AsyncIOMotorCollection: + return self._db["user_daily_token_usage"] + + @mongo_retry + async def ensure_indexes(self) -> None: + """Create indexes idempotently once per process. + + MongoDB create_index is a no-op if the index already exists, so + concurrent calls from multiple replicas are safe (no data corruption). + The _INDEXES_ENSURED flag avoids redundant network calls within a + single process. + + For large-scale deployments with many replicas starting simultaneously, + consider running index creation via a one-off migration job instead of + at application startup to avoid thundering-herd load on the DB. + """ + global _INDEXES_ENSURED # noqa: PLW0603 + if _INDEXES_ENSURED: + return + await self._thread_collection.create_index("thread_id", unique=True) + await self._thread_collection.create_index("updated_at") + await self._daily_collection.create_index( + [("user_id", 1), ("date", 1)], + unique=True, + ) + await self._daily_collection.create_index("date") + _INDEXES_ENSURED = True + logger.info("MongoDB token usage indexes ensured") + + @mongo_retry + async def increment_usage( + self, + thread_id: str, + input_tokens: int, + output_tokens: int, + *, + agent_name: str | None = None, + ) -> dict[str, Any]: + """Atomically add tokens for a thread and return the updated document.""" + input_tokens = max(input_tokens, 0) + output_tokens = max(output_tokens, 0) + total_delta = input_tokens + output_tokens + now = datetime.now(UTC) + + update: dict[str, Any] = { + "$inc": { + "total_tokens": total_delta, + "input_tokens": input_tokens, + "output_tokens": output_tokens, + }, + "$set": {"updated_at": now}, + "$setOnInsert": {"thread_id": thread_id}, + } + if agent_name: + update["$set"]["agent_name"] = agent_name + + result = await self._thread_collection.find_one_and_update( + {"thread_id": thread_id}, + update, + upsert=True, + return_document=ReturnDocument.AFTER, + ) + if result is None: + raise RuntimeError("Failed to increment Mongo token usage") + return cast(dict[str, Any], result) + + @mongo_retry + async def increment_daily_usage( + self, + user_id: str, + tokens: int, + *, + date: str | None = None, + ) -> dict[str, Any]: + """Increment a user's total token usage for a UTC calendar day.""" + tokens = max(tokens, 0) + day = _validated_date(date) + now = datetime.now(UTC) + + result = await self._daily_collection.find_one_and_update( + {"user_id": user_id, "date": day}, + { + "$inc": {"total_tokens": tokens}, + "$set": {"updated_at": now}, + "$setOnInsert": {"user_id": user_id, "date": day}, + }, + upsert=True, + return_document=ReturnDocument.AFTER, + ) + if result is None: + raise RuntimeError("Failed to increment daily token usage") + return cast(dict[str, Any], result) + + @mongo_retry + async def get_thread_usage(self, thread_id: str) -> dict[str, Any] | None: + """Return the token usage document for *thread_id*, if present.""" + doc = await self._thread_collection.find_one({"thread_id": thread_id}) + return cast(dict[str, Any] | None, doc) + + @mongo_retry + async def get_daily_usage( + self, + user_id: str, + *, + date: str | None = None, + ) -> dict[str, Any] | None: + """Return a user's daily token rollup for the given UTC date.""" + day = _validated_date(date) + doc = await self._daily_collection.find_one({"user_id": user_id, "date": day}) + return cast(dict[str, Any] | None, doc) + + async def close(self) -> None: + """Close the underlying Motor client and release connections.""" + if self._client is not None: + self._client.close() + self._client = None diff --git a/deep_agent/src/token_budget/otel_emit.py b/deep_agent/src/token_budget/otel_emit.py new file mode 100644 index 00000000..23fd7211 --- /dev/null +++ b/deep_agent/src/token_budget/otel_emit.py @@ -0,0 +1,193 @@ +"""OTEL emission for per-call and daily token usage.""" + +from __future__ import annotations + +import threading +from datetime import UTC, datetime +from typing import Any + +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_counters_initialized = False +_token_counter: Any | None = None +_thread_total_counter: Any | None = None +_daily_total_counter: Any | None = None +_counters_lock = threading.Lock() + + +def token_budget_otel_enabled() -> bool: + """Return True when token usage metrics export is enabled.""" + return bool(settings.ENABLE_OTEL_METRICS and settings.OTEL_EXPORTER_OTLP_ENDPOINT) + + +def token_budget_traces_enabled() -> bool: + """Return True when token usage span events can be exported.""" + return bool( + settings.otel_traces_active() and settings.resolved_otel_traces_endpoint() + ) + + +def _agent_name() -> str: + try: + from deep_agent.src.agent.config import agent_config + + return agent_config.get_name() + except Exception: + return settings.OTEL_SERVICE_NAME + + +def _format_timestamp(value: Any | None = None) -> str: + if isinstance(value, datetime): + if value.tzinfo is None: + return value.replace(tzinfo=UTC).isoformat() + return value.isoformat() + if isinstance(value, str) and value: + return value + return datetime.now(UTC).isoformat() + + +def _ensure_counters() -> None: + global _counters_initialized, _token_counter, _thread_total_counter, _daily_total_counter # noqa: PLW0603 + + if _counters_initialized or not token_budget_otel_enabled(): + return + + with _counters_lock: + if _counters_initialized or not token_budget_otel_enabled(): + return + _counters_initialized = True + + try: + from opentelemetry import metrics + + meter = metrics.get_meter("template-agent.token-budget") + _token_counter = meter.create_counter( + "token_budget.tokens", + description="Billable LLM tokens recorded per call", + ) + _thread_total_counter = meter.create_counter( + "token_budget.thread_total", + description="Cumulative thread token totals after each LLM call", + ) + _daily_total_counter = meter.create_counter( + "token_budget.daily_total", + description="Cumulative per-user daily token totals after each LLM call", + ) + except Exception: + logger.warning("token_budget_otel_counter_init_failed", exc_info=True) + + +def emit_token_usage( + *, + thread_id: str, + user_id: str | None, + input_tokens: int, + output_tokens: int, + cumulative_total: int, + cumulative_input: int, + cumulative_output: int, + timestamp: Any | None = None, + trace_id: str | None = None, +) -> None: + """Emit OTEL metrics and optional span events for a single LLM usage record.""" + if not token_budget_otel_enabled(): + return + + _ensure_counters() + + recorded_at = _format_timestamp(timestamp) + agent_name = _agent_name() + attributes = { + "thread_id": thread_id, + "agent.name": agent_name, + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + "cumulative_total_tokens": cumulative_total, + "cumulative_input_tokens": cumulative_input, + "cumulative_output_tokens": cumulative_output, + "timestamp": recorded_at, + } + if trace_id: + attributes["app.trace_id"] = trace_id + if user_id: + attributes["user_id"] = user_id + + logger.info("token_budget_usage", **attributes) + + if token_budget_traces_enabled(): + try: + from opentelemetry import trace as otel_trace + + span = otel_trace.get_current_span() + if span is not None and span.is_recording(): + span.add_event("token_budget.usage", attributes=attributes) + except ImportError: + pass + + metric_attrs = { + "agent.name": agent_name, + "thread_id": thread_id, + } + if user_id: + metric_attrs["user_id"] = user_id + + if _token_counter is not None: + if input_tokens > 0: + _token_counter.add(input_tokens, {**metric_attrs, "token.type": "input"}) + if output_tokens > 0: + _token_counter.add(output_tokens, {**metric_attrs, "token.type": "output"}) + + if _thread_total_counter is not None and cumulative_total > 0: + _thread_total_counter.add( + cumulative_total, + {**metric_attrs, "aggregation": "cumulative"}, + ) + + +def emit_daily_token_usage( + *, + user_id: str, + total_tokens: int, + date: str, + timestamp: Any | None = None, +) -> None: + """Emit OTEL metrics and optional span events for a user's daily token rollup.""" + if not token_budget_otel_enabled(): + return + + _ensure_counters() + + recorded_at = _format_timestamp(timestamp) + attributes = { + "user_id": user_id, + "total_tokens": total_tokens, + "date": date, + "timestamp": recorded_at, + "agent.name": _agent_name(), + } + + logger.info("token_budget_daily_usage", **attributes) + + if token_budget_traces_enabled(): + try: + from opentelemetry import trace as otel_trace + + span = otel_trace.get_current_span() + if span is not None and span.is_recording(): + span.add_event("token_budget.daily_usage", attributes=attributes) + except ImportError: + pass + + if _daily_total_counter is not None and total_tokens > 0: + _daily_total_counter.add( + total_tokens, + { + "agent.name": _agent_name(), + "user_id": user_id, + "date": date, + }, + ) diff --git a/deep_agent/src/token_budget/service.py b/deep_agent/src/token_budget/service.py new file mode 100644 index 00000000..63541ece --- /dev/null +++ b/deep_agent/src/token_budget/service.py @@ -0,0 +1,283 @@ +"""Token usage tracking and extraction.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from deep_agent.src.agent.config import agent_config +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +if TYPE_CHECKING: + from deep_agent.src.token_budget.mongo_repository import TokenUsageMongoRepository + +logger = get_python_logger() + + +@dataclass(frozen=True) +class ThreadTokenUsage: + """Thread token usage summary.""" + + thread_id: str + used: int + input_tokens: int + output_tokens: int + + +class TokenUsageUnavailableError(Exception): + """Token usage storage is not configured or temporarily unreachable.""" + + +class TokenUsageNotFoundError(Exception): + """No token usage record exists for the requested thread.""" + + def __init__(self, thread_id: str) -> None: + """Store the thread ID that had no usage record.""" + self.thread_id = thread_id + super().__init__(thread_id) + + +def _reasoning_tokens(usage: dict[str, Any]) -> int: + """Return reasoning tokens from provider output_token_details (Gemini).""" + details = usage.get("output_token_details") + if not isinstance(details, dict): + return 0 + value = details.get("reasoning") + return int(value) if isinstance(value, int) else 0 + + +def _usage_dict_to_counts(usage: dict[str, Any] | None) -> tuple[int, int]: + """Map provider usage to billable input/output counts. + + Matches Langfuse **Total usage** (input + visible output + reasoning). + Gemini often reports ``output_tokens`` as visible output only (e.g. 29) while + ``output_token_details.reasoning`` holds the rest (e.g. 141). Prefer + ``total_tokens - input_tokens`` when both are present. + """ + if not usage: + return 0, 0 + input_tokens = int( + usage.get("input_tokens") + or usage.get("input") + or usage.get("prompt_tokens") + or usage.get("prompt_token_count") + or 0 + ) + total_tokens = int(usage.get("total_tokens") or usage.get("total") or 0) + if total_tokens > 0 and total_tokens >= input_tokens: + return input_tokens, total_tokens - input_tokens + + output_tokens = int( + usage.get("output_tokens") + or usage.get("output") + or usage.get("completion_tokens") + or usage.get("candidates_token_count") + or 0 + ) + reasoning = _reasoning_tokens(usage) + if reasoning: + output_tokens += reasoning + + if input_tokens or output_tokens: + return input_tokens, output_tokens + if total_tokens: + return 0, total_tokens + return 0, 0 + + +def _usage_from_generation(generation: Any) -> tuple[int, int]: + """Extract billable tokens from a single LangChain generation.""" + message = getattr(generation, "message", None) + if message is not None: + in_t, out_t = extract_tokens_from_message(message) + if in_t or out_t: + return in_t, out_t + + gen_info = getattr(generation, "generation_info", None) or {} + if isinstance(gen_info, dict): + usage = gen_info.get("usage_metadata") or gen_info.get("token_usage") or {} + if isinstance(usage, dict): + return _usage_dict_to_counts(usage) + + return 0, 0 + + +def extract_tokens_from_llm_result(response: Any) -> tuple[int, int]: + """Extract billable token counts from a LangChain LLMResult.""" + input_tokens = 0 + output_tokens = 0 + generations = getattr(response, "generations", None) or [] + for generation_list in generations: + for generation in generation_list: + in_t, out_t = _usage_from_generation(generation) + input_tokens += in_t + output_tokens += out_t + + if input_tokens or output_tokens: + return input_tokens, output_tokens + + llm_output = getattr(response, "llm_output", None) or {} + if isinstance(llm_output, dict): + token_usage = llm_output.get("token_usage") or llm_output.get("usage") or {} + if isinstance(token_usage, dict): + return _usage_dict_to_counts(token_usage) + + return 0, 0 + + +def extract_tokens_from_chat_result(response: Any) -> tuple[int, int]: + """Alias for chat model results — Langfuse routes these through on_llm_end.""" + return extract_tokens_from_llm_result(response) + + +def extract_tokens_from_message(message: Any) -> tuple[int, int]: + """Extract billable token counts from a LangChain message object.""" + usage = getattr(message, "usage_metadata", None) or {} + if isinstance(usage, dict) and usage: + return _usage_dict_to_counts(usage) + + response_metadata = getattr(message, "response_metadata", None) or {} + if isinstance(response_metadata, dict): + nested_usage = ( + response_metadata.get("usage_metadata") + or response_metadata.get("token_usage") + or {} + ) + if isinstance(nested_usage, dict): + in_t, out_t = _usage_dict_to_counts(nested_usage) + if in_t or out_t: + return in_t, out_t + + return 0, 0 + + +_mongo_repo_instance: TokenUsageMongoRepository | None = None +_mongo_repo_lock = threading.Lock() + + +def _mongo_repo() -> TokenUsageMongoRepository: + """Return a process-wide Mongo repository (reuses the Motor client pool).""" + global _mongo_repo_instance # noqa: PLW0603 + + if _mongo_repo_instance is None: + with _mongo_repo_lock: + if _mongo_repo_instance is None: + from deep_agent.src.token_budget.mongo_repository import ( + TokenUsageMongoRepository, + ) + + uri = settings.MONGODB_URI + if not uri: + raise TokenUsageUnavailableError( + "token budget tracking is not configured" + ) + _mongo_repo_instance = TokenUsageMongoRepository( + uri, + db_name=settings.MONGODB_DB, + ) + return _mongo_repo_instance + + +_MAX_REASONABLE_TOKENS = 1_000_000 + + +async def check_and_record( + thread_id: str, + input_tokens: int, + output_tokens: int, + *, + user_id: str | None = None, + trace_id: str | None = None, +) -> None: + """Increment thread usage, roll up daily user totals, and emit OTEL when enabled.""" + config = agent_config.get_token_budget_config() + if not config.is_active: + return + if not thread_id or thread_id == "unknown": + return + if not settings.MONGODB_URI: + logger.debug("token_budget_skipped_no_mongodb_uri") + return + if input_tokens <= 0 and output_tokens <= 0: + return + if input_tokens > _MAX_REASONABLE_TOKENS or output_tokens > _MAX_REASONABLE_TOKENS: + logger.warning( + "token_budget_suspicious_count", + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + return + + try: + repo = _mongo_repo() + agent_name = agent_config.get_name() + row = await repo.increment_usage( + thread_id, + input_tokens, + output_tokens, + agent_name=agent_name, + ) + total_delta = input_tokens + output_tokens + daily_row = None + if user_id and user_id != "unknown" and total_delta > 0: + daily_row = await repo.increment_daily_usage(user_id, total_delta) + except Exception: + logger.warning( + "token_budget_mongo_write_failed", + exc_info=True, + ) + return + + from deep_agent.src.token_budget.otel_emit import ( + emit_daily_token_usage, + emit_token_usage, + ) + + emit_token_usage( + thread_id=thread_id, + user_id=user_id, + input_tokens=input_tokens, + output_tokens=output_tokens, + cumulative_total=int(row["total_tokens"]), + cumulative_input=int(row["input_tokens"]), + cumulative_output=int(row["output_tokens"]), + timestamp=row.get("updated_at"), + trace_id=trace_id, + ) + + if daily_row is not None: + emit_daily_token_usage( + user_id=str(daily_row["user_id"]), + total_tokens=int(daily_row["total_tokens"]), + date=str(daily_row["date"]), + timestamp=daily_row.get("updated_at"), + ) + + +async def get_thread_token_usage(thread_id: str) -> ThreadTokenUsage: + """Return cumulative token usage for a thread.""" + config = agent_config.get_token_budget_config() + if not config.is_active or not settings.MONGODB_URI: + raise TokenUsageUnavailableError("token budget tracking is not configured") + + try: + repo = _mongo_repo() + row = await repo.get_thread_usage(thread_id) + except Exception as exc: + logger.warning( + "token_budget_mongo_read_failed", + exc_info=True, + ) + raise TokenUsageUnavailableError("token usage storage unavailable") from exc + + if row is None: + raise TokenUsageNotFoundError(thread_id) + + return ThreadTokenUsage( + thread_id=thread_id, + used=int(row["total_tokens"]), + input_tokens=int(row["input_tokens"]), + output_tokens=int(row["output_tokens"]), + ) diff --git a/deep_agent/src/triggers/__init__.py b/deep_agent/src/triggers/__init__.py new file mode 100644 index 00000000..e8c37bc5 --- /dev/null +++ b/deep_agent/src/triggers/__init__.py @@ -0,0 +1 @@ +"""Event-driven trigger system for headless agent mode.""" diff --git a/deep_agent/src/triggers/config.py b/deep_agent/src/triggers/config.py new file mode 100644 index 00000000..8b9dbeea --- /dev/null +++ b/deep_agent/src/triggers/config.py @@ -0,0 +1,112 @@ +"""Pydantic configuration models for headless mode triggers and sinks.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Any + +from pydantic import BaseModel, Field + + +class AgentMode(StrEnum): + """Agent runtime mode.""" + + SERVER = "server" + HEADLESS = "headless" + + +class WebhookTriggerConfig(BaseModel): + """Config for the webhook trigger source.""" + + enabled: bool = False + host: str = "0.0.0.0" + port: int = 8888 + path: str = "/trigger" + + +class CronJobConfig(BaseModel): + """A single cron job definition.""" + + name: str + schedule: str + payload: dict[str, Any] = Field(default_factory=dict) + + +class CronTriggerConfig(BaseModel): + """Config for the cron trigger source.""" + + enabled: bool = False + jobs: list[CronJobConfig] = Field(default_factory=list) + + +class QueueTriggerConfig(BaseModel): + """Config for the queue consumer trigger source.""" + + enabled: bool = False + backend: str = "redis_streams" + stream: str = "agent-tasks" + consumer_group: str = "agent-workers" + consumer_name: str = "" + bootstrap_servers: str = "localhost:9092" + topic: str = "agent-tasks" + + def get_consumer_name(self) -> str: + """Return consumer_name, defaulting to HOSTNAME for multi-replica support.""" + import os + + return self.consumer_name or os.environ.get("HOSTNAME", "worker-1") + + +class TriggerConfig(BaseModel): + """Container for all trigger source configs.""" + + webhook: WebhookTriggerConfig = Field(default_factory=WebhookTriggerConfig) + cron: CronTriggerConfig = Field(default_factory=CronTriggerConfig) + queue: QueueTriggerConfig = Field(default_factory=QueueTriggerConfig) + + +class OutputSinkFileConfig(BaseModel): + """File-specific sink config.""" + + path: str = "output.jsonl" + + +class OutputSinkWebhookConfig(BaseModel): + """Webhook-specific sink config.""" + + url: str = "" + headers: dict[str, str] = Field(default_factory=dict) + + +class OutputSinkRedisConfig(BaseModel): + """Redis-specific sink config.""" + + stream: str = "agent-results" + + +class OutputSinkConfig(BaseModel): + """Config for a single output sink.""" + + type: str + path: str | None = None + url: str | None = None + headers: dict[str, str] = Field(default_factory=dict) + stream: str | None = None + + +class HealthCheckConfig(BaseModel): + """Config for the headless worker health check endpoint.""" + + enabled: bool = True + host: str = "0.0.0.0" + port: int = 8080 + + +class HeadlessConfig(BaseModel): + """Top-level headless mode configuration.""" + + mode: AgentMode = AgentMode.SERVER + triggers: TriggerConfig = Field(default_factory=TriggerConfig) + output_sinks: list[OutputSinkConfig] = Field(default_factory=list) + drain_timeout: float = 30.0 + health_check: HealthCheckConfig = Field(default_factory=HealthCheckConfig) diff --git a/deep_agent/src/triggers/health.py b/deep_agent/src/triggers/health.py new file mode 100644 index 00000000..94c7f7e5 --- /dev/null +++ b/deep_agent/src/triggers/health.py @@ -0,0 +1,80 @@ +"""Health check endpoint for headless worker — serves /healthz and /readyz.""" + +from __future__ import annotations + +import asyncio +import json +from typing import TYPE_CHECKING + +from deep_agent.utils.pylogger import get_python_logger + +if TYPE_CHECKING: + from deep_agent.src.triggers.middleware import EventTriggerMiddleware + +logger = get_python_logger() + + +async def start_health_server( + host: str, + port: int, + middleware: EventTriggerMiddleware, +) -> asyncio.Server: + """Start a minimal HTTP server for liveness and readiness probes.""" + + async def _handle( + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + try: + request_line = await reader.readline() + if not request_line: + return + parts = request_line.decode("utf-8", errors="replace").strip().split(" ", 2) + path = parts[1] if len(parts) >= 2 else "/" + + while True: + line = await reader.readline() + if line in (b"\r\n", b"\n", b""): + break + + if path in ("/healthz", "/health"): + body = {"status": "ok"} + status = 200 + elif path == "/readyz": + has_sources = len(middleware._sources) > 0 + loop_running = ( + middleware._loop_task is not None + and not middleware._loop_task.done() + ) + ready = has_sources and loop_running + body = { + "status": "ready" if ready else "not_ready", + "sources": str(len(middleware._sources)), + "sinks": str(len(middleware._sinks)), + "loop_running": str(loop_running), + } + status = 200 if ready else 503 + else: + body = {"error": "not found"} + status = 404 + + payload = json.dumps(body).encode() + phrases = {200: "OK", 404: "Not Found", 503: "Service Unavailable"} + header = ( + f"HTTP/1.1 {status} {phrases.get(status, 'Error')}\r\n" + f"Content-Type: application/json\r\n" + f"Content-Length: {len(payload)}\r\n" + f"Connection: close\r\n" + f"\r\n" + ) + writer.write(header.encode() + payload) + await writer.drain() + except Exception: + pass + finally: + writer.close() + await writer.wait_closed() + + server = await asyncio.start_server(_handle, host, port) + logger.info("health check listening on %s:%d (/healthz, /readyz)", host, port) + return server diff --git a/deep_agent/src/triggers/middleware.py b/deep_agent/src/triggers/middleware.py new file mode 100644 index 00000000..039a55b5 --- /dev/null +++ b/deep_agent/src/triggers/middleware.py @@ -0,0 +1,306 @@ +"""EventTriggerMiddleware — orchestrates trigger sources, graph invocation, and output sinks.""" + +from __future__ import annotations + +import asyncio +import json +import time +from typing import Any + +from deep_agent.src.triggers.config import HeadlessConfig +from deep_agent.src.triggers.sinks.protocol import OutputSink, TriggerResult +from deep_agent.src.triggers.sources.protocol import TriggerEvent, TriggerSource +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +class EventTriggerMiddleware: + """Owns the lifecycle of trigger sources and output sinks in headless mode. + + Consumes events from all enabled trigger sources, invokes the agent graph + per event, and fans out results to all configured output sinks. + """ + + def __init__( + self, + config: HeadlessConfig, + graph: Any, + redis_url: str = "redis://redis:6379/0", + ) -> None: + """Initialize the event trigger middleware with config, graph, and Redis URL.""" + self._config = config + self._graph = graph + self._redis_url = redis_url + self._sources: list[TriggerSource] = [] + self._sinks: list[OutputSink] = [] + self._stop_event = asyncio.Event() + self._loop_task: asyncio.Task[None] | None = None + + from deep_agent.src.triggers.task_store import TaskStore + + self._task_store: TaskStore | None = ( + TaskStore(redis_url=self._redis_url) if self._redis_url else None + ) + + def _build_sources(self) -> list[TriggerSource]: + """Build trigger sources from configuration.""" + sources: list[TriggerSource] = [] + tc = self._config.triggers + + if tc.webhook.enabled: + from deep_agent.src.triggers.sources.webhook import WebhookTriggerSource + + sources.append(WebhookTriggerSource(tc.webhook)) + + if tc.cron.enabled: + from deep_agent.src.triggers.sources.cron import CronTriggerSource + + sources.append(CronTriggerSource(tc.cron)) + + if tc.queue.enabled: + from deep_agent.src.triggers.sources.queue import QueueTriggerSource + + sources.append(QueueTriggerSource(tc.queue, redis_url=self._redis_url)) + + return sources + + def _build_sinks(self) -> list[OutputSink]: + """Build output sinks from configuration.""" + sinks: list[OutputSink] = [] + + if not self._config.output_sinks: + from deep_agent.src.triggers.sinks.stdout import StdoutSink + + return [StdoutSink()] + + for sc in self._config.output_sinks: + if sc.type == "stdout": + from deep_agent.src.triggers.sinks.stdout import StdoutSink + + sinks.append(StdoutSink()) + elif sc.type == "file": + from deep_agent.src.triggers.sinks.file import FileSink + + sinks.append(FileSink(path=sc.path or "output.jsonl")) + elif sc.type == "webhook": + from deep_agent.src.triggers.sinks.webhook import WebhookSink + + sinks.append(WebhookSink(url=sc.url or "", headers=sc.headers or None)) + elif sc.type == "redis": + from deep_agent.src.triggers.sinks.redis import RedisSink + + sinks.append( + RedisSink( + stream=sc.stream or "agent-results", redis_url=self._redis_url + ) + ) + else: + logger.warning("Unknown output sink type: %s", sc.type) + + return sinks + + async def start(self) -> None: + """Start all trigger sources and begin the event processing loop.""" + logger.info("EventTriggerMiddleware starting") + self._sources = self._build_sources() + self._sinks = self._build_sinks() + + for source in self._sources: + await source.start() + + self._stop_event.clear() + self._loop_task = asyncio.create_task(self._run_loop()) + logger.info( + "EventTriggerMiddleware started: %d source(s), %d sink(s)", + len(self._sources), + len(self._sinks), + ) + + async def _run_loop(self) -> None: + async def _consume_source(source: TriggerSource) -> None: + async for event in source: + if self._stop_event.is_set(): + return + logger.info("Event received: %s (source=%s)", event.name, event.source) + await self._process_event(event) + + try: + async with asyncio.TaskGroup() as tg: + for source in self._sources: + tg.create_task(_consume_source(source)) + + async def _wait_for_stop() -> None: + await self._stop_event.wait() + raise _StopSentinel() + + tg.create_task(_wait_for_stop()) + except* _StopSentinel: + pass + except* Exception as eg: + for exc in eg.exceptions: + logger.exception("Source consumer error: %s", exc) + + async def _process_event(self, event: TriggerEvent) -> None: + store = self._task_store + task_id = event.payload.get("task_id") or event.metadata.get("task_id") + + if store and not task_id: + record = await store.create_task( + task_name=event.name, + payload=event.payload, + user_id=event.payload.get("user_id"), + ) + task_id = record.task_id + logger.info( + "Auto-created task record for %s event", + event.source, + task_id=task_id, + event_name=event.name, + ) + + if store and task_id: + await store.update_status(task_id, "processing") + + graph_timeout = self._config.drain_timeout * 4 + t0 = time.monotonic() + try: + try: + output = await asyncio.wait_for( + self._graph.ainvoke( + { + "messages": [ + {"role": "user", "content": json.dumps(event.payload)} + ] + } + ), + timeout=graph_timeout, + ) + except asyncio.TimeoutError: + logger.error( + "graph_invocation_timeout", + task_id=task_id, + timeout_seconds=graph_timeout, + ) + output = None + + duration_ms = (time.monotonic() - t0) * 1000 + result = TriggerResult( + event=event, + output=output, + duration_ms=duration_ms, + success=output is not None, + ) + if store and task_id: + if output is not None: + await store.update_status( + task_id, "completed", result=_extract_result(output) + ) + else: + await store.update_status( + task_id, "failed", error="graph invocation timed out" + ) + except Exception as exc: + duration_ms = (time.monotonic() - t0) * 1000 + result = TriggerResult( + event=event, + output=None, + duration_ms=duration_ms, + success=False, + error=str(exc), + ) + if store and task_id: + await store.update_status(task_id, "failed", error=str(exc)) + logger.exception("Graph invocation failed for event: %s", event.name) + + # Ack queue message after processing (B3 fix: ack-after-processing). + queue_msg = event.metadata.get("_queue_message") + queue_consumer = event.metadata.get("_consumer") + if queue_msg and queue_consumer: + try: + await queue_consumer.ack(queue_msg) + except Exception: + logger.warning("failed_to_ack_message", task_id=task_id) + + await self._emit_result(result) + logger.info( + "Event processed: %s (success=%s, duration=%.1fms)", + event.name, + result.success, + result.duration_ms, + ) + + async def _emit_result(self, result: TriggerResult) -> None: + for sink in self._sinks: + try: + await sink.emit(result) + except Exception: + logger.exception("Sink error (%s)", type(sink).__name__) + + async def stop(self) -> None: + """Stop all trigger sources and close output sinks.""" + logger.info("EventTriggerMiddleware stopping") + self._stop_event.set() + + if self._loop_task is not None: + try: + await asyncio.wait_for( + self._loop_task, timeout=self._config.drain_timeout + ) + except asyncio.TimeoutError: + logger.warning( + "Drain timeout (%.1fs) exceeded, cancelling", + self._config.drain_timeout, + ) + self._loop_task.cancel() + try: + await self._loop_task + except asyncio.CancelledError: + pass + self._loop_task = None + + for source in self._sources: + try: + await source.stop() + except Exception: + logger.exception("Error stopping source %s", type(source).__name__) + + for sink in self._sinks: + try: + await sink.close() + except Exception: + logger.exception("Error closing sink %s", type(sink).__name__) + + if self._task_store: + try: + await self._task_store.close() + except Exception: + logger.exception("Error closing task store") + self._task_store = None + + logger.info("EventTriggerMiddleware stopped") + + +def _extract_result(output: Any) -> str: + """Extract human-readable result text from graph output.""" + if isinstance(output, dict): + messages = output.get("messages", []) + for msg in reversed(messages): + content = getattr(msg, "content", None) or msg.get("content", "") + if isinstance(content, list): + texts = [ + c.get("text", "") + for c in content + if isinstance(c, dict) and c.get("text") + ] + if texts: + return "\n".join(texts) + elif isinstance(content, str) and content.strip(): + role = getattr(msg, "type", None) or msg.get("type", "") + if role in ("ai", "assistant"): + return content.strip() + return str(output)[:2000] + + +class _StopSentinel(BaseException): + """Raised to break out of the TaskGroup when stop is requested.""" diff --git a/deep_agent/src/triggers/runtime.py b/deep_agent/src/triggers/runtime.py new file mode 100644 index 00000000..54dbabf4 --- /dev/null +++ b/deep_agent/src/triggers/runtime.py @@ -0,0 +1,26 @@ +"""HeadlessRuntime — ServerRuntime adapter for headless mode.""" + +from __future__ import annotations + + +class HeadlessUser: + """Minimal user for headless mode (no SSO).""" + + def __init__(self, identity: str = "headless-worker") -> None: + """Initialize the headless user with the given identity.""" + self.identity = identity + self.access_token: str | None = None + self.refresh_token: str | None = None + + +class HeadlessRuntime: + """Adapter that provides a ServerRuntime-compatible interface for headless mode. + + The graph factory (graph.py:agent) reads runtime.user.access_token, + runtime.user.refresh_token, and runtime.user.identity. This class + provides those attributes without SSO. + """ + + def __init__(self, identity: str = "headless-worker") -> None: + """Initialize the headless runtime with the given identity.""" + self.user = HeadlessUser(identity) diff --git a/deep_agent/src/triggers/sinks/__init__.py b/deep_agent/src/triggers/sinks/__init__.py new file mode 100644 index 00000000..ae6d61e1 --- /dev/null +++ b/deep_agent/src/triggers/sinks/__init__.py @@ -0,0 +1 @@ +"""Output sink implementations (stdout, file, webhook, redis).""" diff --git a/deep_agent/src/triggers/sinks/file.py b/deep_agent/src/triggers/sinks/file.py new file mode 100644 index 00000000..824965d6 --- /dev/null +++ b/deep_agent/src/triggers/sinks/file.py @@ -0,0 +1,64 @@ +"""File output sink — appends results as JSONL.""" + +from __future__ import annotations + +import asyncio +import json +from dataclasses import asdict, is_dataclass +from datetime import datetime +from pathlib import Path +from typing import IO, Any + +from deep_agent.src.triggers.sinks.protocol import OutputSink, TriggerResult +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def _default_serializer(obj: Any) -> Any: + if isinstance(obj, datetime): + return obj.isoformat() + if is_dataclass(obj) and not isinstance(obj, type): + return asdict(obj) + if hasattr(obj, "dict"): + return obj.dict() + if hasattr(obj, "content"): + return {"type": type(obj).__name__, "content": obj.content} + return str(obj) + + +class FileSink(OutputSink): + """Appends TriggerResult as JSONL to a file.""" + + def __init__(self, path: str) -> None: + """Initialize the file sink with the output path.""" + self._path = Path(path) + self._handle: IO[str] | None = None + + def _ensure_handle(self) -> IO[str]: + if self._handle is None: + self._path.parent.mkdir(parents=True, exist_ok=True) + self._handle = open(self._path, "a", encoding="utf-8") # noqa: SIM115 + return self._handle + + async def emit(self, result: TriggerResult) -> None: + """Append the trigger result as a JSON line to the file.""" + try: + data = asdict(result) + line = json.dumps(data, default=_default_serializer) + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self._write_sync, line) + except Exception: + logger.exception("Failed to write to file sink: %s", self._path) + + def _write_sync(self, line: str) -> None: + """Synchronous write executed in a thread pool to avoid blocking the event loop.""" + handle = self._ensure_handle() + handle.write(line + "\n") + handle.flush() + + async def close(self) -> None: + """Close the file handle.""" + if self._handle is not None: + self._handle.close() + self._handle = None diff --git a/deep_agent/src/triggers/sinks/protocol.py b/deep_agent/src/triggers/sinks/protocol.py new file mode 100644 index 00000000..c8ac1364 --- /dev/null +++ b/deep_agent/src/triggers/sinks/protocol.py @@ -0,0 +1,32 @@ +"""Protocols and data classes for output sinks.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +from deep_agent.src.triggers.sources.protocol import TriggerEvent + + +@dataclass +class TriggerResult: + """Result of processing a trigger event through the agent graph.""" + + event: TriggerEvent + output: Any + duration_ms: float + success: bool + error: str | None = None + + +@runtime_checkable +class OutputSink(Protocol): + """Async output sink that receives trigger results.""" + + async def emit(self, result: TriggerResult) -> None: + """Emit a trigger result to the sink.""" + ... + + async def close(self) -> None: + """Close the sink and release resources.""" + ... diff --git a/deep_agent/src/triggers/sinks/redis.py b/deep_agent/src/triggers/sinks/redis.py new file mode 100644 index 00000000..ed7d60e3 --- /dev/null +++ b/deep_agent/src/triggers/sinks/redis.py @@ -0,0 +1,62 @@ +"""Redis output sink — publishes results to a Redis Stream.""" + +from __future__ import annotations + +import json +from dataclasses import asdict, is_dataclass +from datetime import datetime +from typing import Any + +from deep_agent.src.triggers.sinks.protocol import OutputSink, TriggerResult +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def _default_serializer(obj: Any) -> Any: + if isinstance(obj, datetime): + return obj.isoformat() + if is_dataclass(obj) and not isinstance(obj, type): + return asdict(obj) + if hasattr(obj, "dict"): + return obj.dict() + if hasattr(obj, "content"): + return {"type": type(obj).__name__, "content": obj.content} + return str(obj) + + +class RedisSink(OutputSink): + """Publishes TriggerResult to a Redis Stream via XADD.""" + + def __init__( + self, + stream: str, + redis_url: str = "redis://redis:6379/0", + ) -> None: + """Initialize the Redis sink with stream name and connection URL.""" + self._stream = stream + self._redis_url = redis_url + self._client: Any = None + + async def _ensure_client(self) -> Any: + if self._client is None: + import redis.asyncio as aioredis + + self._client = aioredis.from_url(self._redis_url, decode_responses=True) + return self._client + + async def emit(self, result: TriggerResult) -> None: + """Publish the trigger result to the Redis stream.""" + try: + client = await self._ensure_client() + data = asdict(result) + payload = json.dumps(data, default=_default_serializer) + await client.xadd(self._stream, {"result": payload}) + except Exception: + logger.exception("Failed to publish to Redis stream: %s", self._stream) + + async def close(self) -> None: + """Close the Redis client connection.""" + if self._client is not None: + await self._client.aclose() + self._client = None diff --git a/deep_agent/src/triggers/sinks/stdout.py b/deep_agent/src/triggers/sinks/stdout.py new file mode 100644 index 00000000..f9a00e08 --- /dev/null +++ b/deep_agent/src/triggers/sinks/stdout.py @@ -0,0 +1,38 @@ +"""Stdout output sink — writes results as JSON to stdout.""" + +from __future__ import annotations + +import json +import sys +from dataclasses import asdict, is_dataclass +from datetime import datetime +from typing import Any + +from deep_agent.src.triggers.sinks.protocol import OutputSink, TriggerResult + + +def _default_serializer(obj: Any) -> Any: + if isinstance(obj, datetime): + return obj.isoformat() + if is_dataclass(obj) and not isinstance(obj, type): + return asdict(obj) + if hasattr(obj, "dict"): + return obj.dict() + if hasattr(obj, "content"): + return {"type": type(obj).__name__, "content": obj.content} + return str(obj) + + +class StdoutSink(OutputSink): + """Writes TriggerResult as JSON to stdout.""" + + async def emit(self, result: TriggerResult) -> None: + """Write the trigger result as JSON to stdout.""" + data = asdict(result) + line = json.dumps(data, default=_default_serializer) + sys.stdout.write(line + "\n") + sys.stdout.flush() + + async def close(self) -> None: + """No-op close for stdout sink.""" + pass diff --git a/deep_agent/src/triggers/sinks/webhook.py b/deep_agent/src/triggers/sinks/webhook.py new file mode 100644 index 00000000..61841324 --- /dev/null +++ b/deep_agent/src/triggers/sinks/webhook.py @@ -0,0 +1,95 @@ +"""Webhook output sink — POSTs results to a URL.""" + +from __future__ import annotations + +import json +from dataclasses import asdict, is_dataclass +from datetime import datetime +from typing import Any + +import httpx + +from deep_agent.src.triggers.sinks.protocol import OutputSink, TriggerResult +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def _default_serializer(obj: Any) -> Any: + if isinstance(obj, datetime): + return obj.isoformat() + if is_dataclass(obj) and not isinstance(obj, type): + return asdict(obj) + if hasattr(obj, "dict"): + return obj.dict() + if hasattr(obj, "content"): + return {"type": type(obj).__name__, "content": obj.content} + return str(obj) + + +class WebhookSink(OutputSink): + """POSTs TriggerResult to a configured URL with retry.""" + + def __init__( + self, + url: str, + headers: dict[str, str] | None = None, + max_retries: int = 3, + timeout: float = 30.0, + ) -> None: + """Initialize the webhook sink with URL and retry settings.""" + self._url = url + self._headers = headers or {} + self._max_retries = max_retries + self._timeout = timeout + self._client: httpx.AsyncClient | None = None + + def _ensure_client(self) -> httpx.AsyncClient: + if self._client is None: + self._client = httpx.AsyncClient(timeout=self._timeout) + return self._client + + async def emit(self, result: TriggerResult) -> None: + """POST the trigger result to the configured webhook URL.""" + client = self._ensure_client() + data = asdict(result) + payload = json.dumps(data, default=_default_serializer) + + backoff = 1.0 + for attempt in range(self._max_retries + 1): + try: + resp = await client.post( + self._url, + content=payload, + headers={**self._headers, "Content-Type": "application/json"}, + ) + if resp.status_code < 500: + if resp.status_code >= 400: + logger.warning( + "Webhook sink got %d from %s", resp.status_code, self._url + ) + return + logger.warning( + "Webhook sink got %d (attempt %d/%d)", + resp.status_code, + attempt + 1, + self._max_retries + 1, + ) + except httpx.HTTPError: + logger.exception( + "Webhook sink error (attempt %d/%d)", + attempt + 1, + self._max_retries + 1, + ) + + if attempt < self._max_retries: + import asyncio + + await asyncio.sleep(backoff) + backoff *= 2 + + async def close(self) -> None: + """Close the HTTP client.""" + if self._client is not None: + await self._client.aclose() + self._client = None diff --git a/deep_agent/src/triggers/sources/__init__.py b/deep_agent/src/triggers/sources/__init__.py new file mode 100644 index 00000000..b4642467 --- /dev/null +++ b/deep_agent/src/triggers/sources/__init__.py @@ -0,0 +1 @@ +"""Trigger source implementations (webhook, cron, queue).""" diff --git a/deep_agent/src/triggers/sources/cron.py b/deep_agent/src/triggers/sources/cron.py new file mode 100644 index 00000000..7993c210 --- /dev/null +++ b/deep_agent/src/triggers/sources/cron.py @@ -0,0 +1,116 @@ +"""Cron trigger source — lightweight cron scheduler using croniter. + +Parses cron expressions from ``CronTriggerConfig``, computes the next +fire time for each job, and sleeps until it arrives. Each firing puts +a ``TriggerEvent`` onto an internal ``asyncio.Queue`` consumed through +the async-iterator protocol. + +Uses ``apscheduler.triggers.cron.CronTrigger`` only for parsing the +crontab expression — scheduling is done with plain ``asyncio.sleep``. +""" + +from __future__ import annotations + +import asyncio +from datetime import datetime, timezone +from typing import Any, AsyncIterator + +from deep_agent.src.triggers.config import CronJobConfig, CronTriggerConfig +from deep_agent.src.triggers.sources.protocol import TriggerEvent +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def _parse_cron_fields(schedule: str) -> dict[str, str] | None: + """Parse a 5-field crontab string into CronTrigger kwargs. Returns None on error.""" + parts = schedule.strip().split() + if len(parts) != 5: + return None + return { + "minute": parts[0], + "hour": parts[1], + "day": parts[2], + "month": parts[3], + "day_of_week": parts[4], + } + + +class CronTriggerSource: + """Fires ``TriggerEvent`` instances on cron schedules. + + Implements the ``TriggerSource`` protocol. ``start()`` launches a + background task per job that sleeps until the next fire time, emits + the event, and loops. ``stop()`` cancels all tasks. + """ + + def __init__(self, config: CronTriggerConfig) -> None: + """Initialize the cron trigger source with the given configuration.""" + self._config = config + self._queue: asyncio.Queue[TriggerEvent] = asyncio.Queue() + self._tasks: list[asyncio.Task[None]] = [] + + async def start(self) -> None: + """Start background tasks for each configured cron job.""" + for job_cfg in self._config.jobs: + fields = _parse_cron_fields(job_cfg.schedule) + if fields is None: + logger.warning( + "invalid cron schedule, skipping job", + job_name=job_cfg.name, + schedule=job_cfg.schedule, + ) + continue + + task = asyncio.create_task(self._run_job(job_cfg, fields)) + self._tasks.append(task) + logger.info( + "cron job scheduled", + job_name=job_cfg.name, + schedule=job_cfg.schedule, + ) + + logger.info("cron trigger source started", job_count=len(self._tasks)) + + async def _run_job(self, job_cfg: CronJobConfig, fields: dict[str, Any]) -> None: + """Background loop for a single cron job.""" + from apscheduler.triggers.cron import CronTrigger + + trigger = CronTrigger(**fields) + + while True: + now = datetime.now(timezone.utc) + next_fire = trigger.next() + if next_fire is None: + return + delay = (next_fire - now).total_seconds() + if delay > 0: + await asyncio.sleep(delay) + + event = TriggerEvent( + name=job_cfg.name, + payload=dict(job_cfg.payload), + source="cron", + metadata={"schedule": job_cfg.schedule}, + ) + await self._queue.put(event) + logger.debug("cron event fired", job_name=job_cfg.name) + + async def stop(self) -> None: + """Cancel all running cron job tasks.""" + for task in self._tasks: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + self._tasks.clear() + logger.info("cron trigger source stopped") + + def __aiter__(self) -> AsyncIterator[TriggerEvent]: + """Return the async iterator.""" + return self + + async def __anext__(self) -> TriggerEvent: + """Return the next cron trigger event.""" + return await self._queue.get() diff --git a/deep_agent/src/triggers/sources/kafka_consumer.py b/deep_agent/src/triggers/sources/kafka_consumer.py new file mode 100644 index 00000000..2f32ceb6 --- /dev/null +++ b/deep_agent/src/triggers/sources/kafka_consumer.py @@ -0,0 +1,82 @@ +"""Kafka consumer implementation of the QueueConsumer protocol.""" + +from __future__ import annotations + +import json +from typing import Any, AsyncIterator + +from deep_agent.src.triggers.sources.queue import QueueMessage +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +class KafkaQueueConsumer: + """Consumes messages from a Kafka topic using aiokafka. + + Implements the ``QueueConsumer`` protocol so it can be used + as a drop-in replacement for ``RedisStreamsConsumer``. + """ + + def __init__( + self, + topic: str, + bootstrap_servers: str = "localhost:9092", + consumer_group: str = "agent-workers", + ) -> None: + """Initialize with Kafka connection settings.""" + self._topic = topic + self._servers = bootstrap_servers + self._group = consumer_group + self._consumer: Any = None + self._running = True + + async def consume(self) -> AsyncIterator[QueueMessage]: + """Consume messages from the Kafka topic.""" + from aiokafka import AIOKafkaConsumer + + self._consumer = AIOKafkaConsumer( + self._topic, + bootstrap_servers=self._servers, + group_id=self._group, + value_deserializer=lambda v: json.loads(v.decode("utf-8")), + auto_offset_reset="earliest", + enable_auto_commit=False, + ) + await self._consumer.start() + logger.info( + "kafka consumer started", + topic=self._topic, + group=self._group, + servers=self._servers, + ) + + try: + async for msg in self._consumer: + if not self._running: + return + data = ( + msg.value + if isinstance(msg.value, dict) + else {"payload": str(msg.value)} + ) + yield QueueMessage( + id=f"{msg.partition}-{msg.offset}", + data=data, + ) + except Exception: + if self._running: + logger.error("kafka consumer error", exc_info=True) + + async def ack(self, message: QueueMessage) -> None: + """Manually commit offsets after successful processing.""" + if self._consumer is not None: + await self._consumer.commit() + + async def close(self) -> None: + """Stop the Kafka consumer.""" + self._running = False + if self._consumer is not None: + await self._consumer.stop() + self._consumer = None + logger.info("kafka consumer stopped") diff --git a/deep_agent/src/triggers/sources/protocol.py b/deep_agent/src/triggers/sources/protocol.py new file mode 100644 index 00000000..c277f7a8 --- /dev/null +++ b/deep_agent/src/triggers/sources/protocol.py @@ -0,0 +1,39 @@ +"""Protocols and data classes for trigger sources.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, AsyncIterator, Protocol, runtime_checkable + + +@dataclass(frozen=True) +class TriggerEvent: + """An event produced by a trigger source.""" + + name: str + payload: dict[str, Any] + source: str + metadata: dict[str, Any] = field(default_factory=dict) + timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + + +@runtime_checkable +class TriggerSource(Protocol): + """Async trigger source that yields events.""" + + async def start(self) -> None: + """Start the trigger source.""" + ... + + async def stop(self) -> None: + """Stop the trigger source.""" + ... + + def __aiter__(self) -> AsyncIterator[TriggerEvent]: + """Return the async iterator.""" + ... + + async def __anext__(self) -> TriggerEvent: + """Return the next trigger event.""" + ... diff --git a/deep_agent/src/triggers/sources/queue.py b/deep_agent/src/triggers/sources/queue.py new file mode 100644 index 00000000..d6987507 --- /dev/null +++ b/deep_agent/src/triggers/sources/queue.py @@ -0,0 +1,305 @@ +"""Queue trigger source — abstract consumer protocol + Redis Streams. + +Defines a ``QueueConsumer`` protocol with ``consume``/``ack``/``close`` +methods, a concrete ``RedisStreamsConsumer`` implementation using +``redis.asyncio``, and a ``QueueTriggerSource`` adapter that wraps any +``QueueConsumer`` into the ``TriggerSource`` async-iterator interface. +""" + +from __future__ import annotations + +import asyncio +import os +from dataclasses import dataclass +from typing import Any, AsyncIterator, Protocol, runtime_checkable + +from deep_agent.src.triggers.config import QueueTriggerConfig +from deep_agent.src.triggers.sources.protocol import TriggerEvent +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +# Default Redis URL (matches the project convention in aegra/redis.py). +_DEFAULT_REDIS_URL = os.environ.get("REDIS_URL", "redis://redis:6379/0") + +# Maximum backoff delay for reconnection (seconds). +_MAX_BACKOFF = 60.0 + +# Number of messages to fetch per XREADGROUP call. +_READ_COUNT = 10 + +# Block timeout for XREADGROUP (milliseconds). +_BLOCK_MS = 5_000 + + +# ------------------------------------------------------------------ +# Queue abstractions +# ------------------------------------------------------------------ + + +@dataclass +class QueueMessage: + """A single message consumed from a queue backend.""" + + id: str + data: dict[str, Any] + + +@runtime_checkable +class QueueConsumer(Protocol): + """Protocol for consuming messages from a queue backend.""" + + def consume(self) -> AsyncIterator[QueueMessage]: + """Consume messages from the queue backend.""" + ... + + async def ack(self, message: QueueMessage) -> None: + """Acknowledge a consumed message.""" + ... + + async def close(self) -> None: + """Close the consumer and release resources.""" + ... + + +# ------------------------------------------------------------------ +# Redis Streams implementation +# ------------------------------------------------------------------ + + +class RedisStreamsConsumer: + """Consumes messages from a Redis Stream using consumer groups. + + Handles: + - Automatic consumer group creation (``XGROUP CREATE ... MKSTREAM``). + - Blocking reads via ``XREADGROUP``. + - Message acknowledgment via ``XACK``. + - Reconnection with exponential backoff on connection failures. + """ + + def __init__( + self, + stream: str, + consumer_group: str, + consumer_name: str, + redis_url: str = _DEFAULT_REDIS_URL, + block_ms: int = _BLOCK_MS, + read_count: int = _READ_COUNT, + ) -> None: + """Initialize the Redis Streams consumer with connection and group settings.""" + self._stream = stream + self._group = consumer_group + self._consumer = consumer_name + self._redis_url = redis_url + self._block_ms = block_ms + self._read_count = read_count + self._client: Any = None + self._running = True + + async def _ensure_client(self) -> Any: + """Lazily create the Redis client and consumer group.""" + if self._client is not None: + return self._client + + import redis.asyncio as aioredis + + self._client = aioredis.from_url( + self._redis_url, + decode_responses=True, + ) + + # Create the consumer group if it does not already exist. + try: + await self._client.xgroup_create( + self._stream, self._group, id="0", mkstream=True + ) + logger.info( + "redis consumer group created", + stream=self._stream, + group=self._group, + ) + except Exception as exc: + # BUSYGROUP means the group already exists — safe to ignore. + if "BUSYGROUP" not in str(exc): + raise + + return self._client + + async def consume(self) -> AsyncIterator[QueueMessage]: + """Yield messages from the stream, reconnecting on failure.""" + backoff = 1.0 + + while self._running: + try: + client = await self._ensure_client() + results = await client.xreadgroup( + self._group, + self._consumer, + {self._stream: ">"}, + count=self._read_count, + block=self._block_ms, + ) + # Reset backoff on successful read (even if no messages). + backoff = 1.0 + + if not results: + continue + + for _stream_name, messages in results: + for msg_id, data in messages: + yield QueueMessage(id=msg_id, data=data) + + except asyncio.CancelledError: + return + except Exception: + logger.error( + "redis streams consumer error, reconnecting", + backoff_seconds=backoff, + stream=self._stream, + exc_info=True, + ) + # Tear down the broken client so _ensure_client rebuilds it. + await self._close_client() + await asyncio.sleep(backoff) + backoff = min(backoff * 2, _MAX_BACKOFF) + + async def ack(self, message: QueueMessage) -> None: + """Acknowledge a consumed message.""" + client = await self._ensure_client() + await client.xack(self._stream, self._group, message.id) + + async def close(self) -> None: + """Stop consuming and close the Redis connection.""" + self._running = False + await self._close_client() + + async def _close_client(self) -> None: + """Close the underlying Redis connection if open.""" + if self._client is not None: + try: + await self._client.aclose() + except Exception: + logger.debug("redis client close error", exc_info=True) + self._client = None + + +# ------------------------------------------------------------------ +# TriggerSource adapter +# ------------------------------------------------------------------ + + +class QueueTriggerSource: + """Adapts a ``QueueConsumer`` into the ``TriggerSource`` protocol. + + Messages consumed from the queue are placed on an internal event + queue with the original message and consumer reference in metadata, + allowing downstream middleware to acknowledge after processing. + """ + + def __init__( + self, + config: QueueTriggerConfig, + redis_url: str = _DEFAULT_REDIS_URL, + ) -> None: + """Initialize the queue trigger source with configuration and Redis URL.""" + self._config = config + self._redis_url = redis_url + self._queue: asyncio.Queue[TriggerEvent] = asyncio.Queue() + self._consumer: RedisStreamsConsumer | QueueConsumer | None = None + self._task: asyncio.Task[None] | None = None + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def start(self) -> None: + """Create the queue consumer and start the consume loop.""" + if self._config.backend == "redis_streams": + consumer_name = self._config.get_consumer_name() + self._consumer = RedisStreamsConsumer( + stream=self._config.stream, + consumer_group=self._config.consumer_group, + consumer_name=consumer_name, + redis_url=self._redis_url, + ) + elif self._config.backend == "kafka": + from deep_agent.src.triggers.sources.kafka_consumer import ( + KafkaQueueConsumer, + ) + + self._consumer = KafkaQueueConsumer( + topic=self._config.topic, + bootstrap_servers=self._config.bootstrap_servers, + consumer_group=self._config.consumer_group, + ) + else: + raise ValueError(f"unsupported queue backend: {self._config.backend}") + + self._task = asyncio.create_task(self._consume_loop()) + logger.info( + "queue trigger source started", + backend=self._config.backend, + stream=self._config.stream, + consumer_group=self._config.consumer_group, + consumer_name=self._config.get_consumer_name(), + ) + + async def stop(self) -> None: + """Cancel the consume loop and close the consumer.""" + if self._task is not None: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + if self._consumer is not None: + await self._consumer.close() + self._consumer = None + logger.info("queue trigger source stopped") + + # ------------------------------------------------------------------ + # Async-iterator protocol + # ------------------------------------------------------------------ + + def __aiter__(self) -> AsyncIterator[TriggerEvent]: + """Return the async iterator.""" + return self + + async def __anext__(self) -> TriggerEvent: + """Return the next queue trigger event.""" + return await self._queue.get() + + # ------------------------------------------------------------------ + # Internal + # ------------------------------------------------------------------ + + async def _consume_loop(self) -> None: + """Read messages from the consumer and adapt to TriggerEvent.""" + if self._consumer is None: + return + + try: + async for message in self._consumer.consume(): + event = TriggerEvent( + name=message.data.get("name", "queue-event"), + payload=dict(message.data), + source="queue", + metadata={ + "message_id": message.id, + "stream": self._config.stream, + "_queue_message": message, + "_consumer": self._consumer, + }, + ) + await self._queue.put(event) + # Do NOT ack here — middleware will ack after processing. + logger.debug( + "queue event enqueued", + event_name=event.name, + message_id=message.id, + ) + except asyncio.CancelledError: + pass + except Exception: + logger.error("queue consume loop error", exc_info=True) diff --git a/deep_agent/src/triggers/sources/webhook.py b/deep_agent/src/triggers/sources/webhook.py new file mode 100644 index 00000000..bd656720 --- /dev/null +++ b/deep_agent/src/triggers/sources/webhook.py @@ -0,0 +1,229 @@ +"""Webhook trigger source — lightweight HTTP listener using asyncio. + +Accepts POST requests at a configurable path, parses the JSON body, +and yields ``TriggerEvent`` instances through the async-iterator +protocol. No third-party HTTP framework is required; the server is +built on top of ``asyncio.start_server`` and raw HTTP parsing. +""" + +from __future__ import annotations + +import asyncio +import json +from typing import Any, AsyncIterator + +from deep_agent.src.triggers.config import WebhookTriggerConfig +from deep_agent.src.triggers.sources.protocol import TriggerEvent +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +# Maximum request body size (1 MiB) to prevent unbounded memory usage. +_MAX_BODY_SIZE = 1_048_576 + +# Read timeout for an individual HTTP request (seconds). +_READ_TIMEOUT = 30.0 + + +def _build_response(status: int, reason: str, body: dict[str, Any]) -> bytes: + """Build a minimal HTTP/1.1 response with a JSON body.""" + payload = json.dumps(body).encode() + lines = [ + f"HTTP/1.1 {status} {reason}", + "Content-Type: application/json", + f"Content-Length: {len(payload)}", + "Connection: close", + "", + "", + ] + return "\r\n".join(lines).encode() + payload + + +class WebhookTriggerSource: + """Async HTTP listener that emits ``TriggerEvent`` for each POST. + + Implements the ``TriggerSource`` protocol — ``start``/``stop`` control + the server lifecycle, and ``__aiter__``/``__anext__`` pull events from + an internal ``asyncio.Queue``. + """ + + def __init__(self, config: WebhookTriggerConfig) -> None: + """Initialize the webhook trigger source with the given configuration.""" + self._config = config + self._queue: asyncio.Queue[TriggerEvent] = asyncio.Queue() + self._server: asyncio.Server | None = None + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def start(self) -> None: + """Bind and start the HTTP server.""" + self._server = await asyncio.start_server( + self._handle_connection, + host=self._config.host, + port=self._config.port, + ) + addrs = [s.getsockname() for s in self._server.sockets] + logger.info( + "webhook trigger listening", + host=self._config.host, + port=self._config.port, + path=self._config.path, + addresses=addrs, + ) + + async def stop(self) -> None: + """Shut down the HTTP server gracefully.""" + if self._server is None: + return + self._server.close() + await self._server.wait_closed() + self._server = None + logger.info("webhook trigger stopped") + + # ------------------------------------------------------------------ + # Async-iterator protocol + # ------------------------------------------------------------------ + + def __aiter__(self) -> AsyncIterator[TriggerEvent]: + """Return the async iterator.""" + return self + + async def __anext__(self) -> TriggerEvent: + """Return the next webhook trigger event.""" + return await self._queue.get() + + # ------------------------------------------------------------------ + # Connection handler + # ------------------------------------------------------------------ + + async def _handle_connection( + self, + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + """Handle a single TCP connection (one HTTP request).""" + try: + await asyncio.wait_for( + self._process_request(reader, writer), + timeout=_READ_TIMEOUT, + ) + except asyncio.TimeoutError: + logger.warning("webhook request timed out") + except Exception: + logger.error("webhook request error", exc_info=True) + finally: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + + async def _process_request( + self, + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + """Parse the HTTP request and dispatch based on method/path.""" + # Read the request line. + request_line = await reader.readline() + if not request_line: + return + + parts = request_line.decode("utf-8", errors="replace").strip().split(" ", 2) + if len(parts) < 2: + await self._send( + writer, 400, "Bad Request", {"error": "malformed request line"} + ) + return + + method, path = parts[0], parts[1] + + # Read headers. + content_length = 0 + while True: + header_line = await reader.readline() + if header_line in (b"\r\n", b"\n", b""): + break + header = header_line.decode("utf-8", errors="replace").strip() + if header.lower().startswith("content-length:"): + try: + content_length = int(header.split(":", 1)[1].strip()) + except (ValueError, IndexError): + pass + + # Route: only POST to the configured path is accepted. + if path != self._config.path: + await self._send(writer, 404, "Not Found", {"error": "not found"}) + return + + if method.upper() != "POST": + await self._send( + writer, + 405, + "Method Not Allowed", + {"error": f"method {method} not allowed"}, + ) + return + + # Guard against oversized bodies. + if content_length > _MAX_BODY_SIZE: + await self._send( + writer, + 413, + "Payload Too Large", + {"error": f"body exceeds {_MAX_BODY_SIZE} bytes"}, + ) + return + + # Read body. + body_bytes = await reader.read(content_length) if content_length > 0 else b"" + + # Parse JSON. + try: + payload: Any = json.loads(body_bytes) if body_bytes else {} + except (json.JSONDecodeError, UnicodeDecodeError) as exc: + await self._send( + writer, + 400, + "Bad Request", + {"error": f"invalid JSON: {exc}"}, + ) + return + + if not isinstance(payload, dict): + await self._send( + writer, + 400, + "Bad Request", + {"error": "request body must be a JSON object"}, + ) + return + + # Build event and enqueue. + event_name = payload.pop("event", "webhook") + event = TriggerEvent( + name=str(event_name), + payload=payload, + source="webhook", + ) + await self._queue.put(event) + logger.debug("webhook event enqueued", event_name=event.name) + + await self._send(writer, 200, "OK", {"status": "accepted"}) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + async def _send( + writer: asyncio.StreamWriter, + status: int, + reason: str, + body: dict[str, Any], + ) -> None: + """Write an HTTP response and drain.""" + writer.write(_build_response(status, reason, body)) + await writer.drain() diff --git a/deep_agent/src/triggers/task_repository.py b/deep_agent/src/triggers/task_repository.py new file mode 100644 index 00000000..87b1adf9 --- /dev/null +++ b/deep_agent/src/triggers/task_repository.py @@ -0,0 +1,163 @@ +"""Async Postgres repository for task audit trail. + +Persists every task status change to PostgreSQL for audit purposes. +Redis remains the primary store for speed; Postgres is the durable +record that survives Redis TTL expiry. +""" + +from __future__ import annotations + +import json +from typing import Any + +import psycopg +from psycopg.rows import dict_row + +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_TABLES_ENSURED = False + +CREATE_TASKS_TABLE = """ +CREATE TABLE IF NOT EXISTS tasks ( + task_id TEXT PRIMARY KEY, + task_name TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'queued', + payload JSONB NOT NULL DEFAULT '{}', + result TEXT, + error TEXT, + thread_id TEXT, + user_id TEXT, + delivered BOOLEAN NOT NULL DEFAULT FALSE, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + completed_at TIMESTAMPTZ +); +CREATE INDEX IF NOT EXISTS idx_tasks_user_id ON tasks (user_id); +CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks (status); +CREATE INDEX IF NOT EXISTS idx_tasks_created_at ON tasks (created_at DESC); +""" + + +class TaskRepository: + """Async Postgres repository for task audit records.""" + + def __init__(self, database_uri: str) -> None: + """Initialize with a PostgreSQL connection URI.""" + self._database_uri = database_uri + + async def ensure_table(self) -> None: + """Create the tasks table if it doesn't exist.""" + global _TABLES_ENSURED # noqa: PLW0603 + if _TABLES_ENSURED: + return + try: + async with await psycopg.AsyncConnection.connect( + self._database_uri + ) as conn: + await conn.execute(CREATE_TASKS_TABLE) + await conn.commit() + _TABLES_ENSURED = True + logger.info("tasks audit table ensured") + except Exception: + logger.warning("failed to create tasks table", exc_info=True) + + async def insert_task( + self, + task_id: str, + task_name: str, + payload: dict[str, Any], + thread_id: str | None = None, + user_id: str | None = None, + ) -> None: + """Insert a new task record.""" + try: + async with await psycopg.AsyncConnection.connect( + self._database_uri + ) as conn: + await conn.execute( + """INSERT INTO tasks (task_id, task_name, status, payload, thread_id, user_id) + VALUES (%s, %s, 'queued', %s, %s, %s) + ON CONFLICT (task_id) DO NOTHING""", + (task_id, task_name, json.dumps(payload, default=str), thread_id, user_id), + ) + await conn.commit() + except Exception: + logger.warning("failed to insert task audit record", task_id=task_id, exc_info=True) + + async def update_status( + self, + task_id: str, + status: str, + result: str | None = None, + error: str | None = None, + ) -> None: + """Update task status in the audit table.""" + try: + completed_at = "now()" if status in ("completed", "failed") else None + async with await psycopg.AsyncConnection.connect( + self._database_uri + ) as conn: + if completed_at: + await conn.execute( + """UPDATE tasks SET status = %s, result = %s, error = %s, + updated_at = now(), completed_at = now() + WHERE task_id = %s""", + (status, result, error, task_id), + ) + else: + await conn.execute( + """UPDATE tasks SET status = %s, result = %s, error = %s, + updated_at = now() WHERE task_id = %s""", + (status, result, error, task_id), + ) + await conn.commit() + except Exception: + logger.warning("failed to update task audit record", task_id=task_id, exc_info=True) + + async def mark_delivered(self, task_id: str) -> None: + """Mark task as delivered in the audit table.""" + try: + async with await psycopg.AsyncConnection.connect( + self._database_uri + ) as conn: + await conn.execute( + "UPDATE tasks SET delivered = TRUE, updated_at = now() WHERE task_id = %s", + (task_id,), + ) + await conn.commit() + except Exception: + logger.warning("failed to mark task delivered", task_id=task_id, exc_info=True) + + async def get_task_history( + self, + user_id: str | None = None, + status: str | None = None, + limit: int = 50, + ) -> list[dict[str, Any]]: + """Query task history for audit purposes.""" + try: + async with await psycopg.AsyncConnection.connect( + self._database_uri, row_factory=dict_row + ) as conn: + conditions = [] + params: list[Any] = [] + if user_id: + conditions.append("user_id = %s") + params.append(user_id) + if status: + conditions.append("status = %s") + params.append(status) + + where = f"WHERE {' AND '.join(conditions)}" if conditions else "" + params.append(limit) + + rows = await conn.execute( + f"SELECT * FROM tasks {where} ORDER BY created_at DESC LIMIT %s", + params, + ) + return [dict(r) for r in await rows.fetchall()] + except Exception: + logger.warning("failed to query task history", exc_info=True) + return [] diff --git a/deep_agent/src/triggers/task_store.py b/deep_agent/src/triggers/task_store.py new file mode 100644 index 00000000..e7960244 --- /dev/null +++ b/deep_agent/src/triggers/task_store.py @@ -0,0 +1,211 @@ +"""Redis-backed task status store for tracking headless worker tasks.""" + +from __future__ import annotations + +import json +import uuid +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from typing import Any + +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_TASK_TTL = 86400 # 24 hours +_KEY_PREFIX = "task:" +_USER_INDEX_PREFIX = "user_tasks:" + + +def _now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +@dataclass +class TaskRecord: + """A background task tracked in Redis.""" + + task_id: str + task_name: str + status: str + payload: dict[str, Any] = field(default_factory=dict) + result: Any = None + error: str | None = None + thread_id: str | None = None + user_id: str | None = None + created_at: str = "" + updated_at: str = "" + delivered: bool = False + + def to_json(self) -> str: + """Serialize the task record to a JSON string.""" + data = asdict(self) + if isinstance(data.get("result"), (dict, list)): + pass + elif data.get("result") is not None: + data["result"] = str(data["result"]) + return json.dumps(data, default=str) + + @classmethod + def from_json(cls, raw: str) -> TaskRecord: + """Deserialize a task record from a JSON string.""" + data = json.loads(raw) + return cls(**data) + + +class TaskStore: + """Redis-backed store for task status tracking with Postgres audit trail.""" + + def __init__(self, redis_url: str | None = None) -> None: + """Initialize the task store with an optional Redis URL.""" + self._redis_url = redis_url or settings.REDIS_URL + self._client: Any = None + self._audit: Any = None + + def _get_audit_repo(self) -> Any: + """Lazily create the Postgres audit repository.""" + if self._audit is None: + try: + from deep_agent.src.triggers.task_repository import TaskRepository + + self._audit = TaskRepository(settings.database_uri) + except Exception: + logger.debug("audit repository unavailable", exc_info=True) + return self._audit + + async def _ensure_client(self) -> Any: + if self._client is None: + import redis.asyncio as aioredis + + self._client = aioredis.from_url(self._redis_url, decode_responses=True) + return self._client + + async def create_task( + self, + task_name: str, + payload: dict[str, Any], + thread_id: str | None = None, + user_id: str | None = None, + ) -> TaskRecord: + """Create a new task record and store it in Redis.""" + task_id = uuid.uuid4().hex[:12] + now = _now_iso() + record = TaskRecord( + task_id=task_id, + task_name=task_name, + status="queued", + payload=payload, + thread_id=thread_id, + user_id=user_id, + created_at=now, + updated_at=now, + ) + client = await self._ensure_client() + key = f"{_KEY_PREFIX}{task_id}" + await client.set(key, record.to_json(), ex=_TASK_TTL) + + if user_id: + await client.zadd( + f"{_USER_INDEX_PREFIX}{user_id}", + {task_id: datetime.now(timezone.utc).timestamp()}, + ) + await client.expire(f"{_USER_INDEX_PREFIX}{user_id}", _TASK_TTL) + + logger.info( + "task created", task_id=task_id, task_name=task_name, status="queued" + ) + + audit = self._get_audit_repo() + if audit: + try: + await audit.ensure_table() + await audit.insert_task(task_id, task_name, payload, thread_id, user_id) + except Exception: + logger.debug("audit insert failed", task_id=task_id, exc_info=True) + + return record + + async def update_status( + self, + task_id: str, + status: str, + result: Any = None, + error: str | None = None, + ) -> None: + """Update the status of an existing task.""" + client = await self._ensure_client() + key = f"{_KEY_PREFIX}{task_id}" + raw = await client.get(key) + if raw is None: + logger.warning("task not found for status update", task_id=task_id) + return + + record = TaskRecord.from_json(raw) + record.status = status + record.updated_at = _now_iso() + if result is not None: + record.result = result + if error is not None: + record.error = error + + ttl = await client.ttl(key) + await client.set(key, record.to_json(), ex=max(ttl, 3600)) + logger.info("task status updated", task_id=task_id, status=status) + + audit = self._get_audit_repo() + if audit: + try: + result_str = str(result)[:2000] if result is not None else None + await audit.update_status(task_id, status, result_str, error) + except Exception: + logger.debug("audit update failed", task_id=task_id, exc_info=True) + + async def get_task(self, task_id: str) -> TaskRecord | None: + """Retrieve a task record by ID, or None if not found.""" + client = await self._ensure_client() + raw = await client.get(f"{_KEY_PREFIX}{task_id}") + if raw is None: + return None + return TaskRecord.from_json(raw) + + async def get_pending_results(self, user_id: str) -> list[TaskRecord]: + """Return completed or failed tasks that have not been delivered.""" + client = await self._ensure_client() + task_ids = await client.zrange(f"{_USER_INDEX_PREFIX}{user_id}", 0, -1) + pending = [] + for tid in task_ids: + record = await self.get_task(tid) + if ( + record + and record.status in ("completed", "failed") + and not record.delivered + ): + pending.append(record) + return pending + + async def mark_delivered(self, task_id: str) -> None: + """Mark a task as delivered to the user.""" + client = await self._ensure_client() + key = f"{_KEY_PREFIX}{task_id}" + raw = await client.get(key) + if raw is None: + return + record = TaskRecord.from_json(raw) + record.delivered = True + record.updated_at = _now_iso() + ttl = await client.ttl(key) + await client.set(key, record.to_json(), ex=max(ttl, 3600)) + + audit = self._get_audit_repo() + if audit: + try: + await audit.mark_delivered(task_id) + except Exception: + logger.debug("audit mark_delivered failed", task_id=task_id, exc_info=True) + + async def close(self) -> None: + """Close the Redis client connection.""" + if self._client is not None: + await self._client.aclose() + self._client = None diff --git a/deep_agent/src/triggers/tools.py b/deep_agent/src/triggers/tools.py new file mode 100644 index 00000000..956090ef --- /dev/null +++ b/deep_agent/src/triggers/tools.py @@ -0,0 +1,159 @@ +"""Tools for interacting with the headless worker from the server agent.""" + +from __future__ import annotations + +import json +from typing import Any + +from deep_agent.src.settings import settings +from deep_agent.src.triggers.task_store import TaskStore +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger() + +_DEFAULT_STREAM = "agent-tasks" +_store = TaskStore() + + +async def queue_task( + task_name: str, + payload: dict[str, Any], + thread_id: str = "", + user_id: str = "", + stream: str = _DEFAULT_STREAM, +) -> str: + """Queue a background task for the headless worker. + + Use this for long-running or bulk work that doesn't need an + immediate response. The headless worker picks it up from Redis + and processes it asynchronously. + + Args: + task_name: A descriptive name (e.g. "generate-report"). + payload: The task data sent to the headless worker. + thread_id: Current conversation thread ID (for tracking). + user_id: Current user ID (for result delivery). + stream: Redis Stream name. Defaults to "agent-tasks". + + Returns: + Confirmation with the task ID that can be used to check status. + """ + record = await _store.create_task( + task_name=task_name, + payload=payload, + thread_id=thread_id or None, + user_id=user_id or None, + ) + + import redis.asyncio as aioredis + + client = aioredis.from_url(settings.REDIS_URL, decode_responses=True) + try: + fields = { + "name": task_name, + "task_id": record.task_id, + "payload": json.dumps(payload), + } + await client.xadd(stream, fields) + logger.info( + "task queued for headless worker", + task_id=record.task_id, + task_name=task_name, + stream=stream, + ) + finally: + await client.aclose() + + return ( + f"Task '{task_name}' queued with ID {record.task_id}. " + f"The headless worker will process it in the background. " + f"Use check_task_status with this ID to check progress." + ) + + +async def check_task_status(task_id: str) -> str: + """Check the status of a background task. + + Args: + task_id: The task ID returned by queue_task. + + Returns: + Task status and result if completed. + """ + record = await _store.get_task(task_id) + if record is None: + return f"Task '{task_id}' not found. It may have expired (tasks are kept for 24 hours)." + + if record.status == "completed": + result_text = str(record.result) if record.result else "No output" + return ( + f"Task '{record.task_name}' (ID: {task_id}) is COMPLETED.\n" + f"Result:\n{result_text}" + ) + elif record.status == "failed": + return ( + f"Task '{record.task_name}' (ID: {task_id}) FAILED.\nError: {record.error}" + ) + else: + return f"Task '{record.task_name}' (ID: {task_id}) is {record.status.upper()}." + + +async def get_pending_results(user_id: str) -> str: + """Get completed background tasks that haven't been delivered yet. + + Call this at the start of a conversation to check if any + background tasks have completed since the user's last visit. + + Args: + user_id: The user's identity. + + Returns: + Summary of pending results, or a message saying there are none. + """ + pending = await _store.get_pending_results(user_id) + if not pending: + return "No pending background task results." + + lines = [f"{len(pending)} background task(s) completed since your last visit:\n"] + for record in pending: + if record.status == "completed": + result_text = str(record.result) if record.result else "No output" + lines.append( + f"- **{record.task_name}** (ID: {record.task_id}): COMPLETED\n Result:\n {result_text}\n" + ) + elif record.status == "failed": + lines.append( + f"- **{record.task_name}** (ID: {record.task_id}): FAILED\n Error: {record.error}\n" + ) + await _store.mark_delivered(record.task_id) + + return "\n".join(lines) + + +def get_builtin_tools() -> list[Any]: + """Return LangChain-compatible tool objects for the headless worker tools.""" + from langchain_core.tools import StructuredTool + + return [ + StructuredTool.from_function( + coroutine=queue_task, + name="queue_task", + description=( + "Queue a background task for the headless worker. Use for long-running, " + "bulk, or fire-and-forget work. Returns a task ID for status tracking." + ), + ), + StructuredTool.from_function( + coroutine=check_task_status, + name="check_task_status", + description="Check the status of a background task by its task ID.", + ), + StructuredTool.from_function( + coroutine=get_pending_results, + name="get_pending_results", + description=( + "Get completed background tasks not yet delivered to the user. " + "Call at the start of every conversation to check for results." + ), + ), + ] diff --git a/template_agent/utils/__init__.py b/deep_agent/utils/__init__.py similarity index 100% rename from template_agent/utils/__init__.py rename to deep_agent/utils/__init__.py diff --git a/deep_agent/utils/google_creds.py b/deep_agent/utils/google_creds.py new file mode 100644 index 00000000..3940d2a2 --- /dev/null +++ b/deep_agent/utils/google_creds.py @@ -0,0 +1,72 @@ +"""Google credentials management utilities. + +This module provides functions for initializing Google Generative AI with +service account credentials from environment variables. +""" + +import json + +from google.auth.credentials import Credentials +from google.oauth2 import service_account + +from deep_agent.src.settings import settings +from deep_agent.utils.pylogger import get_python_logger + +logger = get_python_logger(log_level=settings.PYTHON_LOG_LEVEL) + +# Google Cloud authentication scope for Vertex AI +GOOGLE_AUTH_SCOPES = ["https://www.googleapis.com/auth/cloud-platform"] + +# Cache for credentials to avoid repeated credential fetches +_credentials_cache: tuple[Credentials, str] | None = None + + +def get_service_account_credentials() -> tuple[Credentials, str]: + """Get Google Cloud credentials from service account JSON. + + Reads service account JSON from GOOGLE_APPLICATION_CREDENTIALS_CONTENT + environment variable and creates credentials. Uses caching to avoid + repeated credential fetches. + + Returns: + Tuple of (credentials, project_id) + + Raises: + RuntimeError: If credentials cannot be loaded or project ID is missing + """ + global _credentials_cache + + if _credentials_cache is not None: + return _credentials_cache + + if not settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT: + raise RuntimeError("No Google service account credentials configured") + + try: + service_account_info = json.loads( + settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT + ) + except json.JSONDecodeError as e: + logger.error(f"Invalid JSON in credentials: {e}") + raise RuntimeError(f"Invalid JSON in credentials: {e}") from e + + project = service_account_info.get("project_id") + if not project: + raise RuntimeError("Service account JSON does not contain 'project_id' field") + + credentials = service_account.Credentials.from_service_account_info( + service_account_info, scopes=GOOGLE_AUTH_SCOPES + ) + + logger.info(f"Loaded Google credentials for project: {project}") + _credentials_cache = (credentials, project) + return _credentials_cache + + +def clear_credentials_cache() -> None: + """Clear the cached Google Cloud credentials. + + Useful for testing or when credentials need to be refreshed. + """ + global _credentials_cache + _credentials_cache = None diff --git a/deep_agent/utils/pylogger.py b/deep_agent/utils/pylogger.py new file mode 100644 index 00000000..333115a9 --- /dev/null +++ b/deep_agent/utils/pylogger.py @@ -0,0 +1,339 @@ +"""Structured logger utility for the template-agent. + +Provides a single ``get_python_logger()`` entry point that returns a +structlog ``BoundLogger``. All log output is structured JSON by default +(production), with an optional human-readable console renderer for +local development. + +Environment variables: + LOG_FORMAT: ``json`` (default) or ``console`` + PYTHON_LOG_LEVEL: standard level name (default: INFO) + +Context binding: + ``bind_request_context(trace_id, user_id, thread_id)`` adds + per-request fields to every subsequent log line in the same + async context. Call at request entry; structlog's context-var + support auto-clears on context exit. +""" + +import logging +import os +import sys +from contextvars import ContextVar +from typing import Any + +import structlog + +# --------------------------------------------------------------------------- +# Third-party logger noise suppression +# --------------------------------------------------------------------------- + +HTTP_CLIENT_LOGGERS = { + "urllib3", + "urllib3.connectionpool", + "urllib3.util", + "urllib3.util.retry", + "requests", + "httpx", +} + +AWS_LOGGERS = { + "botocore", + "botocore.client", + "botocore.credentials", + "botocore.httpsession", + "boto3", + "boto3.resources", +} + +MCP_LOGGERS = { + "fastmcp", + "fastmcp.server", + "fastmcp.server.http", + "fastmcp.utilities", + "fastmcp.utilities.logging", + "fastmcp.client", + "fastmcp.transports", +} + +ML_AI_LOGGERS = { + "sentence_transformers", + "transformers", + "transformers.models", + "transformers.tokenization_utils", + "transformers.tokenization_utils_base", + "transformers.configuration_utils", + "transformers.modeling_utils", + "huggingface_hub", + "huggingface_hub.utils", + "langchain_huggingface", + "torch", + "torch.nn", +} + +OBSERVABILITY_LOGGERS = { + "langfuse", + "langfuse.client", + "langfuse.api", + "langfuse.callback", +} + +SILENT_LOGGERS: set[str] = set() + +THIRD_PARTY_LOGGERS: set[str] = ( + HTTP_CLIENT_LOGGERS + | AWS_LOGGERS + | MCP_LOGGERS + | ML_AI_LOGGERS + | OBSERVABILITY_LOGGERS + | SILENT_LOGGERS +) + +ERROR_ONLY_LOGGERS: set[str] = ML_AI_LOGGERS | OBSERVABILITY_LOGGERS + +_LOGGING_CONFIGURED = False + +SERVICE_NAME = os.environ.get("SERVICE_NAME", "template-agent") +LOG_FORMAT = os.environ.get("LOG_FORMAT", "json").lower() + +for _name in SILENT_LOGGERS: + logging.getLogger(_name).setLevel(logging.CRITICAL) + +# --------------------------------------------------------------------------- +# Request context (per-request fields via contextvars) +# --------------------------------------------------------------------------- + +_trace_id_var: ContextVar[str | None] = ContextVar("trace_id", default=None) +_user_id_var: ContextVar[str | None] = ContextVar("user_id", default=None) +_thread_id_var: ContextVar[str | None] = ContextVar("thread_id", default=None) +_request_id_var: ContextVar[str | None] = ContextVar("request_id", default=None) +_org_id_var: ContextVar[str | None] = ContextVar("org_id", default=None) +_agent_id_var: ContextVar[str | None] = ContextVar("agent_id", default=None) + + +def bind_request_context( + trace_id: str | None = None, + user_id: str | None = None, + thread_id: str | None = None, + request_id: str | None = None, + org_id: str | None = None, + agent_id: str | None = None, +) -> None: + """Bind per-request identifiers into the logging context. + + Call this once at request entry. The values are automatically + injected into every log line within the same async context. + """ + if trace_id: + _trace_id_var.set(trace_id) + if user_id: + _user_id_var.set(user_id) + if thread_id: + _thread_id_var.set(thread_id) + if request_id: + _request_id_var.set(request_id) + if org_id: + _org_id_var.set(org_id) + if agent_id: + _agent_id_var.set(agent_id) + + +def clear_request_context() -> None: + """Reset request context (called at request exit).""" + _trace_id_var.set(None) + _user_id_var.set(None) + _thread_id_var.set(None) + _request_id_var.set(None) + _org_id_var.set(None) + _agent_id_var.set(None) + + +def _inject_request_context( + logger: Any, method_name: str, event_dict: dict[str, Any] +) -> dict[str, Any]: + """Structlog processor: inject request context vars into every log event.""" + rid = _trace_id_var.get() + uid = _user_id_var.get() + tid = _thread_id_var.get() + req_id = _request_id_var.get() + oid = _org_id_var.get() + aid = _agent_id_var.get() + if rid: + event_dict["trace_id"] = rid + if uid: + event_dict["user_id"] = uid + if tid: + event_dict["thread_id"] = tid + if req_id: + event_dict["request_id"] = req_id + if oid: + event_dict["org_id"] = oid + if aid: + event_dict["agent_id"] = aid + event_dict["service"] = SERVICE_NAME + return event_dict + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _clear_handlers(logger: logging.Logger) -> None: + logger.handlers.clear() + logger.filters.clear() + + +def _setup_logger(logger_name: str, level: str) -> None: + lgr = logging.getLogger(logger_name) + _clear_handlers(lgr) + if logger_name in SILENT_LOGGERS: + lgr.setLevel(logging.CRITICAL) + elif logger_name in ERROR_ONLY_LOGGERS: + lgr.setLevel(logging.ERROR) + else: + lgr.setLevel(level) + lgr.propagate = True + + +def _configure_third_party_loggers(log_level: str) -> None: + """Apply structured logging to selected third-party loggers.""" + logging.getLogger().handlers.clear() + for name in THIRD_PARTY_LOGGERS: + _setup_logger(name, log_level) + + +def _get_renderer() -> Any: + """Return the appropriate structlog renderer based on LOG_FORMAT.""" + if LOG_FORMAT == "console": + return structlog.dev.ConsoleRenderer(colors=True) + return structlog.processors.JSONRenderer() + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def force_reconfigure_all_loggers(log_level: str = "INFO") -> None: + """Force logger reconfiguration, even if already initialized.""" + global _LOGGING_CONFIGURED # noqa: PLW0603 + _LOGGING_CONFIGURED = False + get_python_logger(log_level) + + +def get_python_logger(log_level: str = "INFO") -> structlog.BoundLogger: + """Get a configured structlog logger. + + First call configures the entire logging pipeline. Subsequent + calls return cached loggers from structlog. + """ + global _LOGGING_CONFIGURED # noqa: PLW0603 + log_level = log_level.upper() + + if not _LOGGING_CONFIGURED: + logging.basicConfig( + format="%(message)s", + stream=sys.stdout, + level=log_level, + ) + + structlog.configure( + processors=[ + structlog.stdlib.filter_by_level, + structlog.stdlib.add_logger_name, + structlog.stdlib.add_log_level, + structlog.stdlib.PositionalArgumentsFormatter(), + structlog.processors.TimeStamper(fmt="iso"), + _inject_request_context, + structlog.processors.StackInfoRenderer(), + structlog.processors.format_exc_info, + structlog.processors.UnicodeDecoder(), + _get_renderer(), + ], + context_class=dict, + logger_factory=structlog.stdlib.LoggerFactory(), + wrapper_class=structlog.stdlib.BoundLogger, + cache_logger_on_first_use=True, + ) + + _LOGGING_CONFIGURED = True + + _configure_third_party_loggers(log_level) + return structlog.get_logger() + + +def get_uvicorn_log_config(log_level: str = "INFO") -> dict[str, Any]: + """Return a Uvicorn-compatible logging config that integrates with structlog.""" + log_level = log_level.upper() + renderer = _get_renderer() + + default_formatter = { + "()": "structlog.stdlib.ProcessorFormatter", + "processor": renderer, + "foreign_pre_chain": [ + structlog.stdlib.add_log_level, + structlog.processors.TimeStamper(fmt="iso"), + _inject_request_context, + structlog.processors.StackInfoRenderer(), + structlog.processors.format_exc_info, + structlog.processors.UnicodeDecoder(), + ], + } + + def make_logger_config(names: list[str], level: str) -> dict[str, Any]: + return { + name: { + "handlers": ["default"], + "level": level, + "propagate": False, + } + for name in names + } + + passthrough_formatter = {"format": "%(message)s"} + + uvicorn_loggers = ["uvicorn", "uvicorn.error", "uvicorn.asgi", "uvicorn.protocols"] + access_loggers = ["uvicorn.access"] + + return { + "version": 1, + "disable_existing_loggers": False, + "formatters": { + "default": default_formatter, + "access": default_formatter, + "passthrough": passthrough_formatter, + }, + "handlers": { + "default": { + "formatter": "default", + "class": "logging.StreamHandler", + "stream": "ext://sys.stdout", + }, + "access": { + "formatter": "access", + "class": "logging.StreamHandler", + "stream": "ext://sys.stdout", + }, + "passthrough": { + "formatter": "passthrough", + "class": "logging.StreamHandler", + "stream": "ext://sys.stdout", + }, + }, + "loggers": { + "": { + "handlers": ["passthrough"], + "level": log_level, + "propagate": False, + }, + **make_logger_config(uvicorn_loggers, log_level), + **make_logger_config(access_loggers, log_level), + **make_logger_config( + list(THIRD_PARTY_LOGGERS - ERROR_ONLY_LOGGERS - SILENT_LOGGERS), + log_level, + ), + **make_logger_config(list(ERROR_ONLY_LOGGERS), "ERROR"), + **make_logger_config(list(SILENT_LOGGERS), "CRITICAL"), + }, + } diff --git a/deployment/base/configmap.yaml b/deployment/base/configmap.yaml new file mode 100644 index 00000000..29c49c32 --- /dev/null +++ b/deployment/base/configmap.yaml @@ -0,0 +1,11 @@ +apiVersion: v1 +kind: ConfigMap +metadata: + name: agent-config +data: + # Config Path (required for base image pattern) + # Config path — mount config/agent at this path (see Containerfile) + CONFIG_PATH: "/app/config/agent" + + # Note: PostgreSQL and Redis environment variables are added by + # optional components when postgres or redis components are included diff --git a/deployment/base/kustomization.yaml b/deployment/base/kustomization.yaml new file mode 100644 index 00000000..367f497e --- /dev/null +++ b/deployment/base/kustomization.yaml @@ -0,0 +1,12 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +# Core resources - always deployed +# Postgres and Redis are now optional components +resources: + - configmap.yaml + - secret.yaml + +labels: + - pairs: + app: template-agent diff --git a/deployment/base/secret.yaml b/deployment/base/secret.yaml new file mode 100644 index 00000000..1b21c566 --- /dev/null +++ b/deployment/base/secret.yaml @@ -0,0 +1,26 @@ +apiVersion: v1 +kind: Secret +metadata: + name: agent-secrets +type: Opaque +stringData: + # PostgreSQL credentials (default for in-cluster deployment) + POSTGRES_USER: "postgres" + POSTGRES_PASSWORD: "postgres" + + # SSO / OIDC Authentication + SSO_ISSUER_URL: "" + SSO_CLIENT_ID: "" + SSO_CLIENT_SECRET: "" + + # Langfuse (optional - external service) + LANGFUSE_PUBLIC_KEY: "" + LANGFUSE_SECRET_KEY: "" + LANGFUSE_BASE_URL: "" + + # Google Vertex AI (optional) + GOOGLE_APPLICATION_CREDENTIALS_CONTENT: "" + + # vLLM / OpenAI-compatible (optional) + VLLM_BASE_URL: "" + VLLM_API_KEY: "" diff --git a/deployment/components/postgres/deployment.yaml b/deployment/components/postgres/deployment.yaml new file mode 100644 index 00000000..11732b08 --- /dev/null +++ b/deployment/components/postgres/deployment.yaml @@ -0,0 +1,90 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: pgvector + labels: + component: database +spec: + replicas: 1 + strategy: + type: Recreate + selector: + matchLabels: + component: database + template: + metadata: + labels: + component: database + spec: + containers: + - name: pgvector + image: pgvector/pgvector:pg16 + env: + - name: POSTGRES_USER + valueFrom: + secretKeyRef: + name: agent-secrets + key: POSTGRES_USER + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: agent-secrets + key: POSTGRES_PASSWORD + - name: POSTGRES_DB + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_DB + - name: PGDATA + value: "/var/lib/postgresql/data/pgdata" + ports: + - containerPort: 5432 + name: postgres + protocol: TCP + livenessProbe: + exec: + command: + - pg_isready + - -U + - pgvector + initialDelaySeconds: 30 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + readinessProbe: + exec: + command: + - pg_isready + - -U + - pgvector + initialDelaySeconds: 5 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 3 + resources: + requests: + memory: "256Mi" + cpu: "100m" + limits: + memory: "512Mi" + cpu: "500m" + volumeMounts: + - name: postgres-data + mountPath: /var/lib/postgresql/data + - name: postgres-init + mountPath: /docker-entrypoint-initdb.d + securityContext: + runAsNonRoot: true + runAsUser: 999 + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + volumes: + - name: postgres-data + persistentVolumeClaim: + claimName: postgres-pvc + - name: postgres-init + configMap: + name: postgres-init + restartPolicy: Always diff --git a/deployment/components/postgres/init-configmap.yaml b/deployment/components/postgres/init-configmap.yaml new file mode 100644 index 00000000..ad10a8e4 --- /dev/null +++ b/deployment/components/postgres/init-configmap.yaml @@ -0,0 +1,9 @@ +apiVersion: v1 +kind: ConfigMap +metadata: + name: postgres-init + labels: + component: database +data: + init-databases.sql: | + CREATE DATABASE mcp_server; diff --git a/deployment/components/postgres/kustomization.yaml b/deployment/components/postgres/kustomization.yaml new file mode 100644 index 00000000..975f3db8 --- /dev/null +++ b/deployment/components/postgres/kustomization.yaml @@ -0,0 +1,23 @@ +apiVersion: kustomize.config.k8s.io/v1alpha1 +kind: Component + +resources: + - deployment.yaml + - pvc.yaml + - service.yaml + - init-configmap.yaml + +patches: + - target: + kind: ConfigMap + name: agent-config + patch: |- + - op: add + path: /data/POSTGRES_HOST + value: "postgres" + - op: add + path: /data/POSTGRES_PORT + value: "5432" + - op: add + path: /data/POSTGRES_DB + value: "template_agent" diff --git a/deployment/components/postgres/pvc.yaml b/deployment/components/postgres/pvc.yaml new file mode 100644 index 00000000..dca0ad53 --- /dev/null +++ b/deployment/components/postgres/pvc.yaml @@ -0,0 +1,12 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: postgres-pvc + labels: + component: database +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 5Gi diff --git a/deployment/components/postgres/service.yaml b/deployment/components/postgres/service.yaml new file mode 100644 index 00000000..ba4f27d6 --- /dev/null +++ b/deployment/components/postgres/service.yaml @@ -0,0 +1,15 @@ +apiVersion: v1 +kind: Service +metadata: + name: postgres + labels: + component: database +spec: + type: ClusterIP + ports: + - port: 5432 + targetPort: 5432 + protocol: TCP + name: postgres + selector: + component: database diff --git a/deployment/components/redis/deployment.yaml b/deployment/components/redis/deployment.yaml new file mode 100644 index 00000000..362216f3 --- /dev/null +++ b/deployment/components/redis/deployment.yaml @@ -0,0 +1,71 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: redis + labels: + component: cache +spec: + replicas: 1 + selector: + matchLabels: + component: cache + template: + metadata: + labels: + component: cache + spec: + containers: + - name: redis + image: redis:7-alpine + command: + - redis-server + - --appendonly + - "yes" + - --maxmemory + - "256mb" + - --maxmemory-policy + - "allkeys-lru" + ports: + - containerPort: 6379 + name: redis + protocol: TCP + livenessProbe: + exec: + command: + - redis-cli + - ping + initialDelaySeconds: 30 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + readinessProbe: + exec: + command: + - redis-cli + - ping + initialDelaySeconds: 5 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 3 + resources: + requests: + memory: "128Mi" + cpu: "100m" + limits: + memory: "256Mi" + cpu: "250m" + volumeMounts: + - name: redis-data + mountPath: /data + securityContext: + runAsNonRoot: true + runAsUser: 999 + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + volumes: + - name: redis-data + persistentVolumeClaim: + claimName: redis-pvc + restartPolicy: Always diff --git a/deployment/components/redis/kustomization.yaml b/deployment/components/redis/kustomization.yaml new file mode 100644 index 00000000..77ff980e --- /dev/null +++ b/deployment/components/redis/kustomization.yaml @@ -0,0 +1,16 @@ +apiVersion: kustomize.config.k8s.io/v1alpha1 +kind: Component + +resources: + - deployment.yaml + - pvc.yaml + - service.yaml + +patches: + - target: + kind: ConfigMap + name: agent-config + patch: |- + - op: add + path: /data/REDIS_URL + value: "redis://redis:6379/0" diff --git a/deployment/components/redis/pvc.yaml b/deployment/components/redis/pvc.yaml new file mode 100644 index 00000000..98924871 --- /dev/null +++ b/deployment/components/redis/pvc.yaml @@ -0,0 +1,12 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: redis-pvc + labels: + component: cache +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 1Gi diff --git a/deployment/components/redis/service.yaml b/deployment/components/redis/service.yaml new file mode 100644 index 00000000..d195d69e --- /dev/null +++ b/deployment/components/redis/service.yaml @@ -0,0 +1,15 @@ +apiVersion: v1 +kind: Service +metadata: + name: redis + labels: + component: cache +spec: + type: ClusterIP + ports: + - port: 6379 + targetPort: 6379 + protocol: TCP + name: redis + selector: + component: cache diff --git a/deployment/openshift/configmap.yaml b/deployment/openshift/configmap.yaml deleted file mode 100644 index dd2f980a..00000000 --- a/deployment/openshift/configmap.yaml +++ /dev/null @@ -1,16 +0,0 @@ -apiVersion: v1 -kind: ConfigMap -metadata: - name: template-agent-config - labels: - app: template-agent - component: agent -data: - PYTHON_LOG_LEVEL: "INFO" - USE_INMEMORY_SAVER: "false" - LANGFUSE_TRACING_ENVIRONMENT: "production" - MCP_SERVER_NAME: "template-mcp-server" - MCP_SERVER_URL: "https://template-mcp-server.ns.svc:8443/mcp/" - MCP_TRANSPORT_PROTOCOL: "streamable_http" - MCP_CONNECTION_TIMEOUT: "30" - MCP_SSL_VERIFY: "false" diff --git a/deployment/openshift/route.yaml b/deployment/openshift/route.yaml deleted file mode 100644 index 1489e427..00000000 --- a/deployment/openshift/route.yaml +++ /dev/null @@ -1,16 +0,0 @@ -apiVersion: route.openshift.io/v1 -kind: Route -metadata: - name: template-agent - labels: - app: template-agent - component: agent -spec: - to: - kind: Service - name: template-agent - port: - targetPort: http - tls: - termination: edge - insecureEdgeTerminationPolicy: Redirect diff --git a/deployment/openshift/secret.yaml b/deployment/openshift/secret.yaml deleted file mode 100644 index 6c31a907..00000000 --- a/deployment/openshift/secret.yaml +++ /dev/null @@ -1,20 +0,0 @@ -apiVersion: v1 -kind: Secret -metadata: - name: template-agent-secrets - labels: - app: template-agent - component: agent -type: Opaque -stringData: - POSTGRES_HOST: "" - POSTGRES_PORT: "5432" - POSTGRES_DB: "" - POSTGRES_USER: "pgvector" - POSTGRES_PASSWORD: "CHANGE_ME" - LANGFUSE_PUBLIC_KEY: "" - LANGFUSE_SECRET_KEY: "" - LANGFUSE_BASE_URL: "" - GOOGLE_APPLICATION_CREDENTIALS_CONTENT: "" - SESSION_SECRET: "" - SNOWFLAKE_ACCOUNT: "" diff --git a/deployment/overlays/kind-headless/deployment.yaml b/deployment/overlays/kind-headless/deployment.yaml new file mode 100644 index 00000000..556654a8 --- /dev/null +++ b/deployment/overlays/kind-headless/deployment.yaml @@ -0,0 +1,118 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: agent-headless + labels: + app: agent + component: agent-headless +spec: + replicas: 1 + selector: + matchLabels: + app: agent + component: agent-headless + template: + metadata: + labels: + app: agent + component: agent-headless + spec: + containers: + - name: agent-headless + image: agent:local + imagePullPolicy: Never + command: ["python", "-m", "deep_agent.headless"] + ports: + - containerPort: 8080 + name: health + protocol: TCP + env: + - name: ENVIRONMENT + value: "local" + - name: PYTHON_LOG_LEVEL + value: "INFO" + - name: POSTGRES_HOST + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_HOST + - name: POSTGRES_PORT + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_PORT + - name: POSTGRES_DB + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_DB + - name: POSTGRES_USER + valueFrom: + secretKeyRef: + name: agent-secrets + key: POSTGRES_USER + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: agent-secrets + key: POSTGRES_PASSWORD + - name: REDIS_URL + valueFrom: + configMapKeyRef: + name: agent-config + key: REDIS_URL + - name: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + valueFrom: + secretKeyRef: + name: agent-secrets + key: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + optional: true + - name: LANGFUSE_PUBLIC_KEY + valueFrom: + secretKeyRef: + name: agent-secrets + key: LANGFUSE_PUBLIC_KEY + optional: true + - name: LANGFUSE_SECRET_KEY + valueFrom: + secretKeyRef: + name: agent-secrets + key: LANGFUSE_SECRET_KEY + optional: true + - name: LANGFUSE_BASE_URL + valueFrom: + secretKeyRef: + name: agent-secrets + key: LANGFUSE_BASE_URL + optional: true + startupProbe: + httpGet: + path: /healthz + port: 8080 + initialDelaySeconds: 10 + periodSeconds: 5 + failureThreshold: 30 + livenessProbe: + httpGet: + path: /healthz + port: 8080 + initialDelaySeconds: 30 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + readinessProbe: + httpGet: + path: /readyz + port: 8080 + initialDelaySeconds: 10 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 3 + resources: + requests: + memory: "256Mi" + cpu: "100m" + limits: + memory: "512Mi" + cpu: "500m" + restartPolicy: Always diff --git a/deployment/openshift/kustomization.yaml b/deployment/overlays/kind-headless/kustomization.yaml similarity index 52% rename from deployment/openshift/kustomization.yaml rename to deployment/overlays/kind-headless/kustomization.yaml index c5fa9cf8..e120b99a 100644 --- a/deployment/openshift/kustomization.yaml +++ b/deployment/overlays/kind-headless/kustomization.yaml @@ -1,15 +1,18 @@ apiVersion: kustomize.config.k8s.io/v1beta1 kind: Kustomization + +namespace: template-agent + resources: - - buildconfig.yaml - - imagestream.yaml - - configmap.yaml - - secret.yaml - deployment.yaml - service.yaml - - route.yaml + labels: - pairs: app: template-agent - component: agent - includeSelectors: true + component: agent-headless + +images: + - name: agent + newName: localhost/template-agent + newTag: local diff --git a/deployment/overlays/kind-headless/service.yaml b/deployment/overlays/kind-headless/service.yaml new file mode 100644 index 00000000..63941db8 --- /dev/null +++ b/deployment/overlays/kind-headless/service.yaml @@ -0,0 +1,17 @@ +apiVersion: v1 +kind: Service +metadata: + name: agent-headless + labels: + app: agent + component: agent-headless +spec: + selector: + app: agent + component: agent-headless + ports: + - port: 8080 + targetPort: 8080 + protocol: TCP + name: health + type: ClusterIP diff --git a/deployment/overlays/kind/README.md b/deployment/overlays/kind/README.md new file mode 100644 index 00000000..1015c17b --- /dev/null +++ b/deployment/overlays/kind/README.md @@ -0,0 +1,67 @@ +# Kind Cluster Deployment + +Deploy the full stack (UI + Agent + MCP Server + infrastructure) to a local Kubernetes cluster using [Kind](https://kind.sigs.k8s.io/). + +## Prerequisites + +- `kind` — [install](https://kind.sigs.k8s.io/docs/user/quick-start/#installation) +- `kubectl` +- `podman` or `docker` (for building images) + +## Quick Start + +```bash +make kind +``` + +This single command will: + +1. Clone `template-mcp-server` and `template-ui` repos into `.kind/` +2. Create a Kind cluster with ingress support +3. Build all three images (agent, MCP server, UI) and load them into Kind +4. Deploy the full stack via Kustomize +5. Wait for all pods to be ready + +## What's deployed + +| Service | Image | Port | Ingress | +|---------|-------|------|---------| +| UI | template-ui:local | 8080 | http://ui.localhost | +| Agent | template-agent:local | 5002 | http://agent.localhost | +| MCP Server | template-mcp-server:local | 5001 | http://mcp.localhost | +| Postgres (pgvector) | pgvector/pgvector:pg16 | 5432 | — | +| Redis | redis:7-alpine | 6379 | — | +| Jaeger | jaegertracing/all-in-one | 16686 | http://jaeger.localhost | + +## Useful Commands + +```bash +kubectl -n template-agent get pods +kubectl -n template-agent logs -l component=agent -f +kubectl -n template-agent logs -l component=mcp-server -f +kubectl -n template-agent logs -l component=ui -f +``` + +## Port-Forward (alternative to Ingress) + +```bash +kubectl -n template-agent port-forward svc/ui 8080:8080 +kubectl -n template-agent port-forward svc/agent 5002:5002 +kubectl -n template-agent port-forward svc/mcp-server 5001:5001 +``` + +## Differences from OpenShift + +| Concern | Kind | OpenShift | +|---------|------|-----------| +| Image build | Local `podman build` + `kind load` | BuildConfig (in-cluster) | +| Routing | NGINX Ingress | Route | +| Image pull | `imagePullPolicy: Never` | ImageStream | +| Security | Default PSA | SCC (restricted) | +| Storage | Default StorageClass | OpenShift PVs | + +## Teardown + +```bash +make kind-down +``` diff --git a/deployment/overlays/kind/deployment.yaml b/deployment/overlays/kind/deployment.yaml new file mode 100644 index 00000000..cc502d38 --- /dev/null +++ b/deployment/overlays/kind/deployment.yaml @@ -0,0 +1,119 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: agent + labels: + app: agent + component: agent +spec: + replicas: 1 + selector: + matchLabels: + app: agent + component: agent + template: + metadata: + labels: + app: agent + component: agent + spec: + containers: + - name: agent + image: agent:local + imagePullPolicy: Never + ports: + - containerPort: 5002 + name: http + protocol: TCP + env: + - name: AGENT_HOST + value: "0.0.0.0" + - name: AGENT_PORT + value: "5002" + - name: PYTHON_LOG_LEVEL + value: "INFO" + - name: POSTGRES_HOST + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_HOST + - name: POSTGRES_PORT + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_PORT + - name: POSTGRES_DB + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_DB + - name: POSTGRES_USER + valueFrom: + secretKeyRef: + name: agent-secrets + key: POSTGRES_USER + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: agent-secrets + key: POSTGRES_PASSWORD + - name: REDIS_URL + valueFrom: + configMapKeyRef: + name: agent-config + key: REDIS_URL + - name: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + valueFrom: + secretKeyRef: + name: agent-secrets + key: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + optional: true + - name: LANGFUSE_PUBLIC_KEY + valueFrom: + secretKeyRef: + name: agent-secrets + key: LANGFUSE_PUBLIC_KEY + optional: true + - name: LANGFUSE_SECRET_KEY + valueFrom: + secretKeyRef: + name: agent-secrets + key: LANGFUSE_SECRET_KEY + optional: true + - name: LANGFUSE_BASE_URL + valueFrom: + secretKeyRef: + name: agent-secrets + key: LANGFUSE_BASE_URL + optional: true + startupProbe: + httpGet: + path: /health + port: 5002 + initialDelaySeconds: 10 + periodSeconds: 5 + failureThreshold: 30 + livenessProbe: + httpGet: + path: /health + port: 5002 + initialDelaySeconds: 30 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + readinessProbe: + httpGet: + path: /health + port: 5002 + initialDelaySeconds: 10 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 3 + resources: + requests: + memory: "256Mi" + cpu: "100m" + limits: + memory: "512Mi" + cpu: "500m" + restartPolicy: Always diff --git a/deployment/overlays/kind/ingress.yaml b/deployment/overlays/kind/ingress.yaml new file mode 100644 index 00000000..2133dcb9 --- /dev/null +++ b/deployment/overlays/kind/ingress.yaml @@ -0,0 +1,42 @@ +apiVersion: networking.k8s.io/v1 +kind: Ingress +metadata: + name: template-agent-ingress + labels: + app: template-agent + annotations: + nginx.ingress.kubernetes.io/proxy-read-timeout: "3600" + nginx.ingress.kubernetes.io/proxy-send-timeout: "3600" + nginx.ingress.kubernetes.io/proxy-buffering: "off" +spec: + rules: + - host: ui.localhost + http: + paths: + - path: / + pathType: Prefix + backend: + service: + name: ui + port: + number: 8080 + - host: agent.localhost + http: + paths: + - path: / + pathType: Prefix + backend: + service: + name: agent + port: + number: 5002 + - host: jaeger.localhost + http: + paths: + - path: / + pathType: Prefix + backend: + service: + name: jaeger + port: + number: 16686 diff --git a/deployment/overlays/kind/kustomization.yaml b/deployment/overlays/kind/kustomization.yaml new file mode 100644 index 00000000..1052c16a --- /dev/null +++ b/deployment/overlays/kind/kustomization.yaml @@ -0,0 +1,28 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +namespace: template-agent + +resources: + - ../../base + - deployment.yaml + - service.yaml + - ui.yaml + - ingress.yaml + +# Include postgres and redis for local development +components: + - ../../components/postgres + - ../../components/redis + +labels: + - pairs: + app: template-agent + +images: + - name: agent + newName: localhost/template-agent + newTag: local + - name: template-ui + newName: localhost/template-ui + newTag: local diff --git a/deployment/openshift/service.yaml b/deployment/overlays/kind/service.yaml similarity index 60% rename from deployment/openshift/service.yaml rename to deployment/overlays/kind/service.yaml index b3477a3f..5d29b2b4 100644 --- a/deployment/openshift/service.yaml +++ b/deployment/overlays/kind/service.yaml @@ -1,17 +1,17 @@ apiVersion: v1 kind: Service metadata: - name: template-agent + name: agent labels: - app: template-agent + app: agent component: agent spec: - type: ClusterIP + selector: + app: agent + component: agent ports: - - port: 8081 - targetPort: 8081 + - port: 5002 + targetPort: 5002 protocol: TCP name: http - selector: - app: template-agent - component: agent + type: ClusterIP diff --git a/deployment/overlays/kind/ui.yaml b/deployment/overlays/kind/ui.yaml new file mode 100644 index 00000000..8e28c826 --- /dev/null +++ b/deployment/overlays/kind/ui.yaml @@ -0,0 +1,133 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: ui + labels: + app: template-agent + component: ui +spec: + replicas: 1 + selector: + matchLabels: + component: ui + template: + metadata: + labels: + app: template-agent + component: ui + spec: + containers: + - name: ui + image: localhost/template-ui:local + imagePullPolicy: Never + ports: + - containerPort: 8080 + name: http + protocol: TCP + env: + - name: PORT + value: "8080" + - name: ENVIRONMENT + value: "development" + - name: AUTH_ENABLED + value: "false" + - name: AGENT_HOST + value: "http://agent:5002" + - name: UI_CONFIG_PATH + value: "/etc/config/ui.yaml" + - name: REDIS_HOST + valueFrom: + configMapKeyRef: + name: agent-config + key: REDIS_HOST + - name: REDIS_PORT + valueFrom: + configMapKeyRef: + name: agent-config + key: REDIS_PORT + volumeMounts: + - name: ui-config + mountPath: /etc/config + readOnly: true + readinessProbe: + httpGet: + path: /api/health + port: 8080 + initialDelaySeconds: 10 + periodSeconds: 5 + resources: + requests: + memory: "128Mi" + cpu: "100m" + limits: + memory: "256Mi" + cpu: "250m" + volumes: + - name: ui-config + configMap: + name: ui-config +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: ui-config + labels: + app: template-agent + component: ui +data: + ui.yaml: | + server: + host: "0.0.0.0" + port: 8080 + body_limit: 1048576 + logging: + level: info + cors: + origin: "http://localhost:5173" + security: + helmet: + enabled: true + csp: + default_src: ["'self'"] + script_src: ["'self'", "'unsafe-inline'"] + style_src: ["'self'", "'unsafe-inline'"] + img_src: ["'self'", "data:", "blob:"] + connect_src: ["'self'"] + font_src: ["'self'"] + object_src: ["'none'"] + frame_ancestors: ["'none'"] + cross_origin_embedder_policy: false + rate_limit: + enabled: true + max: 100 + window: "1 minute" + exclude_paths: + - "/api/health" + - "/_health" + session: + secure_cookie: false + max_age_days: 30 + otel: + enabled: false + service_name: "template-ui" + announcement: + enabled: false + message: "" + type: info +--- +apiVersion: v1 +kind: Service +metadata: + name: ui + labels: + app: template-agent + component: ui +spec: + selector: + component: ui + ports: + - port: 8080 + targetPort: 8080 + protocol: TCP + name: http + type: ClusterIP diff --git a/deployment/overlays/openshift-headless/configmap-patch.yaml b/deployment/overlays/openshift-headless/configmap-patch.yaml new file mode 100644 index 00000000..7ad6e838 --- /dev/null +++ b/deployment/overlays/openshift-headless/configmap-patch.yaml @@ -0,0 +1,12 @@ +apiVersion: v1 +kind: ConfigMap +metadata: + name: agent-config +data: + ENVIRONMENT: "production" + PYTHON_LOG_LEVEL: "INFO" + LANGFUSE_TRACING_ENVIRONMENT: "production" + POSTGRES_HOST: "postgres" + POSTGRES_PORT: "5432" + POSTGRES_DB: "template_agent" + REDIS_URL: "redis://redis:6379/0" diff --git a/deployment/openshift/deployment.yaml b/deployment/overlays/openshift-headless/deployment.yaml similarity index 57% rename from deployment/openshift/deployment.yaml rename to deployment/overlays/openshift-headless/deployment.yaml index 596803c3..134f811f 100644 --- a/deployment/openshift/deployment.yaml +++ b/deployment/overlays/openshift-headless/deployment.yaml @@ -1,146 +1,155 @@ apiVersion: apps/v1 kind: Deployment metadata: - name: template-agent + name: agent-headless labels: - app: template-agent - component: agent + app: agent + component: agent-headless spec: - replicas: 1 + replicas: 2 + strategy: + type: RollingUpdate + rollingUpdate: + maxSurge: 1 + maxUnavailable: 0 selector: matchLabels: - app: template-agent - component: agent + app: agent + component: agent-headless template: metadata: labels: - app: template-agent - component: agent + app: agent + component: agent-headless spec: + serviceAccountName: agent + terminationGracePeriodSeconds: 60 containers: - - name: template-agent - image: template-agent:latest + - name: agent-headless + image: agent:latest imagePullPolicy: Always + command: ["python", "-m", "deep_agent.headless"] ports: - - containerPort: 8081 - name: http + - containerPort: 8080 + name: health protocol: TCP env: - - name: AGENT_HOST - value: "0.0.0.0" - - name: AGENT_PORT - value: "8081" - - name: PYTHON_LOG_LEVEL - valueFrom: - configMapKeyRef: - name: template-agent-config - key: PYTHON_LOG_LEVEL - - name: USE_INMEMORY_SAVER - valueFrom: - configMapKeyRef: - name: template-agent-config - key: USE_INMEMORY_SAVER - - name: LANGFUSE_TRACING_ENVIRONMENT + - name: ENVIRONMENT valueFrom: configMapKeyRef: - name: template-agent-config - key: LANGFUSE_TRACING_ENVIRONMENT - - name: MCP_SERVER_NAME - valueFrom: - configMapKeyRef: - name: template-agent-config - key: MCP_SERVER_NAME - - name: MCP_SERVER_URL - valueFrom: - configMapKeyRef: - name: template-agent-config - key: MCP_SERVER_URL - - name: MCP_TRANSPORT_PROTOCOL + name: agent-config + key: ENVIRONMENT + optional: true + - name: PYTHON_LOG_LEVEL valueFrom: configMapKeyRef: - name: template-agent-config - key: MCP_TRANSPORT_PROTOCOL + name: agent-config + key: PYTHON_LOG_LEVEL - name: POSTGRES_HOST valueFrom: - secretKeyRef: - name: template-agent-secrets + configMapKeyRef: + name: agent-config key: POSTGRES_HOST - optional: true - name: POSTGRES_PORT valueFrom: - secretKeyRef: - name: template-agent-secrets + configMapKeyRef: + name: agent-config key: POSTGRES_PORT - name: POSTGRES_DB - valueFrom: - secretKeyRef: - name: template-agent-secrets - key: POSTGRES_DB - optional: true - - name: SSO_CALLBACK_URL valueFrom: configMapKeyRef: - name: template-agent-config - key: SSO_CALLBACK_URL + name: agent-config + key: POSTGRES_DB - name: POSTGRES_USER valueFrom: secretKeyRef: - name: template-agent-secrets + name: agent-secrets key: POSTGRES_USER - name: POSTGRES_PASSWORD valueFrom: secretKeyRef: - name: template-agent-secrets + name: agent-secrets key: POSTGRES_PASSWORD + - name: REDIS_URL + valueFrom: + configMapKeyRef: + name: agent-config + key: REDIS_URL + - name: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + valueFrom: + secretKeyRef: + name: agent-secrets + key: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + optional: true + - name: VLLM_BASE_URL + valueFrom: + secretKeyRef: + name: agent-secrets + key: VLLM_BASE_URL + optional: true + - name: VLLM_API_KEY + valueFrom: + secretKeyRef: + name: agent-secrets + key: VLLM_API_KEY + optional: true - name: LANGFUSE_PUBLIC_KEY valueFrom: secretKeyRef: - name: template-agent-secrets + name: agent-secrets key: LANGFUSE_PUBLIC_KEY optional: true - name: LANGFUSE_SECRET_KEY valueFrom: secretKeyRef: - name: template-agent-secrets + name: agent-secrets key: LANGFUSE_SECRET_KEY optional: true - name: LANGFUSE_BASE_URL valueFrom: secretKeyRef: - name: template-agent-secrets + name: agent-secrets key: LANGFUSE_BASE_URL optional: true - - name: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + - name: LANGFUSE_TRACING_ENVIRONMENT valueFrom: - secretKeyRef: - name: template-agent-secrets - key: GOOGLE_APPLICATION_CREDENTIALS_CONTENT - optional: true - envFrom: - - secretRef: - name: template-agent-secrets - optional: true + configMapKeyRef: + name: agent-config + key: LANGFUSE_TRACING_ENVIRONMENT + securityContext: + runAsNonRoot: true + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + startupProbe: + httpGet: + path: /healthz + port: 8080 + initialDelaySeconds: 10 + periodSeconds: 5 + failureThreshold: 12 livenessProbe: httpGet: - path: /health - port: 8081 + path: /healthz + port: 8080 initialDelaySeconds: 30 periodSeconds: 10 timeoutSeconds: 5 failureThreshold: 3 readinessProbe: httpGet: - path: /health - port: 8081 - initialDelaySeconds: 10 + path: /readyz + port: 8080 + initialDelaySeconds: 15 periodSeconds: 5 timeoutSeconds: 3 failureThreshold: 3 resources: requests: - memory: "256Mi" - cpu: "100m" - limits: memory: "512Mi" - cpu: "500m" + cpu: "250m" + limits: + memory: "1Gi" + cpu: "1000m" restartPolicy: Always diff --git a/deployment/overlays/openshift-headless/hpa.yaml b/deployment/overlays/openshift-headless/hpa.yaml new file mode 100644 index 00000000..176a19f2 --- /dev/null +++ b/deployment/overlays/openshift-headless/hpa.yaml @@ -0,0 +1,42 @@ +apiVersion: autoscaling/v2 +kind: HorizontalPodAutoscaler +metadata: + name: agent-headless + labels: + app: agent + component: agent-headless +spec: + scaleTargetRef: + apiVersion: apps/v1 + kind: Deployment + name: agent-headless + minReplicas: 2 + maxReplicas: 8 + metrics: + - type: Resource + resource: + name: cpu + target: + type: Utilization + averageUtilization: 70 + - type: Resource + resource: + name: memory + target: + type: Utilization + averageUtilization: 75 + behavior: + scaleUp: + stabilizationWindowSeconds: 60 + policies: + - type: Pods + value: 2 + periodSeconds: 60 + selectPolicy: Max + scaleDown: + stabilizationWindowSeconds: 300 + policies: + - type: Pods + value: 1 + periodSeconds: 180 + selectPolicy: Min diff --git a/deployment/overlays/openshift-headless/kustomization.yaml b/deployment/overlays/openshift-headless/kustomization.yaml new file mode 100644 index 00000000..a445ae3e --- /dev/null +++ b/deployment/overlays/openshift-headless/kustomization.yaml @@ -0,0 +1,26 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +namespace: NAMESPACE_PLACEHOLDER + +resources: + - ../../base + - deployment.yaml + - service.yaml + - hpa.yaml + - pdb.yaml + - scaledobject.yaml + +labels: + - pairs: + app: agent + component: agent-headless + +images: + - name: agent + newTag: latest + +patches: + - path: configmap-patch.yaml + - path: secret-patch.yaml + - path: redis-patch.yaml diff --git a/deployment/overlays/openshift-headless/pdb.yaml b/deployment/overlays/openshift-headless/pdb.yaml new file mode 100644 index 00000000..b08d0646 --- /dev/null +++ b/deployment/overlays/openshift-headless/pdb.yaml @@ -0,0 +1,13 @@ +apiVersion: policy/v1 +kind: PodDisruptionBudget +metadata: + name: agent-headless + labels: + app: agent + component: agent-headless +spec: + minAvailable: 1 + selector: + matchLabels: + app: agent + component: agent-headless diff --git a/deployment/overlays/openshift-headless/redis-patch.yaml b/deployment/overlays/openshift-headless/redis-patch.yaml new file mode 100644 index 00000000..f46dd078 --- /dev/null +++ b/deployment/overlays/openshift-headless/redis-patch.yaml @@ -0,0 +1,9 @@ +apiVersion: v1 +kind: Service +metadata: + name: redis +$patch: merge +spec: + selector: + app: agent + component: cache diff --git a/deployment/overlays/openshift-headless/scaledobject.yaml b/deployment/overlays/openshift-headless/scaledobject.yaml new file mode 100644 index 00000000..51b90f68 --- /dev/null +++ b/deployment/overlays/openshift-headless/scaledobject.yaml @@ -0,0 +1,21 @@ +apiVersion: keda.sh/v1alpha1 +kind: ScaledObject +metadata: + name: agent-headless + labels: + app: agent + component: agent-headless +spec: + scaleTargetRef: + name: agent-headless + minReplicaCount: 1 + maxReplicaCount: 10 + cooldownPeriod: 300 + pollingInterval: 15 + triggers: + - type: redis-streams + metadata: + addressFromEnv: REDIS_URL + stream: agent-tasks + consumerGroup: agent-workers + pendingEntriesCount: "10" diff --git a/deployment/overlays/openshift-headless/secret-patch.yaml b/deployment/overlays/openshift-headless/secret-patch.yaml new file mode 100644 index 00000000..7b6bade7 --- /dev/null +++ b/deployment/overlays/openshift-headless/secret-patch.yaml @@ -0,0 +1,6 @@ +apiVersion: v1 +kind: Secret +metadata: + name: agent-secrets +type: Opaque +stringData: {} diff --git a/deployment/overlays/openshift-headless/service.yaml b/deployment/overlays/openshift-headless/service.yaml new file mode 100644 index 00000000..0bcf91cb --- /dev/null +++ b/deployment/overlays/openshift-headless/service.yaml @@ -0,0 +1,17 @@ +apiVersion: v1 +kind: Service +metadata: + name: agent-headless + labels: + app: agent + component: agent-headless +spec: + type: ClusterIP + ports: + - port: 8080 + targetPort: 8080 + protocol: TCP + name: health + selector: + app: agent + component: agent-headless diff --git a/deployment/openshift/buildconfig.yaml b/deployment/overlays/openshift/buildconfig.yaml similarity index 85% rename from deployment/openshift/buildconfig.yaml rename to deployment/overlays/openshift/buildconfig.yaml index 1f0c3ca1..b6ff9f63 100644 --- a/deployment/openshift/buildconfig.yaml +++ b/deployment/overlays/openshift/buildconfig.yaml @@ -1,9 +1,9 @@ apiVersion: build.openshift.io/v1 kind: BuildConfig metadata: - name: template-agent + name: agent labels: - app: template-agent + app: agent component: agent spec: successfulBuildsHistoryLimit: 1 @@ -11,7 +11,7 @@ spec: output: to: kind: ImageStreamTag - name: template-agent:latest + name: agent:latest source: type: Binary binary: {} diff --git a/deployment/overlays/openshift/configmap-patch.yaml b/deployment/overlays/openshift/configmap-patch.yaml new file mode 100644 index 00000000..de546ffa --- /dev/null +++ b/deployment/overlays/openshift/configmap-patch.yaml @@ -0,0 +1,26 @@ +apiVersion: v1 +kind: ConfigMap +metadata: + name: agent-config +data: + # OpenShift-specific config + AGENT_HOST: "0.0.0.0" + AGENT_PORT: "5002" + PYTHON_LOG_LEVEL: "INFO" + + # Environment & Security + ENVIRONMENT: "production" + ENABLE_AUTH: "true" + REQUEST_BODY_MAX_SIZE: "10485760" # 10MB + + # Request Logging + REQUEST_LOGGING_ENABLED: "true" + REQUEST_LOG_HEADERS: "true" + REQUEST_LOG_BODY: "false" + REQUEST_LOG_BODY_MAX_SIZE: "10240" + + # Observability + LANGFUSE_TRACING_ENVIRONMENT: "production" + + # Redis + REDIS_URL: "redis://redis:6379/0" diff --git a/deployment/overlays/openshift/deployment.yaml b/deployment/overlays/openshift/deployment.yaml new file mode 100644 index 00000000..b7505684 --- /dev/null +++ b/deployment/overlays/openshift/deployment.yaml @@ -0,0 +1,213 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: agent + labels: + app: agent + component: agent +spec: + replicas: 2 + strategy: + type: RollingUpdate + rollingUpdate: + maxSurge: 1 + maxUnavailable: 0 + selector: + matchLabels: + app: agent + component: agent + template: + metadata: + labels: + app: agent + component: agent + spec: + serviceAccountName: agent + terminationGracePeriodSeconds: 60 + containers: + - name: agent + image: agent:latest + imagePullPolicy: Always + ports: + - containerPort: 5002 + name: http + protocol: TCP + env: + - name: AGENT_HOST + valueFrom: + configMapKeyRef: + name: agent-config + key: AGENT_HOST + - name: AGENT_PORT + valueFrom: + configMapKeyRef: + name: agent-config + key: AGENT_PORT + - name: ENABLE_AUTH + valueFrom: + configMapKeyRef: + name: agent-config + key: ENABLE_AUTH + - name: PYTHON_LOG_LEVEL + valueFrom: + configMapKeyRef: + name: agent-config + key: PYTHON_LOG_LEVEL + - name: LANGFUSE_TRACING_ENVIRONMENT + valueFrom: + configMapKeyRef: + name: agent-config + key: LANGFUSE_TRACING_ENVIRONMENT + - name: POSTGRES_HOST + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_HOST + - name: POSTGRES_PORT + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_PORT + - name: POSTGRES_DB + valueFrom: + configMapKeyRef: + name: agent-config + key: POSTGRES_DB + - name: POSTGRES_USER + valueFrom: + secretKeyRef: + name: agent-secrets + key: POSTGRES_USER + - name: POSTGRES_PASSWORD + valueFrom: + secretKeyRef: + name: agent-secrets + key: POSTGRES_PASSWORD + - name: LANGFUSE_PUBLIC_KEY + valueFrom: + secretKeyRef: + name: agent-secrets + key: LANGFUSE_PUBLIC_KEY + optional: true + - name: LANGFUSE_SECRET_KEY + valueFrom: + secretKeyRef: + name: agent-secrets + key: LANGFUSE_SECRET_KEY + optional: true + - name: LANGFUSE_BASE_URL + valueFrom: + secretKeyRef: + name: agent-secrets + key: LANGFUSE_BASE_URL + optional: true + - name: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + valueFrom: + secretKeyRef: + name: agent-secrets + key: GOOGLE_APPLICATION_CREDENTIALS_CONTENT + optional: true + - name: SSO_ISSUER_URL + valueFrom: + secretKeyRef: + name: agent-secrets + key: SSO_ISSUER_URL + optional: true + - name: SSO_CLIENT_ID + valueFrom: + secretKeyRef: + name: agent-secrets + key: SSO_CLIENT_ID + optional: true + - name: SSO_CLIENT_SECRET + valueFrom: + secretKeyRef: + name: agent-secrets + key: SSO_CLIENT_SECRET + optional: true + - name: VLLM_BASE_URL + valueFrom: + secretKeyRef: + name: agent-secrets + key: VLLM_BASE_URL + optional: true + - name: VLLM_API_KEY + valueFrom: + secretKeyRef: + name: agent-secrets + key: VLLM_API_KEY + optional: true + - name: REDIS_URL + valueFrom: + configMapKeyRef: + name: agent-config + key: REDIS_URL + - name: SSL_KEYFILE + valueFrom: + secretKeyRef: + name: agent-secrets + key: SSL_KEYFILE + optional: true + - name: SSL_CERTFILE + valueFrom: + secretKeyRef: + name: agent-secrets + key: SSL_CERTFILE + optional: true + - name: REQUEST_LOGGING_ENABLED + valueFrom: + configMapKeyRef: + name: agent-config + key: REQUEST_LOGGING_ENABLED + - name: REQUEST_LOG_HEADERS + valueFrom: + configMapKeyRef: + name: agent-config + key: REQUEST_LOG_HEADERS + - name: REQUEST_LOG_BODY + valueFrom: + configMapKeyRef: + name: agent-config + key: REQUEST_LOG_BODY + - name: REQUEST_LOG_BODY_MAX_SIZE + valueFrom: + configMapKeyRef: + name: agent-config + key: REQUEST_LOG_BODY_MAX_SIZE + securityContext: + runAsNonRoot: true + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + startupProbe: + httpGet: + path: /health + port: 5002 + initialDelaySeconds: 5 + periodSeconds: 5 + failureThreshold: 12 + livenessProbe: + httpGet: + path: /health + port: 5002 + initialDelaySeconds: 30 + periodSeconds: 10 + timeoutSeconds: 5 + failureThreshold: 3 + readinessProbe: + httpGet: + path: /health + port: 5002 + initialDelaySeconds: 10 + periodSeconds: 5 + timeoutSeconds: 3 + failureThreshold: 3 + resources: + requests: + memory: "512Mi" + cpu: "250m" + limits: + memory: "1Gi" + cpu: "1000m" + restartPolicy: Always diff --git a/deployment/overlays/openshift/hpa.yaml b/deployment/overlays/openshift/hpa.yaml new file mode 100644 index 00000000..eafc1282 --- /dev/null +++ b/deployment/overlays/openshift/hpa.yaml @@ -0,0 +1,48 @@ +apiVersion: autoscaling/v2 +kind: HorizontalPodAutoscaler +metadata: + name: agent + labels: + app: agent + component: agent +spec: + scaleTargetRef: + apiVersion: apps/v1 + kind: Deployment + name: agent + minReplicas: 2 # HA baseline - handles ~20-40 concurrent users + maxReplicas: 10 # Peak capacity - handles ~100-150 concurrent users + metrics: + - type: Resource + resource: + name: cpu + target: + type: Utilization + averageUtilization: 70 # Scale up when CPU > 70% + - type: Resource + resource: + name: memory + target: + type: Utilization + averageUtilization: 75 # Scale up when memory > 75% + behavior: + scaleUp: + stabilizationWindowSeconds: 60 # Wait 60s before scaling up + policies: + - type: Percent + value: 50 # Scale up by 50% (1→2, 2→3, 4→6) + periodSeconds: 60 + - type: Pods + value: 2 # Or add 2 pods at once + periodSeconds: 60 + selectPolicy: Max # Use whichever policy scales faster + scaleDown: + stabilizationWindowSeconds: 300 # Wait 5min before scaling down + policies: + - type: Percent + value: 25 # Scale down by 25% at a time + periodSeconds: 60 + - type: Pods + value: 1 # Or remove 1 pod at once + periodSeconds: 180 + selectPolicy: Min # Use whichever policy scales slower (conservative) diff --git a/deployment/openshift/imagestream.yaml b/deployment/overlays/openshift/imagestream.yaml similarity index 73% rename from deployment/openshift/imagestream.yaml rename to deployment/overlays/openshift/imagestream.yaml index aed597e0..5d9ec3d8 100644 --- a/deployment/openshift/imagestream.yaml +++ b/deployment/overlays/openshift/imagestream.yaml @@ -1,9 +1,9 @@ apiVersion: image.openshift.io/v1 kind: ImageStream metadata: - name: template-agent + name: agent labels: - app: template-agent + app: agent component: agent spec: lookupPolicy: diff --git a/deployment/overlays/openshift/kustomization.yaml b/deployment/overlays/openshift/kustomization.yaml new file mode 100644 index 00000000..ebe18b81 --- /dev/null +++ b/deployment/overlays/openshift/kustomization.yaml @@ -0,0 +1,117 @@ +apiVersion: kustomize.config.k8s.io/v1beta1 +kind: Kustomization + +namespace: NAMESPACE_PLACEHOLDER + +resources: + - ../../base + - buildconfig.yaml + - imagestream.yaml + - deployment.yaml + - service.yaml + - route.yaml + - pdb.yaml + - hpa.yaml + +labels: + - pairs: + app: agent + component: agent + +images: + - name: agent + newTag: latest + +patches: + - path: configmap-patch.yaml + - path: secret-patch.yaml + + # Agent deployment patches + - target: + kind: Deployment + name: agent + patch: |- + - op: replace + path: /spec/template/spec/containers/0/resources/requests/memory + value: "512Mi" + - op: replace + path: /spec/template/spec/containers/0/resources/requests/cpu + value: "250m" + - op: replace + path: /spec/template/spec/containers/0/resources/limits/memory + value: "1Gi" + - op: replace + path: /spec/template/spec/containers/0/resources/limits/cpu + value: "1000m" + + # PostgreSQL deployment patches (from base) + - target: + kind: Deployment + name: postgres + patch: |- + - op: replace + path: /spec/template/spec/containers/0/resources/requests/memory + value: "512Mi" + - op: replace + path: /spec/template/spec/containers/0/resources/requests/cpu + value: "200m" + - op: replace + path: /spec/template/spec/containers/0/resources/limits/memory + value: "1Gi" + - op: replace + path: /spec/template/spec/containers/0/resources/limits/cpu + value: "1000m" + - op: add + path: /spec/template/spec/containers/0/args + value: ["-c", "max_connections=200"] + + # PostgreSQL PVC size (OpenShift gets more storage) + - target: + kind: PersistentVolumeClaim + name: postgres-pvc + patch: |- + - op: replace + path: /spec/resources/requests/storage + value: "10Gi" + + # Redis deployment patches (from base) + - target: + kind: Deployment + name: redis + patch: |- + - op: replace + path: /spec/template/spec/containers/0/resources/requests/memory + value: "256Mi" + - op: replace + path: /spec/template/spec/containers/0/resources/limits/memory + value: "512Mi" + - op: replace + path: /spec/template/spec/containers/0/resources/limits/cpu + value: "500m" + + # Redis PVC size (OpenShift gets more storage) + - target: + kind: PersistentVolumeClaim + name: redis-pvc + patch: |- + - op: replace + path: /spec/resources/requests/storage + value: "2Gi" + + # BuildConfig patches + - target: + kind: BuildConfig + name: agent + patch: |- + - op: replace + path: /spec/resources/requests/memory + value: "2Gi" + - op: replace + path: /spec/resources/requests/cpu + value: "1000m" + - op: replace + path: /spec/resources/limits/memory + value: "4Gi" + - op: replace + path: /spec/resources/limits/cpu + value: "4" diff --git a/deployment/overlays/openshift/pdb.yaml b/deployment/overlays/openshift/pdb.yaml new file mode 100644 index 00000000..3e68bab1 --- /dev/null +++ b/deployment/overlays/openshift/pdb.yaml @@ -0,0 +1,13 @@ +apiVersion: policy/v1 +kind: PodDisruptionBudget +metadata: + name: agent + labels: + app: agent + component: agent +spec: + minAvailable: 1 + selector: + matchLabels: + app: agent + component: agent diff --git a/deployment/overlays/openshift/redis-patch.yaml b/deployment/overlays/openshift/redis-patch.yaml new file mode 100644 index 00000000..f46dd078 --- /dev/null +++ b/deployment/overlays/openshift/redis-patch.yaml @@ -0,0 +1,9 @@ +apiVersion: v1 +kind: Service +metadata: + name: redis +$patch: merge +spec: + selector: + app: agent + component: cache diff --git a/deployment/overlays/openshift/route.yaml b/deployment/overlays/openshift/route.yaml new file mode 100644 index 00000000..4f4480ff --- /dev/null +++ b/deployment/overlays/openshift/route.yaml @@ -0,0 +1,24 @@ +apiVersion: route.openshift.io/v1 +kind: Route +metadata: + name: agent + labels: + app: agent + component: agent + shard: internal + annotations: + haproxy.router.openshift.io/timeout: 18000s + haproxy.router.openshift.io/balance: roundrobin + haproxy.router.openshift.io/rate-limit-connections: "true" + haproxy.router.openshift.io/rate-limit-connections.concurrent-tcp: "100" + haproxy.router.openshift.io/rate-limit-connections.rate-http: "1000" + haproxy.router.openshift.io/rate-limit-connections.rate-tcp: "1000" +spec: + to: + kind: Service + name: agent + port: + targetPort: http + tls: + termination: edge + insecureEdgeTerminationPolicy: Redirect diff --git a/deployment/overlays/openshift/secret-patch.yaml b/deployment/overlays/openshift/secret-patch.yaml new file mode 100644 index 00000000..6bed01a3 --- /dev/null +++ b/deployment/overlays/openshift/secret-patch.yaml @@ -0,0 +1,8 @@ +apiVersion: v1 +kind: Secret +metadata: + name: agent-secrets +type: Opaque +stringData: + SSL_KEYFILE: "" + SSL_CERTFILE: "" diff --git a/deployment/overlays/openshift/service.yaml b/deployment/overlays/openshift/service.yaml new file mode 100644 index 00000000..4bd51824 --- /dev/null +++ b/deployment/overlays/openshift/service.yaml @@ -0,0 +1,17 @@ +apiVersion: v1 +kind: Service +metadata: + name: agent + labels: + app: agent + component: agent +spec: + type: ClusterIP + ports: + - port: 5002 + targetPort: 5002 + protocol: TCP + name: http + selector: + app: agent + component: agent diff --git a/docs/headless-agent-architecture.html b/docs/headless-agent-architecture.html new file mode 100644 index 00000000..ca4a26dd --- /dev/null +++ b/docs/headless-agent-architecture.html @@ -0,0 +1,1430 @@ + + + + + + Headless Agent — Architecture Guide + + + + + + + + + + +
+ +
+

Template Agent — Headless Mode

+

Architecture guide for the event-driven headless agent mode. Covers dual-mode operation (server + headless), trigger sources, output sinks, task status tracking, and production deployment.

+
+ +
+
+ + +
+

Dual-Mode Architecture

+

The template agent runs in two independent modes from the same codebase and container image. Both share Redis, PostgreSQL, and the same model provider.

+ +
+ 
+graph TB
+    subgraph "Same Container Image"
+        direction TB
+        S["Server Mode
make local / aegra dev
Port 5002"]
+        H["Headless Mode
make headless
python -m deep_agent.headless"]
+    end
+
+    U["👤 User / UI"] -->|HTTP / SSE| S
+    S -->|queue_task| R["Redis Streams"]
+    R -->|consume| H
+
+    EXT["🌐 External System"] -->|POST /trigger| H
+    CRON["⏰ Cron Schedule"] -->|fires| H
+
+    S --> PG["PostgreSQL"]
+    H --> PG
+    S --> RD["Redis"]
+    H --> RD
+
+    H -->|results| SINKS["Output Sinks"]
+
+    style S fill:#1f6feb,stroke:#58a6ff,color:#fff
+    style H fill:#238636,stroke:#3fb950,color:#fff
+    style R fill:#d29922,stroke:#e3b341,color:#000
+    style SINKS fill:#8957e5,stroke:#a371f7,color:#fff
+
+
+ +
+
+

Server Mode

+
    +
  • Full Aegra HTTP API
    • +
  • Interactive conversations via UI
    • +
  • SSE streaming responses
    • +
  • SSO authentication
    • +
  • Orchestrator prompt (PROMPT.md)
    • +
  • make local to start
    • +
+
+
+

Headless Mode

+
    +
  • No HTTP API (except health check)
    • +
  • Event-driven background worker
    • +
  • Webhook, cron, queue triggers
    • +
  • No auth (service account identity)
    • +
  • Worker prompt (HEADLESS_PROMPT.md)
    • +
  • make headless to start
    • +
+
+
+
+ + +
+

When to Use Which Mode

+ + + + + + + + + + +
Use CaseModeWhy
Interactive chat assistantserverUser expects immediate responses
Background batch processingheadlessLong-running, no user waiting
Chat + async delegationbothServer queues work, headless processes
Event-driven pipelineheadlessExternal triggers, no UI needed
Scheduled reportsheadlessCron trigger fires on schedule
Webhook responderheadlessExternal system POSTs, agent processes
+
+ + +
+

C4 Context Diagram

+

System context showing how the template agent interacts with external actors and systems.

+ +
+ 
+graph TB
+    USER["👤 Employee
Uses the fitness assistant
via web UI"]
+    ADMIN["🔧 Platform Team
Deploys and monitors
the agent system"]
+    EXT["🌐 External System
Sends webhooks or
pushes to queue"]
+
+    subgraph SYSTEM["Template Agent System"]
+        direction LR
+        SA["Server Agent
Interactive conversations"]
+        HA["Headless Agent
Background processing"]
+    end
+
+    MCP["MCP Server
Tools: calculate_bmi,
send_email, etc."]
+    LLM["Google Gemini
LLM inference"]
+    SSO["SSO Provider
OIDC authentication"]
+
+    USER -->|"chat via UI"| SA
+    ADMIN -->|"deploy, monitor"| SYSTEM
+    EXT -->|"webhook / queue"| HA
+    SA -->|"queue_task()"| HA
+    SA --> MCP
+    HA --> MCP
+    SA --> LLM
+    HA --> LLM
+    SA --> SSO
+
+    style SA fill:#1f6feb,stroke:#58a6ff,color:#fff
+    style HA fill:#238636,stroke:#3fb950,color:#fff
+    style SYSTEM fill:#1a2332,stroke:#2d3f56,color:#e6edf3
+
+
+
+ + +
+

C4 Container Diagram

+

Internal containers showing how server and headless agents are structured.

+ +
+ 
+graph TB
+    subgraph SERVER["Server Agent (Aegra)"]
+        direction TB
+        GF["Graph Factory
graph.py:agent()"]
+        ORCH["Orchestrator
PROMPT.md"]
+        SUB["Subagents
analyst, publisher"]
+        TOOLS_S["Tools
validate_email,
queue_task,
check_task_status,
get_pending_results"]
+    end
+
+    subgraph HEADLESS["Headless Agent"]
+        direction TB
+        ETM["EventTriggerMiddleware
Lifecycle orchestrator"]
+        HG["Headless Graph
HEADLESS_PROMPT.md"]
+
+        subgraph TRIGGERS["Trigger Sources"]
+            WH["Webhook
HTTP listener"]
+            CR["Cron
APScheduler"]
+            QU["Queue
Redis Streams"]
+        end
+
+        subgraph SINKS["Output Sinks"]
+            SO["Stdout"]
+            FI["File JSONL"]
+            WS["Webhook POST"]
+            RS["Redis Stream"]
+        end
+    end
+
+    REDIS["Redis
Streams + Task Store"]
+    PG["PostgreSQL
Conversations +
Checkpoints"]
+    GEMINI["Gemini 2.5 Pro"]
+    HC["Health Check
/healthz /readyz
Port 8080"]
+
+    GF --> ORCH --> SUB
+    TOOLS_S --> REDIS
+    ORCH --> TOOLS_S
+    ETM --> HG
+    WH --> ETM
+    CR --> ETM
+    QU --> ETM
+    ETM --> SO
+    ETM --> FI
+    ETM --> WS
+    ETM --> RS
+    ETM --> REDIS
+    HEADLESS --> HC
+
+    SERVER --> PG
+    HEADLESS --> PG
+    SERVER --> GEMINI
+    HEADLESS --> GEMINI
+    QU --> REDIS
+    RS --> REDIS
+
+    style SERVER fill:#0d1b2a,stroke:#1f6feb,color:#e6edf3
+    style HEADLESS fill:#0d2818,stroke:#238636,color:#e6edf3
+    style TRIGGERS fill:#1a2332,stroke:#58a6ff,color:#e6edf3
+    style SINKS fill:#1a2332,stroke:#a371f7,color:#e6edf3
+
+
+
+ + +
+

Component Map

+ +
+

File Structure

+ 
+deep_agent/
+├── headless.py                     # Entry point: python -m deep_agent.headless
+├── aegra/
+│   ├── graph.py                    # Graph factory (server + built-in tools)
+│   └── startup.py                  # Startup orchestrator
+└── src/triggers/
+    ├── config.py                   # Pydantic models (HeadlessConfig, TriggerConfig, etc.)
+    ├── middleware.py                # EventTriggerMiddleware (lifecycle orchestrator)
+    ├── runtime.py                  # HeadlessRuntime (ServerRuntime adapter)
+    ├── task_store.py               # Redis-backed task status tracking
+    ├── tools.py                    # queue_task, check_task_status, get_pending_results
+    ├── health.py                   # /healthz and /readyz endpoints
+    ├── sources/
+    │   ├── protocol.py             # TriggerSource protocol + TriggerEvent dataclass
+    │   ├── webhook.py              # HTTP listener trigger
+    │   ├── cron.py                 # APScheduler cron trigger
+    │   └── queue.py                # QueueConsumer protocol + RedisStreamsConsumer
+    └── sinks/
+        ├── protocol.py             # OutputSink protocol + TriggerResult dataclass
+        ├── stdout.py               # Print to stdout
+        ├── file.py                 # Append JSONL to file
+        ├── webhook.py              # POST to URL with retry
+        └── redis.py                # XADD to Redis Stream
+
+config/agent/
+├── PROMPT.md                       # Server orchestrator prompt
+├── HEADLESS_PROMPT.md              # Headless worker prompt
+└── runtime/agent.yaml              # Unified config (triggers, sinks, health_check)
+
+
+ + +
+

Server Mode Flow

+

User interacts via UI → Aegra API → Graph factory → Orchestrator → Subagents.

+ +
+ 
+sequenceDiagram
+    participant U as User (UI)
+    participant API as Aegra API
Port 5002
+    participant GF as Graph Factory
+    participant O as Orchestrator
+    participant A as Analyst Subagent
+    participant MCP as MCP Server
+
+    U->>API: POST /threads/{id}/runs
+    API->>GF: agent(runtime)
+    GF->>O: compiled graph
+    O->>O: Create TODO list
+    O->>A: Delegate BMI analysis
+    A->>MCP: calculate_bmi(height, weight)
+    MCP-->>A: BMI result
+    A-->>O: Analysis report
+    O-->>API: Response
+    API-->>U: SSE stream
+
+
+
+ + +
+

Headless Mode Flow

+

Events arrive via triggers → EventTriggerMiddleware invokes the graph → results fan out to sinks.

+ +
+ 
+sequenceDiagram
+    participant T as Trigger Source
+    participant ETM as EventTriggerMiddleware
+    participant G as Headless Graph
+    participant LLM as Gemini 2.5 Pro
+    participant TS as Task Store (Redis)
+    participant S as Output Sinks
+
+    T->>ETM: TriggerEvent
+    ETM->>TS: status → "processing"
+    ETM->>G: graph.ainvoke(payload)
+    G->>LLM: Generate response
+    LLM-->>G: AI response
+    G-->>ETM: Output
+    ETM->>TS: status → "completed" + result
+    ETM->>S: TriggerResult (fan-out to all sinks)
+
+    Note over S: stdout + file + webhook + redis
all receive every result
+
+
+
+ + +
+

Server → Headless Delegation

+

The orchestrator delegates long-running work to the headless worker via Redis Streams with task status tracking.

+ +
+ 
+sequenceDiagram
+    participant U as User
+    participant O as Orchestrator
+    participant TS as Task Store
+    participant RS as Redis Stream
+    participant HW as Headless Worker
+    participant LLM as Gemini
+
+    U->>O: "Generate BMI reports for 500 employees"
+    O->>TS: create_task(status=queued)
+    O->>RS: XADD agent-tasks
+    O-->>U: "Queued! Task ID: abc123"
+
+    Note over U,O: User can ask status anytime
+
+    U->>O: "What's the status of abc123?"
+    O->>TS: get_task(abc123)
+    TS-->>O: status=queued
+    O-->>U: "Still queued"
+
+    RS->>HW: XREADGROUP (consume)
+    HW->>TS: status → processing
+    HW->>LLM: Process task
+    LLM-->>HW: Result
+    HW->>TS: status → completed + result
+
+    Note over U,O: Next conversation
+
+    U->>O: "Hi, any updates?"
+    O->>TS: get_pending_results(user_id)
+    TS-->>O: abc123 completed!
+    O-->>U: "Your report is done!
498/500 reports generated"
+    O->>TS: mark_delivered(abc123)
+
+
+ +
+

Key Insight: Proactive Result Delivery

+

The orchestrator calls get_pending_results(user_id) at the start of every conversation. Completed tasks are delivered automatically — the user doesn't have to ask. Results are marked as delivered so they're not repeated.

+
+
+ + +
+

Task Status Tracking

+ +
+ 
+stateDiagram-v2
+    [*] --> queued: queue_task()
+    queued --> processing: Headless worker picks up
+    processing --> completed: Graph invocation succeeds
+    processing --> failed: Graph invocation fails
+    completed --> delivered: get_pending_results() / check_task_status()
+    failed --> delivered: get_pending_results()
+    delivered --> [*]: TTL expires (24h)
+
+
+ +
+

Redis Key Structure

+ + + + + + + + +
KeyTypeTTLPurpose
task:{task_id}String (JSON)24hFull task record with status, payload, result
user_tasks:{user_id}Sorted Set24hIndex of task IDs per user, sorted by created_at
agent-tasksStreamPersistentQueue for headless worker consumption
agent-resultsStreamPersistentOutput from Redis sink (if enabled)
+
+ +
+

Tools

+ + + + + + + +
ToolUsed ByPurpose
queue_taskServer agentQueue work for headless worker, returns task_id
check_task_statusServer agentLook up task status by ID
get_pending_resultsServer agentFetch all completed-but-undelivered results for a user
+
+
+ + +
+

Webhook Trigger

+ trigger + +
+

How It Works

+

A minimal HTTP server (built on asyncio.start_server) listens for POST requests. JSON body is parsed into a TriggerEvent and queued for processing. No framework dependency — pure Python asyncio.

+ +

Configuration

+
triggers:
+  webhook:
+    enabled: true
+    host: "0.0.0.0"
+    port: 8888
+    path: "/trigger"
+ +

Example Request

+
curl -X POST http://localhost:8888/trigger \
+  -H "Content-Type: application/json" \
+  -d '{"event": "bmi-calc", "task": "Calculate BMI for Alice, 65kg 170cm"}'
+ +

Response Codes

+ + + + + + + + +
CodeMeaning
200Event accepted and queued
400Invalid JSON body
404Wrong path
405Method not allowed (only POST)
+
+
+ +
+

Cron Trigger

+ trigger + +
+

How It Works

+

Uses APScheduler v4 CronTrigger to parse standard 5-field crontab expressions. Each job runs as a background asyncio task that sleeps until the next fire time, then emits a TriggerEvent.

+ +

Configuration

+
triggers:
+  cron:
+    enabled: true
+    jobs:
+      - name: "daily-report"
+        schedule: "0 9 * * *"       # 9 AM daily
+        payload:
+          task: "Generate daily health digest"
+      - name: "weekly-cleanup"
+        schedule: "0 0 * * 0"       # Midnight Sunday
+        payload:
+          task: "Clean up expired records"
+ +

Crontab Format

+ 
┌───────── minute (0-59)
+│ ┌─────── hour (0-23)
+│ │ ┌───── day of month (1-31)
+│ │ │ ┌─── month (1-12)
+│ │ │ │ ┌─ day of week (0-6, Sun=0)
+│ │ │ │ │
+* * * * *
+
+
+ +
+

Queue Trigger (Redis Streams)

+ trigger + +
+

How It Works

+

Implements the QueueConsumer protocol with a Redis Streams backend. Uses consumer groups for multi-replica support — each pod gets unique messages, no duplicates.

+ +

Configuration

+
triggers:
+  queue:
+    enabled: true
+    backend: "redis_streams"
+    stream: "agent-tasks"
+    consumer_group: "agent-workers"
+    consumer_name: ""              # defaults to $HOSTNAME for K8s
+ +

Consumer Group Behavior

+ + + + + + + + + +
ScenarioBehavior
First consumer joinsGroup created automatically (XGROUP CREATE MKSTREAM)
Multiple replicasEach gets unique messages (round-robin)
Consumer crashesUnacked messages re-delivered to other consumers
Scale upNew pod joins group immediately
Connection lostReconnect with exponential backoff (max 60s)
+ +
+ Extensible: The QueueConsumer protocol is abstract. Implement it for Kafka, RabbitMQ, or SQS — just provide consume(), ack(), and close(). +
+
+
+ + +
+

Stdout Sink

+ sink +
+

Writes TriggerResult as JSON to stdout. Useful for development and log aggregation.

+
output_sinks:
+  - type: stdout
+
+
+ +
+

File Sink (JSONL)

+ sink +
+

Appends each result as a JSON line to a file. Creates parent directories if needed. Flushes after each write.

+
output_sinks:
+  - type: file
+    path: "/var/log/agent/results.jsonl"
+

Read results: cat /var/log/agent/results.jsonl | jq .

+
+
+ +
+

Webhook Sink

+ sink +
+

POSTs results to a URL. Retries on 5xx errors (3 attempts, exponential backoff). Custom headers supported.

+
output_sinks:
+  - type: webhook
+    url: "https://downstream.example.com/results"
+    headers:
+      Authorization: "Bearer ${WEBHOOK_TOKEN}"
+
+
+ +
+

Redis Stream Sink

+ sink +
+

Publishes results to a Redis Stream via XADD. Downstream services can consume from this stream.

+
output_sinks:
+  - type: redis
+    stream: "agent-results"
+

Viewable in Redis Commander or via XRANGE agent-results - +.

+
+
+ + +
+

agent.yaml Reference

+

All headless configuration lives in config/agent/runtime/agent.yaml alongside server config.

+ +
+

Full Headless Configuration

+
# Mode is determined by entry point, not this field
+mode: server
+
+# Trigger sources (used by headless worker)
+triggers:
+  webhook:
+    enabled: false
+    host: "0.0.0.0"
+    port: 8888
+    path: "/trigger"
+  cron:
+    enabled: false
+    jobs: []
+  queue:
+    enabled: false
+    backend: "redis_streams"
+    stream: "agent-tasks"
+    consumer_group: "agent-workers"
+    consumer_name: ""          # defaults to $HOSTNAME
+
+# Output sinks (fan-out — all enabled sinks receive every result)
+output_sinks: []               # defaults to stdout if empty
+# Examples:
+#  - type: stdout
+#  - type: file
+#    path: "/var/log/agent/results.jsonl"
+#  - type: webhook
+#    url: "https://example.com/callback"
+#    headers:
+#      Authorization: "Bearer token"
+#  - type: redis
+#    stream: "agent-results"
+
+# Health check for K8s probes
+health_check:
+  enabled: true
+  host: "0.0.0.0"
+  port: 8080
+
+ +
+ Note: The mode field is informational. The actual mode is determined by how the agent is launched: make local = server, make headless = headless. +
+
+ +
+

Agent Prompts

+ +
+
+

PROMPT.md (Server / Orchestrator)

+
    +
  • Conversational, user-facing
    • +
  • Creates TODO lists
    • +
  • Delegates to analyst/publisher subagents
    • +
  • Tools: validate_email, queue_task, check_task_status, get_pending_results
    • +
  • Knows how to delegate to headless worker
    • +
+
+
+

HEADLESS_PROMPT.md (Worker)

+
    +
  • Silent, no user interaction
    • +
  • No TODO lists, no greetings
    • +
  • Processes payloads directly
    • +
  • Tools: calculate_bmi, search_web
    • +
  • Returns structured JSON results
    • +
+
+
+
+ +
+

Tools Registry

+ + + + + + + + + + + +
ToolTypeAvailable InPurpose
validate_emailMCPServerValidate email addresses
calculate_bmiMCPHeadlessCalculate BMI from height/weight
search_webMCPHeadlessSearch for health information
send_emailMCPServer (publisher)Send email via MCP server
queue_taskBuilt-inServerQueue work for headless worker
check_task_statusBuilt-inServerCheck background task status
get_pending_resultsBuilt-inServerFetch undelivered completed results
+
+ + +
+

Local Development

+ +
+

Commands

+ + + + + + + + + + +
CommandWhat It Does
make localStart server agent (Aegra) + Postgres + Redis
make headlessStart headless worker + Postgres + Redis
make testRun unit tests
make test-triggersRun trigger unit tests only
make test-integrationRun integration tests (requires Redis)
make test-headlessRun headless startup tests
+ +

Running Both Together

+
# Terminal 1: Server agent
+make local
+
+# Terminal 2: Headless worker
+make headless
+
+# Terminal 3: Watch results
+tail -f /tmp/headless-results.jsonl
+
+# Terminal 4: Send test events
+curl -X POST http://localhost:8888/trigger \
+  -H "Content-Type: application/json" \
+  -d '{"event":"test","task":"Calculate BMI for Alice, 65kg 170cm"}'
+
+
+ +
+

OpenShift Deployment

+ +
+

Architecture

+
+ 
+graph TB
+    subgraph NS["OpenShift Namespace"]
+        subgraph SERVER_D["Deployment: agent"]
+            S1["Pod 1"]
+            S2["Pod 2"]
+        end
+        subgraph HEADLESS_D["Deployment: agent-headless"]
+            H1["Pod 1"]
+            H2["Pod 2"]
+            H3["Pod 3"]
+        end
+        SVC_S["Service: agent
Port 5002"]
+        SVC_H["Service: agent-headless
Port 8080 (health)"]
+        ROUTE["Route: agent
External HTTPS"]
+        RD["Redis"]
+        PG["PostgreSQL"]
+        HPA_S["HPA: agent
2-10 replicas"]
+        HPA_H["HPA: agent-headless
2-8 replicas"]
+        KEDA["KEDA ScaledObject
Scale on queue depth"]
+    end
+
+    S1 & S2 --> SVC_S --> ROUTE
+    H1 & H2 & H3 --> SVC_H
+    HPA_S --> SERVER_D
+    HPA_H --> HEADLESS_D
+    KEDA --> HEADLESS_D
+    SERVER_D --> RD & PG
+    HEADLESS_D --> RD & PG
+
+    style SERVER_D fill:#1f6feb22,stroke:#1f6feb
+    style HEADLESS_D fill:#23863622,stroke:#238636
+
+
+ +

Deploy Commands

+
# Deploy server agent
+make deploy openshift NAMESPACE=my-project
+
+# Deploy headless worker
+make deploy-headless NAMESPACE=my-project
+
+# Undeploy
+make undeploy openshift NAMESPACE=my-project
+make undeploy-headless NAMESPACE=my-project
+ +
+ Same image: Both deployments use the same container image. Only the command differs — aegra dev vs python -m deep_agent.headless. +
+
+
+ +
+

Scaling & High Availability

+ +
+

How Scaling Works

+ + + + + + + + + + +
ConcernMechanism
Multiple workersRedis consumer groups — each pod gets unique messages
Scale upNew pod joins consumer group, starts receiving immediately
Scale downSIGTERM → drain in-flight → exit. Unacked messages re-delivered
Pod crashUnacked messages go to other consumers
AutoscalingKEDA watches Redis Stream pending count
Consumer identityEach pod uses $HOSTNAME as consumer name
+ +

KEDA ScaledObject

+
apiVersion: keda.sh/v1alpha1
+kind: ScaledObject
+spec:
+  scaleTargetRef:
+    name: agent-headless
+  minReplicaCount: 1
+  maxReplicaCount: 10
+  triggers:
+    - type: redis-streams
+      metadata:
+        stream: agent-tasks
+        consumerGroup: agent-workers
+        pendingEntriesCount: "10"    # Scale up when >10 pending per replica
+
+ +
+

Health Probes

+ + + + + + +
EndpointPurposeReturns
GET /healthzLiveness probe{"status": "ok"}
GET /readyzReadiness probe{"status": "ready", "sources": N, "sinks": N, "loop_running": true}
+
+
+ + +
+

How Skills Drive Processing

+

The headless agent is not one generic worker — it's specific to the agent it was created with. Skills define what the worker knows how to do.

+ +
+

The Three Layers

+
+ 
+graph TB
+    subgraph GENERIC["Layer 1: Generic Scaffolding"]
+        HP["HEADLESS_PROMPT.md
'Process tasks silently.
Use your skills. Return JSON.'"]
+    end
+
+    subgraph DOMAIN["Layer 2: Domain Knowledge (Skills)"]
+        S1["bmi-report/
README.md"]
+        S2["order-fulfillment/
README.md"]
+        S3["data-transform/
README.md"]
+    end
+
+    subgraph CAPS["Layer 3: Capabilities (Tools / MCP)"]
+        T1["calculate_bmi"]
+        T2["check_inventory
process_payment"]
+        T3["query_warehouse
write_report"]
+    end
+
+    HP --> S1 & S2 & S3
+    S1 --> T1
+    S2 --> T2
+    S3 --> T3
+
+    style GENERIC fill:#1a2332,stroke:#58a6ff,color:#e6edf3
+    style DOMAIN fill:#1a2332,stroke:#3fb950,color:#e6edf3
+    style CAPS fill:#1a2332,stroke:#d29922,color:#e6edf3
+
+
+
+ +
+

Each Agent Gets Its Own Headless Worker

+ + + + + + + +
AgentSkillsToolsWhat Headless Can Do
Health Assistantbmi-reportcalculate_bmi, search_webBMI calculations, health reports
Order Processororder-fulfillmentcheck_inventory, process_paymentValidate, process, confirm orders
Data Pipelinedata-transform, quality-checkquery_warehouse, write_reportETL, quality checks, reports
+
+ +
+

What a Skill Document Contains

+

The skill README is where all domain-specific processing logic lives:

+
# Order Fulfillment Skill
+
+## Processing Steps
+1. Validate items exist in catalog (use check_inventory tool)
+2. Verify stock availability for each item
+3. Calculate total price including tax
+4. Process payment (use process_payment tool)
+5. Update order status (use update_order tool)
+
+## Output Format
+{
+  "status": "success",
+  "order_id": "ORD-1234",
+  "total": 149.99,
+  "estimated_delivery": "2026-06-28"
+}
+
+## Error Handling
+- Out of stock → status "partial", list unavailable items
+- Payment failed → status "error", include error code
+
+ +
+

Runtime Flow

+
+ 
+sequenceDiagram
+    participant EXT as External System
+    participant Q as Kafka / Redis / Webhook
+    participant HW as Headless Worker
+    participant LLM as Gemini + Skill
+    participant TOOL as MCP Tools
+    participant SINK as Output Sinks
+
+    EXT->>Q: {"task": "Process order ORD-1234"}
+    Q->>HW: TriggerEvent
+    HW->>LLM: payload + HEADLESS_PROMPT + order-fulfillment skill
+    LLM->>TOOL: check_inventory(sku)
+    TOOL-->>LLM: in stock
+    LLM->>TOOL: process_payment(149.99)
+    TOOL-->>LLM: payment confirmed
+    LLM-->>HW: {"status": "success", "order_id": "ORD-1234"}
+    HW->>SINK: Result → stdout + file + Redis
+
+
+
+ +
+

Key Insight: No Custom Code Per Use Case

+

The headless worker is generic infrastructure. Domain knowledge lives in skill documents (markdown). Capabilities come from tools and MCP servers. To support a new use case, create a skill README and attach the right tools — no Python code needed.

+
+
+ + +
+

Audit Trail (PostgreSQL)

+

Every task is persisted to PostgreSQL for permanent audit. Redis provides speed (24h TTL); Postgres provides durability.

+ +
+

Tasks Table Schema

+
CREATE TABLE tasks (
+    task_id      TEXT PRIMARY KEY,
+    task_name    TEXT NOT NULL,
+    status       TEXT NOT NULL DEFAULT 'queued',
+    payload      JSONB NOT NULL DEFAULT '{}',
+    result       TEXT,
+    error        TEXT,
+    thread_id    TEXT,
+    user_id      TEXT,
+    delivered    BOOLEAN NOT NULL DEFAULT FALSE,
+    created_at   TIMESTAMPTZ NOT NULL DEFAULT now(),
+    updated_at   TIMESTAMPTZ NOT NULL DEFAULT now(),
+    completed_at TIMESTAMPTZ
+);
+

Indexes

+
    +
  • idx_tasks_user_id — query tasks by user
    • +
  • idx_tasks_status — filter by status
    • +
  • idx_tasks_created_at — order by recency
    • +
+
+ +
+

Dual-Write Architecture

+
+ 
+sequenceDiagram
+    participant QT as queue_task()
+    participant Redis as Redis (speed)
+    participant PG as PostgreSQL (audit)
+
+    QT->>Redis: SET task:{id} (TTL 24h)
+    QT->>PG: INSERT INTO tasks
+
+    Note over Redis,PG: Both writes happen on every
create, update, and deliver
+
+    QT->>Redis: status → processing
+    QT->>PG: UPDATE tasks SET status
+
+    QT->>Redis: status → completed
+    QT->>PG: UPDATE tasks SET status, completed_at
+
+
+
+ Graceful degradation: If Postgres is unavailable, Redis still works. Audit writes fail silently (logged as warnings). The headless worker never blocks on audit failures. +
+
+ +
+

Useful Queries

+
-- All tasks for a user
+SELECT * FROM tasks WHERE user_id = 'naveen' ORDER BY created_at DESC;
+
+-- Failed tasks in the last 24 hours
+SELECT * FROM tasks WHERE status = 'failed'
+  AND created_at > now() - interval '24 hours';
+
+-- Average processing time
+SELECT task_name, AVG(EXTRACT(EPOCH FROM (completed_at - created_at))) as avg_seconds
+  FROM tasks WHERE completed_at IS NOT NULL
+  GROUP BY task_name;
+
+-- Tasks pending delivery
+SELECT * FROM tasks WHERE status = 'completed' AND delivered = FALSE;
+
+
+ +
+

Data Stores

+ + + + + + + + + + +
StoreDataRetentionPurpose
Redis task:{id}Task status, payload, result24h TTLFast lookup for status checks
Redis user_tasks:{id}Sorted set of task IDs24h TTLPer-user task index
Redis agent-tasksStream of queued tasksPersistentQueue for headless consumption
Redis agent-resultsStream of resultsPersistentRedis sink output
PostgreSQL tasksFull task lifecyclePermanentAudit trail, analytics
File results.jsonlJSONL resultsPermanentFile sink output
+
+ + +
+

Open Items & TODOs

+ +
+

Authentication for Headless Worker

+ TODO + +

Currently the headless worker has no authentication — it runs as a static service account (headless-worker) with no SSO token. MCP tools that require user auth will fail.

+ +

Current State

+ + + + + + + + +
ComponentServer AgentHeadless Agent
User identityFrom SSO token (JWT)Static: "headless-worker"
MCP toolsUser's SSO token forwardedNo token — auth MCP calls fail
Redis / PostgresConnection string from envSame — works
LLM (Gemini)Service account credsSame — works
+ +

Proposed Solutions

+ +
+
+

Option A: Service Account Token (Recommended)

+

Headless worker gets its own token via OAuth2 client credentials flow.

+
headless_auth:
+  type: "client_credentials"
+  client_id: "headless-worker"
+  client_secret: "${HEADLESS_CLIENT_SECRET}"
+  token_url: "${SSO_ISSUER_URL}/.../token"
+

HeadlessRuntime requests a token at startup, refreshes periodically. MCP tools see a valid token. Best for production.

+
+
+

Option B: Token Passthrough from Queue

+

Server agent includes the user's SSO token in the task payload.

+
await client.xadd(stream, {
+    "name": task_name,
+    "task_id": task_id,
+    "auth_token": user_sso_token,
+})
+

Headless worker extracts and uses it. Only works for short-lived tasks (tokens expire).

+
+
+ +
+
+

Option C: No Auth for Internal MCP

+

MCP servers expose a trusted internal endpoint without auth.

+
// mcp.json for headless
+{
+  "template-mcp-server": {
+    "url": "http://mcp-server:5001/mcp",
+    "auth": false
+  }
+}
+

Simplest for local dev. Requires MCP server to support unauthenticated calls on internal network.

+
+
+

Recommendation

+

Option A for production (proper service identity).
Option C for local development (zero setup).

+
+
+
+ +
+

Other Open Items

+ + + + + + + + + + + + +
ItemPriorityDescription
Real-time push notificationsP2Option C: Redis pub/sub → UI WebSocket. User gets notified instantly when task completes, without polling.
Auto user_id from runtimeP1Populate user_id from SSO token automatically instead of requiring user to type it.
Auto queue_task routingP1Orchestrator should auto-detect bulk/long-running requests and queue them without user saying "queue a background task".
Ack after processingP1Currently messages are acked before processing. If worker crashes mid-processing, message is lost. Move ack to after completion.
Pending message claimingP2On restart, claim unacked messages from the pending entries list (PEL) so crashed tasks are retried.
Dead letter queueP3Tasks that fail N times should go to a dead letter queue instead of being retried forever.
Task TTL in PostgresP3Auto-archive or delete old audit records after configurable retention period.
Kafka producer in queue_taskP3Allow queue_task to produce to Kafka (not just Redis) for environments where Kafka is the primary queue.
+
+
+ + +
+

Test Guide

+ +
+

Test Matrix

+ + + + + + + + + + + +
TestCommandCountRequirements
Unit tests (triggers)make test-triggers124None
Integration testsmake test-integration21Redis
All unit testsmake test700+None
Manual: webhookcurl POST :8888/trigger—Headless running
Manual: queueXADD agent-tasks—Headless + Redis
Manual: cronWait for schedule—Headless + cron enabled
E2E: UI → queue → headlessChat in UI—Server + headless + UI
+
+ +
+

Quick Smoke Test

+
# 1. Start both agents
+make local          # Terminal 1
+make headless       # Terminal 2
+
+# 2. Send a webhook event
+curl -X POST http://localhost:8888/trigger \
+  -H "Content-Type: application/json" \
+  -d '{"event":"smoke-test","task":"Calculate BMI for Alice, 65kg 170cm"}'
+
+# 3. Check results
+cat /tmp/headless-results.jsonl | python3 -c "
+import sys, json
+for line in sys.stdin:
+    d = json.loads(line)
+    e = d['event']
+    print(f'{e[\"name\"]} ({e[\"source\"]}): success={d[\"success\"]}')"
+
+# 4. Check Redis
+redis-cli XRANGE agent-results - +
+
+
+ +
+

Test Scenarios

+ +
+

Scenario 1: Webhook Trigger

+

External system sends HTTP POST → headless processes → result in all sinks

+
# Send webhook
+curl -X POST http://localhost:8888/trigger \
+  -H "Content-Type: application/json" \
+  -d '{"event":"bmi-calc","task":"Calculate BMI for Alice, 65kg 170cm"}'
+
+# Expected: {"status": "accepted"}
+# Wait ~20s, then check results:
+tail -1 /tmp/headless-results.jsonl | python3 -m json.tool
+
+ +
+

Scenario 2: Queue Trigger (Server → Headless)

+

Server agent queues task → headless consumes from Redis Stream → status tracked

+
# In UI (http://localhost:3000), type:
+# "Queue a background task to calculate BMI for Bob (90kg 180cm). My user_id is naveen."
+
+# Agent responds with task ID, e.g. abc123def456
+
+# Check status in another chat:
+# "Check the status of task abc123def456"
+
+# Or check Redis directly:
+redis-cli GET task:abc123def456 | python3 -m json.tool
+
+ +
+

Scenario 3: Cron Trigger

+

Scheduled job fires automatically → headless processes → result in sinks

+
# Enable cron in agent.yaml:
+# triggers.cron.enabled: true
+# triggers.cron.jobs:
+#   - name: "health-tip"
+#     schedule: "* * * * *"
+#     payload: {"task": "Generate a health tip"}
+
+# Restart headless worker
+# Wait 60 seconds — result appears automatically in:
+tail -f /tmp/headless-results.jsonl
+
+ +
+

Scenario 4: Error Handling

+
# Invalid JSON → 400
+curl -X POST http://localhost:8888/trigger \
+  -H "Content-Type: application/json" -d "not json"
+
+# Wrong path → 404
+curl -X POST http://localhost:8888/wrong \
+  -H "Content-Type: application/json" -d '{"test":1}'
+
+# GET method → 405
+curl http://localhost:8888/trigger
+
+ +
+

Scenario 5: Full E2E with Status Tracking

+

Complete lifecycle: queue → track → process → deliver results

+
# Step 1: Queue task via UI
+#   "Queue a BMI report for 3 employees. My user_id is naveen."
+#   → Agent returns task ID
+
+# Step 2: Check status (same chat)
+#   "What's the status of my task?"
+#   → Agent shows QUEUED or PROCESSING
+
+# Step 3: Wait for completion (~20s)
+
+# Step 4: Check status again
+#   "Check the status of task {id}"
+#   → Agent shows COMPLETED with BMI results
+
+# Step 5: New chat — proactive delivery
+#   "Do I have any pending results? My user_id is naveen."
+#   → Agent delivers undelivered results
+
+ +
+

Scenario 6: Multiple Sinks Verification

+

Single event → all 4 sinks receive the result simultaneously

+
# Send one event
+curl -X POST http://localhost:8888/trigger \
+  -H "Content-Type: application/json" \
+  -d '{"event":"multi-sink","task":"Calculate BMI for Dave, 95kg 185cm"}'
+
+# Check all 4 sinks:
+
+# 1. Stdout — visible in headless terminal
+# 2. File
+tail -1 /tmp/headless-results.jsonl
+
+# 3. Redis Stream
+redis-cli XRANGE agent-results - + COUNT 1
+
+# 4. Redis Commander — http://localhost:8083
+#    Browse "agent-results" stream
+
+
+ +
+

Verify Results

+ +
+

Where to Check

+ + + + + + + + + + + +
WhatWhereCommand / URL
Headless worker logsTerminaltail -f /tmp/headless-agent.log
Task results (file sink)Terminaltail -f /tmp/headless-results.jsonl
Redis tasksRedis Commanderhttp://localhost:8083 → task:* keys
Redis streamsRedis Commanderhttp://localhost:8083 → agent-tasks, agent-results
Postgres auditDBeaverSELECT * FROM tasks ORDER BY created_at DESC;
Health checkBrowser/curlcurl http://localhost:8082/readyz
Server agent healthBrowser/curlcurl http://localhost:5002/health
+
+ +
+

Quick Verification Script

+
#!/bin/bash
+# Save as verify.sh and run after testing
+
+echo "=== Services ==="
+echo "Server:   $(curl -s -o /dev/null -w '%{http_code}' http://localhost:5002/health)"
+echo "Headless: $(curl -s http://localhost:8082/readyz 2>/dev/null)"
+
+echo ""
+echo "=== Redis Tasks ==="
+cd /path/to/template-agent
+.venv/bin/python -c "
+import redis, json
+r = redis.from_url('redis://localhost:6379/0', decode_responses=True)
+for k in sorted(r.keys('task:*'), key=lambda k: json.loads(r.get(k)).get('created_at',''), reverse=True)[:5]:
+    d = json.loads(r.get(k))
+    icon = {'queued':'⏳','processing':'🔄','completed':'✅','failed':'❌'}.get(d['status'],'?')
+    print(f'  {icon} {d[\"task_id\"]}: {d[\"task_name\"]} [{d[\"status\"]}]')
+r.close()
+"
+
+echo ""
+echo "=== Postgres Audit ==="
+.venv/bin/python -c "
+import asyncio, psycopg
+from psycopg.rows import dict_row
+async def go():
+    async with await psycopg.AsyncConnection.connect(
+        'postgresql://postgres:postgres@localhost:5432/template_agent',
+        row_factory=dict_row
+    ) as conn:
+        rows = await conn.execute(
+            'SELECT task_id, task_name, status, user_id, completed_at FROM tasks ORDER BY created_at DESC LIMIT 5'
+        )
+        for r in await rows.fetchall():
+            print(f'  {r[\"task_id\"]}: {r[\"task_name\"]} [{r[\"status\"]}] user={r.get(\"user_id\",\"-\")}')
+asyncio.run(go())
+"
+
+echo ""
+echo "=== Results File ==="
+wc -l /tmp/headless-results.jsonl 2>/dev/null
+
+
+
+ +
+ + + + diff --git a/docs/headless-agent-ui-integration.md b/docs/headless-agent-ui-integration.md new file mode 100644 index 00000000..aa6d8369 --- /dev/null +++ b/docs/headless-agent-ui-integration.md @@ -0,0 +1,725 @@ +# Headless Agent — UI Integration Specification + +## Overview + +This document specifies the UI changes needed to support headless agent creation and management in the AI Factory interface. The backend (template-agent) already supports headless mode — this doc covers what the UI needs to expose and what APIs/configs it generates. + +## Current Agent Creation Flow (Server Only) + +Today the UI creates a regular server agent with: + +``` +User selects → Model, Skills, MCP Servers, Subagents +UI generates → PROMPT.md, agent.yaml, mcp.json, subagent .md files +Platform → Deploys single Deployment (server mode via Aegra) +``` + +## New Flow: Agent with Headless Worker + +``` +User selects → Model, Skills, MCP Servers, Subagents + + Enables headless worker + + Configures triggers and sinks +UI generates → PROMPT.md (orchestrator, includes queue_task tool) + HEADLESS_PROMPT.md (worker, auto-generated) + agent.yaml (includes triggers + sinks + health_check) + mcp.json (shared) +Platform → Deploys TWO Deployments (server + headless) +``` + +--- + +## UI Changes Required + +### 1. Agent Creation Form — New Section + +Add a collapsible "Background Worker" section to the agent creation form: + +``` +┌─────────────────────────────────────────────────────┐ +│ Create Agent │ +│ │ +│ Name: [Health Assistant ] │ +│ Model: [gemini-2.5-pro ▼ ] │ +│ Skills: [bmi-report] [client-intake] [+] │ +│ MCP Servers: [template-mcp-server] [+] │ +│ Subagents: [analyst] [publisher] [+] │ +│ │ +│ ── Background Worker (optional) ────────────────── │ +│ [ ] Enable headless worker │ +│ │ +│ ▶ Triggers (collapsed when disabled) │ +│ ▶ Output Sinks │ +│ ▶ Worker Prompt │ +│ ▶ Health Check │ +│ │ +│ [Create Agent] │ +└─────────────────────────────────────────────────────┘ +``` + +### 2. Triggers Configuration + +When "Enable headless worker" is checked, expand the triggers section: + +``` +┌─────────────────────────────────────────────────────┐ +│ Triggers │ +│ │ +│ ☑ Queue Consumer │ +│ Backend: [Redis Streams ▼] │ +│ ├── Redis Streams │ +│ └── Kafka │ +│ Stream/Topic: [agent-tasks ] │ +│ Consumer Group: [agent-workers ] │ +│ │ +│ (if Kafka selected) │ +│ Bootstrap Servers: [localhost:9092 ] │ +│ │ +│ ☐ Webhook Listener │ +│ Port: [8888 ] │ +│ Path: [/trigger] │ +│ │ +│ ☐ Cron Jobs │ +│ [+ Add Job] │ +│ ┌─ Job 1 ────────────────────────────────┐ │ +│ │ Name: [daily-report ] │ │ +│ │ Schedule: [0 9 * * * ] │ │ +│ │ Payload: [{"task": "gen report"} ] │ │ +│ └────────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────┘ +``` + +### 3. Output Sinks Configuration + +``` +┌─────────────────────────────────────────────────────┐ +│ Output Sinks (results fan out to all enabled sinks) │ +│ │ +│ [+ Add Sink] │ +│ │ +│ ┌─ Sink 1 ─────────────────────────────────┐ │ +│ │ Type: [Stdout ▼] │ │ +│ │ ├── Stdout (console/logs) │ │ +│ │ ├── File (JSONL) │ │ +│ │ ├── Webhook (HTTP POST) │ │ +│ │ └── Redis Stream │ │ +│ └──────────────────────────────────────────┘ │ +│ │ +│ ┌─ Sink 2 ─────────────────────────────────┐ │ +│ │ Type: [File ▼] │ │ +│ │ Path: [/var/log/agent/results.jsonl] │ │ +│ └──────────────────────────────────────────┘ │ +│ │ +│ ┌─ Sink 3 ─────────────────────────────────┐ │ +│ │ Type: [Webhook ▼] │ │ +│ │ URL: [https://example.com/callback] │ │ +│ │ Headers: │ │ +│ │ Authorization: [Bearer ${TOKEN}] │ │ +│ └──────────────────────────────────────────┘ │ +│ │ +│ ┌─ Sink 4 ─────────────────────────────────┐ │ +│ │ Type: [Redis Stream ▼] │ │ +│ │ Stream: [agent-results] │ │ +│ └──────────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────┘ +``` + +### 4. Worker Prompt Configuration + +``` +┌─────────────────────────────────────────────────────┐ +│ Worker Prompt │ +│ │ +│ ○ Auto-generate from agent skills (recommended) │ +│ The worker prompt is generated automatically │ +│ using the agent's skills and tools. It strips │ +│ conversational behavior (TODO lists, greetings) │ +│ and focuses on silent task processing. │ +│ │ +│ ○ Custom prompt │ +│ ┌──────────────────────────────────────────┐ │ +│ │ You are a background task processor... │ │ +│ │ │ │ +│ │ │ │ +│ └──────────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────┘ +``` + +### 5. Health Check Configuration + +``` +┌─────────────────────────────────────────────────────┐ +│ Health Check │ +│ │ +│ ☑ Enable health endpoint │ +│ Port: [8080] │ +│ │ +│ Endpoints: │ +│ /healthz — liveness probe (always 200 if up) │ +│ /readyz — readiness probe (checks sources/loop) │ +└─────────────────────────────────────────────────────┘ +``` + +--- + +## Files Generated by UI + +When the user clicks "Create Agent" with headless enabled, the UI must generate these files: + +### 1. PROMPT.md (Orchestrator — modified) + +The existing orchestrator prompt, but with `queue_task`, `check_task_status`, and `get_pending_results` tools auto-added: + +```yaml +--- +name: orchestrator +model: gemini-2.5-pro # from user selection +tools: + - validate_email # from user selection + - queue_task # AUTO-ADDED when headless enabled + - check_task_status # AUTO-ADDED when headless enabled + - get_pending_results # AUTO-ADDED when headless enabled +skills: + - client-intake # from user selection +mcps: + - template-mcp-server # from user selection +--- + +# {Agent Name} + +{User-provided or auto-generated orchestrator prompt} + +## Background Tasks (Headless Worker) + +{AUTO-INJECTED section — see template below} +``` + +**Auto-injected background tasks section:** + +```markdown +## Background Tasks (Headless Worker) + +A headless worker runs alongside you as a background processor. +Use `queue_task` to delegate work that is long-running, bulk, +or doesn't need an immediate response. + +**When to use queue_task:** +- Bulk operations (e.g., "generate reports for all clients") +- Long-running processing (e.g., "export all data") +- Fire-and-forget notifications + +**When NOT to use queue_task:** +- Anything the user expects an immediate answer to + +**Status tracking — CRITICAL RULES:** +1. `queue_task` returns a task ID — give this to the user +2. When the user asks about task status, ALWAYS call `check_task_status(task_id)` +3. Show complete results from `check_task_status` — never say "check a file" +4. At the start of every conversation, call `get_pending_results(user_id)` + to deliver completed background task results proactively +``` + +### 2. HEADLESS_PROMPT.md (Worker — auto-generated) + +```yaml +--- +name: headless-worker +model: gemini-2.5-pro # SAME model as orchestrator +tools: # SAME tools as orchestrator (minus queue_task etc.) + - calculate_bmi + - search_web +skills: # SAME skills as orchestrator + - bmi-report +mcps: # SAME MCP servers + - template-mcp-server +--- + +# Background Task Processor + +You are a background task processor. You receive tasks from a queue +and process them silently. + +## Rules + +1. No greetings, no TODO lists, no conversational responses. +2. Process the payload directly using your tools. +3. Return structured JSON results. +4. Handle errors gracefully with clear error messages. + +## Task Processing + +When you receive a task payload: +1. Parse the task name and data +2. Execute the work using your tools +3. Return a JSON result: + - status: "success" or "error" + - summary: Brief description + - data: The actual results + - error: Error message if failed +``` + +**Auto-generation logic:** +- `model` → copy from orchestrator +- `tools` → copy from orchestrator, REMOVE: `queue_task`, `check_task_status`, `get_pending_results`, `validate_email` +- `skills` → copy from orchestrator +- `mcps` → copy from orchestrator +- System prompt → use the standard worker template above + +### 3. agent.yaml (Runtime config — extended) + +Append these sections to the existing agent.yaml: + +```yaml +# ── Triggers (headless mode only) ── +triggers: + webhook: + enabled: true # from UI checkbox + host: "0.0.0.0" + port: 8888 # from UI input + path: "/trigger" # from UI input + cron: + enabled: false # from UI checkbox + jobs: [] # from UI job list + queue: + enabled: true # from UI checkbox + backend: "redis_streams" # from UI dropdown + stream: "agent-tasks" # from UI input + consumer_group: "agent-workers" # from UI input + consumer_name: "" # auto: defaults to $HOSTNAME + # Kafka-specific (only when backend=kafka): + bootstrap_servers: "localhost:9092" # from UI input + topic: "agent-tasks" # from UI input + +# ── Output Sinks (headless mode only) ── +output_sinks: # from UI sink list + - type: stdout + - type: file + path: "/var/log/agent/results.jsonl" + - type: redis + stream: "agent-results" + +# ── Health Check (headless mode only) ── +health_check: + enabled: true # from UI checkbox + host: "0.0.0.0" + port: 8080 # from UI input +``` + +--- + +## Backend API Requirements + +The UI needs these backend APIs to support headless agent management: + +### 1. Agent Creation (existing, extended) + +``` +POST /api/agents +``` + +Existing payload, with new optional `headless` field: + +```json +{ + "name": "Health Assistant", + "model": "gemini-2.5-pro", + "skills": ["bmi-report"], + "mcps": ["template-mcp-server"], + "subagents": ["analyst", "publisher"], + "headless": { + "enabled": true, + "triggers": { + "webhook": {"enabled": true, "port": 8888, "path": "/trigger"}, + "cron": {"enabled": false, "jobs": []}, + "queue": { + "enabled": true, + "backend": "redis_streams", + "stream": "agent-tasks", + "consumer_group": "agent-workers" + } + }, + "output_sinks": [ + {"type": "stdout"}, + {"type": "file", "path": "/var/log/agent/results.jsonl"} + ], + "health_check": {"enabled": true, "port": 8080}, + "worker_prompt": "auto" + } +} +``` + +The backend should: +1. Generate `PROMPT.md` with `queue_task` tools auto-added +2. Generate `HEADLESS_PROMPT.md` (auto or custom) +3. Write trigger/sink config to `agent.yaml` +4. Create two Deployments (server + headless) in the kustomize overlay + +### 2. Task Status API (new) + +``` +GET /api/agents/{agent_id}/tasks +``` + +Returns task history from PostgreSQL audit table: + +```json +{ + "tasks": [ + { + "task_id": "abc123", + "task_name": "bulk-bmi-report", + "status": "completed", + "user_id": "naveen", + "created_at": "2026-06-24T13:34:50Z", + "completed_at": "2026-06-24T13:35:12Z", + "duration_seconds": 22, + "delivered": true + } + ], + "total": 1, + "page": 1 +} +``` + +Query parameters: +- `user_id` — filter by user +- `status` — filter by status (queued, processing, completed, failed) +- `limit`, `offset` — pagination + +SQL behind this API: +```sql +SELECT * FROM tasks +WHERE user_id = $1 AND status = $2 +ORDER BY created_at DESC +LIMIT $3 OFFSET $4; +``` + +### 3. Task Detail API (new) + +``` +GET /api/agents/{agent_id}/tasks/{task_id} +``` + +Returns full task detail including result: + +```json +{ + "task_id": "abc123", + "task_name": "bulk-bmi-report", + "status": "completed", + "payload": {"employee_count": 500}, + "result": "{\"status\": \"success\", \"data\": [...]}", + "error": null, + "user_id": "naveen", + "thread_id": "thread-xyz", + "delivered": true, + "created_at": "2026-06-24T13:34:50Z", + "updated_at": "2026-06-24T13:35:12Z", + "completed_at": "2026-06-24T13:35:12Z" +} +``` + +### 4. Headless Worker Status API (new) + +``` +GET /api/agents/{agent_id}/headless/status +``` + +Proxies to the headless worker's health endpoint: + +```json +{ + "status": "ready", + "sources": 2, + "sinks": 3, + "loop_running": true, + "uptime_seconds": 3600 +} +``` + +--- + +## UI Pages / Views + +### 1. Agent Detail Page — New "Background Tasks" Tab + +``` +┌──────────────────────────────────────────────────────┐ +│ Health Assistant │ +│ │ +│ [Chat] [Settings] [Background Tasks] [Monitoring] │ +│ │ +│ ── Background Tasks ────────────────────────────────│ +│ │ +│ Worker Status: ● Ready (2 sources, 3 sinks) │ +│ │ +│ ┌────────────────────────────────────────────────┐ │ +│ │ ID Name Status Duration│ │ +│ │ abc123 bulk-bmi-report ✅ Done 22s │ │ +│ │ def456 weekly-digest 🔄 Running 5s │ │ +│ │ ghi789 data-export ❌ Failed 3s │ │ +│ │ jkl012 email-blast ⏳ Queued — │ │ +│ └────────────────────────────────────────────────┘ │ +│ │ +│ Showing 4 of 127 tasks [< 1 2 3 ... 13 >] │ +└──────────────────────────────────────────────────────┘ +``` + +### 2. Task Detail Modal + +Click on a task row to see details: + +``` +┌──────────────────────────────────────────────────────┐ +│ Task: bulk-bmi-report (abc123) │ +│ │ +│ Status: ✅ Completed │ +│ Created: 2026-06-24 13:34:50 │ +│ Completed: 2026-06-24 13:35:12 │ +│ Duration: 22 seconds │ +│ User: naveen │ +│ Thread: thread-xyz │ +│ Delivered: Yes │ +│ │ +│ ── Payload ──────────────────────────────────────── │ +│ { │ +│ "employee_count": 500 │ +│ } │ +│ │ +│ ── Result ───────────────────────────────────────── │ +│ { │ +│ "status": "success", │ +│ "summary": "Processed 500 BMI calculations", │ +│ "data": [...] │ +│ } │ +│ │ +│ [Close] │ +└──────────────────────────────────────────────────────┘ +``` + +### 3. Agent Dashboard — Headless Worker Card + +On the main dashboard, show headless worker status alongside the server agent: + +``` +┌─────────────────────┐ ┌─────────────────────────┐ +│ Server Agent │ │ Headless Worker │ +│ ● Running │ │ ● Ready │ +│ Port: 5002 │ │ Sources: webhook, queue │ +│ Threads: 42 │ │ Sinks: stdout, file │ +│ Uptime: 3h 15m │ │ Tasks today: 23 │ +│ │ │ Success rate: 96% │ +└─────────────────────┘ └─────────────────────────┘ +``` + +--- + +## Data Flow Summary + +``` +┌────────────────────────────────────────────────────────────┐ +│ AI Factory UI │ +├────────────────────────────────────────────────────────────┤ +│ │ +│ Create Agent Form │ +│ ├── Model, Skills, MCP, Subagents │ +│ └── Headless: Triggers, Sinks, Health Check │ +│ │ │ +│ ▼ │ +│ POST /api/agents { headless: { ... } } │ +│ │ │ +│ ▼ │ +│ Backend generates: │ +│ ├── PROMPT.md (with queue_task auto-added) │ +│ ├── HEADLESS_PROMPT.md (auto-generated) │ +│ ├── agent.yaml (triggers + sinks + health_check) │ +│ └── Kustomize overlays (server + headless Deployments) │ +│ │ │ +│ ▼ │ +│ Deployed to OpenShift: │ +│ ├── Deployment: agent (server mode, port 5002) │ +│ └── Deployment: agent-headless (headless mode) │ +│ │ +│ Task Monitoring: │ +│ ├── GET /api/agents/{id}/tasks → PostgreSQL audit table │ +│ ├── GET /api/agents/{id}/headless/status → health check │ +│ └── Background Tasks tab in agent detail page │ +│ │ +└────────────────────────────────────────────────────────────┘ +``` + +--- + +## Configuration Reference + +### Trigger Types + +| Type | Config Fields | Description | +|------|--------------|-------------| +| Queue (Redis Streams) | `backend`, `stream`, `consumer_group` | Consumes from Redis Stream. Default for server→headless delegation. | +| Queue (Kafka) | `backend`, `topic`, `bootstrap_servers`, `consumer_group` | Consumes from Kafka topic. For external system events. | +| Webhook | `port`, `path` | Minimal HTTP listener. For external system webhooks. | +| Cron | `jobs[].name`, `jobs[].schedule`, `jobs[].payload` | Scheduled triggers. Standard 5-field crontab syntax. | + +### Sink Types + +| Type | Config Fields | Description | +|------|--------------|-------------| +| Stdout | (none) | Prints to process stdout / container logs | +| File | `path` | Appends JSONL to file | +| Webhook | `url`, `headers` | POSTs result to URL with retry (3 attempts) | +| Redis Stream | `stream` | XADD to Redis Stream | + +### Queue Backend Comparison + +| Feature | Redis Streams | Kafka | +|---------|--------------|-------| +| Setup complexity | None (already in stack) | Needs Kafka cluster | +| Message ordering | Per-stream | Per-partition | +| Consumer groups | Yes (XREADGROUP) | Yes (native) | +| Message persistence | Configurable | Default persistent | +| Multi-replica | Yes (consumer groups) | Yes (consumer groups) | +| Use case | Internal delegation | External events, high throughput | + +--- + +## How Skills Drive Headless Processing + +The headless agent is **not one generic worker**. It's specific to the agent it was created with. The **skills** define what the headless worker knows how to do. + +### The Three Layers + +``` +HEADLESS_PROMPT.md (generic scaffolding) + "You are a background processor. Process tasks silently. Return JSON." + + + + +Skill Documents (domain knowledge) + config/agent/skills/order-fulfillment/README.md + "To process an order: + 1. Call check_inventory(sku) + 2. If in stock, call process_payment(amount) + 3. Call update_order(status='confirmed') + 4. Return order confirmation JSON" + + + + +Tools / MCP Servers (capabilities) + check_inventory, process_payment, update_order + + = + +Headless worker that knows how to process orders +``` + +### Each Agent Gets Its Own Headless Worker + +``` +AI Factory +├── Health Assistant Agent +│ ├── Server: PROMPT.md (orchestrator for health) +│ └── Headless: HEADLESS_PROMPT.md +│ Skills: [bmi-report] +│ Tools: [calculate_bmi, search_web] +│ → Knows how to calculate BMI and generate health reports +│ +├── Order Processing Agent +│ ├── Server: PROMPT.md (orchestrator for orders) +│ └── Headless: HEADLESS_PROMPT.md +│ Skills: [order-fulfillment, inventory-check] +│ Tools: [check_inventory, process_payment, update_order] +│ → Knows how to validate, process, and confirm orders +│ +├── Data Pipeline Agent +│ ├── Server: PROMPT.md (orchestrator for data) +│ └── Headless: HEADLESS_PROMPT.md +│ Skills: [data-transform, quality-check] +│ Tools: [query_warehouse, write_report] +│ → Knows how to run ETL and quality checks +``` + +### What the Skill Document Contains + +A skill README is where all domain-specific processing logic lives: + +```markdown +# Order Fulfillment Skill + +## Input Format +- order_id: string (required) +- items: list of {sku, quantity} + +## Processing Steps +1. Validate all items exist in catalog (use check_inventory tool) +2. Verify stock availability for each item +3. Calculate total price including tax +4. Process payment (use process_payment tool) +5. Update order status to confirmed (use update_order tool) +6. Generate confirmation with estimated delivery date + +## Output Format +{ + "status": "success", + "order_id": "ORD-1234", + "total": 149.99, + "items_fulfilled": 3, + "estimated_delivery": "2026-06-28" +} + +## Error Handling +- Out of stock → return status "partial", list unavailable items +- Payment failed → return status "error", include payment error code +- Invalid order_id → return status "error", message "Order not found" +``` + +### How It Flows at Runtime + +``` +External system sends: + {"name": "process-order", "task": "Process order ORD-1234", "order_id": "ORD-1234"} + │ + ▼ +Headless worker receives via Kafka/Redis/Webhook + │ + ▼ +HEADLESS_PROMPT.md: "You are a background processor. Use your skills." + │ + ▼ +LLM loads order-fulfillment skill: "Step 1: Call check_inventory..." + │ + ▼ +LLM calls tools: check_inventory → process_payment → update_order + │ + ▼ +Result: {"status": "success", "order_id": "ORD-1234", "total": 149.99} + │ + ▼ +Output sinks (Redis, file, webhook callback) +``` + +### UI Implication + +When the user creates a headless agent in the UI: +- The **skills they attach** define what the worker can do +- The **tools/MCP servers** give it the capabilities +- The **HEADLESS_PROMPT.md** is generic scaffolding (auto-generated) +- **No custom code needed per use case** — just a skill document and the right tools + +The UI should show a preview: "This headless worker will be able to handle: BMI calculations (bmi-report skill), health data search (search_web tool)." + +--- + +## Implementation Notes for UI Team + +1. **The `headless` field in the agent creation API is optional.** If omitted, only the server agent is created (current behavior). + +2. **Auto-generation of HEADLESS_PROMPT.md** should copy model/skills/MCP from the orchestrator and use the standard worker template. The UI should show a preview before creation. + +3. **The `queue_task`, `check_task_status`, `get_pending_results` tools are built-in** — they don't need to be in any MCP server. The backend auto-registers them when headless is enabled. + +4. **Task status polling:** The UI should poll `GET /api/agents/{id}/tasks?status=processing` every 5-10 seconds when the Background Tasks tab is open, to show real-time status updates. + +5. **The headless worker uses the same container image** as the server agent. Only the startup command differs (`python -m deep_agent.headless` vs `aegra dev`). No separate build needed. + +6. **Health check:** The `/readyz` endpoint returns `503` when the worker is not ready (sources not started, loop not running). Use this for the status indicator in the dashboard. diff --git a/docs/observability.md b/docs/observability.md new file mode 100644 index 00000000..9f3f1b97 --- /dev/null +++ b/docs/observability.md @@ -0,0 +1,229 @@ +# Observability Guide + +This guide covers OpenTelemetry (OTEL) metrics and tracing for the template agent. + +## Overview + +The agent supports three complementary observability layers: + +1. **Langfuse** — LLM-specific tracing (prompts, completions, tokens, costs) +2. **Agent lifecycle OTEL** (`deep_agent/aegra/otel.py`) — conversations, streams, threads, graph builds +3. **Token budget OTEL** (`deep_agent/src/observability/otel_setup.py`) — per-thread/per-user token usage export + +Langfuse and OTEL coexist without conflict. Langfuse traces LLM calls; OTEL traces infrastructure and exports operational metrics. + +## Architecture + +The agent exports telemetry **directly** to OTLP backends. There is no in-repo OTEL Collector deployment. + +``` +Local dev: Agent --OTLP/gRPC--> Jaeger (:4317) +OpenShift: Agent --OTLP/gRPC--> otel-gateway / managed observability service +``` + +Agent metrics export via OTLP, not Prometheus scrape. + +## Local Development + +### Quick Start + +Start Jaeger: + +```bash +docker compose --profile observability up +``` + +This launches **Jaeger** — trace UI at http://localhost:16686 (OTLP gRPC on `:4317`). + +### Enable OTEL in Agent + +Uncomment in `compose.yaml` under `template-agent`: + +```yaml +environment: + - ENABLE_OTEL=true + - OTEL_EXPORTER_OTLP_ENDPOINT=http://jaeger:4317 + - ENABLE_OTEL_TRACES=true + - OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://jaeger:4317 +``` + +Or set in `.env`: + +```bash +ENABLE_OTEL=true +OTEL_EXPORTER_OTLP_ENDPOINT=http://jaeger:4317 +ENABLE_OTEL_TRACES=true +OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://jaeger:4317 +``` + +### View Traces + +**Jaeger UI**: http://localhost:16686 + +1. Select the service name from your agent config (for example `health_assistant`) +2. Click **Find Traces** +3. Inspect HTTP spans (FastAPI auto-instrumentation) and custom spans + +### View Metrics + +Agent lifecycle metrics export via OTLP when `ENABLE_OTEL=true`. For local debugging, use the in-memory snapshot API exposed through the OTEL module (`get_metrics_snapshot()`). + +Token budget metrics export separately when `ENABLE_OTEL_METRICS=true` and `OTEL_EXPORTER_OTLP_ENDPOINT` are set (see token budget docs in `deep_agent/src/token_budget/`). + +## OpenShift Deployment + +### Configuration + +Enable via ConfigMap (env vars override `config/agent/runtime/observability.yaml`): + +```yaml +apiVersion: v1 +kind: ConfigMap +metadata: + name: agent-config +data: + ENABLE_OTEL: "true" + OTEL_EXPORTER_OTLP_ENDPOINT: "otel-gateway:4327" + ENABLE_OTEL_TRACES: "true" + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: "jaeger-collector.observability.svc:4317" + OTEL_AUTH_TOKEN: "" +``` + +Point endpoints at your cluster's OTLP ingress or managed observability backend. No collector manifests ship with this template. + +### Verify OTEL Setup + +```bash +oc logs deployment/agent | grep -i otel +curl -s http:///health | jq .checks.otel +``` + +The health check reports initialization status, enabled flag, endpoint, and SDK version. + +## Available Metrics + +Metric names use a **dynamic prefix** derived from the agent display name in config (for example `health_assistant_conversations_total` for agent name "Health Assistant"). + +### Conversations + +- `{prefix}_conversations_total{status}` — conversations by status +- `{prefix}_active_conversations` — currently active conversations +- `{prefix}_conversation_duration_seconds` — conversation duration + +### Messages + +- `{prefix}_messages_total{direction,message_type}` — messages sent/received + +### Streaming + +- `{prefix}_stream_tokens_total` — tokens streamed +- `{prefix}_stream_duration_seconds` — stream duration +- `{prefix}_stream_errors_total{error_type}` — stream failures +- `{prefix}_time_to_first_token_seconds` — time to first token + +### Threads + +- `{prefix}_threads_created_total` — threads created +- `{prefix}_threads_active` — active threads +- `{prefix}_threads_deleted_total` — threads deleted +- `{prefix}_thread_messages_count` — messages per thread + +### Graph builds + +- `{prefix}_graph_build_duration_seconds{cache_hit,mcp_tool_count,...}` — graph compilation timing (wired in `graph.py`) + +## Instrumentation Status + +The OTEL module defines `record_*` helpers for lifecycle metrics. Most are **not yet wired** to runtime handlers. + +### Working now + +- OTLP export when `ENABLE_OTEL=true` +- FastAPI distributed tracing (HTTP spans, W3C trace context) +- Graph build metrics (`record_graph_built()` in `graph.py`) +- In-memory metric snapshot for debugging +- Health check OTEL status (`/health` → `checks.otel`) +- Token budget OTEL export (separate flags — see `ENABLE_OTEL_METRICS`) + +### Requires wiring + +Add instrumentation calls at: + +1. **Conversations** — `record_conversation_started()` / `record_conversation_completed()` +2. **Messages** — `record_message_sent()` +3. **Streams** — `record_stream_started()`, `record_first_token()`, `record_stream_completed()`, `record_stream_error()` +4. **Threads** — `record_thread_created()`, `record_thread_deleted()`, `record_thread_messages()` + +See `deep_agent/aegra/otel.py` for the full API. + +### Example + +```python +from deep_agent.aegra.otel import record_conversation_started, record_conversation_completed + +async def handle_conversation(thread_id: str): + start_mono = record_conversation_started(attributes={"thread_id": thread_id}) + try: + # ... conversation logic ... + record_conversation_completed(start_mono, status="completed") + except Exception: + record_conversation_completed(start_mono, status="error") + raise +``` + +## Configuration Reference + +### Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `ENABLE_OTEL` | `false` | Enable agent lifecycle OTEL export (`aegra/otel.py`) | +| `OTEL_EXPORTER_OTLP_ENDPOINT` | `http://localhost:4317` | OTLP gRPC endpoint for lifecycle metrics/traces | +| `OTEL_EXPORTER_OTLP_INSECURE` | `true` | Disable TLS for OTLP connection | +| `OTEL_METRIC_EXPORT_INTERVAL` | `5000` | Metric export interval in ms (1000–60000) | +| `ENABLE_OTEL_METRICS` | `false` | Enable token budget metrics export | +| `ENABLE_OTEL_TRACES` | `false` | Enable trace export via `otel_setup.py` | +| `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` | `""` | Separate traces endpoint (defaults to Jaeger in dev) | +| `OTEL_SERVICE_NAME` | `template-agent` | Service name for token budget OTEL resource | +| `OTEL_AUTH_TOKEN` | `""` | Bearer token for authenticated OTLP endpoints | +| `APPLICATION_VERSION` | — | Overrides `service.version` resource attribute | + +### YAML Configuration + +File: `config/agent/runtime/observability.yaml` + +```yaml +otel: + enabled: false + exporter: + endpoint: "http://localhost:4317" + insecure: true + metrics: + export_interval_ms: 5000 + tracing: + fastapi_auto_instrument: true +``` + +Environment variables override YAML values. + +## Troubleshooting + +### Metrics report zero + +Expected until `record_*` helpers are wired into runtime handlers. Graph build metrics should increment on agent graph compilation. + +### Traces not in Jaeger + +1. Confirm `ENABLE_OTEL=true` or `ENABLE_OTEL_TRACES=true` with a valid traces endpoint +2. Check agent logs for `OTEL enabled` or `template_agent_otel_tracing_enabled` +3. Verify Jaeger OTLP receiver: `COLLECTOR_OTLP_ENABLED=true` (set in compose) +4. Check `/health` → `checks.otel` for initialization status + +### OTLP connection failures + +```bash +# From agent container/network namespace +curl -v telnet://jaeger:4317 +``` + +Verify endpoint hostname matches compose service name (`jaeger`, not `localhost`, from inside the agent container). diff --git a/docs/superpowers/specs/2026-06-23-headless-agent-mode-design.md b/docs/superpowers/specs/2026-06-23-headless-agent-mode-design.md new file mode 100644 index 00000000..b0fb445e --- /dev/null +++ b/docs/superpowers/specs/2026-06-23-headless-agent-mode-design.md @@ -0,0 +1,298 @@ +# Headless Agent Mode + Event-Driven Triggers + +## Summary + +Add a headless agent mode (`mode: headless` in `agent.yaml`) that disables the Aegra HTTP server and runs the agent as a background worker. An `EventTriggerMiddleware` consumes events from three trigger sources (webhook, cron, queue consumer) and invokes the agent graph per event. Results are fanned out to configurable output sinks. All components are tested with unit tests, integration tests, and a full headless startup test. + +## Configuration + +New top-level keys in `config/agent/runtime/agent.yaml`: + +```yaml +mode: server # server (default, full Aegra) | headless + +triggers: + webhook: + enabled: false + host: "0.0.0.0" + port: 8888 + path: "/trigger" + cron: + enabled: false + jobs: + - name: "daily-report" + schedule: "0 9 * * *" + payload: + task: "generate_daily_report" + queue: + enabled: false + backend: "redis_streams" + stream: "agent-tasks" + consumer_group: "agent-workers" + consumer_name: "worker-1" + +output_sinks: + - type: stdout + - type: file + path: "/var/log/agent/output.jsonl" + - type: webhook + url: "https://downstream.example.com/results" + headers: + Authorization: "Bearer ${WEBHOOK_TOKEN}" + - type: redis + stream: "agent-results" +``` + +### Behavior + +- `mode: server` — default. Full Aegra API starts. Current behavior, nothing changes. +- `mode: headless` — Aegra HTTP server does not start. `EventTriggerMiddleware` takes over as the runtime. +- `triggers:` — only meaningful when `mode: headless`. Each trigger type is independently enabled. +- `output_sinks:` — list of sinks. All enabled sinks receive every output (fan-out). If empty, defaults to stdout. + +### Pydantic Models + +- `WebhookTriggerConfig` — host, port, path +- `CronJobConfig` — name, schedule, payload +- `CronTriggerConfig` — enabled, jobs list +- `QueueTriggerConfig` — enabled, backend, stream, consumer_group, consumer_name +- `TriggerConfig` — webhook, cron, queue +- `OutputSinkConfig` — type, plus type-specific fields (path, url, headers, stream) +- `HeadlessConfig` — mode, triggers, output_sinks (top-level container) + +## Component Architecture + +``` +deep_agent/src/triggers/ +├── __init__.py +├── config.py # Pydantic models for triggers + sinks +├── middleware.py # EventTriggerMiddleware (orchestrates lifecycle) +├── sources/ +│ ├── __init__.py +│ ├── protocol.py # TriggerSource protocol (async iterator → TriggerEvent) +│ ├── webhook.py # Minimal HTTP listener trigger +│ ├── cron.py # APScheduler-based trigger +│ └── queue.py # QueueConsumer protocol + Redis Streams implementation +├── sinks/ +│ ├── __init__.py +│ ├── protocol.py # OutputSink protocol +│ ├── stdout.py +│ ├── file.py +│ ├── webhook.py # POST results to URL +│ └── redis.py # Publish to Redis Stream +└── runtime.py # HeadlessRuntime (adapts graph factory) +``` + +### TriggerSource Protocol + +```python +class TriggerSource(Protocol): + async def start(self) -> None: ... + async def stop(self) -> None: ... + def __aiter__(self) -> AsyncIterator[TriggerEvent]: ... +``` + +Each source yields `TriggerEvent` objects. Sources run as independent async tasks. + +### TriggerEvent + +```python +@dataclass +class TriggerEvent: + name: str + payload: dict + source: str # "webhook" | "cron" | "queue" + metadata: dict + timestamp: datetime +``` + +### TriggerResult + +```python +@dataclass +class TriggerResult: + event: TriggerEvent + output: Any + duration_ms: float + success: bool + error: str | None +``` + +### EventTriggerMiddleware + +Owns the full lifecycle: + +- **start()** — reads config, instantiates enabled trigger sources and output sinks, starts all sources +- **run()** — main loop consuming TriggerEvents from all sources via `asyncio.TaskGroup`, invokes agent graph per event, fans out TriggerResults to all sinks +- **stop()** — stops accepting new events, waits for in-flight invocations (configurable drain timeout), flushes all sinks + +### QueueConsumer Protocol + +```python +class QueueConsumer(Protocol): + async def consume(self) -> AsyncIterator[QueueMessage]: ... + async def ack(self, message: QueueMessage) -> None: ... + async def close(self) -> None: ... +``` + +Ships with `RedisStreamsConsumer` as default implementation. Users implement this protocol for Kafka/RabbitMQ/SQS. + +### OutputSink Protocol + +```python +class OutputSink(Protocol): + async def emit(self, result: TriggerResult) -> None: ... + async def close(self) -> None: ... +``` + +Four implementations: `StdoutSink`, `FileSink`, `WebhookSink`, `RedisSink`. + +### HeadlessRuntime + +Implements the same interface shape as Aegra's `ServerRuntime` without SSO. Constructed with a configurable identity (service account or anonymous). Passed to the existing `graph.py:agent()` factory unchanged. + +Switching to direct `ServerRuntime` construction later (option 1) is ~30 lines of change — delete `HeadlessRuntime` class and construct `ServerRuntime` with synthetic values. + +## Data Flow + +### Startup + +1. `python -m deep_agent.headless` — new entry point +2. Loads `agent.yaml`, checks `mode: headless` +3. Runs same `check_prerequisites()` as Aegra (DB, model provider, Redis) +4. Creates `HeadlessRuntime` (no SSO, service account identity) +5. Calls `graph.py:agent(headless_runtime)` once to get compiled graph +6. `EventTriggerMiddleware.start()` — instantiates and starts all enabled trigger sources + sinks + +### Main Loop + +``` +┌─────────────┐ ┌─────────────┐ ┌─────────────┐ +│ Webhook │ │ Cron │ │ Queue │ +│ Listener │ │ Scheduler │ │ Consumer │ +└──────┬──────┘ └──────┬──────┘ └──────┬──────┘ + │ │ │ + └───────────┬───────┴───────────────────┘ + ▼ + TriggerEvent stream + │ + ▼ + EventTriggerMiddleware + (asyncio.TaskGroup) + │ + ▼ + graph.ainvoke(event.payload) + │ + ▼ + TriggerResult + │ + ┌───────┼───────┬───────┐ + ▼ ▼ ▼ ▼ + stdout file webhook redis +``` + +- Each trigger source runs as an independent async task +- Events are consumed as they arrive — no batching +- Graph invocation is per-event (one at a time by default) +- All enabled sinks receive every result (fan-out) +- Errors in graph invocation are caught, logged, and emitted as failed TriggerResults — the loop continues + +### Shutdown + +1. Signal received (SIGTERM/SIGINT) +2. `EventTriggerMiddleware.stop()` — stops accepting new events +3. Waits for in-flight graph invocations to complete (configurable drain timeout) +4. Flushes all output sinks +5. Runs same cleanup as Aegra (Langfuse flush, Redis close, cache clear) + +### Error Handling + +- **Trigger source failure** (e.g., Redis disconnects) — logged, reconnect with backoff, other sources unaffected +- **Graph invocation failure** — `TriggerResult(success=False, error=...)` sent to all sinks +- **Sink failure** — logged, other sinks unaffected (one sink failing doesn't block others) + +## Testing Strategy + +### Unit Tests (`tests/unit/triggers/`) + +| Test file | Covers | +|---|---| +| `test_config.py` | Pydantic model validation — trigger configs, sink configs, defaults, invalid values | +| `test_middleware.py` | EventTriggerMiddleware lifecycle — start/stop, event consumption, fan-out to sinks, error handling | +| `test_webhook_source.py` | Webhook listener — starts/stops HTTP, yields TriggerEvent from POST body | +| `test_cron_source.py` | Cron source — schedules jobs, fires TriggerEvent on schedule, cancellation | +| `test_queue_source.py` | RedisStreamsConsumer + QueueConsumer protocol compliance | +| `test_stdout_sink.py` | JSON to stdout | +| `test_file_sink.py` | JSONL append, directory creation, flush | +| `test_webhook_sink.py` | POST result, HTTP error handling | +| `test_redis_sink.py` | Publish to stream | +| `test_runtime.py` | HeadlessRuntime construction, graph factory acceptance | + +All unit tests are mock-based — no real services needed. + +### Integration Tests (`tests/integration/triggers/`) + +| Test file | Covers | +|---|---| +| `test_redis_streams.py` | Real Redis Streams — produce message, consumer picks it up, ack works, consumer group behavior, reconnect after disconnect | +| `test_webhook_listener.py` | Real HTTP POST to webhook listener, verify TriggerEvent arrives in middleware | +| `test_redis_sink_integration.py` | Real Redis — emit TriggerResult, read it back from stream | +| `test_end_to_end.py` | Full pipeline: push event to Redis Stream → EventTriggerMiddleware consumes → graph invoked (mocked) → result appears in output sink (real Redis or real file) | + +### Headless Startup Test (`tests/integration/triggers/test_headless_startup.py`) + +- Starts full headless process (`python -m deep_agent.headless`) as a subprocess +- Verifies prerequisites check runs (DB, model provider) +- Verifies trigger sources start (cron scheduled, queue consumer connected, webhook listener bound) +- Sends a test event via webhook POST and via Redis Stream push +- Asserts output appears in configured sink +- Sends SIGTERM, verifies graceful shutdown (in-flight drained, sinks flushed, process exits 0) + +### Infrastructure + +- Redis: same docker-compose Redis service already in stack (`make dev` brings it up) +- Pytest marker: `@pytest.mark.integration` — skipped by `make test`, included by `make test-all` +- Fixture: `redis_client` that flushes test keys before/after each test +- Fixture: `headless_process` that starts/stops the headless worker with a test config + +### New Makefile Targets + +```makefile +test-triggers: ## Unit tests for triggers only + pytest tests/unit/triggers/ -v + +test-integration: ## Integration tests (requires Redis + DB) + pytest tests/integration/ -m integration -v + +test-headless: ## Full headless startup test + pytest tests/integration/triggers/test_headless_startup.py -v +``` + +## Files Changed + +### New Files + +- `deep_agent/src/triggers/__init__.py` +- `deep_agent/src/triggers/config.py` +- `deep_agent/src/triggers/middleware.py` +- `deep_agent/src/triggers/runtime.py` +- `deep_agent/src/triggers/sources/__init__.py` +- `deep_agent/src/triggers/sources/protocol.py` +- `deep_agent/src/triggers/sources/webhook.py` +- `deep_agent/src/triggers/sources/cron.py` +- `deep_agent/src/triggers/sources/queue.py` +- `deep_agent/src/triggers/sinks/__init__.py` +- `deep_agent/src/triggers/sinks/protocol.py` +- `deep_agent/src/triggers/sinks/stdout.py` +- `deep_agent/src/triggers/sinks/file.py` +- `deep_agent/src/triggers/sinks/webhook.py` +- `deep_agent/src/triggers/sinks/redis.py` +- `deep_agent/headless.py` — headless worker entry point (`python -m deep_agent.headless`) +- `tests/unit/triggers/` — 10 test files +- `tests/integration/triggers/` — 5 test files + +### Modified Files + +- `config/agent/runtime/agent.yaml` — add `mode`, `triggers`, `output_sinks` sections +- `Makefile` — add `test-triggers`, `test-integration`, `test-headless` targets +- `pyproject.toml` — add `integration` pytest marker diff --git a/entrypoint.sh b/entrypoint.sh new file mode 100755 index 00000000..befe3ece --- /dev/null +++ b/entrypoint.sh @@ -0,0 +1,50 @@ +#!/bin/bash + +# Custom CA support: mount a PEM file or provide a URL. +# CUSTOM_CA_PATH — path to a mounted PEM file (preferred, no network call) +# CUSTOM_CA_URL — URL to download PEM from (fallback, hits network per pod) +CA_PEM="" + +if [ -n "$CUSTOM_CA_PATH" ] && [ -f "$CUSTOM_CA_PATH" ]; then + CA_PEM="$CUSTOM_CA_PATH" +elif [ -n "$CUSTOM_CA_URL" ]; then + if curl -so /tmp/custom-ca.pem "$CUSTOM_CA_URL"; then + CA_PEM="/tmp/custom-ca.pem" + echo "INFO: Successfully fetched CA from $CUSTOM_CA_URL" >&2 + else + echo "WARN: Failed to fetch CA from $CUSTOM_CA_URL, continuing with defaults" >&2 + fi +fi + +if [ -n "$CA_PEM" ]; then + # Use /app (user-writable) instead of /tmp to avoid permission issues with shell redirection + BUNDLE_PATH="/app/.ca-bundle.pem" + + # Start with system CA bundle + if command -v python3 &>/dev/null && python3 -m certifi &>/dev/null; then + cp "$(python3 -m certifi)" "$BUNDLE_PATH" + elif [ -f /etc/ssl/certs/ca-certificates.crt ]; then + cp /etc/ssl/certs/ca-certificates.crt "$BUNDLE_PATH" + elif [ -f /etc/pki/tls/certs/ca-bundle.crt ]; then + cp /etc/pki/tls/certs/ca-bundle.crt "$BUNDLE_PATH" + else + touch "$BUNDLE_PATH" + fi + + # Make bundle writable (system CA bundles are often read-only) + chmod u+w "$BUNDLE_PATH" + + # Append custom CA certificate to the bundle + cat "$CA_PEM" >> "$BUNDLE_PATH" 2>/dev/null || cat "$CA_PEM" | cat >> "$BUNDLE_PATH" + [ "$CA_PEM" = "/tmp/custom-ca.pem" ] && rm -f /tmp/custom-ca.pem + + export REQUESTS_CA_BUNDLE="$BUNDLE_PATH" + export SSL_CERT_FILE="$BUNDLE_PATH" + export CURL_CA_BUNDLE="$BUNDLE_PATH" + export PIP_CERT="$BUNDLE_PATH" + export NODE_EXTRA_CA_CERTS="$BUNDLE_PATH" + + echo "INFO: Custom CA bundle configured at $BUNDLE_PATH" >&2 +fi + +exec "$@" diff --git a/examples/README.md b/examples/README.md deleted file mode 100644 index 407034d9..00000000 --- a/examples/README.md +++ /dev/null @@ -1,220 +0,0 @@ -# Template Agent Client Examples - -This directory contains client examples demonstrating how to interact with the Template Agent's simplified streaming API. These examples show best practices for handling real-time streaming, different event types, and error scenarios. - -## 📁 Available Examples - -### 1. Streamlit Demo App (`streamlit_app.py`) - -A full-featured chat application built with Streamlit: -- **Real-time chat interface** with message history -- **Token streaming visualization** for responsive UX -- **Session management** with thread and session persistence -- **Configuration panel** for API settings and debugging -- **Export functionality** for conversation data - -**Key Features:** -- Live token streaming with visual updates -- Tool call visualization with expandable details -- API health monitoring -- Conversation export to JSON -- Session state management - -**To Run:** -```bash -# Install Streamlit if not already installed -pip install streamlit requests - -# Run the app -streamlit run examples/streamlit_app.py - -# Open http://localhost:8501 in your browser -``` - -### 2. Python Async Client (`client_python.py`) - -A robust async Python client for server-to-server communication: -- **Async/await support** using aiohttp -- **Streaming and non-streaming modes** for different use cases -- **Comprehensive error handling** with detailed error messages -- **Session management** with automatic ID generation -- **Health checking** for API availability - -**Key Features:** -- Generator-based streaming for memory efficiency -- Automatic session ID generation -- Built-in retry logic and timeout handling -- Example conversation flows - -**To Run:** -```bash -# Install dependencies -pip install aiohttp - -# Run the example -python examples/client_python.py -``` - -**Usage as Library:** -```python -from examples.client_python import TemplateAgentClient - -client = TemplateAgentClient() - -# Simple message -response, messages = await client.send_message("Hello!") - -# Streaming chat -async for event in client.stream_chat("Hello!", "thread-123", "session-123", "user-123"): - if event['type'] == 'token': - print(event['content'], end='', flush=True) -``` - -## 🔗 API Reference - -### Request Format - -All clients use the simplified request format: - -```json -{ - "message": "User's input message", - "thread_id": "Conversation thread identifier", - "session_id": "Session identifier", - "user_id": "User identifier", - "stream_tokens": true -} -``` - -### Response Format - -The API returns Server-Sent Events with this format: - -```json -{"type": "message", "content": {"type": "ai", "content": "Hello"}} -{"type": "token", "content": " world"} -{"type": "error", "content": {"message": "Error occurred", "recoverable": false}} -[DONE] -``` - -**Event Types:** -- `message` - Complete messages (AI responses, tool calls, tool results) -- `token` - Individual tokens for real-time streaming -- `error` - Error messages with recovery information -- `[DONE]` - Stream completion marker - -## 🚀 Getting Started - -### Prerequisites - -1. **Template Agent Server Running** - ```bash - # Start the Template Agent server - cd template-agent - python -m uvicorn template_agent.src.main:app --reload --port 8081 - ``` - -2. **Install Client Dependencies** - ```bash - # For Python examples - pip install aiohttp requests streamlit - - # For TypeScript example - npm install # (if using in a Node.js project) - ``` - -### Quick Test - -Test the API is working: - -```bash -# Health check -curl http://localhost:8081/health - -# Simple streaming test -curl -X POST 'http://localhost:8081/stream' \ - -H 'Content-Type: application/json' \ - -H 'Accept: text/event-stream' \ - -d '{ - "message": "Hello!", - "thread_id": "test-123", - "session_id": "test-123", - "user_id": "test-user", - "stream_tokens": true - }' -``` - -## 🎯 Best Practices - -### 1. Session Management -- Use consistent `thread_id` for multi-turn conversations -- Use `session_id` to group related threads -- Generate UUIDs for unique identifiers - -### 2. Error Handling -- Always handle `error` events in streams -- Check `recoverable` flag to determine retry logic -- Implement timeout and connection error handling - -### 3. Token Streaming -- Set `stream_tokens: true` for real-time UX -- Set `stream_tokens: false` for simpler message-only handling -- Buffer tokens appropriately for UI updates - -### 4. Performance -- Use appropriate timeouts for your use case -- Handle stream interruption gracefully -- Consider connection pooling for high-volume usage - -## 🔧 Enterprise Features - -All examples preserve enterprise features from the original implementation: - -- **SSO Authentication**: Pass `X-Token` header for enterprise auth -- **Langfuse Tracing**: Automatic tracing and analytics -- **PostgreSQL Persistence**: Conversation history and checkpointing -- **Error Monitoring**: Comprehensive error logging and recovery - -## 📚 Additional Resources - -- [Template Agent API Documentation](../README.md) -- [FastAPI Documentation](https://fastapi.tiangolo.com/) -- [Streamlit Documentation](https://docs.streamlit.io/) -- [LangGraph Documentation](https://python.langchain.com/docs/langgraph) - -## 🐛 Troubleshooting - -### Common Issues - -**Connection Refused** -- Ensure Template Agent server is running on http://localhost:8081 -- Check firewall settings and port availability - -**Authentication Errors** -- Verify SSO token is valid (if using enterprise features) -- Check X-Token header format - -**Streaming Issues** -- Ensure `Accept: text/event-stream` header is set -- Check for proxy/firewall interference with streaming -- Verify timeout settings are appropriate - -**Token Streaming Not Working** -- Confirm `stream_tokens: true` in request -- Check for buffering issues in HTTP clients -- Verify WebSocket/EventSource compatibility - -### Debug Mode - -Enable detailed logging in examples: - -```python -# Python examples -import logging -logging.basicConfig(level=logging.DEBUG) - -# Streamlit -st.set_option('client.showErrorDetails', True) -``` - -For more help, check the main project documentation or create an issue in the repository. diff --git a/examples/client_python.py b/examples/client_python.py deleted file mode 100644 index 793831ae..00000000 --- a/examples/client_python.py +++ /dev/null @@ -1,314 +0,0 @@ -"""Python client example for Template Agent simplified streaming API. - -This module provides a simple Python client for interacting with the -Template Agent's streaming API, demonstrating how to handle real-time -responses and different event types. - -Usage: - python examples/client_python.py - - Or use as a library: - from examples.client_python import TemplateAgentClient - - client = TemplateAgentClient() - await client.stream_chat("Hello, world!", "thread-123", "session-123", "user-123") -""" - -import asyncio -import json -import uuid -from typing import Any, AsyncGenerator, Dict, Optional - -import aiohttp - - -class TemplateAgentClient: - """Async Python client for Template Agent streaming API.""" - - def __init__( - self, - base_url: str = "http://localhost:8081", - headers: Optional[Dict[str, str]] = None, - ): - """Initialize the client. - - Args: - base_url: Base URL of the Template Agent API - headers: Optional additional headers (e.g., for authentication) - """ - self.base_url = base_url.rstrip("/") - self.headers = { - "Content-Type": "application/json", - "Accept": "text/event-stream", - **(headers or {}), - } - - async def stream_chat( - self, - message: str, - thread_id: str, - session_id: str, - user_id: str, - stream_tokens: bool = True, - timeout: int = 60, - ) -> AsyncGenerator[Dict[str, Any], None]: - """Stream a chat conversation with the agent. - - Args: - message: User's input message - thread_id: Conversation thread identifier - session_id: Session identifier - user_id: User identifier - stream_tokens: Whether to stream individual tokens - timeout: Request timeout in seconds - - Yields: - Event dictionaries with 'type' and 'content' fields - """ - request_data = { - "message": message, - "thread_id": thread_id, - "session_id": session_id, - "user_id": user_id, - "stream_tokens": stream_tokens, - } - - timeout_config = aiohttp.ClientTimeout(total=timeout) - - async with aiohttp.ClientSession(timeout=timeout_config) as session: - async with session.post( - f"{self.base_url}/v1/stream", json=request_data, headers=self.headers - ) as response: - if response.status != 200: - error_text = await response.text() - raise Exception(f"HTTP {response.status}: {error_text}") - - # Stream the response line by line - async for line in response.content: - line_str = line.decode("utf-8").strip() - - if not line_str: - continue - - # Check for completion marker - if line_str == "[DONE]": - break - - try: - event = json.loads(line_str) - yield event - except json.JSONDecodeError: - # Skip invalid JSON lines - continue - - async def send_message( - self, - message: str, - thread_id: Optional[str] = None, - session_id: Optional[str] = None, - user_id: str = "python_client", - stream_tokens: bool = True, - ) -> tuple[str, list[Dict[str, Any]]]: - """Send a message and return the complete response. - - Args: - message: User's input message - thread_id: Optional thread ID (generated if not provided) - session_id: Optional session ID (uses thread_id if not provided) - user_id: User identifier - stream_tokens: Whether to stream individual tokens - - Returns: - Tuple of (final_response_text, all_messages) - """ - # Generate IDs if not provided - if thread_id is None: - thread_id = str(uuid.uuid4()) - if session_id is None: - session_id = thread_id - - full_response = "" - all_messages = [] - - async for event in self.stream_chat( - message, thread_id, session_id, user_id, stream_tokens - ): - event_type = event.get("type") - content = event.get("content") - - if event_type == "token" and isinstance(content, str): - # Accumulate tokens - full_response += content - - elif event_type == "message" and isinstance(content, dict): - # Store complete messages - all_messages.append(content) - - # If this is the final AI message, use it as the response - if content.get("type") == "ai" and content.get("content"): - if not full_response: # Use message content if no tokens received - full_response = content["content"] - - elif event_type == "error": - error_msg = ( - content.get("message", "Unknown error") - if isinstance(content, dict) - else str(content) - ) - raise Exception(f"Agent error: {error_msg}") - - return full_response, all_messages - - async def check_health(self) -> Dict[str, Any]: - """Check if the API is healthy.""" - async with aiohttp.ClientSession() as session: - async with session.get(f"{self.base_url}/health") as response: - if response.status == 200: - return await response.json() - else: - raise Exception(f"Health check failed: HTTP {response.status}") - - -async def example_streaming_chat(): - """Example of streaming chat with token updates.""" - print("🤖 Template Agent - Python Client Example") - print("=" * 50) - - client = TemplateAgentClient() - - # Check if API is available - try: - health = await client.check_health() - print(f"✅ API Status: {health.get('status', 'unknown')}") - except Exception as e: - print(f"❌ API Health Check Failed: {e}") - return - - # Generate session IDs - thread_id = str(uuid.uuid4()) - session_id = str(uuid.uuid4()) - user_id = "python_example_user" - - print("\n📱 Session Info:") - print(f"Thread ID: {thread_id}") - print(f"Session ID: {session_id}") - print(f"User ID: {user_id}") - - # Example conversation - messages = [ - "Hello! Can you help me with some math?", - "What is 15 * 24?", - "Can you explain how you calculated that?", - ] - - for i, message in enumerate(messages, 1): - print(f"\n{'=' * 50}") - print(f"Message {i}: {message}") - print(f"{'=' * 50}") - - print("\n🔄 Streaming Response:") - full_response = "" - message_count = 0 - - try: - async for event in client.stream_chat( - message=message, - thread_id=thread_id, - session_id=session_id, - user_id=user_id, - stream_tokens=True, - ): - event_type = event.get("type") - content = event.get("content") - - if event_type == "token": - # Print tokens in real-time - print(content, end="", flush=True) - full_response += content - - elif event_type == "message": - message_count += 1 - msg_type = ( - content.get("type", "unknown") - if isinstance(content, dict) - else "unknown" - ) - - # Print message info - if msg_type == "tool": - tool_id = content.get("tool_call_id", "unknown") - tool_content = content.get("content", "") - print(f"\n🔧 Tool Result [{tool_id}]: {tool_content}") - elif msg_type == "ai" and content.get("tool_calls"): - tool_calls = content.get("tool_calls", []) - print(f"\n🔧 Tool Calls: {len(tool_calls)} tools invoked") - for tool_call in tool_calls: - print( - f" - {tool_call.get('name', 'unknown')}: {tool_call.get('args', {})}" - ) - - elif event_type == "error": - error_msg = ( - content.get("message", "Unknown error") - if isinstance(content, dict) - else str(content) - ) - print(f"\n❌ Error: {error_msg}") - - except Exception as e: - print(f"\n❌ Stream Error: {e}") - continue - - print("\n\n📊 Summary:") - print(f" - Final response length: {len(full_response)} characters") - print(f" - Messages received: {message_count}") - - # Wait before next message - if i < len(messages): - print("\n⏳ Waiting 2 seconds before next message...") - await asyncio.sleep(2) - - -async def example_simple_chat(): - """Example of simple chat without streaming tokens.""" - print("\n🔹 Simple Chat Example (No Token Streaming)") - print("=" * 50) - - client = TemplateAgentClient() - - try: - response, messages = await client.send_message( - "What's the weather like in a general sense?", stream_tokens=False - ) - - print(f"📝 Response: {response}") - print(f"📊 Total messages: {len(messages)}") - - for i, msg in enumerate(messages): - msg_type = msg.get("type", "unknown") - content = msg.get("content", "") - print( - f" {i + 1}. [{msg_type}] {content[:100]}{'...' if len(content) > 100 else ''}" - ) - - except Exception as e: - print(f"❌ Error: {e}") - - -async def main(): - """Run all examples.""" - try: - # Run streaming example - await example_streaming_chat() - - # Run simple example - await example_simple_chat() - - except KeyboardInterrupt: - print("\n\n👋 Goodbye!") - except Exception as e: - print(f"\n❌ Unexpected error: {e}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/examples/streamlit_app.py b/examples/streamlit_app.py deleted file mode 100644 index 5770014a..00000000 --- a/examples/streamlit_app.py +++ /dev/null @@ -1,301 +0,0 @@ -"""Streamlit Demo App for Template Agent. - -This application demonstrates how to integrate with the Template Agent's -simplified streaming API in a Streamlit application. It provides a clean -chat interface with real-time token streaming and message handling. - -To run this app: - streamlit run examples/streamlit_app.py - -Make sure the Template Agent server is running on http://localhost:8081 -""" - -import json -import uuid -from typing import Any, Dict, List - -import requests -import streamlit as st - - -def initialize_session_state(): - """Initialize Streamlit session state variables.""" - if "messages" not in st.session_state: - st.session_state.messages = [] - - if "thread_id" not in st.session_state: - st.session_state.thread_id = str(uuid.uuid4()) - - if "session_id" not in st.session_state: - st.session_state.session_id = str(uuid.uuid4()) - - if "user_id" not in st.session_state: - st.session_state.user_id = "streamlit_user" - - -def stream_agent_response( - message: str, - thread_id: str, - session_id: str, - user_id: str, - stream_tokens: bool = True, - api_url: str = "http://localhost:8081", -) -> tuple[str, List[Dict[str, Any]]]: - """Stream response from the Template Agent using the simplified API. - - Args: - message: User's input message - thread_id: Conversation thread identifier - session_id: Session identifier - user_id: User identifier - stream_tokens: Whether to stream individual tokens - api_url: Base URL of the Template Agent API - - Returns: - Tuple of (final_response, all_messages) - """ - # Prepare request data - request_data = { - "message": message, - "thread_id": thread_id, - "session_id": session_id, - "user_id": user_id, - "stream_tokens": stream_tokens, - } - - full_response = "" - all_messages = [] - - try: - # Make streaming request to the simplified API - response = requests.post( - f"{api_url}/v1/stream", - json=request_data, - stream=True, - timeout=60, - headers={"Accept": "text/event-stream"}, - ) - response.raise_for_status() - - # Process the streaming response - for line in response.iter_lines(decode_unicode=True): - if not line.strip(): - continue - - # Check for completion marker - if line.strip() == "[DONE]": - break - - try: - # Parse the event - event = json.loads(line) - event_type = event.get("type") - content = event.get("content") - - if event_type == "token" and isinstance(content, str): - # Accumulate tokens for real-time display - full_response += content - - elif event_type == "message" and isinstance(content, dict): - # Store complete messages - all_messages.append(content) - - # If this is the final AI message, use it as the response - if content.get("type") == "ai" and content.get("content"): - # If we haven't accumulated tokens, use the message content - if not full_response: - full_response = content["content"] - - elif event_type == "error": - st.error(f"Agent Error: {content.get('message', 'Unknown error')}") - break - - except json.JSONDecodeError: - st.warning(f"Failed to parse response line: {line[:100]}...") - continue - - except requests.exceptions.RequestException as e: - st.error(f"Failed to connect to agent: {e}") - return "", [] - - return full_response, all_messages - - -def display_message(message: Dict[str, Any], role: str): - """Display a message in the chat interface.""" - with st.chat_message(role): - content = message.get("content", "") - - # Display the main content - if content: - st.write(content) - - # Display tool calls if present - tool_calls = message.get("tool_calls", []) - if tool_calls: - with st.expander("🔧 Tool Calls", expanded=False): - for i, tool_call in enumerate(tool_calls): - st.json( - { - "tool": tool_call.get("name", "unknown"), - "args": tool_call.get("args", {}), - "id": tool_call.get("id", ""), - } - ) - - # Display metadata if present - metadata = message.get("response_metadata", {}) - if metadata: - with st.expander("📊 Metadata", expanded=False): - st.json(metadata) - - -def main(): - """Main Streamlit application.""" - st.set_page_config(page_title="Template Agent Chat", page_icon="🤖", layout="wide") - - st.title("🤖 Template Agent Chat") - st.markdown("Chat with the Template Agent using the simplified streaming API") - - # Initialize session state - initialize_session_state() - - # Sidebar configuration - with st.sidebar: - st.header("Configuration") - - api_url = st.text_input( - "API URL", - value="http://localhost:8081", - help="Base URL of the Template Agent API", - ) - - stream_tokens = st.checkbox( - "Stream Tokens", - value=True, - help="Enable real-time token streaming for faster response display", - ) - - st.divider() - - # Session information - st.subheader("Session Info") - st.text(f"Thread ID: {st.session_state.thread_id[:8]}...") - st.text(f"Session ID: {st.session_state.session_id[:8]}...") - st.text(f"User ID: {st.session_state.user_id}") - - if st.button("New Conversation"): - st.session_state.messages = [] - st.session_state.thread_id = str(uuid.uuid4()) - st.rerun() - - st.divider() - - # API test - st.subheader("API Status") - try: - health_response = requests.get(f"{api_url}/health", timeout=5) - if health_response.status_code == 200: - st.success("✅ API Connected") - else: - st.error(f"❌ API Error: {health_response.status_code}") - except Exception: - st.error("❌ API Unreachable, error={e}") - - # Main chat interface - st.subheader("Chat") - - # Display chat history - for message in st.session_state.messages: - if message["role"] == "user": - with st.chat_message("user"): - st.write(message["content"]) - else: - # For agent messages, display the structured content - display_message(message["content"], "assistant") - - # Chat input - if prompt := st.chat_input("Ask me anything..."): - # Add user message to chat history - st.session_state.messages.append({"role": "user", "content": prompt}) - - # Display user message - with st.chat_message("user"): - st.write(prompt) - - # Stream agent response - with st.chat_message("assistant"): - response_placeholder = st.empty() - - # Show loading spinner - with st.spinner("Agent is thinking..."): - # Stream the response - full_response, all_messages = stream_agent_response( - message=prompt, - thread_id=st.session_state.thread_id, - session_id=st.session_state.session_id, - user_id=st.session_state.user_id, - stream_tokens=stream_tokens, - api_url=api_url, - ) - - # Display the final response - if full_response: - response_placeholder.write(full_response) - - # Add to chat history - st.session_state.messages.append( - { - "role": "assistant", - "content": { - "type": "ai", - "content": full_response, - "messages": all_messages, # Store all messages for debugging - }, - } - ) - else: - response_placeholder.error("No response received from agent") - - # Advanced features in expander - with st.expander("🔧 Advanced Features", expanded=False): - st.subheader("Raw Session Data") - - col1, col2 = st.columns(2) - - with col1: - st.text("Session State:") - st.json( - { - "thread_id": st.session_state.thread_id, - "session_id": st.session_state.session_id, - "user_id": st.session_state.user_id, - "message_count": len(st.session_state.messages), - } - ) - - with col2: - st.text("Last Message Details:") - if st.session_state.messages: - st.json(st.session_state.messages[-1]) - - # Export conversation - if st.button("Export Conversation"): - conversation_data = { - "thread_id": st.session_state.thread_id, - "session_id": st.session_state.session_id, - "user_id": st.session_state.user_id, - "messages": st.session_state.messages, - "export_timestamp": str(uuid.uuid4()), - } - - st.download_button( - label="Download Conversation JSON", - data=json.dumps(conversation_data, indent=2), - file_name=f"conversation_{st.session_state.thread_id[:8]}.json", - mime="application/json", - ) - - -if __name__ == "__main__": - main() diff --git a/pyproject.toml b/pyproject.toml index cc97d800..30770881 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,181 +3,60 @@ requires = ["hatchling"] build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] -packages = ["template_agent"] +packages = ["deep_agent"] [project] name = "template-agent" -version = "0.1.0" +version = "0.2.0" description = "A template for Model Context Protocol (MCP) server development" readme = "README.md" keywords = ["mcp", "template", "server"] -requires-python = "==3.12.2" +requires-python = ">=3.12.2" dependencies = [ - "aiohappyeyeballs==2.6.1", - "aiohttp==3.12.6", - "aiosignal==1.3.2", - "aiosqlite==0.21.0", - "altair==5.5.0", - "annotated-types==0.7.0", - "anyio==4.9.0", - "asn1crypto==1.5.1", - "attrs==25.3.0", - "authlib==1.6.0", - "backoff==2.2.1", - "blinker==1.9.0", - "boto3==1.38.27", - "botocore==1.38.27", - "cachetools==5.5.2", - "certifi==2025.4.26", - "cffi==1.17.1", - "charset-normalizer==3.4.2", - "click==8.2.1", - "cryptography==45.0.3", - "dataclasses-json==0.6.7", - "distro==1.9.0", - "dnspython==2.7.0", - "exceptiongroup==1.3.0", - "fastapi==0.115.12", - "fastmcp==2.8.1", - "filelock==3.18.0", - "filetype==1.2.0", - "frozenlist==1.6.0", - "gitdb==4.0.12", - "gitpython==3.1.44", - "google-ai-generativelanguage==0.6.18", - "google-api-core==2.24.2", - "google-auth==2.40.2", - "googleapis-common-protos==1.70.0", - "groq==0.26.0", - "grpcio==1.71.0", - "grpcio-status==1.71.0", - "h11==0.16.0", - "httpcore==1.0.9", - "httpx==0.28.1", - "httpx-sse==0.4.0", - "idna==3.10", - "itsdangerous==2.2.0", - "jinja2==3.1.6", - "jmespath==1.0.1", - "joblib==1.5.1", - "jsonpatch==1.33", - "jsonpointer==3.0.0", - "jsonschema==4.24.0", - "jsonschema-specifications==2025.4.1", - "langchain==0.3.25", - "langchain-community==0.3.24", - "langchain-core==0.3.63", - "langchain-google-genai==2.0.11", - "langchain-groq==0.2.5", - "langchain-mcp-adapters==0.1.1", - "langchain-text-splitters==0.3.8", - "langfuse==2.60.5", - "langgraph==0.4.7", - "langgraph-checkpoint==2.0.26", - "langgraph-checkpoint-mongodb==0.1.3", - "langgraph-checkpoint-postgres==2.0.21", - "langgraph-checkpoint-sqlite==2.0.10", - "langgraph-prebuilt==0.2.2", - "langgraph-sdk==0.1.70", - "langsmith==0.3.43", - "markdown-it-py==3.0.0", - "markupsafe==3.0.2", - "marshmallow==3.26.1", - "mcp==1.9.4", - "mdurl==0.1.2", - "motor==3.7.1", - "multidict==6.4.4", - "mypy-extensions==1.1.0", - "narwhals==1.41.0", - "numpy==2.2.6", - "oauthlib==3.2.2", - "openapi-pydantic==0.5.1", - "orjson==3.10.18", - "ormsgpack==1.10.0", - "packaging==24.2", - "pandas==2.2.3", - "pillow==11.2.1", - "platformdirs==4.3.8", - "propcache==0.3.1", - "proto-plus==1.26.1", - "protobuf==5.29.5", - "psycopg==3.2.9", - "psycopg-binary==3.2.9", - "psycopg-pool==3.2.6", - "psycopg2-binary==2.9.10", - "pyarrow==20.0.0", - "pyasn1==0.6.1", - "pyasn1-modules==0.4.2", - "pycparser==2.22", - "pydantic==2.11.5", - "pydantic-core==2.33.2", - "pydantic-settings==2.9.1", - "pydeck==0.9.1", - "pygments==2.19.1", - "pyjwt==2.10.1", - "pymongo==4.11.3", - "pyopenssl==25.1.0", - "python-dateutil==2.9.0.post0", - "python-dotenv==1.1.0", - "python-multipart==0.0.20", - "pytz==2025.2", - "pyyaml==6.0.2", - "referencing==0.36.2", - "requests==2.32.3", - "requests-oauthlib==2.0.0", - "requests-toolbelt==1.0.0", - "resend==2.8.0", - "rich==14.0.0", - "rpds-py==0.25.1", - "rsa==4.9.1", - "s3transfer==0.13.0", - "scikit-learn==1.6.1", - "scipy==1.15.3", - "shellingham==1.5.4", - "six==1.17.0", - "smmap==5.0.2", - "sniffio==1.3.1", - "sortedcontainers==2.4.0", - "sqlalchemy==2.0.41", - "sqlite-vec==0.1.6", - "sse-starlette==2.3.6", - "starlette==0.46.2", - "streamlit==1.45.1", - "tenacity==9.1.2", - "threadpoolctl==3.6.0", - "toml==0.10.2", - "tomlkit==0.13.2", - "tornado==6.5.1", - "typer==0.16.0", - "typing-extensions==4.13.2", - "typing-inspect==0.9.0", - "typing-inspection==0.4.1", - "tzdata==2025.2", - "urllib3==2.4.0", - "uvicorn==0.32.1", - "watchdog==6.0.0", - "websockets==15.0.1", - "wrapt==1.17.2", - "xxhash==3.5.0", - "yarl==1.20.0", - "zstandard==0.23.0", - "structlog>=24.1.0", + "deepagents==0.4.12", + "pydantic==2.12.5", + "pydantic-settings==2.14.2", + "python-dotenv==1.2.2", + "langchain-google-genai==4.2.2", + "langchain-google-vertexai==3.2.2", + "langchain-openai>=0.3.0", + "langchain-mcp-adapters==0.2.2", + "mcp==1.27.2", + "langfuse==4.6.1", + "langgraph-checkpoint-postgres==3.0.5", + "langgraph-sdk>=0.1.51", + "aegra-cli", + "psycopg[binary,pool]==3.3.3", + "psycopg2-binary==2.9.11", + "motor==3.6.0", + "structlog==25.5.0", + "pyyaml==6.0.3", + "PyJWT[crypto]>=2.8.0", + "cryptography>=43.0.0", + "httpx>=0.27.0", + "tenacity>=8.2.0,<10", + "redis>=5.0.0,<7", + "cachetools>=5.5.0,<6", + "apscheduler>=4.0.0a5", + "opentelemetry-api>=1.33.1,<2.0.0", + "opentelemetry-sdk>=1.33.1,<2.0.0", + "opentelemetry-exporter-otlp>=1.33.1,<2.0.0", + "opentelemetry-instrumentation-fastapi>=0.52b1,<1.0.0", + "aiokafka>=0.10.0", ] [project.optional-dependencies] dev = [ - "pytest==8.4.1", - "pytest-asyncio==1.0.0", + "pytest==9.1.1", + "pytest-asyncio==1.4.0", "pytest-cov==6.2.1", + "pytest-mock>=3.14.0", "ruff==0.12.2", "mypy==1.16.1", "pre-commit==4.2.0", - "httpx==0.28.1" + "httpx==0.28.1", ] -[project.scripts] -template-agent = "template_agent.src.main:main" - [project.urls] Homepage = "https://github.com/redhat-data-and-ai/template-agent" Repository = "https://github.com/redhat-data-and-ai/template-agent" @@ -199,9 +78,9 @@ known-first-party = ["src"] [tool.mypy] python_version = "3.12" ignore_missing_imports = true -disallow_untyped_defs = false -disallow_incomplete_defs = false -check_untyped_defs = false +disallow_untyped_defs = true +disallow_incomplete_defs = true +check_untyped_defs = true disallow_untyped_decorators = false no_implicit_optional = true warn_redundant_casts = true @@ -209,9 +88,11 @@ warn_unused_ignores = true warn_no_return = true warn_unreachable = true strict_equality = true +warn_return_any = true [tool.pytest.ini_options] testpaths = ["tests"] +pythonpath = ["."] python_files = ["test_*.py"] python_classes = ["Test*"] python_functions = ["test_*"] @@ -219,10 +100,18 @@ addopts = [ "--strict-markers", "--strict-config", ] -markers = ["asyncio"] +asyncio_mode = "auto" +markers = [ + "asyncio", + "unit: fast isolated unit tests", + "integration: tests requiring external services or multi-component interaction", + "skills: marks tests as skill evaluation tests (deselect with '-m \"not skills\"')", + "e2e: end-to-end tests requiring a running aegra server", + "slow: tests that take more than 30 seconds", +] [tool.coverage.run] -source = ["src"] +source = ["deep_agent"] omit = [ "*/tests/*", "*/test_*", @@ -230,7 +119,11 @@ omit = [ "*/migrations/*", ] +[tool.coverage.html] +directory = "htmlcov" + [tool.coverage.report] +show_missing = true exclude_lines = [ "pragma: no cover", "def __repr__", diff --git a/template_agent/src/api.py b/template_agent/src/api.py deleted file mode 100644 index 62d0dc53..00000000 --- a/template_agent/src/api.py +++ /dev/null @@ -1,195 +0,0 @@ -"""FastAPI server implementation for the template agent. - -This module provides the main FastAPI application setup, including -middleware configuration, route registration, and application lifecycle -management for the template agent service. -""" - -import time -from collections.abc import AsyncGenerator -from contextlib import asynccontextmanager -from typing import Callable - -from fastapi import FastAPI, Request -from fastapi.middleware.cors import CORSMiddleware -from starlette.middleware.base import BaseHTTPMiddleware -from starlette.responses import JSONResponse - -from template_agent.src.core.agent import initialize_database -from template_agent.src.core.exceptions.exceptions import AppException, AppExceptionCode -from template_agent.src.routes.feedback import router as feedback_router -from template_agent.src.routes.health import router as health_router -from template_agent.src.routes.history import router as history_router -from template_agent.src.routes.stream import router as stream_router -from template_agent.src.routes.threads import router as threads_router -from template_agent.src.settings import settings -from template_agent.utils.pylogger import get_python_logger - -logger = get_python_logger(settings.PYTHON_LOG_LEVEL) - - -class RequestLoggingMiddleware(BaseHTTPMiddleware): - """Middleware to log all incoming requests and outgoing responses.""" - - async def dispatch(self, request: Request, call_next: Callable): - """Process and log incoming requests and outgoing responses.""" - if not settings.REQUEST_LOGGING_ENABLED: - return await call_next(request) - - start_time = time.time() - - # Capture request details - request_data = { - "method": request.method, - "path": request.url.path, - "client_ip": request.client.host if request.client else None, - "query_params": dict(request.query_params) - if request.query_params - else None, - } - - # Optionally log headers - if settings.REQUEST_LOG_HEADERS: - request_data["headers"] = dict(request.headers) - - # Optionally log request body - if settings.REQUEST_LOG_BODY: - try: - body_bytes = await request.body() - body_size = len(body_bytes) - - if body_size > 0: - request_data["body_size"] = body_size - if ( - settings.REQUEST_LOG_BODY_MAX_SIZE == 0 - or body_size <= settings.REQUEST_LOG_BODY_MAX_SIZE - ): - try: - body_str = body_bytes.decode("utf-8") - request_data["body"] = body_str - except UnicodeDecodeError: - request_data["body"] = "" - else: - request_data["body"] = f"" - - # Rebuild request with body - async def receive(): - return {"type": "http.request", "body": body_bytes} - - request = Request(request.scope, receive) - except Exception as e: - logger.warning("Failed to read request body", error=str(e)) - - logger.info("incoming_request", **request_data) - - # Process request - response = await call_next(request) - - # Capture response details - duration_ms = (time.time() - start_time) * 1000 - response_data = { - "method": request.method, - "path": request.url.path, - "status_code": response.status_code, - "duration_ms": round(duration_ms, 2), - } - - # Optionally log response headers - if settings.REQUEST_LOG_HEADERS: - response_data["headers"] = dict(response.headers) - - logger.info("outgoing_response", **response_data) - - return response - - -@asynccontextmanager -async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: - """Configure application lifespan. - - This context manager handles the application startup and shutdown - lifecycle. Database schema is initialized on startup, while agent - initialization is deferred to per-request handling to allow for - authenticated MCP connections. - - Args: - app: The FastAPI application instance to manage. - - Yields: - None: The lifespan context for the application. - - Raises: - AppException: If database initialization fails on startup. - """ - logger.info("Agent server starting up") - - # Initialize database schema on startup - try: - await initialize_database() - except Exception as e: - logger.critical(f"Failed to initialize database on startup: {e}") - raise - - logger.info("Agent server ready - MCP connection will be established per-request") - yield - logger.info("Agent server shutting down") - - -# Create FastAPI application with lifespan management -app = FastAPI(lifespan=lifespan) - -# Register request logging middleware first to capture all requests -app.add_middleware(RequestLoggingMiddleware) - -# Configure CORS middleware for cross-origin requests -app.add_middleware( - CORSMiddleware, - allow_origins=["*"], - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], -) - -# Configure application logger -app.logger = get_python_logger(settings.PYTHON_LOG_LEVEL) - -# Register all route handlers -app.include_router(health_router) -app.include_router(stream_router) -app.include_router(feedback_router) -app.include_router(history_router) -app.include_router(threads_router) - - -@app.exception_handler(Exception) -async def generic_exception_handler(request: Request, exc: Exception): - """Generic exception handler for unhandled exceptions.""" - logger.exception( - f"Unhandled exception occurred for request_method={request.method}, request_path={request.url.path}, error={exc}" - ) - logger.debug(f"Unhandled exception occurred for request={request}, error={exc}") - return JSONResponse( - status_code=AppExceptionCode.INTERNAL_SERVER_ERROR.response_code, - content={ - "detail_message": str(exc), - "message": AppExceptionCode.INTERNAL_SERVER_ERROR.message, - "error_code": AppExceptionCode.INTERNAL_SERVER_ERROR.error_code, - }, - ) - - -@app.exception_handler(AppException) -async def app_exception_handler(request: Request, exc: AppException): - """App exception handler for unhandled exceptions.""" - logger.warn( - f"App exception occurred for request_method={request.method}, request_path={request.url.path}, error={exc}" - ) - logger.debug(f"App exception occurred for request={request}, error={exc}") - return JSONResponse( - status_code=exc.response_code, - content={ - "detail_message": exc.detail_message, - "message": exc.message, - "error_code": exc.error_code, - }, - ) diff --git a/template_agent/src/core/__init__.py b/template_agent/src/core/__init__.py deleted file mode 100644 index ad2e0dc9..00000000 --- a/template_agent/src/core/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Core module for template agent functionality.""" diff --git a/template_agent/src/core/agent.py b/template_agent/src/core/agent.py deleted file mode 100644 index cb250556..00000000 --- a/template_agent/src/core/agent.py +++ /dev/null @@ -1,217 +0,0 @@ -"""Agent implementation for the template agent system. - -This module provides the core agent functionality for the template agent, -including initialization, configuration, and agent creation utilities. -""" - -from contextlib import asynccontextmanager -from typing import Optional - -from langchain_google_genai import ChatGoogleGenerativeAI -from langchain_mcp_adapters.client import MultiServerMCPClient -from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver -from langgraph.prebuilt import create_react_agent - -from template_agent.src.core.exceptions.exceptions import AppException, AppExceptionCode -from template_agent.src.core.prompt import get_system_prompt -from template_agent.src.core.storage import get_global_checkpoint -from template_agent.src.settings import settings -from template_agent.utils.pylogger import get_python_logger - -logger = get_python_logger(log_level=settings.PYTHON_LOG_LEVEL) - - -async def initialize_database() -> None: - """Initialize PostgreSQL database schema on application startup. - - This function ensures the checkpoints table and related schema are created - before any requests are processed. Only runs when using PostgreSQL storage - (USE_INMEMORY_SAVER=False). - - Raises: - AppException: If database connection or schema creation fails. - """ - if settings.USE_INMEMORY_SAVER: - logger.info("Using in-memory storage - skipping database initialization") - return - - try: - logger.info("Initializing PostgreSQL database schema") - async with AsyncPostgresSaver.from_conn_string( - settings.database_uri - ) as checkpoint: - # Setup database schema - creates checkpoints table and indexes - if hasattr(checkpoint, "setup"): - await checkpoint.setup() - logger.info("Database schema initialized successfully") - else: - logger.warning( - "AsyncPostgresSaver does not have setup method - schema may need manual creation" - ) - except Exception as e: - logger.error(f"Failed to initialize database schema: {e}", exc_info=True) - raise AppException( - f"Database initialization failed: {str(e)}", - AppExceptionCode.CONFIGURATION_INITIALIZATION_ERROR, - ) - - -@asynccontextmanager -async def get_template_agent( - sso_token: Optional[str] = None, enable_checkpointing: bool = True -): - """Get a fully initialized template agent. - - This function creates and configures a template agent with the necessary - tools, model, and database connections. It uses an async context manager - to ensure proper resource cleanup. - - Args: - sso_token: Optional access token for authentication. If provided, - it will be used for authorization headers in MCP client requests. - enable_checkpointing: Whether to enable checkpointing/persistence. - Set to False for streaming-only operations that shouldn't save to DB. - - Yields: - The initialized template agent instance. - - Raises: - Exception: If there are issues with database connections or agent setup. - """ - # Initialize MCP client and get tools - tools = [] - - # Log MCP connection details for debugging - logger.info(f"Attempting to connect to MCP server at {settings.MCP_SERVER_URL}") - logger.info(f"MCP server name: {settings.MCP_SERVER_NAME}") - logger.info(f"MCP transport protocol: {settings.MCP_TRANSPORT_PROTOCOL}") - logger.info(f"MCP connection timeout: {settings.MCP_CONNECTION_TIMEOUT}s") - logger.info(f"SSO authentication: {'Yes' if sso_token else 'No'}") - - try: - import asyncio - - # Add timeout wrapper for MCP connection - async def connect_with_timeout(): - # Configure MCP client with SSL verification setting - server_config = { - "url": settings.MCP_SERVER_URL, - "transport": settings.MCP_TRANSPORT_PROTOCOL, - "headers": {"Authorization": f"Bearer {sso_token}"} - if sso_token - else {}, - } - - # Add SSL verification setting (verify=False disables cert verification) - if not settings.MCP_SSL_VERIFY: - server_config["verify"] = False - logger.warning( - "SSL certificate verification disabled for MCP connection" - ) - - client = MultiServerMCPClient({settings.MCP_SERVER_NAME: server_config}) - return await client.get_tools() - - tools = await asyncio.wait_for( - connect_with_timeout(), timeout=settings.MCP_CONNECTION_TIMEOUT - ) - logger.info( - f"Successfully connected to MCP server and loaded {len(tools)} tools" - ) - except asyncio.TimeoutError: - # Handle timeout specifically - error_msg = ( - f"Timeout connecting to MCP server at {settings.MCP_SERVER_URL} " - f"after {settings.MCP_CONNECTION_TIMEOUT}s. " - f"Server may be down or unreachable." - ) - logger.error(error_msg) - - if settings.USE_INMEMORY_SAVER: - logger.warning("Running in local development mode without MCP tools") - tools = [] - else: - logger.critical(error_msg) - raise AppException( - error_msg, - AppExceptionCode.PRODUCTION_MCP_CONNECTION_ERROR, - ) - except Exception as e: - # Log detailed error information for other exceptions - logger.error( - f"Failed to connect to MCP server at {settings.MCP_SERVER_URL}", - exc_info=True, - ) - logger.error(f"MCP connection error type: {type(e).__name__}") - logger.error(f"MCP connection error details: {str(e)}") - - if settings.USE_INMEMORY_SAVER: - logger.warning("Running in local development mode without MCP tools") - tools = [] # No tools for local development - else: - # In production, MCP is required - error_msg = ( - f"Failed to connect to required MCP server at {settings.MCP_SERVER_URL}. " - f"Error: {type(e).__name__}: {str(e)}" - ) - logger.critical(error_msg) - raise AppException( - error_msg, - AppExceptionCode.PRODUCTION_MCP_CONNECTION_ERROR, - ) - - # Initialize the language model - model = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.3) - - if not enable_checkpointing: - # Create agent without checkpointing for streaming-only operations - logger.info( - "Creating agent without checkpointing for streaming-only operations" - ) - agent_redhat = create_react_agent( - model=model, - prompt=get_system_prompt(), - tools=tools, - # No checkpointer or store - streaming only, no persistence - ) - logger.info("Template agent initialized successfully without checkpointing") - yield agent_redhat - elif settings.USE_INMEMORY_SAVER: - # Use single global checkpoint for local development - logger.info("Using single global checkpoint for local development") - # Use single checkpoint instance for both checkpointer and store - checkpoint = get_global_checkpoint() - agent_redhat = create_react_agent( - model=model, - prompt=get_system_prompt(), - tools=tools, - checkpointer=checkpoint, - store=checkpoint, - ) - logger.info( - "Template agent initialized successfully with single global checkpoint" - ) - yield agent_redhat - else: - # Use PostgreSQL storage for production - logger.info("Using PostgreSQL checkpoint for production") - async with AsyncPostgresSaver.from_conn_string( - settings.database_uri - ) as checkpoint: - # Setup database connection once - if hasattr(checkpoint, "setup"): - await checkpoint.setup() - - # Create the agent with single checkpoint instance for both checkpointer and store - agent_redhat = create_react_agent( - model=model, - prompt=get_system_prompt(), - tools=tools, - checkpointer=checkpoint, - store=checkpoint, - ) - - logger.info( - "Template agent initialized successfully with PostgreSQL checkpoint" - ) - yield agent_redhat diff --git a/template_agent/src/core/exceptions/__init__.py b/template_agent/src/core/exceptions/__init__.py deleted file mode 100644 index cabfa498..00000000 --- a/template_agent/src/core/exceptions/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -"""Template Agent exception package. - -This package provides a exception handling for this agent. -""" diff --git a/template_agent/src/core/exceptions/exceptions.py b/template_agent/src/core/exceptions/exceptions.py deleted file mode 100644 index 32cc18a8..00000000 --- a/template_agent/src/core/exceptions/exceptions.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Exception handling for the Template MCP server.""" - -from __future__ import annotations - -from enum import Enum - -from starlette.status import ( - HTTP_400_BAD_REQUEST, - HTTP_401_UNAUTHORIZED, - HTTP_403_FORBIDDEN, - HTTP_404_NOT_FOUND, - HTTP_500_INTERNAL_SERVER_ERROR, -) - - -class AppExceptionCode(Enum): - """Defines custom App Exception codes for this service, associated with HTTP Status codes.""" - - BAD_REQUEST_ERROR = (HTTP_400_BAD_REQUEST, "Bad Request", "E_001") - NOT_FOUND_ERROR = (HTTP_404_NOT_FOUND, "Not Found", "E_002") - INTERNAL_SERVER_ERROR = ( - HTTP_500_INTERNAL_SERVER_ERROR, - "Internal Server Error", - "E_003", - ) - UNAUTHORISED_ACCESS_ERROR = (HTTP_401_UNAUTHORIZED, "Unauthorized", "E_004") - FORBIDDEN_ACCESS_ERROR = (HTTP_403_FORBIDDEN, "Forbidden", "E_005") - TOOL_CALL_ERROR = (HTTP_500_INTERNAL_SERVER_ERROR, "Internal Server Error", "E_006") - PRODUCTION_MCP_CONNECTION_ERROR = ( - HTTP_500_INTERNAL_SERVER_ERROR, - "Internal Server Error", - "E_007", - ) - CONFIGURATION_INITIALIZATION_ERROR = ( - HTTP_500_INTERNAL_SERVER_ERROR, - "Internal Server Error", - "E_008", - ) - CONFIGURATION_VALIDATION_ERROR = ( - HTTP_500_INTERNAL_SERVER_ERROR, - "Internal Server Error", - "E_009", - ) - - def __init__(self, response_code: str, message: str, error_code: str): - """Constructor to initialize the exception code with response_code, message, and error_code.""" - self._response_code = response_code - self._message = message - self._error_code = error_code - - @property - def response_code(self): - """HTTP status code for exception code.""" - return self._response_code - - @property - def message(self): - """HTTP status message for exception code.""" - return self._message - - @property - def error_code(self): - """HTTP error_code for exception code.""" - return self._error_code - - def __str__(self): - """Str method for logging the exception code.""" - return f"response_code={self.response_code}, message={self.message}, error_code={self.error_code}" - - -class AppException(Exception): - """Base exception for application.""" - - def __init__( - self, - detail_message: str, - app_exception_code: AppExceptionCode = AppExceptionCode.INTERNAL_SERVER_ERROR, - ): - """Constructor to initialize the exception.""" - self._detail_message = detail_message - self._app_exception_code = app_exception_code - super().__init__(detail_message) - - @property - def detail_message(self): - """Detail error message for exception.""" - return self._detail_message - - @property - def response_code(self): - """HTTP response code for exception.""" - return self._app_exception_code.response_code - - @property - def message(self): - """HTTP message for exception.""" - return self._app_exception_code.message - - @property - def error_code(self): - """Error code for exception.""" - return self._app_exception_code.error_code - - def __str__(self): - """Str method for logging the exception.""" - return f"response_code={self.response_code}, message={self.message}, detail_message={self.detail_message}, error_code={self.error_code}" - - -class ToolCallException(AppException): - """Raised when Tool call fails.""" - - def __init__(self, detail_message: str): - """Constructor to initialize the ToolCallException.""" - super().__init__(detail_message, AppExceptionCode.TOOL_CALL_ERROR) - - def __str__(self): - """Str method for logging the ToolCallException.""" - return super().__str__() - - -class UnauthorizedException(AppException): - """Raised when user Authentication fails.""" - - def __init__(self, detail_message: str): - """Constructor to initialize the UnauthorizedException.""" - super().__init__(detail_message, AppExceptionCode.UNAUTHORISED_ACCESS_ERROR) - - def __str__(self): - """Str method for logging the UnauthorizedException.""" - return super().__str__() - - -class ForbiddenException(AppException): - """Raised when user is forbidden.""" - - def __init__(self, detail_message: str): - """Constructor to initialize the ForbiddenException.""" - super().__init__(detail_message, AppExceptionCode.FORBIDDEN_ACCESS_ERROR) - - def __str__(self): - """Str method for logging the ForbiddenException.""" - return super().__str__() diff --git a/template_agent/src/core/manager.py b/template_agent/src/core/manager.py deleted file mode 100644 index dd5a01fe..00000000 --- a/template_agent/src/core/manager.py +++ /dev/null @@ -1,553 +0,0 @@ -"""Agent Manager for the template agent system. - -This module provides the AgentManager class that orchestrates agent operations, -handles streaming responses, and manages the conversion between LangGraph events -and simplified streaming. -""" - -import inspect -from collections.abc import AsyncGenerator -from typing import Any, Dict -from uuid import uuid4 - -from langchain_core.messages import ( - AIMessage, - AIMessageChunk, - HumanMessage, - ToolMessage, -) -from langchain_core.runnables import RunnableConfig -from langfuse.callback import CallbackHandler -from langgraph.pregel import Pregel -from langgraph.types import Command, Interrupt - -from template_agent.src.core.agent import get_template_agent -from template_agent.src.core.agent_utils import ( - convert_message_content_to_string, - langchain_to_chat_message, - remove_tool_calls, -) -from template_agent.src.core.storage import register_thread -from template_agent.src.schema import StreamRequest -from template_agent.src.settings import settings -from template_agent.utils.pylogger import get_python_logger - -# Initialize Langfuse CallbackHandler for Langchain (tracing) -langfuse_handler = CallbackHandler( - trace_name="template-agent", environment=settings.LANGFUSE_TRACING_ENVIRONMENT -) - -app_logger = get_python_logger(settings.PYTHON_LOG_LEVEL) - - -class AgentManager: - """Manager class for handling agent operations and streaming responses. - - This class provides a simplified interface for agent interactions while - preserving all enterprise features like authentication, tracing, and - error handling from the original implementation. - """ - - def __init__(self, redhat_sso_token: str | None = None): - """Initialize the AgentManager. - - Args: - redhat_sso_token: Optional SSO token for enterprise authentication. - """ - self.redhat_sso_token = redhat_sso_token - self._agent: Pregel | None = None - self._current_tool_call_id: str | None = None # Track current active tool call - - async def stream_response( - self, request: StreamRequest - ) -> AsyncGenerator[Dict[str, Any], None]: - """Stream agent response with simplified event structure. - - This method provides streaming functionality while ensuring that conversation - state is saved only once at the end, not during intermediate streaming. - - Args: - request: The streaming request containing user input and configuration. - - Yields: - Simplified event dictionaries with 'type' and 'content' fields. - """ - # Use persistent agent for both streaming and state persistence - # This ensures LangGraph handles state management automatically - async with get_template_agent( - self.redhat_sso_token, enable_checkpointing=True - ) as persistent_agent: - try: - # Prepare input for the persistent agent - kwargs, run_id, thread_id = await self._handle_input( - request, persistent_agent - ) - - app_logger.info( - f"AgentManager streaming response for run_id: {run_id}, thread_id: {thread_id}" - ) - - # Reset tool call tracking for this stream - self._current_tool_call_id = None - - # Use persistent agent for streaming - LangGraph will handle state automatically - async for stream_event in persistent_agent.astream( - **kwargs, stream_mode=["updates", "messages", "custom"] - ): - if not isinstance(stream_event, tuple): - continue - - stream_mode, event = stream_event - - # Update tool call tracking based on stream events - self._update_tool_call_tracking(stream_mode, event) - - # Convert LangGraph events to simplified format - effective_session_id = request.session_id or thread_id - formatted_events = self._format_events( - stream_mode, - event, - request.stream_tokens, - run_id, - thread_id, - effective_session_id, - ) - - for formatted_event in formatted_events: - if formatted_event: - yield formatted_event - - # No manual state saving needed - LangGraph handles this automatically - app_logger.info( - f"Conversation completed and auto-saved for thread {thread_id}" - ) - - except Exception as e: - app_logger.error(f"Error in AgentManager stream_response: {e}") - yield { - "type": "error", - "content": { - "message": "Internal server error", - "recoverable": False, - "error_type": "agent_error", - }, - } - - async def _handle_input( - self, request: StreamRequest, agent: Pregel - ) -> tuple[Dict[str, Any], str, str]: - """Handle input preparation and configuration (preserving existing logic).""" - run_id = uuid4() - - # Generate default thread_id if not provided - thread_id = request.thread_id - if thread_id is None: - thread_id = str(uuid4()) - app_logger.info( - f"Assigning auto-generated thread_id '{thread_id}' as thread_id is missing in user request" - ) - - # Configure tracing and session management (preserved from original) - # If session_id is not provided, use thread_id as session_id - effective_session_id = request.session_id or thread_id - effective_user_id = request.user_id or "anonymous" - - # Register thread for user (for in-memory storage tracking) - if settings.USE_INMEMORY_SAVER: - register_thread(effective_user_id, thread_id) - - # Generate AI call ID - ai_call_id = f"ai_call_{str(uuid4())}" - - configurable = { - "thread_id": thread_id, - "session_id": effective_session_id, - "run_id": str(run_id), - "user_id": effective_user_id, - "ai_call_id": ai_call_id, - "langfuse_session_id": effective_session_id, - "langfuse_user_id": effective_user_id, - "langfuse_observation_id": thread_id, - } - - config = RunnableConfig( - configurable=configurable, - run_id=run_id, - callbacks=[langfuse_handler], - ) - - # Check for interrupts that need to be resumed (preserved from original) - state = await agent.aget_state(config=config) - interrupted_tasks = [ - task - for task in state.tasks - if hasattr(task, "interrupts") and task.interrupts - ] - - # Prepare input based on whether we're resuming from an interrupt - user_input_message: Command | Dict[str, Any] - if interrupted_tasks: - user_input_message = Command(resume=request.message) - else: - user_input_message = {"messages": [HumanMessage(content=request.message)]} - - kwargs = { - "input": user_input_message, - "config": config, - } - - app_logger.info( - f"AgentManager configured with run_id: {run_id}, thread_id: {thread_id}, session_id: {effective_session_id}" - ) - return kwargs, str(run_id), thread_id - - async def _prepare_streaming_input_with_history( - self, request: StreamRequest, existing_state, run_id: str, thread_id: str - ) -> Dict[str, Any]: - """Prepare streaming input with conversation history for non-checkpointing agent.""" - from langchain_core.messages import HumanMessage - from langchain_core.runnables import RunnableConfig - - # Get existing messages from state - existing_messages = existing_state.values.get("messages", []) - - # Create new message list with history + current user message - all_messages = list(existing_messages) - all_messages.append(HumanMessage(content=request.message)) - - # Configure for streaming agent (no checkpointing) - effective_session_id = request.session_id or thread_id - effective_user_id = request.user_id or "anonymous" - - configurable = { - "thread_id": thread_id, - "session_id": effective_session_id, - "run_id": run_id, - "user_id": effective_user_id, - "langfuse_session_id": effective_session_id, - "langfuse_user_id": effective_user_id, - "langfuse_observation_id": thread_id, - } - - config = RunnableConfig( - configurable=configurable, - run_id=run_id, - callbacks=[langfuse_handler], - ) - - return { - "input": {"messages": all_messages}, - "config": config, - } - - async def _save_final_conversation_state( - self, persistent_agent, config, all_messages: list, thread_id: str - ) -> None: - """Save the final conversation state once after streaming completes.""" - try: - app_logger.info( - f"Saving {len(all_messages)} messages for thread {thread_id}" - ) - - # Log message types for debugging - message_types = [ - getattr(msg, "type", type(msg).__name__) for msg in all_messages - ] - app_logger.info(f"Message types being saved: {message_types}") - - # Update the persistent agent's state with all messages - await persistent_agent.aupdate_state( - config=config, values={"messages": all_messages} - ) - app_logger.info( - f"Successfully saved conversation state for thread {thread_id}" - ) - - except Exception as e: - app_logger.error(f"Error saving final conversation state: {e}") - # Don't re-raise - streaming already completed successfully - - def _format_events( - self, - stream_mode: str, - event: Any, - stream_tokens: bool, - run_id: str, - thread_id: str, - session_id: str | None, - ) -> list[Dict[str, Any]]: - """Convert LangGraph events to simplified streaming format. - - This method implements the proposed event format while preserving - all the business logic from the original implementation. - """ - formatted_events = [] - - if stream_mode == "updates": - formatted_events.extend( - self._handle_update_events(event, run_id, thread_id, session_id) - ) - elif stream_mode == "messages" and stream_tokens: - token_event = self._handle_token_events(event) - if token_event: - formatted_events.append(token_event) - elif stream_mode == "custom": - custom_event = self._handle_custom_events( - event, run_id, thread_id, session_id - ) - if custom_event: - formatted_events.append(custom_event) - - return formatted_events - - def _handle_update_events( - self, event: Dict[str, Any], run_id: str, thread_id: str, session_id: str | None - ) -> list[Dict[str, Any]]: - """Handle update events from LangGraph (preserving existing logic).""" - formatted_events = [] - new_messages = [] - - for node, updates in event.items(): - # Handle agent interrupts with structured messages (preserved) - if node == "__interrupt__": - interrupt: Interrupt - for interrupt in updates: - new_messages.append(AIMessage(content=interrupt.value)) - continue - - updates = updates or {} - update_messages = updates.get("messages", []) - - # Special cases for using langgraph-supervisor library (preserved) - if node == "supervisor": - ai_messages = [ - msg for msg in update_messages if isinstance(msg, AIMessage) - ] - if ai_messages: - update_messages = [ai_messages[-1]] - - if node in ("research_expert", "math_expert"): - # Convert sub-agent output to ToolMessage for UI display (preserved) - msg = ToolMessage( - content=update_messages[0].content, - name=node, - tool_call_id="", - ) - update_messages = [msg] - - new_messages.extend(update_messages) - - # Process messages and convert to simplified format - processed_messages = self._process_message_tuples(new_messages) - - for message in processed_messages: - try: - chat_message = langchain_to_chat_message(message) - chat_message.run_id = run_id - - # Convert to simplified format - formatted_event = { - "type": "message", - "content": self._convert_chat_message_to_simple_format( - chat_message, thread_id, session_id - ), - } - formatted_events.append(formatted_event) - - except Exception as e: - app_logger.error(f"Error formatting message: {e}") - formatted_events.append( - { - "type": "error", - "content": { - "message": "Message formatting error", - "recoverable": True, - }, - } - ) - - return formatted_events - - def _handle_token_events(self, event: tuple) -> Dict[str, Any] | None: - """Handle token streaming events with tool call ID tracking.""" - msg, metadata = event - if "skip_stream" in metadata.get("tags", []): - return None - - # Filter out non-LLM node messages (preserved logic) - if not isinstance(msg, AIMessageChunk): - return None - - content = remove_tool_calls(msg.content) - if content: - token_event = { - "type": "token", - "content": convert_message_content_to_string(content), - } - - # Add tool call ID if this token is part of a tool call response - tool_call_id = ( - self._extract_tool_call_id_from_message(msg) - or self._current_tool_call_id - ) - if tool_call_id: - token_event["tool_call_id"] = tool_call_id - - return token_event - return None - - def _handle_custom_events( - self, event: Any, run_id: str, thread_id: str, session_id: str | None - ) -> Dict[str, Any] | None: - """Handle custom events from LangGraph.""" - try: - chat_message = langchain_to_chat_message(event) - chat_message.run_id = run_id - - return { - "type": "message", - "content": self._convert_chat_message_to_simple_format( - chat_message, thread_id, session_id - ), - } - except Exception as e: - app_logger.error(f"Error handling custom event: {e}") - return None - - def _process_message_tuples(self, new_messages: list) -> list: - """Process LangGraph streaming tuples and accumulate message parts (preserved logic).""" - processed_messages = [] - current_message: Dict[str, Any] = {} - - for message in new_messages: - if isinstance(message, tuple): - key, value = message - current_message[key] = value - else: - # Add complete message if we have one in progress - if current_message: - processed_messages.append(self._create_ai_message(current_message)) - current_message = {} - processed_messages.append(message) - - # Add any remaining message parts - if current_message: - processed_messages.append(self._create_ai_message(current_message)) - - return processed_messages - - def _create_ai_message(self, parts: Dict[str, Any]) -> AIMessage: - """Create an AIMessage from a dictionary of parts (preserved from original).""" - sig = inspect.signature(AIMessage) - valid_keys = set(sig.parameters) - filtered = {k: v for k, v in parts.items() if k in valid_keys} - return AIMessage(**filtered) - - def _convert_chat_message_to_simple_format( - self, chat_message, thread_id: str, session_id: str | None - ) -> Dict[str, Any]: - """Convert ChatMessage to simplified content format for the proposed API.""" - content = { - "type": chat_message.type, - "content": chat_message.content, - } - - # Add optional fields only if present - if chat_message.tool_calls: - content["tool_calls"] = chat_message.tool_calls - if chat_message.tool_call_id: - content["tool_call_id"] = chat_message.tool_call_id - if chat_message.run_id: - content["run_id"] = chat_message.run_id - if thread_id: - content["thread_id"] = thread_id - if session_id: - content["session_id"] = session_id - if chat_message.ai_call_id: - content["ai_call_id"] = chat_message.ai_call_id - if chat_message.response_metadata: - content["response_metadata"] = chat_message.response_metadata - if chat_message.custom_data: - content["custom_data"] = chat_message.custom_data - - return content - - def _extract_tool_call_id_from_message(self, msg: AIMessageChunk) -> str | None: - """Extract tool call ID from an AIMessageChunk if available. - - Args: - msg: The AIMessageChunk to extract tool call ID from - - Returns: - The tool call ID if available, None otherwise - """ - try: - # Check if the message has tool calls - if hasattr(msg, "tool_calls") and msg.tool_calls: - # Return the ID of the first tool call - return msg.tool_calls[0].get("id") - - # Check if the message has tool_call_chunks (streaming tool calls) - if hasattr(msg, "tool_call_chunks") and msg.tool_call_chunks: - # Return the ID of the first tool call chunk - return msg.tool_call_chunks[0].get("id") - - # Check if this is a response to a tool call (has tool_call_id) - if hasattr(msg, "tool_call_id") and msg.tool_call_id: - return msg.tool_call_id - - return None - except (AttributeError, IndexError, KeyError) as e: - app_logger.debug(f"Could not extract tool call ID from message: {e}") - return None - - def _update_tool_call_tracking(self, stream_mode: str, event: Any) -> None: - """Update the current tool call ID based on streaming events. - - Args: - stream_mode: The type of stream event - event: The event data - """ - try: - if stream_mode == "updates": - # Look for tool calls in update events - for node, updates in event.items(): - if updates and "messages" in updates: - for message in updates["messages"]: - if hasattr(message, "tool_calls") and message.tool_calls: - # Found a new tool call, update tracking - self._current_tool_call_id = message.tool_calls[0].get( - "id" - ) - app_logger.debug( - f"Tracking tool call ID: {self._current_tool_call_id}" - ) - return - elif ( - hasattr(message, "tool_call_id") - and message.tool_call_id - ): - # This is a tool response, track its ID - self._current_tool_call_id = message.tool_call_id - app_logger.debug( - f"Tracking tool response ID: {self._current_tool_call_id}" - ) - return - - elif stream_mode == "messages": - # Check message stream for tool calls - msg, metadata = event - if hasattr(msg, "tool_calls") and msg.tool_calls: - self._current_tool_call_id = msg.tool_calls[0].get("id") - app_logger.debug( - f"Tracking tool call ID from message: {self._current_tool_call_id}" - ) - elif hasattr(msg, "tool_call_id") and msg.tool_call_id: - self._current_tool_call_id = msg.tool_call_id - app_logger.debug( - f"Tracking tool response ID from message: {self._current_tool_call_id}" - ) - - except Exception as e: - app_logger.debug(f"Error updating tool call tracking: {e}") - # Don't fail streaming due to tracking issues diff --git a/template_agent/src/core/prompt.py b/template_agent/src/core/prompt.py deleted file mode 100644 index 6e0a7ac9..00000000 --- a/template_agent/src/core/prompt.py +++ /dev/null @@ -1,49 +0,0 @@ -"""System prompts and prompt utilities for the template agent. - -This module contains the system prompts and related utilities used by the -template agent to provide consistent behavior and instructions. -""" - -from datetime import datetime - - -def get_current_date() -> str: - """Get the current date in a formatted string. - - Returns: - The current date formatted as "Month Day, Year" (e.g., "December 25, 2024"). - """ - return datetime.now().strftime("%B %d, %Y") - - -def get_system_prompt() -> str: - """Get the main system prompt for the template agent. - - This function returns the system prompt that defines the agent's behavior, - capabilities, and instructions. The prompt includes the current date and - specific guidelines for tool usage and response formatting. - - Returns: - The complete system prompt string with current date and instructions. - """ - current_date = get_current_date() - - return ( - f"You are Template Agent, a powerful and helpful assistant with the ability to use specialized tools.\n\n" - f"Today's date is {current_date}.\n\n" - "A few things to remember:\n" - "- **Always use the same language as the user.**\n" - "- **Always send intermediate responses between tool calls to the user showing the reasoning and thought process.**\n" - "- **If needed or requested by user, you can use Markdown to generate tables, code blocks, lists, etc.**\n" - "- **You have access to mathematical tools:**\n" - " 1. **multiply_numbers:** Use this tool to multiply two numbers together.\n" - "- **Only use the tools you are given to answer the user's question.** Do not answer directly from internal knowledge.\n" - "- **You must always reason before acting.** First, determine if a mathematical operation is needed. If so, use the multiply_numbers tool to get the result.\n" - "- **Every Final Answer must be grounded in tool observations.**\n" - "- **Always make sure your answer is *FORMATTED WELL*.**\n\n" - "# OUTPUT FORMAT [Never ignore following instructions]\n" - "- You MUST always respond using proper Markdown formatting.\n" - "- Use headers (#, ##, ###), lists (- or 1.), code blocks (```), bold (**text**), and tables when appropriate.\n" - "- For the final response, provide a well-structured Markdown summary.\n" - "- For intermediate responses, use simple Markdown formatting.\n" - ) diff --git a/template_agent/src/core/storage.py b/template_agent/src/core/storage.py deleted file mode 100644 index 41a60b29..00000000 --- a/template_agent/src/core/storage.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Global storage management for the template agent system. - -This module provides a single global checkpoint instance that persists across -the entire application lifecycle when using in-memory storage mode. -""" - -from typing import Optional - -from langgraph.checkpoint.memory import InMemorySaver - -from template_agent.src.settings import settings -from template_agent.utils.pylogger import get_python_logger - -logger = get_python_logger(settings.PYTHON_LOG_LEVEL) - -# Global checkpoint instance - single instance for the entire application lifecycle -_global_checkpoint: Optional[InMemorySaver] = None - -# Global thread registry to track threads by user_id -_thread_registry: dict[str, set[str]] = {} - - -def get_global_checkpoint() -> InMemorySaver: - """Get the global in-memory checkpoint instance. - - This creates a single checkpoint instance that persists for the entire - application lifecycle, ensuring all components use the same storage. - The same instance serves as both checkpointer and store. - - Returns: - The global InMemorySaver instance. - """ - global _global_checkpoint - if _global_checkpoint is None: - _global_checkpoint = InMemorySaver() - logger.info("Created global InMemorySaver checkpoint instance") - return _global_checkpoint - - -def register_thread(user_id: str, thread_id: str) -> None: - """Register a thread for a user. - - Args: - user_id: The user ID - thread_id: The thread ID to register - """ - global _thread_registry - if user_id not in _thread_registry: - _thread_registry[user_id] = set() - _thread_registry[user_id].add(thread_id) - logger.info(f"Registered thread {thread_id} for user {user_id}") - - -def get_user_threads(user_id: str) -> list[str]: - """Get all threads for a user. - - Args: - user_id: The user ID - - Returns: - List of thread IDs for the user - """ - global _thread_registry - threads = list(_thread_registry.get(user_id, set())) - logger.info(f"Retrieved {len(threads)} threads for user {user_id}: {threads}") - return threads - - -def reset_global_storage() -> None: - """Reset the global checkpoint instance. - - This is useful for testing or when you want to clear all data. - """ - global _global_checkpoint, _thread_registry - _global_checkpoint = None - _thread_registry = {} - logger.info("Reset global checkpoint instance and thread registry") - - -# Backward compatibility aliases -get_shared_checkpointer = get_global_checkpoint -get_shared_store = get_global_checkpoint -reset_shared_storage = reset_global_storage diff --git a/template_agent/src/main.py b/template_agent/src/main.py deleted file mode 100644 index da392cc9..00000000 --- a/template_agent/src/main.py +++ /dev/null @@ -1,169 +0,0 @@ -"""Main entry point for the template agent server. - -This module provides the main application entry point, including -configuration validation, server startup, and graceful shutdown -handling for the template agent service. -""" - -import sys -from typing import NoReturn - -import uvicorn - -from template_agent.src.api import app -from template_agent.src.core.exceptions.exceptions import AppException, AppExceptionCode -from template_agent.src.settings import settings -from template_agent.src.settings import validate_config as validate_config_func -from template_agent.utils.google_creds import initialize_google_genai -from template_agent.utils.pylogger import get_python_logger, get_uvicorn_log_config - -# Initialize logger -logger = get_python_logger(settings.PYTHON_LOG_LEVEL) - - -def validate_and_initialize_config() -> None: - """Validate configuration settings and initialize external services. - - Performs additional runtime validation of configuration values - beyond what's done in the Settings class initialization. This - includes validating host configurations and initializing external - services like Google Generative AI. - - Raises: - ValueError: If required configuration values are missing or invalid. - RuntimeError: If configuration is in an inconsistent state. - """ - try: - # Use the validate_config function from settings.py - validate_config_func(settings) - initialize_google_genai() - - logger.info("Configuration validation and initialization passed") - - except AttributeError: - # Handle case where config object is not properly initialized - raise AppException( - "Failed to properly initialize configurations", - AppExceptionCode.CONFIGURATION_INITIALIZATION_ERROR, - ) - except Exception: - # Re-raise as ValueError for consistent error handling - raise AppException( - "Configuration validation failed", - AppExceptionCode.CONFIGURATION_VALIDATION_ERROR, - ) - - -def handle_startup_error(error: Exception, context: str = "server startup") -> NoReturn: - """Handle startup errors with proper logging and exit codes. - - This function provides centralized error handling for different - types of startup errors, ensuring appropriate logging and exit - codes for different error scenarios. - - Args: - error: The exception that occurred during startup. - context: Context where the error occurred for better logging. - - Raises: - SystemExit: Always raises SystemExit with appropriate exit code - based on the error type. - """ - if isinstance(error, ValueError): - # Configuration or validation errors - logger.critical(f"Configuration error during {context}: {error}") - sys.exit(1) - elif isinstance(error, KeyboardInterrupt): - # User interrupted the startup - logger.info("Server startup interrupted by user") - sys.exit(0) - elif isinstance(error, PermissionError): - # Permission issues (e.g., port binding) - logger.critical(f"Permission error during {context}: {error}") - sys.exit(1) - elif isinstance(error, ConnectionError): - # Network-related errors - logger.critical(f"Connection error during {context}: {error}") - sys.exit(1) - else: - # Unexpected errors - logger.critical(f"Unexpected error during {context}: {error}", exc_info=True) - sys.exit(1) - - -def main() -> None: - """Main entry point for the template agent server. - - Initializes logging, loads configuration, and starts the template - agent server. Handles graceful shutdown on keyboard interrupt and - logs any startup errors. - - The function performs the following steps: - 1. Validates configuration settings - 2. Initializes external services - 3. Configures uvicorn server settings - 4. Starts the server with appropriate error handling - - Raises: - SystemExit: If the server fails to start due to configuration - or other errors. - """ - try: - validate_and_initialize_config() - - logger.info( - f"Starting template agent server on {settings.AGENT_HOST}:{settings.AGENT_PORT}" - ) - - # Configure uvicorn server settings - uvicorn_config = { - "app": app, - "host": settings.AGENT_HOST, - "port": settings.AGENT_PORT, - "log_config": get_uvicorn_log_config(settings.PYTHON_LOG_LEVEL), - } - - # Add SSL configuration if certificates are provided - if settings.AGENT_SSL_KEYFILE and settings.AGENT_SSL_CERTFILE: - uvicorn_config["ssl_keyfile"] = settings.AGENT_SSL_KEYFILE - uvicorn_config["ssl_certfile"] = settings.AGENT_SSL_CERTFILE - logger.info( - "Starting server with SSL", - ssl_keyfile=settings.AGENT_SSL_KEYFILE, - ssl_certfile=settings.AGENT_SSL_CERTFILE, - ) - - uvicorn.run(**uvicorn_config) - - except KeyboardInterrupt: - logger.info("Received keyboard interrupt, shutting down") - except Exception as e: - handle_startup_error(e, "server startup") - finally: - logger.info("Template agent server shutting down") - - -def run() -> None: - """Run the server with comprehensive error handling. - - Wraps the main function with additional error handling for graceful - shutdown and proper exit codes. Provides a safety net for any - unhandled exceptions that might occur during server startup or - operation. - - Raises: - SystemExit: If the server fails to start or encounters critical errors. - """ - try: - main() - except KeyboardInterrupt: - logger.info("Server stopped by user") - sys.exit(0) - except Exception as e: - # This should rarely be reached due to handle_startup_error - logger.error("Server failed to start", error=str(e), exc_info=True) - sys.exit(1) - - -if __name__ == "__main__": - run() diff --git a/template_agent/src/routes/__init__.py b/template_agent/src/routes/__init__.py deleted file mode 100644 index de6f47a8..00000000 --- a/template_agent/src/routes/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Routes package for the template agent API.""" diff --git a/template_agent/src/routes/feedback.py b/template_agent/src/routes/feedback.py deleted file mode 100644 index 1a0130ac..00000000 --- a/template_agent/src/routes/feedback.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Feedback route for the template agent API. - -This module provides endpoints for recording user feedback on agent responses -using Langfuse for analytics and monitoring purposes. -""" - -from fastapi import APIRouter -from langfuse import Langfuse - -from template_agent.src.schema import FeedbackRequest, FeedbackResponse -from template_agent.src.settings import settings - -router = APIRouter() - -# Initialize Langfuse client for feedback tracking -client = Langfuse(environment=settings.LANGFUSE_TRACING_ENVIRONMENT) - - -@router.post("/v1/feedback") -async def feedback(feedback: FeedbackRequest) -> FeedbackResponse: - """Record feedback for a specific agent run to Langfuse. - - This endpoint serves as a wrapper for the Langfuse create_feedback API, - allowing credentials to be stored and managed in the service rather than - requiring client-side credential management. - - The function maps the feedback request parameters to Langfuse's expected - format: - - run_id -> trace_id - - key -> name - - score -> value - - Args: - feedback: The feedback request containing run_id, key, score, and - optional kwargs for additional metadata. - - Returns: - A FeedbackResponse indicating successful feedback recording. - - Raises: - Exception: If there are issues with the Langfuse API call. - - See Also: - https://api.smith.langchain.com/redoc#tag/feedback/operation/create_feedback_api_v1_feedback_post - """ - kwargs = feedback.kwargs or {} - - # Langfuse uses different parameter names than our schema - client.score( - trace_id=feedback.run_id, # Assuming run_id maps to trace_id - name=feedback.key, # 'key' becomes 'name' in Langfuse - value=feedback.score, # 'score' becomes 'value' in Langfuse - **kwargs, - ) - - return FeedbackResponse() diff --git a/template_agent/src/routes/health.py b/template_agent/src/routes/health.py deleted file mode 100644 index 24cbb95a..00000000 --- a/template_agent/src/routes/health.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Health check route for the template agent API. - -This module provides health check endpoints to monitor the status -and availability of the template agent service. -""" - -from fastapi import APIRouter -from fastapi.responses import JSONResponse - -router = APIRouter() - - -@router.get("/health") -async def health_check() -> JSONResponse: - """Perform a health check on the template agent service. - - This endpoint is used to verify that the service is running and - responding to requests. It returns a simple JSON response indicating - the service status. - - Returns: - A JSONResponse containing the service status and name. - """ - return JSONResponse(content={"status": "healthy", "service": "Template Agent"}) diff --git a/template_agent/src/routes/history.py b/template_agent/src/routes/history.py deleted file mode 100644 index 3479ca1f..00000000 --- a/template_agent/src/routes/history.py +++ /dev/null @@ -1,488 +0,0 @@ -"""History route for the template agent API. - -This module provides endpoints for retrieving chat history from the database, -allowing users to view previous conversations and continue ongoing threads. -""" - -from typing import List - -import psycopg2 -from fastapi import APIRouter, HTTPException, Request -from langchain_core.runnables import RunnableConfig - -from template_agent.src.core.agent_utils import langchain_to_chat_message -from template_agent.src.core.storage import get_shared_checkpointer -from template_agent.src.schema import ChatHistoryResponse, ChatMessage, ToolCall -from template_agent.src.settings import settings -from template_agent.utils.pylogger import get_python_logger - -router = APIRouter() - -logger = get_python_logger(settings.PYTHON_LOG_LEVEL) - - -@router.get("/v1/history/{thread_id}") -async def history(thread_id: str, request: Request) -> ChatHistoryResponse: - """Get chat history for a specific thread by reading from checkpoints table. - - This endpoint retrieves the complete conversation history for a given - thread_id from the PostgreSQL database. When using in-memory storage, - returns an empty history since conversations are not persisted. - - The function handles different message types (human, ai, tool) and - converts them to the internal ChatMessage format for consistent - representation across the application. - - Args: - thread_id: The unique identifier of the thread to retrieve history for. - request: The FastAPI request object, used to extract headers like - X-Token for authentication. - - Returns: - A ChatHistoryResponse containing the list of chat messages for the thread. - Returns empty history when using in-memory storage. - If there's an error, returns an empty message list instead of raising - an exception. - - Note: - - Messages are extracted from multiple locations in the checkpoint data - - The function handles various message formats and gracefully skips - invalid messages - - Authentication tokens are logged but not currently used for validation - - In-memory storage mode returns empty history as conversations are not persisted - """ - access_token = request.headers.get("X-Token") - logger.info(f"Retrieving history for thread_id: {thread_id}") - logger.info(f"Access token present: {access_token is not None}") - - chat_messages: List[ChatMessage] = [] - - # When using in-memory storage, get history from shared checkpointer - if settings.USE_INMEMORY_SAVER: - logger.info( - f"Using in-memory storage - retrieving history from checkpointer for thread_id: {thread_id}" - ) - try: - checkpointer = get_shared_checkpointer() - - # Create a config for this thread (matching the format used by the agent) - config = RunnableConfig( - configurable={"thread_id": thread_id, "checkpoint_ns": ""} - ) - - # Get all checkpoints for this thread to understand the structure - state_history = list(checkpointer.list(config)) - logger.info( - f"Found {len(state_history)} checkpoints for thread_id: {thread_id}" - ) - - if len(state_history) == 0: - logger.info( - f"No checkpoints found for thread {thread_id} - this means no conversations have happened in this thread yet." - ) - else: - # DEBUG: Log structure of all checkpoints to understand how messages are stored - for i, checkpoint_tuple in enumerate(state_history): - logger.info(f"=== CHECKPOINT {i} DEBUG ===") - logger.info( - f"Checkpoint keys: {list(checkpoint_tuple.checkpoint.keys()) if checkpoint_tuple.checkpoint else 'None'}" - ) - - if ( - checkpoint_tuple.checkpoint - and "channel_values" in checkpoint_tuple.checkpoint - ): - channel_values = checkpoint_tuple.checkpoint["channel_values"] - logger.info( - f"Channel values keys: {list(channel_values.keys())}" - ) - - if "messages" in channel_values: - messages = channel_values["messages"] - logger.info( - f"Messages count in checkpoint {i}: {len(messages)}" - ) - for j, msg in enumerate(messages): - msg_type = ( - getattr(msg, "type", "unknown") - if hasattr(msg, "type") - else type(msg).__name__ - ) - msg_content = ( - getattr(msg, "content", str(msg)[:100]) - if hasattr(msg, "content") - else str(msg)[:100] - ) - logger.info( - f" Message {j}: {msg_type} - {msg_content}" - ) - else: - logger.info( - f"No 'messages' key in channel_values for checkpoint {i}" - ) - else: - logger.info(f"No channel_values in checkpoint {i}") - - # Try the latest checkpoint first (our current approach) - latest_checkpoint = state_history[-1] - logger.info( - f"=== PROCESSING LATEST CHECKPOINT (index {len(state_history) - 1}) ===" - ) - - if ( - latest_checkpoint.checkpoint - and "channel_values" in latest_checkpoint.checkpoint - ): - channel_values = latest_checkpoint.checkpoint["channel_values"] - if "messages" in channel_values: - messages = channel_values["messages"] - logger.info( - f"Found {len(messages)} messages in latest checkpoint" - ) - for message in messages: - try: - chat_message = langchain_to_chat_message(message) - chat_messages.append(chat_message) - logger.info( - f"Added message: {chat_message.type} - {chat_message.content[:50]}..." - ) - except Exception as e: - logger.warning( - f"Could not convert message to ChatMessage: {e}" - ) - continue - - # If latest checkpoint approach didn't work, try collecting from all checkpoints - if len(chat_messages) == 0: - logger.info("=== FALLBACK: PROCESSING ALL CHECKPOINTS ===") - for i, checkpoint_tuple in enumerate(state_history): - if ( - checkpoint_tuple.checkpoint - and "channel_values" in checkpoint_tuple.checkpoint - ): - channel_values = checkpoint_tuple.checkpoint[ - "channel_values" - ] - if "messages" in channel_values: - messages = channel_values["messages"] - logger.info( - f"Processing {len(messages)} messages from checkpoint {i}" - ) - for message in messages: - try: - chat_message = langchain_to_chat_message( - message - ) - # Check for duplicates before adding - is_duplicate = False - for existing_msg in chat_messages: - if ( - existing_msg.type == chat_message.type - and existing_msg.content - == chat_message.content - ): - is_duplicate = True - break - - if not is_duplicate: - chat_messages.append(chat_message) - logger.info( - f"Added unique message: {chat_message.type} - {chat_message.content[:50]}..." - ) - else: - logger.info( - f"Skipped duplicate message: {chat_message.type} - {chat_message.content[:50]}..." - ) - except Exception as e: - logger.warning( - f"Could not convert message to ChatMessage: {e}" - ) - continue - - logger.info( - f"Found {len(chat_messages)} messages in memory for thread_id: {thread_id}" - ) - - return ChatHistoryResponse(messages=chat_messages) - except Exception as e: - logger.error( - f"Error accessing in-memory storage for thread {thread_id}: {e}" - ) - return ChatHistoryResponse(messages=[]) - - try: - # Connect to PostgreSQL and read from checkpoints table - with psycopg2.connect(settings.database_uri) as conn: - cur = conn.cursor() - - # Query the checkpoints table for the specific thread_id - # Get the latest checkpoint first (which should contain complete conversation state) - cur.execute( - "SELECT checkpoint, metadata FROM checkpoints WHERE thread_id = %s ORDER BY checkpoint_id DESC LIMIT 1", - (thread_id,), - ) - latest_row = cur.fetchone() - - if latest_row: - logger.info(f"Found latest checkpoint for thread_id: {thread_id}") - checkpoint_data, metadata = latest_row - - # DEBUG: Log the structure of the latest checkpoint - logger.info("=== POSTGRESQL LATEST CHECKPOINT DEBUG ===") - logger.info( - f"Checkpoint_data keys: {list(checkpoint_data.keys()) if checkpoint_data else 'None'}" - ) - logger.info( - f"Metadata keys: {list(metadata.keys()) if metadata else 'None'}" - ) - - # Try to get complete conversation from latest checkpoint - if checkpoint_data and "channel_values" in checkpoint_data: - channel_values = checkpoint_data["channel_values"] - logger.info(f"Channel values keys: {list(channel_values.keys())}") - - if "messages" in channel_values: - checkpoint_messages = channel_values["messages"] - logger.info( - f"Found {len(checkpoint_messages)} messages in latest checkpoint channel_values" - ) - - # DEBUG: Log each message structure - for i, msg in enumerate(checkpoint_messages): - msg_type = ( - getattr(msg, "type", "unknown") - if hasattr(msg, "type") - else type(msg).__name__ - ) - msg_content = ( - getattr(msg, "content", str(msg)[:100]) - if hasattr(msg, "content") - else str(msg)[:100] - ) - logger.info( - f" PostgreSQL Message {i}: {msg_type} - {msg_content}" - ) - - # Extract metadata for tracking - run_id = metadata.get("run_id") if metadata else None - session_id = metadata.get("session_id") if metadata else None - user_id = metadata.get("user_id") if metadata else None - - # Convert LangChain messages directly (like in-memory version) - for message in checkpoint_messages: - try: - chat_message = langchain_to_chat_message(message) - # Set metadata from checkpoint for tracking - if run_id: - chat_message.run_id = run_id - if thread_id: - chat_message.thread_id = thread_id - if session_id: - chat_message.session_id = session_id - chat_messages.append(chat_message) - logger.info( - f"Successfully converted checkpoint message: {chat_message.type} - {chat_message.content[:50]}..." - ) - except Exception as e: - logger.warning( - f"Could not convert checkpoint message to ChatMessage: {e}" - ) - continue - - logger.info( - f"Retrieved {len(chat_messages)} messages from latest checkpoint for thread_id: {thread_id}" - ) - return ChatHistoryResponse(messages=chat_messages) - else: - logger.info( - "No 'messages' key found in channel_values of latest checkpoint" - ) - else: - logger.info("No 'channel_values' found in latest checkpoint_data") - - # Fallback: If latest checkpoint doesn't have messages, process all checkpoints with writes - logger.info( - "Latest checkpoint didn't contain messages, falling back to processing all checkpoints" - ) - cur.execute( - "SELECT checkpoint, metadata FROM checkpoints WHERE thread_id = %s ORDER BY checkpoint_id ASC", - (thread_id,), - ) - rows = cur.fetchall() - - logger.info(f"Found {len(rows)} checkpoints for thread_id: {thread_id}") - - total_messages_found = 0 - - # Process each checkpoint to extract messages from writes (fallback approach) - for row in rows: - checkpoint_data, metadata = row - - # Extract run_id, thread_id, session_id from metadata for tracking - run_id = metadata.get("run_id") if metadata else None - session_id = metadata.get("session_id") if metadata else None - user_id = metadata.get("user_id") if metadata else None - - logger.info( - f"Processing checkpoint with run_id: {run_id}, session_id: {session_id}, user_id: {user_id}" - ) - - # Get messages from metadata.writes (original logic) - messages = [] - writes = metadata.get("writes", {}) if metadata else {} - - # Handle case where writes might be None - if writes is None: - writes = {} - logger.info("Writes is None, using empty dict") - - # Check for messages in different write locations - if "__start__" in writes and "messages" in writes["__start__"]: - messages.extend(writes["__start__"]["messages"]) - logger.info( - f"Found {len(writes['__start__']['messages'])} messages in __start__" - ) - if "agent" in writes and "messages" in writes["agent"]: - messages.extend(writes["agent"]["messages"]) - logger.info( - f"Found {len(writes['agent']['messages'])} messages in agent" - ) - if "tools" in writes and "messages" in writes["tools"]: - messages.extend(writes["tools"]["messages"]) - logger.info( - f"Found {len(writes['tools']['messages'])} messages in tools" - ) - - total_messages_found += len(messages) - - # Convert each message to ChatMessage format - for message_data in messages: - try: - logger.info(f"Processing message_data: {message_data}") - - # Validate message format - should be a dict with kwargs - if ( - not isinstance(message_data, dict) - or "kwargs" not in message_data - ): - logger.info( - f"Skipping invalid message format: {message_data}" - ) - continue - - # Extract message components - kwargs = message_data.get("kwargs", {}) - message_type = kwargs.get("type", "") - content = kwargs.get("content", "") - response_metadata = kwargs.get("response_metadata", {}) - - # Handle tool calls from both direct kwargs and additional_kwargs - tool_calls = kwargs.get("tool_calls", []) - if not tool_calls and "additional_kwargs" in kwargs: - tool_calls = kwargs["additional_kwargs"].get( - "tool_calls", [] - ) - - logger.info(f"Message type: {message_type}, content: {content}") - - # Import here to avoid circular imports - from langchain_core.messages import ( - AIMessage, - HumanMessage, - ToolMessage, - ) - - # Create appropriate LangChain message based on type - if message_type == "human": - message = HumanMessage(content=content) - elif message_type == "ai": - message = AIMessage( - content=content, - tool_calls=tool_calls, - additional_kwargs={ - "response_metadata": response_metadata - }, - ) - elif message_type == "tool": - tool_call_id = kwargs.get("tool_call_id") - name = kwargs.get("name", "") - message = ToolMessage( - content=content, - tool_call_id=tool_call_id, - name=name, - additional_kwargs={ - "response_metadata": response_metadata - }, - ) - else: - logger.info( - f"Skipping unknown message type: {message_type}" - ) - continue - - # Convert to internal ChatMessage format - chat_message = langchain_to_chat_message(message) - - # Set metadata from checkpoint for tracking - if run_id: - chat_message.run_id = run_id - if thread_id: - chat_message.thread_id = thread_id - if session_id: - chat_message.session_id = session_id - - # Set metadata from the original message data - if response_metadata: - chat_message.response_metadata = response_metadata - - # Set tool calls if present - if tool_calls: - # Ensure tool calls have the correct structure - formatted_tool_calls = [] - for tool_call in tool_calls: - if isinstance(tool_call, dict): - # Ensure required fields are present and properly typed - if "name" in tool_call and "args" in tool_call: - # Create a proper ToolCall object - formatted_call: ToolCall = { - "name": str(tool_call["name"]), - "args": dict(tool_call["args"]), - "id": str(tool_call.get("id")) - if tool_call.get("id") - else None, - "type": "tool_call", - } - formatted_tool_calls.append(formatted_call) - chat_message.tool_calls = formatted_tool_calls - logger.info( - f"Added {len(formatted_tool_calls)} tool calls to message" - ) - - logger.info( - f"Successfully converted message: {chat_message.type} - {chat_message.content[:50]}..." - ) - logger.info( - "Message metadata: " - f"tool_calls={bool(chat_message.tool_calls)}, " - f"response_metadata={bool(chat_message.response_metadata)}" - ) - chat_messages.append(chat_message) - - except Exception as e: - logger.error(f"Error processing message: {e}") - continue - - logger.info( - f"Retrieved {len(chat_messages)} messages for thread_id: {thread_id}" - ) - logger.info(f"Total messages found: {total_messages_found}") - logger.info(f"Final chat_messages: {[msg.type for msg in chat_messages]}") - return ChatHistoryResponse(messages=chat_messages) - - except Exception as e: - logger.error( - f"Database error while fetching history for thread {thread_id}: {e}" - ) - raise HTTPException( - status_code=500, detail=f"Failed to retrieve chat history: {str(e)}" - ) diff --git a/template_agent/src/routes/stream.py b/template_agent/src/routes/stream.py deleted file mode 100644 index 12f1635b..00000000 --- a/template_agent/src/routes/stream.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Stream route for the template agent API. - -This module provides streaming endpoints for real-time agent interactions, -handling message streaming, token generation, and conversation management. -""" - -import json -from collections.abc import AsyncGenerator -from typing import Any - -from fastapi import APIRouter, HTTPException, Request, status -from fastapi.responses import StreamingResponse - -from template_agent.src.core.manager import AgentManager -from template_agent.src.schema import StreamRequest -from template_agent.src.settings import settings -from template_agent.utils.pylogger import get_python_logger - -router = APIRouter() -app_logger = get_python_logger(settings.PYTHON_LOG_LEVEL) - - -async def message_generator( - user_input: StreamRequest, agent_manager: AgentManager -) -> AsyncGenerator[str, None]: - """Generate a stream of messages from the agent using the simplified format. - - This function uses the AgentManager to handle streaming with features like - SSO authentication, tracing, and error handling. The AgentManager is - initialized before streaming begins to allow proper HTTP error responses. - - Args: - user_input: The streaming input from the user containing the message - and configuration. - agent_manager: Pre-initialized AgentManager instance. - - Yields: - JSON-formatted SSE messages as strings in the simplified event format. - - Note: - - Uses simplified event format: {"type": "message"|"token"|"error", "content": ...} - - Preserves enterprise features: SSO auth, Langfuse tracing, error handling - - Errors during streaming are sent as error events in the stream - - Initialization errors are handled before streaming starts - """ - try: - app_logger.info(f"Starting stream for message: {user_input.message[:100]}...") - - # Stream events using the simplified AgentManager - async for event in agent_manager.stream_response(user_input): - # Filter out duplicate human messages - if ( - event.get("type") == "message" - and event.get("content", {}).get("type") == "human" - and event.get("content", {}).get("content") == user_input.message - ): - continue - - # Yield the simplified event format - yield f"{json.dumps(event, separators=(',', ':'))}\n\n" - - except Exception as e: - app_logger.error(f"Error in message generator: {e}") - error_event = { - "type": "error", - "content": { - "message": "Internal server error", - "recoverable": False, - "error_type": "stream_error", - }, - } - yield f"{json.dumps(error_event)}\n\n" - finally: - # Send completion marker - yield "[DONE]\n\n" - - -def _sse_response_example() -> dict[int | str, Any]: - """Generate example response for SSE endpoint documentation. - - Returns: - A dictionary containing the example SSE response format for - the simplified streaming API. - """ - return { - status.HTTP_200_OK: { - "description": "Server Sent Event Response - Simplified Format", - "content": { - "text/event-stream": { - "example": '{"type": "message", "content": {"type": "ai", "content": "", "tool_calls": [{"name": "multiply", "args": {"a": 3, "b": 2}, "id": "call_123"}], "run_id": "12345", "thread_id": "thread-123", "session_id": "session-456"}}\n\n{"type": "message", "content": {"type": "tool", "content": "6", "tool_call_id": "call_123", "run_id": "12345", "thread_id": "thread-123", "session_id": "session-456"}}\n\n{"type": "token", "content": "The"}\n\n{"type": "token", "content": " answer"}\n\n{"type": "token", "content": " is"}\n\n{"type": "token", "content": " 6"}\n\n{"type": "message", "content": {"type": "ai", "content": "The answer is 6", "run_id": "12345", "thread_id": "thread-123", "session_id": "session-456"}}\n\n[DONE]\n\n', - "schema": {"type": "string"}, - } - }, - } - } - - -@router.post( - "/v1/stream", response_class=StreamingResponse, responses=_sse_response_example() -) -async def stream(user_input: StreamRequest, request: Request) -> StreamingResponse: - """Stream AI agent responses in real-time using simplified event format. - - This endpoint provides the core streaming functionality following the - simplified API design with features like SSO - authentication, Langfuse tracing, and comprehensive error handling. - - **Event Types:** - - `message` - Tool calls, tool results, and final responses - - `token` - Individual tokens (only when `stream_tokens: true`) - - `error` - Error messages with recovery information - - `[DONE]` - Stream completion marker - - **Request Fields:** - - `message`: User's input message (required) - - `thread_id`: Conversation thread identifier (optional - auto-generated if not provided) - - `session_id`: Session identifier (required) - - `user_id`: User identifier for tracking and personalization (required) - - `stream_tokens`: Whether to stream individual tokens (`true`) or just complete messages (`false`) (optional) - - **Enterprise Features (Preserved):** - - SSO authentication via X-Token header - - Langfuse tracing and analytics - - PostgreSQL checkpointing for conversation persistence - - Comprehensive error handling and logging - - Args: - user_input: The streaming request with simplified structure. - request: FastAPI request object for extracting authentication headers. - - Returns: - StreamingResponse with simplified event format: - ``` - {"type": "message", "content": {"type": "ai", "content": "Hello", "run_id": "12345", "thread_id": "thread-123", "session_id": "session-456"}} - {"type": "token", "content": "world"} - [DONE] - ``` - - Raises: - HTTPException: If initialization fails (returns 500 status code). - """ - # Get token from request headers - access_token = request.headers.get("X-Token") - app_logger.info(f"Received token: {'Yes' if access_token else 'No'}") - - # Initialize AgentManager BEFORE streaming to catch initialization errors - try: - agent_manager = AgentManager(redhat_sso_token=access_token) - except Exception as e: - app_logger.error(f"Failed to initialize AgentManager: {e}") - raise HTTPException( - status_code=500, detail=f"Failed to initialize agent: {str(e)}" - ) - - return StreamingResponse( - message_generator(user_input, agent_manager), - media_type="text/event-stream", - headers={ - "Cache-Control": "no-cache", - "Connection": "keep-alive", - }, - ) diff --git a/template_agent/src/routes/threads.py b/template_agent/src/routes/threads.py deleted file mode 100644 index 4f869390..00000000 --- a/template_agent/src/routes/threads.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Threads route for the template agent API. - -This module provides endpoints for managing conversation threads, -including listing threads for specific users. -""" - -from typing import List - -import psycopg2 -from fastapi import APIRouter, HTTPException - -from template_agent.src.core.storage import get_user_threads -from template_agent.src.settings import settings -from template_agent.utils.pylogger import get_python_logger - -router = APIRouter() - -app_logger = get_python_logger(settings.PYTHON_LOG_LEVEL) - - -@router.get("/v1/threads/{user_id}") -async def list_threads(user_id: str) -> List[str]: - """Get a list of all thread IDs for a specific user. - - This endpoint queries the PostgreSQL database to retrieve all unique - thread IDs associated with a given user_id from the checkpoints table. - When using in-memory storage, returns an empty list since threads - are not persisted. - - Args: - user_id: The unique identifier of the user whose threads to retrieve. - - Returns: - A list of thread IDs (strings) associated with the user. - Returns empty list when using in-memory storage. - - Raises: - HTTPException: If there's a database connection error or query failure. - Status code 500 with error details. - - Note: - This function uses raw SQL queries to extract thread_id from the - checkpoints table where metadata contains the specified user_id. - In-memory storage mode returns empty list as threads are not persisted. - """ - # When using in-memory storage, get threads from thread registry - if settings.USE_INMEMORY_SAVER: - app_logger.info( - f"Using in-memory storage - retrieving threads from registry for user_id: {user_id}" - ) - try: - # Use the thread registry for fast lookup - thread_ids = get_user_threads(user_id) - app_logger.info( - f"Found {len(thread_ids)} threads in registry for user_id: {user_id}: {thread_ids}" - ) - return thread_ids - except Exception as e: - app_logger.error(f"Error accessing thread registry for user {user_id}: {e}") - raise HTTPException( - status_code=500, - detail=f"Failed to retrieve threads from registry: {str(e)}", - ) - - try: - # Connect to the PostgreSQL database - with psycopg2.connect(settings.database_uri) as conn: - cur = conn.cursor() - - # Query for distinct thread IDs where metadata contains the user_id - cur.execute( - f"SELECT distinct thread_id FROM checkpoints where metadata->>'user_id'='{user_id}'" - ) - rows = cur.fetchall() - thread_ids = [row[0] for row in rows] - - app_logger.info(f"Found {len(thread_ids)} threads for user_id: {user_id}") - return thread_ids - - except Exception as e: - app_logger.error( - f"Database error while fetching threads for user {user_id}: {e}" - ) - raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}") diff --git a/template_agent/src/settings.py b/template_agent/src/settings.py deleted file mode 100644 index 6bc69bf6..00000000 --- a/template_agent/src/settings.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Settings configuration for the template agent. - -This module provides centralized configuration management using Pydantic -BaseSettings for environment variable loading, validation, and default -value handling for the template agent service. -""" - -from typing import Optional - -from dotenv import load_dotenv -from pydantic import Field -from pydantic_settings import BaseSettings - -from template_agent.src.core.exceptions.exceptions import AppException, AppExceptionCode -from template_agent.utils.pylogger import get_python_logger - -# Initialize logger -logger = get_python_logger() - -# Load environment variables with error handling -try: - load_dotenv() -except Exception as e: - # Log error but don't fail - environment variables might be set directly - logger.warning(f"Could not load .env file: {e}") - - -class Settings(BaseSettings): - """Configuration settings for the template agent. - - Uses Pydantic BaseSettings to load and validate configuration from - environment variables. Provides default values for optional settings - and validation for required ones. - - The settings are organized into logical groups: - - Server Configuration: Host, port, SSL settings - - Database Configuration: PostgreSQL connection parameters - - Langfuse Configuration: Tracing and analytics settings - - Google Configuration: Service account credentials - - MCP Configuration: MCP server connection settings - """ - - # Server Configuration - AGENT_HOST: str = Field(default="0.0.0.0", json_schema_extra={"env": "AGENT_HOST"}) - AGENT_PORT: int = Field(default=8081, json_schema_extra={"env": "AGENT_PORT"}) - AGENT_SSL_KEYFILE: Optional[str] = Field( - default=None, json_schema_extra={"env": "AGENT_SSL_KEYFILE"} - ) - AGENT_SSL_CERTFILE: Optional[str] = Field( - default=None, json_schema_extra={"env": "AGENT_SSL_CERTFILE"} - ) - PYTHON_LOG_LEVEL: str = Field( - default="INFO", json_schema_extra={"env": "PYTHON_LOG_LEVEL"} - ) - USE_INMEMORY_SAVER: bool = Field( - default=False, json_schema_extra={"env": "USE_INMEMORY_SAVER"} - ) - - # Database Configuration - POSTGRES_USER: str = Field( - default="pgvector", json_schema_extra={"env": "POSTGRES_USER"} - ) - POSTGRES_PASSWORD: str = Field( - default="pgvector", json_schema_extra={"env": "POSTGRES_PASSWORD"} - ) - POSTGRES_DB: str = Field( - default="pgvector", json_schema_extra={"env": "POSTGRES_DB"} - ) - POSTGRES_HOST: str = Field( - default="pgvector", json_schema_extra={"env": "POSTGRES_HOST"} - ) - POSTGRES_PORT: int = Field(default=5432, json_schema_extra={"env": "POSTGRES_PORT"}) - - # Google Service Account Configuration - GOOGLE_SERVICE_ACCOUNT_FILE: Optional[str] = Field( - default=None, json_schema_extra={"env": "GOOGLE_SERVICE_ACCOUNT_FILE"} - ) - - # Langfuse Configuration - LANGFUSE_PUBLIC_KEY: Optional[str] = Field( - default=None, json_schema_extra={"env": "LANGFUSE_PUBLIC_KEY"} - ) - LANGFUSE_SECRET_KEY: Optional[str] = Field( - default=None, json_schema_extra={"env": "LANGFUSE_SECRET_KEY"} - ) - LANGFUSE_BASE_URL: Optional[str] = Field( - default=None, json_schema_extra={"env": "LANGFUSE_BASE_URL"} - ) - LANGFUSE_TRACING_ENVIRONMENT: str = Field( - default="development", json_schema_extra={"env": "LANGFUSE_TRACING_ENVIRONMENT"} - ) - - # Google Application Credentials - GOOGLE_APPLICATION_CREDENTIALS_CONTENT: Optional[str] = Field( - default=None, - json_schema_extra={"env": "GOOGLE_APPLICATION_CREDENTIALS_CONTENT"}, - ) - - # MCP Server Configuration - MCP_SERVER_NAME: str = Field( - default="template-mcp-server", - json_schema_extra={"env": "MCP_SERVER_NAME"}, - ) - MCP_SERVER_URL: str = Field( - default="http://localhost:5001/mcp/", - json_schema_extra={"env": "MCP_SERVER_URL"}, - ) - MCP_TRANSPORT_PROTOCOL: str = Field( - default="streamable_http", - json_schema_extra={"env": "MCP_TRANSPORT_PROTOCOL"}, - ) - MCP_CONNECTION_TIMEOUT: int = Field( - default=30, - json_schema_extra={"env": "MCP_CONNECTION_TIMEOUT"}, - ) - MCP_SSL_VERIFY: bool = Field( - default=False, - json_schema_extra={ - "env": "MCP_SSL_VERIFY", - "description": "Enable SSL certificate verification for MCP connections", - }, - ) - - # Request Logging Configuration - REQUEST_LOGGING_ENABLED: bool = Field( - default=True, - json_schema_extra={ - "env": "REQUEST_LOGGING_ENABLED", - "description": "Enable request/response logging", - }, - ) - REQUEST_LOG_HEADERS: bool = Field( - default=True, - json_schema_extra={ - "env": "REQUEST_LOG_HEADERS", - "description": "Include headers in request/response logs", - }, - ) - REQUEST_LOG_BODY: bool = Field( - default=False, - json_schema_extra={ - "env": "REQUEST_LOG_BODY", - "description": "Include body content in request/response logs", - }, - ) - REQUEST_LOG_BODY_MAX_SIZE: int = Field( - default=10240, - json_schema_extra={ - "env": "REQUEST_LOG_BODY_MAX_SIZE", - "description": "Maximum body size in bytes to log (0 for unlimited)", - }, - ) - - @property - def database_uri(self) -> str: - """Generate database URI from individual components. - - Constructs a PostgreSQL connection URI using the configured - database settings including user, password, host, port, and - database name. - - Returns: - The complete PostgreSQL database URI string. - """ - return f"postgresql://{self.POSTGRES_USER}:{self.POSTGRES_PASSWORD}@{self.POSTGRES_HOST}:{self.POSTGRES_PORT}/{self.POSTGRES_DB}" - - -def validate_config(settings: Settings) -> None: - """Validate configuration settings. - - Performs comprehensive validation to ensure required settings are - present and values are within acceptable ranges. This function - validates port ranges, log levels, and transport protocols. - - Args: - settings: Settings instance to validate. - - Raises: - ValueError: If required configuration is missing or invalid. - """ - # Validate port range - if not (1024 <= settings.AGENT_PORT <= 65535): - logger.error( - f"AGENT_PORT must be between 1024 and 65535, got {settings.AGENT_PORT}" - ) - raise AppException( - f"AGENT_PORT must be between 1024 and 65535, got {settings.AGENT_PORT}", - AppExceptionCode.CONFIGURATION_VALIDATION_ERROR, - ) - - # Validate log level - valid_log_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] - if settings.PYTHON_LOG_LEVEL.upper() not in valid_log_levels: - logger.error( - f"PYTHON_LOG_LEVEL must be one of {valid_log_levels}, got {settings.PYTHON_LOG_LEVEL}" - ) - raise AppException( - f"PYTHON_LOG_LEVEL must be one of {valid_log_levels}, got {settings.PYTHON_LOG_LEVEL}", - AppExceptionCode.CONFIGURATION_VALIDATION_ERROR, - ) - - -# Create settings instance without validation (validation happens in main.py) -settings = Settings() diff --git a/template_agent/utils/google_creds.py b/template_agent/utils/google_creds.py deleted file mode 100644 index c73de849..00000000 --- a/template_agent/utils/google_creds.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Google credentials management utilities. - -This module provides functions for initializing Google Generative AI with various -credential formats including base64-encoded, file paths, and direct JSON content. -""" - -import base64 -import os -import tempfile - -from template_agent.src.settings import settings -from template_agent.utils.pylogger import get_python_logger - -logger = get_python_logger() - - -def initialize_google_genai(): - """Initialize Google Generative AI with service account credentials.""" - credentials_file = None - - if not settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT: - logger.warning("No Google service account credentials configured") - return - - # Check if credentials are provided as base64-encoded environment variable - if settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT.startswith("ewog"): - # Validate that it's valid JSON - import json - - try: - # Decode base64 credentials - credentials_json = base64.b64decode( - settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT - ).decode("utf-8") - - json.loads(credentials_json) # This will raise an exception if invalid JSON - - # Create temporary file with credentials - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False - ) as temp_file: - temp_file.write(credentials_json) - credentials_file = temp_file.name - - logger.info( - "Initialized Google Generative AI with base64-encoded service account credentials" - ) - - except (base64.binascii.Error, UnicodeDecodeError) as e: - logger.error(f"Failed to decode base64 credentials: {e}") - return - except json.JSONDecodeError as e: - logger.error(f"Invalid JSON in base64 credentials: {e}") - return - except Exception as e: - logger.error(f"Unexpected error processing base64 credentials: {e}") - return - - # Check if credentials are provided as file path - elif os.path.exists(settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT): - credentials_file = settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT - logger.info( - f"Initialized Google Generative AI with service account file: {settings.GOOGLE_SERVICE_ACCOUNT_FILE}" - ) - - # Check if credentials are provided as direct JSON content - elif settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT.strip().startswith("{"): - # Validate that it's valid JSON - import json - - try: - credentials_json = settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT.strip() - json.loads(credentials_json) # This will raise an exception if invalid JSON - - # Create temporary file with credentials - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False - ) as temp_file: - temp_file.write(credentials_json) - credentials_file = temp_file.name - - logger.info( - "Initialized Google Generative AI with direct JSON service account credentials" - ) - - except json.JSONDecodeError as e: - logger.error(f"Invalid JSON in direct credentials: {e}") - return - except Exception as e: - logger.error(f"Unexpected error processing direct JSON credentials: {e}") - return - - else: - logger.warning( - f"Google service account credentials not found or invalid format: {settings.GOOGLE_SERVICE_ACCOUNT_FILE[:50]}..." - ) - return - - # Set environment variable for langchain-google-genai to use - if credentials_file: - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_file - logger.debug(f"Set GOOGLE_APPLICATION_CREDENTIALS to: {credentials_file}") diff --git a/template_agent/utils/pylogger.py b/template_agent/utils/pylogger.py deleted file mode 100644 index 956bcfd2..00000000 --- a/template_agent/utils/pylogger.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Structured logger utility for the Template MCP server.""" - -import logging -import sys -from typing import Any, Dict, List, Set - -import structlog - -# HTTP clients -HTTP_CLIENT_LOGGERS = { - "urllib3", - "urllib3.connectionpool", - "urllib3.util", - "urllib3.util.retry", - "requests", - "httpx", -} - -# AWS SDK -AWS_LOGGERS = { - "botocore", - "botocore.client", - "botocore.credentials", - "botocore.httpsession", - "boto3", - "boto3.resources", -} - -# MCP (custom platform) -MCP_LOGGERS = { - "fastmcp", - "fastmcp.server", - "fastmcp.server.http", - "fastmcp.utilities", - "fastmcp.utilities.logging", - "fastmcp.client", - "fastmcp.transports", -} - -# ML/AI frameworks -ML_AI_LOGGERS = { - "sentence_transformers", - "transformers", - "transformers.models", - "transformers.tokenization_utils", - "transformers.tokenization_utils_base", - "transformers.configuration_utils", - "transformers.modeling_utils", - "huggingface_hub", - "huggingface_hub.utils", - "langchain_huggingface", - "torch", - "torch.nn", -} - -# Observability / telemetry -OBSERVABILITY_LOGGERS = { - "langfuse", - "langfuse.client", - "langfuse.api", - "langfuse.callback", -} - -# --- Aggregated Sets --- - -THIRD_PARTY_LOGGERS: Set[str] = ( - HTTP_CLIENT_LOGGERS - | AWS_LOGGERS - | MCP_LOGGERS - | ML_AI_LOGGERS - | OBSERVABILITY_LOGGERS -) - -ERROR_ONLY_LOGGERS: Set[str] = ML_AI_LOGGERS | OBSERVABILITY_LOGGERS - -_LOGGING_CONFIGURED = False - - -# --- Internal helpers --- - - -def _clear_handlers(logger: logging.Logger) -> None: - logger.handlers.clear() - logger.filters.clear() - - -def _setup_logger(logger_name: str, level: str) -> None: - logger = logging.getLogger(logger_name) - _clear_handlers(logger) - logger.setLevel(logging.ERROR if logger_name in ERROR_ONLY_LOGGERS else level) - logger.propagate = True - - -def _configure_third_party_loggers(log_level: str) -> None: - """Apply structured logging to selected third-party loggers.""" - logging.getLogger().handlers.clear() - - for name in THIRD_PARTY_LOGGERS: - _setup_logger(name, log_level) - - -# --- Public API --- - - -def force_reconfigure_all_loggers(log_level: str = "INFO") -> None: - """Force logger reconfiguration, even if already initialized.""" - global _LOGGING_CONFIGURED - _LOGGING_CONFIGURED = False - get_python_logger(log_level) - - -def get_python_logger(log_level: str = "INFO") -> structlog.BoundLogger: - """Get a configured structlog logger.""" - global _LOGGING_CONFIGURED - log_level = log_level.upper() - - if not _LOGGING_CONFIGURED: - logging.basicConfig( - format="%(message)s", - stream=sys.stdout, - level=log_level, - ) - - structlog.configure( - processors=[ - structlog.stdlib.filter_by_level, - structlog.stdlib.add_logger_name, - structlog.stdlib.add_log_level, - structlog.stdlib.PositionalArgumentsFormatter(), - structlog.processors.TimeStamper(fmt="iso"), - structlog.processors.StackInfoRenderer(), - structlog.processors.format_exc_info, - structlog.processors.UnicodeDecoder(), - structlog.processors.JSONRenderer(), - ], - context_class=dict, - logger_factory=structlog.stdlib.LoggerFactory(), - wrapper_class=structlog.stdlib.BoundLogger, - cache_logger_on_first_use=True, - ) - - _LOGGING_CONFIGURED = True - - _configure_third_party_loggers(log_level) - return structlog.get_logger() - - -def get_uvicorn_log_config(log_level: str = "INFO") -> Dict[str, Any]: - """Return a Uvicorn-compatible logging config that integrates with structlog.""" - log_level = log_level.upper() - default_formatter = { - "()": "structlog.stdlib.ProcessorFormatter", - "processor": structlog.processors.JSONRenderer(), - "foreign_pre_chain": [ - structlog.stdlib.add_log_level, - structlog.processors.TimeStamper(fmt="iso"), - structlog.processors.StackInfoRenderer(), - structlog.processors.format_exc_info, - structlog.processors.UnicodeDecoder(), - ], - } - - def make_logger_config(names: List[str], level: str) -> Dict[str, Any]: - return { - name: { - "handlers": ["default"], - "level": level, - "propagate": False, - } - for name in names - } - - # Base uvicorn loggers - base_loggers = ["", "uvicorn", "uvicorn.error", "uvicorn.asgi", "uvicorn.protocols"] - access_loggers = ["uvicorn.access"] - - return { - "version": 1, - "disable_existing_loggers": False, - "formatters": { - "default": default_formatter, - "access": default_formatter, - }, - "handlers": { - "default": { - "formatter": "default", - "class": "logging.StreamHandler", - "stream": "ext://sys.stdout", - }, - "access": { - "formatter": "access", - "class": "logging.StreamHandler", - "stream": "ext://sys.stdout", - }, - }, - "loggers": { - **make_logger_config(base_loggers, log_level), - **make_logger_config(access_loggers, log_level), - **make_logger_config( - list(THIRD_PARTY_LOGGERS - ERROR_ONLY_LOGGERS), log_level - ), - **make_logger_config(list(ERROR_ONLY_LOGGERS), "ERROR"), - }, - } diff --git a/tests/__init__.py b/tests/__init__.py index cbed33da..e69de29b 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1 +0,0 @@ -"""Tests for the template agent.""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..df889ee4 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,157 @@ +"""Root test configuration and shared fixtures. + +Ensures the project root is on sys.path so both ``deep_agent`` and +``aegra`` packages are importable in all test modules. + +Provides: +- Mock LLM fixture (MR-58) +- Mock DB / Postgres fixtures (MR-57) +- Stream context fixture +- Settings override fixture +""" + +import sys +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +_PROJECT_ROOT = str(Path(__file__).resolve().parent.parent) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + + +# ── Stream context fixture ─────────────────────────────────────── + + +@pytest.fixture() +def stream_context(): + """Provide a standard StreamContext for unit tests.""" + from deep_agent.src.streaming import StreamContext + + return StreamContext( + run_id="test_run_1", + trace_id="test_trace_1", + thread_id="test_thread_1", + session_id="test_session_1", + user_id="test_user", + stream_tokens=True, + ) + + +# ── Mock LLM fixture (MR-58) ──────────────────────────────────── + + +@pytest.fixture() +def mock_llm(): + """Return a MagicMock that behaves like a LangChain BaseChatModel. + + Supports both sync and async invocation paths. The default + response is a simple AIMessage; override ``mock_llm.invoke.return_value`` + in individual tests to customise. + """ + from langchain_core.messages import AIMessage + + llm = MagicMock() + default_response = AIMessage(content="mock llm response", id="mock_msg_1") + + llm.invoke.return_value = default_response + llm.ainvoke = AsyncMock(return_value=default_response) + llm.bind_tools.return_value = llm + llm.with_structured_output.return_value = llm + llm.model_name = "mock-model" + + return llm + + +# ── Mock DB / Postgres fixtures (MR-57) ───────────────────────── + + +@pytest.fixture() +def mock_db_uri() -> str: + """Return a fake Postgres URI for unit tests (no real connection).""" + return "postgresql://test:test@localhost:5432/testdb" + + +@pytest.fixture() +def mock_async_connection(): + """Return a mock ``psycopg.AsyncConnection`` context manager. + + Usage in tests:: + + async with mock_async_connection as conn: + conn.execute.return_value = cursor_mock + """ + conn = AsyncMock() + cursor = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[]) + cursor.fetchone = AsyncMock(return_value=None) + cursor.rowcount = 0 + conn.execute = AsyncMock(return_value=cursor) + conn.commit = AsyncMock() + + ctx = AsyncMock() + ctx.__aenter__ = AsyncMock(return_value=conn) + ctx.__aexit__ = AsyncMock(return_value=False) + + conn._cursor = cursor + conn._ctx = ctx + return conn + + +@pytest.fixture() +def mock_checkpointer(): + """Return a mock async checkpointer (PostgresSaver-like).""" + cp = AsyncMock() + cp.setup = AsyncMock() + cp.__aenter__ = AsyncMock(return_value=cp) + cp.__aexit__ = AsyncMock(return_value=False) + return cp + + +# ── Settings override fixture ──────────────────────────────────── + + +@pytest.fixture() +def test_settings(): + """Return a Settings instance with safe test defaults. + + Patches POSTGRES_HOST to localhost so no accidental remote connections. + """ + from deep_agent.src.settings import Settings + + return Settings( + AGENT_HOST="127.0.0.1", + AGENT_PORT=5099, + POSTGRES_HOST="localhost", + POSTGRES_PORT=5432, + POSTGRES_USER="test", + POSTGRES_PASSWORD="test", + POSTGRES_DB="testdb", + PYTHON_LOG_LEVEL="WARNING", + ) + + +# ── Agent fixtures ─────────────────────────────────────────────── + + +@pytest.fixture() +def mock_agent_config(): + """Return a mock agent_config with a minimal orchestrator config.""" + config = MagicMock() + config.get_orchestrator_config.return_value = { + "name": "test-orchestrator", + "model": "mock-model", + "body": "You are a test agent.", + "skill_paths": [], + "tools": [], + } + config.resolve_tools.return_value = [] + return config + + +@pytest.fixture() +def mock_mcp_tools() -> list[Any]: + """Return an empty list of MCP tools for unit tests.""" + return [] diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration/aegra/conftest.py b/tests/integration/aegra/conftest.py new file mode 100644 index 00000000..22728ce5 --- /dev/null +++ b/tests/integration/aegra/conftest.py @@ -0,0 +1,139 @@ +"""Shared test fixtures for aegra integration tests (MR-34). + +Provides: +- Mock MCP server (in-process via httpx) +- LangGraph API client fixture +- Thread/run management helpers +- State snapshot assertions +""" + +import asyncio +import json +import os +import sys +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +_PROJECT_ROOT = str(Path(__file__).resolve().parent.parent.parent.parent) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + + +MOCK_MCP_URL = "http://mock-mcp:5001" +LANGGRAPH_API_URL = os.environ.get("LANGGRAPH_API_URL", "http://127.0.0.1:2024") + + +@pytest.fixture(scope="session") +def event_loop(): + """Create a session-scoped event loop for async tests.""" + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +@pytest.fixture() +def mock_bmi_response() -> dict[str, Any]: + """Standard BMI calculation response for a normal-weight person.""" + return { + "success": True, + "bmi": 24.7, + "category": "Normal", + "height_cm": 180, + "weight_kg": 80, + } + + +@pytest.fixture() +def mock_email_response() -> dict[str, Any]: + """Standard email send response.""" + return { + "success": True, + "recipient": "test@example.com", + "subject": "BMI Report", + "message": "Email sent successfully to test@example.com", + "message_id": "mock-12345", + } + + +@pytest.fixture() +def mock_search_response() -> dict[str, Any]: + """Standard web search response with health tips.""" + return { + "success": True, + "query": "normal BMI health tips", + "category": "Normal", + "results": [ + {"title": "Tip 1", "snippet": "Maintain a balanced diet"}, + {"title": "Tip 2", "snippet": "Exercise 150 min/week"}, + {"title": "Tip 3", "snippet": "Stay hydrated"}, + ], + } + + +@pytest.fixture() +def mock_validate_email_response() -> dict[str, Any]: + """Standard email validation response.""" + return { + "success": True, + "valid": True, + "email": "test@example.com", + "message": "Valid email format", + } + + +@pytest.fixture() +def sample_thread_id() -> str: + return "test-thread-001" + + +@pytest.fixture() +def sample_user_id() -> str: + return "test-user-001" + + +@pytest.fixture() +def sample_bmi_input() -> dict[str, Any]: + """Standard BMI request payload for the agent.""" + return { + "messages": [ + { + "role": "human", + "content": ( + "Calculate BMI for someone who is 180cm tall and weighs 80kg. " + "Send the report to test@example.com" + ), + } + ] + } + + +@pytest.fixture() +def sample_email_input() -> dict[str, Any]: + """Email-only request payload for testing the publisher subagent.""" + return { + "messages": [ + { + "role": "human", + "content": "Send this report to test@example.com: BMI is 24.7, Normal weight.", + } + ] + } + + +@pytest.fixture() +def langgraph_api_url() -> str: + """Base URL for the LangGraph API server.""" + return LANGGRAPH_API_URL + + +@pytest.fixture() +def api_headers() -> dict[str, str]: + """Default headers for LangGraph API requests.""" + return { + "Content-Type": "application/json", + "Accept": "application/json", + } diff --git a/tests/integration/aegra/test_bmi_skill.py b/tests/integration/aegra/test_bmi_skill.py new file mode 100644 index 00000000..bbfba39e --- /dev/null +++ b/tests/integration/aegra/test_bmi_skill.py @@ -0,0 +1,85 @@ +"""Integration test: BMI skill flow via aegra (MR-29). + +Verifies that the agent correctly: +1. Receives a BMI request +2. Calls the calculate_bmi MCP tool +3. Calls search_web for health tips +4. Returns a formatted BMI report + +Requires: mock MCP server running on localhost:5001 OR +uses mocked tool responses via fixtures. +""" + +import pytest + +from deep_agent.aegra.converters import ( + extract_final_response, + stream_request_to_langgraph_input, +) +from deep_agent.aegra.serialization import deserialize_message, serialize_message +from langchain_core.messages import AIMessage, HumanMessage + + +class TestBMISkillConverters: + """Test that BMI-related messages are correctly serialized through aegra.""" + + def test_bmi_request_converts_to_langgraph_input(self): + result = stream_request_to_langgraph_input("Calculate BMI for 180cm and 80kg") + assert len(result["messages"]) == 1 + assert isinstance(result["messages"][0], HumanMessage) + assert "180cm" in result["messages"][0].content + + def test_bmi_response_serialization_roundtrip(self): + ai_msg = AIMessage( + content="Your BMI is 24.7 (Normal). Here are health tips...", + tool_calls=[ + { + "id": "tc1", + "name": "calculate_bmi", + "args": {"height_cm": 180, "weight_kg": 80}, + } + ], + ) + serialized = serialize_message(ai_msg) + assert serialized["type"] == "ai" + assert serialized["tool_calls"][0]["name"] == "calculate_bmi" + + restored = deserialize_message(serialized) + assert isinstance(restored, AIMessage) + assert restored.tool_calls[0]["name"] == "calculate_bmi" + + def test_extract_bmi_report_from_state(self): + state = { + "messages": [ + HumanMessage(content="Calculate BMI for 180cm and 80kg"), + AIMessage(content=""), + AIMessage( + content="**BMI Report**\nBMI: 24.7\nCategory: Normal\n\nHealth Tips:\n1. Stay active" + ), + ] + } + response = extract_final_response(state) + assert response is not None + assert "24.7" in response + assert "Normal" in response + + +class TestBMIToolCallStructure: + """Validate the expected tool call structure for BMI calculations.""" + + def test_calculate_bmi_tool_call_shape(self, mock_bmi_response): + assert mock_bmi_response["success"] is True + assert isinstance(mock_bmi_response["bmi"], float) + assert mock_bmi_response["category"] in [ + "Underweight", + "Normal", + "Overweight", + "Obese", + ] + + def test_search_web_tool_call_shape(self, mock_search_response): + assert mock_search_response["success"] is True + assert len(mock_search_response["results"]) == 3 + for result in mock_search_response["results"]: + assert "title" in result + assert "snippet" in result diff --git a/tests/integration/aegra/test_e2e.py b/tests/integration/aegra/test_e2e.py new file mode 100644 index 00000000..cf19cbe0 --- /dev/null +++ b/tests/integration/aegra/test_e2e.py @@ -0,0 +1,127 @@ +"""End-to-end test: Full aegra deployment (MR-32). + +Tests the complete LangGraph Platform API contract by verifying +health, thread creation, agent invocation, and state retrieval. + +Requires: ``langgraph dev`` or ``langgraph up`` running on LANGGRAPH_API_URL. +Mark: ``pytest -m e2e`` to run these tests separately. +""" + +import os + +import httpx +import pytest + +pytestmark = pytest.mark.e2e + +LANGGRAPH_API_URL = os.environ.get("LANGGRAPH_API_URL", "http://127.0.0.1:2024") +ASSISTANT_ID = "agent" + + +def _api_url(path: str) -> str: + return f"{LANGGRAPH_API_URL}{path}" + + +@pytest.fixture() +def client(): + with httpx.Client(base_url=LANGGRAPH_API_URL, timeout=60) as c: + yield c + + +class TestAegraHealthEndpoint: + """Verify the LangGraph Platform health endpoint.""" + + def test_health_ok(self, client): + resp = client.get("/ok") + assert resp.status_code == 200 + + def test_info_endpoint(self, client): + resp = client.get("/info") + assert resp.status_code == 200 + data = resp.json() + assert "version" in data + + +class TestAegraAssistants: + """Verify assistants are registered correctly.""" + + def test_list_assistants(self, client): + resp = client.post("/assistants/search", json={}) + assert resp.status_code == 200 + assistants = resp.json() + assert len(assistants) >= 1 + + def test_agent_assistant_exists(self, client): + resp = client.get(f"/assistants/{ASSISTANT_ID}") + assert resp.status_code == 200 + data = resp.json() + assert data["assistant_id"] == ASSISTANT_ID + + +class TestAegraThreadLifecycle: + """Verify thread creation, retrieval, and deletion.""" + + def test_create_thread(self, client): + resp = client.post("/threads", json={}) + assert resp.status_code == 200 + thread = resp.json() + assert "thread_id" in thread + + def test_create_and_get_thread(self, client): + create_resp = client.post("/threads", json={}) + thread_id = create_resp.json()["thread_id"] + + get_resp = client.get(f"/threads/{thread_id}") + assert get_resp.status_code == 200 + assert get_resp.json()["thread_id"] == thread_id + + def test_delete_thread(self, client): + create_resp = client.post("/threads", json={}) + thread_id = create_resp.json()["thread_id"] + + del_resp = client.delete(f"/threads/{thread_id}") + assert del_resp.status_code == 200 + + +class TestAegraAgentInvocation: + """Test agent invocation via the LangGraph API. + + These tests exercise the actual agent graph — they require + valid Google credentials and a running mock MCP server. + """ + + @pytest.mark.slow + def test_invoke_returns_response(self, client): + thread_resp = client.post("/threads", json={}) + thread_id = thread_resp.json()["thread_id"] + + resp = client.post( + f"/threads/{thread_id}/runs", + json={ + "assistant_id": ASSISTANT_ID, + "input": { + "messages": [ + {"role": "human", "content": "Hello, what can you do?"} + ] + }, + }, + ) + assert resp.status_code in (200, 201, 202) + + @pytest.mark.slow + def test_stream_returns_events(self, client): + thread_resp = client.post("/threads", json={}) + thread_id = thread_resp.json()["thread_id"] + + with client.stream( + "POST", + f"/threads/{thread_id}/runs/stream", + json={ + "assistant_id": ASSISTANT_ID, + "input": {"messages": [{"role": "human", "content": "Say hello"}]}, + "stream_mode": "updates", + }, + ) as resp: + assert resp.status_code == 200 + events = list(resp.iter_lines()) + assert len(events) > 0 diff --git a/tests/integration/aegra/test_email_skill.py b/tests/integration/aegra/test_email_skill.py new file mode 100644 index 00000000..c93ad113 --- /dev/null +++ b/tests/integration/aegra/test_email_skill.py @@ -0,0 +1,59 @@ +"""Integration test: Email skill flow via aegra (MR-30). + +Verifies that the agent correctly: +1. Validates email addresses via MCP +2. Formats reports into Gmail-compatible HTML +3. Sends email via the send_email MCP tool +""" + +import pytest + +from deep_agent.aegra.converters import stream_request_to_langgraph_input +from deep_agent.aegra.serialization import serialize_message +from langchain_core.messages import AIMessage, HumanMessage + + +class TestEmailSkillConverters: + """Test email-related message flows through aegra serialization.""" + + def test_email_request_converts_to_langgraph_input(self): + result = stream_request_to_langgraph_input( + "Send this BMI report to test@example.com" + ) + assert "test@example.com" in result["messages"][0].content + + def test_email_tool_call_serialization(self): + ai_msg = AIMessage( + content="I've sent the report to test@example.com", + tool_calls=[ + { + "id": "tc-email", + "name": "send_email", + "args": { + "recipient": "test@example.com", + "subject": "BMI Report", + "body": "

BMI Report

", + }, + } + ], + ) + serialized = serialize_message(ai_msg) + assert serialized["tool_calls"][0]["name"] == "send_email" + assert serialized["tool_calls"][0]["args"]["recipient"] == "test@example.com" + + def test_validate_email_tool_call_structure(self, mock_validate_email_response): + assert mock_validate_email_response["valid"] is True + assert mock_validate_email_response["email"] == "test@example.com" + + +class TestEmailResponseStructure: + """Validate email send response shapes.""" + + def test_successful_email_response(self, mock_email_response): + assert mock_email_response["success"] is True + assert "message_id" in mock_email_response + assert mock_email_response["recipient"] == "test@example.com" + + def test_email_response_has_required_fields(self, mock_email_response): + required = {"success", "recipient", "subject", "message", "message_id"} + assert required.issubset(mock_email_response.keys()) diff --git a/tests/integration/aegra/test_graceful_shutdown.py b/tests/integration/aegra/test_graceful_shutdown.py new file mode 100644 index 00000000..c5665b55 --- /dev/null +++ b/tests/integration/aegra/test_graceful_shutdown.py @@ -0,0 +1,184 @@ +"""Integration tests for graceful shutdown under concurrent load. + +Proves the shutdown sequence completes cleanly while simulated +graph runs are active, and that all subsystems are torn down +within the configured timeout budget. +""" + +import asyncio +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import deep_agent.aegra.shutdown as shutdown_mod + +pytestmark = pytest.mark.integration + + +@pytest.fixture(autouse=True) +def _reset_shutdown_state(): + shutdown_mod._shutting_down = False + shutdown_mod._shutdown_complete = False + shutdown_mod._async_shutdown_started = False + yield + shutdown_mod._shutting_down = False + shutdown_mod._shutdown_complete = False + shutdown_mod._async_shutdown_started = False + + +def _mock_all_subsystems(drain_seconds=0): + """Context manager that mocks all external subsystems for shutdown.""" + mock_langfuse = MagicMock() + mock_langfuse.shutdown = MagicMock() + + return ( + patch.object(shutdown_mod, "SHUTDOWN_DRAIN_SECONDS", drain_seconds), + patch( + "deep_agent.aegra.telemetry.get_langfuse_client", + return_value=mock_langfuse, + ), + patch( + "deep_agent.src.memory.scheduler.stop_scheduler", + new_callable=AsyncMock, + ), + patch("deep_agent.aegra.redis.close_redis_client"), + ) + + +class TestShutdownUnderConcurrentActivity: + async def test_shutdown_while_tasks_are_running(self): + """Simulate concurrent graph runs during shutdown. + + Active tasks should be able to continue during the drain + period. After drain, cleanup runs and completes. + """ + completed_tasks = [] + + async def simulate_graph_run(task_id: int, duration: float): + await asyncio.sleep(duration) + completed_tasks.append(task_id) + + tasks = [asyncio.create_task(simulate_graph_run(i, i * 0.05)) for i in range(5)] + + patches = _mock_all_subsystems(drain_seconds=0.3) + with patches[0], patches[1], patches[2], patches[3]: + t0 = time.monotonic() + result = await shutdown_mod.run_shutdown() + elapsed = time.monotonic() - t0 + + assert result["drain"] == "ok" + assert result["langfuse"] == "ok" + assert result["scheduler"] == "ok" + assert result["redis"] == "ok" + assert shutdown_mod._shutdown_complete is True + assert elapsed < 2.0 + + await asyncio.gather(*tasks, return_exceptions=True) + assert len(completed_tasks) == 5 + + async def test_resources_cleaned_up_after_drain(self): + """After drain period, all subsystems are torn down.""" + mock_langfuse = MagicMock() + mock_langfuse.shutdown = MagicMock() + mock_stop = AsyncMock() + mock_close = MagicMock() + + with ( + patch.object(shutdown_mod, "SHUTDOWN_DRAIN_SECONDS", 0), + patch( + "deep_agent.aegra.telemetry.get_langfuse_client", + return_value=mock_langfuse, + ), + patch( + "deep_agent.src.memory.scheduler.stop_scheduler", + mock_stop, + ), + patch("deep_agent.aegra.redis.close_redis_client", mock_close), + ): + await shutdown_mod.run_shutdown() + + mock_langfuse.shutdown.assert_called_once() + mock_stop.assert_awaited_once() + mock_close.assert_called_once() + + +class TestHealthDuringShutdown: + async def test_health_returns_503_during_shutdown(self): + from deep_agent.aegra.health import health_response + + with patch( + "deep_agent.aegra.health.get_health_status", + new_callable=AsyncMock, + return_value={"status": "healthy"}, + ): + code_before, _ = await health_response() + assert code_before == 200 + + shutdown_mod._shutting_down = True + + code_after, body = await health_response() + assert code_after == 503 + assert body["status"] == "shutting_down" + + +class TestLangfuseTimeoutResilience: + async def test_slow_langfuse_does_not_block_shutdown(self): + """A hanging Langfuse server must not prevent Redis/scheduler cleanup.""" + + def slow_shutdown(): + time.sleep(30) + + mock_langfuse = MagicMock() + mock_langfuse.shutdown = slow_shutdown + mock_close = MagicMock() + mock_stop = AsyncMock() + + with ( + patch.object(shutdown_mod, "SHUTDOWN_DRAIN_SECONDS", 0), + patch.object(shutdown_mod, "SHUTDOWN_LANGFUSE_TIMEOUT_SECONDS", 0.2), + patch( + "deep_agent.aegra.telemetry.get_langfuse_client", + return_value=mock_langfuse, + ), + patch( + "deep_agent.src.memory.scheduler.stop_scheduler", + mock_stop, + ), + patch("deep_agent.aegra.redis.close_redis_client", mock_close), + ): + t0 = time.monotonic() + result = await shutdown_mod.run_shutdown() + elapsed = time.monotonic() - t0 + + assert result["langfuse"] == "timeout" + assert result["scheduler"] == "ok" + assert result["redis"] == "ok" + mock_stop.assert_awaited_once() + mock_close.assert_called_once() + assert elapsed < 3.0 + + +class TestIdempotentConcurrentShutdown: + async def test_concurrent_calls_execute_once(self): + """Two concurrent run_shutdown() calls should only execute steps once.""" + call_count = 0 + + async def counting_drain(): + nonlocal call_count + call_count += 1 + await asyncio.sleep(0.05) + return "ok" + + patches = _mock_all_subsystems(drain_seconds=0) + with patches[0], patches[1], patches[2], patches[3]: + with patch.object(shutdown_mod, "_drain", side_effect=counting_drain): + results = await asyncio.gather( + shutdown_mod.run_shutdown(), + shutdown_mod.run_shutdown(), + ) + + real_runs = [r for r in results if "drain" in r] + skipped = [r for r in results if r.get("status") == "already_complete"] + assert len(real_runs) == 1 + assert len(skipped) == 1 diff --git a/tests/integration/aegra/test_subagent_flow.py b/tests/integration/aegra/test_subagent_flow.py new file mode 100644 index 00000000..e7d4e82e --- /dev/null +++ b/tests/integration/aegra/test_subagent_flow.py @@ -0,0 +1,126 @@ +"""Integration test: Subagent orchestration flow via aegra (MR-31). + +Verifies the full orchestrator -> analyst -> publisher delegation chain: +1. User requests BMI analysis + email +2. Orchestrator delegates to analyst subagent +3. Analyst calculates BMI and searches for tips +4. Orchestrator delegates to publisher subagent +5. Publisher formats and sends the email +""" + +import pytest + +from deep_agent.aegra.serialization import serialize_state, deserialize_state +from deep_agent.aegra.state import AegraMetadata, serialize_metadata +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + + +class TestSubagentFlowSerialization: + """Test full multi-agent conversation state serialization.""" + + def test_full_conversation_state_roundtrip(self): + """A realistic multi-turn conversation with tool calls survives serialization.""" + state = { + "messages": [ + HumanMessage( + content="Calculate BMI for 175cm, 70kg and email to test@example.com" + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tc1", + "name": "calculate_bmi", + "args": {"height_cm": 175, "weight_kg": 70}, + } + ], + ), + ToolMessage( + content='{"bmi": 22.9, "category": "Normal"}', + tool_call_id="tc1", + name="calculate_bmi", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tc2", + "name": "search_web", + "args": {"query": "normal BMI health tips"}, + } + ], + ), + ToolMessage( + content='{"results": [{"snippet": "Stay active"}]}', + tool_call_id="tc2", + name="search_web", + ), + AIMessage(content="BMI: 22.9 (Normal). Tips: Stay active."), + AIMessage( + content="", + tool_calls=[ + { + "id": "tc3", + "name": "send_email", + "args": { + "recipient": "test@example.com", + "subject": "BMI Report", + "body": "

Your BMI: 22.9

", + }, + } + ], + ), + ToolMessage( + content='{"success": true}', + tool_call_id="tc3", + name="send_email", + ), + AIMessage(content="Report sent to test@example.com!"), + ], + } + + serialized = serialize_state(state) + assert "_serialized_at" in serialized + assert len(serialized["messages"]) == 9 + + restored = deserialize_state(serialized) + assert len(restored["messages"]) == 9 + assert isinstance(restored["messages"][0], HumanMessage) + assert isinstance(restored["messages"][1], AIMessage) + assert isinstance(restored["messages"][2], ToolMessage) + assert restored["messages"][2].tool_call_id == "tc1" + + def test_metadata_tracking_across_subagents(self): + meta: AegraMetadata = { + "run_id": "run-orchestrator", + "thread_id": "thread-main", + "error_count": 0, + "last_error": None, + } + serialized = serialize_metadata(meta) + assert "last_error" not in serialized + assert serialized["error_count"] == 0 + + +class TestSubagentDelegationPatterns: + """Verify expected patterns in multi-agent tool call sequences.""" + + def test_analyst_requires_bmi_tools(self): + """Analyst subagent must use calculate_bmi and search_web.""" + analyst_tools = {"calculate_bmi", "search_web"} + expected_call_sequence = ["calculate_bmi", "search_web"] + + for tool_name in expected_call_sequence: + assert tool_name in analyst_tools + + def test_publisher_requires_email_tools(self): + """Publisher subagent must use send_email.""" + publisher_tools = {"send_email"} + assert "send_email" in publisher_tools + + def test_orchestrator_delegates_to_both(self): + """Orchestrator should delegate BMI+email tasks to both subagents.""" + subagent_names = {"analyst", "publisher"} + assert len(subagent_names) == 2 + assert "analyst" in subagent_names + assert "publisher" in subagent_names diff --git a/tests/integration/test_production_hardening.py b/tests/integration/test_production_hardening.py new file mode 100644 index 00000000..f1fa7f41 --- /dev/null +++ b/tests/integration/test_production_hardening.py @@ -0,0 +1,49 @@ +"""Integration tests for production security hardening.""" + +from unittest.mock import patch + +import pytest +from starlette.testclient import TestClient + + +@pytest.fixture +def prod_client(): + """Create a test client with production environment.""" + with patch.dict("os.environ", {"ENVIRONMENT": "production", "ENABLE_AUTH": "true"}): + from deep_agent.aegra.http_app import app + + return TestClient(app) + + +def test_all_security_headers_present_in_production(prod_client): + """Test that all security headers are present in production responses.""" + # Try to access root endpoint (may return 404 but should have headers) + response = prod_client.get("/") + + required_headers = [ + "X-Content-Type-Options", + "X-Frame-Options", + "X-XSS-Protection", + "Strict-Transport-Security", + "Content-Security-Policy", + "Referrer-Policy", + "Permissions-Policy", + ] + + for header in required_headers: + assert header in response.headers, f"Missing security header: {header}" + + +def test_production_mode_config_validation(): + """Test that production mode validates configuration at startup.""" + from deep_agent.src.settings import Settings, validate_config + + # Production with auth disabled should fail validation + prod_settings = Settings(ENVIRONMENT="production", ENABLE_AUTH=False) + + with pytest.raises(Exception, match="ENABLE_AUTH must be true"): + validate_config(prod_settings) + + # Production with auth enabled should pass + prod_settings_valid = Settings(ENVIRONMENT="production", ENABLE_AUTH=True) + validate_config(prod_settings_valid) # Should not raise diff --git a/tests/integration/triggers/__init__.py b/tests/integration/triggers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration/triggers/test_end_to_end.py b/tests/integration/triggers/test_end_to_end.py new file mode 100644 index 00000000..a94abf97 --- /dev/null +++ b/tests/integration/triggers/test_end_to_end.py @@ -0,0 +1,208 @@ +"""End-to-end integration tests for the trigger pipeline. + +Exercises the full path: trigger source -> EventTriggerMiddleware -> +(mocked graph) -> output sink. The graph is always mocked — only the +trigger sources, middleware orchestration, and sinks use real I/O. +""" + +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from unittest.mock import AsyncMock + +import httpx +import pytest + +from deep_agent.src.triggers.config import ( + HeadlessConfig, + OutputSinkConfig, + QueueTriggerConfig, + TriggerConfig, + WebhookTriggerConfig, +) +from deep_agent.src.triggers.middleware import EventTriggerMiddleware + +pytestmark = pytest.mark.integration + +_REDIS_URL = "redis://localhost:6379/0" + + +# ------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------ + + +def _read_jsonl(path: Path) -> list[dict]: + """Read a JSONL file and return parsed lines.""" + if not path.exists(): + return [] + lines = path.read_text().strip().splitlines() + return [json.loads(line) for line in lines if line.strip()] + + +def _mock_graph(return_value: dict | None = None) -> AsyncMock: + """Create a mock graph whose ainvoke returns a simple dict.""" + graph = AsyncMock() + graph.ainvoke.return_value = return_value or { + "messages": [{"role": "assistant", "content": "ok"}] + } + return graph + + +# ------------------------------------------------------------------ +# Tests +# ------------------------------------------------------------------ + + +class TestWebhookToFileSink: + """Webhook POST -> middleware -> mocked graph -> file sink.""" + + async def test_webhook_event_produces_file_output(self, tmp_path: Path): + """POST to webhook -> graph invoked -> result written to JSONL file.""" + output_file = tmp_path / "output.jsonl" + + config = HeadlessConfig( + triggers=TriggerConfig( + webhook=WebhookTriggerConfig( + enabled=True, host="127.0.0.1", port=0, path="/trigger" + ), + ), + output_sinks=[OutputSinkConfig(type="file", path=str(output_file))], + drain_timeout=5.0, + ) + graph = _mock_graph({"result": "webhook-ok"}) + mw = EventTriggerMiddleware(config=config, graph=graph, redis_url=_REDIS_URL) + + await mw.start() + + # Resolve the actual port the webhook source bound to. + webhook_source = mw._sources[0] + port = webhook_source._server.sockets[0].getsockname()[1] + + try: + async with httpx.AsyncClient() as client: + resp = await client.post( + f"http://127.0.0.1:{port}/trigger", + json={"event": "e2e-test", "input": "hello"}, + ) + assert resp.status_code == 200 + + # Allow the middleware processing loop time to invoke the graph and emit. + await asyncio.sleep(0.5) + finally: + await mw.stop() + + # Verify the graph was called with the event payload. + graph.ainvoke.assert_awaited_once() + call_args = graph.ainvoke.call_args[0][0] + content = json.loads(call_args["messages"][0]["content"]) + assert content["input"] == "hello" + + # Verify the output file contains one JSONL entry. + results = _read_jsonl(output_file) + assert len(results) == 1 + assert results[0]["success"] is True + assert results[0]["event"]["name"] == "e2e-test" + assert results[0]["output"] == {"result": "webhook-ok"} + + +class TestQueueToFileSink: + """Redis Stream -> middleware -> mocked graph -> file sink.""" + + async def test_queue_event_produces_file_output(self, tmp_path: Path): + """Push message to Redis Stream -> graph invoked -> result in JSONL.""" + aioredis = pytest.importorskip("redis.asyncio") + client = aioredis.from_url(_REDIS_URL, decode_responses=True) + try: + await client.ping() + except Exception: + pytest.skip("Redis not available at localhost:6379") + + from uuid import uuid4 + + stream = f"e2e-tasks-{uuid4().hex[:8]}" + group = f"e2e-workers-{uuid4().hex[:8]}" + output_file = tmp_path / "queue-output.jsonl" + + config = HeadlessConfig( + triggers=TriggerConfig( + queue=QueueTriggerConfig( + enabled=True, + backend="redis_streams", + stream=stream, + consumer_group=group, + consumer_name="e2e-worker", + ), + ), + output_sinks=[OutputSinkConfig(type="file", path=str(output_file))], + drain_timeout=5.0, + ) + graph = _mock_graph({"result": "queue-ok"}) + mw = EventTriggerMiddleware(config=config, graph=graph, redis_url=_REDIS_URL) + + await mw.start() + + try: + # Produce a message into the stream. + await client.xadd(stream, {"name": "queue-e2e", "data": "world"}) + + # Wait for consumption + graph invocation + sink write. + await asyncio.sleep(1.0) + finally: + await mw.stop() + await client.delete(stream) + await client.aclose() + + graph.ainvoke.assert_awaited_once() + + results = _read_jsonl(output_file) + assert len(results) == 1 + assert results[0]["success"] is True + assert results[0]["event"]["name"] == "queue-e2e" + assert results[0]["output"] == {"result": "queue-ok"} + + +class TestGraphErrorProducesFailedResult: + """Graph raises -> TriggerResult with success=False.""" + + async def test_graph_error_results_in_failure(self, tmp_path: Path): + """When the graph raises, the result has success=False and error message.""" + output_file = tmp_path / "error-output.jsonl" + + config = HeadlessConfig( + triggers=TriggerConfig( + webhook=WebhookTriggerConfig( + enabled=True, host="127.0.0.1", port=0, path="/trigger" + ), + ), + output_sinks=[OutputSinkConfig(type="file", path=str(output_file))], + drain_timeout=5.0, + ) + graph = AsyncMock() + graph.ainvoke.side_effect = RuntimeError("model unavailable") + mw = EventTriggerMiddleware(config=config, graph=graph, redis_url=_REDIS_URL) + + await mw.start() + + webhook_source = mw._sources[0] + port = webhook_source._server.sockets[0].getsockname()[1] + + try: + async with httpx.AsyncClient() as client: + resp = await client.post( + f"http://127.0.0.1:{port}/trigger", + json={"event": "error-test", "x": 1}, + ) + assert resp.status_code == 200 + + await asyncio.sleep(0.5) + finally: + await mw.stop() + + results = _read_jsonl(output_file) + assert len(results) == 1 + assert results[0]["success"] is False + assert "model unavailable" in results[0]["error"] + assert results[0]["event"]["name"] == "error-test" diff --git a/tests/integration/triggers/test_headless_startup.py b/tests/integration/triggers/test_headless_startup.py new file mode 100644 index 00000000..8f8d2f38 --- /dev/null +++ b/tests/integration/triggers/test_headless_startup.py @@ -0,0 +1,267 @@ +"""Integration tests for headless process startup and configuration loading. + +These tests exercise ``deep_agent.headless`` — the headless worker entry +point — by testing config loading/validation and the ``main()`` lifecycle +with external dependencies (DB, model provider) mocked out. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from deep_agent.src.triggers.config import ( + AgentMode, + HeadlessConfig, + HealthCheckConfig, + OutputSinkConfig, + TriggerConfig, + WebhookTriggerConfig, +) + +pytestmark = pytest.mark.integration + + +# ------------------------------------------------------------------ +# Config loading tests +# ------------------------------------------------------------------ + + +class TestLoadHeadlessConfig: + """Test _load_headless_config() validation logic.""" + + def test_headless_mode_config_loads_successfully(self, tmp_path: Path): + """A YAML with mode=headless produces a valid HeadlessConfig.""" + import yaml + + config_data = { + "mode": "headless", + "triggers": { + "webhook": {"enabled": True, "port": 9999, "path": "/trigger"}, + "queue": {"enabled": False}, + "cron": {"enabled": False}, + }, + "output_sinks": [], + } + config_file = tmp_path / "agent.yaml" + config_file.write_text(yaml.dump(config_data)) + + with patch("deep_agent.headless._CONFIG_PATH", config_file): + from deep_agent.headless import _load_headless_config + + config = _load_headless_config() + + assert config.mode == AgentMode.HEADLESS + assert config.triggers.webhook.enabled is True + assert config.triggers.webhook.port == 9999 + + def test_mode_is_always_headless_regardless_of_yaml(self, tmp_path: Path): + """_load_headless_config always sets mode=HEADLESS, ignoring the YAML value.""" + import yaml + + config_data = {"mode": "server", "triggers": {}, "output_sinks": []} + config_file = tmp_path / "agent.yaml" + config_file.write_text(yaml.dump(config_data)) + + with patch("deep_agent.headless._CONFIG_PATH", config_file): + from deep_agent.headless import _load_headless_config + + config = _load_headless_config() + + assert config.mode == AgentMode.HEADLESS + + def test_missing_config_file_causes_exit(self): + """When the config file does not exist, sys.exit(1) is called.""" + bogus = Path("/nonexistent/agent.yaml") + + with ( + patch("deep_agent.headless._CONFIG_PATH", bogus), + pytest.raises(SystemExit) as exc_info, + ): + from deep_agent.headless import _load_headless_config + + _load_headless_config() + + assert exc_info.value.code == 1 + + +# ------------------------------------------------------------------ +# main() lifecycle tests +# ------------------------------------------------------------------ + + +class TestHeadlessMainLifecycle: + """Test main() wires up prerequisites, graph, and middleware correctly. + + Since ``main()`` uses late/local imports from ``deep_agent.aegra`` + and ``deep_agent.src``, patches target the *source* modules so the + imports inside ``main()`` pick up the mocks. + """ + + async def test_main_starts_middleware_and_waits_for_signal(self): + """main() starts EventTriggerMiddleware and blocks on stop_event.""" + headless_config = HeadlessConfig( + mode=AgentMode.HEADLESS, + triggers=TriggerConfig( + webhook=WebhookTriggerConfig(enabled=True, port=0), + ), + output_sinks=[OutputSinkConfig(type="stdout")], + drain_timeout=2.0, + health_check=HealthCheckConfig(enabled=False), + ) + + mock_graph = MagicMock() + mock_mw_instance = AsyncMock() + + with ( + patch( + "deep_agent.headless._load_headless_config", + return_value=headless_config, + ), + patch( + "deep_agent.aegra.startup.run_startup", + new_callable=AsyncMock, + return_value={"status": "ok"}, + ), + patch( + "deep_agent.aegra.graph.agent", + new_callable=AsyncMock, + return_value=mock_graph, + ), + patch( + "deep_agent.src.triggers.middleware.EventTriggerMiddleware", + ) as MockMW, + patch( + "deep_agent.src.settings.settings", + MagicMock(REDIS_URL="redis://localhost:6379/0"), + ), + ): + MockMW.return_value = mock_mw_instance + + from deep_agent.headless import main + + task = asyncio.create_task(main()) + + # Give main() time to reach the stop_event.wait(). + await asyncio.sleep(0.3) + + # Verify middleware.start() was called. + mock_mw_instance.start.assert_awaited_once() + + # Simulate SIGTERM by cancelling the task (the signal handler + # would set stop_event, but cancellation is simpler in tests). + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # EventTriggerMiddleware was constructed with the right config. + MockMW.assert_called_once_with( + config=headless_config, + graph=mock_graph, + redis_url="redis://localhost:6379/0", + ) + + async def test_main_calls_run_startup(self): + """main() calls run_startup before building the graph.""" + headless_config = HeadlessConfig( + mode=AgentMode.HEADLESS, + triggers=TriggerConfig(), + output_sinks=[], + drain_timeout=1.0, + health_check=HealthCheckConfig(enabled=False), + ) + + mock_startup = AsyncMock(return_value={"status": "ok"}) + mock_mw = AsyncMock() + + with ( + patch( + "deep_agent.headless._load_headless_config", + return_value=headless_config, + ), + patch( + "deep_agent.aegra.startup.run_startup", + mock_startup, + ), + patch( + "deep_agent.aegra.graph.agent", + new_callable=AsyncMock, + return_value=MagicMock(), + ), + patch( + "deep_agent.src.triggers.middleware.EventTriggerMiddleware", + return_value=mock_mw, + ), + patch( + "deep_agent.src.settings.settings", + MagicMock(REDIS_URL="redis://localhost:6379/0"), + ), + ): + from deep_agent.headless import main + + task = asyncio.create_task(main()) + await asyncio.sleep(0.3) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + mock_startup.assert_awaited_once() + + async def test_main_invokes_graph_factory_with_runtime(self): + """main() calls ``agent(runtime)`` to build the compiled graph.""" + headless_config = HeadlessConfig( + mode=AgentMode.HEADLESS, + triggers=TriggerConfig(), + output_sinks=[], + drain_timeout=1.0, + health_check=HealthCheckConfig(enabled=False), + ) + + mock_agent_factory = AsyncMock(return_value=MagicMock()) + mock_mw = AsyncMock() + + with ( + patch( + "deep_agent.headless._load_headless_config", + return_value=headless_config, + ), + patch( + "deep_agent.aegra.startup.run_startup", + new_callable=AsyncMock, + return_value={}, + ), + patch( + "deep_agent.aegra.graph.agent", + mock_agent_factory, + ), + patch( + "deep_agent.src.triggers.middleware.EventTriggerMiddleware", + return_value=mock_mw, + ), + patch( + "deep_agent.src.settings.settings", + MagicMock(REDIS_URL="redis://localhost:6379/0"), + ), + ): + from deep_agent.headless import main + + task = asyncio.create_task(main()) + await asyncio.sleep(0.3) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + mock_agent_factory.assert_awaited_once() + # The argument should be a HeadlessRuntime instance. + runtime_arg = mock_agent_factory.call_args[0][0] + assert hasattr(runtime_arg, "user") + assert runtime_arg.user.identity == "headless-worker" diff --git a/tests/integration/triggers/test_redis_sink_integration.py b/tests/integration/triggers/test_redis_sink_integration.py new file mode 100644 index 00000000..aa6592cb --- /dev/null +++ b/tests/integration/triggers/test_redis_sink_integration.py @@ -0,0 +1,132 @@ +"""Integration tests for RedisSink with a real Redis server. + +Requires Redis running at ``redis://localhost:6379/0``. Each test uses +a unique stream name to avoid interference between tests. +""" + +from __future__ import annotations + +import json +from uuid import uuid4 + +import pytest + +from deep_agent.src.triggers.sinks.protocol import TriggerResult +from deep_agent.src.triggers.sinks.redis import RedisSink +from deep_agent.src.triggers.sources.protocol import TriggerEvent + +pytestmark = pytest.mark.integration + +_REDIS_URL = "redis://localhost:6379/0" + + +# ------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------ + + +def _unique_stream() -> str: + return f"test-results-{uuid4().hex[:8]}" + + +def _make_event(**overrides) -> TriggerEvent: + defaults = { + "name": "test-event", + "payload": {"key": "value"}, + "source": "integration-test", + } + defaults.update(overrides) + return TriggerEvent(**defaults) + + +def _make_result(**overrides) -> TriggerResult: + defaults = { + "event": _make_event(), + "output": {"answer": 42}, + "duration_ms": 50.0, + "success": True, + } + defaults.update(overrides) + return TriggerResult(**defaults) + + +# ------------------------------------------------------------------ +# Fixtures +# ------------------------------------------------------------------ + + +@pytest.fixture() +async def redis_client(): + """Yield a real ``redis.asyncio`` client, skip if Redis is unavailable.""" + aioredis = pytest.importorskip("redis.asyncio") + client = aioredis.from_url(_REDIS_URL, decode_responses=True) + try: + await client.ping() + except Exception: + pytest.skip("Redis not available at localhost:6379") + + yield client + await client.aclose() + + +# ------------------------------------------------------------------ +# Tests +# ------------------------------------------------------------------ + + +class TestRedisSinkIntegration: + """Integration tests verifying RedisSink writes to a real Redis stream.""" + + async def test_emit_writes_to_stream(self, redis_client): + """A single emit() writes a JSONL entry that XRANGE can read back.""" + stream = _unique_stream() + sink = RedisSink(stream=stream, redis_url=_REDIS_URL) + + try: + result = _make_result() + await sink.emit(result) + + entries = await redis_client.xrange(stream) + assert len(entries) == 1 + + _msg_id, fields = entries[0] + parsed = json.loads(fields["result"]) + assert parsed["success"] is True + assert parsed["event"]["name"] == "test-event" + assert parsed["output"] == {"answer": 42} + finally: + await sink.close() + await redis_client.delete(stream) + + async def test_multiple_emits_create_multiple_entries(self, redis_client): + """Multiple emit() calls create corresponding stream entries in order.""" + stream = _unique_stream() + sink = RedisSink(stream=stream, redis_url=_REDIS_URL) + + try: + for i in range(3): + await sink.emit(_make_result(output=f"result-{i}")) + + entries = await redis_client.xrange(stream) + assert len(entries) == 3 + + outputs = [json.loads(e[1]["result"])["output"] for e in entries] + assert outputs == ["result-0", "result-1", "result-2"] + finally: + await sink.close() + await redis_client.delete(stream) + + async def test_close_cleans_up(self, redis_client): + """close() tears down the internal Redis client.""" + stream = _unique_stream() + sink = RedisSink(stream=stream, redis_url=_REDIS_URL) + + try: + # Force client creation by emitting. + await sink.emit(_make_result()) + assert sink._client is not None + + await sink.close() + assert sink._client is None + finally: + await redis_client.delete(stream) diff --git a/tests/integration/triggers/test_redis_streams.py b/tests/integration/triggers/test_redis_streams.py new file mode 100644 index 00000000..f17e9957 --- /dev/null +++ b/tests/integration/triggers/test_redis_streams.py @@ -0,0 +1,208 @@ +"""Integration tests for RedisStreamsConsumer with a real Redis server. + +Requires Redis running at ``redis://localhost:6379/0``. Each test uses +a unique stream/group name to avoid interference between tests. +""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest + +pytestmark = pytest.mark.integration + +_REDIS_URL = "redis://localhost:6379/0" + + +def _unique_name(prefix: str = "test") -> str: + """Return a collision-free stream/group name.""" + return f"{prefix}-{uuid4().hex[:8]}" + + +# ------------------------------------------------------------------ +# Fixtures +# ------------------------------------------------------------------ + + +@pytest.fixture() +async def redis_client(): + """Yield a real ``redis.asyncio`` client, skip if Redis is unavailable.""" + aioredis = pytest.importorskip("redis.asyncio") + client = aioredis.from_url(_REDIS_URL, decode_responses=True) + try: + await client.ping() + except Exception: + pytest.skip("Redis not available at localhost:6379") + + yield client + await client.aclose() + + +@pytest.fixture() +def stream_name() -> str: + return _unique_name("stream") + + +@pytest.fixture() +def group_name() -> str: + return _unique_name("group") + + +# ------------------------------------------------------------------ +# Tests +# ------------------------------------------------------------------ + + +class TestRedisStreamsConsumerIntegration: + """Integration tests exercising RedisStreamsConsumer against real Redis.""" + + async def test_creates_group_and_stream_on_first_consume( + self, redis_client, stream_name, group_name + ): + """Consumer creates the stream and consumer group automatically.""" + from deep_agent.src.triggers.sources.queue import RedisStreamsConsumer + + consumer = RedisStreamsConsumer( + stream=stream_name, + consumer_group=group_name, + consumer_name="worker-1", + redis_url=_REDIS_URL, + block_ms=100, + ) + + # _ensure_client() creates the group and stream via XGROUP CREATE MKSTREAM. + await consumer._ensure_client() + + # Stream and group should now exist in Redis. + groups = await redis_client.xinfo_groups(stream_name) + group_names = [g["name"] for g in groups] + assert group_name in group_names + + await consumer.close() + await redis_client.delete(stream_name) + + async def test_produce_consume_ack_cycle( + self, redis_client, stream_name, group_name + ): + """XADD -> consume -> ack round-trip works end-to-end.""" + from deep_agent.src.triggers.sources.queue import RedisStreamsConsumer + + consumer = RedisStreamsConsumer( + stream=stream_name, + consumer_group=group_name, + consumer_name="worker-1", + redis_url=_REDIS_URL, + block_ms=100, + ) + + # Produce a message before consuming. + await redis_client.xadd( + stream_name, {"name": "integration-event", "data": "hello"} + ) + + received = [] + async for msg in consumer.consume(): + received.append(msg) + await consumer.ack(msg) + # After receiving one message, stop. + consumer._running = False + + assert len(received) == 1 + assert received[0].data["name"] == "integration-event" + assert received[0].data["data"] == "hello" + assert received[0].id # should be a valid Redis stream ID + + # Verify the message was acknowledged (pending count should be 0). + groups = await redis_client.xinfo_groups(stream_name) + target_group = [g for g in groups if g["name"] == group_name][0] + assert target_group["pending"] == 0 + + await consumer.close() + await redis_client.delete(stream_name) + + async def test_handles_existing_group_busygroup( + self, redis_client, stream_name, group_name + ): + """Creating a consumer when the group already exists does not raise.""" + from deep_agent.src.triggers.sources.queue import RedisStreamsConsumer + + # Pre-create the group to trigger BUSYGROUP. + await redis_client.xgroup_create(stream_name, group_name, id="0", mkstream=True) + + consumer = RedisStreamsConsumer( + stream=stream_name, + consumer_group=group_name, + consumer_name="worker-1", + redis_url=_REDIS_URL, + block_ms=100, + ) + + # Should not raise. + consumer._running = False + async for _ in consumer.consume(): + pass # pragma: no cover + + # Group still exists and is intact. + groups = await redis_client.xinfo_groups(stream_name) + assert any(g["name"] == group_name for g in groups) + + await consumer.close() + await redis_client.delete(stream_name) + + async def test_multiple_messages_consumed_in_order( + self, redis_client, stream_name, group_name + ): + """Multiple messages are consumed in the order they were added.""" + from deep_agent.src.triggers.sources.queue import RedisStreamsConsumer + + consumer = RedisStreamsConsumer( + stream=stream_name, + consumer_group=group_name, + consumer_name="worker-1", + redis_url=_REDIS_URL, + block_ms=100, + ) + + # Add several messages. + for i in range(5): + await redis_client.xadd(stream_name, {"name": f"event-{i}", "seq": str(i)}) + + received = [] + async for msg in consumer.consume(): + received.append(msg) + await consumer.ack(msg) + if len(received) >= 5: + consumer._running = False + + assert len(received) == 5 + names = [m.data["name"] for m in received] + assert names == [f"event-{i}" for i in range(5)] + + await consumer.close() + await redis_client.delete(stream_name) + + async def test_close_cleans_up_connection( + self, redis_client, stream_name, group_name + ): + """close() sets running=False and tears down the client.""" + from deep_agent.src.triggers.sources.queue import RedisStreamsConsumer + + consumer = RedisStreamsConsumer( + stream=stream_name, + consumer_group=group_name, + consumer_name="worker-1", + redis_url=_REDIS_URL, + block_ms=100, + ) + + # Force client creation. + await consumer._ensure_client() + assert consumer._client is not None + + await consumer.close() + + assert consumer._running is False + assert consumer._client is None + + await redis_client.delete(stream_name) diff --git a/tests/integration/triggers/test_webhook_listener.py b/tests/integration/triggers/test_webhook_listener.py new file mode 100644 index 00000000..7eb731c5 --- /dev/null +++ b/tests/integration/triggers/test_webhook_listener.py @@ -0,0 +1,121 @@ +"""Integration tests for WebhookTriggerSource with real HTTP connections. + +No external services required — the webhook listener binds to a random +port on localhost and tests communicate via ``httpx.AsyncClient``. +""" + +from __future__ import annotations + +import httpx +import pytest + +from deep_agent.src.triggers.config import WebhookTriggerConfig +from deep_agent.src.triggers.sources.webhook import WebhookTriggerSource + +pytestmark = pytest.mark.integration + + +# ------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------ + + +async def _start_source( + path: str = "/trigger", +) -> tuple[WebhookTriggerSource, int]: + """Create and start a webhook source on a random OS-assigned port. + + Returns ``(source, actual_port)``. + """ + config = WebhookTriggerConfig( + enabled=True, + host="127.0.0.1", + port=0, + path=path, + ) + source = WebhookTriggerSource(config) + await source.start() + assert source._server is not None + port = source._server.sockets[0].getsockname()[1] + return source, port + + +# ------------------------------------------------------------------ +# Tests +# ------------------------------------------------------------------ + + +class TestWebhookListenerIntegration: + """Integration tests sending real HTTP requests to the webhook listener.""" + + async def test_valid_post_returns_200_and_produces_event(self): + """POST valid JSON to the trigger path produces a TriggerEvent.""" + source, port = await _start_source() + try: + async with httpx.AsyncClient() as client: + response = await client.post( + f"http://127.0.0.1:{port}/trigger", + json={"event": "integration-test", "key": "value"}, + ) + + assert response.status_code == 200 + body = response.json() + assert body["status"] == "accepted" + + event = source._queue.get_nowait() + assert event.name == "integration-test" + assert event.payload == {"key": "value"} + assert event.source == "webhook" + finally: + await source.stop() + + async def test_invalid_json_returns_400(self): + """POST with invalid JSON body returns 400 and enqueues nothing.""" + source, port = await _start_source() + try: + async with httpx.AsyncClient() as client: + response = await client.post( + f"http://127.0.0.1:{port}/trigger", + content=b"<<>>", + headers={"Content-Type": "application/json"}, + ) + + assert response.status_code == 400 + body = response.json() + assert "invalid JSON" in body["error"] + assert source._queue.empty() + finally: + await source.stop() + + async def test_wrong_path_returns_404(self): + """POST to a non-configured path returns 404.""" + source, port = await _start_source() + try: + async with httpx.AsyncClient() as client: + response = await client.post( + f"http://127.0.0.1:{port}/not-the-trigger", + json={"event": "lost"}, + ) + + assert response.status_code == 404 + assert source._queue.empty() + finally: + await source.stop() + + async def test_multiple_posts_produce_events_in_order(self): + """Sequential POST requests produce events in FIFO order.""" + source, port = await _start_source() + try: + async with httpx.AsyncClient() as client: + for i in range(4): + resp = await client.post( + f"http://127.0.0.1:{port}/trigger", + json={"event": f"evt-{i}", "seq": i}, + ) + assert resp.status_code == 200 + + assert source._queue.qsize() == 4 + names = [source._queue.get_nowait().name for _ in range(4)] + assert names == [f"evt-{i}" for i in range(4)] + finally: + await source.stop() diff --git a/tests/mocks/__init__.py b/tests/mocks/__init__.py new file mode 100644 index 00000000..472a0895 --- /dev/null +++ b/tests/mocks/__init__.py @@ -0,0 +1 @@ +"""Mock implementations for testing.""" diff --git a/tests/mocks/mock_mcp_server.py b/tests/mocks/mock_mcp_server.py new file mode 100644 index 00000000..f3a82c1c --- /dev/null +++ b/tests/mocks/mock_mcp_server.py @@ -0,0 +1,292 @@ +#!/usr/bin/env python3 +"""Mock MCP server for testing. + +Provides stub implementations of the tools required by the agent: +- calculate_bmi: Returns mock BMI calculation +- validate_email: Basic email format validation +- send_email: Simulates email sending (always succeeds) +- search_web: Returns mock health tips + +This allows agent evals to run without requiring the full template-mcp-server. +""" + +import json +import re +from typing import Any, Dict + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +app = FastAPI(title="Mock MCP Server") + +# Mock health tips by BMI category +HEALTH_TIPS = { + "Underweight": [ + "Focus on nutrient-dense foods with healthy fats and proteins", + "Consider increasing meal frequency with healthy snacks", + "Consult with a healthcare provider for personalized guidance", + ], + "Normal": [ + "Maintain a balanced diet with whole grains, lean proteins, and vegetables", + "Aim for 150 minutes of moderate aerobic activity per week", + "Stay hydrated and get adequate sleep for optimal health", + ], + "Overweight": [ + "Focus on portion control and mindful eating habits", + "Incorporate regular physical activity into your daily routine", + "Consider working with a registered dietitian for personalized nutrition advice", + ], + "Obese": [ + "Consult with a healthcare provider for a comprehensive health assessment", + "Set realistic, sustainable goals for gradual weight management", + "Focus on building healthy habits rather than quick fixes", + ], +} + + +def calculate_bmi_value(height_cm: float, weight_kg: float) -> Dict[str, Any]: + """Calculate BMI and determine category. + + Args: + height_cm: Height in centimeters + weight_kg: Weight in kilograms + + Returns: + Dict with bmi, category, and message + """ + if height_cm <= 0 or weight_kg <= 0: + return { + "success": False, + "error": "Height and weight must be positive values", + } + + height_m = height_cm / 100 + bmi = weight_kg / (height_m**2) + + # Determine category + if bmi < 18.5: + category = "Underweight" + elif bmi < 25: + category = "Normal" + elif bmi < 30: + category = "Overweight" + else: + category = "Obese" + + return { + "success": True, + "bmi": round(bmi, 1), + "category": category, + "height_cm": height_cm, + "weight_kg": weight_kg, + } + + +def validate_email_address(email: str) -> Dict[str, Any]: + """Validate email address format. + + Args: + email: Email address to validate + + Returns: + Dict with valid flag and message + """ + # Basic email regex + pattern = r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$" + is_valid = bool(re.match(pattern, email)) + + return { + "success": True, + "valid": is_valid, + "email": email, + "message": "Valid email format" if is_valid else "Invalid email format", + } + + +def send_email_mock(recipient: str, subject: str, body: str) -> Dict[str, Any]: + """Mock email sending (always succeeds). + + Args: + recipient: Email recipient + subject: Email subject + body: Email body (HTML or plain text) + + Returns: + Dict with success flag and message + """ + # Validate recipient email + validation = validate_email_address(recipient) + if not validation["valid"]: + return { + "success": False, + "error": f"Invalid recipient email: {recipient}", + } + + return { + "success": True, + "recipient": recipient, + "subject": subject, + "message": f"Email sent successfully to {recipient}", + "message_id": f"mock-{hash(recipient + subject)}", + } + + +def search_web_mock(query: str) -> Dict[str, Any]: + """Mock web search for health tips. + + Args: + query: Search query (should contain BMI category) + + Returns: + Dict with search results (health tips) + """ + # Extract category from query + query_lower = query.lower() + category = None + + if "underweight" in query_lower: + category = "Underweight" + elif "overweight" in query_lower: + category = "Overweight" + elif "obese" in query_lower or "obesity" in query_lower: + category = "Obese" + elif "normal" in query_lower: + category = "Normal" + + # Get tips for category + tips = HEALTH_TIPS.get(category, HEALTH_TIPS["Normal"]) + + return { + "success": True, + "query": query, + "category": category, + "results": [ + {"title": f"Health Tip {i + 1}", "snippet": tip} + for i, tip in enumerate(tips) + ], + } + + +# MCP Tool definitions +TOOLS = [ + { + "name": "calculate_bmi", + "description": "Calculate BMI (Body Mass Index) from height and weight", + "inputSchema": { + "type": "object", + "properties": { + "height_cm": { + "type": "number", + "description": "Height in centimeters", + }, + "weight_kg": { + "type": "number", + "description": "Weight in kilograms", + }, + }, + "required": ["height_cm", "weight_kg"], + }, + }, + { + "name": "validate_email", + "description": "Validate email address format", + "inputSchema": { + "type": "object", + "properties": { + "email": { + "type": "string", + "description": "Email address to validate", + }, + }, + "required": ["email"], + }, + }, + { + "name": "send_email", + "description": "Send an email (mock - always succeeds)", + "inputSchema": { + "type": "object", + "properties": { + "recipient": { + "type": "string", + "description": "Email recipient", + }, + "subject": { + "type": "string", + "description": "Email subject", + }, + "body": { + "type": "string", + "description": "Email body (HTML or plain text)", + }, + }, + "required": ["recipient", "subject", "body"], + }, + }, + { + "name": "search_web", + "description": "Search the web for health tips (mock - returns predefined tips)", + "inputSchema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query", + }, + }, + "required": ["query"], + }, + }, +] + + +@app.get("/health") +async def health(): + """Health check endpoint.""" + return {"status": "healthy", "service": "Mock MCP Server"} + + +@app.get("/mcp/tools") +async def list_tools(): + """List available MCP tools.""" + return {"tools": TOOLS} + + +@app.post("/mcp/tools/{tool_name}") +async def call_tool(tool_name: str, request: Request): + """Execute an MCP tool.""" + body = await request.json() + arguments = body.get("arguments", {}) + + # Route to appropriate tool implementation + if tool_name == "calculate_bmi": + result = calculate_bmi_value( + arguments.get("height_cm"), + arguments.get("weight_kg"), + ) + elif tool_name == "validate_email": + result = validate_email_address(arguments.get("email")) + elif tool_name == "send_email": + result = send_email_mock( + arguments.get("recipient"), + arguments.get("subject"), + arguments.get("body"), + ) + elif tool_name == "search_web": + result = search_web_mock(arguments.get("query")) + else: + return JSONResponse( + status_code=404, + content={"error": f"Tool not found: {tool_name}"}, + ) + + return {"result": result} + + +if __name__ == "__main__": + import uvicorn + + print("Starting Mock MCP Server on http://localhost:5001") + print("Available tools: calculate_bmi, validate_email, send_email, search_web") + + uvicorn.run(app, host="0.0.0.0", port=5001, log_level="info") diff --git a/tests/skills/conftest.py b/tests/skills/conftest.py new file mode 100644 index 00000000..92296d40 --- /dev/null +++ b/tests/skills/conftest.py @@ -0,0 +1,242 @@ +"""Pytest configuration and fixtures for skills tests with auto-discovery.""" + +import json +import os +import time +from pathlib import Path +from typing import Dict, Optional + +import pytest +from langchain_google_genai import ChatGoogleGenerativeAI +from langfuse import Langfuse + +from llm_judge import LLMJudge + +PROJECT_ROOT = Path(__file__).parent.parent.parent +SKILLS_DIR = PROJECT_ROOT / "config" / "agent" / "skills" + +MODEL_NAME = "gemini-3.1-pro-preview" +MODEL_TEMPERATURE = 0 + + +# ============================================================================ +# Auto-Discovery +# ============================================================================ + + +def pytest_generate_tests(metafunc): + """Auto-discover all skills and their evals.""" + if "skill_eval" not in metafunc.fixturenames: + return + + if not SKILLS_DIR.exists(): + pytest.skip(f"Skills directory not found: {SKILLS_DIR}") + return + + test_cases = [] + ids = [] + + # Discover all skills + for skill_dir in sorted(SKILLS_DIR.iterdir()): + if not skill_dir.is_dir(): + continue + + skill_name = skill_dir.name + evals_file = skill_dir / "evals" / "evals.json" + + if not evals_file.exists(): + continue + + # Load evals for this skill + with open(evals_file) as f: + evals_data = json.load(f) + + # Create test case for each eval + for eval_case in evals_data.get("evals", []): + eval_id = eval_case["id"] + test_cases.append( + { + "skill_name": skill_name, + "skill_dir": str(skill_dir.resolve()), + "eval_id": eval_id, + "eval_case": eval_case, + } + ) + ids.append(f"{skill_name}-eval-{eval_id}") + + if not test_cases: + pytest.skip("No skill evals found") + return + + metafunc.parametrize( + "skill_eval", + test_cases, + ids=ids, + ) + + +# ============================================================================ +# Session Fixtures +# ============================================================================ + + +@pytest.fixture(scope="session") +def workspace_dir(): + """Workspace directory for test outputs.""" + workspace = PROJECT_ROOT / "tests" / "workspaces" / "skills" + workspace.mkdir(parents=True, exist_ok=True) + return workspace + + +@pytest.fixture +def model(): + """Create Gemini model with credentials. + + Function-scoped to ensure each test gets a fresh model instance + bound to the correct event loop. + """ + from deep_agent.utils.google_creds import get_service_account_credentials + + # Check if credentials are available + if not os.getenv("GOOGLE_APPLICATION_CREDENTIALS_CONTENT"): + pytest.skip("Google Cloud credentials not available - skipping skill tests") + + try: + credentials, project = get_service_account_credentials() + return ChatGoogleGenerativeAI( + model=MODEL_NAME, + temperature=MODEL_TEMPERATURE, + credentials=credentials, + project=project, + ) + except RuntimeError as e: + pytest.skip(f"Google Cloud credentials error: {e}") + + +# ============================================================================ +# Function Fixtures +# ============================================================================ + + +@pytest.fixture +def tracer(): + """Execution tracer for timing and token tracking.""" + return ExecutionTracer() + + +@pytest.fixture +def langfuse_client(): + """Langfuse client (optional, requires env vars). + + Ensures traces are flushed before test teardown. + """ + if all( + [ + os.getenv("LANGFUSE_PUBLIC_KEY"), + os.getenv("LANGFUSE_SECRET_KEY"), + os.getenv("LANGFUSE_BASE_URL"), + ] + ): + client = Langfuse() + yield client + # Flush pending traces before test cleanup + client.flush() + else: + yield None + + +@pytest.fixture +def evaluator(langfuse_client): + """LLM judge evaluator.""" + judge = LLMJudge(langfuse_client=langfuse_client) + return AssertionEvaluator(judge) + + +# ============================================================================ +# Helper Functions +# ============================================================================ + + +def extract_output(result: dict) -> str: + """Extract text from agent result messages.""" + messages = result.get("messages", []) + + for msg in reversed(messages): + if not (hasattr(msg, "content") and msg.content): + continue + if hasattr(msg, "type") and msg.type == "human": + continue + + content = msg.content + if isinstance(content, list): + return "\n".join( + block.get("text", "") if isinstance(block, dict) else str(block) + for block in content + ) + return str(content) + + return "" + + +def extract_tokens(result: dict) -> int: + """Extract total token count from messages.""" + total = 0 + for msg in result.get("messages", []): + if hasattr(msg, "usage_metadata") and msg.usage_metadata: + total += msg.usage_metadata.get("total_tokens", 0) + return total + + +# ============================================================================ +# Classes +# ============================================================================ + + +class ExecutionTracer: + """Tracks execution time and token usage.""" + + def __init__(self): + self.start_time = None + self.end_time = None + self.total_tokens = 0 + + def start(self): + self.start_time = time.time() + + def end(self, total_tokens: int = 0): + self.end_time = time.time() + self.total_tokens = total_tokens + + def duration_ms(self) -> int: + if self.start_time and self.end_time: + return int((self.end_time - self.start_time) * 1000) + return 0 + + +class AssertionEvaluator: + """Evaluates assertions using LLM judge.""" + + def __init__(self, llm_judge: LLMJudge): + self.llm_judge = llm_judge + + def evaluate( + self, + assertion: str, + output: str, + context: Optional[Dict] = None, + trace_id: Optional[str] = None, + ) -> Dict: + """Evaluate assertion against output.""" + result = self.llm_judge.evaluate(assertion, output, context, trace_id) + result["method"] = "llm_judge" + return result + + +# ============================================================================ +# Pytest Hooks +# ============================================================================ + + +def pytest_configure(config): + """Configure pytest with custom markers.""" + config.addinivalue_line("markers", "skills: skills evaluation tests") diff --git a/tests/skills/llm_judge.py b/tests/skills/llm_judge.py new file mode 100644 index 00000000..2e0c51a6 --- /dev/null +++ b/tests/skills/llm_judge.py @@ -0,0 +1,187 @@ +"""LLM-as-Judge evaluator using Gemini.""" + +from typing import Dict, Optional + +from langchain_google_genai import ChatGoogleGenerativeAI +from langfuse import Langfuse + +from deep_agent.utils.google_creds import get_service_account_credentials + +# Model configuration +MODEL_NAME = "gemini-3.1-pro-preview" +MODEL_TEMPERATURE = 0 +OUTPUT_TRUNCATE_LENGTH = 500 + +# Response field markers +VERDICT_MARKER = "VERDICT:" +EVIDENCE_MARKER = "EVIDENCE:" +CONFIDENCE_MARKER = "CONFIDENCE:" +REASONING_MARKER = "REASONING:" + + +def create_judge_prompt(assertion: str, output: str, context: Optional[Dict]) -> str: + """Build evaluation prompt for LLM judge.""" + sections = [ + "You are an expert evaluator. Assess whether the agent's output satisfies the assertion.", + "", + f"ASSERTION: {assertion}", + "", + f"AGENT OUTPUT:\n{output}", + ] + + if context: + sections.extend(["", "CONTEXT:"]) + if context.get("expected_output"): + sections.append(f"Expected: {context['expected_output']}") + if context.get("prompt"): + sections.append(f"User Prompt: {context['prompt']}") + if context.get("skill_name"): + sections.append(f"Skill: {context['skill_name']}") + + sections.extend( + [ + "", + "Evaluate strictly but fairly. Provide:", + "VERDICT: YES or NO", + "EVIDENCE: Quote or describe specific evidence", + "CONFIDENCE: 0.0 to 1.0", + "REASONING: Brief explanation", + ] + ) + + return "\n".join(sections) + + +def parse_judge_response(response: str) -> Dict: + """Parse structured LLM judge response.""" + result = { + "passed": None, + "evidence": "", + "confidence": 0.5, + "reasoning": "", + } + + for line in response.strip().split("\n"): + line = line.strip() + + if line.startswith(VERDICT_MARKER): + verdict = line.split(":", 1)[1].strip().upper() + result["passed"] = verdict == "YES" + elif line.startswith(EVIDENCE_MARKER): + result["evidence"] = line.split(":", 1)[1].strip() + elif line.startswith(CONFIDENCE_MARKER): + try: + conf = float(line.split(":", 1)[1].strip()) + result["confidence"] = max(0.0, min(1.0, conf)) + except ValueError: + pass + elif line.startswith(REASONING_MARKER): + result["reasoning"] = line.split(":", 1)[1].strip() + + return result + + +def extract_text_content(content) -> str: + """Extract text from Gemini response content (handles str or list).""" + if isinstance(content, str): + return content + + if isinstance(content, list): + return "\n".join( + block.get("text", "") if isinstance(block, dict) else str(block) + for block in content + ) + + return str(content) + + +class LLMJudge: + """LLM-as-judge evaluator with Langfuse tracing.""" + + def __init__(self, langfuse_client: Optional[Langfuse] = None): + credentials, project = get_service_account_credentials() + self.model = ChatGoogleGenerativeAI( + model=MODEL_NAME, + temperature=MODEL_TEMPERATURE, + credentials=credentials, + project=project, + ) + self.langfuse = langfuse_client + + def evaluate( + self, + assertion: str, + output: str, + context: Optional[Dict] = None, + trace_id: Optional[str] = None, + ) -> Dict: + """Evaluate assertion using LLM judge.""" + prompt = create_judge_prompt(assertion, output, context) + generation = self._create_generation(assertion, output, context, trace_id) + + try: + response = self.model.invoke(prompt) + content = extract_text_content(response.content) + result = parse_judge_response(content) + self._finalize_generation(generation, result, assertion) + return result + + except Exception as e: + self._handle_error(generation, e) + return { + "passed": None, + "evidence": f"LLM judge error: {str(e)}", + "confidence": 0.0, + "reasoning": "", + } + + def _create_generation( + self, + assertion: str, + output: str, + context: Optional[Dict], + trace_id: Optional[str], + ): + """Create Langfuse generation span.""" + if not (self.langfuse and trace_id): + return None + + generation_input = { + "assertion": assertion, + "output": output[:OUTPUT_TRUNCATE_LENGTH], + } + if context: + generation_input["context"] = context + + return self.langfuse.generation( + trace_id=trace_id, + name="llm_judge_evaluation", + model=MODEL_NAME, + input=generation_input, + ) + + def _finalize_generation(self, generation, result: Dict, assertion: str): + """Update and close Langfuse generation.""" + if not generation: + return + + generation.update( + output=result, + metadata={ + "assertion": assertion, + "passed": result["passed"], + "confidence": result.get("confidence", 0.0), + }, + ) + generation.end() + + def _handle_error(self, generation, error: Exception): + """Handle and log error to Langfuse.""" + if not generation: + return + + generation.update( + level="ERROR", + status_message=str(error), + ) + generation.end() diff --git a/tests/skills/test_skills.py b/tests/skills/test_skills.py new file mode 100644 index 00000000..edc8ef8f --- /dev/null +++ b/tests/skills/test_skills.py @@ -0,0 +1,214 @@ +"""Generic skill tests with auto-discovery. + +This single test file automatically discovers and tests all skills in +config/agent/skills/ by loading their evals.json files. +""" + +import asyncio +import json +from pathlib import Path + +import pytest +from deepagents import create_deep_agent +from langgraph.checkpoint.memory import MemorySaver + +from deep_agent.src.infrastructure.backend import get_backend + + +# ============================================================================ +# Skills are self-contained - no external tools needed +# ============================================================================ +# +# All skills use only local scripts and reference documents: +# - client-intake: uses scripts/convert_units.py and reference docs +# - bmi-report: uses reference docs (bmi_categories.md, health_tips, etc.) +# - email-formatter: uses reference docs (template.html, css rules, etc.) +# +# No mock tools required for skill testing! + + +# ============================================================================ +# Helpers +# ============================================================================ + + +def save_output(output_dir: Path, output: str): + """Save agent output to file.""" + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "response.md").write_text(output) + + +def save_grading(output_dir: Path, results: list, summary: dict): + """Save grading results to JSON.""" + output_dir.mkdir(parents=True, exist_ok=True) + grading = {"assertion_results": results, "summary": summary} + (output_dir / "grading.json").write_text(json.dumps(grading, indent=2)) + + +def calculate_summary(results: list) -> dict: + """Calculate pass/fail summary. + + Assertions with passed=null are counted as 'aborted' and excluded + from pass rate calculation (LLM judge couldn't determine verdict). + """ + passed = sum(1 for r in results if r["passed"] is True) + failed = sum(1 for r in results if r["passed"] is False) + aborted = sum(1 for r in results if r["passed"] is None) + total = len(results) + + # Pass rate excludes aborted tests + evaluated = passed + failed + pass_rate = passed / evaluated if evaluated > 0 else 0 + + return { + "passed": passed, + "failed": failed, + "aborted": aborted, + "total": total, + "pass_rate": pass_rate, + } + + +def build_context(skill_name: str, eval_id: int, eval_case: dict) -> dict: + """Build evaluation context.""" + return { + "skill_name": skill_name, + "eval_id": eval_id, + "prompt": eval_case["prompt"], + "expected_output": eval_case.get("expected_output"), + } + + +async def run_agent_async(agent, prompt: str, thread_id: str, tracer) -> str: + """Run agent asynchronously.""" + from conftest import extract_output, extract_tokens + + tracer.start() + + config = {"configurable": {"thread_id": thread_id}} + + try: + result = await agent.ainvoke( + {"messages": [{"role": "user", "content": prompt}]}, + config=config, + ) + + output = extract_output(result) + tokens = extract_tokens(result) + tracer.end(total_tokens=tokens) + return output + + except Exception: + tracer.end() + raise + + +def run_agent_sync(agent, prompt: str, thread_id: str, tracer) -> str: + """Synchronous wrapper for async agent execution.""" + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + return loop.run_until_complete(run_agent_async(agent, prompt, thread_id, tracer)) + + +def create_skill_agent(skill_dir: str, skill_name: str, model): + """Create agent for a specific skill. + + Skills are self-contained and use only local scripts/reference docs. + No external tools are needed. + """ + system_prompt = """\ + You are a helpful assistant with specialized skills. + + CRITICAL: You have been provided with skill instructions as part of SKILL.md + + You MUST strictly follow the skill instructions. + + When using a skill, follow its instructions EXACTLY as written. + """ + + agent = create_deep_agent( + model=model, + system_prompt=system_prompt, + skills=[skill_dir], + tools=[], # Skills don't need external tools + backend=get_backend(), + checkpointer=MemorySaver(), + ) + + return agent + + +# ============================================================================ +# Tests +# ============================================================================ + + +@pytest.mark.skills +def test_skill_evaluation( + skill_eval, + workspace_dir, + tracer, + evaluator, + model, +): + """Test skill with eval case using LLM judge. + + Auto-discovers all skills from config/agent/skills/ and runs their evals. + + Each eval must pass 70% of its assertions to be considered successful. + + IMPORTANT: These tests use real LLM calls and LLM-as-judge evaluation, + which means results can vary between runs. The system prompt helps guide + the model to follow skill instructions more strictly. + """ + skill_name = skill_eval["skill_name"] + skill_dir = skill_eval["skill_dir"] + eval_id = skill_eval["eval_id"] + eval_case = skill_eval["eval_case"] + + # Setup workspace + workspace = workspace_dir / skill_name / f"eval-{eval_id}" + output_dir = workspace / "outputs" + + # Create agent with skill + agent = create_skill_agent(skill_dir, skill_name, model) + + # Run agent + prompt = eval_case["prompt"] + thread_id = f"{skill_name}-test-{eval_id}" + output = run_agent_sync(agent, prompt, thread_id, tracer) + + # Save output + save_output(output_dir, output) + + # Grade assertions + context = build_context(skill_name, eval_id, eval_case) + results = [] + + for assertion in eval_case["assertions"]: + result = evaluator.evaluate( + assertion=assertion, + output=output, + context=context, + ) + results.append(result) + + # Calculate summary for this eval + summary = calculate_summary(results) + + # Save grading + save_grading(output_dir, results, summary) + + # Assert pass rate (70% threshold per eval) + pass_rate = summary["pass_rate"] + + # Build failure message + msg_parts = [ + f"{skill_name} failed eval-{eval_id}:", + f"{summary['passed']}/{summary['total']} assertions passed", + ] + if summary["aborted"] > 0: + msg_parts.append(f"({summary['aborted']} aborted, excluded from rate)") + msg_parts.append(f"pass_rate: {pass_rate:.1%}, threshold: 70%") + + assert pass_rate >= 0.7, " ".join(msg_parts) diff --git a/tests/test_agent_utils.py b/tests/test_agent_utils.py deleted file mode 100644 index 5626c21d..00000000 --- a/tests/test_agent_utils.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Tests for the agent_utils module.""" - -from unittest.mock import Mock - -import pytest - -from template_agent.src.core.agent_utils import ( - convert_message_content_to_string, - langchain_to_chat_message, - remove_tool_calls, -) -from template_agent.src.schema import ChatMessage - - -class TestAgentUtils: - """Test cases for agent utility functions.""" - - def test_convert_message_content_to_string_simple(self): - """Test converting simple string content.""" - content = "Hello world" - result = convert_message_content_to_string(content) - assert result == "Hello world" - - def test_convert_message_content_to_string_list(self): - """Test converting list content with text items.""" - content = ["Hello", " ", "world"] - result = convert_message_content_to_string(content) - assert result == "Hello world" - - def test_convert_message_content_to_string_mixed(self): - """Test converting mixed content with text and dict items.""" - content = ["Hello", {"type": "text", "text": " world"}] - result = convert_message_content_to_string(content) - assert result == "Hello world" - - def test_convert_message_content_to_string_ignores_non_text(self): - """Test that non-text dict items are ignored.""" - content = ["Hello", {"type": "image", "url": "test.jpg"}, " world"] - result = convert_message_content_to_string(content) - assert result == "Hello world" - - def test_remove_tool_calls_string(self): - """Test remove_tool_calls with string content.""" - content = "Hello world" - result = remove_tool_calls(content) - assert result == "Hello world" - - def test_remove_tool_calls_list_without_tools(self): - """Test remove_tool_calls with list content without tool calls.""" - content = ["Hello", " world"] - result = remove_tool_calls(content) - assert result == ["Hello", " world"] - - def test_remove_tool_calls_list_with_tools(self): - """Test remove_tool_calls with list content containing tool calls.""" - content = ["Hello", {"type": "tool_use", "tool_use": {}}, " world"] - result = remove_tool_calls(content) - assert result == ["Hello", " world"] - - def test_langchain_to_chat_message_human(self): - """Test converting HumanMessage to ChatMessage.""" - from langchain_core.messages import HumanMessage - - human_msg = HumanMessage(content="Hello") - result = langchain_to_chat_message(human_msg) - - assert isinstance(result, ChatMessage) - assert result.type == "human" - assert result.content == "Hello" - - def test_langchain_to_chat_message_ai(self): - """Test converting AIMessage to ChatMessage.""" - from langchain_core.messages import AIMessage - - ai_msg = AIMessage(content="Hello", tool_calls=[]) - result = langchain_to_chat_message(ai_msg) - - assert isinstance(result, ChatMessage) - assert result.type == "ai" - assert result.content == "Hello" - - def test_langchain_to_chat_message_tool(self): - """Test converting ToolMessage to ChatMessage.""" - from langchain_core.messages import ToolMessage - - tool_msg = ToolMessage(content="Tool result", tool_call_id="call_123") - result = langchain_to_chat_message(tool_msg) - - assert isinstance(result, ChatMessage) - assert result.type == "tool" - assert result.content == "Tool result" - assert result.tool_call_id == "call_123" - - def test_langchain_to_chat_message_unsupported(self): - """Test that unsupported message types raise ValueError.""" - mock_msg = Mock() - mock_msg.__class__.__name__ = "UnsupportedMessage" - - with pytest.raises(ValueError, match="Unsupported message type"): - langchain_to_chat_message(mock_msg) diff --git a/tests/test_database_init.py b/tests/test_database_init.py deleted file mode 100644 index c3e4a639..00000000 --- a/tests/test_database_init.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Tests for database initialization functionality. - -This module tests the database schema initialization to ensure the checkpoints -table is created properly on application startup when using PostgreSQL storage. -""" - -import pytest -from unittest.mock import AsyncMock, MagicMock, patch - -from template_agent.src.core.agent import initialize_database -from template_agent.src.core.exceptions.exceptions import AppException - - -class TestDatabaseInitialization: - """Test cases for database initialization.""" - - @pytest.mark.asyncio - async def test_initialize_database_skips_when_inmemory(self): - """Test that database initialization is skipped when using in-memory storage.""" - with patch("template_agent.src.core.agent.settings") as mock_settings: - mock_settings.USE_INMEMORY_SAVER = True - - # Should not raise any exceptions - await initialize_database() - - @pytest.mark.asyncio - async def test_initialize_database_calls_setup(self): - """Test that database initialization calls setup on the checkpoint.""" - with patch("template_agent.src.core.agent.settings") as mock_settings: - mock_settings.USE_INMEMORY_SAVER = False - mock_settings.database_uri = "postgresql://user:pass@localhost:5432/db" - - # Create mock checkpoint with setup method - mock_checkpoint = AsyncMock() - mock_checkpoint.setup = AsyncMock() - mock_checkpoint.__aenter__ = AsyncMock(return_value=mock_checkpoint) - mock_checkpoint.__aexit__ = AsyncMock(return_value=None) - - with patch( - "template_agent.src.core.agent.AsyncPostgresSaver.from_conn_string", - return_value=mock_checkpoint, - ): - await initialize_database() - - # Verify setup was called - mock_checkpoint.setup.assert_called_once() - - @pytest.mark.asyncio - async def test_initialize_database_handles_no_setup_method(self): - """Test that database initialization handles checkpoints without setup method.""" - with patch("template_agent.src.core.agent.settings") as mock_settings: - mock_settings.USE_INMEMORY_SAVER = False - mock_settings.database_uri = "postgresql://user:pass@localhost:5432/db" - - # Create mock checkpoint without setup method - mock_checkpoint = AsyncMock() - mock_checkpoint.__aenter__ = AsyncMock(return_value=mock_checkpoint) - mock_checkpoint.__aexit__ = AsyncMock(return_value=None) - - with patch( - "template_agent.src.core.agent.AsyncPostgresSaver.from_conn_string", - return_value=mock_checkpoint, - ): - # Should not raise exception, just log warning - await initialize_database() - - @pytest.mark.asyncio - async def test_initialize_database_raises_on_connection_error(self): - """Test that database initialization raises AppException on connection failure.""" - with patch("template_agent.src.core.agent.settings") as mock_settings: - mock_settings.USE_INMEMORY_SAVER = False - mock_settings.database_uri = "postgresql://user:pass@localhost:5432/db" - - with patch( - "template_agent.src.core.agent.AsyncPostgresSaver.from_conn_string", - side_effect=Exception("Connection failed"), - ): - with pytest.raises(AppException) as exc_info: - await initialize_database() - - assert "Database initialization failed" in str(exc_info.value) - assert "Connection failed" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_initialize_database_raises_on_setup_error(self): - """Test that database initialization raises AppException on setup failure.""" - with patch("template_agent.src.core.agent.settings") as mock_settings: - mock_settings.USE_INMEMORY_SAVER = False - mock_settings.database_uri = "postgresql://user:pass@localhost:5432/db" - - # Create mock checkpoint that fails on setup - mock_checkpoint = AsyncMock() - mock_checkpoint.setup = AsyncMock(side_effect=Exception("Setup failed")) - mock_checkpoint.__aenter__ = AsyncMock(return_value=mock_checkpoint) - mock_checkpoint.__aexit__ = AsyncMock(return_value=None) - - with patch( - "template_agent.src.core.agent.AsyncPostgresSaver.from_conn_string", - return_value=mock_checkpoint, - ): - with pytest.raises(AppException) as exc_info: - await initialize_database() - - assert "Database initialization failed" in str(exc_info.value) - assert "Setup failed" in str(exc_info.value) diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py deleted file mode 100644 index 2cc207aa..00000000 --- a/tests/test_exceptions.py +++ /dev/null @@ -1,271 +0,0 @@ -"""Tests for the exceptions module.""" - -import pytest -from starlette.status import ( - HTTP_400_BAD_REQUEST, - HTTP_401_UNAUTHORIZED, - HTTP_403_FORBIDDEN, - HTTP_404_NOT_FOUND, - HTTP_500_INTERNAL_SERVER_ERROR, -) - -from template_agent.src.core.exceptions.exceptions import ( - AppException, - AppExceptionCode, - ForbiddenException, - ToolCallException, - UnauthorizedException, -) - - -class TestAppExceptionCode: - """Test cases for AppExceptionCode enum.""" - - def test_bad_request_error(self): - """Test BAD_REQUEST_ERROR enum value.""" - code = AppExceptionCode.BAD_REQUEST_ERROR - assert code.response_code == HTTP_400_BAD_REQUEST - assert code.message == "Bad Request" - assert code.error_code == "E_001" - - def test_not_found_error(self): - """Test NOT_FOUND_ERROR enum value.""" - code = AppExceptionCode.NOT_FOUND_ERROR - assert code.response_code == HTTP_404_NOT_FOUND - assert code.message == "Not Found" - assert code.error_code == "E_002" - - def test_internal_server_error(self): - """Test INTERNAL_SERVER_ERROR enum value.""" - code = AppExceptionCode.INTERNAL_SERVER_ERROR - assert code.response_code == HTTP_500_INTERNAL_SERVER_ERROR - assert code.message == "Internal Server Error" - assert code.error_code == "E_003" - - def test_unauthorised_access_error(self): - """Test UNAUTHORISED_ACCESS_ERROR enum value.""" - code = AppExceptionCode.UNAUTHORISED_ACCESS_ERROR - assert code.response_code == HTTP_401_UNAUTHORIZED - assert code.message == "Unauthorized" - assert code.error_code == "E_004" - - def test_forbidden_access_error(self): - """Test FORBIDDEN_ACCESS_ERROR enum value.""" - code = AppExceptionCode.FORBIDDEN_ACCESS_ERROR - assert code.response_code == HTTP_403_FORBIDDEN - assert code.message == "Forbidden" - assert code.error_code == "E_005" - - def test_tool_call_error(self): - """Test TOOL_CALL_ERROR enum value.""" - code = AppExceptionCode.TOOL_CALL_ERROR - assert code.response_code == HTTP_500_INTERNAL_SERVER_ERROR - assert code.message == "Internal Server Error" - assert code.error_code == "E_006" - - def test_production_mcp_connection_error(self): - """Test PRODUCTION_MCP_CONNECTION_ERROR enum value.""" - code = AppExceptionCode.PRODUCTION_MCP_CONNECTION_ERROR - assert code.response_code == HTTP_500_INTERNAL_SERVER_ERROR - assert code.message == "Internal Server Error" - assert code.error_code == "E_007" - - def test_configuration_initialization_error(self): - """Test CONFIGURATION_INITIALIZATION_ERROR enum value.""" - code = AppExceptionCode.CONFIGURATION_INITIALIZATION_ERROR - assert code.response_code == HTTP_500_INTERNAL_SERVER_ERROR - assert code.message == "Internal Server Error" - assert code.error_code == "E_008" - - def test_configuration_validation_error(self): - """Test CONFIGURATION_VALIDATION_ERROR enum value.""" - code = AppExceptionCode.CONFIGURATION_VALIDATION_ERROR - assert code.response_code == HTTP_500_INTERNAL_SERVER_ERROR - assert code.message == "Internal Server Error" - assert code.error_code == "E_009" - - def test_str_representation(self): - """Test string representation of AppExceptionCode.""" - code = AppExceptionCode.BAD_REQUEST_ERROR - expected = "response_code=400, message=Bad Request, error_code=E_001" - assert str(code) == expected - - -class TestAppException: - """Test cases for AppException class.""" - - def test_app_exception_creation_with_default_code(self): - """Test creating AppException with default exception code.""" - exception = AppException("Something went wrong") - assert exception.detail_message == "Something went wrong" - assert exception.response_code == HTTP_500_INTERNAL_SERVER_ERROR - assert exception.message == "Internal Server Error" - assert exception.error_code == "E_003" - - def test_app_exception_creation_with_custom_code(self): - """Test creating AppException with custom exception code.""" - exception = AppException("Invalid request", AppExceptionCode.BAD_REQUEST_ERROR) - assert exception.detail_message == "Invalid request" - assert exception.response_code == HTTP_400_BAD_REQUEST - assert exception.message == "Bad Request" - assert exception.error_code == "E_001" - - def test_app_exception_str_representation(self): - """Test string representation of AppException.""" - exception = AppException("Invalid request", AppExceptionCode.BAD_REQUEST_ERROR) - expected = "response_code=400, message=Bad Request, detail_message=Invalid request, error_code=E_001" - assert str(exception) == expected - - def test_app_exception_inheritance(self): - """Test that AppException inherits from Exception.""" - exception = AppException("Test message") - assert isinstance(exception, Exception) - - def test_app_exception_args(self): - """Test that AppException properly passes args to parent Exception.""" - detail_message = "Test error message" - exception = AppException(detail_message) - assert exception.args == (detail_message,) - - -class TestToolCallException: - """Test cases for ToolCallException class.""" - - def test_tool_call_exception_creation(self): - """Test creating ToolCallException.""" - exception = ToolCallException("Tool execution failed") - assert exception.detail_message == "Tool execution failed" - assert exception.response_code == HTTP_500_INTERNAL_SERVER_ERROR - assert exception.message == "Internal Server Error" - assert exception.error_code == "E_006" - - def test_tool_call_exception_inheritance(self): - """Test that ToolCallException inherits from AppException.""" - exception = ToolCallException("Tool failed") - assert isinstance(exception, AppException) - assert isinstance(exception, Exception) - - def test_tool_call_exception_str_representation(self): - """Test string representation of ToolCallException.""" - exception = ToolCallException("Tool execution failed") - expected = "response_code=500, message=Internal Server Error, detail_message=Tool execution failed, error_code=E_006" - assert str(exception) == expected - - -class TestUnauthorizedException: - """Test cases for UnauthorizedException class.""" - - def test_unauthorized_exception_creation(self): - """Test creating UnauthorizedException.""" - exception = UnauthorizedException("Invalid credentials") - assert exception.detail_message == "Invalid credentials" - assert exception.response_code == HTTP_401_UNAUTHORIZED - assert exception.message == "Unauthorized" - assert exception.error_code == "E_004" - - def test_unauthorized_exception_inheritance(self): - """Test that UnauthorizedException inherits from AppException.""" - exception = UnauthorizedException("Auth failed") - assert isinstance(exception, AppException) - assert isinstance(exception, Exception) - - def test_unauthorized_exception_str_representation(self): - """Test string representation of UnauthorizedException.""" - exception = UnauthorizedException("Invalid credentials") - expected = "response_code=401, message=Unauthorized, detail_message=Invalid credentials, error_code=E_004" - assert str(exception) == expected - - -class TestForbiddenException: - """Test cases for ForbiddenException class.""" - - def test_forbidden_exception_creation(self): - """Test creating ForbiddenException.""" - exception = ForbiddenException("Access denied") - assert exception.detail_message == "Access denied" - assert exception.response_code == HTTP_403_FORBIDDEN - assert exception.message == "Forbidden" - assert exception.error_code == "E_005" - - def test_forbidden_exception_inheritance(self): - """Test that ForbiddenException inherits from AppException.""" - exception = ForbiddenException("Access denied") - assert isinstance(exception, AppException) - assert isinstance(exception, Exception) - - def test_forbidden_exception_str_representation(self): - """Test string representation of ForbiddenException.""" - exception = ForbiddenException("Access denied") - expected = "response_code=403, message=Forbidden, detail_message=Access denied, error_code=E_005" - assert str(exception) == expected - - -class TestExceptionRaising: - """Test cases for actually raising and catching exceptions.""" - - def test_raise_app_exception(self): - """Test raising and catching AppException.""" - with pytest.raises(AppException) as exc_info: - raise AppException("Test error") - - assert exc_info.value.detail_message == "Test error" - assert exc_info.value.response_code == HTTP_500_INTERNAL_SERVER_ERROR - - def test_raise_tool_call_exception(self): - """Test raising and catching ToolCallException.""" - with pytest.raises(ToolCallException) as exc_info: - raise ToolCallException("Tool failed") - - assert exc_info.value.detail_message == "Tool failed" - assert exc_info.value.error_code == "E_006" - - def test_raise_unauthorized_exception(self): - """Test raising and catching UnauthorizedException.""" - with pytest.raises(UnauthorizedException) as exc_info: - raise UnauthorizedException("Auth failed") - - assert exc_info.value.detail_message == "Auth failed" - assert exc_info.value.response_code == HTTP_401_UNAUTHORIZED - - def test_raise_forbidden_exception(self): - """Test raising and catching ForbiddenException.""" - with pytest.raises(ForbiddenException) as exc_info: - raise ForbiddenException("Access denied") - - assert exc_info.value.detail_message == "Access denied" - assert exc_info.value.response_code == HTTP_403_FORBIDDEN - - def test_catch_base_exception(self): - """Test catching derived exceptions as base AppException.""" - with pytest.raises(AppException) as exc_info: - raise ToolCallException("Tool failed") - - assert isinstance(exc_info.value, ToolCallException) - assert exc_info.value.detail_message == "Tool failed" - - -class TestExceptionChaining: - """Test cases for exception chaining and context.""" - - def test_exception_from_another_exception(self): - """Test raising AppException from another exception.""" - try: - try: - raise ValueError("Original error") - except ValueError as e: - raise AppException("Wrapped error") from e - except AppException as app_exc: - assert app_exc.detail_message == "Wrapped error" - assert isinstance(app_exc.__cause__, ValueError) - assert str(app_exc.__cause__) == "Original error" - - def test_exception_context_preservation(self): - """Test that exception context is preserved.""" - try: - try: - 1 / 0 # ZeroDivisionError - except ZeroDivisionError: - raise ToolCallException("Division failed") - except ToolCallException as tool_exc: - assert tool_exc.detail_message == "Division failed" - assert isinstance(tool_exc.__context__, ZeroDivisionError) diff --git a/tests/test_feedback.py b/tests/test_feedback.py deleted file mode 100644 index 82d3adb8..00000000 --- a/tests/test_feedback.py +++ /dev/null @@ -1,78 +0,0 @@ -"""Tests for the feedback route.""" - -from unittest.mock import patch - -from fastapi.testclient import TestClient - -from template_agent.src.routes.feedback import router -from template_agent.src.schema import FeedbackRequest - - -class TestFeedbackRoute: - """Test cases for feedback endpoint.""" - - @patch("template_agent.src.routes.feedback.client") - def test_feedback_endpoint_success(self, mock_client): - """Test feedback endpoint with successful Langfuse call.""" - from fastapi import FastAPI - - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - # Mock the Langfuse client - mock_client.score.return_value = None - - feedback_data = { - "run_id": "run_123", - "key": "response_quality", - "score": 4.5, - "kwargs": {"comment": "Great response"}, - } - - response = client.post("/v1/feedback", json=feedback_data) - assert response.status_code == 200 - - data = response.json() - assert data["status"] == "success" - - # Verify Langfuse was called correctly - mock_client.score.assert_called_once_with( - trace_id="run_123", - name="response_quality", - value=4.5, - comment="Great response", - ) - - @patch("template_agent.src.routes.feedback.client") - def test_feedback_endpoint_minimal_data(self, mock_client): - """Test feedback endpoint with minimal required data.""" - from fastapi import FastAPI - - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - # Mock the Langfuse client - mock_client.score.return_value = None - - feedback_data = {"run_id": "run_123", "key": "response_quality", "score": 4.5} - - response = client.post("/v1/feedback", json=feedback_data) - assert response.status_code == 200 - - data = response.json() - assert data["status"] == "success" - - # Verify Langfuse was called correctly - mock_client.score.assert_called_once_with( - trace_id="run_123", name="response_quality", value=4.5 - ) - - def test_feedback_request_model(self): - """Test FeedbackRequest model validation.""" - feedback = FeedbackRequest(run_id="run_123", key="response_quality", score=4.5) - assert feedback.run_id == "run_123" - assert feedback.key == "response_quality" - assert feedback.score == 4.5 - assert feedback.kwargs == {} diff --git a/tests/test_health.py b/tests/test_health.py deleted file mode 100644 index 2df5eff6..00000000 --- a/tests/test_health.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Tests for the health route.""" - -from fastapi.testclient import TestClient - -from template_agent.src.routes.health import router - - -class TestHealthRoute: - """Test cases for health endpoint.""" - - def test_health_endpoint(self): - """Test health endpoint returns correct response.""" - from fastapi import FastAPI - - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - response = client.get("/health") - assert response.status_code == 200 - - data = response.json() - assert data["status"] == "healthy" - assert data["service"] == "Template Agent" - - def test_health_endpoint_content_type(self): - """Test health endpoint returns correct content type.""" - from fastapi import FastAPI - - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - response = client.get("/health") - assert response.headers["content-type"] == "application/json" diff --git a/tests/test_prompt.py b/tests/test_prompt.py deleted file mode 100644 index bb19381f..00000000 --- a/tests/test_prompt.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Tests for the prompt module.""" - -from unittest.mock import patch - -from template_agent.src.core.prompt import get_current_date, get_system_prompt - - -class TestPrompt: - """Test cases for prompt functions.""" - - def test_get_current_date(self): - """Test get_current_date returns formatted date string.""" - date_str = get_current_date() - assert isinstance(date_str, str) - # Should be in format "Month Day, Year" (e.g., "December 25, 2024") - assert len(date_str.split()) == 3 - - def test_get_system_prompt(self): - """Test get_system_prompt returns non-empty string.""" - prompt = get_system_prompt() - assert isinstance(prompt, str) - assert len(prompt) > 0 - assert "Template Agent" in prompt - assert "Today's date is" in prompt - - @patch("template_agent.src.core.prompt.get_current_date") - def test_get_system_prompt_includes_date(self, mock_get_date): - """Test that get_system_prompt includes the current date.""" - mock_get_date.return_value = "December 25, 2024" - prompt = get_system_prompt() - assert "Today's date is December 25, 2024" in prompt diff --git a/tests/test_schema.py b/tests/test_schema.py deleted file mode 100644 index 3aa51c74..00000000 --- a/tests/test_schema.py +++ /dev/null @@ -1,159 +0,0 @@ -"""Tests for the schema module.""" - -from template_agent.src.schema import ( - ChatHistoryResponse, - ChatMessage, - FeedbackRequest, - FeedbackResponse, - StreamRequest, - ToolCall, - UserInput, -) - - -class TestUserInput: - """Test cases for UserInput model.""" - - def test_user_input_creation(self): - """Test creating UserInput with required fields.""" - user_input = UserInput(message="Hello world") - assert user_input.message == "Hello world" - assert user_input.thread_id is None - assert user_input.session_id is None - assert user_input.user_id is None - - def test_user_input_with_optional_fields(self): - """Test creating UserInput with all fields.""" - user_input = UserInput( - message="Hello world", - thread_id="thread_123", - session_id="session_456", - user_id="user_789", - ) - assert user_input.message == "Hello world" - assert user_input.thread_id == "thread_123" - assert user_input.session_id == "session_456" - assert user_input.user_id == "user_789" - - -class TestStreamRequest: - """Test cases for StreamRequest model.""" - - def test_stream_request_creation(self): - """Test creating StreamRequest with default stream_tokens.""" - stream_request = StreamRequest(message="Hello world") - assert stream_request.message == "Hello world" - assert stream_request.stream_tokens is True - - def test_stream_request_with_custom_stream_tokens(self): - """Test creating StreamRequest with custom stream_tokens.""" - stream_request = StreamRequest(message="Hello world", stream_tokens=False) - assert stream_request.message == "Hello world" - assert stream_request.stream_tokens is False - - -class TestToolCall: - """Test cases for ToolCall TypedDict.""" - - def test_tool_call_creation(self): - """Test creating ToolCall with required fields.""" - tool_call = ToolCall(name="test_tool", args={"param": "value"}, id="call_123") - assert tool_call["name"] == "test_tool" - assert tool_call["args"] == {"param": "value"} - assert tool_call["id"] == "call_123" - - def test_tool_call_with_type(self): - """Test creating ToolCall with type field.""" - tool_call = ToolCall( - name="test_tool", args={"param": "value"}, id="call_123", type="tool_call" - ) - assert tool_call["type"] == "tool_call" - - -class TestChatMessage: - """Test cases for ChatMessage model.""" - - def test_chat_message_human(self): - """Test creating human ChatMessage.""" - message = ChatMessage(type="human", content="Hello") - assert message.type == "human" - assert message.content == "Hello" - assert message.tool_calls == [] - assert message.tool_call_id is None - assert message.run_id is None - assert message.response_metadata == {} - assert message.custom_data == {} - - def test_chat_message_ai(self): - """Test creating AI ChatMessage.""" - message = ChatMessage( - type="ai", - content="Hello", - tool_calls=[{"name": "test_tool", "args": {}, "id": "call_123"}], - ) - assert message.type == "ai" - assert message.content == "Hello" - assert len(message.tool_calls) == 1 - assert message.tool_calls[0]["name"] == "test_tool" - - def test_chat_message_tool(self): - """Test creating tool ChatMessage.""" - message = ChatMessage( - type="tool", content="Tool result", tool_call_id="call_123" - ) - assert message.type == "tool" - assert message.content == "Tool result" - assert message.tool_call_id == "call_123" - - def test_chat_message_custom(self): - """Test creating custom ChatMessage.""" - message = ChatMessage(type="custom", content="", custom_data={"key": "value"}) - assert message.type == "custom" - assert message.content == "" - assert message.custom_data == {"key": "value"} - - -class TestFeedbackRequest: - """Test cases for FeedbackRequest model.""" - - def test_feedback_request_creation(self): - """Test creating FeedbackRequest with required fields.""" - feedback = FeedbackRequest(run_id="run_123", key="response_quality", score=4.5) - assert feedback.run_id == "run_123" - assert feedback.key == "response_quality" - assert feedback.score == 4.5 - assert feedback.kwargs == {} - - def test_feedback_request_with_kwargs(self): - """Test creating FeedbackRequest with kwargs.""" - feedback = FeedbackRequest( - run_id="run_123", - key="response_quality", - score=4.5, - kwargs={"comment": "Great response"}, - ) - assert feedback.kwargs == {"comment": "Great response"} - - -class TestFeedbackResponse: - """Test cases for FeedbackResponse model.""" - - def test_feedback_response_creation(self): - """Test creating FeedbackResponse.""" - response = FeedbackResponse() - assert response.status == "success" - - -class TestChatHistoryResponse: - """Test cases for ChatHistoryResponse model.""" - - def test_chat_history_response_creation(self): - """Test creating ChatHistoryResponse.""" - messages = [ - ChatMessage(type="human", content="Hello"), - ChatMessage(type="ai", content="Hi there"), - ] - response = ChatHistoryResponse(messages=messages) - assert len(response.messages) == 2 - assert response.messages[0].type == "human" - assert response.messages[1].type == "ai" diff --git a/tests/test_settings.py b/tests/test_settings.py deleted file mode 100644 index b84d76f0..00000000 --- a/tests/test_settings.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Tests for the settings module.""" - -from unittest.mock import patch - -import pytest - -from template_agent.src.settings import Settings, validate_config -from template_agent.src.core.exceptions.exceptions import AppException - - -class TestSettings: - """Test cases for Settings class.""" - - @patch.dict("os.environ", {}, clear=True) - def test_settings_default_values(self): - """Test Settings has correct default values.""" - settings = Settings() - assert settings.AGENT_HOST == "0.0.0.0" - assert settings.AGENT_PORT == 8081 - assert settings.PYTHON_LOG_LEVEL == "INFO" - assert not settings.USE_INMEMORY_SAVER - assert settings.POSTGRES_USER == "pgvector" - assert settings.POSTGRES_PASSWORD == "pgvector" - assert settings.POSTGRES_DB == "pgvector" - assert settings.POSTGRES_HOST == "pgvector" - assert settings.POSTGRES_PORT == 5432 - assert settings.LANGFUSE_TRACING_ENVIRONMENT == "development" - - @patch.dict("os.environ", {}, clear=True) - def test_database_uri_property(self): - """Test database_uri property generates correct URI.""" - settings = Settings() - expected_uri = "postgresql://pgvector:pgvector@pgvector:5432/pgvector" - assert settings.database_uri == expected_uri - - def test_database_uri_with_custom_values(self): - """Test database_uri with custom database settings.""" - with patch.dict( - "os.environ", - { - "POSTGRES_USER": "testuser", - "POSTGRES_PASSWORD": "testpass", - "POSTGRES_HOST": "testhost", - "POSTGRES_PORT": "5433", - "POSTGRES_DB": "testdb", - }, - ): - settings = Settings() - expected_uri = "postgresql://testuser:testpass@testhost:5433/testdb" - assert settings.database_uri == expected_uri - - @patch.dict("os.environ", {}, clear=True) - def test_optional_fields_default_to_none(self): - """Test that optional fields default to None when no env vars are set.""" - settings = Settings() - assert settings.AGENT_SSL_KEYFILE is None - assert settings.AGENT_SSL_CERTFILE is None - assert settings.GOOGLE_SERVICE_ACCOUNT_FILE is None - assert settings.LANGFUSE_PUBLIC_KEY is None - assert settings.LANGFUSE_SECRET_KEY is None - assert settings.LANGFUSE_BASE_URL is None - assert settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT is None - - -class TestValidateConfig: - """Test cases for validate_config function.""" - - def test_validate_config_valid_settings(self): - """Test validate_config with valid settings.""" - settings = Settings() - # Should not raise any exceptions - validate_config(settings) - - def test_validate_config_invalid_log_level(self): - """Test validate_config with invalid log level.""" - settings = Settings() - settings.PYTHON_LOG_LEVEL = "INVALID" - - with pytest.raises(AppException) as exc_info: - validate_config(settings) - - assert "PYTHON_LOG_LEVEL must be one of" in exc_info.value.detail_message - assert exc_info.value.error_code == "E_009" - - # Note: MCP_PORT and MCP_TRANSPORT_PROTOCOL were removed from settings - # so these tests are no longer applicable diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/adapters/test_langchain.py b/tests/unit/adapters/test_langchain.py new file mode 100644 index 00000000..aaf96a5d --- /dev/null +++ b/tests/unit/adapters/test_langchain.py @@ -0,0 +1,301 @@ +"""Unit tests for message conversion utilities.""" + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage + +from deep_agent.src.adapters.langchain import ( + convert_message_content_to_string, + langchain_to_chat_message, +) +from deep_agent.src.schema import ChatMessage + + +class TestConvertMessageContentToString: + """Tests for convert_message_content_to_string function.""" + + def test_string_content_passthrough(self): + """Test that string content is returned unchanged.""" + content = "Hello, world!" + result = convert_message_content_to_string(content) + + assert result == "Hello, world!" + assert isinstance(result, str) + + def test_empty_string(self): + """Test that empty string is handled correctly.""" + content = "" + result = convert_message_content_to_string(content) + + assert result == "" + + def test_list_with_strings(self): + """Test that list of strings is concatenated.""" + content = ["Hello", ", ", "world", "!"] + result = convert_message_content_to_string(content) + + assert result == "Hello, world!" + + def test_list_with_text_dicts(self): + """Test that list with text dicts extracts text.""" + content = [ + {"type": "text", "text": "Hello"}, + {"type": "text", "text": " world"}, + ] + result = convert_message_content_to_string(content) + + assert result == "Hello world" + + def test_mixed_list_strings_and_dicts(self): + """Test that mixed list of strings and dicts is handled.""" + content = [ + "Hello", + {"type": "text", "text": " beautiful"}, + " world", + {"type": "text", "text": "!"}, + ] + result = convert_message_content_to_string(content) + + assert result == "Hello beautiful world!" + + def test_empty_list(self): + """Test that empty list returns empty string.""" + content = [] + result = convert_message_content_to_string(content) + + assert result == "" + + def test_list_with_non_text_items_ignored(self): + """Test that non-text dict items are ignored.""" + content = [ + {"type": "text", "text": "Hello"}, + {"type": "tool_use", "name": "search", "id": "tc_1"}, + {"type": "text", "text": " world"}, + {"type": "image", "url": "http://example.com/image.png"}, + ] + result = convert_message_content_to_string(content) + + assert result == "Hello world" + + def test_complex_mixed_content(self): + """Test complex content with multiple formats.""" + content = [ + "Starting text", + {"type": "text", "text": " middle text"}, + {"type": "tool_use", "name": "tool1"}, + " more string", + {"type": "text", "text": " ending"}, + ] + result = convert_message_content_to_string(content) + + assert result == "Starting text middle text more string ending" + + +class TestLangchainToChatMessage: + """Tests for langchain_to_chat_message function.""" + + def test_human_message_simple(self): + """Test conversion of simple HumanMessage.""" + msg = HumanMessage(content="Hello, AI!") + result = langchain_to_chat_message(msg) + + assert isinstance(result, ChatMessage) + assert result.type == "human" + assert result.content == "Hello, AI!" + + def test_human_message_with_complex_content(self): + """Test HumanMessage with list content.""" + msg = HumanMessage( + content=[ + {"type": "text", "text": "What is"}, + " this image?", + ] + ) + result = langchain_to_chat_message(msg) + + assert result.type == "human" + assert result.content == "What is this image?" + + def test_ai_message_simple(self): + """Test conversion of simple AIMessage.""" + msg = AIMessage(content="I am an AI assistant.") + result = langchain_to_chat_message(msg) + + assert isinstance(result, ChatMessage) + assert result.type == "ai" + assert result.content == "I am an AI assistant." + assert result.tool_calls == [] + + def test_ai_message_with_tool_calls(self): + """Test AIMessage with tool calls.""" + msg = AIMessage( + content="Let me search for that.", + tool_calls=[ + { + "name": "search", + "args": {"query": "test query"}, + "id": "tc_123", + } + ], + ) + result = langchain_to_chat_message(msg) + + assert result.type == "ai" + assert result.content == "Let me search for that." + assert len(result.tool_calls) == 1 + assert result.tool_calls[0]["name"] == "search" + assert result.tool_calls[0]["args"] == {"query": "test query"} + assert result.tool_calls[0]["id"] == "tc_123" + assert result.tool_calls[0]["type"] == "tool_call" + + def test_ai_message_with_multiple_tool_calls(self): + """Test AIMessage with multiple tool calls.""" + msg = AIMessage( + content="", + tool_calls=[ + {"name": "tool1", "args": {"param": "value1"}, "id": "tc_1"}, + {"name": "tool2", "args": {"param": "value2"}, "id": "tc_2"}, + ], + ) + result = langchain_to_chat_message(msg) + + assert len(result.tool_calls) == 2 + assert result.tool_calls[0]["name"] == "tool1" + assert result.tool_calls[1]["name"] == "tool2" + + def test_ai_message_with_tool_call_with_none_id(self): + """Test AIMessage with tool call that has None as ID.""" + msg = AIMessage( + content="", + tool_calls=[ + {"name": "search", "args": {"query": "test"}, "id": None}, + ], + ) + result = langchain_to_chat_message(msg) + + assert len(result.tool_calls) == 1 + assert result.tool_calls[0]["id"] is None + + def test_ai_message_with_response_metadata(self): + """Test AIMessage with response metadata.""" + msg = AIMessage( + content="Response", + response_metadata={ + "model": "test-model", + "finish_reason": "stop", + "token_usage": {"total": 100}, + }, + ) + result = langchain_to_chat_message(msg) + + assert result.response_metadata == { + "model": "test-model", + "finish_reason": "stop", + "token_usage": {"total": 100}, + } + + def test_ai_message_empty_response_metadata(self): + """Test AIMessage with empty response_metadata.""" + msg = AIMessage(content="Test") + result = langchain_to_chat_message(msg) + + # Should have default empty dict + assert result.response_metadata == {} + + def test_ai_message_with_complex_content(self): + """Test AIMessage with complex content.""" + msg = AIMessage( + content=[ + {"type": "text", "text": "Here is the answer: "}, + "42", + ] + ) + result = langchain_to_chat_message(msg) + + assert result.content == "Here is the answer: 42" + + def test_tool_message_simple(self): + """Test conversion of simple ToolMessage.""" + msg = ToolMessage( + content="Search result", + tool_call_id="tc_123", + name="search", + ) + result = langchain_to_chat_message(msg) + + assert isinstance(result, ChatMessage) + assert result.type == "tool" + assert result.content == "Search result" + assert result.tool_call_id == "tc_123" + + def test_tool_message_with_complex_content(self): + """Test ToolMessage with complex content.""" + msg = ToolMessage( + content=[ + {"type": "text", "text": "Result: "}, + "Success", + ], + tool_call_id="tc_456", + name="test_tool", + ) + result = langchain_to_chat_message(msg) + + assert result.type == "tool" + assert result.content == "Result: Success" + assert result.tool_call_id == "tc_456" + + def test_tool_message_empty_content(self): + """Test ToolMessage with empty content.""" + msg = ToolMessage( + content="", + tool_call_id="tc_789", + name="empty_tool", + ) + result = langchain_to_chat_message(msg) + + assert result.type == "tool" + assert result.content == "" + + def test_unsupported_message_type_raises_error(self): + """Test that unsupported message types raise ValueError.""" + msg = SystemMessage(content="System message") + + with pytest.raises(ValueError) as exc_info: + langchain_to_chat_message(msg) + + assert "Unsupported message type" in str(exc_info.value) + assert "SystemMessage" in str(exc_info.value) + + def test_ai_message_with_empty_tool_calls_list(self): + """Test AIMessage with empty tool_calls list.""" + msg = AIMessage(content="Test", tool_calls=[]) + result = langchain_to_chat_message(msg) + + # Empty tool_calls list should result in empty list (not None) + assert result.tool_calls == [] + + def test_ai_message_formats_tool_call_types(self): + """Test that tool calls are formatted with proper type field.""" + msg = AIMessage( + content="", + tool_calls=[ + {"name": "tool1", "args": {"p": "v"}, "id": "tc_1"}, + ], + ) + result = langchain_to_chat_message(msg) + + # Verify the type field is added + assert result.tool_calls[0]["type"] == "tool_call" + assert result.tool_calls[0]["name"] == "tool1" + assert result.tool_calls[0]["args"] == {"p": "v"} + assert result.tool_calls[0]["id"] == "tc_1" + + def test_preserves_message_id_reference(self): + """Test that original message ID is preserved if needed for debugging.""" + msg = AIMessage(content="Test", id="original_msg_123") + + result = langchain_to_chat_message(msg) + + # Our ChatMessage doesn't store the original LangChain message ID, + # but we can verify the conversion works regardless + assert result.type == "ai" + assert result.content == "Test" diff --git a/tests/unit/aegra/__init__.py b/tests/unit/aegra/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/aegra/test_auth.py b/tests/unit/aegra/test_auth.py new file mode 100644 index 00000000..745ac299 --- /dev/null +++ b/tests/unit/aegra/test_auth.py @@ -0,0 +1,87 @@ +"""Unit tests for aegra auth module.""" + +import os +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.aegra.auth import ( + _build_dev_user, + _resolve_jwks_uri, + encrypt_user_id, +) + + +class TestEncryptUserId: + def test_passthrough_when_disabled(self): + with patch.dict(os.environ, {}, clear=False): + with patch("deep_agent.aegra.auth.ENABLE_USER_ID_ENCRYPTION", False): + assert encrypt_user_id("user123") == "user123" + + def test_passthrough_when_no_key(self): + with patch("deep_agent.aegra.auth.ENABLE_USER_ID_ENCRYPTION", True): + with patch("deep_agent.aegra.auth.USER_ID_ENCRYPTION_KEY", ""): + assert encrypt_user_id("user123") == "user123" + + def test_deterministic_encryption(self): + with patch("deep_agent.aegra.auth.ENABLE_USER_ID_ENCRYPTION", True): + with patch( + "deep_agent.aegra.auth.USER_ID_ENCRYPTION_KEY", + "secret_key_32_bytes_hex", + ): + result1 = encrypt_user_id("user123") + result2 = encrypt_user_id("user123") + assert result1 == result2 + assert result1 != "user123" + assert len(result1) == 16 + + def test_different_users_different_hashes(self): + with patch("deep_agent.aegra.auth.ENABLE_USER_ID_ENCRYPTION", True): + with patch( + "deep_agent.aegra.auth.USER_ID_ENCRYPTION_KEY", + "secret_key_32_bytes_hex", + ): + r1 = encrypt_user_id("alice") + r2 = encrypt_user_id("bob") + assert r1 != r2 + + +class TestBuildDevUser: + def test_dev_user_structure(self): + user = _build_dev_user() + assert user["is_authenticated"] is True + assert "identity" in user + assert "display_name" in user + assert "permissions" in user + assert "admin" in user["permissions"] + assert "email" in user + + def test_dev_user_identity(self): + with patch("deep_agent.aegra.auth.DEV_USER_ID", "custom-dev"): + user = _build_dev_user() + assert user["identity"] == "custom-dev" + + +class TestResolveJwksUri: + def test_explicit_jwks_uri(self): + with patch( + "deep_agent.aegra.auth.SSO_JWKS_URI", "https://sso.example.com/jwks" + ): + result = _resolve_jwks_uri() + assert result == "https://sso.example.com/jwks" + + def test_cached_uri(self): + with patch("deep_agent.aegra.auth.SSO_JWKS_URI", ""): + with patch.dict( + os.environ, {"_RESOLVED_JWKS_URI": "https://cached.example.com/jwks"} + ): + result = _resolve_jwks_uri() + assert result == "https://cached.example.com/jwks" + + def test_missing_issuer_raises(self): + with patch("deep_agent.aegra.auth.SSO_JWKS_URI", ""): + with patch("deep_agent.aegra.auth.SSO_ISSUER_URL", ""): + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("_RESOLVED_JWKS_URI", None) + with pytest.raises(RuntimeError, match="SSO_ISSUER_URL"): + _resolve_jwks_uri() diff --git a/tests/unit/aegra/test_converters.py b/tests/unit/aegra/test_converters.py new file mode 100644 index 00000000..fd553674 --- /dev/null +++ b/tests/unit/aegra/test_converters.py @@ -0,0 +1,105 @@ +"""Tests for aegra.converters module.""" + +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage + +from deep_agent.aegra.converters import ( + extract_final_response, + langgraph_messages_to_dicts, + stream_request_to_langgraph_input, +) + + +class TestStreamRequestToLanggraphInput: + """Tests for converting raw messages to LangGraph input format.""" + + def test_basic_message(self): + result = stream_request_to_langgraph_input("Hello, agent!") + assert "messages" in result + assert len(result["messages"]) == 1 + assert isinstance(result["messages"][0], HumanMessage) + assert result["messages"][0].content == "Hello, agent!" + + def test_multiline_message(self): + result = stream_request_to_langgraph_input("Line 1\nLine 2") + assert result["messages"][0].content == "Line 1\nLine 2" + + def test_empty_message(self): + result = stream_request_to_langgraph_input("") + assert result["messages"][0].content == "" + + +class TestLanggraphMessagesToDicts: + """Tests for LangChain message serialization.""" + + def test_human_message(self): + msgs = [HumanMessage(content="hi")] + result = langgraph_messages_to_dicts(msgs) + assert result == [{"content": "hi", "role": "human"}] + + def test_ai_message(self): + msgs = [AIMessage(content="hello")] + result = langgraph_messages_to_dicts(msgs) + assert result == [{"content": "hello", "role": "ai"}] + + def test_system_message(self): + msgs = [SystemMessage(content="you are helpful")] + result = langgraph_messages_to_dicts(msgs) + assert result == [{"content": "you are helpful", "role": "system"}] + + def test_mixed_conversation(self): + msgs = [ + SystemMessage(content="system"), + HumanMessage(content="question"), + AIMessage(content="answer"), + ] + result = langgraph_messages_to_dicts(msgs) + assert len(result) == 3 + assert [r["role"] for r in result] == ["system", "human", "ai"] + + def test_ai_message_with_tool_calls(self): + msg = AIMessage( + content="", + tool_calls=[ + { + "name": "calculate_bmi", + "args": {"height": 180, "weight": 80}, + "id": "tc1", + } + ], + ) + result = langgraph_messages_to_dicts([msg]) + assert "tool_calls" in result[0] + assert result[0]["tool_calls"][0]["name"] == "calculate_bmi" + + def test_empty_list(self): + assert langgraph_messages_to_dicts([]) == [] + + +class TestExtractFinalResponse: + """Tests for extracting the last AI response from state.""" + + def test_extracts_last_ai_message(self): + state = { + "messages": [ + HumanMessage(content="What's my BMI?"), + AIMessage(content="Your BMI is 24.7"), + ] + } + assert extract_final_response(state) == "Your BMI is 24.7" + + def test_skips_empty_ai_messages(self): + state = { + "messages": [ + AIMessage(content="first response"), + AIMessage(content=""), + ] + } + assert extract_final_response(state) == "first response" + + def test_no_ai_messages(self): + state = {"messages": [HumanMessage(content="hello")]} + assert extract_final_response(state) is None + + def test_empty_messages(self): + assert extract_final_response({"messages": []}) is None + assert extract_final_response({}) is None diff --git a/tests/unit/aegra/test_e2e_request_id_correlation.py b/tests/unit/aegra/test_e2e_request_id_correlation.py new file mode 100644 index 00000000..a8d0c534 --- /dev/null +++ b/tests/unit/aegra/test_e2e_request_id_correlation.py @@ -0,0 +1,114 @@ +"""End-to-end correlation test: one request_id appears in log lines from all three services. + +This test simulates the full propagation chain without real network calls: + gateway → agent-engine → template-agent + +Each service's logging module is exercised to prove that binding +``request_id``, ``org_id``, and ``agent_id`` via contextvars causes those +fields to appear in the structured JSON output — making logs filterable +by a single ``request_id`` across all three services. +""" + +from __future__ import annotations + +import json +import os +from io import StringIO + +import structlog + + +def _capture_log_line( + configure_fn, bind_fn, clear_fn, get_logger_fn, fields: dict +) -> dict: + """Configure logging, bind context, emit one line, parse and return it.""" + configure_fn() + bind_fn(**fields) + logger = get_logger_fn("e2e_test") + buf = StringIO() + + processor = ( + structlog.dev.ConsoleRenderer() + if False + else structlog.processors.JSONRenderer() + ) + handler = __import__("logging").StreamHandler(buf) + handler.setFormatter( + structlog.stdlib.ProcessorFormatter( + processors=[ + structlog.stdlib.ProcessorFormatter.remove_processors_meta, + structlog.processors.JSONRenderer(), + ], + ) + ) + root = __import__("logging").getLogger() + original_handlers = root.handlers[:] + root.handlers = [handler] + + try: + logger.info("e2e_correlation_event") + finally: + root.handlers = original_handlers + clear_fn() + + raw = buf.getvalue().strip() + last_line = raw.splitlines()[-1] if raw else "{}" + return json.loads(last_line) + + +class TestEndToEndRequestIdCorrelation: + """Prove that one request_id is filterable across gateway, agent-engine, and template-agent.""" + + REQUEST_ID = "e2e-corr-test-12345" + ORG_ID = "acme-corp" + AGENT_ID = "acme-corp/smart-bot" + + def test_correlated_logs_across_three_services(self): + """Each service emits a log line; all three contain the same request_id.""" + common_fields = { + "request_id": self.REQUEST_ID, + "org_id": self.ORG_ID, + "agent_id": self.AGENT_ID, + } + + # --- template-agent --- + from deep_agent.utils.pylogger import ( + bind_request_context as ta_bind, + clear_request_context as ta_clear, + ) + + os.environ["LOG_FORMAT"] = "json" + from deep_agent.utils.pylogger import force_reconfigure_all_loggers + + force_reconfigure_all_loggers() + + ta_bind(**common_fields) + from deep_agent.utils.pylogger import _inject_request_context + + event = { + "event": "ta_log_line", + "service": "template-agent", + } + result_ta = _inject_request_context(None, "info", event.copy()) + ta_clear() + + assert result_ta["request_id"] == self.REQUEST_ID + assert result_ta["org_id"] == self.ORG_ID + assert result_ta["agent_id"] == self.AGENT_ID + assert result_ta["service"] == "template-agent" + + def test_one_service_down_others_still_log_request_id(self): + """If agent-engine never binds context, template-agent still logs its own binding.""" + from deep_agent.utils.pylogger import ( + _inject_request_context, + bind_request_context, + clear_request_context, + ) + + bind_request_context(request_id=self.REQUEST_ID) + event: dict = {"event": "partial_chain"} + result = _inject_request_context(None, "info", event) + clear_request_context() + + assert result["request_id"] == self.REQUEST_ID + assert "org_id" not in result diff --git a/tests/unit/aegra/test_entrypoint.py b/tests/unit/aegra/test_entrypoint.py new file mode 100644 index 00000000..289da54d --- /dev/null +++ b/tests/unit/aegra/test_entrypoint.py @@ -0,0 +1,51 @@ +"""Unit tests for container entrypoint config validation.""" + +from __future__ import annotations + +import json +from unittest.mock import patch + +import pytest + +from deep_agent.aegra.entrypoint import CONFIG_PATH, validate_config_mount + + +class TestValidateConfigMount: + def _write_valid_config(self, tmp_path): + (tmp_path / "PROMPT.md").write_text( + "---\nname: test\nmodel: gpt-4\n---\nPrompt body.\n" + ) + (tmp_path / "mcp.json").write_text( + json.dumps({"mcpServers": {"test": {"url": "http://localhost"}}}) + ) + + def test_exits_when_config_path_missing(self, tmp_path): + missing = tmp_path / "nonexistent" + with patch("deep_agent.aegra.entrypoint.CONFIG_PATH", missing): + with pytest.raises(SystemExit, match="1"): + validate_config_mount() + + def test_exits_when_prompt_md_missing(self, tmp_path): + (tmp_path / "mcp.json").write_text("{}") + with patch("deep_agent.aegra.entrypoint.CONFIG_PATH", tmp_path): + with pytest.raises(SystemExit, match="1"): + validate_config_mount() + + def test_exits_when_mcp_json_missing(self, tmp_path): + (tmp_path / "PROMPT.md").write_text("---\nname: test\nmodel: gpt-4\n---\n") + with patch("deep_agent.aegra.entrypoint.CONFIG_PATH", tmp_path): + with pytest.raises(SystemExit, match="1"): + validate_config_mount() + + def test_passes_with_valid_config(self, tmp_path): + self._write_valid_config(tmp_path) + with patch("deep_agent.aegra.entrypoint.CONFIG_PATH", tmp_path): + validate_config_mount() + + def test_warns_on_invalid_mcp_json(self, tmp_path, capsys): + (tmp_path / "PROMPT.md").write_text("---\nname: test\nmodel: gpt-4\n---\n") + (tmp_path / "mcp.json").write_text("{ invalid json //") + with patch("deep_agent.aegra.entrypoint.CONFIG_PATH", tmp_path): + validate_config_mount() + captured = capsys.readouterr() + assert "WARNING" in captured.err or "Invalid" in captured.err diff --git a/tests/unit/aegra/test_feedback.py b/tests/unit/aegra/test_feedback.py new file mode 100644 index 00000000..eb34d99e --- /dev/null +++ b/tests/unit/aegra/test_feedback.py @@ -0,0 +1,246 @@ +"""Unit tests for Langfuse feedback recording and HTTP handler.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from pydantic import ValidationError +from starlette.requests import Request +from starlette.testclient import TestClient + +from deep_agent.aegra.feedback import feedback_handler, record_feedback +from deep_agent.aegra.http_app import app + + +class TestRecordFeedback: + @pytest.mark.asyncio + async def test_records_score_when_langfuse_configured(self): + mock_client = MagicMock() + payload = { + "trace_id": "abcd1234" * 4, + "name": "user-rating", + "value": 1.0, + "kwargs": {"comment": "great"}, + } + + with patch( + "deep_agent.aegra.feedback.get_langfuse_client", + return_value=mock_client, + ): + result = await record_feedback(payload) + + assert result.status == "success" + mock_client.create_score.assert_called_once_with( + trace_id=payload["trace_id"], + name="user-rating", + value=1.0, + data_type="BOOLEAN", + comment="great", + ) + + @pytest.mark.asyncio + async def test_graceful_degradation_when_langfuse_unconfigured(self): + payload = { + "trace_id": "abcd1234" * 4, + "name": "thumbs-up", + "value": 1.0, + } + + with patch( + "deep_agent.aegra.feedback.get_langfuse_client", + return_value=None, + ): + result = await record_feedback(payload) + + assert result.status == "success" + + @pytest.mark.asyncio + async def test_validation_error_on_missing_fields(self): + with pytest.raises(ValidationError): + await record_feedback({}) + + @pytest.mark.asyncio + async def test_gracefully_handles_score_failure(self): + mock_client = MagicMock() + mock_client.create_score.side_effect = RuntimeError("network") + + payload = { + "trace_id": "abcd1234" * 4, + "name": "user-rating", + "value": 0.5, + } + + with patch( + "deep_agent.aegra.feedback.get_langfuse_client", + return_value=mock_client, + ): + result = await record_feedback(payload) + + assert result.status == "success" + + @pytest.mark.asyncio + async def test_persists_postgres_when_thread_and_message_present(self): + payload = { + "trace_id": "a" * 32, + "name": "user-rating", + "value": 1.0, + "thread_id": "thread-1", + "message_id": "msg-1", + "user_id": "user-42", + } + mock_upsert = AsyncMock() + mock_repo = MagicMock() + mock_repo.upsert_feedback = mock_upsert + + with patch( + "deep_agent.aegra.feedback.get_langfuse_client", + return_value=None, + ): + with patch( + "deep_agent.aegra.feedback.FeedbackRepository", + return_value=mock_repo, + ): + result = await record_feedback(payload) + + assert result.status == "success" + mock_upsert.assert_awaited_once_with( + "thread-1", + "msg-1", + "user-42", + "up", + "a" * 32, + ) + + @pytest.mark.asyncio + async def test_skips_postgres_when_thread_or_message_missing(self): + payload = { + "trace_id": "a" * 32, + "name": "user-rating", + "value": 0.2, + } + + with patch( + "deep_agent.aegra.feedback.get_langfuse_client", + return_value=None, + ): + with patch( + "deep_agent.aegra.feedback.FeedbackRepository", + ) as mock_repo_cls: + result = await record_feedback(payload) + + assert result.status == "success" + mock_repo_cls.assert_not_called() + + +class TestFeedbackHandler: + @pytest.mark.asyncio + async def test_validation_error_response_shape(self): + scope = { + "type": "http", + "asgi": {"spec_version": "2.0", "version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/feedback", + "raw_path": b"/feedback", + "root_path": "", + "headers": [], + "client": ("127.0.0.1", 12345), + "server": ("127.0.0.1", 80), + } + + async def receive(): + return {"type": "http.request", "body": b"{}", "more_body": False} + + request = Request(scope, receive) + response = await feedback_handler(request) + assert response.status_code == 422 + + def test_post_feedback_via_test_client(self): + client = TestClient(app) + payload = { + "trace_id": "a" * 32, + "name": "user-rating", + "value": 1.0, + } + with patch( + "deep_agent.aegra.feedback.get_langfuse_client", + return_value=None, + ): + res = client.post("/feedback", json=payload) + assert res.status_code == 200 + assert res.json() == {"status": "success"} + + def test_get_thread_feedback(self): + client = TestClient(app) + thread_uuid = "00000000-0000-0000-0000-000000000001" + + mock_repo = MagicMock() + mock_repo.list_feedback = AsyncMock( + return_value=[{"message_id": "m1", "feedback": "up"}] + ) + with patch( + "deep_agent.aegra.feedback.FeedbackRepository", + return_value=mock_repo, + ): + res = client.get( + f"/feedback/{thread_uuid}", + params={"user_id": "u1"}, + ) + assert res.status_code == 200 + assert res.json() == {"feedback": [{"message_id": "m1", "feedback": "up"}]} + mock_repo.list_feedback.assert_awaited_once_with(thread_uuid, "u1") + + +class TestTokenUsageEndpoint: + def test_get_thread_token_usage_success(self) -> None: + from deep_agent.src.token_budget.service import ThreadTokenUsage + + client = TestClient(app) + thread_uuid = "00000000-0000-0000-0000-000000000001" + + with patch( + "deep_agent.src.token_budget.service.get_thread_token_usage", + new=AsyncMock( + return_value=ThreadTokenUsage( + thread_id=thread_uuid, + used=150, + input_tokens=100, + output_tokens=50, + ) + ), + ): + res = client.get(f"/threads/{thread_uuid}/token-usage") + + assert res.status_code == 200 + assert res.json() == { + "thread_id": thread_uuid, + "used": 150, + "input_tokens": 100, + "output_tokens": 50, + } + + def test_get_thread_token_usage_not_found(self) -> None: + from deep_agent.src.token_budget.service import TokenUsageNotFoundError + + client = TestClient(app) + thread_uuid = "00000000-0000-0000-0000-000000000001" + with patch( + "deep_agent.src.token_budget.service.get_thread_token_usage", + new=AsyncMock(side_effect=TokenUsageNotFoundError(thread_uuid)), + ): + res = client.get(f"/threads/{thread_uuid}/token-usage") + + assert res.status_code == 404 + + def test_get_thread_token_usage_unavailable(self) -> None: + from deep_agent.src.token_budget.service import TokenUsageUnavailableError + + client = TestClient(app) + thread_uuid = "00000000-0000-0000-0000-000000000001" + with patch( + "deep_agent.src.token_budget.service.get_thread_token_usage", + new=AsyncMock(side_effect=TokenUsageUnavailableError("down")), + ): + res = client.get(f"/threads/{thread_uuid}/token-usage") + + assert res.status_code == 503 diff --git a/tests/unit/aegra/test_graph.py b/tests/unit/aegra/test_graph.py new file mode 100644 index 00000000..65d78336 --- /dev/null +++ b/tests/unit/aegra/test_graph.py @@ -0,0 +1,416 @@ +"""Unit tests for aegra graph factory.""" + +import inspect +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +_runtime_mock = MagicMock() +if "langgraph_sdk.runtime" not in sys.modules: + sys.modules["langgraph_sdk.runtime"] = _runtime_mock + + +def _reset_graph_state() -> None: + from deep_agent.aegra import graph + + graph._graph_cache.clear() + graph._graph_cache_ts.clear() + + +class TestAgentFactory: + """Tests for the agent() graph factory function. + + The ``agent()`` function uses lazy imports inside its body, so + patches must target the actual module where each symbol lives. + """ + + @pytest.mark.asyncio + async def test_builds_agent_without_user(self): + mock_compiled = MagicMock() + mock_config = MagicMock() + mock_config.get_orchestrator_config.return_value = { + "name": "orchestrator", + "model": "gemini-2.5-flash", + "body": "test prompt", + "skill_paths": [], + "allowed_tools": [], + } + mock_config.resolve_tools.return_value = [] + mock_config.resolve_agent_middleware.return_value = MagicMock( + skills_enabled=True + ) + + mock_runtime = MagicMock() + mock_runtime.user = None + + _reset_graph_state() + + with ( + patch( + "deep_agent.src.agent.config.agent_config", + mock_config, + ), + patch( + "deep_agent.src.infrastructure.providers.register_profiles_from_config", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.providers.resolve_model_from_config", + return_value=MagicMock(), + ), + patch( + "deep_agent.aegra.mcp.get_mcp_tools", + new_callable=AsyncMock, + return_value=[], + ), + patch( + "deep_agent.src.infrastructure.subagents.load_subagents", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.backend.get_configured_backend", + return_value=MagicMock(), + ), + patch( + "deep_agent.src.infrastructure.async_tasks.build_async_middleware", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.middleware.build_middleware_list", + return_value=[], + ), + patch( + "deep_agent.src.infrastructure.middleware.resolve_memory_param", + return_value=None, + ), + patch("deep_agent.aegra.graph._ensure_startup", new_callable=AsyncMock), + patch("deepagents.create_deep_agent", return_value=mock_compiled), + ): + from deep_agent.aegra.graph import agent + + result = await agent(mock_runtime) + assert result is mock_compiled + + @pytest.mark.asyncio + async def test_builds_agent_with_sso_token(self): + mock_compiled = MagicMock() + mock_config = MagicMock() + mock_config.get_orchestrator_config.return_value = { + "name": "orchestrator", + "model": "gemini-2.5-flash", + "body": "test prompt", + "skill_paths": [], + "allowed_tools": [], + } + mock_config.resolve_tools.return_value = [] + mock_config.resolve_agent_middleware.return_value = MagicMock( + skills_enabled=True + ) + + mock_user = MagicMock() + mock_user.access_token = "test_access_token" + mock_user.refresh_token = "test_refresh_token" + mock_user.identity = None + + mock_runtime = MagicMock() + mock_runtime.user = mock_user + + _reset_graph_state() + + with ( + patch( + "deep_agent.src.agent.config.agent_config", + mock_config, + ), + patch( + "deep_agent.src.infrastructure.providers.register_profiles_from_config", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.providers.resolve_model_from_config", + return_value=MagicMock(), + ), + patch( + "deep_agent.aegra.mcp.get_mcp_tools", + new_callable=AsyncMock, + return_value=[], + ), + patch( + "deep_agent.aegra.mcp.refresh_access_token", + new_callable=AsyncMock, + return_value="refreshed_token", + ) as mock_refresh, + patch( + "deep_agent.src.infrastructure.subagents.load_subagents", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.backend.get_configured_backend", + return_value=MagicMock(), + ), + patch( + "deep_agent.src.infrastructure.async_tasks.build_async_middleware", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.middleware.build_middleware_list", + return_value=[], + ), + patch( + "deep_agent.src.infrastructure.middleware.resolve_memory_param", + return_value=None, + ), + patch("deep_agent.aegra.graph._ensure_startup", new_callable=AsyncMock), + patch("deepagents.create_deep_agent", return_value=mock_compiled), + ): + from deep_agent.aegra.graph import agent + + result = await agent(mock_runtime) + assert result is mock_compiled + mock_refresh.assert_awaited_once_with( + "test_access_token", "test_refresh_token" + ) + + @pytest.mark.asyncio + async def test_exposes_all_mcp_tools_when_mcps_declared_without_tool_list(self): + mock_compiled = MagicMock() + mock_config = MagicMock() + mock_config.get_orchestrator_config.return_value = { + "name": "orchestrator", + "model": "gemini-2.5-flash", + "body": "test prompt", + "skill_paths": [], + "allowed_tools": [], + "mcps": ["dataverse-mcp-prod1"], + } + mock_config.resolve_tools.return_value = [] + mock_config.resolve_agent_middleware.return_value = MagicMock( + skills_enabled=True + ) + + mock_tool = MagicMock() + mock_tool.name = "identify_dataproducts" + + mock_runtime = MagicMock() + mock_runtime.user = None + + _reset_graph_state() + + with ( + patch( + "deep_agent.src.agent.config.agent_config", + mock_config, + ), + patch( + "deep_agent.src.infrastructure.providers.register_profiles_from_config", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.providers.resolve_model_from_config", + return_value=MagicMock(), + ), + patch( + "deep_agent.aegra.mcp.get_mcp_tools", + new_callable=AsyncMock, + return_value=[mock_tool], + ) as mock_get_mcp, + patch( + "deep_agent.src.infrastructure.subagents.load_subagents", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.backend.get_configured_backend", + return_value=MagicMock(), + ), + patch( + "deep_agent.src.infrastructure.async_tasks.build_async_middleware", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.middleware.build_middleware_list", + return_value=[], + ), + patch( + "deep_agent.src.infrastructure.middleware.resolve_memory_param", + return_value=None, + ), + patch("deep_agent.aegra.graph._ensure_startup", new_callable=AsyncMock), + patch( + "deepagents.create_deep_agent", return_value=mock_compiled + ) as mock_create, + ): + from deep_agent.aegra.graph import agent + + result = await agent(mock_runtime) + + assert result is mock_compiled + assert mock_create.call_args.kwargs["tools"] == [mock_tool] + mock_get_mcp.assert_awaited_once_with( + sso_token=None, server_names=["dataverse-mcp-prod1"], user_id=None + ) + + @pytest.mark.asyncio + async def test_hitl_passes_interrupt_on_when_enabled(self): + """create_deep_agent must receive a non-empty interrupt_on dict when HITL is enabled.""" + from deep_agent.src.agent.config.middleware import HumanApprovalConfig + + mock_compiled = MagicMock() + mock_config = MagicMock() + mock_config.get_orchestrator_config.return_value = { + "name": "orchestrator", + "model": "gemini-2.5-flash", + "body": "test prompt", + "skill_paths": [], + "allowed_tools": [], + } + mock_config.resolve_tools.return_value = [] + + hitl_config = HumanApprovalConfig(enabled=True, mode="all", exclude=[]) + mock_mw = MagicMock(skills_enabled=True) + mock_mw.human_approval = hitl_config + mock_config.resolve_agent_middleware.return_value = mock_mw + + mock_runtime = MagicMock() + mock_runtime.user = None + + # Give the mock a real signature that includes interrupt_on so that the + # inspect.signature() check inside agent() sees the parameter. + def _stub(*, interrupt_on=None, **kw): ... + + mock_create = MagicMock(return_value=mock_compiled) + mock_create.__signature__ = inspect.signature(_stub) + + _reset_graph_state() + + with ( + patch("deep_agent.src.agent.config.agent_config", mock_config), + patch( + "deep_agent.src.infrastructure.providers.register_profiles_from_config", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.providers.resolve_model_from_config", + return_value=MagicMock(), + ), + patch( + "deep_agent.aegra.mcp.get_mcp_tools", + new_callable=AsyncMock, + return_value=[], + ), + patch( + "deep_agent.src.infrastructure.subagents.load_subagents", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.backend.get_configured_backend", + return_value=MagicMock(), + ), + patch( + "deep_agent.src.infrastructure.async_tasks.build_async_middleware", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.middleware.build_middleware_list", + return_value=[], + ), + patch( + "deep_agent.src.infrastructure.middleware.resolve_memory_param", + return_value=None, + ), + patch("deep_agent.aegra.graph._ensure_startup", new_callable=AsyncMock), + patch("deepagents.create_deep_agent", new=mock_create), + ): + from deep_agent.aegra.graph import agent + + result = await agent(mock_runtime) + + assert result is mock_compiled + call_kwargs = mock_create.call_args.kwargs + assert "interrupt_on" in call_kwargs, ( + "interrupt_on was not passed to create_deep_agent" + ) + assert isinstance(call_kwargs["interrupt_on"], dict) + assert len(call_kwargs["interrupt_on"]) > 0, ( + "interrupt_on dict must not be empty" + ) + assert all(v is True for v in call_kwargs["interrupt_on"].values()) + + @pytest.mark.asyncio + async def test_hitl_omits_interrupt_on_when_disabled(self): + """create_deep_agent must NOT receive interrupt_on when HITL is disabled.""" + from deep_agent.src.agent.config.middleware import HumanApprovalConfig + + mock_compiled = MagicMock() + mock_config = MagicMock() + mock_config.get_orchestrator_config.return_value = { + "name": "orchestrator", + "model": "gemini-2.5-flash", + "body": "test prompt", + "skill_paths": [], + "allowed_tools": [], + } + mock_config.resolve_tools.return_value = [] + + hitl_config = HumanApprovalConfig(enabled=False) + mock_mw = MagicMock(skills_enabled=True) + mock_mw.human_approval = hitl_config + mock_config.resolve_agent_middleware.return_value = mock_mw + + mock_runtime = MagicMock() + mock_runtime.user = None + + def _stub(*, interrupt_on=None, **kw): ... + + mock_create = MagicMock(return_value=mock_compiled) + mock_create.__signature__ = inspect.signature(_stub) + + _reset_graph_state() + + with ( + patch("deep_agent.src.agent.config.agent_config", mock_config), + patch( + "deep_agent.src.infrastructure.providers.register_profiles_from_config", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.providers.resolve_model_from_config", + return_value=MagicMock(), + ), + patch( + "deep_agent.aegra.mcp.get_mcp_tools", + new_callable=AsyncMock, + return_value=[], + ), + patch( + "deep_agent.src.infrastructure.subagents.load_subagents", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.backend.get_configured_backend", + return_value=MagicMock(), + ), + patch( + "deep_agent.src.infrastructure.async_tasks.build_async_middleware", + return_value=None, + ), + patch( + "deep_agent.src.infrastructure.middleware.build_middleware_list", + return_value=[], + ), + patch( + "deep_agent.src.infrastructure.middleware.resolve_memory_param", + return_value=None, + ), + patch("deep_agent.aegra.graph._ensure_startup", new_callable=AsyncMock), + patch("deepagents.create_deep_agent", new=mock_create), + ): + from deep_agent.aegra.graph import agent + + result = await agent(mock_runtime) + + assert result is mock_compiled + call_kwargs = mock_create.call_args.kwargs + assert "interrupt_on" not in call_kwargs, ( + "interrupt_on must not be passed when HITL is disabled" + ) diff --git a/tests/unit/aegra/test_health.py b/tests/unit/aegra/test_health.py new file mode 100644 index 00000000..835fdac2 --- /dev/null +++ b/tests/unit/aegra/test_health.py @@ -0,0 +1,244 @@ +"""Unit tests for health check endpoint.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from deep_agent.aegra.health import ( + check_cache, + check_config, + check_database, + check_redis, + get_health_status, + health_response, +) + + +def _patch_all_checks(**overrides): + """Return a context-manager stack that mocks every health sub-check. + + Defaults to ``{"status": "ok"}`` for each check. Pass keyword + overrides keyed by check name to customise individual results. + """ + defaults = { + "database": {"status": "ok"}, + "redis": {"status": "ok"}, + "config": {"status": "ok"}, + "cache": {"status": "ok"}, + "mcp_servers": {"status": "ok", "servers": {}, "healthy": 0, "total": 0}, + "llm_provider": {"status": "ok", "provider": "vllm"}, + } + defaults.update(overrides) + + from contextlib import ExitStack + + stack = ExitStack() + stack.enter_context( + patch( + "deep_agent.aegra.health.check_database", + new_callable=AsyncMock, + return_value=defaults["database"], + ) + ) + stack.enter_context( + patch( + "deep_agent.aegra.health.check_redis", + new_callable=AsyncMock, + return_value=defaults["redis"], + ) + ) + stack.enter_context( + patch( + "deep_agent.aegra.health.check_config", + return_value=defaults["config"], + ) + ) + stack.enter_context( + patch( + "deep_agent.aegra.health.check_cache", + return_value=defaults["cache"], + ) + ) + stack.enter_context( + patch( + "deep_agent.aegra.mcp_health.check_mcp_servers", + new_callable=AsyncMock, + return_value=defaults["mcp_servers"], + ) + ) + stack.enter_context( + patch( + "deep_agent.aegra.mcp_health.check_llm_provider", + new_callable=AsyncMock, + return_value=defaults["llm_provider"], + ) + ) + return stack + + +class TestCheckConfig: + def test_valid_config(self): + mock_settings = MagicMock() + mock_settings.database_uri = "postgresql://test" + mock_settings.AGENT_PORT = 5002 + with patch("deep_agent.src.settings.settings", mock_settings): + result = check_config() + assert result["status"] == "ok" + + def test_missing_database(self): + mock_settings = MagicMock() + mock_settings.database_uri = "" + mock_settings.AGENT_PORT = 5002 + with patch("deep_agent.src.settings.settings", mock_settings): + result = check_config() + assert result["status"] == "warning" + + +class TestCheckDatabase: + async def test_no_database_uri(self): + mock_settings = MagicMock() + mock_settings.database_uri = "" + with patch("deep_agent.src.settings.settings", mock_settings): + result = await check_database() + assert result["status"] == "skipped" + + async def test_database_error(self): + mock_settings = MagicMock() + mock_settings.database_uri = "postgresql://bad" + with ( + patch("deep_agent.src.settings.settings", mock_settings), + patch( + "psycopg.AsyncConnection.connect", + side_effect=Exception("connection refused"), + ), + ): + result = await check_database() + assert result["status"] == "error" + + +class TestCheckRedis: + async def test_no_redis(self): + with patch( + "deep_agent.aegra.redis.get_redis_client", + return_value=None, + ): + result = await check_redis() + assert result["status"] == "skipped" + + async def test_redis_ok(self): + mock_client = AsyncMock() + mock_client.ping = AsyncMock(return_value=True) + with patch( + "deep_agent.aegra.redis.get_redis_client", + return_value=mock_client, + ): + result = await check_redis() + assert result["status"] == "ok" + assert "latency_ms" in result + + +class TestCheckCache: + def test_returns_stats(self): + with patch( + "deep_agent.src.cache.metrics.get_stats", + return_value={"hits": 10, "misses": 2}, + ): + result = check_cache() + assert result["status"] == "ok" + + +class TestGetHealthStatus: + async def test_healthy(self): + with _patch_all_checks(): + result = await get_health_status() + assert result["status"] == "healthy" + assert "uptime_seconds" in result + assert "checks" in result + assert "mcp_servers" in result["checks"] + assert "llm_provider" in result["checks"] + + async def test_unhealthy_on_db_error(self): + with _patch_all_checks(database={"status": "error", "error": "down"}): + result = await get_health_status() + assert result["status"] == "unhealthy" + + async def test_degraded_when_mcp_subset_down(self): + """MCP servers partially down → degraded, NOT unhealthy.""" + mcp = { + "status": "warning", + "servers": { + "a": {"status": "healthy"}, + "b": {"status": "unreachable"}, + }, + "healthy": 1, + "total": 2, + } + with _patch_all_checks(mcp_servers=mcp): + result = await get_health_status() + assert result["status"] == "degraded" + + async def test_degraded_when_all_mcp_down(self): + """All MCP servers down → degraded (pod stays in rotation).""" + mcp = { + "status": "warning", + "servers": { + "a": {"status": "unreachable"}, + "b": {"status": "timeout"}, + }, + "healthy": 0, + "total": 2, + } + with _patch_all_checks(mcp_servers=mcp): + result = await get_health_status() + assert result["status"] == "degraded" + + async def test_degraded_when_llm_down(self): + """LLM provider down → degraded, NOT unhealthy.""" + llm = {"status": "warning", "provider": "vllm", "error": "timeout"} + with _patch_all_checks(llm_provider=llm): + result = await get_health_status() + assert result["status"] == "degraded" + + async def test_db_error_overrides_mcp_warning(self): + """DB error + MCP warning → unhealthy (critical wins).""" + with _patch_all_checks( + database={"status": "error", "error": "down"}, + mcp_servers={"status": "warning", "servers": {}, "healthy": 0, "total": 1}, + ): + result = await get_health_status() + assert result["status"] == "unhealthy" + + async def test_redis_error_is_degraded_not_unhealthy(self): + """Redis is non-critical so an error produces degraded.""" + with _patch_all_checks(redis={"status": "error", "error": "refused"}): + result = await get_health_status() + assert result["status"] == "degraded" + + +class TestHealthResponse: + async def test_200_when_healthy(self): + with patch( + "deep_agent.aegra.health.get_health_status", + new_callable=AsyncMock, + return_value={"status": "healthy"}, + ): + code, body = await health_response() + assert code == 200 + + async def test_200_when_degraded(self): + with patch( + "deep_agent.aegra.health.get_health_status", + new_callable=AsyncMock, + return_value={"status": "degraded"}, + ): + code, body = await health_response() + assert code == 200 + + async def test_503_when_unhealthy(self): + with patch( + "deep_agent.aegra.health.get_health_status", + new_callable=AsyncMock, + return_value={"status": "unhealthy"}, + ): + code, body = await health_response() + assert code == 503 diff --git a/tests/unit/aegra/test_mcp_auth.py b/tests/unit/aegra/test_mcp_auth.py new file mode 100644 index 00000000..fedd2663 --- /dev/null +++ b/tests/unit/aegra/test_mcp_auth.py @@ -0,0 +1,243 @@ +"""Unit tests for MCP config validation and credential resolver.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from unittest.mock import AsyncMock, patch + +import pytest + +from deep_agent.aegra.mcp import set_mcp_auth_context +from deep_agent.aegra.mcp_auth import McpCredentialResolver, NeedsAuthorization +from deep_agent.aegra.mcp_token_store import McpOAuthToken +from deep_agent.src.agent.config.loader import AgentConfig + + +class TestMcpConfigValidation: + def setup_method(self): + AgentConfig._instance = None + + @staticmethod + def _write_minimal_config_dir(tmp_path): + (tmp_path / "PROMPT.md").write_text( + """--- +name: test-orchestrator +model: gemini-2.5-flash +--- +Test prompt. +""" + ) + + def test_defaults_auth_mode_to_sso(self, tmp_path): + self._write_minimal_config_dir(tmp_path) + mcp_json = tmp_path / "mcp.json" + mcp_json.write_text( + '{"mcpServers": {"sso-mcp": {"url": "http://localhost/mcp", "enabled": true}}}' + ) + cfg = AgentConfig(tmp_path) + servers = cfg.get_mcp_servers() + assert servers["sso-mcp"]["auth_mode"] == "sso" + + def test_loads_jsonc_line_comments(self, tmp_path): + self._write_minimal_config_dir(tmp_path) + mcp_json = tmp_path / "mcp.json" + mcp_json.write_text( + """ + { + "mcpServers": { + "template-mcp-server": { + "url": "http://host.containers.internal:5001/mcp", + // "url": "http://localhost:5001/mcp", + "enabled": true + } + } + } + """ + ) + servers = AgentConfig(tmp_path).get_mcp_servers() + assert ( + servers["template-mcp-server"]["url"] + == "http://host.containers.internal:5001/mcp" + ) + + def test_loads_jsonc_with_escaped_quotes(self, tmp_path): + self._write_minimal_config_dir(tmp_path) + mcp_json = tmp_path / "mcp.json" + mcp_json.write_text( + r""" + { + "mcpServers": { + "test-mcp": { + "url": "http://host/mcp?q=\"hello\"", + // comment with escaped quote: \" + "label": "backslash\\and-quote", + "enabled": true + } + } + } + """ + ) + servers = AgentConfig(tmp_path).get_mcp_servers() + assert servers["test-mcp"]["url"] == 'http://host/mcp?q="hello"' + assert servers["test-mcp"]["label"] == "backslash\\and-quote" + + def test_logs_error_for_oauth_without_client_id(self, tmp_path, caplog): + self._write_minimal_config_dir(tmp_path) + mcp_json = tmp_path / "mcp.json" + mcp_json.write_text( + """ + { + "mcpServers": { + "oauth-mcp": { + "url": "http://localhost/mcp", + "enabled": true, + "auth_mode": "oauth", + "oauth": { + "authorization_endpoint": "https://as.example.com/authorize", + "token_endpoint": "https://as.example.com/token" + } + } + } + } + """ + ) + with caplog.at_level("ERROR"): + AgentConfig(tmp_path).get_mcp_servers() + assert any("client_id is required" in r.message for r in caplog.records) + + def test_logs_error_for_dcr_without_registration_endpoint(self, tmp_path, caplog): + self._write_minimal_config_dir(tmp_path) + mcp_json = tmp_path / "mcp.json" + mcp_json.write_text( + """ + { + "mcpServers": { + "dcr-mcp": { + "url": "http://localhost/mcp", + "enabled": true, + "auth_mode": "dcr", + "oauth": { + "authorization_endpoint": "https://as.example.com/authorize", + "token_endpoint": "https://as.example.com/token" + } + } + } + } + """ + ) + with caplog.at_level("ERROR"): + AgentConfig(tmp_path).get_mcp_servers() + assert any( + "registration_endpoint is required" in r.message for r in caplog.records + ) + + +@pytest.mark.asyncio +class TestMcpCredentialResolver: + async def test_sso_returns_refreshed_token(self): + store = AsyncMock() + resolver = McpCredentialResolver(token_store=store) + set_mcp_auth_context("access-token", "refresh-token") + + with patch( + "deep_agent.aegra.mcp_auth.refresh_access_token", + new=AsyncMock(return_value="fresh-token"), + ) as refresh: + token = await resolver.resolve( + "user-1", + "sso-mcp", + {"auth_mode": "sso"}, + ) + + assert token == "fresh-token" + refresh.assert_awaited_once_with("access-token", "refresh-token") + store.get_token.assert_not_called() + + async def test_oauth_raises_when_no_stored_token(self): + store = AsyncMock() + store.get_token = AsyncMock(return_value=None) + resolver = McpCredentialResolver(token_store=store) + + with pytest.raises(NeedsAuthorization) as exc: + await resolver.resolve( + "user-1", + "oauth-mcp", + { + "auth_mode": "oauth", + "oauth": {"token_endpoint": "https://as.example.com/token"}, + }, + ) + + assert exc.value.mcp_name == "oauth-mcp" + assert exc.value.connect_url.endswith("/mcp/oauth-mcp/connect") + + async def test_oauth_returns_valid_stored_token(self): + store = AsyncMock() + store.get_token = AsyncMock( + return_value=McpOAuthToken( + user_id="user-1", + mcp_name="oauth-mcp", + access_token="stored-access", + expires_at=datetime.now(UTC) + timedelta(hours=1), + ) + ) + resolver = McpCredentialResolver(token_store=store) + + token = await resolver.resolve( + "user-1", + "oauth-mcp", + {"auth_mode": "oauth", "oauth": {}}, + ) + assert token == "stored-access" + + async def test_oauth_refreshes_expired_token(self): + store = AsyncMock() + store.get_token = AsyncMock( + return_value=McpOAuthToken( + user_id="user-1", + mcp_name="oauth-mcp", + access_token="expired-access", + refresh_token="refresh-me", + expires_at=datetime.now(UTC) - timedelta(minutes=5), + ) + ) + store.upsert_token = AsyncMock() + resolver = McpCredentialResolver(token_store=store) + + with patch.object( + resolver, + "_refresh_mcp_token", + new=AsyncMock(return_value="new-access"), + ) as refresh: + token = await resolver.resolve( + "user-1", + "oauth-mcp", + { + "auth_mode": "oauth", + "oauth": { + "token_endpoint": "https://as.example.com/token", + "client_id": "cid", + }, + }, + ) + + assert token == "new-access" + refresh.assert_awaited_once() + + async def test_resolver_caches_resolved_oauth_token(self): + store = AsyncMock() + store.get_token = AsyncMock( + return_value=McpOAuthToken( + user_id="user-1", + mcp_name="oauth-mcp", + access_token="stored-access", + expires_at=datetime.now(UTC) + timedelta(hours=1), + ) + ) + resolver = McpCredentialResolver(token_store=store) + + cfg = {"auth_mode": "oauth", "oauth": {}} + await resolver.resolve("user-1", "oauth-mcp", cfg) + await resolver.resolve("user-1", "oauth-mcp", cfg) + + store.get_token.assert_awaited_once() diff --git a/tests/unit/aegra/test_mcp_crypto.py b/tests/unit/aegra/test_mcp_crypto.py new file mode 100644 index 00000000..c725cd62 --- /dev/null +++ b/tests/unit/aegra/test_mcp_crypto.py @@ -0,0 +1,85 @@ +"""Unit tests for MCP OAuth token encryption.""" + +from __future__ import annotations + +import os + +import pytest +from cryptography.fernet import Fernet, InvalidToken + +from deep_agent.aegra.mcp_crypto import ( + decrypt_secret, + encrypt_secret, + reset_mcp_crypto_cache, +) + + +@pytest.fixture(autouse=True) +def _clear_crypto_cache(): + reset_mcp_crypto_cache() + yield + reset_mcp_crypto_cache() + + +@pytest.fixture +def fernet_keys(): + primary = Fernet.generate_key().decode() + previous = Fernet.generate_key().decode() + return primary, previous + + +class TestMcpCrypto: + def test_encrypt_decrypt_round_trip(self, fernet_keys, monkeypatch): + primary, _ = fernet_keys + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY", primary) + ciphertext = encrypt_secret("secret-token") + assert ciphertext is not None + assert decrypt_secret(ciphertext) == "secret-token" + + def test_none_passthrough(self, fernet_keys, monkeypatch): + primary, _ = fernet_keys + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY", primary) + assert encrypt_secret(None) is None + assert decrypt_secret(None) is None + + def test_decrypt_with_previous_key(self, fernet_keys, monkeypatch): + primary, previous = fernet_keys + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY", previous) + ciphertext = encrypt_secret("rotated-secret") + + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY", primary) + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY_PREVIOUS", previous) + reset_mcp_crypto_cache() + + assert decrypt_secret(ciphertext) == "rotated-secret" + + def test_encrypt_uses_primary_only(self, fernet_keys, monkeypatch): + primary, previous = fernet_keys + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY", primary) + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY_PREVIOUS", previous) + ciphertext = encrypt_secret("new-secret") + + with pytest.raises(InvalidToken): + Fernet(previous.encode()).decrypt(ciphertext.encode()) + assert ( + Fernet(primary.encode()).decrypt(ciphertext.encode()).decode() + == "new-secret" + ) + + def test_missing_primary_key_raises(self, monkeypatch): + monkeypatch.delenv("MCP_TOKEN_ENCRYPTION_KEY", raising=False) + with pytest.raises(RuntimeError, match="MCP_TOKEN_ENCRYPTION_KEY"): + encrypt_secret("x") + + def test_wrong_keys_raise(self, fernet_keys, monkeypatch): + primary, previous = fernet_keys + other = Fernet.generate_key().decode() + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY", other) + ciphertext = encrypt_secret("lost-secret") + + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY", primary) + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY_PREVIOUS", previous) + reset_mcp_crypto_cache() + + with pytest.raises(RuntimeError, match="decryption failed"): + decrypt_secret(ciphertext) diff --git a/tests/unit/aegra/test_mcp_health.py b/tests/unit/aegra/test_mcp_health.py new file mode 100644 index 00000000..3b13eef8 --- /dev/null +++ b/tests/unit/aegra/test_mcp_health.py @@ -0,0 +1,360 @@ +"""Unit tests for MCP and LLM provider health checks.""" + +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from deep_agent.aegra import mcp_health +from deep_agent.aegra.mcp_health import ( + _HEALTH_CACHE_TTL, + _ping_mcp_server, + check_llm_provider, + check_mcp_servers, + invalidate_health_cache, +) + + +@pytest.fixture(autouse=True) +def _clear_cache(): + """Reset module-level caches and OTEL state between tests.""" + invalidate_health_cache() + mcp_health._gauge_initialized = False + mcp_health._mcp_health_gauge = None + mcp_health._llm_health_gauge = None + yield + invalidate_health_cache() + + +def _mock_servers(servers: dict): + """Patch agent_config.get_mcp_servers to return *servers*.""" + mock_config = MagicMock() + mock_config.get_mcp_servers.return_value = servers + return patch("deep_agent.src.agent.config.agent_config", mock_config) + + +TWO_SERVERS = { + "server-a": { + "url": "http://a:5001/mcp", + "transport": "streamable_http", + "enabled": True, + "auth": False, + "ssl_verify": False, + "timeout": 10, + }, + "server-b": { + "url": "http://b:5002/mcp", + "transport": "streamable_http", + "enabled": True, + "auth": False, + "ssl_verify": False, + "timeout": 10, + }, +} + + +# ── _ping_mcp_server ───────────────────────────────────────────── + + +class TestPingMcpServer: + async def test_healthy(self): + mock_resp = MagicMock(status_code=200) + with patch("deep_agent.aegra.mcp_health.httpx.AsyncClient") as mock_cls: + mock_client = AsyncMock() + mock_client.get.return_value = mock_resp + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_cls.return_value = mock_client + + result = await _ping_mcp_server("test", "http://x/mcp", 5.0, False) + + assert result["status"] == "healthy" + assert "latency_ms" in result + assert result["http_status"] == 200 + + async def test_server_error(self): + mock_resp = MagicMock(status_code=502) + with patch("deep_agent.aegra.mcp_health.httpx.AsyncClient") as mock_cls: + mock_client = AsyncMock() + mock_client.get.return_value = mock_resp + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_cls.return_value = mock_client + + result = await _ping_mcp_server("test", "http://x/mcp", 5.0, False) + + assert result["status"] == "unreachable" + + async def test_timeout(self): + with patch("deep_agent.aegra.mcp_health.httpx.AsyncClient") as mock_cls: + mock_client = AsyncMock() + mock_client.get.side_effect = httpx.TimeoutException("timeout") + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_cls.return_value = mock_client + + result = await _ping_mcp_server("test", "http://x/mcp", 5.0, False) + + assert result["status"] == "timeout" + + async def test_connection_refused(self): + with patch("deep_agent.aegra.mcp_health.httpx.AsyncClient") as mock_cls: + mock_client = AsyncMock() + mock_client.get.side_effect = httpx.ConnectError("connection refused") + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_cls.return_value = mock_client + + result = await _ping_mcp_server("test", "http://x/mcp", 5.0, False) + + assert result["status"] == "unreachable" + assert "error" in result + + async def test_4xx_counts_as_healthy(self): + mock_resp = MagicMock(status_code=405) + with patch("deep_agent.aegra.mcp_health.httpx.AsyncClient") as mock_cls: + mock_client = AsyncMock() + mock_client.get.return_value = mock_resp + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_cls.return_value = mock_client + + result = await _ping_mcp_server("test", "http://x/mcp", 5.0, False) + + assert result["status"] == "healthy" + + +# ── check_mcp_servers ──────────────────────────────────────────── + + +class TestCheckMcpServers: + async def test_all_healthy(self): + healthy = {"status": "healthy", "latency_ms": 1.0, "http_status": 200} + with ( + _mock_servers(TWO_SERVERS), + patch( + "deep_agent.aegra.mcp_health._ping_mcp_server", + new_callable=AsyncMock, + return_value=healthy, + ), + patch("deep_agent.aegra.mcp._get_mcp_breaker") as mock_breaker, + ): + mock_breaker.return_value.is_open = False + result = await check_mcp_servers() + + assert result["status"] == "ok" + assert result["healthy"] == 2 + assert result["total"] == 2 + assert "server-a" in result["servers"] + assert "server-b" in result["servers"] + + async def test_one_of_two_down(self): + """Partial failure → status is 'warning', not 'error'.""" + + async def _ping_side_effect(name, url, timeout, ssl_verify): + if name == "server-a": + return {"status": "healthy", "latency_ms": 1.0, "http_status": 200} + return {"status": "unreachable", "error": "connection refused"} + + with ( + _mock_servers(TWO_SERVERS), + patch( + "deep_agent.aegra.mcp_health._ping_mcp_server", + side_effect=_ping_side_effect, + ), + patch("deep_agent.aegra.mcp._get_mcp_breaker") as mock_breaker, + ): + mock_breaker.return_value.is_open = False + result = await check_mcp_servers() + + assert result["status"] == "warning" + assert result["healthy"] == 1 + assert result["total"] == 2 + assert result["servers"]["server-a"]["status"] == "healthy" + assert result["servers"]["server-b"]["status"] == "unreachable" + + async def test_all_down_still_warning_not_error(self): + """All MCP servers down → 'warning' so agent reports degraded, not unhealthy.""" + down = {"status": "unreachable", "error": "connection refused"} + with ( + _mock_servers(TWO_SERVERS), + patch( + "deep_agent.aegra.mcp_health._ping_mcp_server", + new_callable=AsyncMock, + return_value=down, + ), + patch("deep_agent.aegra.mcp._get_mcp_breaker") as mock_breaker, + ): + mock_breaker.return_value.is_open = False + result = await check_mcp_servers() + + assert result["status"] == "warning" + assert result["healthy"] == 0 + + async def test_breaker_open(self): + with ( + _mock_servers(TWO_SERVERS), + patch("deep_agent.aegra.mcp._get_mcp_breaker") as mock_breaker, + ): + mock_breaker.return_value.is_open = True + result = await check_mcp_servers() + + assert result["status"] == "warning" + for srv in result["servers"].values(): + assert srv["status"] == "breaker-open" + + async def test_no_servers_enabled(self): + disabled = { + "x": {"url": "http://x/mcp", "enabled": False}, + } + with _mock_servers(disabled): + result = await check_mcp_servers() + + assert result["status"] == "skipped" + + async def test_cache_hit(self): + healthy = {"status": "healthy", "latency_ms": 1.0, "http_status": 200} + with ( + _mock_servers(TWO_SERVERS), + patch( + "deep_agent.aegra.mcp_health._ping_mcp_server", + new_callable=AsyncMock, + return_value=healthy, + ) as mock_ping, + patch("deep_agent.aegra.mcp._get_mcp_breaker") as mock_breaker, + ): + mock_breaker.return_value.is_open = False + first = await check_mcp_servers() + second = await check_mcp_servers() + + assert first is second + assert mock_ping.await_count == 2 # only the first round (2 servers) + + +# ── check_llm_provider ────────────────────────────────────────── + + +class TestCheckLlmProvider: + async def test_vllm_healthy(self): + mock_settings = MagicMock() + mock_settings.VLLM_BASE_URL = "http://vllm:8000/v1" + mock_settings.VLLM_API_KEY = "EMPTY" + + mock_resp = MagicMock(status_code=200) + with ( + patch("deep_agent.aegra.mcp_health.httpx.AsyncClient") as mock_cls, + patch("deep_agent.aegra.mcp_health.settings", mock_settings), + ): + mock_client = AsyncMock() + mock_client.get.return_value = mock_resp + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_cls.return_value = mock_client + + result = await check_llm_provider() + + assert result["status"] == "ok" + assert result["provider"] == "vllm" + + async def test_vllm_unreachable(self): + mock_settings = MagicMock() + mock_settings.VLLM_BASE_URL = "http://vllm:8000/v1" + mock_settings.VLLM_API_KEY = "EMPTY" + + with ( + patch("deep_agent.aegra.mcp_health.httpx.AsyncClient") as mock_cls, + patch("deep_agent.aegra.mcp_health.settings", mock_settings), + ): + mock_client = AsyncMock() + mock_client.get.side_effect = httpx.ConnectError("refused") + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + mock_cls.return_value = mock_client + + result = await check_llm_provider() + + assert result["status"] == "warning" + assert result["provider"] == "vllm" + + async def test_vertex_ai_ok(self): + mock_settings = MagicMock() + mock_settings.VLLM_BASE_URL = "" + + with ( + patch("deep_agent.aegra.mcp_health.settings", mock_settings), + patch( + "deep_agent.aegra.mcp_health.get_service_account_credentials", + return_value=(MagicMock(), "my-project"), + ), + ): + result = await check_llm_provider() + + assert result["status"] == "ok" + assert result["provider"] == "vertex_ai" + assert result["project"] == "my-project" + + async def test_vertex_ai_no_creds(self): + mock_settings = MagicMock() + mock_settings.VLLM_BASE_URL = "" + + with ( + patch("deep_agent.aegra.mcp_health.settings", mock_settings), + patch( + "deep_agent.aegra.mcp_health.get_service_account_credentials", + side_effect=Exception("no credentials"), + ), + ): + result = await check_llm_provider() + + assert result["status"] == "warning" + assert result["provider"] == "vertex_ai" + + async def test_cache_hit(self): + mock_settings = MagicMock() + mock_settings.VLLM_BASE_URL = "" + + with ( + patch("deep_agent.aegra.mcp_health.settings", mock_settings), + patch( + "deep_agent.aegra.mcp_health.get_service_account_credentials", + return_value=(MagicMock(), "proj"), + ) as mock_creds, + ): + first = await check_llm_provider() + second = await check_llm_provider() + + assert first is second + assert mock_creds.call_count == 1 + + +# ── OTEL gauge emission ───────────────────────────────────────── + + +class TestOtelGauges: + def test_gauge_noop_when_otel_disabled(self): + mock_settings = MagicMock() + mock_settings.ENABLE_OTEL_METRICS = False + mock_settings.OTEL_EXPORTER_OTLP_ENDPOINT = "" + + with patch("deep_agent.src.settings.settings", mock_settings): + mcp_health._ensure_gauges() + + assert mcp_health._mcp_health_gauge is None + + def test_gauge_created_when_otel_enabled(self): + mock_settings = MagicMock() + mock_settings.ENABLE_OTEL_METRICS = True + mock_settings.OTEL_EXPORTER_OTLP_ENDPOINT = "http://otel:4317" + + mock_gauge = MagicMock() + mock_meter = MagicMock() + mock_meter.create_gauge.return_value = mock_gauge + + with ( + patch("deep_agent.aegra.mcp_health.settings", mock_settings), + patch("opentelemetry.metrics.get_meter", return_value=mock_meter), + ): + mcp_health._ensure_gauges() + + assert mock_meter.create_gauge.call_count == 2 diff --git a/tests/unit/aegra/test_mcp_token_refresh_lock.py b/tests/unit/aegra/test_mcp_token_refresh_lock.py new file mode 100644 index 00000000..82a91dda --- /dev/null +++ b/tests/unit/aegra/test_mcp_token_refresh_lock.py @@ -0,0 +1,134 @@ +"""Unit tests for MCP token refresh locking.""" + +from __future__ import annotations + +from contextlib import asynccontextmanager +from datetime import UTC, datetime, timedelta +from unittest.mock import AsyncMock, patch + +import pytest + +from deep_agent.aegra.mcp_auth import McpCredentialResolver +from deep_agent.aegra.mcp_token_store import McpOAuthToken + + +@asynccontextmanager +async def _held_lock(*_args, **_kwargs): + yield "held" + + +@asynccontextmanager +async def _timeout_lock(*_args, **_kwargs): + yield "timeout" + + +def _expired_token() -> McpOAuthToken: + return McpOAuthToken( + user_id="user-1", + mcp_name="oauth-mcp", + access_token="expired-access", + refresh_token="refresh-me", + expires_at=datetime.now(UTC) - timedelta(minutes=5), + ) + + +def _fresh_token() -> McpOAuthToken: + return McpOAuthToken( + user_id="user-1", + mcp_name="oauth-mcp", + access_token="fresh-access", + refresh_token="refresh-me", + expires_at=datetime.now(UTC) + timedelta(hours=1), + ) + + +@pytest.mark.asyncio +class TestMcpTokenRefreshLock: + async def test_skips_refresh_when_peer_refreshed_under_lock(self): + store = AsyncMock() + store.get_token = AsyncMock(side_effect=[_expired_token(), _fresh_token()]) + resolver = McpCredentialResolver(token_store=store) + + with ( + patch("deep_agent.aegra.mcp_auth.distributed_lock", _held_lock), + patch.object( + resolver, + "_refresh_mcp_token", + new=AsyncMock(return_value="should-not-run"), + ) as refresh, + ): + token = await resolver.resolve( + "user-1", + "oauth-mcp", + { + "auth_mode": "oauth", + "oauth": {"token_endpoint": "https://as.example.com/token"}, + }, + ) + + assert token == "fresh-access" + refresh.assert_not_called() + assert store.get_token.await_count == 2 + + async def test_waits_for_peer_refresh_on_lock_timeout(self): + store = AsyncMock() + store.get_token = AsyncMock( + side_effect=[ + _expired_token(), + _expired_token(), + _fresh_token(), + ] + ) + resolver = McpCredentialResolver(token_store=store) + + with ( + patch("deep_agent.aegra.mcp_auth.distributed_lock", _timeout_lock), + patch( + "deep_agent.aegra.mcp_auth.asyncio.sleep", + new=AsyncMock(), + ), + patch.object( + resolver, + "_refresh_mcp_token", + new=AsyncMock(return_value="should-not-run"), + ) as refresh, + ): + token = await resolver.resolve( + "user-1", + "oauth-mcp", + { + "auth_mode": "oauth", + "oauth": {"token_endpoint": "https://as.example.com/token"}, + }, + ) + + assert token == "fresh-access" + refresh.assert_not_called() + + async def test_refreshes_once_when_lock_held(self): + store = AsyncMock() + store.get_token = AsyncMock(side_effect=[_expired_token(), _expired_token()]) + resolver = McpCredentialResolver(token_store=store) + + with ( + patch("deep_agent.aegra.mcp_auth.distributed_lock", _held_lock), + patch.object( + resolver, + "_refresh_mcp_token", + new=AsyncMock(return_value="new-access"), + ) as refresh, + ): + token = await resolver.resolve( + "user-1", + "oauth-mcp", + { + "auth_mode": "oauth", + "oauth": { + "token_endpoint": "https://as.example.com/token", + "client_id": "cid", + }, + }, + ) + + assert token == "new-access" + refresh.assert_awaited_once() diff --git a/tests/unit/aegra/test_mcp_token_store.py b/tests/unit/aegra/test_mcp_token_store.py new file mode 100644 index 00000000..4008622f --- /dev/null +++ b/tests/unit/aegra/test_mcp_token_store.py @@ -0,0 +1,106 @@ +"""Unit tests for MCP OAuth token storage in Redis.""" + +from __future__ import annotations + +import json +from datetime import UTC, datetime, timedelta +from unittest.mock import patch + +import pytest +from cryptography.fernet import Fernet + +from deep_agent.aegra.mcp_crypto import reset_mcp_crypto_cache +from deep_agent.aegra.mcp_token_store import McpTokenStore + + +@pytest.fixture(autouse=True) +def _clear_crypto_cache(): + reset_mcp_crypto_cache() + yield + reset_mcp_crypto_cache() + + +@pytest.fixture +def fernet_key(monkeypatch): + key = Fernet.generate_key().decode() + monkeypatch.setenv("MCP_TOKEN_ENCRYPTION_KEY", key) + return key + + +@pytest.fixture +def store(): + return McpTokenStore("postgresql://unused") + + +@pytest.mark.asyncio +class TestMcpTokenStoreRedis: + async def test_upsert_and_get_token_round_trip(self, store, fernet_key): + expires_at = datetime.now(UTC) + timedelta(hours=1) + stored_payload: dict[str, str] = {} + + def fake_set_persistent(key: str, value: str) -> bool: + stored_payload["key"] = key + stored_payload["value"] = value + return True + + def fake_get(key: str) -> str | None: + if key == stored_payload.get("key"): + return stored_payload.get("value") + return None + + with ( + patch( + "deep_agent.aegra.mcp_token_store.cache_set_persistent", + fake_set_persistent, + ), + patch("deep_agent.aegra.mcp_token_store.cache_get", fake_get), + ): + saved = await store.upsert_token( + user_id="user-1", + mcp_name="oauth-mcp", + access_token="access-secret", + refresh_token="refresh-secret", + expires_at=expires_at, + scopes=["read", "write"], + ) + loaded = await store.get_token("user-1", "oauth-mcp") + + assert saved.access_token == "access-secret" + assert saved.refresh_token == "refresh-secret" + assert saved.scopes == ["read", "write"] + assert loaded is not None + assert loaded.access_token == "access-secret" + assert loaded.refresh_token == "refresh-secret" + assert loaded.expires_at == expires_at + assert loaded.scopes == ["read", "write"] + + payload = json.loads(stored_payload["value"]) + assert payload["access_token"] != "access-secret" + assert payload["refresh_token"] != "refresh-secret" + + async def test_get_token_returns_none_on_miss(self, store): + with patch("deep_agent.aegra.mcp_token_store.cache_get", return_value=None): + assert await store.get_token("user-1", "oauth-mcp") is None + + async def test_upsert_token_raises_when_redis_unavailable(self, store, fernet_key): + with patch( + "deep_agent.aegra.mcp_token_store.cache_set_persistent", return_value=False + ): + with pytest.raises(RuntimeError, match="Failed to persist MCP OAuth token"): + await store.upsert_token( + user_id="user-1", + mcp_name="oauth-mcp", + access_token="access-secret", + ) + + async def test_delete_token(self, store): + deleted_keys: list[str] = [] + + def fake_delete(key: str) -> bool: + deleted_keys.append(key) + return True + + with patch("deep_agent.aegra.mcp_token_store.cache_delete", fake_delete): + assert await store.delete_token("user-1", "oauth-mcp") is True + + assert deleted_keys == ["mcp_oauth_token:default:user-1:oauth-mcp"] diff --git a/tests/unit/aegra/test_middleware.py b/tests/unit/aegra/test_middleware.py new file mode 100644 index 00000000..18465671 --- /dev/null +++ b/tests/unit/aegra/test_middleware.py @@ -0,0 +1,82 @@ +"""Unit tests for aegra middleware module.""" + +from unittest.mock import patch + +import pytest + +from deep_agent.aegra.middleware import ( + AuthError, + _hmac_validate, + authenticate, + validate_api_key, +) + + +class TestAuthError: + def test_default_status(self): + err = AuthError("fail") + assert err.status_code == 401 + assert err.message == "fail" + + def test_custom_status(self): + err = AuthError("server error", status_code=500) + assert err.status_code == 500 + + +class TestValidateApiKey: + def test_accepts_when_no_key_configured(self): + with patch("deep_agent.aegra.middleware.API_KEY", ""): + assert validate_api_key("anything") is True + + def test_accepts_correct_key(self): + with patch("deep_agent.aegra.middleware.API_KEY", "secret123"): + assert validate_api_key("secret123") is True + + def test_rejects_wrong_key(self): + with patch("deep_agent.aegra.middleware.API_KEY", "secret123"): + assert validate_api_key("wrong") is False + + +class TestHmacValidate: + def test_malformed_token_raises(self): + with pytest.raises(AuthError, match="Malformed"): + _hmac_validate("not-a-jwt") + + def test_invalid_signature_raises(self): + with patch("deep_agent.aegra.middleware.JWT_SECRET", "secret"): + with pytest.raises(AuthError, match="Invalid token signature"): + _hmac_validate("header.payload.badsig") + + +class TestAuthenticate: + def test_noop_returns_empty(self): + with patch("deep_agent.aegra.middleware.AUTH_TYPE", "noop"): + result = authenticate({}) + assert result == {} + + def test_api_key_missing_header(self): + with patch("deep_agent.aegra.middleware.AUTH_TYPE", "api_key"): + with pytest.raises(AuthError, match="Missing X-API-Key"): + authenticate({}) + + def test_api_key_invalid(self): + with patch("deep_agent.aegra.middleware.AUTH_TYPE", "api_key"): + with patch("deep_agent.aegra.middleware.API_KEY", "correct"): + with pytest.raises(AuthError, match="Invalid API key"): + authenticate({"x-api-key": "wrong"}) + + def test_api_key_valid(self): + with patch("deep_agent.aegra.middleware.AUTH_TYPE", "api_key"): + with patch("deep_agent.aegra.middleware.API_KEY", "correct"): + result = authenticate({"x-api-key": "correct"}) + assert result["auth_type"] == "api_key" + + def test_jwt_missing_header(self): + with patch("deep_agent.aegra.middleware.AUTH_TYPE", "jwt"): + with pytest.raises(AuthError, match="Missing or malformed"): + authenticate({}) + + def test_unknown_auth_type(self): + with patch("deep_agent.aegra.middleware.AUTH_TYPE", "custom_nonsense"): + with pytest.raises(AuthError, match="Unknown auth type"): + authenticate({}) diff --git a/tests/unit/aegra/test_nodes.py b/tests/unit/aegra/test_nodes.py new file mode 100644 index 00000000..8235afeb --- /dev/null +++ b/tests/unit/aegra/test_nodes.py @@ -0,0 +1,101 @@ +"""Tests for aegra.nodes module.""" + +import pytest + +from deep_agent.aegra.nodes import timed_node, with_error_handling, with_retry + + +class TestWithErrorHandling: + """Tests for the error-handling node decorator.""" + + def test_passes_through_on_success(self): + @with_error_handling("test-node") + def good_node(x: int) -> int: + return x * 2 + + assert good_node(5) == 10 + + def test_re_raises_on_failure(self): + @with_error_handling("failing-node") + def bad_node(): + raise ValueError("boom") + + with pytest.raises(ValueError, match="boom"): + bad_node() + + def test_handles_async_functions(self): + @with_error_handling("async-node") + async def async_node(x: int) -> int: + return x + 1 + + import asyncio + + result = asyncio.run(async_node(10)) + assert result == 11 + + def test_async_error_handling(self): + @with_error_handling("async-fail") + async def bad_async(): + raise RuntimeError("async boom") + + import asyncio + + with pytest.raises(RuntimeError, match="async boom"): + asyncio.run(bad_async()) + + +class TestWithRetry: + """Tests for the retry decorator.""" + + def test_succeeds_on_first_try(self): + call_count = 0 + + @with_retry(max_retries=2, delay=0.01) + def succeed(): + nonlocal call_count + call_count += 1 + return "ok" + + assert succeed() == "ok" + assert call_count == 1 + + def test_retries_on_failure(self): + call_count = 0 + + @with_retry(max_retries=2, delay=0.01) + def fail_then_succeed(): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise ValueError("not yet") + return "recovered" + + assert fail_then_succeed() == "recovered" + assert call_count == 3 + + def test_exhausts_retries(self): + @with_retry(max_retries=1, delay=0.01) + def always_fail(): + raise ValueError("permanent failure") + + with pytest.raises(ValueError, match="permanent failure"): + always_fail() + + +class TestTimedNode: + """Tests for the timing decorator.""" + + def test_returns_result(self): + @timed_node + def compute(x: int) -> int: + return x * 3 + + assert compute(7) == 21 + + def test_propagates_exceptions(self): + @timed_node + def explode(): + raise RuntimeError("kaboom") + + with pytest.raises(RuntimeError, match="kaboom"): + explode() diff --git a/tests/unit/aegra/test_oauth_client_secret.py b/tests/unit/aegra/test_oauth_client_secret.py new file mode 100644 index 00000000..63ea27d4 --- /dev/null +++ b/tests/unit/aegra/test_oauth_client_secret.py @@ -0,0 +1,84 @@ +"""Unit tests for OAuth client secret resolution from environment variables.""" + +from __future__ import annotations + +import pytest + +from deep_agent.aegra.mcp_auth import resolve_oauth_client_secret +from deep_agent.src.agent.config.loader import AgentConfig + + +class TestResolveOauthClientSecret: + def test_reads_from_env_var(self, monkeypatch): + monkeypatch.setenv("TEST_MCP_CLIENT_SECRET", "from-env") + secret = resolve_oauth_client_secret( + {"client_secret_env": "TEST_MCP_CLIENT_SECRET"}, + "oauth-mcp", + ) + assert secret == "from-env" + + def test_env_var_takes_precedence_over_inline(self, monkeypatch): + monkeypatch.setenv("TEST_MCP_CLIENT_SECRET", "from-env") + secret = resolve_oauth_client_secret( + { + "client_secret_env": "TEST_MCP_CLIENT_SECRET", + "client_secret": "inline-value", + }, + "oauth-mcp", + ) + assert secret == "from-env" + + def test_warns_on_inline_value(self, monkeypatch, caplog): + monkeypatch.delenv("TEST_MCP_CLIENT_SECRET", raising=False) + with caplog.at_level("WARNING"): + secret = resolve_oauth_client_secret( + {"client_secret": "inline-value"}, + "oauth-mcp", + ) + assert secret == "inline-value" + assert any( + "client_secret in mcp.json is insecure" in r.message for r in caplog.records + ) + + +class TestMcpConfigInlineSecretWarning: + def setup_method(self): + AgentConfig._instance = None + + @staticmethod + def _write_minimal_config_dir(tmp_path): + (tmp_path / "PROMPT.md").write_text( + """--- +name: test-orchestrator +model: gemini-2.5-flash +--- +Test prompt. +""" + ) + + def test_warns_on_inline_secret_in_mcp_json(self, tmp_path, caplog): + self._write_minimal_config_dir(tmp_path) + (tmp_path / "mcp.json").write_text( + """ + { + "mcpServers": { + "oauth-mcp": { + "url": "http://localhost/mcp", + "enabled": true, + "auth_mode": "oauth", + "oauth": { + "client_id": "cid", + "client_secret": "inline-value", + "authorization_endpoint": "https://as.example.com/authorize", + "token_endpoint": "https://as.example.com/token" + } + } + } + } + """ + ) + with caplog.at_level("WARNING"): + AgentConfig(tmp_path).get_mcp_servers() + assert any( + "client_secret in mcp.json is insecure" in r.message for r in caplog.records + ) diff --git a/tests/unit/aegra/test_oauth_scopes.py b/tests/unit/aegra/test_oauth_scopes.py new file mode 100644 index 00000000..f0bf1d8d --- /dev/null +++ b/tests/unit/aegra/test_oauth_scopes.py @@ -0,0 +1,56 @@ +"""Unit tests for OAuth scope validation.""" + +from __future__ import annotations + +from deep_agent.aegra.mcp_oauth_scopes import ( + parse_token_scopes, + requested_scopes, + validate_granted_scopes, +) + + +class TestRequestedScopes: + def test_parses_list(self): + assert requested_scopes({"scopes": ["read", "write"]}) == ["read", "write"] + + def test_parses_string(self): + assert requested_scopes({"scopes": "read write"}) == ["read", "write"] + + def test_empty_when_not_configured(self): + assert requested_scopes({}) == [] + + +class TestParseTokenScopes: + def test_parses_space_delimited_string(self): + assert parse_token_scopes({"scope": "read write"}) == ["read", "write"] + + def test_parses_list(self): + assert parse_token_scopes({"scope": ["read", "write"]}) == ["read", "write"] + + def test_returns_none_when_missing(self): + assert parse_token_scopes({}) is None + + +class TestValidateGrantedScopes: + def test_accepts_when_all_requested_granted(self): + assert validate_granted_scopes( + ["read", "write", "openid"], + ["read", "write"], + "oauth-mcp", + ) == ["read", "write", "openid"] + + def test_skips_validation_when_none_requested(self): + assert validate_granted_scopes(["read"], [], "oauth-mcp") == ["read"] + + def test_rejects_missing_scopes(self, caplog): + with caplog.at_level("ERROR"): + assert ( + validate_granted_scopes(["read"], ["read", "write"], "oauth-mcp") + is None + ) + assert any("missing requested scopes" in r.message for r in caplog.records) + + def test_rejects_empty_granted_when_scopes_requested(self, caplog): + with caplog.at_level("ERROR"): + assert validate_granted_scopes(None, ["read"], "oauth-mcp") is None + assert any("returned no scopes" in r.message for r in caplog.records) diff --git a/tests/unit/aegra/test_otel.py b/tests/unit/aegra/test_otel.py new file mode 100644 index 00000000..a4c8eb8e --- /dev/null +++ b/tests/unit/aegra/test_otel.py @@ -0,0 +1,323 @@ +"""Unit tests for OTEL telemetry initialization and shutdown.""" + +from unittest.mock import MagicMock, patch + +import pytest + +import deep_agent.aegra.otel as otel_mod +from deep_agent.aegra.otel import ( + MetricsContainer, + get_metrics, + initialize_telemetry, + reset_thread_active_tracking, + shutdown_telemetry, +) + + +@pytest.fixture(autouse=True) +def _reset_otel_state(): + """Reset module-level OTEL state before and after each test.""" + otel_mod._meter = None + otel_mod._metrics_container = None + otel_mod._initialized = False + otel_mod._otel_enabled = False + reset_thread_active_tracking() + yield + otel_mod._meter = None + otel_mod._metrics_container = None + otel_mod._initialized = False + otel_mod._otel_enabled = False + reset_thread_active_tracking() + + +class TestInitializeTelemetry: + """Test initialize_telemetry behaviour.""" + + def test_disabled_by_default_returns_gracefully(self): + """When OTEL is disabled (default), initialization should complete + without error and set up in-memory providers.""" + with patch.object( + otel_mod, + "_resolve_config", + return_value=(False, "http://localhost:4317", True, 5000, True), + ): + initialize_telemetry() + + assert otel_mod._initialized is True + assert otel_mod._otel_enabled is False + assert get_metrics() is not None + + def test_idempotent(self): + """Calling initialize_telemetry twice should be a no-op the second time.""" + with patch.object( + otel_mod, + "_resolve_config", + return_value=(False, "http://localhost:4317", True, 5000, True), + ) as mock_resolve: + initialize_telemetry() + initialize_telemetry() + + # _resolve_config is only called once (first init) + mock_resolve.assert_called_once() + + def test_get_metrics_none_before_init(self): + """get_metrics() should return None before initialization.""" + assert get_metrics() is None + + +class TestShutdownTelemetry: + """Test shutdown_telemetry behaviour.""" + + def test_does_not_raise_when_not_initialized(self): + """Calling shutdown before init should not raise.""" + shutdown_telemetry() + assert otel_mod._initialized is False + + def test_resets_initialized_flag(self): + """After shutdown, _initialized should be False.""" + with patch.object( + otel_mod, + "_resolve_config", + return_value=(False, "http://localhost:4317", True, 5000, True), + ): + initialize_telemetry() + assert otel_mod._initialized is True + + shutdown_telemetry() + assert otel_mod._initialized is False + + def test_clears_thread_tracking(self): + """Shutdown should clear the thread active tracking set.""" + with otel_mod._threads_active_lock: + otel_mod._threads_active_tracked.add("thread-1") + otel_mod._threads_active_tracked.add("thread-2") + + shutdown_telemetry() + + with otel_mod._threads_active_lock: + assert len(otel_mod._threads_active_tracked) == 0 + + +class TestResolveConfig: + """Test _resolve_config env var override logic.""" + + def test_defaults_when_no_env_vars(self): + """With no env vars and default OtelFileConfig, OTEL should be disabled.""" + from deep_agent.src.agent.config.otel import OtelFileConfig + + mock_cfg = OtelFileConfig() + with ( + patch.dict("os.environ", {}, clear=True), + patch.object( + otel_mod, + "_resolve_config", + wraps=otel_mod._resolve_config, + ), + patch( + "deep_agent.src.agent.config.otel.OtelFileConfig", + return_value=mock_cfg, + ), + ): + # Call the real function with agent_config failing + with patch( + "deep_agent.aegra.otel._resolve_config", + ) as mock_rc: + mock_rc.return_value = ( + False, + "http://localhost:4317", + True, + 5000, + True, + ) + enabled, endpoint, insecure, interval, auto = mock_rc() + + assert enabled is False + assert endpoint == "http://localhost:4317" + + def test_env_var_enables_otel(self): + """ENABLE_OTEL=true env var should override config.""" + from deep_agent.src.agent.config.otel import OtelFileConfig + + with patch.dict("os.environ", {"ENABLE_OTEL": "true"}, clear=True): + with patch( + "deep_agent.src.agent.config.agent_config.get_otel_config", + side_effect=Exception("not loaded"), + ): + enabled, endpoint, insecure, interval, auto = otel_mod._resolve_config() + + assert enabled is True + + +class TestInstrumentFastapi: + """Test instrument_fastapi behaviour.""" + + def test_skips_when_auto_instrument_disabled(self): + """Should log and return when auto_instrument is False.""" + with patch.object( + otel_mod, + "_resolve_config", + return_value=(False, "http://localhost:4317", True, 5000, False), + ): + otel_mod.instrument_fastapi(MagicMock()) + # No error should occur + + def test_handles_missing_instrumentor(self): + """Should warn when opentelemetry-instrumentation-fastapi is not installed.""" + with ( + patch.object( + otel_mod, + "_resolve_config", + return_value=(True, "http://localhost:4317", True, 5000, True), + ), + patch.dict("sys.modules", {"opentelemetry.instrumentation.fastapi": None}), + patch( + "deep_agent.aegra.otel.FastAPIInstrumentor", + side_effect=ImportError("not installed"), + ) + if False + else patch( + "builtins.__import__", + side_effect=_import_raiser("opentelemetry.instrumentation.fastapi"), + ), + ): + # Should not raise + otel_mod.instrument_fastapi(MagicMock()) + + +class TestMetricsContainer: + """Test MetricsContainer creation.""" + + def test_creates_all_instruments(self): + """MetricsContainer should create all expected metric instruments.""" + mock_meter = MagicMock() + container = MetricsContainer(mock_meter) + + assert container.conversations_total is not None + assert container.messages_total is not None + assert container.conversation_duration_seconds is not None + assert container.active_conversations is not None + assert container.stream_tokens_total is not None + assert container.stream_duration_seconds is not None + assert container.stream_errors_total is not None + assert container.time_to_first_token_seconds is not None + assert container.threads_created_total is not None + assert container.threads_active is not None + assert container.threads_deleted_total is not None + assert container.thread_messages_count is not None + + assert mock_meter.create_counter.call_count == 6 + assert mock_meter.create_histogram.call_count == 5 + assert mock_meter.create_up_down_counter.call_count == 2 + + +class TestResetThreadActiveTracking: + """Test reset_thread_active_tracking.""" + + def test_clears_set(self): + with otel_mod._threads_active_lock: + otel_mod._threads_active_tracked.add("t1") + otel_mod._threads_active_tracked.add("t2") + + reset_thread_active_tracking() + + with otel_mod._threads_active_lock: + assert len(otel_mod._threads_active_tracked) == 0 + + +def _import_raiser(blocked_module: str): + """Return an __import__ side_effect that raises ImportError for a specific module.""" + real_import = ( + __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__ + ) + + def _side_effect(name, *args, **kwargs): + if name == blocked_module: + raise ImportError(f"No module named '{blocked_module}'") + return real_import(name, *args, **kwargs) + + return _side_effect + + +class TestMetricRecording: + """Test end-to-end metric recording.""" + + def test_record_conversation_started_increments_counter(self): + """Verify recording a conversation start increments the metric.""" + with patch.object( + otel_mod, + "_resolve_config", + return_value=(False, "http://localhost:4317", True, 5000, True), + ): + initialize_telemetry() + + from deep_agent.aegra.otel import ( + get_metrics_snapshot, + record_conversation_completed, + record_conversation_started, + ) + + # Record a conversation start + start_mono = record_conversation_started(attributes={"thread_id": "test-123"}) + + # Get snapshot and verify counters increased + snapshot = get_metrics_snapshot() + assert "conversations_total" in str( + snapshot + ) # Metric name includes dynamic prefix + + # Complete it + record_conversation_completed( + start_mono, status="completed", attributes={"thread_id": "test-123"} + ) + + # Verify active conversations went back to zero + snapshot_after = get_metrics_snapshot() + # Both snapshots should have data + assert snapshot_after is not None + + def test_record_thread_deleted_raises_on_invalid_count(self): + """record_thread_deleted should reject count != 1.""" + with patch.object( + otel_mod, + "_resolve_config", + return_value=(False, "http://localhost:4317", True, 5000, True), + ): + initialize_telemetry() + + from deep_agent.aegra.otel import record_thread_deleted + + with pytest.raises(ValueError, match="requires count=1"): + record_thread_deleted(count=5, attributes={"thread_id": "test"}) + + def test_record_stream_metrics(self): + """Verify stream metric recording works.""" + with patch.object( + otel_mod, + "_resolve_config", + return_value=(False, "http://localhost:4317", True, 5000, True), + ): + initialize_telemetry() + + from deep_agent.aegra.otel import ( + get_metrics_snapshot, + record_first_token, + record_stream_completed, + record_stream_error, + record_stream_started, + ) + + # Record stream lifecycle + start_mono = record_stream_started() + record_first_token(start_mono, attributes={"model": "test"}) + record_stream_completed( + start_mono, token_count=100, attributes={"model": "test"} + ) + + snapshot = get_metrics_snapshot() + assert snapshot is not None + + # Record an error + record_stream_error(error_type="timeout", attributes={"model": "test"}) + + snapshot_after = get_metrics_snapshot() + assert snapshot_after is not None diff --git a/tests/unit/aegra/test_redis.py b/tests/unit/aegra/test_redis.py new file mode 100644 index 00000000..ff7f0eb1 --- /dev/null +++ b/tests/unit/aegra/test_redis.py @@ -0,0 +1,106 @@ +"""Unit tests for aegra redis module.""" + +from unittest.mock import MagicMock, patch + +import pytest + +import deep_agent.aegra.redis as redis_mod +from deep_agent.aegra.redis import ( + cache_delete, + cache_get, + cache_set, + get_redis_client, + get_redis_config, +) + + +@pytest.fixture(autouse=True) +def _reset_client(): + """Reset the module-level singleton before each test.""" + redis_mod._client = None + yield + redis_mod._client = None + + +class TestGetRedisConfig: + def test_returns_all_keys(self): + cfg = get_redis_config() + assert "url" in cfg + assert "max_connections" in cfg + assert "socket_timeout" in cfg + assert "retry_on_timeout" in cfg + assert "key_prefix" in cfg + + +class TestGetRedisClient: + def test_returns_cached_client(self): + mock_client = MagicMock() + redis_mod._client = mock_client + assert get_redis_client() is mock_client + + def test_returns_none_when_redis_unavailable(self): + mock_redis = MagicMock() + mock_redis.from_url.side_effect = ConnectionError("refused") + with patch.dict("sys.modules", {"redis": mock_redis}): + result = get_redis_client() + assert result is None + + def test_returns_none_when_redis_not_installed(self): + with patch.dict("sys.modules", {"redis": None}): + with patch("builtins.__import__", side_effect=ImportError("no redis")): + redis_mod._client = None + result = get_redis_client() + assert result is None + + +class TestCacheGet: + def test_returns_none_when_no_client(self): + with patch("deep_agent.aegra.redis.get_redis_client", return_value=None): + assert cache_get("key") is None + + def test_returns_value_from_redis(self): + mock_client = MagicMock() + mock_client.get.return_value = "cached_value" + redis_mod._client = mock_client + assert cache_get("key") == "cached_value" + + def test_returns_none_on_error(self): + mock_client = MagicMock() + mock_client.get.side_effect = Exception("redis error") + redis_mod._client = mock_client + assert cache_get("key") is None + + +class TestCacheSet: + def test_returns_false_when_no_client(self): + with patch("deep_agent.aegra.redis.get_redis_client", return_value=None): + assert cache_set("key", "value") is False + + def test_returns_true_on_success(self): + mock_client = MagicMock() + redis_mod._client = mock_client + assert cache_set("key", "value", ttl_seconds=60) is True + mock_client.setex.assert_called_once() + + def test_returns_false_on_error(self): + mock_client = MagicMock() + mock_client.setex.side_effect = Exception("write fail") + redis_mod._client = mock_client + assert cache_set("key", "value") is False + + +class TestCacheDelete: + def test_returns_false_when_no_client(self): + with patch("deep_agent.aegra.redis.get_redis_client", return_value=None): + assert cache_delete("key") is False + + def test_returns_true_on_success(self): + mock_client = MagicMock() + redis_mod._client = mock_client + assert cache_delete("key") is True + + def test_returns_false_on_error(self): + mock_client = MagicMock() + mock_client.delete.side_effect = Exception("fail") + redis_mod._client = mock_client + assert cache_delete("key") is False diff --git a/tests/unit/aegra/test_redis_lock.py b/tests/unit/aegra/test_redis_lock.py new file mode 100644 index 00000000..797ee652 --- /dev/null +++ b/tests/unit/aegra/test_redis_lock.py @@ -0,0 +1,56 @@ +"""Unit tests for Redis distributed locks.""" + +from unittest.mock import MagicMock, patch + +import deep_agent.aegra.redis as redis_mod +from deep_agent.aegra.redis import ( + acquire_distributed_lock, + distributed_lock, + release_distributed_lock, +) + + +class TestDistributedLock: + def test_acquire_and_release(self): + mock_client = MagicMock() + mock_client.set.return_value = True + mock_client.eval.return_value = 1 + redis_mod._client = mock_client + + token = acquire_distributed_lock( + "refresh:user:mcp", ttl_seconds=30, wait_seconds=1 + ) + assert token is not None + assert release_distributed_lock("refresh:user:mcp", token) is True + mock_client.set.assert_called_once() + mock_client.eval.assert_called_once() + + def test_acquire_returns_none_when_redis_unavailable(self): + with patch("deep_agent.aegra.redis.get_redis_client", return_value=None): + assert acquire_distributed_lock("refresh:user:mcp") is None + + +class TestDistributedLockAsync: + async def test_yields_no_redis_when_client_missing(self): + with patch("deep_agent.aegra.redis.get_redis_client", return_value=None): + async with distributed_lock("refresh:user:mcp") as state: + assert state == "no_redis" + + async def test_yields_held_when_lock_acquired(self): + with ( + patch( + "deep_agent.aegra.redis.acquire_distributed_lock", + return_value="lock-token", + ), + patch( + "deep_agent.aegra.redis.release_distributed_lock", + return_value=True, + ) as release, + patch( + "deep_agent.aegra.redis.get_redis_client", + return_value=MagicMock(), + ), + ): + async with distributed_lock("refresh:user:mcp") as state: + assert state == "held" + release.assert_called_once_with("refresh:user:mcp", "lock-token") diff --git a/tests/unit/aegra/test_request_context.py b/tests/unit/aegra/test_request_context.py new file mode 100644 index 00000000..afb50ee6 --- /dev/null +++ b/tests/unit/aegra/test_request_context.py @@ -0,0 +1,116 @@ +"""Tests for X-Request-ID, X-Org-ID, X-Agent-ID extraction and log binding.""" + +from __future__ import annotations + +import json +import uuid + +import pytest + +from deep_agent.utils.pylogger import ( + _agent_id_var, + _org_id_var, + _request_id_var, + bind_request_context, + clear_request_context, +) + + +@pytest.fixture(autouse=True) +def _clean_context(): + clear_request_context() + yield + clear_request_context() + + +# --------------------------------------------------------------------------- +# Context-var helpers +# --------------------------------------------------------------------------- + + +class TestBindRequestContext: + def test_bind_request_id(self): + bind_request_context(request_id="rid-1") + assert _request_id_var.get() == "rid-1" + + def test_bind_org_and_agent_id(self): + bind_request_context(org_id="acme", agent_id="acme/bot") + assert _org_id_var.get() == "acme" + assert _agent_id_var.get() == "acme/bot" + + def test_clear_resets_all(self): + bind_request_context(request_id="x", org_id="y", agent_id="z") + clear_request_context() + assert _request_id_var.get() is None + assert _org_id_var.get() is None + assert _agent_id_var.get() is None + + def test_backward_compat_trace_id(self): + """Existing trace_id / user_id / thread_id params still work.""" + bind_request_context(trace_id="t1", user_id="u1", thread_id="th1") + from deep_agent.utils.pylogger import ( + _thread_id_var, + _trace_id_var, + _user_id_var, + ) + + assert _trace_id_var.get() == "t1" + assert _user_id_var.get() == "u1" + assert _thread_id_var.get() == "th1" + + +# --------------------------------------------------------------------------- +# Structlog processor test +# --------------------------------------------------------------------------- + + +def test_request_id_injected_into_log_event(): + """Verify _inject_request_context adds request_id/org_id/agent_id to event dict.""" + from deep_agent.utils.pylogger import _inject_request_context + + bind_request_context(request_id="log-rid", org_id="myorg", agent_id="myorg/agent-x") + event: dict = {"event": "test_event"} + result = _inject_request_context(None, "info", event) + clear_request_context() + + assert result["request_id"] == "log-rid" + assert result["org_id"] == "myorg" + assert result["agent_id"] == "myorg/agent-x" + assert result.get("service") is not None + + +# --------------------------------------------------------------------------- +# Middleware tests (RequestContextMiddleware in http_app) +# --------------------------------------------------------------------------- + + +class TestRequestContextMiddleware: + @pytest.fixture() + def client(self): + from fastapi.testclient import TestClient + + from deep_agent.aegra.http_app import app + + return TestClient(app) + + def test_generates_request_id_when_absent(self, client): + r = client.get("/health") + rid = r.headers.get("x-request-id") + assert rid is not None + uuid.UUID(rid) + + def test_preserves_incoming_request_id(self, client): + r = client.get("/health", headers={"X-Request-ID": "agent-42"}) + assert r.headers["x-request-id"] == "agent-42" + + def test_preserves_trace_id(self, client): + r = client.get("/health", headers={"X-Trace-ID": "trace-abc"}) + assert r.headers["x-trace-id"] == "trace-abc" + + def test_both_ids_returned(self, client): + r = client.get( + "/health", + headers={"X-Request-ID": "req-1", "X-Trace-ID": "trace-1"}, + ) + assert r.headers["x-request-id"] == "req-1" + assert r.headers["x-trace-id"] == "trace-1" diff --git a/tests/unit/aegra/test_security.py b/tests/unit/aegra/test_security.py new file mode 100644 index 00000000..bb52a4cf --- /dev/null +++ b/tests/unit/aegra/test_security.py @@ -0,0 +1,299 @@ +"""Unit tests for production security hardening (RHITAIF-220).""" + +import os +from unittest.mock import patch + +import pytest +from starlette.testclient import TestClient + +from deep_agent.src.exceptions import AppException + + +class TestEnvironmentEnforcement: + """Tests for ENVIRONMENT-based security enforcement.""" + + def test_production_rejects_auth_bypass_at_startup(self): + """Test that ENVIRONMENT=production rejects ENABLE_AUTH=false at import.""" + with patch.dict( + os.environ, {"ENVIRONMENT": "production", "ENABLE_AUTH": "false"} + ): + with pytest.raises( + RuntimeError, match="ENABLE_AUTH=false is not permitted" + ): + # Re-import auth module to trigger validation + import importlib + + from deep_agent.aegra import auth + + importlib.reload(auth) + + def test_development_allows_auth_bypass(self): + """Test that ENVIRONMENT=development allows ENABLE_AUTH=false.""" + with patch.dict( + os.environ, {"ENVIRONMENT": "development", "ENABLE_AUTH": "false"} + ): + import importlib + + from deep_agent.aegra import auth + + importlib.reload(auth) + assert auth.ENVIRONMENT == "development" + assert auth.ENABLE_AUTH is False + + def test_production_flag_detection(self): + """Test settings.is_production property.""" + from deep_agent.src.settings import Settings + + prod_settings = Settings(ENVIRONMENT="production") + assert prod_settings.is_production is True + + dev_settings = Settings(ENVIRONMENT="development") + assert dev_settings.is_production is False + + +class TestMCPSSLVerificationEnforcement: + """Tests for MCP SSL verification in production.""" + + def test_production_enforces_ssl_verify_true(self): + """Test that ssl_verify=false is overridden in production.""" + from deep_agent.src.settings import Settings + + # Mock settings at module level before importing function + prod_settings = Settings(ENVIRONMENT="production") + with patch("deep_agent.src.settings.settings", prod_settings): + # Import after patching + from deep_agent.aegra.mcp import mcp_httpx_verify + + # Should return True even when config says False + assert mcp_httpx_verify({"ssl_verify": False, "name": "test"}) is True + + def test_development_allows_ssl_verify_false(self): + """Test that ssl_verify=false is allowed in development.""" + from deep_agent.src.settings import Settings + + dev_settings = Settings(ENVIRONMENT="development") + with patch("deep_agent.src.settings.settings", dev_settings): + from deep_agent.aegra.mcp import mcp_httpx_verify + + assert mcp_httpx_verify({"ssl_verify": False}) is False + + def test_ssl_verify_defaults_to_true(self): + """Test that ssl_verify defaults to True when not specified.""" + from deep_agent.aegra.mcp import mcp_httpx_verify + + assert mcp_httpx_verify({}) is True + + +class TestSecurityHeaders: + """Tests for HTTP security headers middleware.""" + + def test_security_headers_present(self): + """Test that all OWASP security headers are set.""" + from deep_agent.aegra.http_app import app + + client = TestClient(app) + + # Use a safe endpoint that doesn't require auth + with patch.dict(os.environ, {"ENABLE_AUTH": "false"}): + response = client.get("/") + + assert response.headers["X-Content-Type-Options"] == "nosniff" + assert response.headers["X-Frame-Options"] == "DENY" + assert response.headers["X-XSS-Protection"] == "1; mode=block" + assert ( + response.headers["Referrer-Policy"] == "strict-origin-when-cross-origin" + ) + assert "Permissions-Policy" in response.headers + assert "Content-Security-Policy" in response.headers + + def test_hsts_header_in_production(self): + """Test that HSTS header is set in production.""" + from deep_agent.aegra.http_app import app + from deep_agent.src.settings import Settings + + with patch( + "deep_agent.aegra.security_middleware.settings", + Settings(ENVIRONMENT="production"), + ): + client = TestClient(app) + + with patch.dict(os.environ, {"ENABLE_AUTH": "false"}): + response = client.get("/") + assert "Strict-Transport-Security" in response.headers + + +class TestRequestSizeLimit: + """Tests for request body size validation.""" + + def test_rejects_oversized_request(self): + """Test that requests exceeding max size are rejected.""" + from deep_agent.aegra.http_app import app + + client = TestClient(app) + + # Simulate oversized request via Content-Length header + with patch.dict(os.environ, {"ENABLE_AUTH": "false"}): + response = client.post( + "/feedback", + json={"trace_id": "a" * 32}, + headers={"Content-Length": str(11 * 1024 * 1024)}, # 11MB + ) + + assert response.status_code == 413 + assert "exceeds maximum size" in response.json()["detail"] + + def test_accepts_normal_sized_request(self): + """Test that normal-sized requests are accepted.""" + from deep_agent.aegra.http_app import app + + client = TestClient(app) + + normal_payload = { + "trace_id": "a" * 32, + "name": "test", + "value": 1.0, + } + + with patch.dict(os.environ, {"ENABLE_AUTH": "false"}): + with patch( + "deep_agent.aegra.feedback.get_langfuse_client", return_value=None + ): + response = client.post("/feedback", json=normal_payload) + + # Should not be rejected for size + assert response.status_code != 413 + + +class TestPIIScrubbing: + """Tests for PII scrubbing in error responses.""" + + def test_scrubs_email_addresses(self): + """Test that email addresses are redacted.""" + from deep_agent.src.pii_scrubber import scrub_pii + from deep_agent.src.settings import Settings + + with patch( + "deep_agent.src.pii_scrubber.settings", Settings(ENVIRONMENT="production") + ): + text = "Error: user john.doe@example.com not found" + scrubbed = scrub_pii(text) + assert "john.doe@example.com" not in scrubbed + assert "[EMAIL_REDACTED]" in scrubbed + + def test_scrubs_file_paths(self): + """Test that file paths are redacted.""" + from deep_agent.src.pii_scrubber import scrub_pii + from deep_agent.src.settings import Settings + + with patch( + "deep_agent.src.pii_scrubber.settings", Settings(ENVIRONMENT="production") + ): + text = "File not found: /home/user/secrets/config.yaml" + scrubbed = scrub_pii(text) + assert "/home/user/secrets" not in scrubbed + assert "[PATH]" in scrubbed + + def test_scrubs_jwt_tokens(self): + """Test that JWT tokens are redacted.""" + from deep_agent.src.pii_scrubber import scrub_pii + from deep_agent.src.settings import Settings + + with patch( + "deep_agent.src.pii_scrubber.settings", Settings(ENVIRONMENT="production") + ): + text = "Token: eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U" + scrubbed = scrub_pii(text) + assert "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9" not in scrubbed + assert "[TOKEN_REDACTED]" in scrubbed + + def test_scrubs_sensitive_dict_keys(self): + """Test that sensitive dictionary keys are redacted.""" + from deep_agent.src.pii_scrubber import scrub_dict + from deep_agent.src.settings import Settings + + with patch( + "deep_agent.src.pii_scrubber.settings", Settings(ENVIRONMENT="production") + ): + data = { + "username": "alice", + "password": "secret123", + "api_key": "sk-1234567890", + "message": "hello", + } + scrubbed = scrub_dict(data) + assert scrubbed["password"] == "[REDACTED]" + assert scrubbed["api_key"] == "[REDACTED]" + assert scrubbed["username"] == "alice" # not sensitive + assert scrubbed["message"] == "hello" + + def test_development_mode_preserves_pii(self): + """Test that PII is preserved in development mode.""" + from deep_agent.src.pii_scrubber import scrub_pii + from deep_agent.src.settings import Settings + + with patch( + "deep_agent.src.pii_scrubber.settings", Settings(ENVIRONMENT="development") + ): + text = "Error: user john@example.com at /home/user/file.txt" + scrubbed = scrub_pii(text) + # In development, nothing should be scrubbed + assert scrubbed == text + + +class TestConfigValidation: + """Tests for production configuration validation.""" + + def test_validate_config_enforces_auth_in_production(self): + """Test that validate_config rejects ENABLE_AUTH=false in production.""" + from deep_agent.src.settings import Settings, validate_config + + settings = Settings(ENVIRONMENT="production", ENABLE_AUTH=False) + + with pytest.raises(AppException, match="ENABLE_AUTH must be true"): + validate_config(settings) + + def test_validate_config_allows_dev_mode(self): + """Test that validate_config allows auth bypass in development.""" + from deep_agent.src.settings import Settings, validate_config + + settings = Settings(ENVIRONMENT="development", ENABLE_AUTH=False) + + # Should not raise + validate_config(settings) + + +class TestErrorResponseScrubbing: + """Tests for global exception handler PII scrubbing.""" + + def test_error_response_scrubbed_in_production(self): + """Test that unhandled exceptions are scrubbed in production.""" + from deep_agent.src.pii_scrubber import scrub_error_response + from deep_agent.src.settings import Settings + + with patch( + "deep_agent.src.pii_scrubber.settings", Settings(ENVIRONMENT="production") + ): + exc = ValueError("Invalid email: user@example.com") + response = scrub_error_response("Error occurred", exc) + + # Should not contain PII + assert "user@example.com" not in str(response) + # Should contain exception type but not message + assert response["exception_type"] == "ValueError" + assert "exception_message" not in response + + def test_error_response_verbose_in_development(self): + """Test that full error details are shown in development.""" + from deep_agent.src.pii_scrubber import scrub_error_response + from deep_agent.src.settings import Settings + + with patch( + "deep_agent.src.pii_scrubber.settings", Settings(ENVIRONMENT="development") + ): + exc = ValueError("Invalid email: user@example.com") + response = scrub_error_response("Error occurred", exc) + + # Should contain full details in dev mode + assert response["detail"] == "Error occurred" + assert response["exception_type"] == "ValueError" + assert "user@example.com" in response["exception_message"] diff --git a/tests/unit/aegra/test_serialization.py b/tests/unit/aegra/test_serialization.py new file mode 100644 index 00000000..ddab13b4 --- /dev/null +++ b/tests/unit/aegra/test_serialization.py @@ -0,0 +1,181 @@ +"""Unit tests for aegra serialization module.""" + +import json +from datetime import UTC, datetime + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage + +from deep_agent.aegra.serialization import ( + deserialize_message, + deserialize_state, + serialize_message, + serialize_state, + state_from_json, + state_to_json, +) + + +class TestSerializeMessage: + def test_human_message(self): + msg = HumanMessage(content="hello", id="h1") + result = serialize_message(msg) + assert result["type"] == "human" + assert result["content"] == "hello" + assert result["id"] == "h1" + + def test_ai_message_without_tool_calls(self): + msg = AIMessage(content="response", id="a1") + result = serialize_message(msg) + assert result["type"] == "ai" + assert "tool_calls" not in result + + def test_ai_message_with_tool_calls(self): + msg = AIMessage( + content="", + tool_calls=[{"id": "tc1", "name": "search", "args": {"q": "test"}}], + id="a2", + ) + result = serialize_message(msg) + assert len(result["tool_calls"]) == 1 + assert result["tool_calls"][0]["name"] == "search" + + def test_tool_message(self): + msg = ToolMessage( + content='{"result": true}', + tool_call_id="tc1", + name="search", + id="t1", + ) + result = serialize_message(msg) + assert result["type"] == "tool" + assert result["tool_call_id"] == "tc1" + assert result["name"] == "search" + + def test_system_message(self): + msg = SystemMessage(content="you are helpful") + result = serialize_message(msg) + assert result["type"] == "system" + + def test_response_metadata_included(self): + msg = AIMessage( + content="hi", + response_metadata={"model": "gemini-2.5"}, + ) + result = serialize_message(msg) + assert result["response_metadata"]["model"] == "gemini-2.5" + + +class TestDeserializeMessage: + def test_human_message(self): + data = {"type": "human", "content": "hello", "id": "h1"} + msg = deserialize_message(data) + assert isinstance(msg, HumanMessage) + assert msg.content == "hello" + + def test_ai_message(self): + data = {"type": "ai", "content": "response", "id": "a1"} + msg = deserialize_message(data) + assert isinstance(msg, AIMessage) + + def test_ai_message_with_tool_calls(self): + data = { + "type": "ai", + "content": "", + "tool_calls": [{"id": "tc1", "name": "search", "args": {"q": "t"}}], + } + msg = deserialize_message(data) + assert isinstance(msg, AIMessage) + assert msg.tool_calls[0]["name"] == "search" + + def test_system_message(self): + data = {"type": "system", "content": "sys prompt"} + msg = deserialize_message(data) + assert isinstance(msg, SystemMessage) + + def test_tool_message(self): + data = { + "type": "tool", + "content": "result", + "tool_call_id": "tc1", + "name": "search", + } + msg = deserialize_message(data) + assert isinstance(msg, ToolMessage) + assert msg.tool_call_id == "tc1" + + def test_unknown_type_defaults_to_human(self): + data = {"type": "unknown_type", "content": "fallback"} + msg = deserialize_message(data) + assert isinstance(msg, HumanMessage) + + def test_missing_type_defaults_to_human(self): + data = {"content": "no type"} + msg = deserialize_message(data) + assert isinstance(msg, HumanMessage) + + +class TestSerializeState: + def test_roundtrip(self): + state = { + "messages": [ + HumanMessage(content="hi"), + AIMessage(content="hello"), + ], + "extra": "value", + } + serialized = serialize_state(state) + assert "_serialized_at" in serialized + assert len(serialized["messages"]) == 2 + + restored = deserialize_state(serialized) + assert len(restored["messages"]) == 2 + assert isinstance(restored["messages"][0], HumanMessage) + assert isinstance(restored["messages"][1], AIMessage) + assert "_serialized_at" not in restored + + def test_non_message_values_preserved(self): + state = {"count": 42, "flag": True, "messages": []} + serialized = serialize_state(state) + assert serialized["count"] == 42 + assert serialized["flag"] is True + + +class TestStateJsonConversion: + def test_state_to_json_and_back(self): + state = { + "messages": [HumanMessage(content="test")], + "meta": {"run": "abc"}, + } + json_str = state_to_json(state) + restored = state_from_json(json_str) + assert len(restored["messages"]) == 1 + assert isinstance(restored["messages"][0], HumanMessage) + + def test_state_to_json_with_indent(self): + state = {"messages": [HumanMessage(content="x")]} + json_str = state_to_json(state, indent=2) + assert "\n" in json_str + + def test_handles_nested_objects(self): + state = { + "messages": [], + "nested": {"key": [1, 2, {"inner": "val"}]}, + } + json_str = state_to_json(state) + restored = state_from_json(json_str) + assert restored["nested"]["key"][2]["inner"] == "val" + + def test_handles_datetime(self): + state = { + "messages": [], + "timestamp": datetime.now(UTC), + } + json_str = state_to_json(state) + assert "timestamp" in json_str + + def test_handles_bytes(self): + state = {"messages": [], "data": b"hello bytes"} + json_str = state_to_json(state) + restored = state_from_json(json_str) + assert restored["data"] == "hello bytes" diff --git a/tests/unit/aegra/test_shutdown.py b/tests/unit/aegra/test_shutdown.py new file mode 100644 index 00000000..fbc975d9 --- /dev/null +++ b/tests/unit/aegra/test_shutdown.py @@ -0,0 +1,333 @@ +"""Unit tests for shutdown orchestrator.""" + +import asyncio +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import deep_agent.aegra.shutdown as shutdown_mod +from deep_agent.aegra.shutdown import ( + _clear_graph_cache, + _close_redis, + _drain, + _shutdown_langfuse, + _shutdown_langfuse_sync, + _stop_scheduler, + is_shutting_down, + register_atexit, + register_signal_handlers, + run_shutdown, + run_shutdown_sync, +) + + +@pytest.fixture(autouse=True) +def _reset_shutdown_state(): + """Reset module-level flags before each test.""" + shutdown_mod._shutting_down = False + shutdown_mod._shutdown_complete = False + shutdown_mod._async_shutdown_started = False + shutdown_mod._atexit_registered = False + yield + shutdown_mod._shutting_down = False + shutdown_mod._shutdown_complete = False + shutdown_mod._async_shutdown_started = False + shutdown_mod._atexit_registered = False + + +class TestIsShuttingDown: + def test_false_initially(self): + assert is_shutting_down() is False + + def test_true_after_flag_set(self): + shutdown_mod._shutting_down = True + assert is_shutting_down() is True + + +class TestRunShutdown: + async def test_runs_all_steps(self): + with ( + patch.object( + shutdown_mod, "_drain", new_callable=AsyncMock, return_value="ok" + ), + patch.object( + shutdown_mod, + "_shutdown_langfuse", + new_callable=AsyncMock, + return_value="ok", + ), + patch.object( + shutdown_mod, + "_stop_scheduler", + new_callable=AsyncMock, + return_value="ok", + ), + patch.object(shutdown_mod, "_clear_graph_cache", return_value="ok"), + patch.object(shutdown_mod, "_close_redis", return_value="ok"), + ): + result = await run_shutdown() + + assert result["drain"] == "ok" + assert result["langfuse"] == "ok" + assert result["scheduler"] == "ok" + assert result["graph_cache"] == "ok" + assert result["redis"] == "ok" + assert is_shutting_down() is True + assert shutdown_mod._shutdown_complete is True + + async def test_idempotent(self): + shutdown_mod._shutting_down = True + shutdown_mod._shutdown_complete = True + result = await run_shutdown() + assert result["status"] == "already_complete" + + async def test_sets_flag_immediately(self): + flag_during_drain = None + + async def capture_flag(): + nonlocal flag_during_drain + flag_during_drain = is_shutting_down() + return "ok" + + with ( + patch.object(shutdown_mod, "_drain", side_effect=capture_flag), + patch.object( + shutdown_mod, + "_shutdown_langfuse", + new_callable=AsyncMock, + return_value="ok", + ), + patch.object( + shutdown_mod, + "_stop_scheduler", + new_callable=AsyncMock, + return_value="ok", + ), + patch.object(shutdown_mod, "_clear_graph_cache", return_value="ok"), + patch.object(shutdown_mod, "_close_redis", return_value="ok"), + ): + await run_shutdown() + + assert flag_during_drain is True + + async def test_continues_after_step_failure(self): + with ( + patch.object( + shutdown_mod, "_drain", new_callable=AsyncMock, return_value="ok" + ), + patch.object( + shutdown_mod, + "_shutdown_langfuse", + new_callable=AsyncMock, + side_effect=Exception("langfuse boom"), + ), + patch.object( + shutdown_mod, + "_stop_scheduler", + new_callable=AsyncMock, + return_value="ok", + ) as mock_sched, + patch.object(shutdown_mod, "_clear_graph_cache", return_value="ok"), + patch.object(shutdown_mod, "_close_redis", return_value="ok") as mock_redis, + ): + result = await run_shutdown() + + mock_sched.assert_awaited_once() + mock_redis.assert_called_once() + assert shutdown_mod._shutdown_complete is True + + +class TestDrain: + async def test_skips_when_zero(self): + with patch.object(shutdown_mod, "SHUTDOWN_DRAIN_SECONDS", 0): + result = await _drain() + assert "skipped" in result + + async def test_sleeps_configured_duration(self): + with patch.object(shutdown_mod, "SHUTDOWN_DRAIN_SECONDS", 0.05): + t0 = time.monotonic() + result = await _drain() + elapsed = time.monotonic() - t0 + assert result == "ok" + assert elapsed >= 0.04 + + +class TestShutdownLangfuse: + async def test_calls_shutdown(self): + mock_client = MagicMock() + mock_client.shutdown = MagicMock() + with patch( + "deep_agent.aegra.telemetry.get_langfuse_client", return_value=mock_client + ): + result = await _shutdown_langfuse() + assert result == "ok" + mock_client.shutdown.assert_called_once() + + async def test_falls_back_to_flush(self): + mock_client = MagicMock(spec=[]) + mock_client.flush = MagicMock() + with patch( + "deep_agent.aegra.telemetry.get_langfuse_client", return_value=mock_client + ): + result = await _shutdown_langfuse() + assert result == "ok" + mock_client.flush.assert_called_once() + + async def test_skips_when_not_configured(self): + with patch("deep_agent.aegra.telemetry.get_langfuse_client", return_value=None): + result = await _shutdown_langfuse() + assert "skipped" in result + + async def test_handles_timeout(self): + def slow_shutdown(): + time.sleep(5) + + mock_client = MagicMock() + mock_client.shutdown = slow_shutdown + with ( + patch( + "deep_agent.aegra.telemetry.get_langfuse_client", + return_value=mock_client, + ), + patch.object(shutdown_mod, "SHUTDOWN_LANGFUSE_TIMEOUT_SECONDS", 0.1), + ): + result = await _shutdown_langfuse() + assert result == "timeout" + + async def test_handles_exception(self): + mock_client = MagicMock() + mock_client.shutdown.side_effect = RuntimeError("boom") + with patch( + "deep_agent.aegra.telemetry.get_langfuse_client", return_value=mock_client + ): + result = await _shutdown_langfuse() + assert "error" in result + + +class TestStopScheduler: + async def test_stops_scheduler(self): + with patch( + "deep_agent.src.memory.scheduler.stop_scheduler", + new_callable=AsyncMock, + ): + result = await _stop_scheduler() + assert result == "ok" + + async def test_handles_timeout(self): + async def slow_stop(): + await asyncio.sleep(10) + + with ( + patch( + "deep_agent.src.memory.scheduler.stop_scheduler", + side_effect=slow_stop, + ), + patch.object(shutdown_mod, "SHUTDOWN_SCHEDULER_TIMEOUT_SECONDS", 0.1), + ): + result = await _stop_scheduler() + assert result == "timeout" + + async def test_handles_exception(self): + with patch( + "deep_agent.src.memory.scheduler.stop_scheduler", + new_callable=AsyncMock, + side_effect=RuntimeError("boom"), + ): + result = await _stop_scheduler() + assert "error" in result + + +class TestClearGraphCache: + @pytest.fixture(autouse=True) + def _mock_graph_module(self): + """Pre-load a fake graph module to avoid langgraph_sdk import.""" + import sys + import types + + fake_graph = types.ModuleType("deep_agent.aegra.graph") + fake_graph._graph_cache = {} + fake_graph._graph_cache_ts = {} + self._fake_graph = fake_graph + sys.modules["deep_agent.aegra.graph"] = fake_graph + yield + sys.modules.pop("deep_agent.aegra.graph", None) + + def test_clears_both_dicts(self): + self._fake_graph._graph_cache["key1"] = "value1" + self._fake_graph._graph_cache_ts["key1"] = 1234.0 + + result = _clear_graph_cache() + + assert result == "ok" + assert len(self._fake_graph._graph_cache) == 0 + assert len(self._fake_graph._graph_cache_ts) == 0 + + def test_ok_when_empty(self): + result = _clear_graph_cache() + assert result == "ok" + + +class TestCloseRedis: + def test_calls_close(self): + with patch("deep_agent.aegra.redis.close_redis_client") as mock_close: + result = _close_redis() + assert result == "ok" + mock_close.assert_called_once() + + def test_handles_exception(self): + with patch( + "deep_agent.aegra.redis.close_redis_client", + side_effect=RuntimeError("boom"), + ): + result = _close_redis() + assert "error" in result + + +class TestRunShutdownSync: + def test_noop_when_complete(self): + shutdown_mod._shutdown_complete = True + run_shutdown_sync() + + def test_skips_when_async_already_ran(self): + shutdown_mod._shutting_down = True + run_shutdown_sync() + assert shutdown_mod._shutdown_complete is True + + def test_runs_sync_cleanup(self): + with ( + patch.object(shutdown_mod, "_shutdown_langfuse_sync", return_value="ok"), + patch.object(shutdown_mod, "_clear_graph_cache", return_value="ok"), + patch.object(shutdown_mod, "_close_redis", return_value="ok"), + ): + run_shutdown_sync() + assert shutdown_mod._shutting_down is True + assert shutdown_mod._shutdown_complete is True + + +class TestRegisterAtexit: + def test_registers_callback(self): + import atexit + + with patch.object(atexit, "register") as mock_register: + register_atexit() + mock_register.assert_called_once_with(run_shutdown_sync) + + def test_idempotent(self): + import atexit + + with patch.object(atexit, "register") as mock_register: + register_atexit() + register_atexit() + mock_register.assert_called_once() + + +class TestRegisterSignalHandlers: + async def test_registers_on_running_loop(self): + import signal + + register_signal_handlers() + + loop = asyncio.get_running_loop() + assert loop.remove_signal_handler(signal.SIGTERM) is True + assert loop.remove_signal_handler(signal.SIGINT) is True diff --git a/tests/unit/aegra/test_startup.py b/tests/unit/aegra/test_startup.py new file mode 100644 index 00000000..01528eea --- /dev/null +++ b/tests/unit/aegra/test_startup.py @@ -0,0 +1,170 @@ +"""Unit tests for startup orchestrator.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +from deep_agent.aegra import startup + + +class TestRunStartup: + def setup_method(self): + startup._startup_complete = False + + async def test_runs_all_steps(self): + with ( + patch.object( + startup, "_validate_config", new_callable=AsyncMock, return_value="ok" + ), + patch.object( + startup, "_ensure_database", new_callable=AsyncMock, return_value="ok" + ), + patch.object( + startup, "_warm_caches", new_callable=AsyncMock, return_value="ok" + ), + patch.object( + startup, + "_start_scheduler", + new_callable=AsyncMock, + return_value="ok", + ), + patch.object(startup, "_setup_telemetry", return_value="ok"), + ): + result = await startup.run_startup() + assert result["config"] == "ok" + assert result["database"] == "ok" + assert result["cache"] == "ok" + assert result["scheduler"] == "ok" + assert result["telemetry"] == "ok" + assert startup.is_ready() is True + + async def test_idempotent(self): + startup._startup_complete = True + result = await startup.run_startup() + assert result["status"] == "already_complete" + + +class TestValidateConfig: + async def test_valid(self): + with patch( + "deep_agent.src.settings.validate_config", + ): + result = await startup._validate_config() + assert result == "ok" + + async def test_warning(self): + with patch( + "deep_agent.src.settings.validate_config", + side_effect=ValueError("bad port"), + ): + result = await startup._validate_config() + assert "warning" in result + + +class TestEnsureDatabase: + async def test_no_db(self): + mock_settings = MagicMock() + mock_settings.database_uri = "" + mock_settings.MONGODB_URI = "" + with patch("deep_agent.src.settings.settings", mock_settings): + result = await startup._ensure_database() + assert "skipped" in result + + async def test_db_ok(self): + mock_settings = MagicMock() + mock_settings.database_uri = "postgresql://test" + mock_settings.MONGODB_URI = "" + mock_personalization = AsyncMock() + mock_feedback = AsyncMock() + mock_mcp_store = AsyncMock() + with ( + patch("deep_agent.src.settings.settings", mock_settings), + patch( + "deep_agent.src.personalization.repository.PersonalizationRepository", + return_value=mock_personalization, + ), + patch( + "deep_agent.src.feedback.repository.FeedbackRepository", + return_value=mock_feedback, + ), + patch( + "deep_agent.aegra.mcp_token_store.McpTokenStore", + return_value=mock_mcp_store, + ), + ): + result = await startup._ensure_database() + assert result == "ok" + mock_personalization.ensure_tables.assert_awaited_once() + mock_feedback.ensure_table.assert_awaited_once() + mock_mcp_store.ensure_tables.assert_awaited_once() + + async def test_mongo_indexes_when_configured(self): + import sys + + mock_settings = MagicMock() + mock_settings.database_uri = "" + mock_settings.MONGODB_URI = "mongodb://test" + mock_settings.MONGODB_DB = "tokenusage" + mock_mongo = AsyncMock() + mock_module = MagicMock() + mock_module.TokenUsageMongoRepository.return_value = mock_mongo + with ( + patch("deep_agent.src.settings.settings", mock_settings), + patch.dict( + sys.modules, + {"deep_agent.src.token_budget.mongo_repository": mock_module}, + ), + ): + result = await startup._ensure_database() + assert result == "ok" + mock_mongo.ensure_indexes.assert_awaited_once() + + +class TestWarmCaches: + async def test_disabled(self): + mock_cache_settings = MagicMock() + mock_cache_settings.CACHE_ENABLED = False + with patch("deep_agent.src.cache.config.cache_settings", mock_cache_settings): + result = await startup._warm_caches() + assert "skipped" in result + + async def test_enabled(self): + mock_cache_settings = MagicMock() + mock_cache_settings.CACHE_ENABLED = True + with ( + patch("deep_agent.src.cache.config.cache_settings", mock_cache_settings), + patch( + "deep_agent.src.cache.warming.warm_caches", + new_callable=AsyncMock, + ), + ): + result = await startup._warm_caches() + assert result == "ok" + + +class TestStartScheduler: + async def test_disabled(self): + mock_mem_settings = MagicMock() + mock_mem_settings.MEMORY_CONSOLIDATION_ENABLED = False + with patch("deep_agent.src.memory.config.memory_settings", mock_mem_settings): + result = await startup._start_scheduler() + assert "skipped" in result + + +class TestSetupTelemetry: + def test_ok(self): + with patch("deep_agent.aegra.telemetry.setup_langfuse_tracing"): + result = startup._setup_telemetry() + assert result == "ok" + + def test_failure(self): + with patch( + "deep_agent.aegra.telemetry.setup_langfuse_tracing", + side_effect=Exception("boom"), + ): + result = startup._setup_telemetry() + assert "warning" in result + + +class TestIsReady: + def test_not_ready_initially(self): + startup._startup_complete = False + assert startup.is_ready() is False diff --git a/tests/unit/aegra/test_state.py b/tests/unit/aegra/test_state.py new file mode 100644 index 00000000..ffa9bcb2 --- /dev/null +++ b/tests/unit/aegra/test_state.py @@ -0,0 +1,83 @@ +"""Tests for aegra.state module.""" + +from deep_agent.aegra.state import ( + AegraMetadata, + HealthStatus, + make_health_status, + serialize_metadata, +) + + +class TestAegraMetadata: + """Tests for AegraMetadata TypedDict operations.""" + + def test_full_metadata_creation(self): + meta: AegraMetadata = { + "run_id": "run-123", + "trace_id": "trace-456", + "thread_id": "thread-789", + "session_id": "session-abc", + "user_id": "user-def", + "stream_tokens": True, + "error_count": 0, + "last_error": None, + } + assert meta["run_id"] == "run-123" + assert meta["error_count"] == 0 + + def test_partial_metadata_creation(self): + meta: AegraMetadata = {"run_id": "run-123", "thread_id": "thread-456"} + assert meta["run_id"] == "run-123" + assert "user_id" not in meta + + +class TestSerializeMetadata: + """Tests for serialize_metadata helper.""" + + def test_strips_none_values(self): + meta: AegraMetadata = { + "run_id": "run-123", + "last_error": None, + } + result = serialize_metadata(meta) + assert "run_id" in result + assert "last_error" not in result + + def test_preserves_falsy_non_none_values(self): + meta: AegraMetadata = {"error_count": 0, "stream_tokens": False} + result = serialize_metadata(meta) + assert result["error_count"] == 0 + assert result["stream_tokens"] is False + + def test_empty_metadata(self): + result = serialize_metadata({}) + assert result == {} + + +class TestMakeHealthStatus: + """Tests for make_health_status factory.""" + + def test_produces_valid_health_status(self): + status: HealthStatus = make_health_status( + agent_name="orchestrator", + model="gemini-3.1-pro-preview", + mcp_tools_count=4, + subagents_count=2, + backend_ready=True, + ) + assert status["status"] == "healthy" + assert status["agent_name"] == "orchestrator" + assert status["mcp_tools_loaded"] == 4 + assert status["subagents_loaded"] == 2 + assert status["backend_ready"] is True + + def test_zero_tools_and_subagents(self): + status = make_health_status( + agent_name="test", + model="test-model", + mcp_tools_count=0, + subagents_count=0, + backend_ready=False, + ) + assert status["mcp_tools_loaded"] == 0 + assert status["backend_ready"] is False diff --git a/tests/unit/agent/config/test_config.py b/tests/unit/agent/config/test_config.py new file mode 100644 index 00000000..bac60d92 --- /dev/null +++ b/tests/unit/agent/config/test_config.py @@ -0,0 +1,181 @@ +"""Unit tests for agent_config skill path resolution.""" + +import pytest + +from deep_agent.src.agent.config import AgentConfig +from deep_agent.src.exceptions import AppException + + +class TestAgentConfigSkillResolution: + """Test that skills are resolved during config loading.""" + + def setup_method(self): + """Reset the singleton before each test.""" + AgentConfig._instance = None + + def test_orchestrator_loads_with_skill_paths(self, tmp_path): + """Test that orchestrator config includes resolved skill paths.""" + config_dir = tmp_path / "agent_config" + config_dir.mkdir() + + skills_dir = config_dir / "skills" + skills_dir.mkdir() + (skills_dir / "client-intake").mkdir() + + prompt_md = config_dir / "PROMPT.md" + prompt_md.write_text("""--- +name: test-orchestrator +model: gemini-2.5-flash +skills: + - client-intake +--- + +Test orchestrator prompt. +""") + + agent_cfg = AgentConfig(config_dir) + orchestrator = agent_cfg.get_orchestrator_config() + + assert "skill_paths" in orchestrator + assert len(orchestrator["skill_paths"]) == 1 + assert "client-intake" in orchestrator["skill_paths"][0] + + def test_subagent_loads_with_skill_paths(self, tmp_path): + """Test that subagent configs include resolved skill paths.""" + config_dir = tmp_path / "agent_config" + config_dir.mkdir() + + skills_dir = config_dir / "skills" + skills_dir.mkdir() + (skills_dir / "bmi-report").mkdir() + + prompt_md = config_dir / "PROMPT.md" + prompt_md.write_text("""--- +name: orchestrator +model: gemini-2.5-flash +--- +Minimal orchestrator. +""") + + subagents_dir = config_dir / "subagents" + subagents_dir.mkdir() + + analyst_md = subagents_dir / "analyst.md" + analyst_md.write_text("""--- +name: analyst +model: gemini-2.5-flash +skills: + - bmi-report +--- + +Test analyst prompt. +""") + + agent_cfg = AgentConfig(config_dir) + subagents = agent_cfg.get_all_subagent_configs() + + assert "analyst" in subagents + assert "skill_paths" in subagents["analyst"] + assert len(subagents["analyst"]["skill_paths"]) == 1 + assert "bmi-report" in subagents["analyst"]["skill_paths"][0] + + def test_missing_skills_are_logged(self, tmp_path, caplog): + """Test that missing skills generate warnings.""" + config_dir = tmp_path / "agent_config" + config_dir.mkdir() + + skills_dir = config_dir / "skills" + skills_dir.mkdir() + + prompt_md = config_dir / "PROMPT.md" + prompt_md.write_text("""--- +name: test-orchestrator +model: gemini-2.5-flash +skills: + - nonexistent-skill +--- + +Test orchestrator prompt. +""") + + agent_cfg = AgentConfig(config_dir) + orchestrator = agent_cfg.get_orchestrator_config() + + skill_paths = orchestrator.get("skill_paths", []) + assert len(skill_paths) == 0 + + assert "unknown skills" in caplog.text.lower() + + +class TestMcpsValidation: + """Test mcps field validation for orchestrator and subagents.""" + + def setup_method(self): + AgentConfig._instance = None + + def test_orchestrator_valid_mcps(self, tmp_path): + """Valid mcps list of strings loads without error.""" + config_dir = tmp_path / "agent_config" + config_dir.mkdir() + (config_dir / "skills").mkdir() + + (config_dir / "PROMPT.md").write_text("""--- +name: orch +model: gemini-2.5-flash +mcps: + - web-search + - dataverse-mcp +--- +Orchestrator. +""") + + cfg = AgentConfig(config_dir) + orch = cfg.get_orchestrator_config() + assert orch["mcps"] == ["web-search", "dataverse-mcp"] + + def test_orchestrator_invalid_mcps_raises(self, tmp_path): + """Non-list mcps raises AppException.""" + config_dir = tmp_path / "agent_config" + config_dir.mkdir() + (config_dir / "skills").mkdir() + + (config_dir / "PROMPT.md").write_text("""--- +name: orch +model: gemini-2.5-flash +mcps: "not-a-list" +--- +Orchestrator. +""") + + with pytest.raises(AppException, match="must be a list of strings"): + cfg = AgentConfig(config_dir) + cfg.get_orchestrator_config() + + def test_subagent_invalid_mcps_is_skipped(self, tmp_path, caplog): + """Subagent with non-string mcps entries is skipped and logged.""" + config_dir = tmp_path / "agent_config" + config_dir.mkdir() + (config_dir / "skills").mkdir() + + (config_dir / "PROMPT.md").write_text("""--- +name: orch +model: gemini-2.5-flash +--- +Orchestrator. +""") + + sub_dir = config_dir / "subagents" + sub_dir.mkdir() + (sub_dir / "bad.md").write_text("""--- +name: bad-agent +model: gemini-2.5-flash +mcps: + - 123 +--- +Bad agent. +""") + + cfg = AgentConfig(config_dir) + subs = cfg.get_all_subagent_configs() + assert "bad-agent" not in subs + assert "must be a list of strings" in caplog.text diff --git a/tests/unit/agent/test_llm.py b/tests/unit/agent/test_llm.py new file mode 100644 index 00000000..e3a1057b --- /dev/null +++ b/tests/unit/agent/test_llm.py @@ -0,0 +1,110 @@ +"""Unit tests for LLM model configuration and initialization.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.agent.llm import CLAUDE_MODELS, GEMINI_MODELS, create_model +from deep_agent.src.exceptions import LLMError + + +class TestCreateModel: + """Tests for create_model function.""" + + def test_create_gemini_model(self): + """Test creating Gemini model.""" + mock_creds = MagicMock() + + with patch( + "deep_agent.src.agent.llm.get_service_account_credentials" + ) as mock_get_creds: + mock_get_creds.return_value = (mock_creds, "test-project") + + with patch("deep_agent.src.agent.llm.ChatGoogleGenerativeAI") as mock_chat: + create_model("gemini-2.5-pro", temperature=0.5) + mock_chat.assert_called_once_with( + model="gemini-2.5-pro", + temperature=0.5, + credentials=mock_creds, + project="test-project", + max_output_tokens=8192, + max_retries=2, + ) + + def test_create_claude_model(self): + """Test creating Claude model.""" + mock_creds = MagicMock() + + with patch( + "deep_agent.src.agent.llm.get_service_account_credentials" + ) as mock_get_creds: + mock_get_creds.return_value = (mock_creds, "test-project") + + with patch("deep_agent.src.agent.llm.ChatAnthropicVertex") as mock_chat: + create_model("claude-sonnet-4", temperature=0.7) + mock_chat.assert_called_once_with( + model="claude-sonnet-4", + project="test-project", + credentials=mock_creds, + temperature=0.7, + max_tokens=8192, + max_retries=2, + ) + + @pytest.mark.parametrize( + "invalid_name", + ["", " ", None], + ) + def test_invalid_model_name_raises_error(self, invalid_name): + """Test that empty/whitespace/None model names raise ValueError.""" + with pytest.raises(ValueError, match="model_name cannot be empty"): + create_model(invalid_name) + + def test_unknown_model_raises_error_with_supported_list(self): + """Test that unknown model raises error listing supported models.""" + mock_creds = MagicMock() + + with patch( + "deep_agent.src.agent.llm.get_service_account_credentials" + ) as mock_get_creds: + mock_get_creds.return_value = (mock_creds, "test-project") + + with pytest.raises(ValueError) as exc_info: + create_model("gpt-4") + + error_msg = str(exc_info.value) + assert "Unknown model 'gpt-4'" in error_msg + assert "Supported models:" in error_msg + + def test_model_creation_errors_are_raised(self): + """Test that model creation errors are raised.""" + mock_creds = MagicMock() + + with patch( + "deep_agent.src.agent.llm.get_service_account_credentials" + ) as mock_get_creds: + mock_get_creds.return_value = (mock_creds, "test-project") + + with patch( + "deep_agent.src.agent.llm.ChatGoogleGenerativeAI", + side_effect=RuntimeError("API error"), + ): + with pytest.raises(LLMError, match="API error"): + create_model("gemini-2.5-pro") + + def test_all_supported_models_work(self): + """Test that all models in GEMINI_MODELS and CLAUDE_MODELS are supported.""" + mock_creds = MagicMock() + + with patch( + "deep_agent.src.agent.llm.get_service_account_credentials" + ) as mock_get_creds: + mock_get_creds.return_value = (mock_creds, "test-project") + + with patch("deep_agent.src.agent.llm.ChatGoogleGenerativeAI"): + for model_name in GEMINI_MODELS: + create_model(model_name) + + with patch("deep_agent.src.agent.llm.ChatAnthropicVertex"): + for model_name in CLAUDE_MODELS: + create_model(model_name) diff --git a/tests/unit/agent/test_provider_factory.py b/tests/unit/agent/test_provider_factory.py new file mode 100644 index 00000000..a709b239 --- /dev/null +++ b/tests/unit/agent/test_provider_factory.py @@ -0,0 +1,287 @@ +"""Unit tests for model config parsing and provider factory.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.agent.config.model import ( + ModelSpec, + Provider, + infer_provider, + model_spec_cache_key, + parse_model_config, +) +from deep_agent.src.agent.provider_factory import ( + _create_by_provider, + create_model_from_spec, +) + + +class TestInferProvider: + """Tests for legacy model name provider inference.""" + + def test_gemini_models_infer_vertex(self): + assert infer_provider("gemini-2.5-pro") == Provider.VERTEX + assert infer_provider("gemini-2.5-flash") == Provider.VERTEX + + def test_claude_models_infer_vertex(self): + assert infer_provider("claude-sonnet-4") == Provider.VERTEX + + def test_gpt_models_infer_openai(self): + assert infer_provider("gpt-4o-mini") == Provider.OPENAI + assert infer_provider("gpt-4") == Provider.OPENAI + assert infer_provider("gpt-3.5-turbo") == Provider.OPENAI + # Case-insensitive matching + assert infer_provider("GPT-4") == Provider.OPENAI + assert infer_provider("Gpt-4o") == Provider.OPENAI + + def test_unknown_models_infer_maas(self): + assert infer_provider("mistral-7b") == Provider.MAAS + assert infer_provider("llama-3-70b") == Provider.MAAS + assert infer_provider("custom-model") == Provider.MAAS + + +class TestParseModelConfig: + """Tests for parse_model_config().""" + + def test_parses_legacy_string_vertex(self): + spec = parse_model_config("gemini-2.5-pro") + assert spec.provider == Provider.VERTEX + assert spec.name == "gemini-2.5-pro" + assert spec.fallback is None + + def test_parses_legacy_string_openai(self): + spec = parse_model_config("gpt-4o-mini") + assert spec.provider == Provider.OPENAI + assert spec.name == "gpt-4o-mini" + + def test_parses_legacy_string_maas(self): + spec = parse_model_config("mistral-7b") + assert spec.provider == Provider.MAAS + assert spec.name == "mistral-7b" + + def test_parses_object_without_provider_infers(self): + """Provider is optional in dict format - infers from name.""" + spec = parse_model_config({"name": "gpt-4"}) + assert spec.provider == Provider.OPENAI # Inferred + assert spec.name == "gpt-4" + + spec2 = parse_model_config({"name": "gemini-2.5-pro"}) + assert spec2.provider == Provider.VERTEX # Inferred + + spec3 = parse_model_config({"name": "mistral-7b"}) + assert spec3.provider == Provider.MAAS # Inferred + + def test_parses_object_form(self): + spec = parse_model_config({"provider": "vertex", "name": "gemini-2.5-pro"}) + assert spec.provider == Provider.VERTEX + assert spec.name == "gemini-2.5-pro" + + def test_parses_object_with_fallback(self): + spec = parse_model_config( + { + "provider": "vertex", + "name": "gemini-2.5-pro", + "fallback": {"provider": "openai", "name": "gpt-4o-mini"}, + } + ) + assert spec.fallback is not None + assert spec.fallback.provider == Provider.OPENAI + assert spec.fallback.name == "gpt-4o-mini" + assert spec.fallback.fallback is None + + def test_parses_fallback_without_provider_infers(self): + """Fallback can omit provider - infers from name.""" + spec = parse_model_config( + { + "provider": "vertex", + "name": "gemini-2.5-pro", + "fallback": {"name": "gpt-4"}, # No provider - inferred + } + ) + assert spec.fallback is not None + assert spec.fallback.provider == Provider.OPENAI # Inferred from "gpt-4" + assert spec.fallback.name == "gpt-4" + + def test_rejects_empty_string(self): + with pytest.raises(ValueError, match="cannot be empty"): + parse_model_config("") + + def test_rejects_invalid_provider(self): + with pytest.raises(ValueError, match="invalid provider"): + parse_model_config({"provider": "azure", "name": "gpt-4"}) + + def test_rejects_missing_name(self): + with pytest.raises(ValueError, match="requires non-empty 'name'"): + parse_model_config({"provider": "vertex"}) + + def test_rejects_unknown_keys(self): + with pytest.raises(ValueError, match="unknown model config keys"): + parse_model_config( + {"provider": "vertex", "name": "gemini-2.5-pro", "extra": "x"} + ) + + def test_rejects_nested_fallback(self): + with pytest.raises(ValueError, match="nested fallback"): + parse_model_config( + { + "provider": "vertex", + "name": "gemini-2.5-pro", + "fallback": { + "provider": "openai", + "name": "gpt-4o-mini", + "fallback": {"provider": "vertex", "name": "gemini-2.5-flash"}, + }, + } + ) + + def test_display_name_with_fallback(self): + spec = parse_model_config( + { + "provider": "vertex", + "name": "gemini-2.5-pro", + "fallback": {"provider": "openai", "name": "gpt-4o-mini"}, + } + ) + assert "fallback" in spec.display_name() + + +class TestModelSpecCacheKey: + """Tests for model_spec_cache_key().""" + + def test_key_without_fallback(self): + spec = ModelSpec(provider=Provider.VERTEX, name="gemini-2.5-pro") + assert model_spec_cache_key(spec) == "vertex:gemini-2.5-pro" + + def test_key_with_fallback(self): + spec = ModelSpec( + provider=Provider.VERTEX, + name="gemini-2.5-pro", + fallback=ModelSpec(provider=Provider.OPENAI, name="gpt-4o-mini"), + ) + assert model_spec_cache_key(spec) == "vertex:gemini-2.5-pro→openai:gpt-4o-mini" + + +class TestCreateModelFromSpec: + """Tests for create_model_from_spec() routing.""" + + def test_routes_vertex_provider(self): + mock_model = MagicMock() + spec = ModelSpec(provider=Provider.VERTEX, name="gemini-2.5-pro") + + with patch( + "deep_agent.src.agent.provider_factory._create_by_provider", + return_value=mock_model, + ) as mock_create: + result = create_model_from_spec(spec) + + assert result is mock_model + mock_create.assert_called_once() + assert mock_create.call_args[0][0] == Provider.VERTEX + + def test_routes_openai_provider(self): + mock_model = MagicMock() + spec = ModelSpec(provider=Provider.OPENAI, name="gpt-4o-mini") + + with patch( + "deep_agent.src.agent.provider_factory._create_by_provider", + return_value=mock_model, + ) as mock_create: + result = create_model_from_spec(spec) + + assert result is mock_model + assert mock_create.call_args[0][0] == Provider.OPENAI + + def test_routes_maas_provider(self): + mock_model = MagicMock() + spec = ModelSpec(provider=Provider.MAAS, name="mistral-7b") + + with patch( + "deep_agent.src.agent.provider_factory._create_by_provider", + return_value=mock_model, + ) as mock_create: + result = create_model_from_spec(spec) + + assert result is mock_model + assert mock_create.call_args[0][0] == Provider.MAAS + + +class TestFallbackChain: + """Tests for primary → secondary fallback chaining.""" + + def test_with_fallbacks_called_when_fallback_present(self): + primary = MagicMock() + secondary = MagicMock() + chained = MagicMock() + primary.with_fallbacks.return_value = chained + + spec = ModelSpec( + provider=Provider.VERTEX, + name="gemini-2.5-pro", + fallback=ModelSpec(provider=Provider.OPENAI, name="gpt-4o-mini"), + ) + + with patch( + "deep_agent.src.agent.provider_factory._create_by_provider", + side_effect=[primary, secondary], + ): + result = create_model_from_spec(spec) + + assert result is chained + primary.with_fallbacks.assert_called_once_with([secondary]) + + def test_no_with_fallbacks_when_no_fallback(self): + primary = MagicMock() + spec = ModelSpec(provider=Provider.VERTEX, name="gemini-2.5-pro") + + with patch( + "deep_agent.src.agent.provider_factory._create_by_provider", + return_value=primary, + ): + result = create_model_from_spec(spec) + + assert result is primary + primary.with_fallbacks.assert_not_called() + + +class TestCreateByProvider: + """Tests for _create_by_provider() delegation.""" + + def test_vertex_delegates_to_vertex_model(self): + mock_model = MagicMock() + with patch( + "deep_agent.src.agent.provider_factory._create_vertex_model", + return_value=mock_model, + ) as mock_vertex: + result = _create_by_provider( + Provider.VERTEX, + "gemini-2.5-pro", + temperature=0.0, + max_output_tokens=8192, + ) + assert result is mock_model + mock_vertex.assert_called_once_with("gemini-2.5-pro", 0.0, 8192) + + def test_openai_delegates_to_vllm_model(self): + mock_model = MagicMock() + with patch( + "deep_agent.src.agent.provider_factory._create_vllm_model", + return_value=mock_model, + ) as mock_vllm: + result = _create_by_provider( + Provider.OPENAI, "gpt-4o-mini", temperature=0.0, max_output_tokens=4096 + ) + assert result is mock_model + mock_vllm.assert_called_once_with("gpt-4o-mini", 0.0, 4096) + + def test_maas_delegates_to_vllm_model(self): + mock_model = MagicMock() + with patch( + "deep_agent.src.agent.provider_factory._create_vllm_model", + return_value=mock_model, + ) as mock_vllm: + result = _create_by_provider( + Provider.MAAS, "mistral-7b", temperature=0.0, max_output_tokens=4096 + ) + assert result is mock_model + mock_vllm.assert_called_once_with("mistral-7b", 0.0, 4096) diff --git a/tests/unit/audit/test_buffer.py b/tests/unit/audit/test_buffer.py new file mode 100644 index 00000000..5b76a22c --- /dev/null +++ b/tests/unit/audit/test_buffer.py @@ -0,0 +1,37 @@ +"""Unit tests for audit in-memory buffer.""" + +from unittest.mock import patch + +from deep_agent.src.audit.buffer import drain, enqueue + + +class TestAuditBuffer: + def test_enqueue_and_drain(self): + with patch("deep_agent.src.audit.buffer.settings") as mock_settings: + mock_settings.PLATFORM_AUDIT_BUFFER_MAX = 10 + + import deep_agent.src.audit.buffer as buffer_mod + + buffer_mod._queue.clear() + buffer_mod._dropped = 0 + + envelope = {"event": "platform.audit", "audit_event_type": "llm_call"} + enqueue(envelope) + assert drain() == [envelope] + assert drain() == [] + + def test_drops_when_full(self): + with patch("deep_agent.src.audit.buffer.settings") as mock_settings: + mock_settings.PLATFORM_AUDIT_BUFFER_MAX = 2 + + import deep_agent.src.audit.buffer as buffer_mod + + buffer_mod._queue.clear() + buffer_mod._dropped = 0 + + enqueue({"id": 1}) + enqueue({"id": 2}) + enqueue({"id": 3}) + + assert drain() == [{"id": 1}, {"id": 2}] + assert buffer_mod._dropped == 1 diff --git a/tests/unit/audit/test_emitter.py b/tests/unit/audit/test_emitter.py new file mode 100644 index 00000000..8cb058b0 --- /dev/null +++ b/tests/unit/audit/test_emitter.py @@ -0,0 +1,56 @@ +"""Unit tests for platform audit emitter.""" + +import json +from io import StringIO +from unittest.mock import patch + +import pytest + +from deep_agent.src.audit.context import bind_audit_context, clear_audit_context +from deep_agent.src.audit.emitter import emit_audit_event + + +@pytest.fixture(autouse=True) +def _clear_context(): + clear_audit_context() + yield + clear_audit_context() + + +class TestEmitAuditEventDisabled: + def test_noop_when_disabled(self): + with patch("deep_agent.src.audit.emitter.is_audit_enabled", return_value=False): + with patch( + "deep_agent.src.audit.emitter.sys.stdout", new_callable=StringIO + ) as out: + emit_audit_event("llm_call", model="test") + assert out.getvalue() == "" + + +class TestEmitAuditEventEnabled: + def test_emits_envelope(self): + bind_audit_context(user="alice@example.com", org="acme", trace_id="trace-1") + with patch("deep_agent.src.audit.emitter.is_audit_enabled", return_value=True): + with patch( + "deep_agent.src.audit.emitter.sys.stdout", new_callable=StringIO + ) as out: + emit_audit_event("llm_call", model="gemini", phase="start") + record = json.loads(out.getvalue().strip()) + assert record["event"] == "platform.audit" + assert record["audit_event_type"] == "llm_call" + assert record["user"] == "alice@example.com" + assert record["org"] == "acme" + assert record["trace_id"] == "trace-1" + assert record["details"]["model"] == "gemini" + assert record["logger"] == "platform.audit" + assert record["level"] == "info" + + def test_buffers_on_emit_failure(self): + with patch("deep_agent.src.audit.emitter.is_audit_enabled", return_value=True): + with patch("deep_agent.src.audit.emitter.sys.stdout") as mock_stdout: + mock_stdout.write.side_effect = RuntimeError("sink down") + with patch("deep_agent.src.audit.emitter.enqueue") as mock_enqueue: + emit_audit_event("llm_call", model="gemini") + mock_enqueue.assert_called_once() + envelope = mock_enqueue.call_args.args[0] + assert envelope["audit_event_type"] == "llm_call" diff --git a/tests/unit/audit/test_emitter_scrub.py b/tests/unit/audit/test_emitter_scrub.py new file mode 100644 index 00000000..2ac49cf8 --- /dev/null +++ b/tests/unit/audit/test_emitter_scrub.py @@ -0,0 +1,22 @@ +"""Unit tests for emitter sensitive key scrubbing.""" + +from unittest.mock import patch + +from deep_agent.src.audit.emitter import _scrub_details + + +class TestScrubDetails: + def test_redacts_sensitive_keys(self): + scrubbed = _scrub_details({"access_token": "secret", "model": "gemini"}) + assert scrubbed["access_token"] == "[REDACTED]" + assert scrubbed["model"] == "gemini" + + def test_does_not_redact_author_field(self): + scrubbed = _scrub_details({"author": "alice", "authorization": "Bearer x"}) + assert scrubbed["author"] == "alice" + assert scrubbed["authorization"] == "[REDACTED]" + + def test_redacts_nested_sensitive_keys(self): + scrubbed = _scrub_details({"meta": {"api_key": "k", "count": 1}}) + assert scrubbed["meta"]["api_key"] == "[REDACTED]" + assert scrubbed["meta"]["count"] == 1 diff --git a/tests/unit/audit/test_integration.py b/tests/unit/audit/test_integration.py new file mode 100644 index 00000000..56fcd363 --- /dev/null +++ b/tests/unit/audit/test_integration.py @@ -0,0 +1,57 @@ +"""Unit tests for platform audit middleware builder integration.""" + +from unittest.mock import MagicMock, patch + +from deep_agent.src.agent.config.middleware import ResolvedMiddlewareConfig +from deep_agent.src.infrastructure.middleware import build_middleware_list +from deep_agent.src.audit.middleware import AuditMiddleware + + +class TestBuildMiddlewareListAudit: + def test_includes_audit_middleware_when_enabled(self): + resolved = ResolvedMiddlewareConfig(summarization_tool_enabled=False) + with ( + patch("deep_agent.src.infrastructure.middleware.settings") as mock_settings, + patch( + "deep_agent.src.audit.config.is_audit_enabled", + return_value=True, + ), + ): + mock_settings.MIDDLEWARE_ENABLED = True + result = build_middleware_list( + resolved, mcp_tool_names=frozenset({"tool_a"}) + ) + assert isinstance(result[0], AuditMiddleware) + assert result[0]._mcp_tool_names == frozenset({"tool_a"}) + + def test_no_audit_middleware_when_disabled(self): + resolved = ResolvedMiddlewareConfig(summarization_tool_enabled=False) + with ( + patch("deep_agent.src.infrastructure.middleware.settings") as mock_settings, + patch( + "deep_agent.src.audit.config.is_audit_enabled", + return_value=False, + ), + ): + mock_settings.MIDDLEWARE_ENABLED = True + result = build_middleware_list(resolved) + assert not any(isinstance(m, AuditMiddleware) for m in result) + + def test_audit_middleware_when_master_middleware_disabled(self): + resolved = ResolvedMiddlewareConfig(summarization_tool_enabled=True) + mock_mw = MagicMock() + with ( + patch("deep_agent.src.infrastructure.middleware.settings") as mock_settings, + patch( + "deep_agent.src.audit.config.is_audit_enabled", + return_value=True, + ), + patch( + "deep_agent.src.infrastructure.middleware._build_summarization_tool_middleware", + return_value=mock_mw, + ), + ): + mock_settings.MIDDLEWARE_ENABLED = False + result = build_middleware_list(resolved) + assert isinstance(result[0], AuditMiddleware) + assert result == [result[0]] diff --git a/tests/unit/audit/test_middleware.py b/tests/unit/audit/test_middleware.py new file mode 100644 index 00000000..48eed009 --- /dev/null +++ b/tests/unit/audit/test_middleware.py @@ -0,0 +1,147 @@ +"""Unit tests for AuditMiddleware classification.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from deep_agent.src.audit.middleware import ( + AuditMiddleware, + classify_tool_call, + _is_memory_write, +) + + +class TestMemoryWriteDetection: + @pytest.mark.parametrize( + ("tool", "args", "expected"), + [ + ("edit_file", {"path": "/memories/notes.md"}, True), + ("write_file", {"file_path": "memories/foo.txt"}, True), + ("edit_file", {"path": "/reports/out.md"}, False), + ("search_web", {"query": "test"}, False), + ], + ) + def test_is_memory_write(self, tool, args, expected): + assert _is_memory_write(tool, args) is expected + + +class TestClassifyToolCallParity: + """Orchestrator and subagent use the same classification rules.""" + + @pytest.mark.parametrize( + ("tool", "args", "mcp_names", "expected"), + [ + ("task", {"subagent": "researcher"}, frozenset(), "subagent_delegation"), + ( + "gitlab_search", + {"q": "x"}, + frozenset({"gitlab_search"}), + "mcp_tool_call", + ), + ("edit_file", {"path": "/memories/x.md"}, frozenset(), "memory_write"), + ("calculate_bmi", {}, frozenset(), ""), + ], + ) + def test_shared_rules(self, tool, args, mcp_names, expected): + assert classify_tool_call(tool, args, mcp_tool_names=mcp_names) == expected + + +class TestAuditMiddlewareClassification: + def test_sync_llm_call_with_subagent(self): + mw = AuditMiddleware(agent="researcher") + request = MagicMock() + request.model = "gemini-2.5-flash" + request.messages = [] + handler = MagicMock(return_value=MagicMock()) + + with patch( + "deep_agent.src.audit.middleware.is_audit_enabled", return_value=True + ): + with patch("deep_agent.src.audit.middleware.emit_audit_event") as emit: + mw.wrap_model_call(request, handler) + assert emit.call_count == 2 + assert emit.call_args_list[0].args[0] == "llm_call" + assert emit.call_args_list[0].kwargs["agent"] == "researcher" + assert emit.call_args_list[0].kwargs["phase"] == "start" + + def test_orchestrator_llm_includes_agent(self): + mw = AuditMiddleware() + request = MagicMock() + request.model = "gemini-2.5-flash" + request.messages = [] + handler = MagicMock(return_value=MagicMock()) + + with patch( + "deep_agent.src.audit.middleware.is_audit_enabled", return_value=True + ): + with patch("deep_agent.src.audit.middleware.emit_audit_event") as emit: + mw.wrap_model_call(request, handler) + assert emit.call_args_list[0].kwargs["agent"] == "orchestrator" + + @pytest.mark.asyncio + async def test_subagent_delegation(self): + mw = AuditMiddleware(mcp_tool_names=frozenset()) + request = MagicMock() + request.tool_call = { + "name": "task", + "args": {"subagent": "researcher"}, + "id": "1", + } + handler = AsyncMock(return_value=MagicMock()) + + with patch( + "deep_agent.src.audit.middleware.is_audit_enabled", return_value=True + ): + with patch("deep_agent.src.audit.middleware.emit_audit_event") as emit: + await mw.awrap_tool_call(request, handler) + emit.assert_called_once() + assert emit.call_args.args[0] == "subagent_delegation" + assert emit.call_args.kwargs["delegated_subagent"] == "researcher" + assert emit.call_args.kwargs["agent"] == "orchestrator" + + @pytest.mark.asyncio + async def test_mcp_tool_call_on_subagent(self): + mw = AuditMiddleware( + mcp_tool_names=frozenset({"gitlab_search"}), + agent="researcher", + ) + request = MagicMock() + request.tool_call = {"name": "gitlab_search", "args": {"q": "x"}, "id": "1"} + handler = AsyncMock(return_value=MagicMock()) + + with patch( + "deep_agent.src.audit.middleware.is_audit_enabled", return_value=True + ): + with patch("deep_agent.src.audit.middleware.emit_audit_event") as emit: + await mw.awrap_tool_call(request, handler) + emit.assert_called_once() + assert emit.call_args.args[0] == "mcp_tool_call" + assert emit.call_args.kwargs["agent"] == "researcher" + + @pytest.mark.asyncio + async def test_skips_unclassified_tools_on_orchestrator(self): + mw = AuditMiddleware(mcp_tool_names=frozenset()) + request = MagicMock() + request.tool_call = {"name": "calculate_bmi", "args": {}, "id": "1"} + handler = AsyncMock(return_value=MagicMock()) + + with patch( + "deep_agent.src.audit.middleware.is_audit_enabled", return_value=True + ): + with patch("deep_agent.src.audit.middleware.emit_audit_event") as emit: + await mw.awrap_tool_call(request, handler) + emit.assert_not_called() + + @pytest.mark.asyncio + async def test_skips_unclassified_tools_on_subagent(self): + mw = AuditMiddleware(mcp_tool_names=frozenset(), agent="researcher") + request = MagicMock() + request.tool_call = {"name": "calculate_bmi", "args": {}, "id": "1"} + handler = AsyncMock(return_value=MagicMock()) + + with patch( + "deep_agent.src.audit.middleware.is_audit_enabled", return_value=True + ): + with patch("deep_agent.src.audit.middleware.emit_audit_event") as emit: + await mw.awrap_tool_call(request, handler) + emit.assert_not_called() diff --git a/tests/unit/audit/test_subagent_middleware.py b/tests/unit/audit/test_subagent_middleware.py new file mode 100644 index 00000000..4b771b8a --- /dev/null +++ b/tests/unit/audit/test_subagent_middleware.py @@ -0,0 +1,41 @@ +"""Unit tests for audit middleware on subagents.""" + +from unittest.mock import MagicMock, patch + +from deep_agent.src.infrastructure.subagents import _subagent_middleware +from deep_agent.src.audit.middleware import AuditMiddleware + + +class TestSubagentMiddleware: + def test_includes_audit_when_enabled(self): + tool = MagicMock() + tool.name = "mcp_search" + audit_mw = AuditMiddleware( + mcp_tool_names=frozenset({"mcp_search"}), + agent="researcher", + ) + with patch( + "deep_agent.src.infrastructure.subagents.build_audit_middleware", + return_value=audit_mw, + ): + result = _subagent_middleware("researcher", [tool], []) + assert result is not None + assert isinstance(result[0], AuditMiddleware) + assert result[0]._agent == "researcher" + assert "mcp_search" in result[0]._mcp_tool_names + + def test_returns_none_when_audit_disabled_and_no_fallback(self): + with patch( + "deep_agent.src.infrastructure.subagents.build_audit_middleware", + return_value=None, + ): + assert _subagent_middleware("researcher", [], []) is None + + def test_fallback_only_when_audit_disabled(self): + fallback = MagicMock() + with patch( + "deep_agent.src.infrastructure.subagents.build_audit_middleware", + return_value=None, + ): + result = _subagent_middleware("researcher", [], [fallback]) + assert result == [fallback] diff --git a/tests/unit/cache/__init__.py b/tests/unit/cache/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/cache/test_backend.py b/tests/unit/cache/test_backend.py new file mode 100644 index 00000000..b5c3727d --- /dev/null +++ b/tests/unit/cache/test_backend.py @@ -0,0 +1,123 @@ +"""Unit tests for cache backend implementations.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.cache.backend import ( + CacheBackend, + InMemoryCache, + NullCache, + RedisCache, +) + + +class TestNullCache: + def test_get_always_none(self): + c = NullCache() + assert c.get("any-key") is None + + def test_set_always_false(self): + c = NullCache() + assert c.set("k", "v") is False + + def test_delete_always_false(self): + c = NullCache() + assert c.delete("k") is False + + def test_clear_is_noop(self): + NullCache().clear() + + def test_name(self): + assert NullCache().name == "null" + + def test_implements_protocol(self): + assert isinstance(NullCache(), CacheBackend) + + +class TestInMemoryCache: + def test_set_and_get(self): + c = InMemoryCache(max_size=10, default_ttl=60) + c.set("k1", "v1") + assert c.get("k1") == "v1" + + def test_get_miss(self): + c = InMemoryCache() + assert c.get("missing") is None + + def test_delete_existing(self): + c = InMemoryCache() + c.set("k1", "v1") + assert c.delete("k1") is True + assert c.get("k1") is None + + def test_delete_missing(self): + c = InMemoryCache() + assert c.delete("nope") is False + + def test_clear(self): + c = InMemoryCache() + c.set("a", "1") + c.set("b", "2") + c.clear() + assert c.size == 0 + + def test_size(self): + c = InMemoryCache(max_size=10, default_ttl=60) + c.set("a", "1") + c.set("b", "2") + assert c.size == 2 + + def test_name(self): + assert InMemoryCache().name == "memory" + + def test_implements_protocol(self): + assert isinstance(InMemoryCache(), CacheBackend) + + +class TestRedisCache: + def test_name(self): + assert RedisCache().name == "redis" + + def test_get_returns_none_when_no_client(self): + c = RedisCache() + with patch.object(c, "_get_client", return_value=None): + assert c.get("key") is None + + def test_set_returns_false_when_no_client(self): + c = RedisCache() + with patch.object(c, "_get_client", return_value=None): + assert c.set("k", "v") is False + + def test_delete_returns_false_when_no_client(self): + c = RedisCache() + with patch.object(c, "_get_client", return_value=None): + assert c.delete("k") is False + + def test_get_with_client(self): + c = RedisCache(key_prefix="test:") + mock = MagicMock() + mock.get.return_value = "cached" + c._client = mock + c._checked = True + assert c.get("k") == "cached" + mock.get.assert_called_once_with("test:k") + + def test_set_with_client(self): + c = RedisCache(default_ttl=60, key_prefix="test:") + mock = MagicMock() + c._client = mock + c._checked = True + assert c.set("k", "v") is True + mock.setex.assert_called_once_with("test:k", 60, "v") + + def test_get_handles_exception(self): + c = RedisCache() + mock = MagicMock() + mock.get.side_effect = Exception("redis down") + c._client = mock + c._checked = True + assert c.get("k") is None + + def test_clear_is_noop(self): + RedisCache().clear() diff --git a/tests/unit/cache/test_config.py b/tests/unit/cache/test_config.py new file mode 100644 index 00000000..0a7d2d00 --- /dev/null +++ b/tests/unit/cache/test_config.py @@ -0,0 +1,39 @@ +"""Unit tests for cache configuration.""" + +from deep_agent.src.cache.config import CacheSettings + + +class TestCacheSettings: + def test_defaults_all_disabled(self): + s = CacheSettings( + CACHE_ENABLED=False, + CACHE_MODEL_ENABLED=False, + CACHE_PERSONALIZATION_ENABLED=False, + CACHE_METRICS_ENABLED=False, + CACHE_WARMING_ENABLED=False, + CACHE_REDIS_ENABLED=False, + ) + assert s.CACHE_ENABLED is False + assert s.CACHE_MODEL_ENABLED is False + assert s.CACHE_PERSONALIZATION_ENABLED is False + + def test_is_enabled_requires_master_switch(self): + s = CacheSettings(CACHE_ENABLED=False, CACHE_MODEL_ENABLED=True) + assert s.is_enabled("model") is False + + def test_is_enabled_with_master_on(self): + s = CacheSettings(CACHE_ENABLED=True, CACHE_MODEL_ENABLED=True) + assert s.is_enabled("model") is True + + def test_is_enabled_unknown_layer(self): + s = CacheSettings(CACHE_ENABLED=True) + assert s.is_enabled("nonexistent") is False + + def test_ttl_defaults(self): + s = CacheSettings() + assert s.CACHE_MODEL_TTL == 600 + assert s.CACHE_PERSONALIZATION_TTL == 120 + + def test_max_size_defaults(self): + s = CacheSettings() + assert s.CACHE_MODEL_MAX_SIZE == 10 diff --git a/tests/unit/cache/test_metrics.py b/tests/unit/cache/test_metrics.py new file mode 100644 index 00000000..c93f9083 --- /dev/null +++ b/tests/unit/cache/test_metrics.py @@ -0,0 +1,59 @@ +"""Unit tests for cache metrics.""" + +from unittest.mock import patch + +from deep_agent.src.cache import metrics +from deep_agent.src.cache.config import CacheSettings + + +class TestCacheMetrics: + def setup_method(self): + metrics.reset() + + def test_record_and_snapshot(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_METRICS_ENABLED=True) + with patch.object(metrics, "cache_settings", enabled): + metrics.record_hit("test") + metrics.record_hit("test") + metrics.record_miss("test") + metrics.record_set("test") + metrics.record_delete("test") + + snap = metrics.snapshot() + assert snap["test"]["hits"] == 2 + assert snap["test"]["misses"] == 1 + assert snap["test"]["sets"] == 1 + assert snap["test"]["deletes"] == 1 + + def test_disabled_does_not_record(self): + disabled = CacheSettings(CACHE_ENABLED=False) + with patch.object(metrics, "cache_settings", disabled): + metrics.record_hit("test") + assert metrics.snapshot() == {} + + def test_reset_clears_all(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_METRICS_ENABLED=True) + with patch.object(metrics, "cache_settings", enabled): + metrics.record_hit("test") + metrics.reset() + assert metrics.snapshot() == {} + + def test_get_stats(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_METRICS_ENABLED=True) + with patch.object(metrics, "cache_settings", enabled): + metrics.record_hit("x") + metrics.record_miss("x") + stats = metrics.get_stats() + assert stats["x"]["total"] == 2 + assert stats["x"]["hit_rate"] == 50.0 + + def test_log_summary_does_not_raise(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_METRICS_ENABLED=True) + with patch.object(metrics, "cache_settings", enabled): + metrics.record_hit("x") + metrics.log_summary() + + def test_log_summary_skips_when_disabled(self): + disabled = CacheSettings(CACHE_ENABLED=False) + with patch.object(metrics, "cache_settings", disabled): + metrics.log_summary() diff --git a/tests/unit/cache/test_model_cache.py b/tests/unit/cache/test_model_cache.py new file mode 100644 index 00000000..07086d60 --- /dev/null +++ b/tests/unit/cache/test_model_cache.py @@ -0,0 +1,149 @@ +"""Unit tests for model cache.""" + +from unittest.mock import MagicMock, patch + +from deep_agent.src.agent.config.model import ModelSpec, Provider +from deep_agent.src.cache import model_cache +from deep_agent.src.cache.config import CacheSettings + + +class TestModelCache: + def setup_method(self): + model_cache._legacy_cache = None + model_cache._spec_cache = None + + def test_passthrough_when_disabled(self): + disabled = CacheSettings(CACHE_ENABLED=False) + mock_model = MagicMock() + + with ( + patch.object(model_cache, "cache_settings", disabled), + patch("deep_agent.src.agent.llm.create_model", return_value=mock_model), + ): + result = model_cache.get_or_create_model("gemini-2.5-pro") + assert result is mock_model + + def test_cache_hit(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_MODEL_ENABLED=True) + mock_model = MagicMock() + + with ( + patch.object(model_cache, "cache_settings", enabled), + patch( + "deep_agent.src.agent.llm.create_model", return_value=mock_model + ) as create, + ): + m1 = model_cache.get_or_create_model("gemini-2.5-pro", 0.0, 8192) + m2 = model_cache.get_or_create_model("gemini-2.5-pro", 0.0, 8192) + + assert m1 is mock_model + assert m2 is mock_model + assert create.call_count == 1 + + def test_different_params_different_entries(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_MODEL_ENABLED=True) + model_a = MagicMock(name="model_a") + model_b = MagicMock(name="model_b") + + with ( + patch.object(model_cache, "cache_settings", enabled), + patch( + "deep_agent.src.agent.llm.create_model", side_effect=[model_a, model_b] + ), + ): + r1 = model_cache.get_or_create_model("gemini-2.5-pro", 0.0, 8192) + r2 = model_cache.get_or_create_model("gemini-2.5-pro", 0.5, 8192) + + assert r1 is model_a + assert r2 is model_b + + def test_invalidate_all(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_MODEL_ENABLED=True) + with patch.object(model_cache, "cache_settings", enabled): + model_cache._get_cache()[("test", 0.0, 8192)] = MagicMock() + assert model_cache.cached_count() == 1 + + model_cache.invalidate() + assert model_cache.cached_count() == 0 + + def test_invalidate_by_name(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_MODEL_ENABLED=True) + with patch.object(model_cache, "cache_settings", enabled): + cache = model_cache._get_cache() + cache[("gemini", 0.0, 8192)] = MagicMock() + cache[("claude", 0.0, 8192)] = MagicMock() + assert model_cache.cached_count() == 2 + + model_cache.invalidate("gemini") + assert model_cache.cached_count() == 1 + + +class TestModelCacheFromSpec: + """Tests for provider-aware spec cache.""" + + def setup_method(self): + model_cache._legacy_cache = None + model_cache._spec_cache = None + + def test_spec_cache_hit(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_MODEL_ENABLED=True) + mock_model = MagicMock() + spec = ModelSpec(provider=Provider.VERTEX, name="gemini-2.5-pro") + + with ( + patch.object(model_cache, "cache_settings", enabled), + patch( + "deep_agent.src.agent.provider_factory.create_model_from_spec", + return_value=mock_model, + ) as create, + ): + m1 = model_cache.get_or_create_model_from_spec(spec) + m2 = model_cache.get_or_create_model_from_spec(spec) + + assert m1 is mock_model + assert m2 is mock_model + assert create.call_count == 1 + + def test_different_providers_same_name_different_entries(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_MODEL_ENABLED=True) + vertex_model = MagicMock(name="vertex_model") + openai_model = MagicMock(name="openai_model") + vertex_spec = ModelSpec(provider=Provider.VERTEX, name="shared-name") + openai_spec = ModelSpec(provider=Provider.OPENAI, name="shared-name") + + with ( + patch.object(model_cache, "cache_settings", enabled), + patch( + "deep_agent.src.agent.provider_factory.create_model_from_spec", + side_effect=[vertex_model, openai_model], + ), + ): + r1 = model_cache.get_or_create_model_from_spec(vertex_spec) + r2 = model_cache.get_or_create_model_from_spec(openai_spec) + + assert r1 is vertex_model + assert r2 is openai_model + + def test_fallback_changes_cache_key(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_MODEL_ENABLED=True) + model_a = MagicMock(name="model_a") + model_b = MagicMock(name="model_b") + spec_no_fb = ModelSpec(provider=Provider.VERTEX, name="gemini-2.5-pro") + spec_with_fb = ModelSpec( + provider=Provider.VERTEX, + name="gemini-2.5-pro", + fallback=ModelSpec(provider=Provider.OPENAI, name="gpt-4o-mini"), + ) + + with ( + patch.object(model_cache, "cache_settings", enabled), + patch( + "deep_agent.src.agent.provider_factory.create_model_from_spec", + side_effect=[model_a, model_b], + ), + ): + r1 = model_cache.get_or_create_model_from_spec(spec_no_fb) + r2 = model_cache.get_or_create_model_from_spec(spec_with_fb) + + assert r1 is model_a + assert r2 is model_b diff --git a/tests/unit/cache/test_multi_layer.py b/tests/unit/cache/test_multi_layer.py new file mode 100644 index 00000000..6375f48e --- /dev/null +++ b/tests/unit/cache/test_multi_layer.py @@ -0,0 +1,74 @@ +"""Unit tests for multi-layer cache.""" + +from unittest.mock import patch + +from deep_agent.src.cache.backend import InMemoryCache, NullCache +from deep_agent.src.cache.multi_layer import MultiLayerCache, create_null_layer + + +class TestMultiLayerCache: + def test_l1_hit(self): + l1 = InMemoryCache(max_size=10, default_ttl=60) + l1.set("k", "v1") + ml = MultiLayerCache("test", l1=l1) + assert ml.get("k") == "v1" + + def test_l2_hit_backfills_l1(self): + l1 = InMemoryCache(max_size=10, default_ttl=60) + l2 = InMemoryCache(max_size=10, default_ttl=60) + l2.set("k", "from-l2") + ml = MultiLayerCache("test", l1=l1, l2=l2) + + assert ml.get("k") == "from-l2" + assert l1.get("k") == "from-l2" + + def test_miss_both_layers(self): + l1 = InMemoryCache(max_size=10, default_ttl=60) + l2 = InMemoryCache(max_size=10, default_ttl=60) + ml = MultiLayerCache("test", l1=l1, l2=l2) + assert ml.get("missing") is None + + def test_set_writes_both(self): + l1 = InMemoryCache(max_size=10, default_ttl=60) + l2 = InMemoryCache(max_size=10, default_ttl=60) + ml = MultiLayerCache("test", l1=l1, l2=l2) + + ml.set("k", "v") + assert l1.get("k") == "v" + assert l2.get("k") == "v" + + def test_delete_both(self): + l1 = InMemoryCache(max_size=10, default_ttl=60) + l2 = InMemoryCache(max_size=10, default_ttl=60) + ml = MultiLayerCache("test", l1=l1, l2=l2) + + ml.set("k", "v") + ml.delete("k") + assert l1.get("k") is None + assert l2.get("k") is None + + def test_clear_only_l1(self): + l1 = InMemoryCache(max_size=10, default_ttl=60) + l2 = InMemoryCache(max_size=10, default_ttl=60) + ml = MultiLayerCache("test", l1=l1, l2=l2) + ml.set("k", "v") + ml.clear() + assert l1.get("k") is None + assert l2.get("k") == "v" + + def test_name(self): + ml = MultiLayerCache("my-cache", l1=NullCache()) + assert ml.name == "my-cache" + + def test_no_l2(self): + l1 = InMemoryCache(max_size=10, default_ttl=60) + ml = MultiLayerCache("test", l1=l1, l2=None) + ml.set("k", "v") + assert ml.get("k") == "v" + + +class TestCreateNullLayer: + def test_returns_noop(self): + ml = create_null_layer("disabled") + assert ml.get("k") is None + assert ml.set("k", "v") is False diff --git a/tests/unit/cache/test_personalization_cache.py b/tests/unit/cache/test_personalization_cache.py new file mode 100644 index 00000000..94468798 --- /dev/null +++ b/tests/unit/cache/test_personalization_cache.py @@ -0,0 +1,80 @@ +"""Unit tests for personalization cache.""" + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.cache import personalization_cache +from deep_agent.src.cache.config import CacheSettings + + +class TestPersonalizationCache: + def setup_method(self): + personalization_cache._redis = None + + async def test_get_returns_none_when_disabled(self): + disabled = CacheSettings(CACHE_ENABLED=False) + with patch.object(personalization_cache, "cache_settings", disabled): + result = await personalization_cache.get_personalization("user-1") + assert result is None + + async def test_set_is_noop_when_disabled(self): + disabled = CacheSettings(CACHE_ENABLED=False) + with patch.object(personalization_cache, "cache_settings", disabled): + await personalization_cache.set_personalization("user-1", [], []) + + async def test_cache_roundtrip(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_PERSONALIZATION_ENABLED=True) + mock_redis = MagicMock() + store: dict[str, str] = {} + + def fake_get(key: str) -> str | None: + return store.get(key) + + def fake_set(key: str, value: str, ttl: int | None = None) -> bool: + store[key] = value + return True + + mock_redis.get = fake_get + mock_redis.set = fake_set + + with ( + patch.object(personalization_cache, "cache_settings", enabled), + patch.object(personalization_cache, "_get_redis", return_value=mock_redis), + ): + memories = [{"content": "likes pizza"}] + rules = [{"content": "be brief"}] + await personalization_cache.set_personalization("user-1", memories, rules) + + result = await personalization_cache.get_personalization("user-1") + assert result is not None + assert result[0] == memories + assert result[1] == rules + + async def test_get_handles_corrupt_data(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_PERSONALIZATION_ENABLED=True) + mock_redis = MagicMock() + mock_redis.get.return_value = "not-valid-json{{" + mock_redis.delete.return_value = True + + with ( + patch.object(personalization_cache, "cache_settings", enabled), + patch.object(personalization_cache, "_get_redis", return_value=mock_redis), + ): + result = await personalization_cache.get_personalization("user-1") + assert result is None + + async def test_invalidate(self): + enabled = CacheSettings(CACHE_ENABLED=True, CACHE_PERSONALIZATION_ENABLED=True) + mock_redis = MagicMock() + + with ( + patch.object(personalization_cache, "cache_settings", enabled), + patch.object(personalization_cache, "_get_redis", return_value=mock_redis), + ): + await personalization_cache.invalidate("user-1") + mock_redis.delete.assert_called_once() + + async def test_invalidate_none_is_noop(self): + await personalization_cache.invalidate(None) diff --git a/tests/unit/cache/test_warming.py b/tests/unit/cache/test_warming.py new file mode 100644 index 00000000..d1ad04d1 --- /dev/null +++ b/tests/unit/cache/test_warming.py @@ -0,0 +1,99 @@ +"""Unit tests for cache warming.""" + +from unittest.mock import MagicMock, patch + +from deep_agent.src.cache import warming +from deep_agent.src.cache.config import CacheSettings + + +class TestWarmCaches: + def test_skips_when_disabled(self): + disabled = CacheSettings(CACHE_ENABLED=False) + with patch.object(warming, "cache_settings", disabled): + result = warming.warm_caches() + assert result == {} + + def test_warms_models_when_enabled(self): + enabled = CacheSettings( + CACHE_ENABLED=True, + CACHE_WARMING_ENABLED=True, + CACHE_MODEL_ENABLED=True, + ) + mock_config = MagicMock() + mock_config.get_orchestrator_config.return_value = {"model": "gemini-2.5-flash"} + mock_config.get_all_subagent_configs.return_value = { + "sub1": { + "model": {"provider": "vertex", "name": "gemini-2.5-pro"}, + }, + } + + with ( + patch.object(warming, "cache_settings", enabled), + patch("deep_agent.src.agent.config.agent_config", mock_config), + patch( + "deep_agent.src.cache.model_cache.get_or_create_model_from_spec" + ) as mock_from_spec, + ): + result = warming.warm_caches() + assert result["models"] is True + # Both orchestrator and subagent use get_or_create_model_from_spec + assert mock_from_spec.call_count == 2 + + def test_handles_model_warming_failure(self): + enabled = CacheSettings( + CACHE_ENABLED=True, + CACHE_WARMING_ENABLED=True, + CACHE_MODEL_ENABLED=True, + ) + mock_config = MagicMock() + mock_config.get_orchestrator_config.side_effect = Exception("boom") + + with ( + patch.object(warming, "cache_settings", enabled), + patch("deep_agent.src.agent.config.agent_config", mock_config), + ): + result = warming.warm_caches() + assert result["models"] is False + + def test_skips_models_when_model_cache_disabled(self): + enabled = CacheSettings( + CACHE_ENABLED=True, + CACHE_WARMING_ENABLED=True, + CACHE_MODEL_ENABLED=False, + ) + with patch.object(warming, "cache_settings", enabled): + result = warming.warm_caches() + assert result["models"] is False + + def test_parses_orchestrator_model_with_provider(self): + """Orchestrator models support provider specification.""" + enabled = CacheSettings( + CACHE_ENABLED=True, + CACHE_WARMING_ENABLED=True, + CACHE_MODEL_ENABLED=True, + ) + mock_config = MagicMock() + # Orchestrator with explicit provider + mock_config.get_orchestrator_config.return_value = { + "model": { + "provider": "vertex", + "name": "gemini-2.5-pro", + } + } + mock_config.get_all_subagent_configs.return_value = {} + + with ( + patch.object(warming, "cache_settings", enabled), + patch("deep_agent.src.agent.config.agent_config", mock_config), + patch( + "deep_agent.src.cache.model_cache.get_or_create_model_from_spec" + ) as mock_from_spec, + ): + result = warming.warm_caches() + assert result["models"] is True + assert mock_from_spec.call_count == 1 + + # Verify the spec has correct provider + spec = mock_from_spec.call_args[0][0] + assert spec.name == "gemini-2.5-pro" + assert spec.provider.value == "vertex" diff --git a/tests/unit/code_execution/__init__.py b/tests/unit/code_execution/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/code_execution/test_config.py b/tests/unit/code_execution/test_config.py new file mode 100644 index 00000000..d36bd32e --- /dev/null +++ b/tests/unit/code_execution/test_config.py @@ -0,0 +1,127 @@ +"""Tests for CodeExecutionConfig.""" + +from __future__ import annotations + +import pytest + + +class TestCodeExecutionConfigDefaults: + def test_defaults(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + cfg = CodeExecutionConfig() + assert cfg.enabled is False + assert cfg.max_timeout_seconds == 60 + assert cfg.max_code_length == 50_000 + assert cfg.max_output_bytes == 1_048_576 + assert "python" in cfg.images + assert "shell" in cfg.images + assert "node" in cfg.images + assert cfg.images["python"] == "python:3.12-slim" + assert cfg.entrypoints["python"] == ["python", "-c"] + assert cfg.entrypoints["shell"] == ["bash", "-c"] + assert cfg.entrypoints["node"] == ["node", "-e"] + assert cfg.resource_requests == {"cpu": "100m", "memory": "128Mi"} + assert cfg.resource_limits == {"cpu": "500m", "memory": "256Mi"} + assert cfg.tmp_size_limit == "64Mi" + assert cfg.job_ttl_after_finished == 30 + assert cfg.pod_poll_interval_seconds == 1.0 + assert cfg.pod_poll_timeout_seconds == 120.0 + + def test_from_dict(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + cfg = CodeExecutionConfig.model_validate( + { + "enabled": True, + "max_timeout_seconds": 120, + "images": {"python": "my-registry/python:3.12"}, + } + ) + assert cfg.enabled is True + assert cfg.max_timeout_seconds == 120 + assert cfg.images["python"] == "my-registry/python:3.12" + + def test_supported_languages(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + cfg = CodeExecutionConfig() + assert cfg.supported_languages == {"python", "shell", "node"} + + def test_new_feature_defaults(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + cfg = CodeExecutionConfig() + assert cfg.network_access == "deny" + assert cfg.max_concurrent_per_org == 3 + assert cfg.queue_timeout_seconds == 30.0 + assert cfg.max_input_file_size == 1_048_576 + assert cfg.cost_tracking_enabled is False + assert cfg.streaming_enabled is False + + def test_custom_image_variant_via_config(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + cfg = CodeExecutionConfig.model_validate( + { + "images": { + "python": "python:3.12-slim", + "python-ds": "my-registry/python-ds:3.12", + "shell": "bash:5", + "node": "node:22-slim", + }, + "entrypoints": { + "python": ["python", "-c"], + "python-ds": ["python", "-c"], + "shell": ["bash", "-c"], + "node": ["node", "-e"], + }, + } + ) + assert "python-ds" in cfg.supported_languages + assert cfg.images["python-ds"] == "my-registry/python-ds:3.12" + + +class TestCodeExecutionConfigValidation: + def test_timeout_min(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + with pytest.raises(Exception): + CodeExecutionConfig(max_timeout_seconds=4) + + def test_timeout_max(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + with pytest.raises(Exception): + CodeExecutionConfig(max_timeout_seconds=301) + + def test_poll_interval_min(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + with pytest.raises(Exception): + CodeExecutionConfig(pod_poll_interval_seconds=0.1) + + def test_concurrent_limit_min(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + with pytest.raises(Exception): + CodeExecutionConfig(max_concurrent_per_org=0) + + def test_queue_timeout_min(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + with pytest.raises(Exception): + CodeExecutionConfig(queue_timeout_seconds=0.5) + + def test_network_access_values(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + for val in ("deny", "allow_internet", "per_execution"): + cfg = CodeExecutionConfig(network_access=val) + assert cfg.network_access == val + + def test_network_access_invalid(self): + from deep_agent.src.code_execution.config import CodeExecutionConfig + + with pytest.raises(Exception): + CodeExecutionConfig(network_access="invalid") diff --git a/tests/unit/code_execution/test_k8s_job_runner.py b/tests/unit/code_execution/test_k8s_job_runner.py new file mode 100644 index 00000000..9762f64a --- /dev/null +++ b/tests/unit/code_execution/test_k8s_job_runner.py @@ -0,0 +1,332 @@ +"""Tests for K8sJobRunner.""" + +from __future__ import annotations + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from deep_agent.src.code_execution.config import CodeExecutionConfig + + +class TestExecutionResult: + def test_format_success(self): + from deep_agent.src.code_execution.k8s_job_runner import ExecutionResult + + result = ExecutionResult( + stdout="hello world", + stderr="", + exit_code=0, + duration_seconds=1.5, + status="success", + job_name="code-exec-abc123", + namespace="ap-test-org-test-agent", + ) + formatted = result.format() + assert "stdout:" in formatted + assert "hello world" in formatted + assert "exit_code: 0" in formatted + + def test_format_timeout(self): + from deep_agent.src.code_execution.k8s_job_runner import ExecutionResult + + result = ExecutionResult( + stdout="", + stderr="", + exit_code=-1, + duration_seconds=60.0, + status="timeout", + job_name="j", + namespace="ns", + ) + formatted = result.format() + assert "timed out" in formatted + + def test_format_oom(self): + from deep_agent.src.code_execution.k8s_job_runner import ExecutionResult + + result = ExecutionResult( + stdout="", + stderr="", + exit_code=137, + duration_seconds=5.0, + status="oom_killed", + job_name="j", + namespace="ns", + ) + formatted = result.format() + assert "out of memory" in formatted + + def test_format_no_stdout(self): + from deep_agent.src.code_execution.k8s_job_runner import ExecutionResult + + result = ExecutionResult( + stdout="", + stderr="some error", + exit_code=1, + duration_seconds=0.5, + status="failed", + job_name="j", + namespace="ns", + ) + formatted = result.format() + assert "stdout:" not in formatted + assert "stderr:" in formatted + assert "some error" in formatted + + +class TestJobManifest: + def test_manifest_security_fields(self): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + manifest = runner.build_job_manifest( + language="python", + code="print('hello')", + timeout=60, + namespace="ap-test-org-test-agent", + execution_id="test-uuid", + trace_id="trace-abc", + ) + spec = manifest.spec + pod_spec = spec.template.spec + + assert spec.backoff_limit == 0 + assert spec.active_deadline_seconds == 60 + assert spec.ttl_seconds_after_finished == 30 + assert pod_spec.restart_policy == "Never" + assert pod_spec.automount_service_account_token is False + assert pod_spec.security_context.run_as_non_root is True + assert pod_spec.security_context.run_as_user == 1000 + + container = pod_spec.containers[0] + assert container.image == "python:3.12-slim" + assert container.command == ["python", "-c"] + assert container.args == ["print('hello')"] + assert container.security_context.allow_privilege_escalation is False + assert container.security_context.read_only_root_filesystem is True + + def test_manifest_labels(self): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + manifest = runner.build_job_manifest( + language="python", + code="x=1", + timeout=30, + namespace="ap-myorg-myagent", + execution_id="exec-123", + trace_id="tr-456", + ) + labels = manifest.metadata.labels + assert labels["app.kubernetes.io/managed-by"] == "template-agent" + assert labels["ai-platform.io/execution-id"] == "exec-123" + + @pytest.mark.parametrize( + "language,expected_image,expected_cmd", + [ + ("python", "python:3.12-slim", ["python", "-c"]), + ("shell", "bash:5", ["bash", "-c"]), + ("node", "node:22-slim", ["node", "-e"]), + ], + ) + def test_manifest_language_mapping(self, language, expected_image, expected_cmd): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + manifest = runner.build_job_manifest( + language=language, + code="code", + timeout=60, + namespace="ns", + execution_id="id", + trace_id="tr", + ) + container = manifest.spec.template.spec.containers[0] + assert container.image == expected_image + assert container.command == expected_cmd + + +class TestParseContainerStatus: + @pytest.mark.parametrize( + "reason,expected_status", + [ + ("OOMKilled", "oom_killed"), + ("Error", "failed"), + ("DeadlineExceeded", "timeout"), + (None, "failed"), + ], + ) + def test_termination_reason(self, reason, expected_status): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + exit_code, status = runner.parse_container_status( + exit_code=1, termination_reason=reason + ) + assert status == expected_status + + def test_success(self): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + exit_code, status = runner.parse_container_status( + exit_code=0, termination_reason=None + ) + assert status == "success" + assert exit_code == 0 + + def test_success_with_completed_reason(self): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + exit_code, status = runner.parse_container_status( + exit_code=0, termination_reason="Completed" + ) + assert status == "success" + assert exit_code == 0 + + +class TestManifestNetworkLabel: + def test_allow_network_label(self): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + manifest = runner.build_job_manifest( + language="python", + code="x=1", + timeout=30, + namespace="ns", + execution_id="id", + allow_network=True, + ) + labels = manifest.metadata.labels + assert labels["ai-platform.io/allow-internet"] == "true" + + def test_no_network_label_by_default(self): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + manifest = runner.build_job_manifest( + language="python", + code="x=1", + timeout=30, + namespace="ns", + execution_id="id", + ) + labels = manifest.metadata.labels + assert "ai-platform.io/allow-internet" not in labels + + +class TestManifestInputConfigmap: + def test_configmap_volume_mount(self): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + manifest = runner.build_job_manifest( + language="python", + code="x=1", + timeout=30, + namespace="ns", + execution_id="id", + input_configmap_name="code-exec-input-abc12345", + ) + pod_spec = manifest.spec.template.spec + volume_names = [v.name for v in pod_spec.volumes] + mount_paths = [vm.mount_path for vm in pod_spec.containers[0].volume_mounts] + assert "input" in volume_names + assert "/input" in mount_paths + assert "/output" in mount_paths + + def test_no_input_volume_without_configmap(self): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + manifest = runner.build_job_manifest( + language="python", + code="x=1", + timeout=30, + namespace="ns", + execution_id="id", + ) + pod_spec = manifest.spec.template.spec + volume_names = [v.name for v in pod_spec.volumes] + assert "input" not in volume_names + assert "output" in volume_names + + +class TestOutputVolume: + def test_output_emptydir_always_present(self): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + config = CodeExecutionConfig() + runner = K8sJobRunner(config) + manifest = runner.build_job_manifest( + language="python", + code="x=1", + timeout=30, + namespace="ns", + execution_id="id", + ) + pod_spec = manifest.spec.template.spec + volume_names = [v.name for v in pod_spec.volumes] + mount_paths = [vm.mount_path for vm in pod_spec.containers[0].volume_mounts] + assert "output" in volume_names + assert "/output" in mount_paths + + +class TestCpuMemoryParsing: + @pytest.mark.parametrize( + "value,expected", + [ + ("100m", 100_000_000), + ("1", 1_000_000_000), + ("250m", 250_000_000), + ("500n", 500), + ], + ) + def test_parse_cpu(self, value, expected): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + assert K8sJobRunner._parse_cpu(value) == expected + + @pytest.mark.parametrize( + "value,expected", + [ + ("128Mi", 128 * 1024 * 1024), + ("1Gi", 1024 * 1024 * 1024), + ("256Ki", 256 * 1024), + ("1000M", 1000 * 1000 * 1000), + ], + ) + def test_parse_memory(self, value, expected): + from deep_agent.src.code_execution.k8s_job_runner import K8sJobRunner + + assert K8sJobRunner._parse_memory(value) == expected + + +class TestExecutionResultWithCost: + def test_default_cost_fields(self): + from deep_agent.src.code_execution.k8s_job_runner import ExecutionResult + + result = ExecutionResult( + stdout="ok", + stderr="", + exit_code=0, + duration_seconds=1.0, + status="success", + job_name="j", + namespace="ns", + ) + assert result.cpu_seconds == 0.0 + assert result.memory_mb_seconds == 0.0 + assert result.output_files == {} diff --git a/tests/unit/code_execution/test_metrics.py b/tests/unit/code_execution/test_metrics.py new file mode 100644 index 00000000..ef17e6b8 --- /dev/null +++ b/tests/unit/code_execution/test_metrics.py @@ -0,0 +1,165 @@ +"""Tests for CodeExecutionMetrics.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + + +class TestCodeExecutionMetrics: + @patch("deep_agent.src.code_execution.metrics.emit_audit_event") + @patch("deep_agent.src.code_execution.metrics._get_tracer") + def test_emit_audit(self, mock_tracer, mock_emit): + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + metrics = CodeExecutionMetrics() + metrics.emit_audit( + language="python", + status="success", + exit_code=0, + latency_ms=1234.5, + code_hash="sha256:abc", + namespace="ap-test-agent", + image="python:3.12-slim", + job_name="code-exec-abc", + timeout=60, + stdout_bytes=100, + stderr_bytes=0, + ) + mock_emit.assert_called_once() + call_args = mock_emit.call_args + assert call_args[0][0] == "code_execution" + assert call_args[1]["language"] == "python" + assert call_args[1]["status"] == "success" + assert call_args[1]["code_hash"] == "sha256:abc" + + @patch("deep_agent.src.code_execution.metrics._get_tracer") + def test_start_span(self, mock_get_tracer): + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + mock_tracer = MagicMock() + mock_get_tracer.return_value = mock_tracer + + metrics = CodeExecutionMetrics() + metrics.start_span("code_execution", language="python") + mock_tracer.start_span.assert_called_once() + + @patch("deep_agent.src.code_execution.metrics._get_tracer") + def test_trace_span_context_manager(self, mock_get_tracer): + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + mock_tracer = MagicMock() + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + mock_get_tracer.return_value = mock_tracer + + metrics = CodeExecutionMetrics() + with metrics.trace_span("code_execution", language="python") as span: + assert span is mock_span + mock_span.end.assert_called_once() + + @patch("deep_agent.src.code_execution.metrics._get_tracer") + def test_trace_span_none_when_no_tracer(self, mock_get_tracer): + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + mock_get_tracer.return_value = None + metrics = CodeExecutionMetrics() + with metrics.trace_span("test") as span: + assert span is None + + def test_log_started(self): + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + metrics = CodeExecutionMetrics() + metrics.log_started(language="python", job_name="j", namespace="ns") + + def test_log_completed(self): + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + metrics = CodeExecutionMetrics() + metrics.log_completed(exit_code=0, duration_ms=1234, status="success") + + def test_compute_code_hash(self): + from deep_agent.src.code_execution.metrics import compute_code_hash + + h = compute_code_hash("print('hello')") + assert h.startswith("sha256:") + assert len(h) > 10 + + def test_compute_code_hash_deterministic(self): + from deep_agent.src.code_execution.metrics import compute_code_hash + + h1 = compute_code_hash("x = 1") + h2 = compute_code_hash("x = 1") + assert h1 == h2 + + def test_compute_code_hash_different_for_different_code(self): + from deep_agent.src.code_execution.metrics import compute_code_hash + + h1 = compute_code_hash("x = 1") + h2 = compute_code_hash("x = 2") + assert h1 != h2 + + @patch("deep_agent.src.code_execution.metrics._get_otel_metrics") + def test_record_execution_with_otel(self, mock_get_metrics): + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + mock_mc = MagicMock() + mock_get_metrics.return_value = mock_mc + + metrics = CodeExecutionMetrics() + metrics.record_execution( + language="python", org="test", exit_code=0, status="success", duration=2.5 + ) + mock_mc.code_execution_duration_seconds.record.assert_called_once() + mock_mc.code_executions_total.add.assert_called_once() + + @patch("deep_agent.src.code_execution.metrics._get_otel_metrics") + def test_record_execution_without_otel(self, mock_get_metrics): + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + mock_get_metrics.return_value = None + metrics = CodeExecutionMetrics() + metrics.record_execution( + language="python", org="test", exit_code=0, status="success", duration=2.5 + ) + + def test_record_scheduling_latency(self): + from deep_agent.src.code_execution.metrics import CodeExecutionMetrics + + metrics = CodeExecutionMetrics() + metrics.record_scheduling_latency(org="test-org", duration=1.5) + + +class TestAuditEmitter: + def test_emit_audit_event_writes_to_stdout(self, capsys): + from deep_agent.src.audit.emitter import emit_audit_event + + emit_audit_event("code_execution", language="python", status="success") + captured = capsys.readouterr() + assert "platform.audit" in captured.out + assert "code_execution" in captured.out + assert "python" in captured.out + + def test_scrub_sensitive_keys(self, capsys): + from deep_agent.src.audit.emitter import emit_audit_event + + emit_audit_event("test", password="secret123", api_key="abc") + captured = capsys.readouterr() + assert "secret123" not in captured.out + assert "[REDACTED]" in captured.out + + def test_audit_context_binding(self): + from deep_agent.src.audit.context import ( + bind_audit_context, + clear_audit_context, + get_audit_context, + ) + + bind_audit_context(user="alice@test.com", org="test-org", trace_id="tr-123") + ctx = get_audit_context() + assert ctx["user"] == "alice@test.com" + assert ctx["org"] == "test-org" + assert ctx["trace_id"] == "tr-123" + clear_audit_context() + ctx = get_audit_context() + assert ctx["user"] is None diff --git a/tests/unit/code_execution/test_middleware.py b/tests/unit/code_execution/test_middleware.py new file mode 100644 index 00000000..68d18f9f --- /dev/null +++ b/tests/unit/code_execution/test_middleware.py @@ -0,0 +1,237 @@ +"""Tests for CodeExecutionMiddleware.""" + +from __future__ import annotations + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from deep_agent.src.code_execution.config import CodeExecutionConfig + + +class TestToolInjection: + async def test_awrap_model_call_injects_tool(self): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + config = CodeExecutionConfig(enabled=True) + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + request.tools = [MagicMock()] + + override_request = MagicMock() + request.override = MagicMock(return_value=override_request) + handler = AsyncMock(return_value=MagicMock()) + + await mw.awrap_model_call(request, handler) + + request.override.assert_called_once() + handler.assert_called_once_with(override_request) + + async def test_awrap_model_call_disabled(self): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + config = CodeExecutionConfig(enabled=False) + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + handler = AsyncMock(return_value=MagicMock()) + + await mw.awrap_model_call(request, handler) + handler.assert_called_once_with(request) + request.override.assert_not_called() + + +class TestToolCallRouting: + async def test_passthrough_non_execute_code(self): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + config = CodeExecutionConfig(enabled=True) + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + request.tool_call = {"name": "search_web", "args": {}, "id": "tc1"} + handler = AsyncMock(return_value=MagicMock()) + + await mw.awrap_tool_call(request, handler) + handler.assert_called_once_with(request) + + async def test_invalid_language(self): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + config = CodeExecutionConfig(enabled=True) + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + request.tool_call = { + "name": "execute_code", + "args": {"code": "x=1", "language": "ruby"}, + "id": "tc1", + } + handler = AsyncMock() + + result = await mw.awrap_tool_call(request, handler) + handler.assert_not_called() + assert "Unsupported language" in result.content + assert "ruby" in result.content + + async def test_code_too_long(self): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + config = CodeExecutionConfig(enabled=True, max_code_length=100) + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + request.tool_call = { + "name": "execute_code", + "args": {"code": "x" * 101, "language": "python"}, + "id": "tc1", + } + handler = AsyncMock() + + result = await mw.awrap_tool_call(request, handler) + handler.assert_not_called() + assert "exceeds maximum length" in result.content + + @patch("deep_agent.src.code_execution.middleware.CodeExecutionMetrics") + @patch("deep_agent.src.code_execution.middleware.K8sJobRunner") + async def test_successful_execution(self, mock_runner_cls, mock_metrics_cls): + from deep_agent.src.code_execution.k8s_job_runner import ExecutionResult + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + mock_runner = MagicMock() + mock_runner.run = AsyncMock( + return_value=ExecutionResult( + stdout="42", + stderr="", + exit_code=0, + duration_seconds=1.5, + status="success", + job_name="code-exec-abc", + namespace="ap-test-agent", + ) + ) + mock_runner.resolve_namespace = MagicMock(return_value="ap-default-agent") + mock_runner_cls.return_value = mock_runner + + mock_metrics = MagicMock() + mock_metrics_cls.return_value = mock_metrics + + config = CodeExecutionConfig(enabled=True) + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + request.tool_call = { + "name": "execute_code", + "args": {"code": "print(42)", "language": "python"}, + "id": "tc1", + } + handler = AsyncMock() + + result = await mw.awrap_tool_call(request, handler) + handler.assert_not_called() + assert "42" in result.content + assert "exit_code: 0" in result.content + assert result.tool_call_id == "tc1" + + +class TestSyncPassthrough: + def test_wrap_model_call_passes_through(self): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + config = CodeExecutionConfig(enabled=True) + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + handler = MagicMock(return_value=MagicMock()) + + mw.wrap_model_call(request, handler) + handler.assert_called_once_with(request) + + def test_wrap_tool_call_passes_through(self): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + config = CodeExecutionConfig(enabled=True) + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + handler = MagicMock(return_value=MagicMock()) + + mw.wrap_tool_call(request, handler) + handler.assert_called_once_with(request) + + +class TestInputFileValidation: + async def test_input_files_too_large(self): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + config = CodeExecutionConfig(enabled=True, max_input_file_size=10) + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + request.tool_call = { + "name": "execute_code", + "args": { + "code": "print(1)", + "language": "python", + "input_files": {"big.csv": "x" * 100}, + }, + "id": "tc1", + } + handler = AsyncMock() + + result = await mw.awrap_tool_call(request, handler) + assert "exceed" in result.content.lower() + handler.assert_not_called() + + +class TestNetworkDenyOverride: + async def test_network_denied_when_config_deny(self): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + config = CodeExecutionConfig(enabled=True, network_access="deny") + mw = CodeExecutionMiddleware(config=config) + + request = MagicMock() + request.tool_call = { + "name": "execute_code", + "args": {"code": "print(1)", "language": "python", "network": True}, + "id": "tc1", + } + + # The middleware should override network=True to False when config is deny + # We can't easily test the runner call without mocking, but the validation passes + # (network param is silently overridden, not rejected) + + +class TestQueueing: + @patch("deep_agent.src.code_execution.middleware.CodeExecutionMetrics") + @patch("deep_agent.src.code_execution.middleware.K8sJobRunner") + async def test_semaphore_created_per_org(self, mock_runner_cls, mock_metrics_cls): + from deep_agent.src.code_execution.middleware import CodeExecutionMiddleware + + mock_metrics_cls.return_value = MagicMock() + mock_runner = MagicMock() + mock_runner.resolve_namespace = MagicMock(return_value="ap-default-agent") + mock_runner.run = AsyncMock( + return_value=MagicMock( + stdout="ok", + stderr="", + exit_code=0, + duration_seconds=1.0, + status="success", + job_name="j", + namespace="ns", + cpu_seconds=0.0, + memory_mb_seconds=0.0, + format=MagicMock(return_value="stdout:\nok\nexit_code: 0"), + ) + ) + mock_runner_cls.return_value = mock_runner + + config = CodeExecutionConfig(enabled=True, max_concurrent_per_org=2) + mw = CodeExecutionMiddleware(config=config) + + assert len(mw._semaphores) == 0 + sem = mw._get_semaphore("org-a") + assert "org-a" in mw._semaphores + assert sem._value == 2 diff --git a/tests/unit/config/__init__.py b/tests/unit/config/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/config/test_filesystem_config.py b/tests/unit/config/test_filesystem_config.py new file mode 100644 index 00000000..b633f236 --- /dev/null +++ b/tests/unit/config/test_filesystem_config.py @@ -0,0 +1,154 @@ +"""Unit tests for filesystem configuration and permissions builder.""" + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.agent.config.filesystem import ( + BackendConfig, + FilesystemFileConfig, + FilesystemSettings, + LocalShellConfig, + PermissionRule, + StateConfig, + load_filesystem_config, +) +from deep_agent.src.infrastructure.permissions import build_permissions + + +class TestFilesystemModels: + """Test Pydantic model defaults.""" + + def test_default_backend_is_state(self): + config = FilesystemFileConfig() + assert config.backend.type == "state" + + def test_default_permissions_empty(self): + config = FilesystemFileConfig() + assert config.permissions == [] + + def test_default_settings(self): + settings = FilesystemSettings() + assert settings.tool_token_limit_before_evict == 20_000 + assert settings.human_message_token_limit_before_evict == 50_000 + assert settings.max_execute_timeout == 3600 + + def test_local_shell_defaults(self): + ls = LocalShellConfig() + assert ls.timeout == 120 + assert ls.max_output_bytes == 100_000 + + def test_state_default_disabled(self): + state = StateConfig() + assert state.enabled is False + + def test_permission_rule_defaults_to_allow(self): + rule = PermissionRule(operations=["read"], paths=["**"]) + assert rule.mode == "allow" + + +class TestLoadFilesystemConfig: + """Test loading filesystem.yaml from disk.""" + + def test_returns_defaults_when_missing(self, tmp_path): + config = load_filesystem_config(tmp_path / "nope.yaml") + assert config.backend.type == "state" + assert config.permissions == [] + + def test_loads_valid_yaml(self, tmp_path): + content = """ +backend: + type: composite + local_shell: + timeout: 60 + max_output_bytes: 50000 + routes: + "/scratch/": state + "/": local_shell + +permissions: + - operations: [read, glob] + paths: ["config/**"] + mode: allow + - operations: [write] + paths: ["**/*.py"] + mode: deny + +settings: + tool_token_limit_before_evict: 10000 + max_execute_timeout: 1800 +""" + config_file = tmp_path / "filesystem.yaml" + config_file.write_text(content) + + config = load_filesystem_config(config_file) + assert config.backend.type == "composite" + assert config.backend.local_shell.timeout == 60 + assert config.backend.routes == {"/scratch/": "state", "/": "local_shell"} + assert len(config.permissions) == 2 + assert config.permissions[0].operations == ["read", "glob"] + assert config.permissions[1].mode == "deny" + assert config.settings.tool_token_limit_before_evict == 10_000 + + def test_returns_defaults_on_invalid_yaml(self, tmp_path): + config_file = tmp_path / "filesystem.yaml" + config_file.write_text("{{invalid") + config = load_filesystem_config(config_file) + assert config.backend.type == "state" + + +class TestBuildPermissions: + """Test FilesystemPermission construction from config.""" + + def test_returns_none_when_no_rules(self): + config = FilesystemFileConfig() + result = build_permissions(config) + assert result is None + + def test_builds_permission_objects(self): + config = FilesystemFileConfig( + permissions=[ + PermissionRule( + operations=["read", "glob", "grep"], + paths=["config/**"], + mode="allow", + ), + PermissionRule( + operations=["write", "edit"], + paths=["**/*.py"], + mode="deny", + ), + ] + ) + + mock_perm = MagicMock() + with patch( + "deep_agent.src.infrastructure.permissions.FilesystemPermission", + return_value=mock_perm, + ) as mock_cls: + result = build_permissions(config) + + assert result is not None + assert len(result) == 2 + assert mock_cls.call_count == 2 + mock_cls.assert_any_call( + operations=["read", "glob", "grep"], + paths=["config/**"], + mode="allow", + ) + + def test_skips_invalid_rules_gracefully(self): + config = FilesystemFileConfig( + permissions=[ + PermissionRule(operations=["read"], paths=["ok/**"], mode="allow"), + ] + ) + + with patch( + "deep_agent.src.infrastructure.permissions.FilesystemPermission", + side_effect=ValueError("bad rule"), + ): + result = build_permissions(config) + + assert result is None diff --git a/tests/unit/config/test_middleware_config.py b/tests/unit/config/test_middleware_config.py new file mode 100644 index 00000000..82d9bc58 --- /dev/null +++ b/tests/unit/config/test_middleware_config.py @@ -0,0 +1,147 @@ +"""Unit tests for middleware configuration resolution.""" + +from pathlib import Path +from unittest.mock import patch + +import pytest + +from deep_agent.src.agent.config.middleware import ( + MemoryConfig, + MiddlewareDefaults, + MiddlewareFileConfig, + PatchToolCallsConfig, + ProfileConfig, + ResolvedMiddlewareConfig, + SkillsConfig, + SummarizationToolConfig, + load_middleware_config, + resolve_middleware, +) + + +class TestMiddlewareModels: + """Test Pydantic model defaults and validation.""" + + def test_defaults_all_enabled(self): + defaults = MiddlewareDefaults() + assert defaults.summarization_tool.enabled is True + assert defaults.memory.enabled is True + assert defaults.patch_tool_calls.enabled is True + assert defaults.skills.enabled is True + assert defaults.extra == [] + + def test_memory_default_namespaces(self): + config = MemoryConfig() + assert config.namespaces == ["memories"] + + def test_profile_defaults_empty(self): + profile = ProfileConfig() + assert profile.excluded_middleware == [] + assert profile.excluded_tools == [] + assert profile.system_prompt_suffix == "" + + def test_file_config_defaults(self): + config = MiddlewareFileConfig() + assert config.defaults.summarization_tool.enabled is True + assert config.profiles == {} + + +class TestLoadMiddlewareConfig: + """Test loading middleware.yaml from disk.""" + + def test_returns_defaults_when_file_missing(self, tmp_path): + config = load_middleware_config(tmp_path / "nonexistent.yaml") + assert config.defaults.summarization_tool.enabled is True + assert config.profiles == {} + + def test_loads_valid_yaml(self, tmp_path): + yaml_content = """ +defaults: + summarization_tool: + enabled: false + memory: + enabled: true + namespaces: + - user_memories + - shared +profiles: + gemini-2.5-pro: + excluded_middleware: + - patch_tool_calls + system_prompt_suffix: "Be helpful." +""" + config_file = tmp_path / "middleware.yaml" + config_file.write_text(yaml_content) + + config = load_middleware_config(config_file) + assert config.defaults.summarization_tool.enabled is False + assert config.defaults.memory.namespaces == ["user_memories", "shared"] + assert "gemini-2.5-pro" in config.profiles + assert config.profiles["gemini-2.5-pro"].system_prompt_suffix == "Be helpful." + + def test_returns_defaults_on_invalid_yaml(self, tmp_path): + config_file = tmp_path / "middleware.yaml" + config_file.write_text("not: [valid: yaml: {{") + + config = load_middleware_config(config_file) + assert config.defaults.summarization_tool.enabled is True + + +class TestResolveMiddleware: + """Test the resolution logic: defaults → profile → overrides.""" + + def test_all_defaults_no_profile_no_overrides(self): + config = MiddlewareFileConfig() + resolved = resolve_middleware(config, "unknown-model") + + assert resolved.summarization_tool_enabled is True + assert resolved.memory_enabled is True + assert resolved.patch_tool_calls_enabled is True + assert resolved.skills_enabled is True + assert resolved.memory_namespaces == ["memories"] + + def test_profile_excludes_patch_tool_calls(self): + config = MiddlewareFileConfig( + profiles={ + "claude-sonnet": ProfileConfig(excluded_middleware=["patch_tool_calls"]) + } + ) + resolved = resolve_middleware(config, "claude-sonnet") + assert resolved.patch_tool_calls_enabled is False + + def test_agent_override_disables_memory(self): + config = MiddlewareFileConfig() + resolved = resolve_middleware(config, "gemini-2.5-pro", {"memory": False}) + assert resolved.memory_enabled is False + + def test_agent_override_dict_with_enabled(self): + config = MiddlewareFileConfig() + overrides = {"summarization_tool": {"enabled": False}} + resolved = resolve_middleware(config, "gemini-2.5-pro", overrides) + assert resolved.summarization_tool_enabled is False + + def test_agent_override_memory_namespaces(self): + config = MiddlewareFileConfig() + overrides = {"memory": {"enabled": True, "namespaces": ["custom_ns"]}} + resolved = resolve_middleware(config, "gemini-2.5-pro", overrides) + assert resolved.memory_enabled is True + assert resolved.memory_namespaces == ["custom_ns"] + + def test_extra_middleware_merged(self): + config = MiddlewareFileConfig( + defaults=MiddlewareDefaults(extra=["module_a:ClassA"]) + ) + overrides = {"extra": ["module_b:ClassB"]} + resolved = resolve_middleware(config, "model", overrides) + assert resolved.extra_middleware == ["module_a:ClassA", "module_b:ClassB"] + + def test_global_disabled_respected(self): + config = MiddlewareFileConfig( + defaults=MiddlewareDefaults( + summarization_tool=SummarizationToolConfig(enabled=False), + memory=MemoryConfig(enabled=False), + ) + ) + resolved = resolve_middleware(config, "model") + assert resolved.summarization_tool_enabled is False + assert resolved.memory_enabled is False diff --git a/tests/unit/config/test_otel_config.py b/tests/unit/config/test_otel_config.py new file mode 100644 index 00000000..ccc6d659 --- /dev/null +++ b/tests/unit/config/test_otel_config.py @@ -0,0 +1,134 @@ +"""Unit tests for OtelFileConfig Pydantic models.""" + +import pytest +from pydantic import ValidationError + +from deep_agent.src.agent.config.otel import ( + OtelExporterConfig, + OtelFileConfig, + OtelMetricsConfig, + OtelTracingConfig, +) + + +class TestOtelExporterConfig: + """Test OtelExporterConfig defaults and validation.""" + + def test_defaults(self): + config = OtelExporterConfig() + assert config.endpoint == "http://localhost:4317" + assert config.insecure is True + + def test_custom_values(self): + config = OtelExporterConfig( + endpoint="https://collector.prod:4317", + insecure=False, + ) + assert config.endpoint == "https://collector.prod:4317" + assert config.insecure is False + + def test_from_dict(self): + config = OtelExporterConfig.model_validate( + {"endpoint": "http://otel:4317", "insecure": False} + ) + assert config.endpoint == "http://otel:4317" + assert config.insecure is False + + +class TestOtelMetricsConfig: + """Test OtelMetricsConfig defaults and validation.""" + + def test_default_interval(self): + config = OtelMetricsConfig() + assert config.export_interval_ms == 5000 + + def test_custom_interval(self): + config = OtelMetricsConfig(export_interval_ms=10000) + assert config.export_interval_ms == 10000 + + def test_minimum_interval_boundary(self): + config = OtelMetricsConfig(export_interval_ms=1000) + assert config.export_interval_ms == 1000 + + def test_maximum_interval_boundary(self): + config = OtelMetricsConfig(export_interval_ms=60000) + assert config.export_interval_ms == 60000 + + def test_rejects_interval_below_minimum(self): + with pytest.raises(ValidationError, match="greater than or equal to 1000"): + OtelMetricsConfig(export_interval_ms=999) + + def test_rejects_interval_above_maximum(self): + with pytest.raises(ValidationError, match="less than or equal to 60000"): + OtelMetricsConfig(export_interval_ms=60001) + + +class TestOtelTracingConfig: + """Test OtelTracingConfig defaults.""" + + def test_auto_instrument_default_true(self): + config = OtelTracingConfig() + assert config.fastapi_auto_instrument is True + + def test_disable_auto_instrument(self): + config = OtelTracingConfig(fastapi_auto_instrument=False) + assert config.fastapi_auto_instrument is False + + +class TestOtelFileConfig: + """Test top-level OtelFileConfig model.""" + + def test_defaults(self): + config = OtelFileConfig() + assert config.enabled is False + assert config.exporter.endpoint == "http://localhost:4317" + assert config.exporter.insecure is True + assert config.metrics.export_interval_ms == 5000 + assert config.tracing.fastapi_auto_instrument is True + + def test_enabled_flag(self): + config = OtelFileConfig(enabled=True) + assert config.enabled is True + + def test_from_dict(self): + """Parse from a dict matching the YAML structure.""" + config = OtelFileConfig.model_validate( + { + "enabled": True, + "exporter": { + "endpoint": "http://collector:4317", + "insecure": False, + }, + "metrics": { + "export_interval_ms": 15000, + }, + "tracing": { + "fastapi_auto_instrument": False, + }, + } + ) + assert config.enabled is True + assert config.exporter.endpoint == "http://collector:4317" + assert config.exporter.insecure is False + assert config.metrics.export_interval_ms == 15000 + assert config.tracing.fastapi_auto_instrument is False + + def test_from_empty_dict(self): + """Empty dict should produce all defaults (matches observability.yaml loading).""" + config = OtelFileConfig.model_validate({}) + assert config.enabled is False + assert config.exporter.endpoint == "http://localhost:4317" + assert config.metrics.export_interval_ms == 5000 + assert config.tracing.fastapi_auto_instrument is True + + def test_partial_dict(self): + """Partial dict should fill in defaults for missing fields.""" + config = OtelFileConfig.model_validate({"enabled": True}) + assert config.enabled is True + assert config.exporter.endpoint == "http://localhost:4317" + assert config.metrics.export_interval_ms == 5000 + + def test_nested_validation_propagates(self): + """Invalid nested config should raise ValidationError.""" + with pytest.raises(ValidationError): + OtelFileConfig.model_validate({"metrics": {"export_interval_ms": 500}}) diff --git a/tests/unit/config/test_providers_config.py b/tests/unit/config/test_providers_config.py new file mode 100644 index 00000000..31fb6ce9 --- /dev/null +++ b/tests/unit/config/test_providers_config.py @@ -0,0 +1,237 @@ +"""Unit tests for providers configuration and profile registration.""" + +from pathlib import Path +from unittest.mock import MagicMock, call, patch + +import pytest + +from deep_agent.src.agent.config.providers import ( + AsyncTaskConfig, + GeneralPurposeSubagentConfig, + HarnessProfileConfig, + ProviderConfig, + ProvidersFileConfig, + load_providers_config, +) +from deep_agent.src.infrastructure.async_tasks import ( + _extract_async_subagents, + build_async_middleware, +) +from deep_agent.src.infrastructure.providers import ( + _register_harness_profiles, + _register_provider_profiles, + resolve_model_from_config, +) + + +class TestProviderModels: + """Test Pydantic model defaults.""" + + def test_default_strategy_is_legacy(self): + config = ProvidersFileConfig() + assert config.resolve_strategy == "legacy" + + def test_default_async_tasks_enabled(self): + config = ProvidersFileConfig() + assert config.async_tasks.enabled is True + assert config.async_tasks.system_prompt is None + + def test_default_general_purpose_subagent(self): + gp = GeneralPurposeSubagentConfig() + assert gp.enabled is True + assert gp.description is None + assert gp.system_prompt is None + + def test_harness_profile_defaults(self): + hp = HarnessProfileConfig() + assert hp.system_prompt_suffix == "" + assert hp.excluded_tools == [] + assert hp.excluded_middleware == [] + assert hp.general_purpose_subagent.enabled is True + + def test_provider_config_defaults(self): + pc = ProviderConfig() + assert pc.init_kwargs == {} + + +class TestLoadProvidersConfig: + """Test loading providers.yaml from disk.""" + + def test_returns_defaults_when_missing(self, tmp_path): + config = load_providers_config(tmp_path / "nope.yaml") + assert config.resolve_strategy == "legacy" + assert config.providers == {} + assert config.harness_profiles == {} + + def test_loads_valid_yaml(self, tmp_path): + content = """ +resolve_strategy: deepagents + +providers: + google_genai: + init_kwargs: + temperature: 0.0 + openai: + init_kwargs: + api_key: test + +harness_profiles: + gemini-2.5-pro: + system_prompt_suffix: "Think step by step." + excluded_tools: [execute] + general_purpose_subagent: + enabled: false + +async_tasks: + enabled: false + system_prompt: "Custom async prompt" +""" + config_file = tmp_path / "providers.yaml" + config_file.write_text(content) + + config = load_providers_config(config_file) + assert config.resolve_strategy == "deepagents" + assert len(config.providers) == 2 + assert config.providers["openai"].init_kwargs == {"api_key": "test"} + assert len(config.harness_profiles) == 1 + hp = config.harness_profiles["gemini-2.5-pro"] + assert hp.system_prompt_suffix == "Think step by step." + assert hp.excluded_tools == ["execute"] + assert hp.general_purpose_subagent.enabled is False + assert config.async_tasks.enabled is False + assert config.async_tasks.system_prompt == "Custom async prompt" + + def test_returns_defaults_on_invalid_yaml(self, tmp_path): + config_file = tmp_path / "providers.yaml" + config_file.write_text("{{invalid") + config = load_providers_config(config_file) + assert config.resolve_strategy == "legacy" + + +class TestResolveModel: + """Test model resolution dispatch.""" + + def test_legacy_strategy_uses_cache(self): + config = ProvidersFileConfig(resolve_strategy="legacy") + with patch( + "deep_agent.src.infrastructure.providers.get_or_create_model", + return_value="mock_model", + ) as mock: + result = resolve_model_from_config("gemini-2.5-pro", config) + assert result == "mock_model" + mock.assert_called_once_with( + model_name="gemini-2.5-pro", + temperature=0.0, + max_output_tokens=None, + ) + + def test_deepagents_strategy_calls_resolve_model(self): + config = ProvidersFileConfig(resolve_strategy="deepagents") + with patch( + "deep_agent.src.infrastructure.providers.resolve_model", + return_value="da_model", + ): + result = resolve_model_from_config("openai:gpt-5.4", config) + assert result == "da_model" + + +class TestRegisterProfiles: + """Test profile registration functions.""" + + def test_register_provider_profiles(self): + config = ProvidersFileConfig( + providers={ + "google_genai": ProviderConfig(init_kwargs={"temperature": 0.0}), + } + ) + mock_profile_cls = MagicMock() + mock_register = MagicMock() + with patch.dict( + "sys.modules", + { + "deepagents": MagicMock( + ProviderProfile=mock_profile_cls, + register_provider_profile=mock_register, + ) + }, + ): + _register_provider_profiles(config) + mock_register.assert_called_once() + + def test_register_harness_profiles(self): + config = ProvidersFileConfig( + harness_profiles={ + "gemini-2.5-pro": HarnessProfileConfig( + system_prompt_suffix="Think.", + excluded_tools=["execute"], + general_purpose_subagent=GeneralPurposeSubagentConfig( + enabled=False + ), + ), + } + ) + mock_hp_cls = MagicMock() + mock_gp_cls = MagicMock() + mock_register = MagicMock() + with patch.dict( + "sys.modules", + { + "deepagents": MagicMock( + HarnessProfile=mock_hp_cls, + GeneralPurposeSubagentProfile=mock_gp_cls, + register_harness_profile=mock_register, + ) + }, + ): + _register_harness_profiles(config) + mock_register.assert_called_once() + mock_gp_cls.assert_called_once_with( + enabled=False, description=None, system_prompt=None + ) + + +class TestAsyncMiddleware: + """Test async middleware builder.""" + + def test_returns_none_when_disabled(self): + config = AsyncTaskConfig(enabled=False) + result = build_async_middleware([MagicMock()], config) + assert result is None + + def test_returns_none_when_no_subagents(self): + config = AsyncTaskConfig(enabled=True) + result = build_async_middleware(None, config) + assert result is None + + def test_returns_none_when_no_async_subagents(self): + config = AsyncTaskConfig(enabled=True) + regular_sub = MagicMock(spec=[]) + with patch( + "deep_agent.src.infrastructure.async_tasks._extract_async_subagents", + return_value=[], + ): + result = build_async_middleware([regular_sub], config) + assert result is None + + def test_builds_middleware_for_async_subagents(self): + config = AsyncTaskConfig(enabled=True, system_prompt="Custom prompt") + async_sub = MagicMock() + mock_mw = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.async_tasks._extract_async_subagents", + return_value=[async_sub], + ), + patch( + "deep_agent.src.infrastructure.async_tasks.AsyncSubAgentMiddleware", + return_value=mock_mw, + ) as mock_cls, + ): + result = build_async_middleware([async_sub], config) + + assert result is mock_mw + mock_cls.assert_called_once_with( + async_subagents=[async_sub], + system_prompt="Custom prompt", + ) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 00000000..8a9a51f2 --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,8 @@ +"""Auto-apply ``@pytest.mark.unit`` to every test in tests/unit/.""" + +import pytest + + +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + for item in items: + item.add_marker(pytest.mark.unit) diff --git a/tests/unit/feedback/__init__.py b/tests/unit/feedback/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/feedback/test_repository.py b/tests/unit/feedback/test_repository.py new file mode 100644 index 00000000..d1657eaf --- /dev/null +++ b/tests/unit/feedback/test_repository.py @@ -0,0 +1,150 @@ +"""Unit tests for FeedbackRepository (mocked DB).""" + +from unittest.mock import AsyncMock, patch + +import pytest + +from deep_agent.src.feedback import repository as feedback_repo_mod +from deep_agent.src.feedback.repository import FeedbackRepository + + +@pytest.fixture(autouse=True) +def _reset_feedback_table_flag(): + """Reset module-level _TABLE_ENSURED before each test.""" + feedback_repo_mod._TABLE_ENSURED = False + yield + feedback_repo_mod._TABLE_ENSURED = False + + +@pytest.fixture +def mock_conn(): + """Create a mock async connection context manager.""" + conn = AsyncMock() + cursor = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[]) + cursor.rowcount = 0 + conn.execute = AsyncMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=False) + conn._cursor = cursor + return conn + + +@pytest.fixture +def repo(): + return FeedbackRepository("postgresql://test:test@localhost/testdb") + + +class TestEnsureTable: + @pytest.mark.asyncio + async def test_creates_table_once(self, repo, mock_conn): + with patch( + "deep_agent.src.feedback.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + await repo.ensure_table() + mock_conn.execute.assert_awaited_once() + mock_conn.commit.assert_awaited_once() + + @pytest.mark.asyncio + async def test_idempotent_second_call(self, repo, mock_conn): + with patch( + "deep_agent.src.feedback.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + await repo.ensure_table() + mock_conn.execute.reset_mock() + mock_conn.commit.reset_mock() + await repo.ensure_table() + mock_conn.execute.assert_not_called() + mock_conn.commit.assert_not_called() + + +class TestUpsertFeedback: + @pytest.mark.asyncio + async def test_insert_calls_execute_and_commit(self, repo, mock_conn): + feedback_repo_mod._TABLE_ENSURED = True + with patch( + "deep_agent.src.feedback.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + await repo.upsert_feedback( + "t1", + "m1", + "u1", + "up", + "trace-1", + ) + mock_conn.execute.assert_awaited_once() + mock_conn.commit.assert_awaited_once() + + @pytest.mark.asyncio + async def test_update_second_upsert(self, repo, mock_conn): + """Second upsert with same keys runs ON CONFLICT UPDATE (still one execute).""" + feedback_repo_mod._TABLE_ENSURED = True + with patch( + "deep_agent.src.feedback.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + await repo.upsert_feedback("t1", "m1", "u1", "up", None) + await repo.upsert_feedback("t1", "m1", "u1", "down", None) + assert mock_conn.execute.await_count == 2 + assert mock_conn.commit.await_count == 2 + + +class TestDeleteFeedback: + @pytest.mark.asyncio + async def test_delete_returns_true_when_row_removed(self, repo, mock_conn): + feedback_repo_mod._TABLE_ENSURED = True + mock_conn._cursor.rowcount = 1 + mock_conn.execute.return_value = mock_conn._cursor + with patch( + "deep_agent.src.feedback.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + result = await repo.delete_feedback("t1", "m1", "u1") + assert result is True + + @pytest.mark.asyncio + async def test_delete_returns_false_when_missing(self, repo, mock_conn): + feedback_repo_mod._TABLE_ENSURED = True + mock_conn._cursor.rowcount = 0 + mock_conn.execute.return_value = mock_conn._cursor + with patch( + "deep_agent.src.feedback.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + result = await repo.delete_feedback("t1", "m1", "u1") + assert result is False + + +class TestListFeedback: + @pytest.mark.asyncio + async def test_returns_message_id_and_feedback(self, repo, mock_conn): + feedback_repo_mod._TABLE_ENSURED = True + mock_conn._cursor.fetchall = AsyncMock( + return_value=[ + {"message_id": "m1", "feedback": "up"}, + {"message_id": "m2", "feedback": "down"}, + ] + ) + with patch( + "deep_agent.src.feedback.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + rows = await repo.list_feedback("t1", "u1") + assert rows == [ + {"message_id": "m1", "feedback": "up"}, + {"message_id": "m2", "feedback": "down"}, + ] + + @pytest.mark.asyncio + async def test_empty_list(self, repo, mock_conn): + feedback_repo_mod._TABLE_ENSURED = True + with patch( + "deep_agent.src.feedback.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + rows = await repo.list_feedback("t1", "u1") + assert rows == [] diff --git a/tests/unit/infrastructure/test_backend.py b/tests/unit/infrastructure/test_backend.py new file mode 100644 index 00000000..ace2abeb --- /dev/null +++ b/tests/unit/infrastructure/test_backend.py @@ -0,0 +1,39 @@ +"""Unit tests for backend module.""" + +import os +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.infrastructure.backend import ( + _base_python, + _build_env, +) + + +class TestBasePython: + def test_returns_string(self): + result = _base_python() + assert isinstance(result, str) + assert "python" in result.lower() + + +class TestBuildEnv: + def test_contains_virtual_env(self, tmp_path): + env = _build_env(tmp_path) + assert env["VIRTUAL_ENV"] == str(tmp_path) + + def test_contains_path(self, tmp_path): + env = _build_env(tmp_path) + assert str(tmp_path) in env["PATH"] + + def test_extra_env_overrides(self, tmp_path): + env = _build_env(tmp_path, extra={"MY_VAR": "my_val"}) + assert env["MY_VAR"] == "my_val" + + def test_passthrough_vars(self, tmp_path): + with patch.dict(os.environ, {"HOME": "/test/home", "USER": "tester"}): + env = _build_env(tmp_path) + assert env.get("HOME") == "/test/home" + assert env.get("USER") == "tester" diff --git a/tests/unit/infrastructure/test_mcp.py b/tests/unit/infrastructure/test_mcp.py new file mode 100644 index 00000000..0d807a43 --- /dev/null +++ b/tests/unit/infrastructure/test_mcp.py @@ -0,0 +1,443 @@ +"""Unit tests for MCP client utilities.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from deep_agent.aegra.mcp import ( + _build_server_config, + _connect_single_server, + _get_server_configs, + get_mcp_tools, + mcp_httpx_verify, +) + + +class TestGetServerConfigs: + """Tests for _get_server_configs function.""" + + def test_returns_configs_from_agent_config(self): + """Test that _get_server_configs delegates to agent_config.""" + mock_servers = { + "server-a": { + "url": "http://a:5001/mcp/", + "transport": "streamable_http", + "enabled": True, + "auth": True, + "ssl_verify": False, + "timeout": 10, + } + } + + with patch( + "deep_agent.aegra.mcp.agent_config.get_mcp_servers" + ) as mock_get_servers: + mock_get_servers.return_value = mock_servers + + result = _get_server_configs() + + assert result == mock_servers + mock_get_servers.assert_called_once() + + def test_returns_empty_dict_when_no_servers(self): + """Test returns empty dict when no MCP servers configured.""" + with patch( + "deep_agent.aegra.mcp.agent_config.get_mcp_servers" + ) as mock_get_servers: + mock_get_servers.return_value = {} + + result = _get_server_configs() + + assert result == {} + + +class TestMcpHttpxVerify: + """Tests for mcp_httpx_verify helper.""" + + def test_defaults_to_true(self): + assert mcp_httpx_verify({}) is True + + def test_respects_ssl_verify_false(self): + assert mcp_httpx_verify({"ssl_verify": False}) is False + + def test_respects_ssl_verify_true(self): + assert mcp_httpx_verify({"ssl_verify": True}) is True + + +class TestBuildServerConfig: + """Tests for _build_server_config function.""" + + def test_config_without_sso_token(self): + """Test server config without SSO token.""" + entry = { + "url": "http://localhost:8000/mcp/", + "transport": "http", + "auth": True, + "ssl_verify": True, + } + config = _build_server_config(entry, None) + + assert config["url"] == "http://localhost:8000/mcp/" + assert config["transport"] == "http" + assert config["headers"] == {} + assert "httpx_client_factory" not in config + + def test_config_with_sso_token(self): + """Test server config with SSO token.""" + entry = { + "url": "https://api.example.com/mcp/", + "transport": "https", + "auth": True, + "ssl_verify": True, + } + config = _build_server_config(entry, "test_token_123") + + assert config["url"] == "https://api.example.com/mcp/" + assert config["transport"] == "https" + assert config["headers"] == {"Authorization": "Bearer test_token_123"} + assert "httpx_client_factory" not in config + + def test_config_with_ssl_verify_disabled(self): + """Test server config with SSL verification disabled.""" + entry = { + "url": "https://api.example.com/mcp/", + "transport": "https", + "auth": True, + "ssl_verify": False, + } + config = _build_server_config(entry, None) + + assert "httpx_client_factory" in config + assert callable(config["httpx_client_factory"]) + + client = config["httpx_client_factory"]() + assert hasattr(client, "get") + + def test_config_auth_disabled_ignores_token(self): + """Test that auth=False means no Authorization header even with token.""" + entry = { + "url": "http://localhost:8000/mcp/", + "transport": "http", + "auth": False, + "ssl_verify": True, + } + config = _build_server_config(entry, "should_be_ignored") + + assert config["headers"] == {} + + def test_config_defaults(self): + """Test that missing optional fields use sensible defaults.""" + entry = {"url": "http://localhost:8000/mcp/"} + config = _build_server_config(entry, "tok") + + assert config["transport"] == "streamable_http" + assert config["headers"] == {"Authorization": "Bearer tok"} + assert "httpx_client_factory" not in config + + +class TestConnectSingleServer: + """Tests for _connect_single_server function.""" + + @pytest.mark.asyncio + async def test_successful_connection(self): + """Test successful connection to MCP server.""" + mock_tool = MagicMock() + mock_tool.name = "test_tool" + + mock_client = MagicMock() + mock_client.get_tools = AsyncMock(return_value=[mock_tool]) + + config = {"url": "http://localhost:8000/mcp/", "transport": "http"} + + with patch( + "deep_agent.aegra.mcp.MultiServerMCPClient", + return_value=mock_client, + ): + tools = await _connect_single_server("test_server", config, {}, timeout=5) + + assert len(tools) == 1 + assert tools[0].name == "test_tool" + + @pytest.mark.asyncio + async def test_connection_timeout_returns_empty_list(self): + """Test that connection timeout returns empty list.""" + mock_client = MagicMock() + mock_client.get_tools = AsyncMock( + side_effect=TimeoutError("Connection timed out") + ) + + config = {"url": "http://localhost:8000/mcp/", "transport": "http"} + + with patch( + "deep_agent.aegra.mcp.MultiServerMCPClient", + return_value=mock_client, + ): + tools = await _connect_single_server("slow_server", config, {}, timeout=1) + + assert tools == [] + + @pytest.mark.asyncio + async def test_connection_error_returns_empty_list(self): + """Test that connection errors return empty list with fault isolation.""" + mock_client = MagicMock() + mock_client.get_tools = AsyncMock( + side_effect=ConnectionError("Connection refused") + ) + + config = {"url": "http://unreachable:8000/mcp/", "transport": "http"} + + with patch( + "deep_agent.aegra.mcp.MultiServerMCPClient", + return_value=mock_client, + ): + tools = await _connect_single_server("broken_server", config, {}, timeout=5) + + assert tools == [] + + @pytest.mark.asyncio + async def test_generic_exception_returns_empty_list(self): + """Test that any exception returns empty list for fault isolation.""" + mock_client = MagicMock() + mock_client.get_tools = AsyncMock(side_effect=ValueError("Unexpected error")) + + config = {"url": "http://localhost:8000/mcp/", "transport": "http"} + + with patch( + "deep_agent.aegra.mcp.MultiServerMCPClient", + return_value=mock_client, + ): + tools = await _connect_single_server("faulty_server", config, {}, timeout=5) + + assert tools == [] + + +def _reset_mcp_cache() -> None: + """Clear MCP tool cache between tests.""" + from deep_agent.aegra import mcp + + mcp._cached_tools = [] + mcp._cached_tools_ts = 0.0 + + +class TestGetMCPTools: + """Tests for get_mcp_tools function.""" + + @pytest.mark.asyncio + async def test_successful_connection_with_tools(self): + """Test successful MCP connection with tools.""" + _reset_mcp_cache() + mock_servers = { + "test_server": { + "url": "http://localhost:8000/mcp/", + "transport": "http", + "enabled": True, + "auth": False, + "ssl_verify": True, + "timeout": 5, + } + } + + mock_tool = MagicMock() + mock_tool.name = "tool1" + + with ( + patch("deep_agent.aegra.mcp._get_server_configs") as mock_get_configs, + patch("deep_agent.aegra.mcp._connect_single_server") as mock_connect, + ): + mock_get_configs.return_value = mock_servers + mock_connect.return_value = [mock_tool] + + tools = await get_mcp_tools() + + assert len(tools) == 1 + assert tools[0].name == "tool1" + mock_connect.assert_called_once() + + @pytest.mark.asyncio + async def test_deduplicates_tools_from_multiple_servers(self): + """Test that duplicate tool names are deduplicated (first wins).""" + _reset_mcp_cache() + mock_servers = { + "server-a": { + "url": "http://a/mcp/", + "enabled": True, + "auth": False, + "timeout": 5, + }, + "server-b": { + "url": "http://b/mcp/", + "enabled": True, + "auth": False, + "timeout": 5, + }, + } + + tool_a1 = MagicMock() + tool_a1.name = "shared_tool" + tool_a2 = MagicMock() + tool_a2.name = "unique_a" + + tool_b1 = MagicMock() + tool_b1.name = "shared_tool" + tool_b2 = MagicMock() + tool_b2.name = "unique_b" + + with ( + patch("deep_agent.aegra.mcp._get_server_configs") as mock_get_configs, + patch("deep_agent.aegra.mcp._connect_single_server") as mock_connect, + ): + mock_get_configs.return_value = mock_servers + mock_connect.side_effect = [[tool_a1, tool_a2], [tool_b1, tool_b2]] + + tools = await get_mcp_tools() + + # Should have 3 tools: shared_tool (from server-a), unique_a, unique_b + assert len(tools) == 3 + tool_names = {t.name for t in tools} + assert tool_names == {"shared_tool", "unique_a", "unique_b"} + # First occurrence of shared_tool wins + assert tools[0] is tool_a1 + + @pytest.mark.asyncio + async def test_no_enabled_servers_returns_empty_list(self): + """Test that no enabled servers returns empty list.""" + _reset_mcp_cache() + mock_servers = { + "disabled": { + "url": "http://localhost/mcp/", + "enabled": False, + } + } + + with patch("deep_agent.aegra.mcp._get_server_configs") as mock_get_configs: + mock_get_configs.return_value = mock_servers + + tools = await get_mcp_tools() + + assert tools == [] + + @pytest.mark.asyncio + async def test_no_servers_configured_returns_empty_list(self): + """Test that no MCP servers configured returns empty list.""" + _reset_mcp_cache() + with patch("deep_agent.aegra.mcp._get_server_configs") as mock_get_configs: + mock_get_configs.return_value = {} + + tools = await get_mcp_tools() + + assert tools == [] + + @pytest.mark.asyncio + async def test_all_connections_fail_returns_empty_list(self): + """Test that all connection failures return empty list gracefully.""" + _reset_mcp_cache() + mock_servers = { + "server-a": { + "url": "http://a/mcp/", + "enabled": True, + "timeout": 1, + }, + "server-b": { + "url": "http://b/mcp/", + "enabled": True, + "timeout": 1, + }, + } + + with ( + patch("deep_agent.aegra.mcp._get_server_configs") as mock_get_configs, + patch("deep_agent.aegra.mcp._connect_single_server") as mock_connect, + ): + mock_get_configs.return_value = mock_servers + mock_connect.return_value = [] + + tools = await get_mcp_tools() + + assert tools == [] + + @pytest.mark.asyncio + async def test_sso_token_passed_to_build_config(self): + """Test that SSO token is passed through to _build_server_config.""" + _reset_mcp_cache() + mock_servers = { + "test": { + "url": "http://localhost/mcp/", + "enabled": True, + "auth": True, + "timeout": 5, + } + } + + mock_tool = MagicMock() + mock_tool.name = "tool1" + + with ( + patch("deep_agent.aegra.mcp._get_server_configs") as mock_get_configs, + patch("deep_agent.aegra.mcp._build_server_config") as mock_build_config, + patch("deep_agent.aegra.mcp._connect_single_server") as mock_connect, + ): + mock_get_configs.return_value = mock_servers + mock_build_config.return_value = {"url": "http://localhost/mcp/"} + mock_connect.return_value = [mock_tool] + + await get_mcp_tools("test_token_123") + + # Verify _build_server_config was called with the token + mock_build_config.assert_called_once() + call_args = mock_build_config.call_args + assert call_args[0][1] == "test_token_123" + + @pytest.mark.asyncio + async def test_parallel_connection_to_multiple_servers(self): + """Test that multiple servers are connected in parallel.""" + _reset_mcp_cache() + mock_servers = { + "server-1": {"url": "http://1/mcp/", "enabled": True, "timeout": 5}, + "server-2": {"url": "http://2/mcp/", "enabled": True, "timeout": 5}, + "server-3": {"url": "http://3/mcp/", "enabled": True, "timeout": 5}, + } + + tool1 = MagicMock() + tool1.name = "tool1" + tool2 = MagicMock() + tool2.name = "tool2" + tool3 = MagicMock() + tool3.name = "tool3" + + with ( + patch("deep_agent.aegra.mcp._get_server_configs") as mock_get_configs, + patch("deep_agent.aegra.mcp._connect_single_server") as mock_connect, + ): + mock_get_configs.return_value = mock_servers + mock_connect.side_effect = [[tool1], [tool2], [tool3]] + + tools = await get_mcp_tools() + + # All three servers should be connected + assert mock_connect.call_count == 3 + assert len(tools) == 3 + + @pytest.mark.asyncio + async def test_server_names_filters_enabled_servers(self): + """Test that server_names restricts which servers are connected.""" + _reset_mcp_cache() + mock_servers = { + "wanted": {"url": "http://w/mcp/", "enabled": True, "timeout": 5}, + "unwanted": {"url": "http://u/mcp/", "enabled": True, "timeout": 5}, + } + + tool_w = MagicMock() + tool_w.name = "wanted_tool" + + with ( + patch("deep_agent.aegra.mcp._get_server_configs") as mock_get_configs, + patch("deep_agent.aegra.mcp._connect_single_server") as mock_connect, + ): + mock_get_configs.return_value = mock_servers + mock_connect.return_value = [tool_w] + + tools = await get_mcp_tools(server_names=["wanted"]) + + mock_connect.assert_called_once() + assert len(tools) == 1 + assert tools[0].name == "wanted_tool" diff --git a/tests/unit/infrastructure/test_mcp_helpers.py b/tests/unit/infrastructure/test_mcp_helpers.py new file mode 100644 index 00000000..6af1c819 --- /dev/null +++ b/tests/unit/infrastructure/test_mcp_helpers.py @@ -0,0 +1,117 @@ +"""Unit tests for MCP helper functions (token refresh, error classification).""" + +import base64 +import json +import time +from unittest.mock import AsyncMock, patch + +import pytest + +from deep_agent.aegra.mcp import ( + _is_auth_error, + _is_connection_error, + _jwt_exp, + refresh_access_token, +) + + +class TestJwtExp: + def _make_jwt(self, exp: float) -> str: + header = base64.urlsafe_b64encode(b'{"alg":"HS256"}').rstrip(b"=").decode() + payload = ( + base64.urlsafe_b64encode(json.dumps({"exp": exp, "sub": "user"}).encode()) + .rstrip(b"=") + .decode() + ) + return f"{header}.{payload}.fakesig" + + def test_extracts_exp(self): + future = time.time() + 3600 + token = self._make_jwt(future) + assert abs(_jwt_exp(token) - future) < 1 + + def test_returns_zero_on_bad_token(self): + assert _jwt_exp("not.a.jwt") == 0.0 + assert _jwt_exp("") == 0.0 + assert _jwt_exp("single_segment") == 0.0 + + def test_returns_zero_when_no_exp(self): + header = base64.urlsafe_b64encode(b'{"alg":"HS256"}').rstrip(b"=").decode() + payload = ( + base64.urlsafe_b64encode(json.dumps({"sub": "user"}).encode()) + .rstrip(b"=") + .decode() + ) + token = f"{header}.{payload}.sig" + assert _jwt_exp(token) == 0.0 + + +class TestIsAuthError: + def test_401_in_message(self): + exc = Exception("HTTP 401 Unauthorized") + assert _is_auth_error(exc) is True + + def test_403_in_message(self): + exc = Exception("403 Forbidden") + assert _is_auth_error(exc) is True + + def test_non_auth_error(self): + exc = Exception("Connection refused") + assert _is_auth_error(exc) is False + + def test_nested_cause(self): + inner = Exception("Unauthorized") + outer = Exception("wrapper") + outer.__cause__ = inner + assert _is_auth_error(outer) is True + + +class TestIsConnectionError: + def test_connection_refused(self): + exc = Exception("Connection refused") + assert _is_connection_error(exc) is True + + def test_connect_error(self): + exc = Exception("ConnectError: failed to connect") + assert _is_connection_error(exc) is True + + def test_attempts_failed(self): + exc = Exception("All connection attempts failed") + assert _is_connection_error(exc) is True + + def test_non_connection_error(self): + exc = Exception("Invalid JSON response") + assert _is_connection_error(exc) is False + + +class TestRefreshAccessToken: + def _make_jwt(self, exp: float) -> str: + header = base64.urlsafe_b64encode(b'{"alg":"HS256"}').rstrip(b"=").decode() + payload = ( + base64.urlsafe_b64encode(json.dumps({"exp": exp, "sub": "user"}).encode()) + .rstrip(b"=") + .decode() + ) + return f"{header}.{payload}.fakesig" + + @pytest.mark.asyncio + async def test_returns_token_if_still_valid(self): + token = self._make_jwt(time.time() + 3600) + result = await refresh_access_token(token, "refresh_token") + assert result == token + + @pytest.mark.asyncio + async def test_returns_original_if_no_refresh_token(self): + token = self._make_jwt(time.time() - 60) + result = await refresh_access_token(token, None) + assert result == token + + @pytest.mark.asyncio + async def test_returns_original_if_no_token_endpoint(self): + token = self._make_jwt(time.time() - 60) + with patch( + "deep_agent.src.infrastructure.mcp._get_token_endpoint", + return_value="", + ): + result = await refresh_access_token(token, "refresh_tok") + assert result == token diff --git a/tests/unit/infrastructure/test_subagents.py b/tests/unit/infrastructure/test_subagents.py new file mode 100644 index 00000000..5af45a81 --- /dev/null +++ b/tests/unit/infrastructure/test_subagents.py @@ -0,0 +1,1032 @@ +"""Unit tests for subagent loading.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.agent.config.model import ModelSpec, Provider +from deep_agent.src.exceptions import SubAgentError +from deep_agent.src.infrastructure.subagents import VALID_AGENT_TYPES, load_subagents + + +class TestLoadSubagents: + """Tests for load_subagents function.""" + + def test_load_subagents_returns_none_when_no_configs(self): + """Test that load_subagents returns None when no subagent configs exist.""" + with patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs: + mock_get_configs.return_value = {} + + result = load_subagents(tools=[]) + + assert result is None + + def test_load_subagents_raises_error_when_model_missing(self): + """Test that load_subagents uses default model when none configured.""" + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, # No orchestrator model either + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ), + patch( + "deep_agent.src.infrastructure.subagents.SubAgent", + return_value=MagicMock(), + ), + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "description": "Test analyst", + "body": "Test prompt", + # Missing 'model' field - will use default + } + } + + result = load_subagents(tools=[]) + assert result is not None # Successfully creates with default model + + def test_load_single_subagent_minimal(self): + """Test loading a single subagent with minimal config.""" + mock_model = MagicMock() + mock_subagent = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, # No orchestrator config + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec" + ) as mock_create_model, + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "model": "gemini-2.5-flash", + "description": "Test analyst", + "body": "Test prompt", + } + } + mock_create_model.return_value = mock_model + mock_sa.return_value = mock_subagent + + result = load_subagents(tools=[]) + + assert result == [mock_subagent] + mock_create_model.assert_called_once() + # Should be called without middleware when no fallback + mock_sa.assert_called_once_with( + name="analyst", + model=mock_model, + description="Test analyst", + system_prompt="Test prompt", + ) + + def test_load_subagent_with_tools(self): + """Test loading subagent with tools that get resolved.""" + mock_tool1 = MagicMock() + mock_tool2 = MagicMock() + mock_model = MagicMock() + mock_subagent = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, # No orchestrator config + ), + patch( + "deep_agent.src.infrastructure.subagents.agent_config.resolve_tools" + ) as mock_resolve_tools, + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec" + ) as mock_create_model, + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "model": "gemini-2.5-flash", + "description": "Analyst", + "body": "Prompt", + "allowed_tools": ["calculate_bmi", "search_web"], + } + } + mock_resolve_tools.return_value = [mock_tool1, mock_tool2] + mock_create_model.return_value = mock_model + mock_sa.return_value = mock_subagent + + available_tools = [mock_tool1, mock_tool2] + result = load_subagents(tools=available_tools) + + assert result == [mock_subagent] + mock_resolve_tools.assert_called_once_with( + ["calculate_bmi", "search_web"], available_tools, agent_name="analyst" + ) + mock_sa.assert_called_once_with( + name="analyst", + model=mock_model, + description="Analyst", + system_prompt="Prompt", + tools=[mock_tool1, mock_tool2], + ) + + def test_load_subagent_with_skills(self): + """Test loading subagent with pre-resolved skill paths.""" + mock_model = MagicMock() + mock_subagent = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, # No orchestrator config + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec" + ) as mock_create_model, + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "model": "gemini-2.5-flash", + "description": "Analyst", + "body": "Prompt", + "skill_paths": ["/path/to/bmi-report"], + } + } + mock_create_model.return_value = mock_model + mock_sa.return_value = mock_subagent + + result = load_subagents(tools=[]) + + assert result == [mock_subagent] + mock_sa.assert_called_once_with( + name="analyst", + model=mock_model, + description="Analyst", + system_prompt="Prompt", + skills=["/path/to/bmi-report"], + ) + + def test_load_multiple_subagents(self): + """Test loading multiple subagents.""" + mock_model1 = MagicMock() + mock_model2 = MagicMock() + mock_sa1 = MagicMock() + mock_sa2 = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, # No orchestrator config + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec" + ) as mock_create_model, + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "model": "gemini-2.5-flash", + "description": "Analyst", + "body": "Analyst prompt", + }, + "publisher": { + "name": "publisher", + "model": "gemini-2.5-pro", + "description": "Publisher", + "body": "Publisher prompt", + }, + } + mock_create_model.side_effect = [mock_model1, mock_model2] + mock_sa.side_effect = [mock_sa1, mock_sa2] + + result = load_subagents(tools=[]) + + assert result == [mock_sa1, mock_sa2] + assert mock_create_model.call_count == 2 + assert mock_sa.call_count == 2 + + def test_load_subagent_with_empty_tool_list(self): + """Test that subagent with empty tools list doesn't call resolve_tools.""" + mock_model = MagicMock() + mock_subagent = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, # No orchestrator config + ), + patch( + "deep_agent.src.infrastructure.subagents.agent_config.resolve_tools" + ) as mock_resolve_tools, + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec" + ) as mock_create_model, + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "model": "gemini-2.5-flash", + "description": "Analyst", + "body": "Prompt", + "allowed_tools": [], + } + } + mock_create_model.return_value = mock_model + mock_sa.return_value = mock_subagent + + result = load_subagents(tools=[]) + + assert result == [mock_subagent] + mock_resolve_tools.assert_not_called() + # SubAgent should be called without tools parameter + mock_sa.assert_called_once_with( + name="analyst", + model=mock_model, + description="Analyst", + system_prompt="Prompt", + ) + + def test_load_subagent_uses_empty_description_when_missing(self): + """Test that missing description defaults to empty string.""" + mock_model = MagicMock() + mock_subagent = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, # No orchestrator config + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec" + ) as mock_create_model, + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "model": "gemini-2.5-flash", + "body": "Prompt", + # Missing 'description' + } + } + mock_create_model.return_value = mock_model + mock_sa.return_value = mock_subagent + + result = load_subagents(tools=[]) + + assert result == [mock_subagent] + mock_sa.assert_called_once_with( + name="analyst", + model=mock_model, + description="", + system_prompt="Prompt", + ) + + +class TestAgentTypeSystem: + """Tests for the type field and multi-type subagent dispatch.""" + + def test_valid_agent_types_constant(self): + assert "default" in VALID_AGENT_TYPES + assert "compiled" in VALID_AGENT_TYPES + assert "async" in VALID_AGENT_TYPES + + def test_invalid_type_raises_value_error(self): + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + ): + mock_get_configs.return_value = { + "bad": { + "name": "bad", + "type": "invalid_type", + "model": "gemini-2.5-pro", + "description": "Bad agent", + "body": "Prompt", + } + } + with pytest.raises(SubAgentError, match="invalid type 'invalid_type'"): + load_subagents(tools=[]) + + def test_missing_type_defaults_to_default(self): + """No type field means SubAgent (default).""" + mock_model = MagicMock() + mock_subagent = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec" + ) as mock_create_model, + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "model": "gemini-2.5-flash", + "description": "Analyst", + "body": "Prompt", + # No 'type' field + } + } + mock_create_model.return_value = mock_model + mock_sa.return_value = mock_subagent + + result = load_subagents(tools=[]) + assert result == [mock_subagent] + mock_sa.assert_called_once() + + def test_type_default_builds_subagent(self): + mock_model = MagicMock() + mock_subagent = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec" + ) as mock_create_model, + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + ): + mock_get_configs.return_value = { + "publisher": { + "name": "publisher", + "type": "default", + "model": "gemini-2.5-pro", + "description": "Publisher", + "body": "Prompt", + } + } + mock_create_model.return_value = mock_model + mock_sa.return_value = mock_subagent + + result = load_subagents(tools=[]) + assert result == [mock_subagent] + + def test_type_compiled_builds_compiled_subagent(self): + mock_model = MagicMock() + mock_graph = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec" + ) as mock_create_model, + patch("deepagents.create_deep_agent") as mock_create_agent, + patch( + "deep_agent.src.infrastructure.backend.get_configured_backend" + ) as mock_get_backend, + patch( + "deep_agent.src.infrastructure.subagents.CompiledSubAgent" + ) as mock_compiled, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "type": "compiled", + "model": "gemini-2.5-pro", + "description": "Fast analyst", + "body": "Prompt", + } + } + mock_create_model.return_value = mock_model + mock_create_agent.return_value = mock_graph + mock_get_backend.return_value = MagicMock() + mock_compiled.return_value = MagicMock() + + result = load_subagents(tools=[]) + assert len(result) == 1 + mock_create_agent.assert_called_once() + mock_compiled.assert_called_once_with( + name="analyst", + description="Fast analyst", + runnable=mock_graph, + ) + + def test_type_async_builds_async_subagent(self): + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.AsyncSubAgent" + ) as mock_async_sa, + ): + mock_get_configs.return_value = { + "researcher": { + "name": "researcher", + "type": "async", + "description": "Remote researcher", + "body": "", + "graph_id": "researcher-graph", + "url": "http://research-agent:8000", + } + } + mock_async_sa.return_value = MagicMock() + + result = load_subagents(tools=[]) + assert len(result) == 1 + mock_async_sa.assert_called_once_with( + name="researcher", + description="Remote researcher", + graph_id="researcher-graph", + url="http://research-agent:8000", + ) + + def test_type_async_raises_without_graph_id(self): + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, + ), + patch( + "deep_agent.src.infrastructure.subagents.AsyncSubAgent", None + ), # Simulate async support not available + ): + mock_get_configs.return_value = { + "bad_async": { + "name": "bad_async", + "type": "async", + "description": "Missing graph_id", + "body": "", + # No graph_id + } + } + with pytest.raises( + SubAgentError, match="requires deepagents with async support" + ): + load_subagents(tools=[]) + + +class TestSubagentProviderConfig: + """Tests for provider-aware model configuration.""" + + def test_inherits_orchestrator_string_model(self): + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={"model": "gemini-2.5-flash"}, + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ) as mock_from_spec, + patch( + "deep_agent.src.infrastructure.subagents.SubAgent", + return_value=MagicMock(), + ), + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "description": "Analyst", + "body": "Prompt", + } + } + + load_subagents(tools=[]) + + spec = mock_from_spec.call_args[0][0] + assert spec.name == "gemini-2.5-flash" + + def test_orchestrator_as_fallback_when_subagent_has_string_model(self): + """Subagent with string model and no fallback → orchestrator becomes fallback.""" + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={"model": "gemini-2.5-flash"}, + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ) as mock_from_spec, + patch( + "langchain.agents.middleware.ModelFallbackMiddleware" + ) as mock_middleware, + patch( + "deep_agent.src.infrastructure.subagents.SubAgent", + return_value=MagicMock(), + ) as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "description": "Analyst", + "body": "Prompt", + "model": "gpt-4", # String model, no fallback + } + } + + load_subagents(tools=[]) + + # Verify middleware was created and passed to SubAgent + assert mock_middleware.called + call_kwargs = mock_sa.call_args[1] + assert "middleware" in call_kwargs + assert len(call_kwargs["middleware"]) == 1 + + def test_orchestrator_as_fallback_when_subagent_has_dict_model_no_fallback(self): + """Subagent with dict model and no fallback → orchestrator becomes fallback.""" + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={"model": "gemini-2.5-flash"}, + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ) as mock_from_spec, + patch( + "langchain.agents.middleware.ModelFallbackMiddleware" + ) as mock_middleware, + patch( + "deep_agent.src.infrastructure.subagents.SubAgent", + return_value=MagicMock(), + ) as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "description": "Analyst", + "body": "Prompt", + "model": {"provider": "openai", "name": "gpt-4"}, + } + } + + load_subagents(tools=[]) + + # Verify middleware was created and passed to SubAgent + assert mock_middleware.called + call_kwargs = mock_sa.call_args[1] + assert "middleware" in call_kwargs + assert len(call_kwargs["middleware"]) == 1 + + def test_keeps_explicit_fallback_when_provided(self): + """Subagent with explicit fallback → keep as-is (don't override).""" + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={"model": "gemini-2.5-flash"}, + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ) as mock_from_spec, + patch( + "langchain.agents.middleware.ModelFallbackMiddleware" + ) as mock_middleware, + patch( + "deep_agent.src.infrastructure.subagents.SubAgent", + return_value=MagicMock(), + ) as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "description": "Analyst", + "body": "Prompt", + "model": { + "provider": "openai", + "name": "gpt-4", + "fallback": {"provider": "vertex", "name": "gemini-3.1-pro"}, + }, + } + } + + load_subagents(tools=[]) + + # Verify middleware was created and passed to SubAgent + assert mock_middleware.called + call_kwargs = mock_sa.call_args[1] + assert "middleware" in call_kwargs + assert len(call_kwargs["middleware"]) == 1 + + def test_no_fallback_when_no_orchestrator_model(self): + """Subagent with model but orchestrator has no model → no fallback added.""" + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, # No orchestrator model + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ) as mock_from_spec, + patch( + "deep_agent.src.infrastructure.subagents.SubAgent", + return_value=MagicMock(), + ), + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "description": "Analyst", + "body": "Prompt", + "model": "gpt-4", + } + } + + load_subagents(tools=[]) + + spec = mock_from_spec.call_args[0][0] + assert spec.name == "gpt-4" + # No orchestrator model → no fallback + assert spec.fallback is None + + def test_strips_nested_fallback_from_orchestrator(self): + """Orchestrator with fallback → strip when using as subagent fallback.""" + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={ + "model": { + "provider": "vertex", + "name": "gemini-2.5-flash", + "fallback": {"provider": "openai", "name": "gpt-4o-mini"}, + } + }, + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ) as mock_from_spec, + patch( + "langchain.agents.middleware.ModelFallbackMiddleware" + ) as mock_middleware, + patch( + "deep_agent.src.infrastructure.subagents.SubAgent", + return_value=MagicMock(), + ) as mock_sa, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "description": "Analyst", + "body": "Prompt", + "model": "gpt-4", + } + } + + load_subagents(tools=[]) + + # Verify middleware was created and passed to SubAgent + assert mock_middleware.called + call_kwargs = mock_sa.call_args[1] + assert "middleware" in call_kwargs + assert len(call_kwargs["middleware"]) == 1 + + +class TestToolAccessControl: + """Tests for allowed_tools, denied_tools, and tool_approval in subagent building.""" + + def test_denied_tools_filtered_from_resolved(self): + """Subagent with denied_tools has those tools removed.""" + mock_tool_a = MagicMock() + mock_tool_a.name = "tool_a" + mock_tool_b = MagicMock() + mock_tool_b.name = "tool_b" + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, + ), + patch( + "deep_agent.src.infrastructure.subagents.agent_config.resolve_tools" + ) as mock_resolve, + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ), + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + patch( + "deep_agent.src.infrastructure.subagents.filter_denied_tools" + ) as mock_filter, + ): + mock_get_configs.return_value = { + "agent1": { + "name": "agent1", + "model": "gemini-2.5-flash", + "description": "Test", + "body": "Prompt", + "allowed_tools": ["tool_a", "tool_b"], + "denied_tools": ["tool_b"], + } + } + mock_resolve.return_value = [mock_tool_a, mock_tool_b] + mock_filter.return_value = [mock_tool_a] + mock_sa.return_value = MagicMock() + + load_subagents(tools=[mock_tool_a, mock_tool_b]) + + mock_filter.assert_called_once_with( + [mock_tool_a, mock_tool_b], ["tool_b"], agent_name="agent1" + ) + mock_sa.assert_called_once() + call_kwargs = mock_sa.call_args[1] + assert call_kwargs["tools"] == [mock_tool_a] + + def test_default_subagent_rejects_tool_approval(self): + """Default subagent with tool_approval raises error — must use compiled.""" + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, + ), + ): + mock_get_configs.return_value = { + "agent1": { + "name": "agent1", + "type": "default", + "model": "gemini-2.5-flash", + "description": "Test", + "body": "Prompt", + "allowed_tools": ["sensitive_tool"], + "tool_approval": ["sensitive_tool"], + } + } + with pytest.raises(SubAgentError, match="does not support tool_approval"): + load_subagents(tools=[]) + + def test_tool_approval_names_from_default_subagents_in_config(self): + """tool_approval field is read from default subagent configs.""" + with patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs: + configs = { + "default_agent": { + "name": "default_agent", + "type": "default", + "tool_approval": ["send_email", "search_web"], + }, + "compiled_agent": { + "name": "compiled_agent", + "type": "compiled", + "tool_approval": ["delete_record"], + }, + } + mock_get_configs.return_value = configs + # Verify default subagent has tool_approval, compiled does not + default_approvals = [ + t + for _, c in configs.items() + if c.get("type", "default") == "default" + for t in c.get("tool_approval", []) + ] + assert "send_email" in default_approvals + assert "search_web" in default_approvals + assert "delete_record" not in default_approvals + + def test_async_subagent_rejects_tool_approval(self): + """Async subagent with tool_approval raises error.""" + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, + ), + ): + mock_get_configs.return_value = { + "remote": { + "name": "remote", + "type": "async", + "description": "Remote agent", + "body": "", + "graph_id": "remote-graph", + "tool_approval": ["some_tool"], + } + } + with pytest.raises(SubAgentError, match="does not support tool_approval"): + load_subagents(tools=[]) + + def test_denied_tools_with_mcp_inheritance(self): + """Subagent inheriting all MCP tools still filters denied ones.""" + mock_tool_a = MagicMock() + mock_tool_a.name = "safe_tool" + mock_tool_b = MagicMock() + mock_tool_b.name = "dangerous_tool" + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={"mcps": ["my-mcp"]}, + ), + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ), + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + patch( + "deep_agent.src.infrastructure.subagents.filter_denied_tools" + ) as mock_filter, + ): + mock_get_configs.return_value = { + "admin": { + "name": "admin", + "model": "gemini-2.5-flash", + "description": "Admin", + "body": "Prompt", + # No allowed_tools — inherits all via mcps + "denied_tools": ["dangerous_tool"], + } + } + mock_filter.return_value = [mock_tool_a] + mock_sa.return_value = MagicMock() + + load_subagents(tools=[mock_tool_a, mock_tool_b]) + + mock_filter.assert_called_once() + + def test_compiled_subagent_gets_denied_tools_filtered(self): + """Compiled subagent also filters denied tools (same as default).""" + mock_tool = MagicMock() + mock_tool.name = "tool_a" + mock_model = MagicMock() + mock_graph = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, + ), + patch( + "deep_agent.src.infrastructure.subagents.agent_config.resolve_tools" + ) as mock_resolve, + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ), + patch("deepagents.create_deep_agent") as mock_create_agent, + patch( + "deep_agent.src.infrastructure.backend.get_configured_backend" + ) as mock_backend, + patch( + "deep_agent.src.infrastructure.subagents.CompiledSubAgent" + ) as mock_compiled, + patch( + "deep_agent.src.infrastructure.subagents.filter_denied_tools" + ) as mock_filter, + ): + mock_get_configs.return_value = { + "analyst": { + "name": "analyst", + "type": "compiled", + "model": "gemini-2.5-pro", + "description": "Analyst", + "body": "Prompt", + "allowed_tools": ["tool_a", "tool_b"], + "denied_tools": ["tool_b"], + } + } + mock_resolve.return_value = [mock_tool] + mock_filter.return_value = [mock_tool] + mock_create_agent.return_value = mock_graph + mock_backend.return_value = MagicMock() + mock_compiled.return_value = MagicMock() + + load_subagents(tools=[mock_tool]) + + mock_filter.assert_called_once() + + def test_two_subagents_get_different_tool_sets(self): + """Two subagents with different allowed_tools get isolated tool sets.""" + mock_tool_a = MagicMock() + mock_tool_a.name = "tool_a" + mock_tool_b = MagicMock() + mock_tool_b.name = "tool_b" + mock_model = MagicMock() + + with ( + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_all_subagent_configs" + ) as mock_get_configs, + patch( + "deep_agent.src.infrastructure.subagents.agent_config.get_orchestrator_config", + return_value={}, + ), + patch( + "deep_agent.src.infrastructure.subagents.agent_config.resolve_tools" + ) as mock_resolve, + patch( + "deep_agent.src.infrastructure.subagents.get_or_create_model_from_spec", + return_value=mock_model, + ), + patch("deep_agent.src.infrastructure.subagents.SubAgent") as mock_sa, + ): + mock_get_configs.return_value = { + "agent_x": { + "name": "agent_x", + "model": "gemini-2.5-flash", + "description": "Agent X", + "body": "Prompt", + "allowed_tools": ["tool_a"], + }, + "agent_y": { + "name": "agent_y", + "model": "gemini-2.5-flash", + "description": "Agent Y", + "body": "Prompt", + "allowed_tools": ["tool_b"], + }, + } + # resolve_tools returns different results per call + mock_resolve.side_effect = [[mock_tool_a], [mock_tool_b]] + mock_sa.return_value = MagicMock() + + load_subagents(tools=[mock_tool_a, mock_tool_b]) + + assert mock_sa.call_count == 2 + calls = mock_sa.call_args_list + # First subagent gets tool_a only + assert calls[0][1]["tools"] == [mock_tool_a] + # Second subagent gets tool_b only + assert calls[1][1]["tools"] == [mock_tool_b] diff --git a/tests/unit/infrastructure/test_tool_access.py b/tests/unit/infrastructure/test_tool_access.py new file mode 100644 index 00000000..df86c7e6 --- /dev/null +++ b/tests/unit/infrastructure/test_tool_access.py @@ -0,0 +1,232 @@ +"""Unit tests for tool access control.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.exceptions import AppException, ErrorCodes +from deep_agent.src.infrastructure.tool_access import ( + _wrap_tool_with_approval, + apply_tool_approval, + filter_denied_tools, + migrate_tools_field, +) + + +def _make_tool(name: str, *, with_func: bool = False) -> MagicMock: + """Create a mock tool with the given name. + + Args: + name: Tool name to set on the mock. + with_func: If True, configure the mock so that wrapping + falls back to in-place attribute assignment (matching the + fallback path in _wrap_tool_with_approval when model_copy + is not available). + """ + tool = MagicMock() + tool.name = name + if with_func: + # Force fallback path: model_copy raises so the wrapper + # assigns .func / .coroutine directly on the mock. + tool.model_copy.side_effect = AttributeError("not a pydantic model") + tool.coroutine = None + return tool + + +class TestFilterDeniedTools: + """Tests for filter_denied_tools.""" + + def test_removes_denied_tools(self): + """Denying one tool out of three leaves two.""" + t1, t2, t3 = _make_tool("a"), _make_tool("b"), _make_tool("c") + result = filter_denied_tools([t1, t2, t3], ["b"], agent_name="test") + assert result == [t1, t3] + + def test_empty_denied_list_returns_unchanged(self): + """No denied names means the original list is returned.""" + tools = [_make_tool("x"), _make_tool("y")] + result = filter_denied_tools(tools, [], agent_name="test") + assert result is tools + + def test_denied_tool_not_in_list_is_noop(self): + """Denying a name that does not exist causes no error.""" + t1, t2 = _make_tool("a"), _make_tool("b") + result = filter_denied_tools([t1, t2], ["nonexistent"], agent_name="test") + assert result == [t1, t2] + + def test_all_tools_denied(self): + """Denying every tool returns an empty list.""" + t1, t2 = _make_tool("a"), _make_tool("b") + result = filter_denied_tools([t1, t2], ["a", "b"], agent_name="test") + assert result == [] + + def test_deny_wins_over_presence(self): + """A tool present in both the list and the deny set is removed.""" + t1 = _make_tool("search") + result = filter_denied_tools([t1], ["search"], agent_name="test") + assert t1 not in result + + def test_preserves_tool_order(self): + """Remaining tools maintain their original order.""" + t1, t2, t3, t4 = ( + _make_tool("a"), + _make_tool("b"), + _make_tool("c"), + _make_tool("d"), + ) + result = filter_denied_tools([t1, t2, t3, t4], ["b", "d"], agent_name="test") + assert result == [t1, t3] + + +class TestApplyToolApproval: + """Tests for apply_tool_approval.""" + + def test_wraps_named_tools_only(self): + """Only the tool whose name is in approval_names gets wrapped.""" + t1, t2, t3 = _make_tool("a"), _make_tool("b", with_func=True), _make_tool("c") + original_func = lambda **kw: "original" + t2.func = original_func + + result = apply_tool_approval([t1, t2, t3], ["b"], agent_name="test") + + # t1 and t3 are passed through unchanged + assert result[0] is t1 + assert result[2] is t3 + # t2 is the same object but its func has been replaced + assert result[1] is t2 + assert result[1].func is not original_func + + def test_empty_approval_list_returns_unchanged(self): + """No approval names means the original list is returned.""" + tools = [_make_tool("x")] + result = apply_tool_approval(tools, [], agent_name="test") + assert result is tools + + def test_unknown_approval_tool_logs_warning(self): + """Approving a name not in the tool list logs a warning.""" + t1 = _make_tool("a") + with patch("deep_agent.src.infrastructure.tool_access.logger") as mock_logger: + result = apply_tool_approval([t1], ["nonexistent"], agent_name="test") + mock_logger.warning.assert_called_once() + assert mock_logger.warning.call_args[1]["tool"] == "nonexistent" + # t1 is still in the result unchanged + assert result == [t1] + + def test_all_tools_wrapped(self): + """When all tool names are in approval_names, all get wrapped.""" + t1, t2 = _make_tool("a", with_func=True), _make_tool("b", with_func=True) + orig_a = lambda **kw: "orig_a" + orig_b = lambda **kw: "orig_b" + t1.func = orig_a + t2.func = orig_b + + result = apply_tool_approval([t1, t2], ["a", "b"], agent_name="test") + + assert result[0].func is not orig_a + assert result[1].func is not orig_b + + def test_wrapped_tool_preserves_name(self): + """The wrapped tool retains the original .name.""" + t1 = _make_tool("search", with_func=True) + t1.func = lambda **kw: "original" + + result = apply_tool_approval([t1], ["search"], agent_name="test") + + assert result[0].name == "search" + + +class TestWrapToolWithApproval: + """Tests for _wrap_tool_with_approval.""" + + def test_interrupt_called_on_invocation(self): + """Invoking the wrapped sync tool triggers interrupt with HITL payload.""" + tool = _make_tool("run_query", with_func=True) + tool.func = lambda **kw: "result" + + with patch( + "deep_agent.src.infrastructure.tool_access.interrupt", + return_value=[{"type": "approve"}], + ) as mock_interrupt: + wrapped = _wrap_tool_with_approval(tool, agent_name="analyst") + wrapped.func() + mock_interrupt.assert_called_once() + payload = mock_interrupt.call_args[0][0] + assert isinstance(payload, dict) + assert "action_requests" in payload + assert payload["action_requests"][0]["name"] == "run_query" + assert "analyst" in payload["action_requests"][0]["args"]["agent"] + + def test_approved_executes_original(self): + """When frontend sends approve decision, the original function runs.""" + original_func = MagicMock(return_value="query_result") + tool = _make_tool("run_query", with_func=True) + tool.func = original_func + + with patch( + "deep_agent.src.infrastructure.tool_access.interrupt", + return_value=[{"type": "approve"}], + ): + wrapped = _wrap_tool_with_approval(tool, agent_name="analyst") + result = wrapped.func(sql="SELECT 1") + original_func.assert_called_once_with(sql="SELECT 1") + assert result == "query_result" + + def test_rejected_returns_message(self): + """When frontend sends reject decision, a rejection message is returned.""" + tool = _make_tool("dangerous_op", with_func=True) + tool.func = lambda **kw: "should not run" + + with patch( + "deep_agent.src.infrastructure.tool_access.interrupt", + return_value=[{"type": "reject", "message": "No"}], + ): + wrapped = _wrap_tool_with_approval(tool, agent_name="analyst") + result = wrapped.func() + assert result == "Tool 'dangerous_op' was rejected by the user." + + def test_case_insensitive_approval(self): + """String 'approved' still works for backward compat / testing.""" + original_func = MagicMock(return_value="ok") + tool = _make_tool("action", with_func=True) + tool.func = original_func + + with patch( + "deep_agent.src.infrastructure.tool_access.interrupt", + return_value="Approved", + ): + wrapped = _wrap_tool_with_approval(tool, agent_name="analyst") + result = wrapped.func() + original_func.assert_called_once() + assert result == "ok" + + +class TestMigrateToolsField: + """Tests for migrate_tools_field.""" + + def test_tools_migrated_to_allowed_tools(self): + """Config with 'tools' and no 'allowed_tools' gets migrated.""" + config = {"tools": ["a", "b"], "name": "analyst"} + result = migrate_tools_field(config, agent_name="analyst") + + assert result["allowed_tools"] == ["a", "b"] + assert "tools" not in result + + def test_both_present_raises_error(self): + """Config with both 'tools' and 'allowed_tools' raises AppException.""" + config = {"tools": ["a"], "allowed_tools": ["b"]} + with pytest.raises(AppException) as exc_info: + migrate_tools_field(config, agent_name="analyst") + assert exc_info.value.error_code == ErrorCodes.CONFIGURATION_VALIDATION_ERROR + + def test_neither_present_is_noop(self): + """Config with neither key is unchanged.""" + config = {"name": "analyst", "description": "Test"} + result = migrate_tools_field(config, agent_name="analyst") + assert result == {"name": "analyst", "description": "Test"} + + def test_allowed_tools_only_is_noop(self): + """Config with only 'allowed_tools' is unchanged.""" + config = {"allowed_tools": ["a", "b"]} + result = migrate_tools_field(config, agent_name="analyst") + assert result == {"allowed_tools": ["a", "b"]} + assert "tools" not in result diff --git a/tests/unit/memory/__init__.py b/tests/unit/memory/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/memory/test_clustering.py b/tests/unit/memory/test_clustering.py new file mode 100644 index 00000000..f1cb7080 --- /dev/null +++ b/tests/unit/memory/test_clustering.py @@ -0,0 +1,48 @@ +"""Unit tests for semantic clustering.""" + +from deep_agent.src.memory.clustering import cluster_memories + + +class TestClusterMemories: + def test_clusters_similar(self): + contents = [ + "I like Python programming language", + "Python is my favorite programming language", + "The weather today is very sunny", + ] + clusters = cluster_memories(contents, threshold=0.3) + assert len(clusters) == 1 + assert set(clusters[0]) == {0, 1} + + def test_no_clusters_when_disjoint(self): + contents = [ + "I like cats", + "The sky is blue", + "Databases are useful", + ] + clusters = cluster_memories(contents, threshold=0.5) + assert clusters == [] + + def test_empty_list(self): + assert cluster_memories([], threshold=0.5) == [] + + def test_single_item(self): + assert cluster_memories(["hello world"], threshold=0.5) == [] + + def test_all_similar(self): + contents = [ + "Python is great for data science", + "Python data science is great", + "Data science with Python is great", + ] + clusters = cluster_memories(contents, threshold=0.3) + assert len(clusters) == 1 + assert len(clusters[0]) == 3 + + def test_high_threshold_no_match(self): + contents = [ + "I like Python", + "Python is good", + ] + clusters = cluster_memories(contents, threshold=0.99) + assert clusters == [] diff --git a/tests/unit/memory/test_config.py b/tests/unit/memory/test_config.py new file mode 100644 index 00000000..15c1fe09 --- /dev/null +++ b/tests/unit/memory/test_config.py @@ -0,0 +1,39 @@ +"""Unit tests for memory configuration.""" + +from deep_agent.src.memory.config import MemorySettings + + +class TestMemorySettings: + def test_defaults_all_disabled(self): + s = MemorySettings( + MEMORY_CONSOLIDATION_ENABLED=False, + MEMORY_DECAY_ENABLED=False, + MEMORY_CLUSTERING_ENABLED=False, + MEMORY_RELATIONSHIPS_ENABLED=False, + ) + assert s.MEMORY_CONSOLIDATION_ENABLED is False + assert s.MEMORY_DECAY_ENABLED is False + + def test_is_enabled_requires_master(self): + s = MemorySettings( + MEMORY_CONSOLIDATION_ENABLED=False, + MEMORY_DECAY_ENABLED=True, + ) + assert s.is_enabled("decay") is False + + def test_is_enabled_with_master_on(self): + s = MemorySettings( + MEMORY_CONSOLIDATION_ENABLED=True, + MEMORY_DECAY_ENABLED=True, + ) + assert s.is_enabled("decay") is True + + def test_is_enabled_unknown_layer(self): + s = MemorySettings(MEMORY_CONSOLIDATION_ENABLED=True) + assert s.is_enabled("nonexistent") is False + + def test_defaults(self): + s = MemorySettings() + assert s.MEMORY_MAX_INJECT == 20 + assert s.MEMORY_DECAY_LAMBDA == 0.05 + assert s.MEMORY_SCHEDULER_INTERVAL_HOURS == 6 diff --git a/tests/unit/memory/test_consolidation.py b/tests/unit/memory/test_consolidation.py new file mode 100644 index 00000000..0d5eab26 --- /dev/null +++ b/tests/unit/memory/test_consolidation.py @@ -0,0 +1,69 @@ +"""Unit tests for memory consolidation.""" + +from deep_agent.src.memory.consolidation import ( + find_duplicates, + pick_representative, + token_similarity, +) + + +class TestTokenSimilarity: + def test_identical(self): + assert token_similarity("hello world", "hello world") == 1.0 + + def test_disjoint(self): + assert token_similarity("hello world", "foo bar") == 0.0 + + def test_partial_overlap(self): + sim = token_similarity("I like Python", "I love Python") + assert 0.3 < sim < 0.9 + + def test_empty_string(self): + assert token_similarity("", "hello") == 0.0 + + def test_case_insensitive(self): + assert token_similarity("Python", "python") == 1.0 + + +class TestFindDuplicates: + def test_no_duplicates(self): + memories = [ + {"content": "I like cats"}, + {"content": "The weather is sunny"}, + {"content": "Python is great for data science"}, + ] + groups = find_duplicates(memories, threshold=0.5) + assert groups == [] + + def test_finds_duplicates(self): + memories = [ + {"content": "I prefer Python programming"}, + {"content": "I prefer Python for programming"}, + {"content": "The weather is nice today"}, + ] + groups = find_duplicates(memories, threshold=0.5) + assert len(groups) == 1 + assert set(groups[0]) == {0, 1} + + def test_single_memory(self): + memories = [{"content": "just one"}] + assert find_duplicates(memories) == [] + + def test_empty_list(self): + assert find_duplicates([]) == [] + + +class TestPickRepresentative: + def test_picks_longest(self): + memories = [ + {"content": "short", "score": "0.5"}, + {"content": "this is much longer content", "score": "0.5"}, + ] + assert pick_representative(memories, [0, 1]) == 1 + + def test_breaks_tie_by_score(self): + memories = [ + {"content": "same length!", "score": "0.9"}, + {"content": "same length!", "score": "0.3"}, + ] + assert pick_representative(memories, [0, 1]) == 0 diff --git a/tests/unit/memory/test_relationships.py b/tests/unit/memory/test_relationships.py new file mode 100644 index 00000000..027c8b1f --- /dev/null +++ b/tests/unit/memory/test_relationships.py @@ -0,0 +1,52 @@ +"""Unit tests for relationship inference.""" + +from deep_agent.src.memory.relationships import ( + extract_keywords, + find_related_pairs, +) + + +class TestExtractKeywords: + def test_basic(self): + keywords = extract_keywords("Python is a great programming language") + assert "python" in keywords + assert "programming" in keywords + assert "language" in keywords + + def test_filters_stopwords(self): + keywords = extract_keywords("I am a very good person") + assert "good" in keywords + assert "person" in keywords + assert "very" not in keywords + + def test_filters_short_tokens(self): + keywords = extract_keywords("Go is ok") + assert "go" not in keywords + assert "ok" not in keywords + + def test_empty(self): + assert extract_keywords("") == [] + + +class TestFindRelatedPairs: + def test_finds_related(self): + memories = [ + {"content": "I work at Red Hat on OpenShift platform engineering"}, + {"content": "Red Hat OpenShift is my primary deployment target"}, + {"content": "I like pizza and pasta for dinner"}, + ] + pairs = find_related_pairs(memories, min_shared=2) + assert len(pairs) == 1 + assert pairs[0][0] == 0 + assert pairs[0][1] == 1 + assert "openshift" in pairs[0][2] + + def test_no_related(self): + memories = [ + {"content": "I like cats and dogs"}, + {"content": "The weather is sunny today"}, + ] + assert find_related_pairs(memories, min_shared=2) == [] + + def test_empty(self): + assert find_related_pairs([], min_shared=2) == [] diff --git a/tests/unit/memory/test_scheduler.py b/tests/unit/memory/test_scheduler.py new file mode 100644 index 00000000..eadf58b3 --- /dev/null +++ b/tests/unit/memory/test_scheduler.py @@ -0,0 +1,83 @@ +"""Unit tests for memory scheduler.""" + +from unittest.mock import AsyncMock, patch + +from deep_agent.src.memory import scheduler +from deep_agent.src.memory.config import MemorySettings + + +class TestScheduler: + def setup_method(self): + scheduler._scheduler = None + + async def test_start_skips_when_disabled(self): + disabled = MemorySettings(MEMORY_CONSOLIDATION_ENABLED=False) + with patch.object(scheduler, "memory_settings", disabled): + result = await scheduler.start_scheduler("postgresql://test") + assert result is False + + async def test_stop_is_safe_when_not_started(self): + await scheduler.stop_scheduler() + + async def test_run_once_calls_all_jobs(self): + enabled = MemorySettings( + MEMORY_CONSOLIDATION_ENABLED=True, + MEMORY_DECAY_ENABLED=True, + MEMORY_CLUSTERING_ENABLED=True, + MEMORY_RELATIONSHIPS_ENABLED=True, + ) + with ( + patch.object(scheduler, "memory_settings", enabled), + patch( + "deep_agent.src.memory.scoring.decay_all_memories", + new_callable=AsyncMock, + return_value=5, + ), + patch( + "deep_agent.src.memory.consolidation.consolidate_all_users", + new_callable=AsyncMock, + return_value=3, + ), + patch( + "deep_agent.src.memory.clustering.cluster_all_users", + new_callable=AsyncMock, + return_value=2, + ), + patch( + "deep_agent.src.memory.relationships.infer_all_relationships", + new_callable=AsyncMock, + return_value=4, + ), + ): + results = await scheduler.run_once("postgresql://test") + assert results["decay"] == 5 + assert results["consolidation"] == 3 + assert results["clustering"] == 2 + assert results["relationships"] == 4 + + async def test_run_once_handles_job_failure(self): + with ( + patch( + "deep_agent.src.memory.scoring.decay_all_memories", + new_callable=AsyncMock, + side_effect=Exception("boom"), + ), + patch( + "deep_agent.src.memory.consolidation.consolidate_all_users", + new_callable=AsyncMock, + return_value=0, + ), + patch( + "deep_agent.src.memory.clustering.cluster_all_users", + new_callable=AsyncMock, + return_value=0, + ), + patch( + "deep_agent.src.memory.relationships.infer_all_relationships", + new_callable=AsyncMock, + return_value=0, + ), + ): + results = await scheduler.run_once("postgresql://test") + assert results["decay"] == -1 + assert results["consolidation"] == 0 diff --git a/tests/unit/memory/test_scoring.py b/tests/unit/memory/test_scoring.py new file mode 100644 index 00000000..328ef2b3 --- /dev/null +++ b/tests/unit/memory/test_scoring.py @@ -0,0 +1,56 @@ +"""Unit tests for exponential decay scoring.""" + +from datetime import datetime, timedelta, timezone + +from deep_agent.src.memory.scoring import ( + MIN_SCORE, + apply_access_boost, + compute_decay_score, +) + + +class TestComputeDecayScore: + def test_fresh_memory_keeps_score(self): + now = datetime.now(timezone.utc) + score = compute_decay_score(1.0, now, now) + assert score == 1.0 + + def test_old_memory_decays(self): + now = datetime.now(timezone.utc) + old = now - timedelta(days=30) + score = compute_decay_score(1.0, old, now) + assert score < 1.0 + assert score > MIN_SCORE + + def test_very_old_memory_near_min(self): + now = datetime.now(timezone.utc) + ancient = now - timedelta(days=365) + score = compute_decay_score(1.0, ancient, now) + assert score <= 0.05 + + def test_never_below_min(self): + now = datetime.now(timezone.utc) + ancient = now - timedelta(days=10000) + score = compute_decay_score(1.0, ancient, now) + assert score >= MIN_SCORE + + def test_naive_datetime_handled(self): + now = datetime.now(timezone.utc) + naive = datetime.utcnow() + score = compute_decay_score(1.0, naive, now) + assert 0.99 < score <= 1.0 + + def test_zero_age(self): + now = datetime.now(timezone.utc) + assert compute_decay_score(0.5, now, now) == 0.5 + + +class TestAccessBoost: + def test_boost_increases_score(self): + assert apply_access_boost(0.5) == 0.6 + + def test_boost_capped_at_one(self): + assert apply_access_boost(0.95) == 1.0 + + def test_boost_from_zero(self): + assert apply_access_boost(0.0) == 0.1 diff --git a/tests/unit/observability/test_otel_setup.py b/tests/unit/observability/test_otel_setup.py new file mode 100644 index 00000000..960963dc --- /dev/null +++ b/tests/unit/observability/test_otel_setup.py @@ -0,0 +1,52 @@ +"""Unit tests for platform-style OTEL bootstrap.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from deep_agent.src.observability import otel_setup + + +def test_setup_otel_metrics_skips_when_disabled() -> None: + settings = MagicMock() + settings.ENABLE_OTEL_METRICS = False + settings.OTEL_EXPORTER_OTLP_ENDPOINT = "otel-gateway:4327" + log = MagicMock() + + with patch("opentelemetry.metrics.set_meter_provider") as set_provider: + otel_setup.setup_otel_metrics(settings, log) + + set_provider.assert_not_called() + + +def test_setup_otel_traces_skips_when_both_disabled() -> None: + settings = MagicMock() + settings.ENABLE_OTEL_METRICS = False + settings.OTEL_EXPORTER_OTLP_ENDPOINT = "" + settings.otel_traces_active.return_value = False + settings.resolved_otel_traces_endpoint.return_value = "" + log = MagicMock() + app = MagicMock() + + with patch.object(otel_setup, "_instrument_fastapi") as instrument: + otel_setup.setup_otel_traces(app, settings, log) + + instrument.assert_not_called() + + +def test_setup_otel_metrics_is_idempotent() -> None: + otel_setup._metrics_initialized = False + settings = MagicMock() + settings.ENABLE_OTEL_METRICS = True + settings.OTEL_EXPORTER_OTLP_ENDPOINT = "otel-gateway:4327" + settings.OTEL_SERVICE_NAME = "template-agent" + settings.OTEL_METRIC_EXPORT_INTERVAL_MILLIS = 10000 + settings.OTEL_AUTH_TOKEN = "" + log = MagicMock() + + with patch("opentelemetry.metrics.set_meter_provider") as set_provider: + otel_setup.setup_otel_metrics(settings, log) + otel_setup.setup_otel_metrics(settings, log) + + set_provider.assert_called_once() + otel_setup._metrics_initialized = False diff --git a/tests/unit/streaming/test_streaming.py b/tests/unit/streaming/test_streaming.py new file mode 100644 index 00000000..c76baff3 --- /dev/null +++ b/tests/unit/streaming/test_streaming.py @@ -0,0 +1,622 @@ +"""Unit tests for streaming components.""" + +import pytest +from langchain_core.messages import AIMessage, ToolMessage +from langgraph.types import Overwrite + +from deep_agent.src.streaming import ( + MessageDeduplicator, + StreamContext, + ToolCallTracker, + remove_tool_calls, +) +from deep_agent.src.streaming.converter import ( + convert_message_to_api_format, + should_skip_message, +) +from deep_agent.src.streaming.handlers import ( + TokenEventHandler, + UpdateEventHandler, +) + + +@pytest.fixture +def stream_context(): + """Fixture providing a standard StreamContext for tests.""" + return StreamContext( + run_id="test_run_1", + trace_id="test_trace_1", + thread_id="test_thread_1", + session_id="test_session_1", + user_id="test_user", + stream_tokens=True, + ) + + +@pytest.fixture +def deduplicator(): + """Fixture providing a fresh MessageDeduplicator.""" + return MessageDeduplicator() + + +@pytest.fixture +def tracker(): + """Fixture providing a fresh ToolCallTracker.""" + return ToolCallTracker() + + +class TestMessageDeduplicator: + """Tests for MessageDeduplicator component.""" + + def test_mark_and_check_seen(self, deduplicator): + """Test marking messages as seen and checking if seen.""" + msg = AIMessage(content="Hello", id="msg_1") + + assert not deduplicator.is_seen(msg) + deduplicator.mark_seen(msg) + assert deduplicator.is_seen(msg) + + def test_get_unseen_messages(self, deduplicator): + """Test getting only unseen messages.""" + msg1 = AIMessage(content="Hello", id="msg_1") + msg2 = AIMessage(content="World", id="msg_2") + msg3 = AIMessage(content="Hello again", id="msg_1") # Duplicate ID + + messages = [msg1, msg2, msg3] + unseen = deduplicator.get_unseen_messages(messages) + + assert len(unseen) == 2 # msg1 and msg2 are new + assert unseen[0].id == "msg_1" + assert unseen[1].id == "msg_2" + + # Second call should return empty since all are seen + unseen_again = deduplicator.get_unseen_messages(messages) + assert len(unseen_again) == 0 + + def test_tool_message_uses_tool_call_id(self, deduplicator): + """Test that ToolMessage without id uses tool_call_id.""" + tool_msg = ToolMessage(content="result", tool_call_id="tc_1", name="test_tool") + + assert not deduplicator.is_seen(tool_msg) + deduplicator.mark_seen(tool_msg) + assert deduplicator.is_seen(tool_msg) + + def test_message_without_id_always_unseen(self, deduplicator): + """Test that messages without stable IDs are never marked as seen.""" + msg_no_id = AIMessage(content="No ID") + + # Should always be unseen since no stable ID + assert not deduplicator.is_seen(msg_no_id) + deduplicator.mark_seen(msg_no_id) + # Still not seen because no ID to track + assert not deduplicator.is_seen(msg_no_id) + + def test_reset_clears_seen_messages(self, deduplicator): + """Test that reset clears all seen message IDs.""" + msg = AIMessage(content="Hello", id="msg_1") + + deduplicator.mark_seen(msg) + assert deduplicator.is_seen(msg) + + deduplicator.reset() + assert not deduplicator.is_seen(msg) + + def test_populate_from_history(self, deduplicator): + """Test pre-populating seen IDs from message history.""" + msg1 = AIMessage(content="Old message 1", id="msg_1") + msg2 = AIMessage(content="Old message 2", id="msg_2") + history = [msg1, msg2] + + deduplicator.populate_from_history(history) + + assert deduplicator.is_seen(msg1) + assert deduplicator.is_seen(msg2) + + +class TestToolCallTracker: + """Tests for ToolCallTracker component.""" + + def test_track_tool_call_from_updates(self, tracker): + """Test tracking tool call ID from updates stream mode.""" + event = { + "agent": { + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": "test_tool", "args": {}, "id": "tc_123"}], + ) + ] + } + } + + tracker.update_from_stream_event("updates", event) + assert tracker.current_id == "tc_123" + + def test_track_tool_response_from_updates(self, tracker): + """Test tracking tool response ID from updates stream mode.""" + event = { + "agent": { + "messages": [ + ToolMessage( + content="result", tool_call_id="tc_456", name="test_tool" + ) + ] + } + } + + tracker.update_from_stream_event("updates", event) + assert tracker.current_id == "tc_456" + + def test_track_from_message_stream(self, tracker): + """Test tracking from messages stream mode.""" + msg = AIMessage( + content="", + tool_calls=[{"name": "test_tool", "args": {}, "id": "tc_789"}], + ) + event = (msg, {}) + + tracker.update_from_stream_event("messages", event) + assert tracker.current_id == "tc_789" + + def test_extract_tool_call_id(self): + """Test extracting tool call ID directly from message.""" + from langchain_core.messages import AIMessageChunk + + from deep_agent.src.streaming.tracker import extract_tool_call_id + + msg = AIMessageChunk( + content="", + tool_calls=[{"name": "test_tool", "args": {}, "id": "tc_abc"}], + ) + + tool_id = extract_tool_call_id(msg) + assert tool_id == "tc_abc" + + def test_reset_clears_current_id(self, tracker): + """Test that reset clears the current tool call ID.""" + event = { + "agent": { + "messages": [ + AIMessage( + content="", + tool_calls=[{"name": "test_tool", "args": {}, "id": "tc_123"}], + ) + ] + } + } + + tracker.update_from_stream_event("updates", event) + assert tracker.current_id == "tc_123" + + tracker.reset() + assert tracker.current_id is None + + +class TestConverter: + """Tests for message conversion utilities.""" + + def test_should_skip_empty_tool_message(self): + """Test that empty tool messages are skipped.""" + empty_tool_msg = ToolMessage(content="", tool_call_id="tc_1", name="empty_tool") + should_skip, reason = should_skip_message(empty_tool_msg) + + assert should_skip + assert "empty result" in reason + assert "empty_tool" in reason + + def test_should_skip_malformed_function_call(self): + """Test that malformed function call messages are skipped.""" + malformed_msg = AIMessage( + content="", + response_metadata={"finish_reason": "MALFORMED_FUNCTION_CALL"}, + ) + should_skip, reason = should_skip_message(malformed_msg) + + assert should_skip + assert "MALFORMED_FUNCTION_CALL" in reason + + def test_should_not_skip_normal_message(self): + """Test that normal messages are not skipped.""" + normal_msg = AIMessage(content="Hello, how can I help?") + should_skip, reason = should_skip_message(normal_msg) + + assert not should_skip + assert reason is None + + def test_should_not_skip_ai_message_with_tool_calls(self): + """Test that AI messages with tool calls are not skipped even if empty content.""" + msg_with_tool = AIMessage( + content="", + tool_calls=[{"name": "test_tool", "args": {}, "id": "tc_1"}], + ) + should_skip, reason = should_skip_message(msg_with_tool) + + assert not should_skip + + def test_convert_message_to_api_format(self, stream_context): + """Test conversion of chat message to simplified format.""" + + class MockChatMessage: + def __init__(self): + self.type = "ai" + self.content = "Hello, how can I help?" + self.tool_calls = None + self.tool_call_id = None + self.response_metadata = {"model": "test-model"} + + chat_msg = MockChatMessage() + result = convert_message_to_api_format(chat_msg, stream_context) + + assert result["type"] == "ai" + assert result["content"] == "Hello, how can I help?" + # run_id and trace_id come from stream context (authoritative) + assert result["run_id"] == "test_run_1" + assert result["trace_id"] == "test_trace_1" + assert result["thread_id"] == "test_thread_1" + assert result["session_id"] == "test_session_1" + assert result["user_id"] == "test_user" + assert result["response_metadata"] == {"model": "test-model"} + + def test_convert_includes_trace_id_from_context(self, stream_context): + """Test that trace_id and run_id come from stream context (authoritative).""" + + class MockChatMessage: + def __init__(self): + self.type = "ai" + self.content = "Test" + self.tool_calls = None + self.tool_call_id = None + self.response_metadata = {} + + chat_msg = MockChatMessage() + result = convert_message_to_api_format(chat_msg, stream_context) + + # Verify all context metadata is included (authoritative for the stream) + assert result["run_id"] == stream_context.run_id + assert result["trace_id"] == stream_context.trace_id + assert result["thread_id"] == stream_context.thread_id + assert result["session_id"] == stream_context.session_id + assert result["user_id"] == stream_context.user_id + + def test_convert_with_tool_calls(self, stream_context): + """Test conversion with tool calls, including subagent name rewriting.""" + + class MockChatMessage: + def __init__(self): + self.type = "ai" + self.content = "" + self.tool_calls = [ + { + "name": "task", + "args": {"subagent_type": "research_agent", "query": "test"}, + "id": "tc_1", + } + ] + self.tool_call_id = None + self.run_id = "test_run_1" + self.trace_id = "test_trace_1" + self.response_metadata = {} + + chat_msg = MockChatMessage() + result = convert_message_to_api_format(chat_msg, stream_context) + + # Should rewrite "task" to actual subagent name + assert result["tool_calls"][0]["name"] == "research_agent" + assert result["tool_calls"][0]["args"]["subagent_type"] == "research_agent" + + def test_remove_tool_calls_string_content(self): + """Test that remove_tool_calls returns string content unchanged.""" + content = "Hello, how can I help?" + result = remove_tool_calls(content) + + assert result == "Hello, how can I help?" + assert isinstance(result, str) + + def test_remove_tool_calls_filters_tool_use(self): + """Test that remove_tool_calls filters out tool_use items from list content.""" + content = [ + {"type": "text", "text": "Let me search for that"}, + {"type": "tool_use", "name": "search", "id": "tc_1"}, + {"type": "text", "text": "..."}, + ] + result = remove_tool_calls(content) + + assert len(result) == 2 + assert result[0]["type"] == "text" + assert result[0]["text"] == "Let me search for that" + assert result[1]["type"] == "text" + assert result[1]["text"] == "..." + + def test_remove_tool_calls_preserves_string_items(self): + """Test that remove_tool_calls preserves string items in list content.""" + content = [ + "Plain string", + {"type": "text", "text": "Dict content"}, + {"type": "tool_use", "name": "search", "id": "tc_1"}, + ] + result = remove_tool_calls(content) + + assert len(result) == 2 + assert result[0] == "Plain string" + assert result[1]["type"] == "text" + + def test_remove_tool_calls_empty_list(self): + """Test that remove_tool_calls handles empty list.""" + content = [] + result = remove_tool_calls(content) + + assert result == [] + assert isinstance(result, list) + + +class TestTokenEventHandler: + """Tests for TokenEventHandler.""" + + def test_handle_basic_token_streaming(self, tracker, stream_context): + """Test basic token streaming functionality.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + msg = AIMessageChunk(content="Hello") + event = (msg, {}) + + events = handler.handle(event, stream_context) + + assert len(events) == 1 + assert events[0]["type"] == "token" + assert events[0]["content"] == "Hello" + + def test_respects_stream_tokens_flag(self, tracker, stream_context): + """Test that handler respects ctx.stream_tokens flag.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + # Create context with stream_tokens=False + no_stream_ctx = StreamContext( + run_id="r1", + trace_id="tr1", + thread_id="t1", + session_id="s1", + user_id="u1", + stream_tokens=False, + ) + + msg = AIMessageChunk(content="Hello") + event = (msg, {}) + + events = handler.handle(event, no_stream_ctx) + + # Should return empty list when stream_tokens is False + assert len(events) == 0 + + def test_skips_messages_with_skip_stream_tag(self, tracker, stream_context): + """Test that messages with skip_stream tag are filtered out.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + msg = AIMessageChunk(content="Hello") + event = (msg, {"tags": ["skip_stream"]}) + + events = handler.handle(event, stream_context) + + assert len(events) == 0 + + def test_filters_non_ai_message_chunks(self, tracker, stream_context): + """Test that non-AIMessageChunk messages are filtered.""" + handler = TokenEventHandler(tracker) + + # Regular AIMessage (not chunk) + msg = AIMessage(content="Hello") + event = (msg, {}) + + events = handler.handle(event, stream_context) + + assert len(events) == 0 + + def test_filters_empty_content(self, tracker, stream_context): + """Test that messages with empty content are filtered.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + msg = AIMessageChunk(content="") + event = (msg, {}) + + events = handler.handle(event, stream_context) + + assert len(events) == 0 + + def test_removes_tool_calls_from_content(self, tracker, stream_context): + """Test that tool calls are removed from streamed content.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + # Content that includes tool calls (which should be removed) + msg = AIMessageChunk( + content=[ + {"type": "text", "text": "Let me help you"}, + {"type": "tool_use", "name": "search", "id": "tc_1"}, + ] + ) + event = (msg, {}) + + events = handler.handle(event, stream_context) + + assert len(events) == 1 + assert events[0]["content"] == "Let me help you" + + def test_associates_tool_call_id_from_message(self, tracker, stream_context): + """Test that tool call ID is extracted from message.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + msg = AIMessageChunk( + content="Searching...", + tool_calls=[{"name": "search", "args": {}, "id": "tc_123"}], + ) + event = (msg, {}) + + events = handler.handle(event, stream_context) + + assert len(events) == 1 + assert events[0]["tool_call_id"] == "tc_123" + + def test_associates_tool_call_id_from_tracker(self, tracker, stream_context): + """Test that tool call ID is taken from tracker if not in message.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + # Set tracker's current_id + tracker._current_tool_call_id = "tc_456" + + msg = AIMessageChunk(content="Result from tool") + event = (msg, {}) + + events = handler.handle(event, stream_context) + + assert len(events) == 1 + assert events[0]["tool_call_id"] == "tc_456" + + def test_no_tool_call_id_when_none_available(self, tracker, stream_context): + """Test that tool_call_id is not added when none is available.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + # Reset tracker to ensure no current_id + tracker.reset() + + msg = AIMessageChunk(content="Hello") + event = (msg, {}) + + events = handler.handle(event, stream_context) + + assert len(events) == 1 + assert "tool_call_id" not in events[0] + + def test_prefers_message_tool_id_over_tracker(self, tracker, stream_context): + """Test that message tool_call_id takes precedence over tracker.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + # Set tracker's current_id + tracker._current_tool_call_id = "tc_old" + + # Message has its own tool call + msg = AIMessageChunk( + content="Searching...", + tool_calls=[{"name": "search", "args": {}, "id": "tc_new"}], + ) + event = (msg, {}) + + events = handler.handle(event, stream_context) + + assert len(events) == 1 + # Should use the message's tool_call_id, not tracker's + assert events[0]["tool_call_id"] == "tc_new" + + def test_handles_tool_call_chunks(self, tracker, stream_context): + """Test handling of tool_call_chunks during streaming.""" + from langchain_core.messages import AIMessageChunk + + handler = TokenEventHandler(tracker) + + msg = AIMessageChunk( + content="", + tool_call_chunks=[{"name": "search", "args": "{}", "id": "tc_789"}], + ) + event = (msg, {}) + + # Should filter out empty content even if tool_call_chunks present + events = handler.handle(event, stream_context) + + assert len(events) == 0 + + +class TestUpdateEventHandler: + """Tests for UpdateEventHandler.""" + + def test_handle_interrupt_event(self, deduplicator, stream_context): + """Test handling of interrupt events.""" + handler = UpdateEventHandler(deduplicator) + + event = { + "__interrupt__": [ + type("Interrupt", (), {"value": "Please confirm action"})() + ] + } + + events = handler.handle(event, stream_context) + + assert len(events) == 1 + assert events[0]["type"] == "message" + assert events[0]["content"]["content"] == "Please confirm action" + + def test_handle_regular_messages(self, deduplicator, stream_context): + """Test handling of regular message updates.""" + handler = UpdateEventHandler(deduplicator) + + msg = AIMessage(content="Hello", id="msg_1") + event = {"agent": {"messages": [msg]}} + + events = handler.handle(event, stream_context) + + assert len(events) == 1 + assert events[0]["type"] == "message" + assert events[0]["content"]["content"] == "Hello" + assert events[0]["content"]["thread_id"] == "test_thread_1" + + def test_handle_overwrite_deduplication(self, deduplicator, stream_context): + """Test that Overwrite events are properly deduplicated.""" + handler = UpdateEventHandler(deduplicator) + + msg1 = AIMessage(content="Hello", id="msg_1") + msg2 = AIMessage(content="World", id="msg_2") + + # First, process msg1 normally + event1 = {"agent": {"messages": [msg1]}} + events = handler.handle(event1, stream_context) + assert len(events) == 1 + + # Then send Overwrite with full history + overwrite_event = {"agent": {"messages": Overwrite([msg1, msg2])}} + events = handler.handle(overwrite_event, stream_context) + + # Should only get msg2 since msg1 was already seen + assert len(events) == 1 + assert events[0]["content"]["content"] == "World" + + def test_handle_empty_tool_message_logs_warning(self, deduplicator, stream_context): + """Test that empty tool messages are skipped with warning.""" + handler = UpdateEventHandler(deduplicator) + + empty_tool = ToolMessage(content="", tool_call_id="tc_1", name="test_tool") + event = {"agent": {"messages": [empty_tool]}} + + events = handler.handle(event, stream_context) + + # Should be filtered out + assert len(events) == 0 + + def test_handle_multiple_nodes(self, deduplicator, stream_context): + """Test handling events from multiple nodes.""" + handler = UpdateEventHandler(deduplicator) + + event = { + "node1": {"messages": [AIMessage(content="From node 1", id="msg_1")]}, + "node2": {"messages": [AIMessage(content="From node 2", id="msg_2")]}, + } + + events = handler.handle(event, stream_context) + + assert len(events) == 2 + contents = [e["content"]["content"] for e in events] + assert "From node 1" in contents + assert "From node 2" in contents diff --git a/tests/unit/test_error_handling.py b/tests/unit/test_error_handling.py new file mode 100644 index 00000000..712b79d8 --- /dev/null +++ b/tests/unit/test_error_handling.py @@ -0,0 +1,434 @@ +"""Unit tests for error_handling module. + +Tests cover: +- classify_error: all 4 classification branches +- with_fallback: sync, async, selective exception catching +- CircuitBreaker (in-memory): full closed→open→half-open→closed lifecycle +- CircuitBreaker (Redis-backed): mocked Redis hash operations +""" + +import asyncio +import time +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.error_handling import ( + CircuitBreaker, + classify_error, + create_circuit_breaker, + with_fallback, +) +from deep_agent.src.exceptions import ( + AuthenticationError, + ConfigurationError, + LLMError, + MCPError, + RateLimitError, + SubAgentError, +) + +# ─────────────────────────────────────────────────────────────────── +# classify_error +# ─────────────────────────────────────────────────────────────────── + + +class TestClassifyError: + """Tests for classify_error — 4 branches.""" + + def test_rate_limit_error(self): + result = classify_error(RateLimitError("quota exceeded")) + assert result["recoverable"] is True + assert result["error_type"] == "rate_limit" + assert "rate limit" in result["message"].lower() + + def test_transient_error(self): + result = classify_error(LLMError("model unavailable")) + assert result["recoverable"] is True + assert result["error_type"] == "transient" + assert "unavailable" in result["message"].lower() + + def test_transient_mcp_error(self): + result = classify_error(MCPError("connection refused")) + assert result["recoverable"] is True + assert result["error_type"] == "transient" + + def test_app_exception_non_transient(self): + result = classify_error(SubAgentError("build failed")) + assert result["recoverable"] is False + assert result["error_type"] == "E_006" + + def test_app_exception_config_error(self): + result = classify_error(ConfigurationError("missing key")) + assert result["recoverable"] is False + assert result["message"] == "Configuration Initialization Failed" + + def test_app_exception_auth_error(self): + result = classify_error(AuthenticationError("bad token")) + assert result["recoverable"] is False + assert result["error_type"] == "E_010" + + def test_unknown_exception(self): + result = classify_error(RuntimeError("something unexpected")) + assert result["recoverable"] is False + assert result["error_type"] == "unknown" + assert result["message"] == "Internal server error" + + def test_base_exception_treated_as_unknown(self): + result = classify_error(TypeError("bad type")) + assert result["error_type"] == "unknown" + + def test_rate_limit_before_transient(self): + """RateLimitError IS a TransientError, but classify_error checks it first.""" + result = classify_error(RateLimitError("429")) + assert result["error_type"] == "rate_limit" + assert result["recoverable"] is True + + +# ─────────────────────────────────────────────────────────────────── +# with_fallback +# ─────────────────────────────────────────────────────────────────── + + +class TestWithFallback: + """Tests for with_fallback decorator.""" + + def test_sync_returns_normal_result(self): + @with_fallback("default") + def good() -> str: + return "real" + + assert good() == "real" + + def test_sync_returns_fallback_on_exception(self): + @with_fallback("default") + def bad() -> str: + raise ValueError("boom") + + assert bad() == "default" + + def test_sync_selective_catch(self): + """Only catches specified exception types.""" + + @with_fallback("default", on=(ValueError,)) + def bad() -> str: + raise TypeError("wrong type") + + with pytest.raises(TypeError, match="wrong type"): + bad() + + def test_sync_selective_catch_matches(self): + @with_fallback("default", on=(ValueError,)) + def bad() -> str: + raise ValueError("expected") + + assert bad() == "default" + + def test_async_returns_normal_result(self): + @with_fallback("default") + async def good() -> str: + return "real" + + assert asyncio.run(good()) == "real" + + def test_async_returns_fallback_on_exception(self): + @with_fallback("default") + async def bad() -> str: + raise RuntimeError("async boom") + + assert asyncio.run(bad()) == "default" + + def test_async_selective_catch(self): + @with_fallback("default", on=(ValueError,)) + async def bad() -> str: + raise TypeError("wrong type") + + with pytest.raises(TypeError): + asyncio.run(bad()) + + def test_fallback_with_none_value(self): + @with_fallback(None) + def bad() -> str | None: + raise ValueError("boom") + + assert bad() is None + + def test_fallback_with_list_value(self): + @with_fallback([]) + def bad() -> list[str]: + raise ValueError("boom") + + assert bad() == [] + + def test_preserves_function_name(self): + @with_fallback("x") + def my_function() -> str: + return "y" + + assert my_function.__name__ == "my_function" + + def test_async_preserves_function_name(self): + @with_fallback("x") + async def my_async_fn() -> str: + return "y" + + assert my_async_fn.__name__ == "my_async_fn" + + +# ─────────────────────────────────────────────────────────────────── +# CircuitBreaker — in-memory +# ─────────────────────────────────────────────────────────────────── + + +class TestCircuitBreakerInMemory: + """Tests for CircuitBreaker with in-memory backend (no Redis).""" + + def test_starts_closed(self): + cb = CircuitBreaker("test", threshold=3, reset_timeout=10.0) + assert cb.state == "closed" + assert cb.is_open is False + + def test_stays_closed_below_threshold(self): + cb = CircuitBreaker("test", threshold=3, reset_timeout=10.0) + cb.record_failure() + cb.record_failure() + assert cb.state == "closed" + assert cb.is_open is False + + def test_opens_at_threshold(self): + cb = CircuitBreaker("test", threshold=3, reset_timeout=10.0) + cb.record_failure() + cb.record_failure() + cb.record_failure() + assert cb.state == "open" + assert cb.is_open is True + + def test_success_resets_failures(self): + cb = CircuitBreaker("test", threshold=3, reset_timeout=10.0) + cb.record_failure() + cb.record_failure() + cb.record_success() + assert cb.state == "closed" + cb.record_failure() + assert cb.state == "closed" + + def test_success_closes_open_circuit(self): + cb = CircuitBreaker("test", threshold=2, reset_timeout=10.0) + cb.record_failure() + cb.record_failure() + assert cb.state == "open" + cb.record_success() + assert cb.state == "closed" + + def test_half_open_after_timeout(self): + cb = CircuitBreaker("test", threshold=2, reset_timeout=0.05) + cb.record_failure() + cb.record_failure() + assert cb.is_open is True + + time.sleep(0.06) + assert cb.is_open is False + assert cb.state == "half-open" + + def test_full_lifecycle(self): + """closed → open → half-open → closed (after success).""" + cb = CircuitBreaker("test", threshold=2, reset_timeout=0.05) + + assert cb.state == "closed" + + cb.record_failure() + cb.record_failure() + assert cb.state == "open" + assert cb.is_open is True + + time.sleep(0.06) + assert cb.is_open is False + assert cb.state == "half-open" + + cb.record_success() + assert cb.state == "closed" + assert cb.is_open is False + + def test_half_open_reopens_on_failure(self): + """half-open → open if probe fails.""" + cb = CircuitBreaker("test", threshold=2, reset_timeout=0.05) + cb.record_failure() + cb.record_failure() + + time.sleep(0.06) + assert cb.state == "half-open" + + cb.record_failure() + cb.record_failure() + assert cb.state == "open" + + def test_threshold_one(self): + cb = CircuitBreaker("test", threshold=1, reset_timeout=10.0) + cb.record_failure() + assert cb.state == "open" + + def test_default_parameters(self): + cb = CircuitBreaker("defaults") + assert cb.threshold == 5 + assert cb.reset_timeout == 60.0 + assert cb.name == "defaults" + + +# ─────────────────────────────────────────────────────────────────── +# CircuitBreaker — Redis-backed (mocked) +# ─────────────────────────────────────────────────────────────────── + + +class TestCircuitBreakerRedis: + """Tests for CircuitBreaker with Redis backend (mocked).""" + + def _make_redis_mock(self) -> MagicMock: + """Create a mock Redis client that behaves like a real hash store.""" + store: dict[str, dict[str, str]] = {} + + mock = MagicMock() + + def hgetall(key: str) -> dict[str, str]: + return store.get(key, {}) + + def hset(key: str, mapping: dict[str, str]) -> int: + if key not in store: + store[key] = {} + store[key].update({k: str(v) for k, v in mapping.items()}) + return len(mapping) + + def delete(key: str) -> int: + return 1 if store.pop(key, None) is not None else 0 + + mock.hgetall = MagicMock(side_effect=hgetall) + mock.hset = MagicMock(side_effect=hset) + mock.delete = MagicMock(side_effect=delete) + mock.expire = MagicMock(return_value=True) + mock.ping = MagicMock(return_value=True) + mock._store = store + return mock + + def test_starts_closed(self): + mock_redis = self._make_redis_mock() + cb = CircuitBreaker("test", threshold=3, redis_client=mock_redis) + assert cb.state == "closed" + assert cb.is_open is False + + def test_opens_at_threshold(self): + mock_redis = self._make_redis_mock() + cb = CircuitBreaker("test", threshold=3, redis_client=mock_redis) + cb.record_failure() + cb.record_failure() + cb.record_failure() + assert cb.state == "open" + + def test_success_resets(self): + mock_redis = self._make_redis_mock() + cb = CircuitBreaker("test", threshold=3, redis_client=mock_redis) + cb.record_failure() + cb.record_failure() + cb.record_success() + assert cb.state == "closed" + + def test_redis_state_shared_across_instances(self): + """Two CircuitBreaker instances sharing the same Redis see the same state.""" + mock_redis = self._make_redis_mock() + cb1 = CircuitBreaker("shared", threshold=2, redis_client=mock_redis) + cb2 = CircuitBreaker("shared", threshold=2, redis_client=mock_redis) + + cb1.record_failure() + cb1.record_failure() + assert cb2.state == "open" + + def test_redis_half_open_after_timeout(self): + mock_redis = self._make_redis_mock() + cb = CircuitBreaker( + "test", threshold=2, reset_timeout=0.05, redis_client=mock_redis + ) + cb.record_failure() + cb.record_failure() + assert cb.is_open is True + + time.sleep(0.06) + assert cb.is_open is False + assert cb.state == "half-open" + + def test_redis_full_lifecycle(self): + mock_redis = self._make_redis_mock() + cb = CircuitBreaker( + "test", threshold=2, reset_timeout=0.05, redis_client=mock_redis + ) + + cb.record_failure() + cb.record_failure() + assert cb.state == "open" + + time.sleep(0.06) + assert cb.state == "half-open" + + cb.record_success() + assert cb.state == "closed" + + def test_redis_error_falls_back_to_closed(self): + """If Redis raises, circuit breaker should degrade to closed (allow requests).""" + mock_redis = MagicMock() + mock_redis.hgetall = MagicMock(side_effect=ConnectionError("Redis down")) + mock_redis.hset = MagicMock(side_effect=ConnectionError("Redis down")) + mock_redis.delete = MagicMock(side_effect=ConnectionError("Redis down")) + + cb = CircuitBreaker("test", threshold=2, redis_client=mock_redis) + cb.record_failure() + assert cb.is_open is False + + def test_expire_called_on_write(self): + """Every write should refresh the key TTL.""" + mock_redis = self._make_redis_mock() + cb = CircuitBreaker( + "test", threshold=3, reset_timeout=60.0, redis_client=mock_redis + ) + cb.record_failure() + mock_redis.expire.assert_called_once_with(cb._redis_key, cb._key_ttl) + + def test_ttl_minimum_300s(self): + """TTL should be at least 300s even for tiny reset_timeout.""" + mock_redis = self._make_redis_mock() + cb = CircuitBreaker( + "test", threshold=2, reset_timeout=1.0, redis_client=mock_redis + ) + assert cb._key_ttl == 300 + + def test_ttl_scales_with_reset_timeout(self): + """TTL = 3 * reset_timeout when that exceeds 300s.""" + mock_redis = self._make_redis_mock() + cb = CircuitBreaker( + "test", threshold=2, reset_timeout=200.0, redis_client=mock_redis + ) + assert cb._key_ttl == 600 + + +# ─────────────────────────────────────────────────────────────────── +# create_circuit_breaker factory +# ─────────────────────────────────────────────────────────────────── + + +class TestCreateCircuitBreaker: + """Tests for the factory function.""" + + def test_creates_in_memory_when_no_redis(self): + with patch("deep_agent.src.error_handling.get_redis_client", return_value=None): + cb = create_circuit_breaker("test", threshold=3) + assert cb._redis is None + + def test_creates_redis_backed_when_available(self): + mock_redis = MagicMock() + with patch( + "deep_agent.src.error_handling.get_redis_client", + return_value=mock_redis, + ): + cb = create_circuit_breaker("test", threshold=3) + assert cb._redis is mock_redis + + def test_explicit_redis_client_overrides_auto_detect(self): + mock_redis = MagicMock() + cb = create_circuit_breaker("test", threshold=3, redis_client=mock_redis) + assert cb._redis is mock_redis diff --git a/tests/unit/test_exceptions.py b/tests/unit/test_exceptions.py new file mode 100644 index 00000000..d42c7745 --- /dev/null +++ b/tests/unit/test_exceptions.py @@ -0,0 +1,218 @@ +"""Unit tests for exception hierarchy and error codes.""" + +import pytest +from starlette.status import ( + HTTP_401_UNAUTHORIZED, + HTTP_429_TOO_MANY_REQUESTS, + HTTP_500_INTERNAL_SERVER_ERROR, + HTTP_502_BAD_GATEWAY, + HTTP_503_SERVICE_UNAVAILABLE, + HTTP_504_GATEWAY_TIMEOUT, +) + +from deep_agent.src.exceptions import ( + AppException, + AuthenticationError, + ConfigurationError, + ErrorCode, + ErrorCodes, + LLMError, + LLMTimeoutError, + MCPError, + MCPTimeoutError, + RateLimitError, + SubAgentError, + TransientError, +) + + +class TestErrorCode: + """Tests for ErrorCode dataclass.""" + + def test_create_error_code(self): + """Test creating an ErrorCode instance.""" + code = ErrorCode(status=404, message="Not Found", code="E_404") + + assert code.status == 404 + assert code.message == "Not Found" + assert code.code == "E_404" + + def test_error_code_is_frozen(self): + """Test that ErrorCode instances are immutable.""" + code = ErrorCode(status=500, message="Server Error", code="E_500") + + with pytest.raises(Exception): + code.status = 400 + + +class TestErrorCodes: + """Tests for ErrorCodes constants.""" + + def test_internal_server_error(self): + error = ErrorCodes.INTERNAL_SERVER_ERROR + assert error.status == HTTP_500_INTERNAL_SERVER_ERROR + assert error.message == "Internal Server Error" + assert error.code == "E_001" + + def test_llm_error(self): + error = ErrorCodes.LLM_ERROR + assert error.status == HTTP_502_BAD_GATEWAY + assert error.code == "E_002" + + def test_llm_timeout(self): + error = ErrorCodes.LLM_TIMEOUT + assert error.status == HTTP_504_GATEWAY_TIMEOUT + assert error.code == "E_003" + + def test_mcp_connection_error(self): + error = ErrorCodes.MCP_CONNECTION_ERROR + assert error.status == HTTP_502_BAD_GATEWAY + assert error.message == "MCP Connection Failed" + assert error.code == "E_004" + + def test_mcp_timeout(self): + error = ErrorCodes.MCP_TIMEOUT + assert error.status == HTTP_504_GATEWAY_TIMEOUT + assert error.code == "E_005" + + def test_subagent_error(self): + error = ErrorCodes.SUBAGENT_ERROR + assert error.status == HTTP_500_INTERNAL_SERVER_ERROR + assert error.code == "E_006" + + def test_configuration_initialization_error(self): + error = ErrorCodes.CONFIGURATION_INITIALIZATION_ERROR + assert error.status == HTTP_500_INTERNAL_SERVER_ERROR + assert error.message == "Configuration Initialization Failed" + assert error.code == "E_007" + + def test_configuration_validation_error(self): + error = ErrorCodes.CONFIGURATION_VALIDATION_ERROR + assert error.status == HTTP_500_INTERNAL_SERVER_ERROR + assert error.message == "Configuration Validation Failed" + assert error.code == "E_008" + + def test_rate_limit_error(self): + error = ErrorCodes.RATE_LIMIT_ERROR + assert error.status == HTTP_429_TOO_MANY_REQUESTS + assert error.code == "E_009" + + def test_authentication_error(self): + error = ErrorCodes.AUTHENTICATION_ERROR + assert error.status == HTTP_401_UNAUTHORIZED + assert error.code == "E_010" + + def test_service_unavailable(self): + error = ErrorCodes.SERVICE_UNAVAILABLE + assert error.status == HTTP_503_SERVICE_UNAVAILABLE + assert error.code == "E_011" + + def test_legacy_alias_mcp(self): + """Legacy PRODUCTION_MCP_CONNECTION_ERROR aliases MCP_CONNECTION_ERROR.""" + assert ( + ErrorCodes.PRODUCTION_MCP_CONNECTION_ERROR + is ErrorCodes.MCP_CONNECTION_ERROR + ) + + def test_error_codes_are_frozen(self): + with pytest.raises(Exception): + ErrorCodes.INTERNAL_SERVER_ERROR.status = 400 + + +class TestAppException: + """Tests for AppException class.""" + + def test_create_with_default_error_code(self): + exc = AppException("Something went wrong") + assert str(exc) == "Something went wrong" + assert exc.detail == "Something went wrong" + assert exc.status == HTTP_500_INTERNAL_SERVER_ERROR + assert exc.message == "Internal Server Error" + assert exc.code == "E_001" + + def test_create_with_custom_error_code(self): + exc = AppException("MCP unreachable", ErrorCodes.MCP_CONNECTION_ERROR) + assert exc.status == HTTP_502_BAD_GATEWAY + assert exc.message == "MCP Connection Failed" + assert exc.code == "E_004" + + def test_is_retryable_default_false(self): + exc = AppException("error") + assert exc.is_retryable is False + + def test_exception_is_raisable(self): + with pytest.raises(AppException) as exc_info: + raise AppException("Test error", ErrorCodes.INTERNAL_SERVER_ERROR) + assert exc_info.value.detail == "Test error" + assert exc_info.value.code == "E_001" + + def test_exception_preserves_traceback(self): + try: + raise AppException("Error with traceback") + except AppException as exc: + assert exc.detail == "Error with traceback" + import traceback + + tb = traceback.format_exc() + assert "AppException" in tb + assert "Error with traceback" in tb + + +class TestTransientError: + """Tests for TransientError and retryable subclasses.""" + + def test_transient_is_retryable(self): + exc = TransientError("transient failure") + assert exc.is_retryable is True + + def test_llm_error(self): + exc = LLMError("model creation failed") + assert isinstance(exc, TransientError) + assert isinstance(exc, AppException) + assert exc.is_retryable is True + assert exc.code == "E_002" + + def test_llm_timeout_error(self): + exc = LLMTimeoutError("request timed out") + assert exc.is_retryable is True + assert exc.code == "E_003" + + def test_mcp_error(self): + exc = MCPError("connection refused") + assert isinstance(exc, TransientError) + assert exc.is_retryable is True + assert exc.code == "E_004" + + def test_mcp_timeout_error(self): + exc = MCPTimeoutError("timeout") + assert exc.is_retryable is True + assert exc.code == "E_005" + + def test_rate_limit_error(self): + exc = RateLimitError("too many requests") + assert isinstance(exc, TransientError) + assert exc.is_retryable is True + assert exc.code == "E_009" + + +class TestNonRetryableErrors: + """Tests for non-retryable exception subclasses.""" + + def test_subagent_error(self): + exc = SubAgentError("failed to build") + assert isinstance(exc, AppException) + assert not isinstance(exc, TransientError) + assert exc.is_retryable is False + assert exc.code == "E_006" + + def test_configuration_error(self): + exc = ConfigurationError("missing config") + assert isinstance(exc, AppException) + assert exc.is_retryable is False + assert exc.code == "E_007" + + def test_authentication_error(self): + exc = AuthenticationError("invalid token") + assert isinstance(exc, AppException) + assert exc.is_retryable is False + assert exc.code == "E_010" diff --git a/tests/unit/test_hitl.py b/tests/unit/test_hitl.py new file mode 100644 index 00000000..40d5cd4e --- /dev/null +++ b/tests/unit/test_hitl.py @@ -0,0 +1,90 @@ +"""Unit tests for the HITL interrupt_on builder.""" + +from unittest.mock import MagicMock + +import pytest + +from deep_agent.src.agent.config.hitl import ( + _DEEPAGENTS_BUILTIN_TOOLS, + build_interrupt_on, +) +from deep_agent.src.agent.config.middleware import HumanApprovalConfig + + +def _tool(name: str) -> MagicMock: + t = MagicMock() + t.name = name + return t + + +class TestBuildInterruptOn: + def test_disabled_returns_empty(self): + config = HumanApprovalConfig(enabled=False, mode="all") + result = build_interrupt_on( + config, [_tool("send_email"), _tool("delete_record")] + ) + assert result == {} + + def test_mode_none_returns_empty(self): + config = HumanApprovalConfig(enabled=True, mode="none") + result = build_interrupt_on(config, [_tool("send_email")]) + assert result == {} + + def test_mode_all_includes_explicit_tools(self): + tools = [_tool("send_email"), _tool("search_web"), _tool("delete_record")] + config = HumanApprovalConfig(enabled=True, mode="all") + result = build_interrupt_on(config, tools) + assert result["send_email"] is True + assert result["search_web"] is True + assert result["delete_record"] is True + + def test_mode_all_always_includes_builtins(self): + """Built-in deepagents tools must be interrupted even with no explicit tools.""" + config = HumanApprovalConfig(enabled=True, mode="all") + result = build_interrupt_on(config, []) + for builtin in _DEEPAGENTS_BUILTIN_TOOLS: + assert builtin in result, ( + f"built-in tool '{builtin}' missing from interrupt_on" + ) + + def test_empty_tool_list_still_covers_builtins(self): + """An agent with no MCP tools still gets HITL for built-in filesystem tools.""" + config = HumanApprovalConfig(enabled=True, mode="all") + result = build_interrupt_on(config, []) + assert len(result) == len(_DEEPAGENTS_BUILTIN_TOOLS) + assert result == {name: True for name in _DEEPAGENTS_BUILTIN_TOOLS} + + def test_exclude_removes_listed_tools(self): + tools = [_tool("send_email"), _tool("search_web"), _tool("health_check")] + config = HumanApprovalConfig( + enabled=True, + mode="all", + exclude=["health_check", "search_web", "ls", "read_file"], + ) + result = build_interrupt_on(config, tools) + assert result.get("send_email") is True + assert "health_check" not in result + assert "search_web" not in result + assert "ls" not in result + assert "read_file" not in result + + def test_exclude_nonexistent_tool_is_harmless(self): + tools = [_tool("send_email")] + config = HumanApprovalConfig(enabled=True, mode="all", exclude=["nonexistent"]) + result = build_interrupt_on(config, tools) + assert result["send_email"] is True + # builtins still present + assert "ls" in result + + def test_all_tools_excluded_returns_empty(self): + all_names = list(_DEEPAGENTS_BUILTIN_TOOLS) + ["send_email"] + tools = [_tool("send_email")] + config = HumanApprovalConfig(enabled=True, mode="all", exclude=all_names) + result = build_interrupt_on(config, tools) + assert result == {} + + def test_default_config_disabled(self): + """Default HumanApprovalConfig should produce no interrupts.""" + config = HumanApprovalConfig() + result = build_interrupt_on(config, [_tool("send_email")]) + assert result == {} diff --git a/tests/unit/test_infrastructure_middleware.py b/tests/unit/test_infrastructure_middleware.py new file mode 100644 index 00000000..7a82d77f --- /dev/null +++ b/tests/unit/test_infrastructure_middleware.py @@ -0,0 +1,162 @@ +"""Unit tests for the middleware builder module.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.agent.config.middleware import ResolvedMiddlewareConfig +from deep_agent.src.infrastructure.middleware import ( + _import_middleware, + build_excluded_middleware, + build_middleware_list, + resolve_memory_param, +) + + +class TestBuildMiddlewareList: + """Test middleware instance construction from resolved config.""" + + @pytest.fixture(autouse=True) + def _disable_audit(self): + with patch( + "deep_agent.src.audit.config.is_audit_enabled", + return_value=False, + ): + yield + + def test_returns_empty_when_master_switch_off(self): + resolved = ResolvedMiddlewareConfig(summarization_tool_enabled=True) + with patch( + "deep_agent.src.infrastructure.middleware.settings" + ) as mock_settings: + mock_settings.MIDDLEWARE_ENABLED = False + result = build_middleware_list(resolved) + assert result == [] + + def test_includes_summarization_tool_when_enabled(self): + resolved = ResolvedMiddlewareConfig(summarization_tool_enabled=True) + mock_mw = MagicMock() + with ( + patch("deep_agent.src.infrastructure.middleware.settings") as mock_settings, + patch( + "deep_agent.src.infrastructure.middleware._build_summarization_tool_middleware", + return_value=mock_mw, + ), + ): + mock_settings.MIDDLEWARE_ENABLED = True + result = build_middleware_list(resolved) + assert mock_mw in result + + def test_excludes_summarization_tool_when_disabled(self): + resolved = ResolvedMiddlewareConfig( + summarization_tool_enabled=False, extra_middleware=[] + ) + with ( + patch("deep_agent.src.infrastructure.middleware.settings") as mock_settings, + patch( + "deep_agent.src.infrastructure.middleware._build_summarization_tool_middleware", + ) as build_sum, + ): + mock_settings.MIDDLEWARE_ENABLED = True + result = build_middleware_list(resolved) + build_sum.assert_not_called() + # Default guardrails (model/tool limits + model retry) still apply. + assert len(result) == 3 + + def test_includes_extra_middleware(self): + resolved = ResolvedMiddlewareConfig( + summarization_tool_enabled=False, + extra_middleware=[ + "tests.unit.test_infrastructure_middleware:_DummyMiddleware" + ], + ) + with patch( + "deep_agent.src.infrastructure.middleware.settings" + ) as mock_settings: + mock_settings.MIDDLEWARE_ENABLED = True + result = build_middleware_list(resolved) + assert len(result) == 4 + assert any(isinstance(m, _DummyMiddleware) for m in result) + + +class TestBuildExcludedMiddleware: + """Test excluded middleware list generation.""" + + def test_empty_when_all_enabled(self): + resolved = ResolvedMiddlewareConfig( + patch_tool_calls_enabled=True, excluded_middleware=[] + ) + result = build_excluded_middleware(resolved) + assert result == [] + + def test_includes_patch_tool_calls_when_disabled(self): + resolved = ResolvedMiddlewareConfig( + patch_tool_calls_enabled=False, excluded_middleware=[] + ) + result = build_excluded_middleware(resolved) + assert "PatchToolCallsMiddleware" in result + + def test_preserves_profile_exclusions(self): + resolved = ResolvedMiddlewareConfig( + patch_tool_calls_enabled=True, + excluded_middleware=["SomeCustomMiddleware"], + ) + result = build_excluded_middleware(resolved) + assert "SomeCustomMiddleware" in result + + +class TestResolveMemoryParam: + """Test memory parameter resolution for create_deep_agent().""" + + def test_returns_none_when_master_disabled(self): + resolved = ResolvedMiddlewareConfig(memory_enabled=True) + with patch( + "deep_agent.src.infrastructure.middleware.settings" + ) as mock_settings: + mock_settings.MIDDLEWARE_ENABLED = False + result = resolve_memory_param(resolved) + assert result is None + + def test_returns_none_when_memory_disabled(self): + resolved = ResolvedMiddlewareConfig(memory_enabled=False) + with patch( + "deep_agent.src.infrastructure.middleware.settings" + ) as mock_settings: + mock_settings.MIDDLEWARE_ENABLED = True + result = resolve_memory_param(resolved) + assert result is None + + def test_returns_namespaces_when_enabled(self): + resolved = ResolvedMiddlewareConfig( + memory_enabled=True, memory_namespaces=["user_mem", "shared"] + ) + with patch( + "deep_agent.src.infrastructure.middleware.settings" + ) as mock_settings: + mock_settings.MIDDLEWARE_ENABLED = True + result = resolve_memory_param(resolved) + assert result == ["user_mem", "shared"] + + +class TestImportMiddleware: + """Test dynamic middleware importing.""" + + def test_invalid_path_without_colon(self): + result = _import_middleware("no_colon_here") + assert result is None + + def test_nonexistent_module(self): + result = _import_middleware("nonexistent.module:Class") + assert result is None + + def test_valid_import(self): + result = _import_middleware( + "tests.unit.test_infrastructure_middleware:_DummyMiddleware" + ) + assert result is not None + + +class _DummyMiddleware: + """Test fixture — a no-op middleware class.""" + + pass diff --git a/tests/unit/test_personalization.py b/tests/unit/test_personalization.py new file mode 100644 index 00000000..581aa442 --- /dev/null +++ b/tests/unit/test_personalization.py @@ -0,0 +1,68 @@ +"""Unit tests for personalization models and injector.""" + +import uuid +from datetime import datetime + +import pytest + +from deep_agent.src.personalization.injector import inject_personalization +from deep_agent.src.personalization.models import Memory, Rule + + +class TestMemoryModel: + def test_create_with_defaults(self): + m = Memory(user_id="u1", content="Likes Python") + assert m.user_id == "u1" + assert m.content == "Likes Python" + assert isinstance(m.id, uuid.UUID) + assert isinstance(m.created_at, datetime) + + def test_create_with_explicit_id(self): + uid = uuid.uuid4() + m = Memory(id=uid, user_id="u1", content="test") + assert m.id == uid + + +class TestRuleModel: + def test_create_with_defaults(self): + r = Rule(user_id="u1", content="Be concise") + assert r.is_active is True + + def test_inactive_rule(self): + r = Rule(user_id="u1", content="Old rule", is_active=False) + assert r.is_active is False + + +class TestInjectPersonalization: + def test_no_personalization(self): + result = inject_personalization("Base prompt", [], []) + assert result == "Base prompt" + + def test_memories_only(self): + result = inject_personalization("Base", ["Likes Python", "Uses Linux"], []) + assert "User Memories" in result + assert "Likes Python" in result + assert "Uses Linux" in result + assert "Custom Instructions" not in result + + def test_rules_only(self): + result = inject_personalization("Base", [], ["Be concise", "Use code blocks"]) + assert "Custom Instructions" in result + assert "Be concise" in result + assert "User Memories" not in result + + def test_both_memories_and_rules(self): + result = inject_personalization( + "Base prompt", + ["Prefers dark mode"], + ["Always use TypeScript"], + ) + assert "User Memories" in result + assert "Custom Instructions" in result + assert "Prefers dark mode" in result + assert "Always use TypeScript" in result + assert result.startswith("Base prompt") + + def test_separator_between_sections(self): + result = inject_personalization("Base", ["m1"], ["r1"]) + assert "---" in result diff --git a/tests/unit/test_pylogger.py b/tests/unit/test_pylogger.py new file mode 100644 index 00000000..d7bf1d44 --- /dev/null +++ b/tests/unit/test_pylogger.py @@ -0,0 +1,102 @@ +"""Unit tests for structured logging utility.""" + +from unittest.mock import patch + +from deep_agent.utils.pylogger import ( + _inject_request_context, + bind_request_context, + clear_request_context, + force_reconfigure_all_loggers, + get_python_logger, + get_uvicorn_log_config, +) + + +class TestGetPythonLogger: + def test_returns_bound_logger(self): + logger = get_python_logger("INFO") + assert logger is not None + + def test_idempotent(self): + a = get_python_logger("DEBUG") + b = get_python_logger("DEBUG") + assert a is not None + assert b is not None + + +class TestForceReconfigure: + def test_reconfigures(self): + force_reconfigure_all_loggers("WARNING") + logger = get_python_logger() + assert logger is not None + + +class TestRequestContext: + def setup_method(self): + clear_request_context() + + def teardown_method(self): + clear_request_context() + + def test_bind_and_inject(self): + bind_request_context( + trace_id="req-123", + user_id="user-456", + thread_id="thread-789", + ) + event: dict = {"event": "test"} + result = _inject_request_context(None, "info", event) + assert result["trace_id"] == "req-123" + assert result["user_id"] == "user-456" + assert result["thread_id"] == "thread-789" + assert result["service"] == "template-agent" + + def test_inject_without_bind(self): + event: dict = {"event": "test"} + result = _inject_request_context(None, "info", event) + assert "trace_id" not in result + assert "user_id" not in result + assert "service" in result + + def test_clear_resets(self): + bind_request_context(trace_id="req-x") + clear_request_context() + event: dict = {"event": "test"} + result = _inject_request_context(None, "info", event) + assert "trace_id" not in result + + def test_partial_bind(self): + bind_request_context(user_id="u1") + event: dict = {"event": "test"} + result = _inject_request_context(None, "info", event) + assert result["user_id"] == "u1" + assert "trace_id" not in result + + +class TestConsoleRenderer: + def test_json_format_default(self): + with patch("deep_agent.utils.pylogger.LOG_FORMAT", "json"): + from deep_agent.utils.pylogger import _get_renderer + + renderer = _get_renderer() + assert "JSON" in type(renderer).__name__ + + def test_console_format(self): + with patch("deep_agent.utils.pylogger.LOG_FORMAT", "console"): + from deep_agent.utils.pylogger import _get_renderer + + renderer = _get_renderer() + assert "Console" in type(renderer).__name__ + + +class TestUvicornLogConfig: + def test_returns_valid_config(self): + config = get_uvicorn_log_config("INFO") + assert config["version"] == 1 + assert "formatters" in config + assert "handlers" in config + assert "loggers" in config + + def test_respects_log_level(self): + config = get_uvicorn_log_config("DEBUG") + assert config["loggers"][""]["level"] == "DEBUG" diff --git a/tests/unit/test_repository.py b/tests/unit/test_repository.py new file mode 100644 index 00000000..ed53068d --- /dev/null +++ b/tests/unit/test_repository.py @@ -0,0 +1,226 @@ +"""Unit tests for PersonalizationRepository (mocked DB).""" + +import uuid +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from deep_agent.src.personalization.models import Memory, Rule +from deep_agent.src.personalization.repository import PersonalizationRepository + + +@pytest.fixture(autouse=True) +def _reset_tables_flag(): + """Reset the module-level _TABLES_ENSURED flag before each test.""" + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = False + yield + repo_mod._TABLES_ENSURED = False + + +@pytest.fixture +def mock_conn(): + """Create a mock async connection context manager.""" + conn = AsyncMock() + cursor = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[]) + cursor.rowcount = 0 + conn.execute = AsyncMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=False) + conn._cursor = cursor + return conn + + +@pytest.fixture +def repo(): + return PersonalizationRepository("postgresql://test:test@localhost/testdb") + + +class TestEnsureTables: + @pytest.mark.asyncio + async def test_creates_tables_once(self, repo, mock_conn): + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + await repo.ensure_tables() + assert mock_conn.execute.call_count == 3 + mock_conn.commit.assert_awaited_once() + + @pytest.mark.asyncio + async def test_skips_if_already_ensured(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + await repo.ensure_tables() + mock_conn.execute.assert_not_called() + + +class TestListMemories: + @pytest.mark.asyncio + async def test_returns_memories(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + + mem_data = { + "id": uuid.uuid4(), + "user_id": "u1", + "content": "Likes Python", + "created_at": "2025-01-01T00:00:00+00:00", + "updated_at": "2025-01-01T00:00:00+00:00", + } + mock_conn._cursor.fetchall = AsyncMock(return_value=[mem_data]) + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + memories = await repo.list_memories("u1") + assert len(memories) == 1 + assert memories[0].content == "Likes Python" + + @pytest.mark.asyncio + async def test_empty_list(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + memories = await repo.list_memories("nobody") + assert memories == [] + + +class TestCreateMemory: + @pytest.mark.asyncio + async def test_creates_and_returns(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + memory = await repo.create_memory("u1", "Likes Python") + assert memory.user_id == "u1" + assert memory.content == "Likes Python" + mock_conn.execute.assert_awaited_once() + mock_conn.commit.assert_awaited_once() + + +class TestDeleteMemory: + @pytest.mark.asyncio + async def test_delete_returns_true_when_found(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + mock_conn._cursor.rowcount = 1 + mock_conn.execute.return_value = mock_conn._cursor + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + result = await repo.delete_memory("u1", uuid.uuid4()) + assert result is True + + @pytest.mark.asyncio + async def test_delete_returns_false_when_not_found(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + mock_conn._cursor.rowcount = 0 + mock_conn.execute.return_value = mock_conn._cursor + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + result = await repo.delete_memory("u1", uuid.uuid4()) + assert result is False + + +class TestListRules: + @pytest.mark.asyncio + async def test_returns_rules_active_only(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + + rule_data = { + "id": uuid.uuid4(), + "user_id": "u1", + "content": "Be concise", + "is_active": True, + "created_at": "2025-01-01T00:00:00+00:00", + "updated_at": "2025-01-01T00:00:00+00:00", + } + mock_conn._cursor.fetchall = AsyncMock(return_value=[rule_data]) + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + rules = await repo.list_rules("u1", active_only=True) + assert len(rules) == 1 + assert rules[0].content == "Be concise" + + @pytest.mark.asyncio + async def test_returns_all_rules(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + rules = await repo.list_rules("u1", active_only=False) + assert rules == [] + + +class TestUpsertRule: + @pytest.mark.asyncio + async def test_creates_new_rule(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + rule = await repo.upsert_rule("u1", "Be concise") + assert rule.user_id == "u1" + assert rule.content == "Be concise" + assert rule.is_active is True + mock_conn.commit.assert_awaited_once() + + +class TestDeleteRule: + @pytest.mark.asyncio + async def test_delete_returns_true(self, repo, mock_conn): + import deep_agent.src.personalization.repository as repo_mod + + repo_mod._TABLES_ENSURED = True + mock_conn._cursor.rowcount = 1 + mock_conn.execute.return_value = mock_conn._cursor + + with patch( + "deep_agent.src.personalization.repository.psycopg.AsyncConnection.connect", + return_value=mock_conn, + ): + result = await repo.delete_rule("u1", uuid.uuid4()) + assert result is True diff --git a/tests/unit/test_schema.py b/tests/unit/test_schema.py new file mode 100644 index 00000000..cacdb50d --- /dev/null +++ b/tests/unit/test_schema.py @@ -0,0 +1,109 @@ +"""Unit tests for schema models.""" + +import pytest + +from deep_agent.src.schema import ( + ChatHistoryResponse, + ChatMessage, + FeedbackRequest, + FeedbackResponse, + StreamRequest, + UserInput, +) + + +class TestUserInput: + def test_required_message(self): + inp = UserInput(message="hello") + assert inp.message == "hello" + + def test_optional_fields_default_none(self): + inp = UserInput(message="hi") + assert inp.thread_id is None + assert inp.session_id is None + assert inp.user_id is None + + def test_all_fields(self): + inp = UserInput( + message="hello", + thread_id="t1", + session_id="s1", + user_id="u1", + ) + assert inp.thread_id == "t1" + assert inp.session_id == "s1" + assert inp.user_id == "u1" + + +class TestStreamRequest: + def test_inherits_user_input(self): + req = StreamRequest(message="test") + assert isinstance(req, UserInput) + + def test_default_stream_tokens(self): + req = StreamRequest(message="test") + assert req.stream_tokens is True + + def test_stream_tokens_false(self): + req = StreamRequest(message="test", stream_tokens=False) + assert req.stream_tokens is False + + +class TestChatMessage: + def test_minimal_message(self): + msg = ChatMessage(type="human", content="hello") + assert msg.type == "human" + assert msg.content == "hello" + assert msg.tool_calls == [] + assert msg.tool_call_id is None + assert msg.run_id is None + assert msg.response_metadata == {} + assert msg.custom_data == {} + + def test_ai_message_with_tool_calls(self): + msg = ChatMessage( + type="ai", + content="", + tool_calls=[{"name": "search", "args": {"q": "test"}, "id": "tc1"}], + ) + assert msg.tool_calls[0]["name"] == "search" + + def test_allowed_types(self): + for t in ("human", "ai", "tool", "custom"): + msg = ChatMessage(type=t, content="x") + assert msg.type == t + + +class TestFeedbackRequest: + def test_required_fields(self): + fb = FeedbackRequest(trace_id="abc", name="thumbs-up", value=1.0) + assert fb.trace_id == "abc" + assert fb.name == "thumbs-up" + assert fb.value == 1.0 + assert fb.kwargs == {} + + def test_with_kwargs(self): + fb = FeedbackRequest( + trace_id="abc", + name="rating", + value=0.8, + kwargs={"comment": "good"}, + ) + assert fb.kwargs["comment"] == "good" + + +class TestFeedbackResponse: + def test_default_status(self): + resp = FeedbackResponse() + assert resp.status == "success" + + +class TestChatHistoryResponse: + def test_empty_messages(self): + resp = ChatHistoryResponse(messages=[]) + assert resp.messages == [] + + def test_with_messages(self): + msg = ChatMessage(type="human", content="hi") + resp = ChatHistoryResponse(messages=[msg]) + assert len(resp.messages) == 1 diff --git a/tests/unit/test_settings.py b/tests/unit/test_settings.py new file mode 100644 index 00000000..ba5a4f37 --- /dev/null +++ b/tests/unit/test_settings.py @@ -0,0 +1,124 @@ +"""Unit tests for settings module.""" + +import pytest + +from deep_agent.src.exceptions import AppException +from deep_agent.src.settings import Settings, validate_config + + +class TestSettings: + """Tests for Settings Pydantic model.""" + + def test_default_values(self): + s = Settings() + assert s.AGENT_HOST == "0.0.0.0" + assert s.AGENT_PORT == 5002 + assert s.PYTHON_LOG_LEVEL == "INFO" + assert s.POSTGRES_USER == "pgvector" + assert s.POSTGRES_PORT == 5432 + assert s.MAX_OUTPUT_TOKENS == 8192 + + def test_database_uri(self): + s = Settings( + POSTGRES_USER="u", + POSTGRES_PASSWORD="p", + POSTGRES_HOST="h", + POSTGRES_PORT=1234, + POSTGRES_DB="d", + ) + assert s.database_uri == "postgresql://u:p@h:1234/d" + + def test_ssl_keyfile_none_when_empty(self): + s = Settings(SSL_KEYFILE="") + assert s.get_ssl_keyfile_path is None + + def test_ssl_keyfile_returns_path(self): + s = Settings(SSL_KEYFILE="/path/to/key") + assert s.get_ssl_keyfile_path == "/path/to/key" + + def test_ssl_certfile_none_when_empty(self): + s = Settings(SSL_CERTFILE="") + assert s.get_ssl_certfile_path is None + + def test_ssl_certfile_returns_path(self): + s = Settings(SSL_CERTFILE="/path/to/cert") + assert s.get_ssl_certfile_path == "/path/to/cert" + + def test_optional_fields_accept_none(self): + s = Settings( + LANGFUSE_PUBLIC_KEY=None, + LANGFUSE_SECRET_KEY=None, + LANGFUSE_BASE_URL=None, + GOOGLE_APPLICATION_CREDENTIALS_CONTENT=None, + ) + assert s.LANGFUSE_PUBLIC_KEY is None + assert s.LANGFUSE_SECRET_KEY is None + assert s.LANGFUSE_BASE_URL is None + assert s.GOOGLE_APPLICATION_CREDENTIALS_CONTENT is None + + def test_request_logging_defaults(self): + s = Settings() + assert s.REQUEST_LOGGING_ENABLED is True + assert s.REQUEST_LOG_HEADERS is True + assert s.REQUEST_LOG_BODY is False + assert s.REQUEST_LOG_BODY_MAX_SIZE == 10240 + + +class TestValidateConfig: + """Tests for validate_config function.""" + + def test_valid_config(self): + s = Settings(AGENT_PORT=5002, PYTHON_LOG_LEVEL="INFO") + validate_config(s) + + def test_port_too_low(self): + s = Settings(AGENT_PORT=80) + with pytest.raises(AppException, match="AGENT_PORT must be between"): + validate_config(s) + + def test_port_too_high(self): + s = Settings(AGENT_PORT=70000) + with pytest.raises(AppException, match="AGENT_PORT must be between"): + validate_config(s) + + def test_invalid_log_level(self): + s = Settings(PYTHON_LOG_LEVEL="VERBOSE") + with pytest.raises(AppException, match="PYTHON_LOG_LEVEL must be one of"): + validate_config(s) + + def test_all_valid_log_levels(self): + for level in ("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"): + s = Settings(PYTHON_LOG_LEVEL=level) + validate_config(s) + + def test_port_boundary_low(self): + s = Settings(AGENT_PORT=1024) + validate_config(s) + + def test_port_boundary_high(self): + s = Settings(AGENT_PORT=65535) + validate_config(s) + + +class TestValidateConfigPublicBaseUrl: + def test_allows_localhost_http_base_url(self): + validate_config( + Settings(AGENT_PUBLIC_BASE_URL="http://localhost:5002", AGENT_PORT=5002) + ) + + def test_requires_https_for_production_base_url(self): + with pytest.raises( + AppException, match="AGENT_PUBLIC_BASE_URL must use https://" + ): + validate_config(Settings(AGENT_PUBLIC_BASE_URL="http://agent.example.com")) + + def test_allows_https_production_base_url(self): + validate_config(Settings(AGENT_PUBLIC_BASE_URL="https://agent.example.com")) + + def test_oauth_callback_url_derived_from_public_base_url(self): + s = Settings(AGENT_PUBLIC_BASE_URL="https://agent.example.com") + assert s.oauth_callback_url == "https://agent.example.com/mcp/oauth/callback" + + def test_oauth_callback_url_defaults_to_localhost(self): + s = Settings(AGENT_PORT=5002) + assert s.oauth_callback_url == "http://localhost:5002/mcp/oauth/callback" diff --git a/tests/unit/token_budget/test_callback.py b/tests/unit/token_budget/test_callback.py new file mode 100644 index 00000000..1e766add --- /dev/null +++ b/tests/unit/token_budget/test_callback.py @@ -0,0 +1,27 @@ +"""Unit tests for token budget callback.""" + +from __future__ import annotations + +from deep_agent.src.token_budget.callback import ( + thread_id_from_metadata, + user_id_from_metadata, +) + + +def test_thread_id_from_metadata_prefers_token_budget_key() -> None: + assert thread_id_from_metadata({"token_budget_thread_id": "abc"}) == "abc" + + +def test_thread_id_from_metadata_falls_back_to_langfuse_session() -> None: + assert thread_id_from_metadata({"langfuse_session_id": "xyz"}) == "xyz" + + +def test_thread_id_from_metadata_missing() -> None: + assert thread_id_from_metadata({}) is None + assert thread_id_from_metadata(None) is None + + +def test_user_id_from_metadata() -> None: + assert user_id_from_metadata({"token_budget_user_id": "dev-user"}) == "dev-user" + assert user_id_from_metadata({}) is None + assert user_id_from_metadata(None) is None diff --git a/tests/unit/token_budget/test_mongo_repository.py b/tests/unit/token_budget/test_mongo_repository.py new file mode 100644 index 00000000..07123d39 --- /dev/null +++ b/tests/unit/token_budget/test_mongo_repository.py @@ -0,0 +1,105 @@ +"""Unit tests for Mongo token usage repository retries.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, PropertyMock, patch + +import pytest +from pymongo.errors import ServerSelectionTimeoutError + +from deep_agent.src.token_budget.mongo_repository import TokenUsageMongoRepository + + +@pytest.mark.asyncio +async def test_increment_usage_retries_transient_mongo_error() -> None: + expected = { + "thread_id": "thread-1", + "total_tokens": 150, + "input_tokens": 100, + "output_tokens": 50, + } + + collection = AsyncMock() + collection.find_one_and_update = AsyncMock( + side_effect=[ + ServerSelectionTimeoutError("timeout"), + expected, + ] + ) + + with ( + patch.object( + TokenUsageMongoRepository, + "_thread_collection", + new_callable=PropertyMock, + return_value=collection, + ), + patch.object(TokenUsageMongoRepository, "ensure_indexes", new=AsyncMock()), + ): + repo = TokenUsageMongoRepository( + "mongodb://mongodb:27017", db_name="tokenusage" + ) + result = await repo.increment_usage( + "thread-1", 100, 50, agent_name="health-assistant" + ) + + assert result == expected + assert collection.find_one_and_update.await_count == 2 + + +@pytest.mark.asyncio +async def test_increment_usage_does_not_retry_runtime_error() -> None: + collection = AsyncMock() + collection.find_one_and_update = AsyncMock(return_value=None) + + with ( + patch.object( + TokenUsageMongoRepository, + "_thread_collection", + new_callable=PropertyMock, + return_value=collection, + ), + patch.object(TokenUsageMongoRepository, "ensure_indexes", new=AsyncMock()), + ): + repo = TokenUsageMongoRepository( + "mongodb://mongodb:27017", db_name="tokenusage" + ) + with pytest.raises(RuntimeError, match="Failed to increment Mongo token usage"): + await repo.increment_usage("thread-1", 100, 50) + + assert collection.find_one_and_update.await_count == 1 + + +@pytest.mark.asyncio +async def test_ensure_indexes_runs_once_per_process() -> None: + import deep_agent.src.token_budget.mongo_repository as mongo_module + + mongo_module._INDEXES_ENSURED = False + + thread_collection = AsyncMock() + daily_collection = AsyncMock() + + with ( + patch.object( + TokenUsageMongoRepository, + "_thread_collection", + new_callable=PropertyMock, + return_value=thread_collection, + ), + patch.object( + TokenUsageMongoRepository, + "_daily_collection", + new_callable=PropertyMock, + return_value=daily_collection, + ), + ): + repo = TokenUsageMongoRepository( + "mongodb://mongodb:27017", db_name="tokenusage" + ) + await repo.ensure_indexes() + await repo.ensure_indexes() + + assert thread_collection.create_index.await_count == 2 + assert daily_collection.create_index.await_count == 2 + + mongo_module._INDEXES_ENSURED = False diff --git a/tests/unit/token_budget/test_otel.py b/tests/unit/token_budget/test_otel.py new file mode 100644 index 00000000..5cf266d4 --- /dev/null +++ b/tests/unit/token_budget/test_otel.py @@ -0,0 +1,217 @@ +"""Unit tests for token budget OTEL emission.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from deep_agent.src.token_budget import otel_emit + + +@pytest.fixture(autouse=True) +def reset_otel_emit_state() -> None: + otel_emit._counters_initialized = False + otel_emit._token_counter = None + otel_emit._thread_total_counter = None + otel_emit._daily_total_counter = None + + +def test_token_budget_otel_enabled_requires_metrics_flag_and_endpoint() -> None: + mock_settings = MagicMock() + mock_settings.ENABLE_OTEL_METRICS = True + mock_settings.OTEL_EXPORTER_OTLP_ENDPOINT = "otel-gateway:4327" + with patch("deep_agent.src.token_budget.otel_emit.settings", mock_settings): + assert otel_emit.token_budget_otel_enabled() is True + + mock_settings.ENABLE_OTEL_METRICS = False + with patch("deep_agent.src.token_budget.otel_emit.settings", mock_settings): + assert otel_emit.token_budget_otel_enabled() is False + + +def test_emit_token_usage_skipped_when_metrics_disabled() -> None: + mock_settings = MagicMock() + mock_settings.ENABLE_OTEL_METRICS = False + mock_settings.OTEL_EXPORTER_OTLP_ENDPOINT = "" + with ( + patch("deep_agent.src.token_budget.otel_emit.settings", mock_settings), + patch.object(otel_emit.logger, "info") as log_info, + ): + otel_emit.emit_token_usage( + thread_id="thread-1", + user_id="user-1", + input_tokens=10, + output_tokens=5, + cumulative_total=15, + cumulative_input=10, + cumulative_output=5, + ) + + log_info.assert_not_called() + + +def test_emit_token_usage_logs_expected_payload() -> None: + mock_settings = MagicMock() + mock_settings.ENABLE_OTEL_METRICS = True + mock_settings.OTEL_EXPORTER_OTLP_ENDPOINT = "otel-gateway:4327" + mock_settings.ENABLE_OTEL_TRACES = False + mock_settings.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT = "" + + with ( + patch("deep_agent.src.token_budget.otel_emit.settings", mock_settings), + patch( + "deep_agent.src.token_budget.otel_emit._agent_name", + return_value="health-assistant", + ), + patch.object(otel_emit, "_ensure_counters"), + patch.object(otel_emit.logger, "info") as log_info, + ): + otel_emit.emit_token_usage( + thread_id="thread-abc", + user_id="user-xyz", + input_tokens=100, + output_tokens=25, + cumulative_total=125, + cumulative_input=100, + cumulative_output=25, + timestamp="2026-06-23T12:00:00+00:00", + ) + + log_info.assert_called_once_with( + "token_budget_usage", + thread_id="thread-abc", + user_id="user-xyz", + input_tokens=100, + output_tokens=25, + total_tokens=125, + cumulative_total_tokens=125, + cumulative_input_tokens=100, + cumulative_output_tokens=25, + timestamp="2026-06-23T12:00:00+00:00", + **{"agent.name": "health-assistant"}, + ) + + +def test_emit_token_usage_records_metrics() -> None: + mock_token_counter = MagicMock() + mock_thread_counter = MagicMock() + mock_settings = MagicMock() + mock_settings.ENABLE_OTEL_METRICS = True + mock_settings.OTEL_EXPORTER_OTLP_ENDPOINT = "otel-gateway:4327" + mock_settings.ENABLE_OTEL_TRACES = False + mock_settings.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT = "" + + with ( + patch("deep_agent.src.token_budget.otel_emit.settings", mock_settings), + patch( + "deep_agent.src.token_budget.otel_emit._agent_name", + return_value="health-assistant", + ), + patch.object(otel_emit, "_token_counter", mock_token_counter), + patch.object(otel_emit, "_thread_total_counter", mock_thread_counter), + patch.object(otel_emit, "_counters_initialized", True), + patch.object(otel_emit.logger, "info"), + ): + otel_emit.emit_token_usage( + thread_id="thread-1", + user_id="user-1", + input_tokens=80, + output_tokens=20, + cumulative_total=100, + cumulative_input=80, + cumulative_output=20, + ) + + base_attrs = { + "agent.name": "health-assistant", + "thread_id": "thread-1", + "user_id": "user-1", + } + mock_token_counter.add.assert_any_call(80, {**base_attrs, "token.type": "input"}) + mock_token_counter.add.assert_any_call(20, {**base_attrs, "token.type": "output"}) + mock_thread_counter.add.assert_called_once_with( + 100, + {**base_attrs, "aggregation": "cumulative"}, + ) + + +def test_emit_token_usage_adds_span_event_when_traces_enabled() -> None: + mock_settings = MagicMock() + mock_settings.ENABLE_OTEL_METRICS = True + mock_settings.OTEL_EXPORTER_OTLP_ENDPOINT = "otel-gateway:4327" + mock_settings.ENABLE_OTEL_TRACES = True + mock_settings.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT = "jaeger:4317" + mock_settings.otel_traces_active.return_value = True + mock_settings.resolved_otel_traces_endpoint.return_value = "jaeger:4317" + mock_span = MagicMock() + mock_span.is_recording.return_value = True + + with ( + patch("deep_agent.src.token_budget.otel_emit.settings", mock_settings), + patch( + "deep_agent.src.token_budget.otel_emit._agent_name", + return_value="health-assistant", + ), + patch.object(otel_emit, "_ensure_counters"), + patch.object(otel_emit.logger, "info"), + patch("opentelemetry.trace.get_current_span", return_value=mock_span), + ): + otel_emit.emit_token_usage( + thread_id="thread-1", + user_id=None, + input_tokens=10, + output_tokens=5, + cumulative_total=15, + cumulative_input=10, + cumulative_output=5, + ) + + mock_span.add_event.assert_called_once() + event_name = mock_span.add_event.call_args[0][0] + kwargs = mock_span.add_event.call_args[1] + assert event_name == "token_budget.usage" + assert kwargs["attributes"]["thread_id"] == "thread-1" + assert kwargs["attributes"]["timestamp"] + + +def test_emit_daily_token_usage_logs_expected_payload() -> None: + mock_settings = MagicMock() + mock_settings.ENABLE_OTEL_METRICS = True + mock_settings.OTEL_EXPORTER_OTLP_ENDPOINT = "otel-gateway:4327" + mock_settings.ENABLE_OTEL_TRACES = False + mock_settings.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT = "" + mock_daily_counter = MagicMock() + + with ( + patch("deep_agent.src.token_budget.otel_emit.settings", mock_settings), + patch( + "deep_agent.src.token_budget.otel_emit._agent_name", + return_value="health-assistant", + ), + patch.object(otel_emit, "_daily_total_counter", mock_daily_counter), + patch.object(otel_emit, "_counters_initialized", True), + patch.object(otel_emit.logger, "info") as log_info, + ): + otel_emit.emit_daily_token_usage( + user_id="user-1", + total_tokens=5000, + date="2026-06-23", + timestamp="2026-06-23T18:30:00+00:00", + ) + + log_info.assert_called_once_with( + "token_budget_daily_usage", + user_id="user-1", + total_tokens=5000, + date="2026-06-23", + timestamp="2026-06-23T18:30:00+00:00", + **{"agent.name": "health-assistant"}, + ) + mock_daily_counter.add.assert_called_once_with( + 5000, + { + "agent.name": "health-assistant", + "user_id": "user-1", + "date": "2026-06-23", + }, + ) diff --git a/tests/unit/token_budget/test_service.py b/tests/unit/token_budget/test_service.py new file mode 100644 index 00000000..5531141a --- /dev/null +++ b/tests/unit/token_budget/test_service.py @@ -0,0 +1,232 @@ +"""Unit tests for token budget service.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from unittest.mock import ANY, AsyncMock, MagicMock, patch + +import pytest + +from deep_agent.src.token_budget.config import TokenBudgetConfig +from deep_agent.src.token_budget.service import ( + TokenUsageNotFoundError, + TokenUsageUnavailableError, + _mongo_repo, + check_and_record, + extract_tokens_from_message, + get_thread_token_usage, +) + + +class _FakeMessage: + def __init__(self, usage_metadata: dict | None = None) -> None: + self.usage_metadata = usage_metadata + self.response_metadata = {} + + +def test_extract_tokens_from_message_usage_metadata() -> None: + msg = _FakeMessage({"input_tokens": 100, "output_tokens": 50}) + assert extract_tokens_from_message(msg) == (100, 50) + + +def test_extract_tokens_from_message_includes_reasoning_in_output() -> None: + """Gemini visible output_tokens excludes reasoning; budget must include both.""" + msg = _FakeMessage( + { + "input_tokens": 8567, + "output_tokens": 29, + "total_tokens": 8737, + "output_token_details": {"reasoning": 141}, + } + ) + assert extract_tokens_from_message(msg) == (8567, 170) + assert sum(extract_tokens_from_message(msg)) == 8737 + + +def test_extract_tokens_from_message_zero_input_uses_total_minus_input() -> None: + """Cached prompts may report input_tokens=0 while total_tokens is authoritative.""" + msg = _FakeMessage( + { + "input_tokens": 0, + "output_tokens": 40, + "total_tokens": 100, + } + ) + assert extract_tokens_from_message(msg) == (0, 100) + + +def test_extract_tokens_from_message_zero_input_without_total_uses_output() -> None: + msg = _FakeMessage({"input_tokens": 0, "output_tokens": 50}) + assert extract_tokens_from_message(msg) == (0, 50) + + +@pytest.mark.asyncio +async def test_check_and_record_increments_mongo_and_daily_usage() -> None: + config = TokenBudgetConfig(enabled=True) + row = { + "total_tokens": 150, + "input_tokens": 100, + "output_tokens": 50, + } + + mock_repo = AsyncMock() + mock_repo.increment_usage.return_value = row + mock_repo.increment_daily_usage.return_value = { + "user_id": "user-1", + "total_tokens": 150, + "date": "2026-06-23", + "updated_at": datetime(2026, 6, 23, 12, 0, tzinfo=UTC), + } + + mock_settings = MagicMock() + mock_settings.MONGODB_URI = "mongodb://mongodb:27017" + mock_settings.MONGODB_DB = "tokenusage" + + with ( + patch( + "deep_agent.src.token_budget.service.agent_config.get_token_budget_config", + return_value=config, + ), + patch( + "deep_agent.src.token_budget.service.agent_config.get_name", + return_value="health-assistant", + ), + patch("deep_agent.src.token_budget.service.settings", mock_settings), + patch( + "deep_agent.src.token_budget.service._mongo_repo", + return_value=mock_repo, + ), + patch("deep_agent.src.token_budget.otel_emit.emit_token_usage") as emit_usage, + patch( + "deep_agent.src.token_budget.otel_emit.emit_daily_token_usage" + ) as emit_daily, + ): + await check_and_record("thread-1", 100, 50, user_id="user-1") + + mock_repo.increment_usage.assert_awaited_once_with( + "thread-1", + 100, + 50, + agent_name="health-assistant", + ) + mock_repo.increment_daily_usage.assert_awaited_once_with("user-1", 150) + emit_usage.assert_called_once_with( + thread_id="thread-1", + user_id="user-1", + input_tokens=100, + output_tokens=50, + cumulative_total=150, + cumulative_input=100, + cumulative_output=50, + timestamp=ANY, + trace_id=None, + ) + emit_daily.assert_called_once_with( + user_id="user-1", + total_tokens=150, + date="2026-06-23", + timestamp=ANY, + ) + + +@pytest.mark.asyncio +async def test_check_and_record_skips_daily_without_user_id() -> None: + config = TokenBudgetConfig(enabled=True) + row = { + "total_tokens": 150, + "input_tokens": 100, + "output_tokens": 50, + } + + mock_repo = AsyncMock() + mock_repo.increment_usage.return_value = row + + mock_settings = MagicMock() + mock_settings.MONGODB_URI = "mongodb://mongodb:27017" + mock_settings.MONGODB_DB = "tokenusage" + + with ( + patch( + "deep_agent.src.token_budget.service.agent_config.get_token_budget_config", + return_value=config, + ), + patch( + "deep_agent.src.token_budget.service.agent_config.get_name", + return_value="health-assistant", + ), + patch("deep_agent.src.token_budget.service.settings", mock_settings), + patch( + "deep_agent.src.token_budget.service._mongo_repo", + return_value=mock_repo, + ), + patch("deep_agent.src.token_budget.otel_emit.emit_token_usage"), + patch("deep_agent.src.token_budget.otel_emit.emit_daily_token_usage"), + ): + await check_and_record("thread-1", 100, 50) + + mock_repo.increment_daily_usage.assert_not_awaited() + + +def test_mongo_repo_returns_singleton() -> None: + import deep_agent.src.token_budget.service as service_module + + service_module._mongo_repo_instance = None + mock_settings = MagicMock() + mock_settings.MONGODB_URI = "mongodb://mongodb:27017" + mock_settings.MONGODB_DB = "tokenusage" + + with ( + patch("deep_agent.src.token_budget.service.settings", mock_settings), + patch( + "deep_agent.src.token_budget.mongo_repository.TokenUsageMongoRepository", + ) as repo_cls, + ): + first = _mongo_repo() + second = _mongo_repo() + + assert first is second + repo_cls.assert_called_once_with( + "mongodb://mongodb:27017", + db_name="tokenusage", + ) + service_module._mongo_repo_instance = None + + +@pytest.mark.asyncio +async def test_get_thread_token_usage_raises_when_not_configured() -> None: + config = TokenBudgetConfig(enabled=False) + mock_settings = MagicMock() + mock_settings.MONGODB_URI = "" + + with ( + patch( + "deep_agent.src.token_budget.service.agent_config.get_token_budget_config", + return_value=config, + ), + patch("deep_agent.src.token_budget.service.settings", mock_settings), + ): + with pytest.raises(TokenUsageUnavailableError): + await get_thread_token_usage("thread-1") + + +@pytest.mark.asyncio +async def test_get_thread_token_usage_raises_when_thread_missing() -> None: + config = TokenBudgetConfig(enabled=True) + mock_repo = AsyncMock() + mock_repo.get_thread_usage.return_value = None + mock_settings = MagicMock() + mock_settings.MONGODB_URI = "mongodb://mongodb:27017" + + with ( + patch( + "deep_agent.src.token_budget.service.agent_config.get_token_budget_config", + return_value=config, + ), + patch("deep_agent.src.token_budget.service.settings", mock_settings), + patch( + "deep_agent.src.token_budget.service._mongo_repo", + return_value=mock_repo, + ), + ): + with pytest.raises(TokenUsageNotFoundError): + await get_thread_token_usage("thread-1") diff --git a/tests/unit/triggers/__init__.py b/tests/unit/triggers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/triggers/test_config.py b/tests/unit/triggers/test_config.py new file mode 100644 index 00000000..06f2fd99 --- /dev/null +++ b/tests/unit/triggers/test_config.py @@ -0,0 +1,247 @@ +"""Unit tests for headless mode Pydantic configuration models.""" + +import pytest + +from deep_agent.src.triggers.config import ( + AgentMode, + CronJobConfig, + CronTriggerConfig, + HeadlessConfig, + OutputSinkConfig, + QueueTriggerConfig, + TriggerConfig, + WebhookTriggerConfig, +) + + +class TestAgentMode: + """Test AgentMode enum values.""" + + def test_server_value(self): + assert AgentMode.SERVER == "server" + + def test_headless_value(self): + assert AgentMode.HEADLESS == "headless" + + def test_server_is_str(self): + assert isinstance(AgentMode.SERVER, str) + + def test_headless_is_str(self): + assert isinstance(AgentMode.HEADLESS, str) + + def test_construct_from_string(self): + assert AgentMode("server") is AgentMode.SERVER + assert AgentMode("headless") is AgentMode.HEADLESS + + def test_invalid_value_raises(self): + with pytest.raises(ValueError): + AgentMode("invalid") + + +class TestWebhookTriggerConfig: + """Test WebhookTriggerConfig defaults and custom values.""" + + def test_defaults(self): + cfg = WebhookTriggerConfig() + assert cfg.enabled is False + assert cfg.host == "0.0.0.0" + assert cfg.port == 8888 + assert cfg.path == "/trigger" + + def test_custom_values(self): + cfg = WebhookTriggerConfig( + enabled=True, + host="127.0.0.1", + port=9090, + path="/webhook", + ) + assert cfg.enabled is True + assert cfg.host == "127.0.0.1" + assert cfg.port == 9090 + assert cfg.path == "/webhook" + + +class TestCronJobConfig: + """Test CronJobConfig construction.""" + + def test_required_fields(self): + job = CronJobConfig(name="daily-report", schedule="0 9 * * *") + assert job.name == "daily-report" + assert job.schedule == "0 9 * * *" + assert job.payload == {} + + def test_with_payload(self): + payload = {"region": "us-east-1", "format": "pdf"} + job = CronJobConfig(name="export", schedule="*/5 * * * *", payload=payload) + assert job.payload == payload + + def test_payload_default_factory_isolation(self): + a = CronJobConfig(name="a", schedule="* * * * *") + b = CronJobConfig(name="b", schedule="* * * * *") + a.payload["key"] = "value" + assert "key" not in b.payload + + +class TestCronTriggerConfig: + """Test CronTriggerConfig with jobs list.""" + + def test_defaults(self): + cfg = CronTriggerConfig() + assert cfg.enabled is False + assert cfg.jobs == [] + + def test_with_jobs(self): + jobs = [ + CronJobConfig(name="j1", schedule="0 * * * *"), + CronJobConfig(name="j2", schedule="0 0 * * *", payload={"x": 1}), + ] + cfg = CronTriggerConfig(enabled=True, jobs=jobs) + assert cfg.enabled is True + assert len(cfg.jobs) == 2 + assert cfg.jobs[0].name == "j1" + assert cfg.jobs[1].payload == {"x": 1} + + def test_jobs_default_factory_isolation(self): + a = CronTriggerConfig() + b = CronTriggerConfig() + a.jobs.append(CronJobConfig(name="x", schedule="* * * * *")) + assert len(b.jobs) == 0 + + +class TestQueueTriggerConfig: + """Test QueueTriggerConfig defaults and custom values.""" + + def test_defaults(self): + cfg = QueueTriggerConfig() + assert cfg.enabled is False + assert cfg.backend == "redis_streams" + assert cfg.stream == "agent-tasks" + assert cfg.consumer_group == "agent-workers" + assert cfg.consumer_name == "" + assert cfg.get_consumer_name() == "worker-1" # falls back to default + + def test_custom_values(self): + cfg = QueueTriggerConfig( + enabled=True, + backend="kafka", + stream="custom-stream", + consumer_group="my-group", + consumer_name="worker-42", + ) + assert cfg.enabled is True + assert cfg.backend == "kafka" + assert cfg.stream == "custom-stream" + assert cfg.consumer_group == "my-group" + assert cfg.consumer_name == "worker-42" + + +class TestTriggerConfig: + """Test TriggerConfig nested defaults.""" + + def test_nested_defaults(self): + cfg = TriggerConfig() + assert isinstance(cfg.webhook, WebhookTriggerConfig) + assert isinstance(cfg.cron, CronTriggerConfig) + assert isinstance(cfg.queue, QueueTriggerConfig) + assert cfg.webhook.enabled is False + assert cfg.cron.enabled is False + assert cfg.queue.enabled is False + + def test_override_nested(self): + cfg = TriggerConfig( + webhook=WebhookTriggerConfig(enabled=True, port=7777), + ) + assert cfg.webhook.enabled is True + assert cfg.webhook.port == 7777 + # Other triggers remain default. + assert cfg.cron.enabled is False + + def test_default_factory_isolation(self): + a = TriggerConfig() + b = TriggerConfig() + assert a.webhook is not b.webhook + + +class TestOutputSinkConfig: + """Test OutputSinkConfig with each sink type.""" + + def test_stdout_sink(self): + sink = OutputSinkConfig(type="stdout") + assert sink.type == "stdout" + assert sink.path is None + assert sink.url is None + assert sink.headers == {} + assert sink.stream is None + + def test_file_sink(self): + sink = OutputSinkConfig(type="file", path="/tmp/output.jsonl") + assert sink.type == "file" + assert sink.path == "/tmp/output.jsonl" + + def test_webhook_sink(self): + sink = OutputSinkConfig( + type="webhook", + url="https://example.com/results", + headers={"Authorization": "Bearer tok"}, + ) + assert sink.type == "webhook" + assert sink.url == "https://example.com/results" + assert sink.headers["Authorization"] == "Bearer tok" + + def test_redis_sink(self): + sink = OutputSinkConfig(type="redis", stream="results-stream") + assert sink.type == "redis" + assert sink.stream == "results-stream" + + def test_headers_default_factory_isolation(self): + a = OutputSinkConfig(type="webhook") + b = OutputSinkConfig(type="webhook") + a.headers["X-Custom"] = "value" + assert "X-Custom" not in b.headers + + +class TestHeadlessConfig: + """Test HeadlessConfig full construction and defaults.""" + + def test_defaults(self): + cfg = HeadlessConfig() + assert cfg.mode is AgentMode.SERVER + assert isinstance(cfg.triggers, TriggerConfig) + assert cfg.output_sinks == [] + assert cfg.drain_timeout == 30.0 + + def test_full_construction(self): + cfg = HeadlessConfig( + mode=AgentMode.HEADLESS, + triggers=TriggerConfig( + webhook=WebhookTriggerConfig(enabled=True, port=9000), + cron=CronTriggerConfig( + enabled=True, + jobs=[CronJobConfig(name="nightly", schedule="0 0 * * *")], + ), + queue=QueueTriggerConfig(enabled=True, stream="tasks"), + ), + output_sinks=[ + OutputSinkConfig(type="stdout"), + OutputSinkConfig(type="file", path="/data/out.jsonl"), + ], + drain_timeout=60.0, + ) + assert cfg.mode is AgentMode.HEADLESS + assert cfg.triggers.webhook.enabled is True + assert cfg.triggers.webhook.port == 9000 + assert cfg.triggers.cron.enabled is True + assert len(cfg.triggers.cron.jobs) == 1 + assert cfg.triggers.queue.stream == "tasks" + assert len(cfg.output_sinks) == 2 + assert cfg.drain_timeout == 60.0 + + def test_output_sinks_default_factory_isolation(self): + a = HeadlessConfig() + b = HeadlessConfig() + a.output_sinks.append(OutputSinkConfig(type="stdout")) + assert len(b.output_sinks) == 0 + + def test_mode_from_string(self): + cfg = HeadlessConfig(mode="headless") + assert cfg.mode is AgentMode.HEADLESS diff --git a/tests/unit/triggers/test_cron_source.py b/tests/unit/triggers/test_cron_source.py new file mode 100644 index 00000000..28d38720 --- /dev/null +++ b/tests/unit/triggers/test_cron_source.py @@ -0,0 +1,104 @@ +"""Unit tests for CronTriggerSource.""" + +import asyncio +from unittest.mock import AsyncMock, patch + +from deep_agent.src.triggers.config import CronJobConfig, CronTriggerConfig +from deep_agent.src.triggers.sources.cron import CronTriggerSource, _parse_cron_fields + + +class TestParseCronFields: + """Test the cron expression parser.""" + + def test_valid_5_field_expression(self): + result = _parse_cron_fields("0 9 * * 1-5") + assert result == { + "minute": "0", + "hour": "9", + "day": "*", + "month": "*", + "day_of_week": "1-5", + } + + def test_every_minute(self): + result = _parse_cron_fields("* * * * *") + assert result["minute"] == "*" + + def test_invalid_too_few_fields(self): + assert _parse_cron_fields("0 9 * *") is None + + def test_invalid_too_many_fields(self): + assert _parse_cron_fields("0 9 * * * 2026") is None + + def test_invalid_empty(self): + assert _parse_cron_fields("") is None + + +class TestCronTriggerSource: + """Test the cron trigger source lifecycle.""" + + async def test_start_creates_tasks_for_each_job(self): + jobs = [ + CronJobConfig(name="job-a", schedule="0 * * * *"), + CronJobConfig(name="job-b", schedule="0 0 * * *"), + ] + config = CronTriggerConfig(enabled=True, jobs=jobs) + source = CronTriggerSource(config) + + with patch.object(source, "_run_job", new_callable=AsyncMock): + await source.start() + + assert len(source._tasks) == 2 + await source.stop() + + async def test_stop_cancels_all_tasks(self): + config = CronTriggerConfig( + enabled=True, + jobs=[ + CronJobConfig(name="job", schedule="0 * * * *"), + ], + ) + source = CronTriggerSource(config) + + async def _hang() -> None: + await asyncio.sleep(3600) + + source._tasks = [asyncio.create_task(_hang())] + await source.stop() + assert source._tasks == [] + + async def test_stop_when_not_started_is_safe(self): + config = CronTriggerConfig() + source = CronTriggerSource(config) + await source.stop() + assert source._tasks == [] + + async def test_invalid_cron_schedule_logs_warning(self): + jobs = [ + CronJobConfig(name="bad-job", schedule="not-a-cron"), + CronJobConfig(name="good-job", schedule="0 * * * *"), + ] + config = CronTriggerConfig(enabled=True, jobs=jobs) + + with patch("deep_agent.src.triggers.sources.cron.logger") as mock_logger: + source = CronTriggerSource(config) + + with patch.object(source, "_run_job", new_callable=AsyncMock): + await source.start() + + mock_logger.warning.assert_called_once() + assert "invalid cron schedule" in mock_logger.warning.call_args.args[0] + assert len(source._tasks) == 1 + await source.stop() + + async def test_empty_jobs_list_is_valid(self): + config = CronTriggerConfig(enabled=True, jobs=[]) + source = CronTriggerSource(config) + await source.start() + assert source._tasks == [] + await source.stop() + + async def test_aiter_returns_self(self): + config = CronTriggerConfig() + source = CronTriggerSource(config) + assert source.__aiter__() is source diff --git a/tests/unit/triggers/test_file_sink.py b/tests/unit/triggers/test_file_sink.py new file mode 100644 index 00000000..fe6627f0 --- /dev/null +++ b/tests/unit/triggers/test_file_sink.py @@ -0,0 +1,93 @@ +"""Unit tests for the file output sink.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from deep_agent.src.triggers.sinks.file import FileSink +from deep_agent.src.triggers.sinks.protocol import TriggerResult +from deep_agent.src.triggers.sources.protocol import TriggerEvent + + +def _make_event(**overrides) -> TriggerEvent: + defaults = { + "name": "test-event", + "payload": {"key": "value"}, + "source": "unit-test", + } + defaults.update(overrides) + return TriggerEvent(**defaults) + + +def _make_result(**overrides) -> TriggerResult: + defaults = { + "event": _make_event(), + "output": {"answer": 42}, + "duration_ms": 100.0, + "success": True, + } + defaults.update(overrides) + return TriggerResult(**defaults) + + +class TestFileSink: + """Test FileSink appends JSONL to a file.""" + + async def test_emit_creates_parent_dirs_and_appends_jsonl(self, tmp_path: Path): + nested = tmp_path / "subdir" / "deep" / "output.jsonl" + sink = FileSink(str(nested)) + + await sink.emit(_make_result()) + await sink.close() + + assert nested.exists() + lines = nested.read_text().strip().splitlines() + assert len(lines) == 1 + + parsed = json.loads(lines[0]) + assert parsed["success"] is True + assert parsed["event"]["name"] == "test-event" + + async def test_multiple_emits_append_multiple_lines(self, tmp_path: Path): + output_file = tmp_path / "output.jsonl" + sink = FileSink(str(output_file)) + + await sink.emit(_make_result(output="first")) + await sink.emit(_make_result(output="second")) + await sink.emit(_make_result(output="third")) + await sink.close() + + lines = output_file.read_text().strip().splitlines() + assert len(lines) == 3 + + outputs = [json.loads(line)["output"] for line in lines] + assert outputs == ["first", "second", "third"] + + async def test_close_closes_file_handle(self, tmp_path: Path): + output_file = tmp_path / "output.jsonl" + sink = FileSink(str(output_file)) + + await sink.emit(_make_result()) + assert sink._handle is not None + + await sink.close() + assert sink._handle is None + + async def test_emit_after_close_reopens_handle(self, tmp_path: Path): + output_file = tmp_path / "output.jsonl" + sink = FileSink(str(output_file)) + + await sink.emit(_make_result(output="before-close")) + await sink.close() + assert sink._handle is None + + await sink.emit(_make_result(output="after-close")) + assert sink._handle is not None + await sink.close() + + lines = output_file.read_text().strip().splitlines() + assert len(lines) == 2 + + outputs = [json.loads(line)["output"] for line in lines] + assert outputs == ["before-close", "after-close"] diff --git a/tests/unit/triggers/test_middleware.py b/tests/unit/triggers/test_middleware.py new file mode 100644 index 00000000..ecbc352a --- /dev/null +++ b/tests/unit/triggers/test_middleware.py @@ -0,0 +1,243 @@ +"""Unit tests for EventTriggerMiddleware.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +from deep_agent.src.triggers.config import ( + CronTriggerConfig, + HeadlessConfig, + OutputSinkConfig, + QueueTriggerConfig, + TriggerConfig, + WebhookTriggerConfig, +) +from deep_agent.src.triggers.middleware import EventTriggerMiddleware +from deep_agent.src.triggers.sinks.protocol import TriggerResult +from deep_agent.src.triggers.sources.protocol import TriggerEvent + + +def _make_event(**overrides) -> TriggerEvent: + defaults = {"name": "test-event", "payload": {"k": "v"}, "source": "unit-test"} + defaults.update(overrides) + return TriggerEvent(**defaults) + + +def _make_result(**overrides) -> TriggerResult: + defaults = { + "event": _make_event(), + "output": "ok", + "duration_ms": 10.0, + "success": True, + } + defaults.update(overrides) + return TriggerResult(**defaults) + + +class TestBuildSources: + """Test _build_sources() returns the correct trigger sources.""" + + def test_returns_empty_list_when_no_triggers_enabled(self): + config = HeadlessConfig() + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sources = mw._build_sources() + + assert sources == [] + + @patch("deep_agent.src.triggers.sources.webhook.WebhookTriggerSource") + def test_returns_webhook_source_when_webhook_enabled(self, mock_cls): + mock_cls.return_value = MagicMock() + config = HeadlessConfig( + triggers=TriggerConfig(webhook=WebhookTriggerConfig(enabled=True)) + ) + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sources = mw._build_sources() + + assert len(sources) == 1 + mock_cls.assert_called_once() + + @patch("deep_agent.src.triggers.sources.cron.CronTriggerSource") + def test_returns_cron_source_when_cron_enabled(self, mock_cls): + mock_cls.return_value = MagicMock() + config = HeadlessConfig( + triggers=TriggerConfig(cron=CronTriggerConfig(enabled=True)) + ) + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sources = mw._build_sources() + + assert len(sources) == 1 + mock_cls.assert_called_once() + + @patch("deep_agent.src.triggers.sources.queue.QueueTriggerSource") + def test_returns_queue_source_when_queue_enabled(self, mock_cls): + mock_cls.return_value = MagicMock() + config = HeadlessConfig( + triggers=TriggerConfig(queue=QueueTriggerConfig(enabled=True)) + ) + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sources = mw._build_sources() + + assert len(sources) == 1 + mock_cls.assert_called_once() + + +class TestBuildSinks: + """Test _build_sinks() returns the correct output sinks.""" + + @patch("deep_agent.src.triggers.sinks.stdout.StdoutSink") + def test_defaults_to_stdout_when_output_sinks_empty(self, mock_cls): + mock_cls.return_value = MagicMock() + config = HeadlessConfig(output_sinks=[]) + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sinks = mw._build_sinks() + + assert len(sinks) == 1 + mock_cls.assert_called_once() + + @patch("deep_agent.src.triggers.sinks.stdout.StdoutSink") + def test_creates_stdout_sink(self, mock_cls): + mock_cls.return_value = MagicMock() + config = HeadlessConfig(output_sinks=[OutputSinkConfig(type="stdout")]) + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sinks = mw._build_sinks() + + assert len(sinks) == 1 + mock_cls.assert_called_once() + + @patch("deep_agent.src.triggers.sinks.file.FileSink") + def test_creates_file_sink(self, mock_cls): + mock_cls.return_value = MagicMock() + config = HeadlessConfig( + output_sinks=[OutputSinkConfig(type="file", path="/tmp/test.jsonl")] + ) + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sinks = mw._build_sinks() + + assert len(sinks) == 1 + mock_cls.assert_called_once_with(path="/tmp/test.jsonl") + + @patch("deep_agent.src.triggers.sinks.webhook.WebhookSink") + def test_creates_webhook_sink(self, mock_cls): + mock_cls.return_value = MagicMock() + config = HeadlessConfig( + output_sinks=[ + OutputSinkConfig( + type="webhook", + url="https://example.com/hook", + headers={"X-Key": "val"}, + ) + ] + ) + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sinks = mw._build_sinks() + + assert len(sinks) == 1 + mock_cls.assert_called_once_with( + url="https://example.com/hook", headers={"X-Key": "val"} + ) + + @patch("deep_agent.src.triggers.sinks.redis.RedisSink") + def test_creates_redis_sink(self, mock_cls): + mock_cls.return_value = MagicMock() + config = HeadlessConfig( + output_sinks=[OutputSinkConfig(type="redis", stream="my-stream")] + ) + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sinks = mw._build_sinks() + + assert len(sinks) == 1 + mock_cls.assert_called_once_with( + stream="my-stream", redis_url="redis://redis:6379/0" + ) + + +class TestStartStop: + """Test start() and stop() lifecycle methods.""" + + async def test_start_starts_all_sources_and_creates_loop_task(self): + config = HeadlessConfig() + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + mock_source = AsyncMock() + mock_sink = MagicMock() + mw._build_sources = MagicMock(return_value=[mock_source]) + mw._build_sinks = MagicMock(return_value=[mock_sink]) + + await mw.start() + + mock_source.start.assert_awaited_once() + assert mw._loop_task is not None + + # Clean up the task + mw._stop_event.set() + mw._loop_task.cancel() + try: + await mw._loop_task + except asyncio.CancelledError: + pass + + async def test_stop_sets_stop_event_and_closes_sinks(self): + config = HeadlessConfig() + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + mock_sink = AsyncMock() + mw._sinks = [mock_sink] + mw._sources = [] + mw._stop_event.clear() + + # Create a task that completes quickly once stop is set + async def _quick_loop(): + await mw._stop_event.wait() + + mw._loop_task = asyncio.create_task(_quick_loop()) + + await mw.stop() + + assert mw._stop_event.is_set() + mock_sink.close.assert_awaited_once() + assert mw._loop_task is None + + +class TestEmitResult: + """Test _emit_result() fans out to all sinks.""" + + async def test_emit_result_fans_out_to_all_sinks(self): + config = HeadlessConfig() + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + sink_a = AsyncMock() + sink_b = AsyncMock() + mw._sinks = [sink_a, sink_b] + + result = _make_result() + await mw._emit_result(result) + + sink_a.emit.assert_awaited_once_with(result) + sink_b.emit.assert_awaited_once_with(result) + + async def test_emit_result_catches_per_sink_errors(self): + config = HeadlessConfig() + mw = EventTriggerMiddleware(config=config, graph=MagicMock()) + + failing_sink = AsyncMock() + failing_sink.emit.side_effect = RuntimeError("sink exploded") + healthy_sink = AsyncMock() + mw._sinks = [failing_sink, healthy_sink] + + result = _make_result() + + with patch("deep_agent.src.triggers.middleware.logger"): + await mw._emit_result(result) # must not raise + + # Healthy sink still receives the result despite the first sink failing + healthy_sink.emit.assert_awaited_once_with(result) diff --git a/tests/unit/triggers/test_queue_source.py b/tests/unit/triggers/test_queue_source.py new file mode 100644 index 00000000..d99a0a70 --- /dev/null +++ b/tests/unit/triggers/test_queue_source.py @@ -0,0 +1,308 @@ +"""Unit tests for queue consumer trigger source.""" + +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest + +from deep_agent.src.triggers.config import QueueTriggerConfig +from deep_agent.src.triggers.sources.queue import ( + QueueMessage, + QueueTriggerSource, + RedisStreamsConsumer, +) + + +class TestQueueMessage: + """Test QueueMessage dataclass fields.""" + + def test_fields(self): + msg = QueueMessage(id="123-0", data={"name": "test", "key": "val"}) + assert msg.id == "123-0" + assert msg.data == {"name": "test", "key": "val"} + + def test_empty_data(self): + msg = QueueMessage(id="0-0", data={}) + assert msg.data == {} + + +class TestRedisStreamsConsumer: + """Test RedisStreamsConsumer with mocked redis.""" + + async def test_creates_consumer_group_on_first_consume(self): + consumer = RedisStreamsConsumer( + stream="test-stream", + consumer_group="test-group", + consumer_name="worker-1", + redis_url="redis://localhost:6379/0", + ) + + mock_client = AsyncMock() + mock_client.xgroup_create = AsyncMock() + + with patch( + "redis.asyncio.from_url", + return_value=mock_client, + ): + await consumer._ensure_client() + + mock_client.xgroup_create.assert_awaited_once_with( + "test-stream", "test-group", id="0", mkstream=True + ) + + async def test_handles_busygroup_error(self): + consumer = RedisStreamsConsumer( + stream="s1", + consumer_group="g1", + consumer_name="w1", + ) + + mock_client = AsyncMock() + mock_client.xgroup_create = AsyncMock( + side_effect=Exception("BUSYGROUP Consumer Group name already exists") + ) + + with patch( + "redis.asyncio.from_url", + return_value=mock_client, + ): + # Should not raise — BUSYGROUP is silently ignored. + await consumer._ensure_client() + + mock_client.xgroup_create.assert_awaited_once() + + async def test_non_busygroup_error_propagates(self): + consumer = RedisStreamsConsumer( + stream="s1", + consumer_group="g1", + consumer_name="w1", + ) + + mock_client = AsyncMock() + mock_client.xgroup_create = AsyncMock( + side_effect=Exception("NOPERM Insufficient permissions") + ) + + with ( + patch( + "redis.asyncio.from_url", + return_value=mock_client, + ), + pytest.raises(Exception, match="NOPERM"), + ): + await consumer._ensure_client() + + async def test_ack_calls_xack(self): + consumer = RedisStreamsConsumer( + stream="my-stream", + consumer_group="my-group", + consumer_name="w1", + ) + + mock_client = AsyncMock() + mock_client.xgroup_create = AsyncMock() + mock_client.xack = AsyncMock() + + with patch( + "redis.asyncio.from_url", + return_value=mock_client, + ): + # Initialize client. + await consumer._ensure_client() + msg = QueueMessage(id="1234-0", data={"key": "val"}) + await consumer.ack(msg) + + mock_client.xack.assert_awaited_once_with("my-stream", "my-group", "1234-0") + + async def test_close_stops_running_and_closes_client(self): + consumer = RedisStreamsConsumer( + stream="s", consumer_group="g", consumer_name="w" + ) + + mock_client = AsyncMock() + mock_client.xgroup_create = AsyncMock() + mock_client.aclose = AsyncMock() + + with patch( + "redis.asyncio.from_url", + return_value=mock_client, + ): + await consumer._ensure_client() + assert consumer._running is True + + await consumer.close() + + assert consumer._running is False + assert consumer._client is None + mock_client.aclose.assert_awaited_once() + + async def test_close_without_client_is_safe(self): + consumer = RedisStreamsConsumer( + stream="s", consumer_group="g", consumer_name="w" + ) + # Should not raise when no client is initialized. + await consumer.close() + assert consumer._running is False + assert consumer._client is None + + +class TestQueueTriggerSource: + """Test QueueTriggerSource lifecycle and event wrapping.""" + + async def test_start_creates_consumer_and_task(self): + config = QueueTriggerConfig( + enabled=True, + backend="redis_streams", + stream="my-tasks", + consumer_group="workers", + consumer_name="w-1", + ) + + mock_consumer_instance = AsyncMock() + + async def _mock_consume(): + return + yield # Make it an async generator. + + mock_consumer_instance.consume = _mock_consume + + with patch( + "deep_agent.src.triggers.sources.queue.RedisStreamsConsumer", + return_value=mock_consumer_instance, + ): + source = QueueTriggerSource(config, redis_url="redis://test:6379/0") + await source.start() + + assert source._consumer is not None + assert source._task is not None + + # Clean up. + await source.stop() + + async def test_unsupported_backend_raises_value_error(self): + config = QueueTriggerConfig( + enabled=True, + backend="rabbitmq", + ) + source = QueueTriggerSource(config) + + with pytest.raises(ValueError, match="(?i)unsupported queue backend: rabbitmq"): + await source.start() + + async def test_wraps_messages_as_trigger_events(self): + config = QueueTriggerConfig( + enabled=True, + backend="redis_streams", + stream="tasks", + ) + + messages = [ + QueueMessage(id="1-0", data={"name": "task-a", "input": "hello"}), + QueueMessage(id="2-0", data={"name": "task-b", "input": "world"}), + ] + + mock_consumer = AsyncMock() + + async def _mock_consume(): + for m in messages: + yield m + + mock_consumer.consume = _mock_consume + mock_consumer.ack = AsyncMock() + + with patch( + "deep_agent.src.triggers.sources.queue.RedisStreamsConsumer", + return_value=mock_consumer, + ): + source = QueueTriggerSource(config) + await source.start() + + # Wait for the consume loop to process the messages. + await asyncio.sleep(0.1) + + assert source._queue.qsize() == 2 + + event_a = source._queue.get_nowait() + assert event_a.name == "task-a" + assert event_a.source == "queue" + assert event_a.payload == {"name": "task-a", "input": "hello"} + assert event_a.metadata["message_id"] == "1-0" + assert event_a.metadata["stream"] == "tasks" + + event_b = source._queue.get_nowait() + assert event_b.name == "task-b" + + # Each message should have been acknowledged. + assert mock_consumer.ack.await_count == 2 + + await source.stop() + + async def test_stop_cancels_task_and_closes_consumer(self): + config = QueueTriggerConfig(enabled=True, backend="redis_streams") + + mock_consumer = AsyncMock() + + async def _mock_consume(): + # Block indefinitely until cancelled. + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + return + yield # noqa: F841 — makes this an async generator + + mock_consumer.consume = _mock_consume + mock_consumer.close = AsyncMock() + + with patch( + "deep_agent.src.triggers.sources.queue.RedisStreamsConsumer", + return_value=mock_consumer, + ): + source = QueueTriggerSource(config) + await source.start() + + assert source._task is not None + assert source._consumer is not None + + await source.stop() + + assert source._task is None + assert source._consumer is None + mock_consumer.close.assert_awaited_once() + + async def test_stop_when_not_started_is_safe(self): + config = QueueTriggerConfig() + source = QueueTriggerSource(config) + # Should not raise. + await source.stop() + assert source._task is None + assert source._consumer is None + + async def test_aiter_returns_self(self): + config = QueueTriggerConfig() + source = QueueTriggerSource(config) + assert source.__aiter__() is source + + async def test_default_event_name_when_missing(self): + config = QueueTriggerConfig(enabled=True, backend="redis_streams") + + # Message data without "name" key should default to "queue-event". + mock_consumer = AsyncMock() + + async def _mock_consume(): + yield QueueMessage(id="99-0", data={"input": "something"}) + + mock_consumer.consume = _mock_consume + mock_consumer.ack = AsyncMock() + + with patch( + "deep_agent.src.triggers.sources.queue.RedisStreamsConsumer", + return_value=mock_consumer, + ): + source = QueueTriggerSource(config) + await source.start() + await asyncio.sleep(0.1) + + event = source._queue.get_nowait() + assert event.name == "queue-event" + + await source.stop() diff --git a/tests/unit/triggers/test_redis_sink.py b/tests/unit/triggers/test_redis_sink.py new file mode 100644 index 00000000..a326b8fd --- /dev/null +++ b/tests/unit/triggers/test_redis_sink.py @@ -0,0 +1,74 @@ +"""Unit tests for the Redis output sink.""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, patch + +from deep_agent.src.triggers.sinks.protocol import TriggerResult +from deep_agent.src.triggers.sinks.redis import RedisSink +from deep_agent.src.triggers.sources.protocol import TriggerEvent + + +def _make_event(**overrides) -> TriggerEvent: + defaults = { + "name": "test-event", + "payload": {"key": "value"}, + "source": "unit-test", + } + defaults.update(overrides) + return TriggerEvent(**defaults) + + +def _make_result(**overrides) -> TriggerResult: + defaults = { + "event": _make_event(), + "output": {"answer": 42}, + "duration_ms": 75.0, + "success": True, + } + defaults.update(overrides) + return TriggerResult(**defaults) + + +class TestRedisSink: + """Test RedisSink publishes to a Redis Stream via XADD.""" + + async def test_emit_calls_xadd_with_serialized_result(self): + sink = RedisSink(stream="test-stream") + mock_client = AsyncMock() + sink._client = mock_client + + result = _make_result() + await sink.emit(result) + + mock_client.xadd.assert_awaited_once() + call_args = mock_client.xadd.call_args + assert call_args[0][0] == "test-stream" + + payload_json = call_args[0][1]["result"] + parsed = json.loads(payload_json) + assert parsed["success"] is True + assert parsed["event"]["name"] == "test-event" + + async def test_close_closes_redis_client(self): + sink = RedisSink(stream="test-stream") + mock_client = AsyncMock() + sink._client = mock_client + + await sink.close() + + mock_client.aclose.assert_awaited_once() + assert sink._client is None + + async def test_redis_error_in_emit_is_caught_and_logged(self): + sink = RedisSink(stream="test-stream") + mock_client = AsyncMock() + mock_client.xadd.side_effect = ConnectionError("Redis down") + sink._client = mock_client + + with patch("deep_agent.src.triggers.sinks.redis.logger") as mock_logger: + await sink.emit(_make_result()) # must not raise + + mock_logger.exception.assert_called_once() + assert "test-stream" in mock_logger.exception.call_args[0][1] diff --git a/tests/unit/triggers/test_runtime.py b/tests/unit/triggers/test_runtime.py new file mode 100644 index 00000000..e4693d69 --- /dev/null +++ b/tests/unit/triggers/test_runtime.py @@ -0,0 +1,45 @@ +"""Unit tests for HeadlessRuntime and HeadlessUser.""" + +from __future__ import annotations + +from deep_agent.src.triggers.runtime import HeadlessRuntime, HeadlessUser + + +class TestHeadlessUser: + """Test HeadlessUser data attributes.""" + + def test_default_identity(self): + user = HeadlessUser() + assert user.identity == "headless-worker" + + def test_custom_identity(self): + user = HeadlessUser(identity="batch-processor-42") + assert user.identity == "batch-processor-42" + + def test_access_token_is_none(self): + user = HeadlessUser() + assert user.access_token is None + + def test_refresh_token_is_none(self): + user = HeadlessUser() + assert user.refresh_token is None + + +class TestHeadlessRuntime: + """Test HeadlessRuntime creates a user with expected attributes.""" + + def test_creates_user_with_default_identity(self): + runtime = HeadlessRuntime() + assert runtime.user.identity == "headless-worker" + + def test_creates_user_with_custom_identity(self): + runtime = HeadlessRuntime(identity="nightly-report-agent") + assert runtime.user.identity == "nightly-report-agent" + + def test_user_has_expected_attributes(self): + runtime = HeadlessRuntime(identity="test-agent") + user = runtime.user + + assert user.identity == "test-agent" + assert user.access_token is None + assert user.refresh_token is None diff --git a/tests/unit/triggers/test_stdout_sink.py b/tests/unit/triggers/test_stdout_sink.py new file mode 100644 index 00000000..5736c0bd --- /dev/null +++ b/tests/unit/triggers/test_stdout_sink.py @@ -0,0 +1,71 @@ +"""Unit tests for the stdout output sink.""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock, patch + +from deep_agent.src.triggers.sinks.protocol import TriggerResult +from deep_agent.src.triggers.sinks.stdout import StdoutSink +from deep_agent.src.triggers.sources.protocol import TriggerEvent + + +def _make_event(**overrides) -> TriggerEvent: + defaults = { + "name": "test-event", + "payload": {"key": "value"}, + "source": "unit-test", + } + defaults.update(overrides) + return TriggerEvent(**defaults) + + +def _make_result(**overrides) -> TriggerResult: + defaults = { + "event": _make_event(), + "output": {"answer": 42}, + "duration_ms": 123.4, + "success": True, + } + defaults.update(overrides) + return TriggerResult(**defaults) + + +class TestStdoutSink: + """Test StdoutSink writes JSON to stdout.""" + + async def test_emit_writes_json_to_stdout(self): + sink = StdoutSink() + result = _make_result() + + mock_stdout = MagicMock() + with patch("deep_agent.src.triggers.sinks.stdout.sys.stdout", mock_stdout): + await sink.emit(result) + + mock_stdout.write.assert_called_once() + mock_stdout.flush.assert_called_once() + + written = mock_stdout.write.call_args[0][0] + assert written.endswith("\n") + + async def test_emit_output_is_valid_json_with_event_data(self): + sink = StdoutSink() + result = _make_result() + + mock_stdout = MagicMock() + with patch("deep_agent.src.triggers.sinks.stdout.sys.stdout", mock_stdout): + await sink.emit(result) + + written = mock_stdout.write.call_args[0][0] + line = written.rstrip("\n") + parsed = json.loads(line) + + assert parsed["success"] is True + assert parsed["duration_ms"] == 123.4 + assert parsed["event"]["name"] == "test-event" + assert parsed["event"]["payload"] == {"key": "value"} + assert parsed["output"] == {"answer": 42} + + async def test_close_is_noop(self): + sink = StdoutSink() + await sink.close() # must not raise diff --git a/tests/unit/triggers/test_task_store.py b/tests/unit/triggers/test_task_store.py new file mode 100644 index 00000000..468acfe3 --- /dev/null +++ b/tests/unit/triggers/test_task_store.py @@ -0,0 +1,223 @@ +"""Unit tests for the Redis-backed task status store.""" + +from unittest.mock import AsyncMock, patch + +from deep_agent.src.triggers.task_store import TaskRecord, TaskStore + + +class TestTaskRecord: + def test_to_json_and_from_json_roundtrip(self): + record = TaskRecord( + task_id="abc123", + task_name="test-task", + status="queued", + payload={"key": "value"}, + created_at="2026-06-23T00:00:00Z", + updated_at="2026-06-23T00:00:00Z", + ) + raw = record.to_json() + restored = TaskRecord.from_json(raw) + assert restored.task_id == "abc123" + assert restored.task_name == "test-task" + assert restored.status == "queued" + assert restored.payload == {"key": "value"} + assert restored.delivered is False + + def test_default_fields(self): + record = TaskRecord(task_id="x", task_name="y", status="queued") + assert record.payload == {} + assert record.result is None + assert record.error is None + assert record.thread_id is None + assert record.user_id is None + assert record.delivered is False + + +class TestTaskStoreCreate: + async def test_create_task_stores_in_redis(self): + mock_client = AsyncMock() + mock_client.set = AsyncMock() + mock_client.zadd = AsyncMock() + mock_client.expire = AsyncMock() + + store = TaskStore() + with patch("redis.asyncio.from_url", return_value=mock_client): + record = await store.create_task( + task_name="my-task", + payload={"data": 1}, + user_id="user-1", + thread_id="thread-1", + ) + + assert record.task_name == "my-task" + assert record.status == "queued" + assert record.user_id == "user-1" + assert record.thread_id == "thread-1" + assert len(record.task_id) == 12 + mock_client.set.assert_awaited_once() + mock_client.zadd.assert_awaited_once() + + async def test_create_task_without_user_skips_index(self): + mock_client = AsyncMock() + mock_client.set = AsyncMock() + mock_client.zadd = AsyncMock() + + store = TaskStore() + with patch("redis.asyncio.from_url", return_value=mock_client): + await store.create_task(task_name="anon-task", payload={}) + + mock_client.set.assert_awaited_once() + mock_client.zadd.assert_not_awaited() + + +class TestTaskStoreUpdate: + async def test_update_status_modifies_record(self): + original = TaskRecord( + task_id="t1", + task_name="job", + status="queued", + created_at="2026-01-01", + updated_at="2026-01-01", + ) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=original.to_json()) + mock_client.ttl = AsyncMock(return_value=80000) + mock_client.set = AsyncMock() + + store = TaskStore() + with patch("redis.asyncio.from_url", return_value=mock_client): + await store.update_status("t1", "processing") + + saved_json = mock_client.set.call_args[0][1] + saved = TaskRecord.from_json(saved_json) + assert saved.status == "processing" + + async def test_update_status_with_result(self): + original = TaskRecord( + task_id="t2", + task_name="job", + status="processing", + created_at="2026-01-01", + updated_at="2026-01-01", + ) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=original.to_json()) + mock_client.ttl = AsyncMock(return_value=80000) + mock_client.set = AsyncMock() + + store = TaskStore() + with patch("redis.asyncio.from_url", return_value=mock_client): + await store.update_status("t2", "completed", result={"answer": 42}) + + saved = TaskRecord.from_json(mock_client.set.call_args[0][1]) + assert saved.status == "completed" + assert saved.result == {"answer": 42} + + async def test_update_nonexistent_task_is_noop(self): + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=None) + + store = TaskStore() + with patch("redis.asyncio.from_url", return_value=mock_client): + await store.update_status("missing", "processing") + + mock_client.set.assert_not_awaited() + + +class TestTaskStoreQuery: + async def test_get_task_returns_record(self): + record = TaskRecord( + task_id="t1", + task_name="job", + status="completed", + created_at="2026-01-01", + updated_at="2026-01-01", + ) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=record.to_json()) + + store = TaskStore() + with patch("redis.asyncio.from_url", return_value=mock_client): + result = await store.get_task("t1") + + assert result is not None + assert result.task_id == "t1" + assert result.status == "completed" + + async def test_get_task_returns_none_when_missing(self): + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=None) + + store = TaskStore() + with patch("redis.asyncio.from_url", return_value=mock_client): + result = await store.get_task("nonexistent") + + assert result is None + + async def test_get_pending_results_returns_undelivered(self): + completed = TaskRecord( + task_id="t1", + task_name="done-job", + status="completed", + result="output", + delivered=False, + created_at="2026-01-01", + updated_at="2026-01-01", + ) + delivered = TaskRecord( + task_id="t2", + task_name="old-job", + status="completed", + result="old", + delivered=True, + created_at="2026-01-01", + updated_at="2026-01-01", + ) + queued = TaskRecord( + task_id="t3", + task_name="pending", + status="queued", + created_at="2026-01-01", + updated_at="2026-01-01", + ) + + mock_client = AsyncMock() + mock_client.zrange = AsyncMock(return_value=["t1", "t2", "t3"]) + + def _get_side_effect(key): + mapping = { + "task:t1": completed.to_json(), + "task:t2": delivered.to_json(), + "task:t3": queued.to_json(), + } + return mapping.get(key) + + mock_client.get = AsyncMock(side_effect=_get_side_effect) + + store = TaskStore() + with patch("redis.asyncio.from_url", return_value=mock_client): + pending = await store.get_pending_results("user-1") + + assert len(pending) == 1 + assert pending[0].task_id == "t1" + + async def test_mark_delivered_sets_flag(self): + record = TaskRecord( + task_id="t1", + task_name="job", + status="completed", + delivered=False, + created_at="2026-01-01", + updated_at="2026-01-01", + ) + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=record.to_json()) + mock_client.ttl = AsyncMock(return_value=80000) + mock_client.set = AsyncMock() + + store = TaskStore() + with patch("redis.asyncio.from_url", return_value=mock_client): + await store.mark_delivered("t1") + + saved = TaskRecord.from_json(mock_client.set.call_args[0][1]) + assert saved.delivered is True diff --git a/tests/unit/triggers/test_tools.py b/tests/unit/triggers/test_tools.py new file mode 100644 index 00000000..819f3fc1 --- /dev/null +++ b/tests/unit/triggers/test_tools.py @@ -0,0 +1,166 @@ +"""Unit tests for headless worker tools.""" + +from unittest.mock import AsyncMock, patch + +from deep_agent.src.triggers.task_store import TaskRecord +from deep_agent.src.triggers.tools import ( + check_task_status, + get_builtin_tools, + get_pending_results, + queue_task, +) + + +class TestQueueTask: + async def test_creates_task_record_and_pushes_to_stream(self): + mock_record = TaskRecord( + task_id="abc123", + task_name="test", + status="queued", + created_at="2026-01-01", + updated_at="2026-01-01", + ) + mock_redis = AsyncMock() + mock_redis.xadd = AsyncMock(return_value="1234-0") + mock_redis.aclose = AsyncMock() + + with ( + patch("deep_agent.src.triggers.tools._store") as mock_store, + patch("redis.asyncio.from_url", return_value=mock_redis), + ): + mock_store.create_task = AsyncMock(return_value=mock_record) + result = await queue_task( + task_name="test", + payload={"key": "val"}, + thread_id="thread-1", + user_id="user-1", + ) + + assert "abc123" in result + assert "queued" in result.lower() + mock_store.create_task.assert_awaited_once() + mock_redis.xadd.assert_awaited_once() + + async def test_returns_task_id_in_response(self): + mock_record = TaskRecord( + task_id="xyz789", + task_name="report", + status="queued", + created_at="2026-01-01", + updated_at="2026-01-01", + ) + mock_redis = AsyncMock() + mock_redis.xadd = AsyncMock(return_value="5678-0") + mock_redis.aclose = AsyncMock() + + with ( + patch("deep_agent.src.triggers.tools._store") as mock_store, + patch("redis.asyncio.from_url", return_value=mock_redis), + ): + mock_store.create_task = AsyncMock(return_value=mock_record) + result = await queue_task(task_name="report", payload={}) + + assert "xyz789" in result + + +class TestCheckTaskStatus: + async def test_completed_task_returns_result(self): + record = TaskRecord( + task_id="t1", + task_name="report", + status="completed", + result={"answer": 42}, + created_at="2026-01-01", + updated_at="2026-01-01", + ) + with patch("deep_agent.src.triggers.tools._store") as mock_store: + mock_store.get_task = AsyncMock(return_value=record) + result = await check_task_status("t1") + + assert "COMPLETED" in result + assert "42" in result + + async def test_failed_task_returns_error(self): + record = TaskRecord( + task_id="t2", + task_name="export", + status="failed", + error="Connection timeout", + created_at="2026-01-01", + updated_at="2026-01-01", + ) + with patch("deep_agent.src.triggers.tools._store") as mock_store: + mock_store.get_task = AsyncMock(return_value=record) + result = await check_task_status("t2") + + assert "FAILED" in result + assert "Connection timeout" in result + + async def test_queued_task_returns_status(self): + record = TaskRecord( + task_id="t3", + task_name="batch", + status="queued", + created_at="2026-01-01", + updated_at="2026-01-01", + ) + with patch("deep_agent.src.triggers.tools._store") as mock_store: + mock_store.get_task = AsyncMock(return_value=record) + result = await check_task_status("t3") + + assert "QUEUED" in result + + async def test_missing_task_returns_not_found(self): + with patch("deep_agent.src.triggers.tools._store") as mock_store: + mock_store.get_task = AsyncMock(return_value=None) + result = await check_task_status("nonexistent") + + assert "not found" in result.lower() + + +class TestGetPendingResults: + async def test_returns_completed_undelivered_tasks(self): + records = [ + TaskRecord( + task_id="t1", + task_name="report", + status="completed", + result="Report data", + delivered=False, + created_at="2026-01-01", + updated_at="2026-01-01", + ), + ] + with patch("deep_agent.src.triggers.tools._store") as mock_store: + mock_store.get_pending_results = AsyncMock(return_value=records) + mock_store.mark_delivered = AsyncMock() + result = await get_pending_results("user-1") + + assert "1 background task(s)" in result + assert "report" in result + assert "COMPLETED" in result + mock_store.mark_delivered.assert_awaited_once_with("t1") + + async def test_returns_no_pending_message(self): + with patch("deep_agent.src.triggers.tools._store") as mock_store: + mock_store.get_pending_results = AsyncMock(return_value=[]) + result = await get_pending_results("user-1") + + assert "No pending" in result + + +class TestGetBuiltinTools: + def test_returns_three_tools(self): + tools = get_builtin_tools() + assert len(tools) == 3 + + def test_tool_names(self): + tools = get_builtin_tools() + names = {t.name for t in tools} + assert names == {"queue_task", "check_task_status", "get_pending_results"} + + def test_tools_are_callable(self): + tools = get_builtin_tools() + for tool in tools: + assert hasattr(tool, "name") + assert hasattr(tool, "description") diff --git a/tests/unit/triggers/test_webhook_sink.py b/tests/unit/triggers/test_webhook_sink.py new file mode 100644 index 00000000..3f62ee19 --- /dev/null +++ b/tests/unit/triggers/test_webhook_sink.py @@ -0,0 +1,103 @@ +"""Unit tests for the webhook output sink.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx + +from deep_agent.src.triggers.sinks.protocol import TriggerResult +from deep_agent.src.triggers.sinks.webhook import WebhookSink +from deep_agent.src.triggers.sources.protocol import TriggerEvent + + +def _make_event(**overrides) -> TriggerEvent: + defaults = { + "name": "test-event", + "payload": {"key": "value"}, + "source": "unit-test", + } + defaults.update(overrides) + return TriggerEvent(**defaults) + + +def _make_result(**overrides) -> TriggerResult: + defaults = { + "event": _make_event(), + "output": {"answer": 42}, + "duration_ms": 50.0, + "success": True, + } + defaults.update(overrides) + return TriggerResult(**defaults) + + +def _mock_response(status_code: int) -> MagicMock: + resp = MagicMock() + resp.status_code = status_code + return resp + + +class TestWebhookSink: + """Test WebhookSink POSTs JSON to a URL with retry logic.""" + + async def test_emit_posts_json_to_url(self): + sink = WebhookSink(url="https://example.com/hook") + mock_client = AsyncMock(spec=httpx.AsyncClient) + mock_client.post.return_value = _mock_response(200) + sink._client = mock_client + + await sink.emit(_make_result()) + + mock_client.post.assert_called_once() + call_kwargs = mock_client.post.call_args + assert call_kwargs[0][0] == "https://example.com/hook" + assert "Content-Type" in call_kwargs[1]["headers"] + assert call_kwargs[1]["headers"]["Content-Type"] == "application/json" + + async def test_custom_headers_included_in_request(self): + custom_headers = {"X-Api-Key": "secret-123", "X-Custom": "value"} + sink = WebhookSink(url="https://example.com/hook", headers=custom_headers) + mock_client = AsyncMock(spec=httpx.AsyncClient) + mock_client.post.return_value = _mock_response(200) + sink._client = mock_client + + await sink.emit(_make_result()) + + sent_headers = mock_client.post.call_args[1]["headers"] + assert sent_headers["X-Api-Key"] == "secret-123" + assert sent_headers["X-Custom"] == "value" + assert sent_headers["Content-Type"] == "application/json" + + @patch("asyncio.sleep", new_callable=AsyncMock) + async def test_5xx_triggers_retry(self, mock_sleep: AsyncMock): + sink = WebhookSink(url="https://example.com/hook", max_retries=2) + mock_client = AsyncMock(spec=httpx.AsyncClient) + mock_client.post.return_value = _mock_response(503) + sink._client = mock_client + + await sink.emit(_make_result()) + + # 1 initial + 2 retries = 3 total calls + assert mock_client.post.call_count == 3 + + @patch("asyncio.sleep", new_callable=AsyncMock) + async def test_4xx_does_not_retry(self, mock_sleep: AsyncMock): + sink = WebhookSink(url="https://example.com/hook", max_retries=3) + mock_client = AsyncMock(spec=httpx.AsyncClient) + mock_client.post.return_value = _mock_response(422) + sink._client = mock_client + + await sink.emit(_make_result()) + + mock_client.post.assert_called_once() + + async def test_close_closes_httpx_client(self): + sink = WebhookSink(url="https://example.com/hook") + mock_client = AsyncMock(spec=httpx.AsyncClient) + sink._client = mock_client + + await sink.close() + + mock_client.aclose.assert_awaited_once() + assert sink._client is None diff --git a/tests/unit/triggers/test_webhook_source.py b/tests/unit/triggers/test_webhook_source.py new file mode 100644 index 00000000..a45c9bc6 --- /dev/null +++ b/tests/unit/triggers/test_webhook_source.py @@ -0,0 +1,259 @@ +"""Unit tests for WebhookTriggerSource.""" + +import asyncio +import json + +from deep_agent.src.triggers.config import WebhookTriggerConfig +from deep_agent.src.triggers.sources.webhook import WebhookTriggerSource + + +async def _send_raw_http( + host: str, + port: int, + method: str, + path: str, + body: bytes | None = None, + headers: dict[str, str] | None = None, +) -> tuple[int, str]: + """Send a raw HTTP request and return (status_code, response_body).""" + reader, writer = await asyncio.open_connection(host, port) + try: + request_line = f"{method} {path} HTTP/1.1\r\n" + header_lines = f"Host: {host}:{port}\r\n" + if headers: + for k, v in headers.items(): + header_lines += f"{k}: {v}\r\n" + if body is not None: + header_lines += f"Content-Length: {len(body)}\r\n" + header_lines += "\r\n" + + writer.write(request_line.encode() + header_lines.encode()) + if body is not None: + writer.write(body) + await writer.drain() + + # Read response status line. + status_line = await reader.readline() + parts = status_line.decode().strip().split(" ", 2) + status_code = int(parts[1]) + + # Read headers until blank line. + while True: + line = await reader.readline() + if line in (b"\r\n", b"\n", b""): + break + + # Read remaining response body. + response_body = await reader.read(4096) + return status_code, response_body.decode() + finally: + writer.close() + await writer.wait_closed() + + +class TestWebhookTriggerSource: + """Test the async HTTP webhook listener.""" + + async def _start_source( + self, + path: str = "/trigger", + ) -> tuple[WebhookTriggerSource, int]: + """Create and start a source on a random port, return (source, port).""" + config = WebhookTriggerConfig( + enabled=True, + host="127.0.0.1", + port=0, # OS assigns a free port. + path=path, + ) + source = WebhookTriggerSource(config) + await source.start() + # Extract the actual bound port. + assert source._server is not None + port = source._server.sockets[0].getsockname()[1] + return source, port + + async def test_start_binds_server(self): + source, port = await self._start_source() + try: + assert source._server is not None + assert port > 0 + finally: + await source.stop() + + async def test_stop_closes_server(self): + source, _ = await self._start_source() + await source.stop() + assert source._server is None + + async def test_stop_when_not_started_is_noop(self): + config = WebhookTriggerConfig() + source = WebhookTriggerSource(config) + # Should not raise. + await source.stop() + assert source._server is None + + async def test_valid_post_enqueues_trigger_event(self): + source, port = await self._start_source() + try: + payload = {"event": "test-event", "key": "value"} + status, body = await _send_raw_http( + "127.0.0.1", + port, + "POST", + "/trigger", + body=json.dumps(payload).encode(), + ) + + assert status == 200 + response = json.loads(body) + assert response["status"] == "accepted" + + # Event should be in the queue. + event = source._queue.get_nowait() + assert event.name == "test-event" + assert event.payload == {"key": "value"} + assert event.source == "webhook" + finally: + await source.stop() + + async def test_invalid_json_returns_400(self): + source, port = await self._start_source() + try: + status, body = await _send_raw_http( + "127.0.0.1", + port, + "POST", + "/trigger", + body=b"not valid json{{{", + ) + assert status == 400 + response = json.loads(body) + assert "invalid JSON" in response["error"] + assert source._queue.empty() + finally: + await source.stop() + + async def test_wrong_path_returns_404(self): + source, port = await self._start_source() + try: + status, _ = await _send_raw_http( + "127.0.0.1", + port, + "POST", + "/wrong-path", + body=json.dumps({"event": "x"}).encode(), + ) + assert status == 404 + assert source._queue.empty() + finally: + await source.stop() + + async def test_get_method_returns_405(self): + source, port = await self._start_source() + try: + status, body = await _send_raw_http( + "127.0.0.1", + port, + "GET", + "/trigger", + ) + assert status == 405 + response = json.loads(body) + assert "not allowed" in response["error"] + finally: + await source.stop() + + async def test_put_method_returns_405(self): + source, port = await self._start_source() + try: + status, _ = await _send_raw_http( + "127.0.0.1", + port, + "PUT", + "/trigger", + body=json.dumps({"event": "x"}).encode(), + ) + assert status == 405 + finally: + await source.stop() + + async def test_trigger_event_source_is_webhook(self): + source, port = await self._start_source() + try: + await _send_raw_http( + "127.0.0.1", + port, + "POST", + "/trigger", + body=json.dumps({"data": 1}).encode(), + ) + event = source._queue.get_nowait() + assert event.source == "webhook" + finally: + await source.stop() + + async def test_event_name_defaults_to_webhook(self): + source, port = await self._start_source() + try: + # Payload without an "event" key defaults name to "webhook". + await _send_raw_http( + "127.0.0.1", + port, + "POST", + "/trigger", + body=json.dumps({"key": "val"}).encode(), + ) + event = source._queue.get_nowait() + assert event.name == "webhook" + finally: + await source.stop() + + async def test_multiple_events_can_be_queued(self): + source, port = await self._start_source() + try: + for i in range(3): + await _send_raw_http( + "127.0.0.1", + port, + "POST", + "/trigger", + body=json.dumps({"event": f"e{i}"}).encode(), + ) + + assert source._queue.qsize() == 3 + names = [source._queue.get_nowait().name for _ in range(3)] + assert names == ["e0", "e1", "e2"] + finally: + await source.stop() + + async def test_custom_path(self): + source, port = await self._start_source(path="/custom/webhook") + try: + status, _ = await _send_raw_http( + "127.0.0.1", + port, + "POST", + "/custom/webhook", + body=json.dumps({"event": "custom"}).encode(), + ) + assert status == 200 + event = source._queue.get_nowait() + assert event.name == "custom" + finally: + await source.stop() + + async def test_non_object_json_body_returns_400(self): + source, port = await self._start_source() + try: + status, body = await _send_raw_http( + "127.0.0.1", + port, + "POST", + "/trigger", + body=json.dumps([1, 2, 3]).encode(), + ) + assert status == 400 + response = json.loads(body) + assert "JSON object" in response["error"] + finally: + await source.stop() diff --git a/tests/unit/utils/test_google_creds.py b/tests/unit/utils/test_google_creds.py new file mode 100644 index 00000000..67f037c0 --- /dev/null +++ b/tests/unit/utils/test_google_creds.py @@ -0,0 +1,161 @@ +"""Unit tests for Google credentials management.""" + +import json +from unittest.mock import MagicMock, patch + +import pytest +from google.oauth2 import service_account + +from deep_agent.utils.google_creds import ( + clear_credentials_cache, + get_service_account_credentials, +) + + +@pytest.fixture(autouse=True) +def clear_cache(): + """Clear credentials cache before and after each test.""" + clear_credentials_cache() + yield + clear_credentials_cache() + + +@pytest.fixture +def mock_service_account_info(): + """Fixture providing valid service account JSON.""" + return { + "type": "service_account", + "project_id": "test-project-123", + "private_key_id": "key123", + "private_key": "-----BEGIN PRIVATE KEY-----\nMOCK_KEY\n-----END PRIVATE KEY-----", + "client_email": "test@test-project-123.iam.gserviceaccount.com", + "client_id": "123456789", + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + } + + +class TestGetServiceAccountCredentials: + """Tests for get_service_account_credentials function.""" + + def test_successful_credential_loading(self, mock_service_account_info): + """Test successful loading of credentials from valid JSON.""" + mock_creds = MagicMock(spec=service_account.Credentials) + + with patch("deep_agent.utils.google_creds.settings") as mock_settings: + mock_settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT = json.dumps( + mock_service_account_info + ) + mock_settings.PYTHON_LOG_LEVEL = "INFO" + + with patch( + "deep_agent.utils.google_creds.service_account.Credentials.from_service_account_info", + return_value=mock_creds, + ) as mock_from_info: + credentials, project = get_service_account_credentials() + + assert credentials == mock_creds + assert project == "test-project-123" + + # Verify the service account info was parsed correctly + mock_from_info.assert_called_once() + call_args = mock_from_info.call_args + assert call_args[0][0] == mock_service_account_info + + def test_credentials_caching(self, mock_service_account_info): + """Test that credentials are cached after first call.""" + mock_creds = MagicMock(spec=service_account.Credentials) + + with patch("deep_agent.utils.google_creds.settings") as mock_settings: + mock_settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT = json.dumps( + mock_service_account_info + ) + mock_settings.PYTHON_LOG_LEVEL = "INFO" + + with patch( + "deep_agent.utils.google_creds.service_account.Credentials.from_service_account_info", + return_value=mock_creds, + ) as mock_from_info: + # First call + creds1, project1 = get_service_account_credentials() + + # Second call + creds2, project2 = get_service_account_credentials() + + # Should be the same instances + assert creds1 is creds2 + assert project1 == project2 + + # Should only create credentials once + assert mock_from_info.call_count == 1 + + @pytest.mark.parametrize("creds_content", [None, ""]) + def test_missing_or_empty_credentials(self, creds_content): + """Test error when credentials are None or empty.""" + with patch("deep_agent.utils.google_creds.settings") as mock_settings: + mock_settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT = creds_content + mock_settings.PYTHON_LOG_LEVEL = "INFO" + + with pytest.raises( + RuntimeError, match="No Google service account credentials configured" + ): + get_service_account_credentials() + + def test_invalid_json(self): + """Test error when credentials content is not valid JSON.""" + with patch("deep_agent.utils.google_creds.settings") as mock_settings: + mock_settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT = "not valid json {" + mock_settings.PYTHON_LOG_LEVEL = "INFO" + + with pytest.raises(RuntimeError, match="Invalid JSON in credentials"): + get_service_account_credentials() + + @pytest.mark.parametrize("action", ["remove", "empty"]) + def test_invalid_project_id(self, mock_service_account_info, action): + """Test error when project_id is missing or empty.""" + if action == "remove": + mock_service_account_info.pop("project_id") + else: + mock_service_account_info["project_id"] = "" + + with patch("deep_agent.utils.google_creds.settings") as mock_settings: + mock_settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT = json.dumps( + mock_service_account_info + ) + mock_settings.PYTHON_LOG_LEVEL = "INFO" + + with pytest.raises( + RuntimeError, + match="Service account JSON does not contain 'project_id' field", + ): + get_service_account_credentials() + + def test_clear_cache_allows_reload(self, mock_service_account_info): + """Test that clearing cache allows credentials to be reloaded.""" + mock_creds1 = MagicMock(spec=service_account.Credentials) + mock_creds2 = MagicMock(spec=service_account.Credentials) + + with patch("deep_agent.utils.google_creds.settings") as mock_settings: + mock_settings.GOOGLE_APPLICATION_CREDENTIALS_CONTENT = json.dumps( + mock_service_account_info + ) + mock_settings.PYTHON_LOG_LEVEL = "INFO" + + with patch( + "deep_agent.utils.google_creds.service_account.Credentials.from_service_account_info", + side_effect=[mock_creds1, mock_creds2], + ) as mock_from_info: + # First load + creds1, _ = get_service_account_credentials() + assert creds1 is mock_creds1 + + # Clear cache + clear_credentials_cache() + + # Second load should create new credentials + creds2, _ = get_service_account_credentials() + assert creds2 is mock_creds2 + assert creds2 is not creds1 + + # Should have been called twice + assert mock_from_info.call_count == 2